fix: fix CPU offloading in FSDP grad clipping and weight updates (#680)

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

@@ -448,6 +448,26 @@ def _init_weight_update_from_distributed(self, meta: WeightUpdateMeta):
 
             fut.result()
 
+    def _get_full_tensor(self, param: nn.Parameter) -> torch.Tensor:
+        """Get full tensor from a parameter, handling DTensor and CPU offloaded tensors."""
+        tensor = param.data
+        if isinstance(tensor, DTensor):
+            # For non-offloaded DTensor, directly call full_tensor()
+            if tensor.device.type != "cpu":
+                return tensor.full_tensor()
+
+            # Handle CPU offloaded DTensor: reconstruct DTensor from local tensor
+            temp_dtensor = DTensor.from_local(
+                tensor.to_local(),
+                device_mesh=tensor.device_mesh,
+                placements=tensor.placements,
+            )
+            return temp_dtensor.full_tensor()
+        else:
+            if tensor.device.type == "cpu":
+                tensor = tensor.to(current_platform.device_type)
+            return 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)."""
@@ -458,18 +478,16 @@ 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()
-            else:
-                tensor = param.data
+            tensor = self._get_full_tensor(param)
 
             # Ranks other than 0 only help to get the full tensor
-            if dist.get_rank() != 0:
+            if not main_rank:
                 continue
 
             tensor_size = tensor.numel() * tensor.element_size()
@@ -481,7 +499,7 @@ def _update_weights_from_distributed(self, meta: WeightUpdateMeta):
             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)
 
@@ -807,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):