deep copy state dict for checkpoint

Summary:
deep copy the state dict for sending checkpoint because if the replica moves to the next step, the state dict can change before the checkpoint is sent

Author: tushar00jain

torchft/checkpointing/http_transport.py

@@ -16,6 +16,7 @@
 from typing import Generator, List, Optional, TypeVar, cast
 
 import torch
+from torch.distributed.tensor import DTensor, distribute_tensor
 from torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten
 
 from torchft.checkpointing._rwlock import RWLock
@@ -278,7 +279,13 @@ def _to_cpu(values: List[T], pin_memory: bool) -> List[T]:
                 else:
                     out.append(v.cpu())
             else:
-                out.append(v)
+                if isinstance(v, DTensor):
+                    clone = distribute_tensor(
+                        v.to_local().clone(), v.device_mesh, v.placements
+                    )
+                else:
+                    clone = v.clone()
+                out.append(clone)
         else:
             out.append(v)
     return out