[DINO] [5/n] Support FSDP

cosmos_rl/policy/model/dino_cradio_v3/__init__.py

@@ -172,7 +172,7 @@ def load_hf_weights(
             )
             for name in ckpt.keys():
                 ckpt_tensor = ckpt.get_tensor(name)
-                dest_name, tensor = convert_weight_from_hf(
+                dest_name, sharded_tensor = convert_weight_from_hf(
                     ckpt_tensor, name, parallel_dims
                 )
                 if dest_name not in backbone_state_dict:
@@ -181,11 +181,16 @@ def load_hf_weights(
                     )
                     continue
                 target_tensor = backbone_state_dict[dest_name]
+                local_view = (
+                    target_tensor.to_local()
+                    if isinstance(target_tensor, torch.distributed.tensor.DTensor)
+                    else target_tensor
+                )
                 assert (
-                    target_tensor.shape == tensor.shape
-                ), f"Shape mismatch: {target_tensor.shape} != {tensor.shape} for {dest_name}"
+                    local_view.shape == sharded_tensor.shape
+                ), f"Shape mismatch: {local_view.shape} != {sharded_tensor.shape} for {dest_name}"
                 with torch.no_grad():
-                    target_tensor.copy_(tensor)
+                    local_view.copy_(sharded_tensor)
                 used_checkpoint_names.add(dest_name)
 
         for name, parameter in backbone.named_parameters():

cosmos_rl/policy/model/dino_cradio_v3/model/model_utils.py

@@ -249,6 +249,26 @@ def gen_sineembed_for_position(pos_tensor):
     return pos
 
 
+def _get_shard_indices(dtensor: torch.distributed.tensor.DTensor):
+    mesh = dtensor.device_mesh
+    placements = dtensor.placements
+    global_shape = dtensor.shape
+    coord = mesh.get_coordinate()
+
+    indices = []
+    for i, p in enumerate(placements):
+        if isinstance(p, torch.distributed.tensor.Shard):
+            dim = p.dim
+            n_chunks = mesh.size(i)
+            chunk_size = (global_shape[dim] + n_chunks - 1) // n_chunks
+            start = coord[i] * chunk_size
+            end = min(start + chunk_size, global_shape[dim])
+            indices.append(slice(start, end))
+        else:
+            indices.append(slice(None))
+    return tuple(indices)
+
+
 class LinearWithCustomInit(nn.Linear):
     """
     Copy of nn.Linear, with custom initialization for weight and bias.
@@ -280,7 +300,12 @@ def reset_parameters(self):
         assert (self.bias_value is not None) != (self.bias_compute_fn is not None)
         if self.bias_compute_fn is not None:
             with torch.no_grad():
-                self.bias.copy_(self.bias_compute_fn().to(self.bias.device))
+                if isinstance(self.bias, torch.distributed.tensor.DTensor):
+                    idx = _get_shard_indices(self.bias)
+                    local = self.bias.to_local()
+                    local.copy_(self.bias_compute_fn()[idx].to(self.bias.device))
+                else:
+                    self.bias.copy_(self.bias_compute_fn().to(self.bias.device))
         else:
             nn.init.constant_(self.bias, self.bias_value)
 

cosmos_rl/policy/model/dino_cradio_v3/parallelize.py

@@ -3,6 +3,7 @@
 import torch.nn as nn
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed._composable.replicate import replicate
+from torch.distributed.fsdp import fully_shard
 
 from cosmos_rl.utils.logging import logger
 from cosmos_rl.utils.parallelism import ParallelDims
@@ -19,8 +20,9 @@ def parallelize_model(
         return None, None
 
     world_mesh = parallel_dims.mesh
-    # DDP
-    if parallel_dims.dp_replicate_enabled:
+    if parallel_dims.dp_shard_enabled:
+        _apply_fsdp(model, world_mesh["dp_shard"])
+    elif parallel_dims.dp_replicate_enabled:
         assert world_mesh.ndim == 1, "DDP does not support > 1D parallelism"
         _apply_ddp(model, world_mesh)
 
@@ -30,3 +32,7 @@ def parallelize_model(
 def _apply_ddp(model: nn.Module, dp_mesh: DeviceMesh):
     replicate(model, device_mesh=dp_mesh)
     logger.info("Applied DDP to the model")
+
+
+def _apply_fsdp(model: nn.Module, dp_mesh: DeviceMesh):
+    fully_shard(model, mesh=dp_mesh, reshard_after_forward=True)

cosmos_rl/policy/model/dino_cradio_v3/weight_converter.py

@@ -14,6 +14,19 @@ def convert_weight_from_hf(
     name: str,
     parallel_dims: ParallelDims,
 ) -> Tuple[str, torch.Tensor]:
+    if parallel_dims.dp_shard_enabled:
+        dp_shard_rank = parallel_dims.mesh["dp_shard"].get_local_rank()
+        dp_shard_size = parallel_dims.mesh["dp_shard"].size()
+    else:
+        dp_shard_rank = 0
+        dp_shard_size = 1
+
     dest_name = map_key_from_hf(name)
 
-    return dest_name, tensor
+    if tensor.shape[0] % dp_shard_size == 0:
+        shard = tensor.tensor_split(dp_shard_size, dim=0)[dp_shard_rank]
+    else:
+        chunk_size = (tensor.shape[0] + dp_shard_size - 1) // dp_shard_size
+        shard = tensor[dp_shard_rank * chunk_size : (dp_shard_rank + 1) * chunk_size]
+
+    return dest_name, shard.contiguous()