fix: fix CPU offloading in FSDP grad clipping and weight updates

Updates the gradient clipping implementation to correctly handle parameters offloaded to CPU, bypassing CUDA-specific optimizations when necessary to prevent runtime errors. Refactors the FSDP engine's weight broadcasting logic to properly materialize and batch DTensors in offloaded scenarios. Additionally, introduces a new test suite to verify gradient normalization and clipping behavior across different device configurations.

Author: rchardx

areal/engine/fsdp_engine.py

@@ -449,6 +449,44 @@ def _init_weight_update_from_distributed(self, meta: WeightUpdateMeta):
 
             fut.result()
 
+    def _dtensor_to_full_tensor(self, dtensor: DTensor) -> torch.Tensor:
+        """Convert a DTensor to a full tensor, handling CPU offloaded tensors."""
+        local_tensor = dtensor.to_local()
+        if local_tensor.device.type != "cpu":
+            return dtensor.full_tensor()
+
+        device_mesh = dtensor.device_mesh
+        placements = dtensor.placements
+        temp_dtensor = DTensor.from_local(
+            local_tensor,
+            device_mesh=device_mesh,
+            placements=placements,
+        )
+        return temp_dtensor.full_tensor()
+
+    def _materialize_and_update_bucket(
+        self,
+        meta: WeightUpdateMeta,
+        named_params: list[tuple[str, nn.Parameter]],
+    ):
+        """Materialize DTensors to full tensors and broadcast to inference engine."""
+        main_rank: bool = dist.get_rank() == 0
+        named_tensors = []
+
+        for name, param in named_params:
+            if isinstance(param.data, DTensor):
+                tensor = self._dtensor_to_full_tensor(param.data)
+            else:
+                tensor = param.data
+                if tensor.device.type == "cpu":
+                    tensor = tensor.to(current_platform.device_type)
+
+            if main_rank:
+                named_tensors.append((name, tensor))
+
+        if named_tensors:
+            self._update_bucket_weights_from_distributed(meta, named_tensors)
+
     @trace_perf("fsdp_engine.update_weights_from_distributed", category="comm")
     def _update_weights_from_distributed(self, meta: WeightUpdateMeta):
         """Broadcast parameters (chunked) from rank 0 (FSDP2 compatible)."""
@@ -459,32 +497,33 @@ def _update_weights_from_distributed(self, meta: WeightUpdateMeta):
         dist.barrier(group=self.cpu_group)
 
         weight_chunked_mem_size = meta.weight_chunked_mem_mb * 1024 * 1024
+        fsdp_world_size = dist.get_world_size(self.world_mesh["dp_sp"].get_group())
 
         buffer_size = 0
-        named_tensors = []
+        named_params_bucket: list[tuple[str, nn.Parameter]] = []
 
         for name, param in self.get_model_name_parameters():
             if isinstance(param.data, DTensor):
-                tensor = param.data.full_tensor()
+                local_tensor = param.data.to_local()
+                tensor_size = local_tensor.numel() * local_tensor.element_size()
+                tensor_size *= fsdp_world_size
             else:
-                tensor = param.data
+                tensor_size = param.data.numel() * param.data.element_size()
 
-            # Ranks other than 0 only help to get the full tensor
-            if dist.get_rank() != 0:
-                continue
-
-            tensor_size = tensor.numel() * tensor.element_size()
-
-            if tensor_size + buffer_size > weight_chunked_mem_size:
-                self._update_bucket_weights_from_distributed(meta, named_tensors)
+            if (
+                tensor_size + buffer_size > weight_chunked_mem_size
+                and named_params_bucket
+            ):
+                self._materialize_and_update_bucket(meta, named_params_bucket)
+                named_params_bucket = []
                 buffer_size = 0
 
-            named_tensors.append((name, tensor))
+            named_params_bucket.append((name, param))
             buffer_size += tensor_size
 
-        # Only rank-0 CAN contain named tensors here
-        if named_tensors:
-            self._update_bucket_weights_from_distributed(meta, named_tensors)
+        # Process remaining parameters
+        if named_params_bucket:
+            self._materialize_and_update_bucket(meta, named_params_bucket)
 
         dist.barrier(group=self.cpu_group)
 
@@ -808,6 +847,7 @@ def train_batch(
             list(self.model.parameters()),
             self.world_mesh,
             max_norm=self.optimizer_config.gradient_clipping,
+            offload_params=self.config.fsdp.offload_params,
         )
 
         if not math.isfinite(grad_norm):