update muon_sharding_optimizer with rebuilding 2d_params.

Author: xxyux

python/paddle/distributed/fleet/meta_optimizers/muon_sharding_optimizer.py

@@ -58,7 +58,6 @@
 g_shard_bypass_dygraph_optimizer = int(
     os.environ.get("FLAGS_shard_bypass_dygraph_optimizer", 0)
 )
-g_shard_fused_gradient = int(os.environ.get("FLAGS_shard_fused_gradient", 0))
 
 
 def _is_trainable(param):
@@ -239,29 +238,25 @@ def __init__(self, optimizer, hcg=None):
                 key=lambda p: self._param2rank_2d_by_color[color_key][p.name]
             )
 
-        # ---- Backward compatibility: expose legacy attributes ----
-        # These are kept for any external code that might reference them
-        self._params_2d = self._params_2d_by_color.get(None, [])
-        self._params_2d_moe = self._params_2d_by_color.get('moe_expert', [])
-        self._rank2params_2d = self._rank2params_2d_by_color.get(None, {0: []})
-        self._param2rank_2d = self._param2rank_2d_by_color.get(None, {})
-        self._rank2params_2d_moe = self._rank2params_2d_by_color.get(
-            'moe_expert', {0: []}
-        )
-        self._param2rank_2d_moe = self._param2rank_2d_by_color.get(
-            'moe_expert', {}
-        )
+        # 2D params owned by this sharding rank
+        self._local_2d = []
+        for color_key, params_2d in self._params_2d_by_color.items():
+            rank2params_2d_by_color = self._rank2params_2d_by_color[color_key]
+
+            group_info = self._color_to_group_info[color_key]
+            sharding_rank = max(group_info['rank'], 0)
+
+            self._local_2d.extend(rank2params_2d_by_color[sharding_rank])
 
         self.sd_release_grads = (
             strategy.hybrid_configs['pp_configs'].release_gradients
             or sharding_configs.release_gradients
         )
-        self._use_fuse_gradients = g_shard_fused_gradient
+        self._use_fuse_gradients = self.comm_buffer_size_MB > 0
         # ---- Build comm buffers for 2D params (V1-style) ----
         if self._use_fuse_gradients:
-            if not hasattr(self, 'comm_buffer_2d'):
-                self.comm_buffer_2d = self._build_2d_comm_buffers()
-                self.comm_buffer_2d.sort(key=lambda x: x._dst)
+            self.comm_buffer_2d = self._build_2d_comm_buffers()
+            self.comm_buffer_2d.sort(key=lambda x: x._dst)
 
         # ---- Step 3: Build comm buffers for 1D params (V2-style) ----
         self._slice_params = {}
@@ -278,15 +273,9 @@ def __init__(self, optimizer, hcg=None):
         # The optimizer should see:
         #   - All 2D params assigned to this rank (all colors, as whole tensors)
         #   - 1D slice_params for all non-2D params (element-wise shards)
-        local_2d_params = []
-        for color_key, rank2params in self._rank2params_2d_by_color.items():
-            group_info = self._color_to_group_info.get(color_key, {})
-            color_rank = group_info.get('rank', 0)
-            world_size = group_info.get('world_size', 1)
-            rank_key = color_rank if world_size > 1 else 0
-            local_2d_params.extend(rank2params.get(rank_key, []))
-
-        local_opt_params = local_2d_params + list(self._local_parameter_list_1d)
+        local_opt_params = list(self._local_2d) + list(
+            self._local_parameter_list_1d
+        )
 
         self._set_inner_opt_attr('_parameter_list', local_opt_params)
         self._set_inner_opt_attr('_param_groups', local_opt_params)
@@ -306,18 +295,10 @@ def __init__(self, optimizer, hcg=None):
         _sg_group = hcg.get_sharding_parallel_group()
         _N = self._sharding_world_size
 
