feat: add Muon optimizer support

Author: xxyux

test/collective/fleet/hybrid_parallel_sharding_mp_model.py

@@ -21,7 +21,6 @@
 import paddle
 import paddle.distributed as dist
 from paddle.distributed import fleet
-from paddle.distributed.fleet.utils.log_util import logger
 from paddle.distributed.fleet.utils.mix_precision_utils import (
     MixPrecisionLayer,
     MixPrecisionOptimizer,
@@ -45,6 +44,8 @@
 
 
 class SimpleMPNet(paddle.nn.Layer):
+    """Tensor-parallel model: linear1 is column-split, linear2 is row-split."""
+
     def __init__(
         self,
         vocab_size,
@@ -70,31 +71,32 @@ def __init__(
             ),
         )
 
+        # Each TP rank owns a contiguous column slice of fc1 and row slice of fc2.
         inner_per_rank = inner_size // mp_degree
         fc1_start = mp_id * inner_per_rank
         fc1_end = fc1_start + inner_per_rank
 
-        init_fc1_slice = np_fc1[:, fc1_start:fc1_end]
-
         self.linear1 = fleet.meta_parallel.ColumnParallelLinear(
             hidden_size,
             inner_size,
             weight_attr=paddle.framework.ParamAttr(
-                initializer=paddle.nn.initializer.Assign(init_fc1_slice)
+                initializer=paddle.nn.initializer.Assign(
+                    np_fc1[:, fc1_start:fc1_end]
+                )
             ),
-            gather_output=False,  # 关键：输出保持切分状态，不聚合，直接喂给下一层 RowParallel
+            gather_output=False,  # keep output sharded for RowParallel input
             has_bias=True,
         )
 
-        init_fc2_slice = np_fc2[fc1_start:fc1_end, :]
-
         self.linear2 = fleet.meta_parallel.RowParallelLinear(
             inner_size,
             hidden_size,
             weight_attr=paddle.framework.ParamAttr(
-                initializer=paddle.nn.initializer.Assign(init_fc2_slice)
+                initializer=paddle.nn.initializer.Assign(
+                    np_fc2[fc1_start:fc1_end, :]
+                )
             ),
-            input_is_parallel=True,  # 关键：告诉这一层，输入已经是切分过的了
+            input_is_parallel=True,  # input already sharded from ColumnParallel
             has_bias=True,
         )
 
@@ -119,6 +121,8 @@ def forward(self, x):
 
 
 class SimpleDPNet(paddle.nn.Layer):
+    """Single-process reference model with identical weight initialisation."""
+
     def __init__(
         self,
         vocab_size,
@@ -209,8 +213,8 @@ def setUp(self):
     def train_batch(self, batch, model, optimizer):
         output = model(batch)
         loss = output.mean()
-        loss.backward()  # do backward
-        optimizer.step()  # update parameters
+        loss.backward()
+        optimizer.step()
         optimizer.clear_grad()
         return loss
 
@@ -239,7 +243,6 @@ def build_optimizer(self, model, strategy=None, Optimizer="adam"):
         return optimizer
 
     def build_model_optimizer(self, Optimizer="adam", amp_level=None):
-
         np_fc1 = np.random.random_sample((hidden_size, inner_size)).astype(
             "float32"
         )
@@ -253,6 +256,7 @@ def build_model_optimizer(self, Optimizer="adam", amp_level=None):
             "float32"
         )
 
+        # model_a: sharded+TP distributed model; model_b: single-process reference
         model_a = SimpleMPNet(
             vocab_size,
             hidden_size,
@@ -302,79 +306,45 @@ def sharding_model(self, Optimizer, sharded_accumulators, amp_level=None):
         )
 
         hcg = fleet.get_hybrid_communicate_group()
-        # degree
         mp_degree = self.strategy.hybrid_configs['mp_degree']
         sharding_degree = self.strategy.hybrid_configs['sharding_degree']
-        # rank
         rank = dist.get_rank()
         sharding_rank = hcg.get_sharding_parallel_rank()
         tp_rank = hcg.get_model_parallel_rank()
-        # data size
         local_batch_size = batch_size // sharding_degree
+        tp_group = hcg.get_model_parallel_group()
 
         for idx in range(STEPS):
-            print(f"Step = {idx}")
-
+            # Each sharding rank processes its own mini-batch slice.
             start_index = sharding_rank * local_batch_size
             end_index = (sharding_rank + 1) * local_batch_size
-
             batch_sharding = paddle.to_tensor(
                 self.data[idx][start_index:end_index]
             )
-            logger.info(
-                f"rank = {rank}, sharding_rank = {sharding_rank}, tp_rank = {tp_rank} start_index = {start_index}, end_index = {end_index}"
-            )
-
             batch_single = paddle.to_tensor(self.data[idx])
+
             loss_a = self.train_batch(batch_sharding, model_a, optimizer_a)
             loss_b = self.train_batch(batch_single, model_b, optimizer_b)
 
-            # Reduce loss
+            # Average loss across all ranks for a fair global comparison.
             loss_a_metric = loss_a.detach().clone()
             dist.all_reduce(loss_a_metric, op=dist.ReduceOp.SUM)
-            world_size = dist.get_world_size()
-            loss_a_global = loss_a_metric / world_size
-
-            abs_err = np.abs(loss_a_global.numpy() - loss_b.numpy())
-            rel_err = abs_err / (np.abs(loss_b.numpy()) + 1e-9)
+            loss_a_global = loss_a_metric / dist.get_world_size()
 
-            print(
-                f"step {idx}, loss_a(local)={loss_a.numpy()}, loss_a(global)={loss_a_global.numpy()}, loss_b={loss_b.numpy()}"
-            )
-            print(f"abs_error = {abs_err}, rel_error = {rel_err}")
-
-            print(f"\n--- Checking Parameters at Step {idx} ---")
-
-            # 获取参数列表 (假设顺序是一致的)
-            params_a = model_a.parameters()
-            params_b = model_b.parameters()
-
-            tp_group = hcg.get_model_parallel_group()
-
-            for i, (param_a, param_b) in enumerate(zip(params_a, params_b)):
-                name = param_a.name
-
-                # 1. 获取本地参数值
+            # Compare each parameter between the distributed and reference model.
+            for param_a, param_b in zip(
+                model_a.parameters(), model_b.parameters()
+            ):
                 val_a_local = param_a.numpy()
                 val_b = param_b.numpy()
 
-                # 2. 判断是否是 TP 参数 (通过形状是否一致判断)
-                is_tp_param = val_a_local.shape != val_b.shape
-
-                if is_tp_param:
-                    # 分布式参数：需要 Gather
+                # TP parameters are split across TP ranks; gather before comparing.
+                if val_a_local.shape != val_b.shape:
                     gathered_list = []
                     paddle.distributed.all_gather(
                         gathered_list, param_a, group=tp_group
                     )
-
-                    # 3. 拼接
-                    # 关键：ColumnParallelLinear 是按列切分 (axis=1)
-                    #       RowParallelLinear 是按行切分 (axis=0)
-                    #       VocabParallelEmbedding 是按行切分 (axis=0)
-
-                    # 简单的启发式判断拼接维度：
-                    # 看看 val_b 的哪个维度是 val_a 的 mp_degree 倍
+                    # Determine the split axis: find which dim is mp_degree times smaller.
                     concat_axis = -1
                     for dim in range(len(val_b.shape)):
                         if (
@@ -383,43 +353,21 @@ def sharding_model(self, Optimizer, sharded_accumulators, amp_level=None):
                         ):
                             concat_axis = dim
                             break
-
                     if concat_axis == -1:
-                        print(
-                            f"[Warning] Param {name} shape mismatch but axis not found. Skip."
-                        )
                         continue
-
-                    # 拼接
                     val_a_global = np.concatenate(
                         [t.numpy() for t in gathered_list], axis=concat_axis
                     )
                 else:
-                    # 非分布式参数 (或者 Sharding Only 参数)，本地即全量
                     val_a_global = val_a_local
 
-                # 4. 计算误差
-                # 注意：由于 O2 精度问题，对比时可能需要 cast 到 float32
-                diff = np.abs(val_a_global - val_b)
-                max_abs_err = diff.max()
-                max_rel_err = max_abs_err / (np.abs(val_b).max() + 1e-9)
-
-                # 5. 打印与断言
-                # 只在 Rank 0 打印，避免刷屏
-                if dist.get_rank() == 0:
-                    print(
-                        f"Param: {name} | Type: {'TP' if is_tp_param else 'Global'} | Shape: {val_a_global.shape}"
-                    )
-                    print(f"  Max Abs Error: {max_abs_err:.2e}")
-                    print(f"  Max Rel Error: {max_rel_err:.2e}")
-
-                # 设置相对宽松的阈值 (因为 O2 累积误差)
+                # Loose tolerance to account for O2 AMP accumulated error.
                 np.testing.assert_allclose(
                     val_a_global,
                     val_b,
                     rtol=1e-4,
                     atol=1e-4,
-                    err_msg=f"Param {name} mismatch!",
+                    err_msg=f"Param {param_a.name} mismatch!",
                 )
 
     def test_sharding_muon(self):