update muon_sharding_optimizer with rebuilding 2d_params.

Author: xxyux

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

@@ -521,6 +521,10 @@ def _build_1d_comm_buffers(self):
         self._comm_buffer_list.sort(key=lambda x: x._dst)
 
     def clear_param_storage(self, color):
+        assert self._multi_precision, (
+            "Muon Sharding Optimizer only support clear param with multi_precision mode"
+        )
+
         self.clear_color.add(color)
         # 1D params
         if color in self._color_to_comm_buffer_list.keys():

test/collective/fleet/test_parallel_dygraph_muon.py

@@ -26,7 +26,10 @@ def test_muon_sharding_optimizer(self):
         Test logic is in hybrid_parallel_sharding_muon_model.py,
         iterating 4 ns_coeff_types. fp32 matmul is auto-selected on V100.
         """
-        self.run_mnist_2accelerators('hybrid_parallel_sharding_muon_model.py')
+        self.run_mnist_2accelerators(
+            'hybrid_parallel_sharding_muon_model.py',
+            need_envs={"MULTI_PRECISION": "1"},
+        )
 
     def test_muon_sharding_fused_gradient(self):
         """MuonSharding test with FLAGS_shard_fused_gradient=1.
@@ -36,7 +39,10 @@ def test_muon_sharding_fused_gradient(self):
         """
         self.run_mnist_2accelerators(
             'hybrid_parallel_sharding_muon_model.py',
-            need_envs={"FLAGS_shard_fused_gradient": "1"},
+            need_envs={
+                "FLAGS_shard_fused_gradient": "1",
+                "MULTI_PRECISION": "1",
+            },
         )
 
     def test_muon_sharding_fuse_optimizer_states(self):
@@ -46,7 +52,10 @@ def test_muon_sharding_fuse_optimizer_states(self):
         """
         self.run_mnist_2accelerators(
             'hybrid_parallel_sharding_muon_model.py',
-            need_envs={"ENABLE_FUSE_OPTIMIZER_STATES": "1"},
+            need_envs={
+                "ENABLE_FUSE_OPTIMIZER_STATES": "1",
+                "MULTI_PRECISION": "1",
+            },
         )
 
     def test_muon_sharding_release_grads_fused(self):
@@ -60,6 +69,7 @@ def test_muon_sharding_release_grads_fused(self):
             need_envs={
                 "FLAGS_shard_fused_gradient": "1",
                 "RELEASE_GRADIENTS": "1",
+                "MULTI_PRECISION": "1",
             },
         )