refactor(mfsdp): clarify Muon update phases and types

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Author: wujingyue

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/muon.py

@@ -52,6 +52,7 @@ class FsdpMuon(Muon):
             including in optimizer parameter-group dictionaries.
     """
 
+    dp_mesh: DeviceMesh
     _fsdp_groups: list[FsdpParameterGroup]
     _owner_maps: list[dict[int, int]]
 
@@ -65,12 +66,12 @@ def __init__(self, params: ParamsT, dp_mesh: DeviceMesh, **kwargs: Any) -> None:
 
         # Preserve discovery order: communication must be posted identically on every rank.
         self._fsdp_groups = []
-        seen = set()
+        seen: set[FsdpParameterGroup] = set()
         for group in self.param_groups:
             for param in group["params"]:
                 if param.ndim != 2:
                     continue
-                fsdp_group = get_containing_parameter_group(param)
+                fsdp_group: FsdpParameterGroup | None = get_containing_parameter_group(param)
                 assert fsdp_group is not None
                 if fsdp_group not in seen:
                     seen.add(fsdp_group)
@@ -79,14 +80,14 @@ def __init__(self, params: ParamsT, dp_mesh: DeviceMesh, **kwargs: Any) -> None:
 
     def _validate_param_groups(self) -> None:
         """Validate optimizer parameters and the FSDP buffers that contain them."""
-        seen = set()
+        seen: set[FsdpParameterGroup] = set()
         for group in self.param_groups:
             if group["weight_decay"] != 0:
                 raise NotImplementedError("FsdpMuon does not support nonzero weight decay yet.")
             for param in group["params"]:
                 if param.ndim != 2:
                     continue
-                fsdp_group = get_containing_parameter_group(param)
+                fsdp_group: FsdpParameterGroup | None = get_containing_parameter_group(param)
                 if fsdp_group is None:
                     raise ValueError("FsdpMuon requires parameters sharded by MFSDP v2.")
                 if fsdp_group in seen:
@@ -117,13 +118,13 @@ def step(self, closure: Callable[[], float] | None = None) -> None:
             for param in group["params"]:
                 if param.ndim != 2 or param.grad is None:
                     continue
-                state = self.state[param]
+                state: dict[str, Any] = self.state[param]
                 if not state:
                     state["momentum_buffer"] = torch.zeros_like(param)
                 param_to_group_index[param] = group_index
 
         for fsdp_group, owners in zip(self._fsdp_groups, self._owner_maps, strict=True):
-            group_indices = {
+            group_indices: dict[int, int] = {
                 i: param_to_group_index[fp.sharded]
                 for i, fp in enumerate(fsdp_group.fsdp_parameters)
                 if fp.sharded in param_to_group_index
@@ -139,56 +140,90 @@ def _step_group(
         group_indices: dict[int, int],
         tensor_to_owner: dict[int, int],
     ) -> None:
-        """Update a buffer using a tensor-index-to-optimizer-group-index mapping."""
-        buffer = fsdp_group.main_weight
+        """Gather momentum, orthogonalize on owners, and scatter updates.
+
+        Args:
+            fsdp_group: FSDP buffers and parameters to update.
+            group_indices: Active indices into ``fsdp_group.fsdp_parameters`` mapped
+                to indices into ``self.param_groups`` for their optimizer settings.
+            tensor_to_owner: The same active tensor indices mapped to global owner ranks.
+        """
+        buffer: DBuffer = fsdp_group.main_weight
+        momentum_inputs: DBuffer = self._compute_momentum_inputs(fsdp_group, group_indices)
+        this_rank: int = self.dp_mesh.get_rank()
+        full_updates: dict[int, torch.Tensor] = {
+            i: torch.empty(buffer.layout.tensor_shapes[i], dtype=buffer.dtype, device=buffer.device)
+            for i, owner in tensor_to_owner.items()
+            if owner == this_rank
+        }
+        gather(momentum_inputs, full_updates, tensor_to_owner=tensor_to_owner)
+        self._orthogonalize_updates(fsdp_group, group_indices, full_updates)
 
-        # Compute the local momentum inputs.
-        inputs = DBuffer(buffer.mesh, buffer.placements, buffer.layout, buffer.dtype, buffer.device)
+        local_updates: DBuffer = DBuffer(
+            buffer.mesh, buffer.placements, buffer.layout, torch.float32, buffer.device
+        )
+        scatter(full_updates, local_updates, tensor_to_owner=tensor_to_owner)
+        self._apply_local_updates(fsdp_group, group_indices, full_updates, local_updates)
+
+    def _compute_momentum_inputs(
+        self, fsdp_group: FsdpParameterGroup, group_indices: dict[int, int]
+    ) -> DBuffer:
+        """Update sharded momentum and build the inputs for orthogonalization."""
+        buffer: DBuffer = fsdp_group.main_weight
+        momentum_inputs: DBuffer = DBuffer(
+            buffer.mesh, buffer.placements, buffer.layout, buffer.dtype, buffer.device
+        )
         for i, group_index in group_indices.items():
-            group = self.param_groups[group_index]
-            param = cast(DTensor, fsdp_group.fsdp_parameters[i].sharded)
+            group: dict[str, Any] = self.param_groups[group_index]
+            param: DTensor = cast(DTensor, fsdp_group.fsdp_parameters[i].sharded)
             assert isinstance(param.grad, DTensor)
-            local = inputs.get_tensor_view(i)
-            if local.numel() == 0:
-                continue
-            momentum = self.state[param]["momentum_buffer"].to_local()
-            grad = param.grad.to_local().to(momentum.dtype).reshape(momentum.shape)
+            local_input: torch.Tensor = momentum_inputs.get_tensor_view(i)
+            momentum: torch.Tensor = self.state[param]["momentum_buffer"].to_local()
+            grad: torch.Tensor = param.grad.to_local().to(momentum.dtype).reshape(momentum.shape)
+            # EO's OrthogonalizedOptimizer.step implements momentum inline, without a helper.
             momentum.lerp_(grad, 1 - group["momentum"])
-            update = grad.lerp(momentum, group["momentum"]) if self.nesterov else momentum
-            local.copy_(update.reshape(local.shape))
+            update: torch.Tensor = (
+                grad.lerp(momentum, group["momentum"]) if self.nesterov else momentum
+            )
+            local_input.copy_(update.reshape(local_input.shape))
+        return momentum_inputs
 
-        # Gather and orthogonalize full matrices on their owners.
-        this_rank = self.dp_mesh.get_rank()
-        full_updates = {
-            i: torch.empty(buffer.layout.tensor_shapes[i], dtype=buffer.dtype, device=buffer.device)
-            for i, owner in tensor_to_owner.items()
-            if owner == this_rank
-        }
-        gather(inputs, full_updates, tensor_to_owner=tensor_to_owner)
+    def _orthogonalize_updates(
+        self,
+        fsdp_group: FsdpParameterGroup,
+        group_indices: dict[int, int],
+        full_updates: dict[int, torch.Tensor],
+    ) -> None:
+        """Replace this rank's gathered inputs with full orthogonalized updates."""
         with eo_utils.fp32_matmul_precision(self.fp32_matmul_prec):
             for i, update in full_updates.items():
