fix: fix CPU offload for FSDP

Updates the FSDP engine to correctly handle DTensors residing on CPU during weight synchronization, ensuring proper materialization to full tensors.
Refactors weight gathering logic to batch parameter processing, improving memory efficiency during distributed updates.
Modifies gradient norm calculation and clipping functions to support CPU-resident gradients, ensuring operations are applied to the correct device types.

Author: rchardx

areal/engine/fsdp_engine.py

@@ -122,6 +122,7 @@ def __init__(self, config: TrainEngineConfig):
         self.parallel_helper: ParallelHelper
         self.world_mesh: DeviceMesh
 
+        self.fsdp_group: dist.ProcessGroup
         self.dp_group: dist.ProcessGroup
         self.sp_group: dist.ProcessGroup
         self.mp_group: dist.ProcessGroup
@@ -192,6 +193,7 @@ def create_process_group(self, parallel_strategy: ParallelStrategy | None = None
 
         self.world_mesh = self.parallel_helper.world_mesh
 
+        self.fsdp_group = self.world_mesh["dp_sp"].get_group()
         self.dp_group = self.world_mesh["dp"].get_group()
         self.sp_group = self.world_mesh["sp"].get_group()
 
@@ -449,6 +451,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 +499,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.fsdp_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 +849,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):

areal/utils/fsdp/grad.py

@@ -50,8 +50,7 @@
 
 
 def to_local_if_dtensor(tensor: Tensor | DTensor) -> Tensor:
-    with torch.no_grad():
-        return tensor.to_local() if isinstance(tensor, DTensor) else tensor
+    return tensor.to_local() if isinstance(tensor, DTensor) else tensor
 
 
 def device_mesh_has_dim(mesh: DeviceMesh, dim_name: str) -> bool:
@@ -90,11 +89,12 @@ def get_grad_norm_fp32(
     data_parallel_group: ProcessGroup,
     model_parallel_group: ProcessGroup,
     norm_type: float = 2.0,
+    offload_params: bool = False,
 ) -> float:
     if isinstance(grads_for_norm, Tensor):
         grads_for_norm = [grads_for_norm]
 
-    grads_for_norm = [to_local_if_dtensor(grad) for grad in grads_for_norm]
+    grads_for_norm = [to_local_if_dtensor(grad).detach() for grad in grads_for_norm]
 
     norm_type = float(norm_type)
     total_norm = 0.0