-        # 2D params owned by this sharding rank (default color, via legacy alias)
+        # 2D params owned by this sharding rank
         _local_2d_numel = sum(
             int(functools_reduce(lambda x, y: x * y, p.shape, 1))
-            for p in self._rank2params_2d.get(self._sharding_rank, [])
-        )
-        # 2D MoE-expert params owned by this rank (moe_expert color, via legacy alias)
-        _moe_rank_key = (
-            self._moe_sharding_rank if self._moe_sharding_world_size > 1 else 0
-        )
-        _local_2d_moe_numel = sum(
-            int(functools_reduce(lambda x, y: x * y, p.shape, 1))
-            for p in self._rank2params_2d_moe.get(_moe_rank_key, [])
+            for p in self._local_2d
         )
         # 1D (AdamW) slice: each rank holds ceil(numel / sharding_world_size) elements.
         _local_1d_numel = sum(
@@ -327,9 +308,7 @@ def __init__(self, optimizer, hcg=None):
             for p in self._params_1d
         )
 
-        _local_total_numel = (
-            _local_2d_numel + _local_2d_moe_numel + _local_1d_numel
-        )
+        _local_total_numel = _local_2d_numel + _local_1d_numel
         _local_total_MB = (
             _local_total_numel * 2 / (1024 * 1024)
         )  # bf16/fp16 = 2 bytes
@@ -545,6 +524,7 @@ def _build_1d_comm_buffers(self):
         self._comm_buffer_list.sort(key=lambda x: x._dst)
 
     def clear_param_storage(self, color):
+        # Only clear param_storage for 1d_params, 2d_params are not added to comm_buffers.
         self.clear_color.add(color)
         if color in self._color_to_comm_buffer_list.keys():
             for comm_buffer in self._color_to_comm_buffer_list[color]:
@@ -671,31 +651,13 @@ def reduce_gradients(self, parameter_list, hcg):
 
     def filter_parameters(self, parameter_list, hcg):
         """Filter parameters: return local 2D params + initialized 1D slices."""
-        sharding_rank = hcg.get_sharding_parallel_rank()
-        local_2d = [
-            p
-            for p in parameter_list
-            if p.name in self._param2rank_2d
-            and self._param2rank_2d[p.name] == sharding_rank
-        ]
-        # Also include MoE 2D params owned by this rank
-        if self._moe_sharding_world_size > 1:
-            moe_rank = self._moe_sharding_rank
-        else:
-            moe_rank = 0
-        local_2d_moe = [
-            p
-            for p in parameter_list
-            if p.name in self._param2rank_2d_moe
-            and self._param2rank_2d_moe[p.name] == moe_rank
-        ]
         local_1d = [
             self._slice_params[p.name]
             for p in parameter_list
             if p.name in self._slice_params
         ]
         local_1d = [p for p in local_1d if p._is_initialized()]
-        return local_2d + local_2d_moe + local_1d
+        return self._local_2d + local_1d
 
     # ------------------------------------------------------------------
     # Parameter sync after optimizer step
@@ -884,18 +846,7 @@ def step(self):
     def set_state_dict(self, state_dict):
         inner_state = {}
         # Collect local parameters: 2D whole-tensor params + 1D original params
-        # (set_state_dict uses legacy aliases; covers default and moe_expert colors)
-        local_2d = list(self._rank2params_2d.get(self._sharding_rank, []))
-        if self._moe_sharding_world_size > 1:
-            local_2d_moe = list(
-                self._rank2params_2d_moe.get(self._moe_sharding_rank, [])
-            )
-        else:
-            local_2d_moe = list(self._rank2params_2d_moe.get(0, []))
-        parameters = local_2d + local_2d_moe
-        # Add 1D params (use original param names for matching)
-        for p in self._params_1d:
-            parameters.append(p)
+        parameters = list(self._local_2d) + list(self._params_1d)
 
         if "LR_Scheduler" in state_dict:
             inner_state["LR_Scheduler"] = state_dict.pop("LR_Scheduler")