-                param = fsdp_group.fsdp_parameters[i].sharded
-                group = self.param_groups[group_indices[i]]
-                kwargs = {k: v for k, v in group.items() if k != "params"}
+                param: torch.nn.Parameter = fsdp_group.fsdp_parameters[i].sharded
+                group: dict[str, Any] = self.param_groups[group_indices[i]]
+                kwargs: dict[str, Any] = {k: v for k, v in group.items() if k != "params"}
                 full_updates[i] = self.orthogonalize(param, update.float(), **kwargs).contiguous()
 
-        # Scatter updates and apply each local shard.
-        local_updates = DBuffer(
-            buffer.mesh, buffer.placements, buffer.layout, torch.float32, buffer.device
-        )
-        scatter(full_updates, local_updates, tensor_to_owner=tensor_to_owner)
+    def _apply_local_updates(
+        self,
+        fsdp_group: FsdpParameterGroup,
+        group_indices: dict[int, int],
+        full_updates: dict[int, torch.Tensor],
+        local_updates: DBuffer,
+    ) -> None:
+        """Apply owned result slices and received updates to local parameter shards."""
+        buffer: DBuffer = fsdp_group.main_weight
         for i, group_index in group_indices.items():
-            param = cast(DTensor, fsdp_group.fsdp_parameters[i].sharded)
-            weight = param.to_local()
+            param: DTensor = cast(DTensor, fsdp_group.fsdp_parameters[i].sharded)
+            weight: torch.Tensor = param.to_local()
             if weight.numel() == 0:
                 continue
             if i in full_updates:
                 # Scatter writes only remote-owned shards; owners keep their full result.
-                offset = max(buffer.offset - buffer.layout.get_tensor_range(i).start, 0)
-                update = full_updates[i].view(-1).narrow(0, offset, weight.numel())
+                offset: int = max(buffer.offset - buffer.layout.get_tensor_range(i).start, 0)
+                update: torch.Tensor = full_updates[i].view(-1).narrow(0, offset, weight.numel())
             else:
                 update = local_updates.get_tensor_view(i)
             update = update.reshape(weight.shape).to(weight.dtype)
+            # These hooks are inherited from EO's OrthogonalizedOptimizer, via Muon.
             self.pre_weight_update_fn_inplace(weight, update)
             weight.add_(update, alpha=-self.param_groups[group_index]["lr"])
             self.post_weight_update_fn_inplace(weight)