[BUG] DeepSpeed checkpointing partitioning modifies input tensor shape which cannot be reused

Describe the bug
Currently, when I enable checkpointing partitioning in stage 3, I find that DeepSpeed's checkpoint() function would modify the input tensor shape when wrapping the function call. Making them unable to be re-used in other forward() calls and throwing tensor shape errors. Stage 3 works fine if I do not enable checkpointing partitioning.

To Reproduce
Run with a conventional network doing for-loop in the forward function like this

from deepspeed.runtime.activation_checkpointing import checkpointing

def forward(self, x, condition):
    for layer in self.residual_layers:
        if checkpointing.is_configured():
            x, skip_connection = checkpointing.checkpoint(layer, x, condition)
        else:
            x, skip_connection = layer(x, condition)
        skip = skip_connection if skip is None else skip_connection + skip

The shape of x and condition would be modified and cannot be reused in the next iteration of the for-loop (while the first will succeed).

Details
I did a bunch of debugging and find that the root cause is that during partitioning, the new input tensors are first flattened here
https://github.com/microsoft/DeepSpeed/blob/25e500e8dd8bfb2fe07fe6e7eb6261bd58e6bd71/deepspeed/runtime/activation_checkpointing/checkpointing.py#L380
And then the original arguments are overridden by their own flattened version here
https://github.com/microsoft/DeepSpeed/blob/25e500e8dd8bfb2fe07fe6e7eb6261bd58e6bd71/deepspeed/runtime/activation_checkpointing/checkpointing.py#L435
I realize a simple fix (although, I have to first apply another fix mentioned in https://github.com/microsoft/DeepSpeed/issues/3779 to get the correct partitioning) may be to replace the above line with arg.data.copy_(inp.view(arg.shape).data). However, it would not work in a multi-GPU context since inp.data only contains one part of the original arg.data.

I think this is a crucial issue if we really want checkpointing partitioning to work in DeepSpeed. Without it, I don't see any benefit in replacing the native Pytorch checkpointing with this non-partitioned one.

### Tasks

Read deepspeed/runtime/activation_checkpointing/checkpointing.py around lines 380 and 435, then reproduce the issue with the provided looping forward method, stage 3, and checkpointing partitioning enabled. Trace how flattened inputs replace the original tensors and verify a fix preserves reusable tensor shapes in both single- and multi-GPU runs.