[DEV] torch_dist fixes, speed improvements and memory reduction for layerwise distopt

megatron/core/optimizer/layer_wise_optimizer.py

@@ -1,37 +1,39 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from typing import Callable, List, Optional
 
 import torch
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
-from megatron.core.dist_checkpointing import ShardedTensor
 from megatron.core.dist_checkpointing.dict_utils import nested_values
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.utils import get_pg_rank, get_pg_size
 
 from .clip_grads import count_zeros_fp32, get_grad_norm_fp32
-from .optimizer import ChainedOptimizer, Float16OptimizerWithFloat16Params, MegatronOptimizer
+from .optimizer import (
+    ChainedOptimizer,
+    Float16OptimizerWithFloat16Params,
+    FP32Optimizer,
+    MegatronOptimizer,
+)
 from .optimizer_config import OptimizerConfig
 
 
 class LayerWiseDistributedOptimizer(ChainedOptimizer):
     """Layer-wise distributed optimizer for Megatron-core models.
 
-    This is a experimental distributed optimizer wrapper that distributes weight to DP ranks
-    by full layer. Implemented as ChainedOptimizer to support different weights use different
-    optimizers (e.g. muon+adam). When using, keep all megatron distributed optimizer related
-    options OFF.
+    Experimental distributed optimizer wrapper that distributes weight to DP ranks by layer.
+    Implemented as ChainedOptimizer to support multiple optimizers (e.g. muon + adamW)
+    When using, keep all megatron distributed-optimizer related options OFF.
 
     How LayerWiseDistributedOptimizer work:
     1. weights are splited into lists and each rank only keep its shard in its optimizer
-    2. Megatron DDP handle allreduce grad for all params, note that each rank have full model
-    and grad.
+    2. Megatron DDP handle allreduce grad, note that each rank have full model and grad
     3. optimizer is already modified so only param belong to this DP rank is updated
-    3. grad_norm and zero counting will reduce metrics globally in step function
-    4. Do regular update with chained optimizers, optimizer is already modified so partial update
-    happens.
-    5. allgather updated params to every rank(currently through broadcast loop)
+    4. grad_norm and zero counting will reduce metrics globally in step function
+    5. Do regular update with chained optimizers, modified optimizer only update shard
+    6. allgather updated params to every rank
     """
 
     def __init__(
@@ -53,40 +55,42 @@ def __init__(
 
         self.pg_collection = pg_collection
         self.shard_params(optimizers)
-        # wrap optimizer after sharding to avoid unnecessary master weight creation
-        # TODO(deyuf): check if underlying optimizer.config need to fixed and if so can use
-        # that instead of passing
-        if init_state_fn_list is None:
-            init_state_fn_list = [None] * len(optimizers)
-        else:
-            assert len(init_state_fn_list) == len(optimizers), (
-                "init_state_fn_list must be the " "same length as optimizers if provided"
-            )
+        if init_state_fn_list:
+            assert len(init_state_fn_list) == len(
+                optimizers
+            ), "init_state_fn_list must be the same length as optimizers if provided"
 
+        # wrap optimizer after sharding to avoid unnecessary master weight creation
+        # for higher precision, optimizers are wrapped with megatron already
         if config.bf16:
-            if isinstance(optimizers[0], Float16OptimizerWithFloat16Params):
-                raise TypeError('LayerWiseDistributedOptimizer received Float16 optimizer already.')
-            optimizers = [
-                Float16OptimizerWithFloat16Params(optim, config, None, init_state_fn_list[idx])
-                for idx, optim in enumerate(optimizers)
-            ]
+            # unwrap FP32 optimizer, possibly from reusing get_megatron_optimizer for adam
+            for i in range(len(optimizers)):
+                opt = optimizers[i]
+                if isinstance(opt, Float16OptimizerWithFloat16Params):
+                    raise TypeError(
+                        'LayerWiseDistributedOptimizer received Float16 optimizer already.'
+                    )
+                # unwrap FP32 optimizer from reusing get_megatron_optimizer for adam
+                if isinstance(opt, FP32Optimizer):
+                    opt = opt.optimizer
+                optimizers[i] = Float16OptimizerWithFloat16Params(
+                    opt, config, None, init_state_fn_list[i] if init_state_fn_list else None
+                )
+
         super().__init__(optimizers)
 
         # TODO(kunlun, deyuf): potential future perf optimization
-        # since allreduce is unchanged and handled by megatron DDP, they're already in contiguous
-        # gbuf, so instead of shard param by layer randomly, we can still shard by buf range but
-        # keep some "extras" to keep boundary weight not sharded. This way each rank do some
-        # duplicated work but we can call single allgather later and all current distopt
-        # optimization can be applied.
+        # since allreduce is unchanged and handled by megatron DDP, they're already in
+        # contiguous gbuf. So instead of shard param by layer randomly, we can shard by
+        # buf range but keep some "extras" to keep boundary weight not sharded.
+        # This way each rank do some duplicated work but allgather_v is no longer needed
+        # All current distopt optimization can also be potentially applied
 
     def shard_params(self, optimizers):
         """Shard all params into lists by rank."""
-        # We'll optimize sharding later if there is perf issue. should be ok since linear are
-        # grouped already.
-        # Key is to create separate sharding for dp/expt parallel, saved in dp_cp_params_list,
-        # expt_dp_params_list.
-        # Example of 4 dp rank and 10 non-expert parameters p0-p9, then dp_cp_params_list will
-        # look like: [[p0, p4, p8], [p1, p5, p9], [p2, p6], [p3, p7]]
+        # list of parameter are sorted by numel and assigned to ranks in ping-pong style
+        # example of 4 ranks and 10 parameters p0-p9 after sorting, then dp_cp_params_list will be
+        # [[p0, p7, p8], [p1, p6, p9], [p2, p5], [p3, p4]]
 
         # simplify when dp_cp group size is 1
         if get_pg_size(self.pg_collection.dp_cp) == 1:
@@ -97,40 +101,87 @@ def shard_params(self, optimizers):
         dp_cp_idx, expt_dp_idx = 0, 0
         dp_cp_size = get_pg_size(self.pg_collection.dp_cp)
         expt_dp_size = get_pg_size(self.pg_collection.expt_dp)
+        # create ping-pong style loop so memory is more balanced
+        dp_cp_loop = list(range(dp_cp_size)) + list(range(dp_cp_size))[::-1]
+        expt_dp_loop = list(range(expt_dp_size)) + list(range(expt_dp_size))[::-1]
         self.dp_cp_params_list = [[] for _ in range(dp_cp_size)]
         self.expt_dp_params_list = [[] for _ in range(expt_dp_size)]
-        # get all param groups, this is called before init so cannot rely on
-        # Chained optimizer method
+        # get all param groups
         param_groups = []
         for optimizer in optimizers:
             param_groups += optimizer.param_groups
-        for group in param_groups:
-            params_this_rank = []
-            if group.get("is_expert_parallel", False):
-                for p in group["params"]:
-                    if expt_dp_idx == get_pg_rank(self.pg_collection.expt_dp):
-                        params_this_rank.append(p)
-                    self.expt_dp_params_list[expt_dp_idx].append(p)
-                    expt_dp_idx = (expt_dp_idx + 1) % expt_dp_size
+
+        # sort param in all groups by param numel and assign to each rank evenly
+        param_list = []
+        for group_index, group in enumerate(param_groups):
+            for p in group["params"]:
+                param_list.append((p, group_index))
+        param_list.sort(key=lambda x: x[0].numel())
+        param_groups_this_rank = [[] for g in param_groups]
+
+        # assign params to rank in ping-pong style loop
+        for p, group_index in param_list:
+            if param_groups[group_index].get("is_expert_parallel", False):
+                if expt_dp_loop[expt_dp_idx] == get_pg_rank(self.pg_collection.expt_dp):
+                    param_groups_this_rank[group_index].append(p)
+                self.expt_dp_params_list[expt_dp_loop[expt_dp_idx]].append(p)
+                expt_dp_idx = (expt_dp_idx + 1) % len(expt_dp_loop)
             else:
-                for p in group["params"]:
-                    if dp_cp_idx == get_pg_rank(self.pg_collection.dp_cp):
-                        params_this_rank.append(p)
-                    self.dp_cp_params_list[dp_cp_idx].append(p)
-                    dp_cp_idx = (dp_cp_idx + 1) % dp_cp_size
-            # now we modify the group to only handle local params
-            group["params"] = params_this_rank
+                if dp_cp_loop[dp_cp_idx] == get_pg_rank(self.pg_collection.dp_cp):
+                    param_groups_this_rank[group_index].append(p)
+                self.dp_cp_params_list[dp_cp_loop[dp_cp_idx]].append(p)
+                dp_cp_idx = (dp_cp_idx + 1) % len(dp_cp_loop)
+
+        # now we modify the group to only handle local params
+        for groups, params in zip(param_groups, param_groups_this_rank):
+            groups["params"] = params
 
         # simplify when expt_dp group size is 1 or expert parallel is off
         if expt_dp_size == 1 or len(self.expt_dp_params_list[0]) == 0:
             self.expt_dp_params_list = None
 
+    @torch.no_grad()
+    def allgather_params(self) -> None:
+        """All-gather updated params from all ranks."""
+
+        # helper function to flatten local params, allgather, unflatten and copy to model params
+        def _allgather_helper(params_list, group):
+            # flatten this rank's params and create empty tensor output list
+            device = params_list[0][0].device
+            dtype = params_list[0][0].dtype
+            rank = get_pg_rank(group)
+            # for rank without params create empty tensor and participate in allgather
+            src = (
+                _flatten_dense_tensors(params_list[rank])
+                if len(params_list[rank]) > 0
+                else torch.empty(0, device=device, dtype=dtype)
+            )
+            output_list = [
+                torch.empty(sum([p.numel() for p in params]), device=device, dtype=dtype)
+                for params in params_list
+            ]
+            # single all_gather_v to collect all updated params
+            torch.distributed.all_gather(output_list, src, group=group)
+            # unflatten and copy gathered params for each rank i
+            for idx, (flat_params, params) in enumerate(zip(output_list, params_list)):
+                # skip local params and empty tensors
+                if len(params) == 0 or idx == rank:
+                    continue
+                updated_params = _unflatten_dense_tensors(flat_params, params)
+                for updated_p, model_p in zip(updated_params, params):
+                    model_p.data.copy_(updated_p)
+
+        if self.pg_collection is None:
+            return
+        if self.dp_cp_params_list:
+            _allgather_helper(self.dp_cp_params_list, self.pg_collection.dp_cp)
+        if self.expt_dp_params_list:
+            _allgather_helper(self.expt_dp_params_list, self.pg_collection.expt_dp)
+
     @torch.no_grad()
     def broadcast_params(self):
-        """All rank broadcast updated local params(allgatherv)."""
-        # Broadcast linear layer weights to all other ranks.
-        # This may not be slower than PyTorch allgatherv which calls broadcast internally.
-        # TODO(skyw): Profile and implement more efficient version.
+        """All rank broadcast updated local params."""
+        # Broadcast linear layer weights to all other ranks. Kept as reference test.
         if self.dp_cp_params_list is None:
             return
         for i, params in enumerate(self.dp_cp_params_list):
@@ -170,7 +221,7 @@ def step(self):  # type: ignore[no-untyped-def]
         update_successful, grad_norm, num_zeros_in_grad = super().step()
 
         # All gather updated params.
-        self.broadcast_params()
+        self.allgather_params()
 
         return update_successful, grad_norm, num_zeros_in_grad
 
@@ -187,10 +238,33 @@ def sharded_state_dict(
 
         # for fixed DP usage only
         for sh_base in nested_values(sharded_state_dict):
-            if isinstance(sh_base, ShardedTensor):
+            if hasattr(sh_base, 'replica_id'):
                 assert (
-                    len(sh_base.replica_id) == 3
-                ), f'Expected replica_id format (PP, TP, DP), got: {sh_base}'
-                sh_base.replica_id = (*sh_base.replica_id[:2], 0)
+                    isinstance(sh_base.replica_id, int) or len(sh_base.replica_id) == 3
+                ), f'Expected replica_id as int or (PP, TP, DP), got: {sh_base}'
+                sh_base.replica_id = (
+                    0 if isinstance(sh_base.replica_id, int) else (*sh_base.replica_id[:2], 0)
+                )
+
+        if len(self.chained_optimizers) == 1:
+            wrapped_sharded_state_dict = {1: sharded_state_dict}
+        else:
+            wrapped_sharded_state_dict = sharded_state_dict
+        # Adjust dict due to possible empty rank 0 which output common_dict
+        for sd in wrapped_sharded_state_dict.values():
+            # Drop empty group state to avoid save in common dict (non-empty rank still save)
+            if 'fp32_from_fp16_params' in sd:
+                sd['fp32_from_fp16_params'][:] = [
+                    group for group in sd['fp32_from_fp16_params'] if group
+                ]
+            # TODO(deyuf): 'common_step' code path is broken and 'step' is saved in 'param_groups'
+            # Find next 'step' if present. note this still break if rank0 adam is fully empty
+            step = next(
+                (group['step'] for group in sd['optimizer']['param_groups'] if 'step' in group),
+                None,
+            )
+            if step is not None:
+                for group in sd['optimizer']['param_groups']:
+                    group['step'] = step
 
         return sharded_state_dict

tests/unit_tests/test_layer_wise_optimizer.py

@@ -1,4 +1,5 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
 import os
 
 import pytest
@@ -401,3 +402,39 @@ def test_parameter_updates_insufficient_parameters(self):
         This will be insufficient when world size > 2.
         """
         self._run_parameter_update_test(model_class=TinyModel)
+
+    def test_broadcast_vs_allgather(self):
+        """Test LayerWiseDistributedOptimizer allgather code agains broadcast code."""
+        model, optimizer, pg_collection = self.create_model_and_optimizer(model_class=SimpleModel)
+
+        # Create reference model and optimizer using the same function
+        reference_model, reference_optimizer, _ = self.create_model_and_optimizer(
+            model_class=SimpleModel, copy_from=model
+        )
+
+        # Set same gradients on both models
+        for param, ref_param in zip(model.parameters(), reference_model.parameters()):
+            assert torch.equal(param.data, ref_param.data)
+            torch.testing.assert_close(param.data, ref_param.data, rtol=0, atol=0)
+            grad_value = torch.randn_like(param)
+            torch.distributed.broadcast(grad_value, src=0, group=pg_collection.dp_cp)
+            param.main_grad = grad_value.clone().detach()
+            ref_param.main_grad = grad_value.clone().detach()
+
+        optimizer.step()
+
+        # Verify at least some parameters were updated
+        params_updated = 0
+        for param, ref_param in zip(model.parameters(), reference_model.parameters()):
+            if not torch.equal(param.data, ref_param.data):
+                params_updated += 1
+
+        assert params_updated > 0, "At least some parameters should be updated"
+
+        # step() internal call allgather_params. replace reference object with bcast
+        reference_optimizer.allgather_params = reference_optimizer.broadcast_params
+        reference_optimizer.step()
+
+        # Verify updated values match reference optimizer
+        for param, ref_param in zip(model.parameters(), reference_model.parameters()):
+            torch.testing.assert_close(param.data, ref_param.data, rtol=0, atol=0)