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,21 @@ 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()
+
     @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,30 +474,32 @@ 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
+        main_rank = dist.get_rank() == 0
 
         buffer_size = 0
-        named_tensors = []
+        named_tensors: list[tuple[str, torch.Tensor]] = []
 
         for name, param in self.get_model_name_parameters():
             if isinstance(param.data, DTensor):
-                tensor = param.data.full_tensor()
+                tensor = self._dtensor_to_full_tensor(param.data)
             else:
                 tensor = param.data
-
-            # Ranks other than 0 only help to get the full tensor
-            if dist.get_rank() != 0:
-                continue
+                if tensor.device.type == "cpu":
+                    tensor = tensor.to(current_platform.device_type)
 
             tensor_size = tensor.numel() * tensor.element_size()
 
-            if tensor_size + buffer_size > weight_chunked_mem_size:
+            if tensor_size + buffer_size > weight_chunked_mem_size and named_tensors:
                 self._update_bucket_weights_from_distributed(meta, named_tensors)
+                named_tensors = []
                 buffer_size = 0
 
-            named_tensors.append((name, tensor))
+            # Only rank 0 collects tensors for broadcasting
+            if main_rank:
+                named_tensors.append((name, tensor))
             buffer_size += tensor_size
 
-        # Only rank-0 CAN contain named tensors here
+        # Process remaining parameters
         if named_tensors:
             self._update_bucket_weights_from_distributed(meta, named_tensors)
 
@@ -808,6 +825,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):