@@ -105,25 +105,26 @@ def get_grad_norm_fp32(
     device = current_platform.current_device()
 
     if norm_type == torch.inf:
-        norms = [grad.abs().max() for grad in grads_for_norm]
-        total_norm = torch.max(torch.stack(norms)) if norms else 0.0
+        norms = [grad.abs().max().item() for grad in grads_for_norm]
+        total_norm = max(norms) if norms else 0.0
         total_norm_cuda = torch.tensor(
             [float(total_norm)], dtype=torch.float, device=device
         )
         if data_parallel_group:
-            torch.distributed.all_reduce(
+            dist.all_reduce(
                 total_norm_cuda,
-                op=torch.distributed.ReduceOp.MAX,
+                op=dist.ReduceOp.MAX,
                 group=data_parallel_group,
             )
-        torch.distributed.all_reduce(
+        dist.all_reduce(
             total_norm_cuda,
-            op=torch.distributed.ReduceOp.MAX,
+            op=dist.ReduceOp.MAX,
             group=model_parallel_group,
         )
-        total_norm = float(total_norm_cuda[0].item())
+        total_norm = float(total_norm_cuda.item())
     else:
-        if norm_type == 2.0:
+        if norm_type == 2.0 and not offload_params:
+            # Use multi_tensor_applier for better performance when grads are on GPU
             dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=device)
             grad_norm, _ = multi_tensor_applier(
                 l2_norm_impl,
@@ -133,20 +134,23 @@ def get_grad_norm_fp32(
             )
             total_norm_cuda = grad_norm**norm_type
         else:
-            total_norm_cuda = torch.tensor([0.0], dtype=torch.float, device=device)
+            total_norm = 0.0
             for grad in grads_for_norm:
-                grad_norm = torch.norm(grad, norm_type)
-                total_norm_cuda += grad_norm**norm_type
+                grad_norm = torch.norm(grad, norm_type).item()
+                total_norm += grad_norm**norm_type
+            total_norm_cuda = torch.tensor(
+                [float(total_norm)], dtype=torch.float, device=device
+            )
 
         if data_parallel_group:
-            torch.distributed.all_reduce(
+            dist.all_reduce(
                 total_norm_cuda,
-                op=torch.distributed.ReduceOp.SUM,
+                op=dist.ReduceOp.SUM,
                 group=data_parallel_group,
             )
-        torch.distributed.all_reduce(
+        dist.all_reduce(
             total_norm_cuda,
-            op=torch.distributed.ReduceOp.SUM,
+            op=dist.ReduceOp.SUM,
             group=model_parallel_group,
         )
         total_norm = float(total_norm_cuda.item()) ** (1.0 / norm_type)
@@ -159,48 +163,68 @@ def clip_grad_by_total_norm_fp32(
     parameters: list[nn.Parameter],
     max_norm: int | float,
     total_norm: float,
-):
+) -> None:
+    clip_coeff = max_norm / (total_norm + 1.0e-6)
+    if clip_coeff >= 1.0:
+        return
+
     # dtype -> grad
     grads = defaultdict(list)
+    cpu_grads = defaultdict(list)
     for param in parameters:
         if param.grad is not None:
-            # For naive FSDP, lm_head has bf16 grad while others have fp32 grad
             grad = to_local_if_dtensor(param.grad).detach()
-            grads[grad.dtype].append(grad)
+            if grad.device.type != "cpu":
+                grads[grad.dtype].append(grad)
+            else:
+                cpu_grads[grad.dtype].append(grad)
 
-    assert len(grads) > 0, len(grads)
-    clip_coeff = max_norm / (total_norm + 1.0e-6)
-    if clip_coeff < 1.0:
-        for dtype, _grads in grads.items():
-            dummy_overflow_buf = torch.tensor(
-                [0], dtype=torch.int, device=current_platform.device_type
+    if len(grads) == 0 and len(cpu_grads) == 0:
+        return
+
+    from .multi_tensor_apply import (
+        local_multi_tensor_applier,
+        local_multi_tensor_scale,
+    )
+
+    # Clip GPU grads
+    for dtype, _grads in grads.items():
+        dummy_overflow_buf = torch.tensor(
+            [0], dtype=torch.int, device=current_platform.device_type
+        )
+        # For naive FSDP, lm_head has bf16 grad while others have fp32 grad
+        if dtype == torch.float32:
+            multi_tensor_applier(
+                multi_tensor_scale_impl,
+                dummy_overflow_buf,
+                [_grads, _grads],
+                clip_coeff,
+            )
+        else:
+            local_multi_tensor_applier(
+                local_multi_tensor_scale,
+                dummy_overflow_buf,
+                [_grads, _grads],
+                clip_coeff,
             )
-            if dtype == torch.float32:
-                multi_tensor_applier(
-                    multi_tensor_scale_impl,
-                    dummy_overflow_buf,
-                    [_grads, _grads],
-                    clip_coeff,
-                )
-            else:
-                from .multi_tensor_apply import (
-                    local_multi_tensor_applier,
-                    local_multi_tensor_scale,
-                )
 
-                local_multi_tensor_applier(
-                    local_multi_tensor_scale,
-                    dummy_overflow_buf,
-                    [_grads, _grads],
-                    clip_coeff,
-                )
+    # Clip CPU grads
+    dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cpu")
+    for _grads in cpu_grads.values():
+        local_multi_tensor_applier(
+            local_multi_tensor_scale,
+            dummy_overflow_buf,
+            [_grads, _grads],
+            clip_coeff,
+        )
 
 
 def fsdp2_clip_grad_norm(
     parameters: list[nn.Parameter],
     nd_device_mesh: DeviceMesh,
     max_norm: float,
     norm_type: float = 2.0,
+    offload_params: bool = False,
 ) -> float:
     if norm_type <= 0 and norm_type != float("inf"):
         raise ValueError(
@@ -215,7 +239,11 @@ def fsdp2_clip_grad_norm(
     grads_for_norm = get_main_grads_for_grad_norm(parameters, tensor_parallel_rank)
 
     grad_norm = get_grad_norm_fp32(
-        grads_for_norm, fsdp_group, tp_group, norm_type=norm_type
+        grads_for_norm,
+        fsdp_group,
+        tp_group,
+        norm_type=norm_type,
+        offload_params=offload_params,
     )
 
     if parameters: