diff --git a/cosmos_rl/policy/model/dino_cradio_v3/__init__.py b/cosmos_rl/policy/model/dino_cradio_v3/__init__.py index 8359757cb..388842e07 100644 --- a/cosmos_rl/policy/model/dino_cradio_v3/__init__.py +++ b/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(): diff --git a/cosmos_rl/policy/model/dino_cradio_v3/model/model_utils.py b/cosmos_rl/policy/model/dino_cradio_v3/model/model_utils.py index ff1f4cea4..a82a255dd 100755 --- a/cosmos_rl/policy/model/dino_cradio_v3/model/model_utils.py +++ b/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) diff --git a/cosmos_rl/policy/model/dino_cradio_v3/parallelize.py b/cosmos_rl/policy/model/dino_cradio_v3/parallelize.py index 10fd8a019..428c84be7 100644 --- a/cosmos_rl/policy/model/dino_cradio_v3/parallelize.py +++ b/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) diff --git a/cosmos_rl/policy/model/dino_cradio_v3/weight_converter.py b/cosmos_rl/policy/model/dino_cradio_v3/weight_converter.py index 1047f2e73..0d76c3127 100644 --- a/cosmos_rl/policy/model/dino_cradio_v3/weight_converter.py +++ b/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()