[BUG] Activation Offloading with Residual-Type Connections

Bug Description
When using deepspeed with pytorch-geometric, I get a shape error after offloading the model inputs (e.g., x) to the CPU. What seems to be happening is that the model inputs are used, sent to the CPU for offloading, and then used again in the forward pass. When they are accessed the second time in the forward pass, it tries to use the inputs, but they have been flattened and the shape information is not easily accessible anywhere (at least I could not find it). This shape information is available somewhere since it is used once the inputs are re-loaded for backprop, but not easily in the forward pass. There should be some way for deepseed to handle offloading some activations and reloading them later in the forward pass. Another example of where this would be useful is in residual connections (which is why the bug is titled so).

To Reproduce
Many possible ways to reproduce with an activation that is offloaded and needs to be used in the forward again. Can share my code if necessary, but it is a large project I am working on so is not easy to navigate.

ds_report output

[2023-10-05 22:39:56,885] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
async_io ............... [YES] ...... [OKAY]
fused_adam ............. [YES] ...... [OKAY]
cpu_adam ............... [YES] ...... [OKAY]
cpu_adagrad ............ [YES] ...... [OKAY]
fused_lamb ............. [YES] ...... [OKAY]
quantizer .............. [YES] ...... [OKAY]
random_ltd ............. [YES] ...... [OKAY]
 [WARNING]  using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [YES] ...... [OKAY]
transformer ............ [YES] ...... [OKAY]
stochastic_transformer . [YES] ...... [OKAY]
transformer_inference .. [YES] ...... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/metabob/nfsdata/workspace/venv/lib/python3.8/site-packages/torch']
torch version .................... 1.13.1+cu117
deepspeed install path ........... ['/home/metabob/nfsdata/workspace/venv/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.10.3, unknown, unknown
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.7
deepspeed wheel compiled w. ...... torch 1.13, cuda 11.7
shared memory (/dev/shm) size .... 55.02 GB

System info

• OS: Ubuntu 20.04
• GPU: 1x v100
• Python 3.8.10

Launcher context
Using pytorch lightning to launch deepspeed

No repository file, entry point, or test is named. Start by constructing a minimal activation-offloading case that reuses an input during the forward pass, using the reported PyTorch 1.13.1 and DeepSpeed 0.10.3 environment; done means the case handles the reused activation without a shape error.