【Dist】Update color strategy

python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py

@@ -722,8 +722,7 @@ def __init__(self, optimizer, hcg):
         self._hcg = hcg
         self._sharding_world_size = self._hcg.get_sharding_parallel_world_size()
         self._sharding_rank = self._hcg.get_sharding_parallel_rank()
-        self.clear_color = None
-
+        self.clear_color = []
         self._parameter_list = optimizer._parameter_list
 
         # param name -> slice_param
@@ -931,7 +930,7 @@ def _build_comm_buffers(
                         self.param2bucket[p.name] = [buffer]
 
     def clear_param_storage(self, color):
-        self.clear_color = color
+        self.clear_color.append(color)
         if color in self._color_to_comm_buffer_list.keys():
             for comm_buffer in self._color_to_comm_buffer_list[color]:
                 for param in comm_buffer.params:
@@ -949,12 +948,13 @@ def clear_param_storage(self, color):
                 comm_buffer._clear_param_storage()
 
     def reset_param_storage(self):
-        color = self.clear_color
-        if color is None:
-            return
-        if color in self._color_to_comm_buffer_list.keys():
-            for comm_buffer in self._color_to_comm_buffer_list[color]:
-                comm_buffer._reset_param_storage()
+        for color in self.clear_color:
+            if color is None:
+                continue
+
+            if color in self._color_to_comm_buffer_list.keys():
+                for comm_buffer in self._color_to_comm_buffer_list[color]:
+                    comm_buffer._reset_param_storage()
 
     def clear_grad(self, set_to_zero=True):
         """

test/auto_parallel/CMakeLists.txt

@@ -212,6 +212,9 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
     NVIDIA_TF32_OVERRIDE=0)
   # End of unittests WITH single card WITHOUT timeout
 
+  py_test_modules(test_clear_param_storage_api MODULES
+                  test_clear_param_storage_api)
+
 endif()
 
 py_test_modules(test_job_schedule_profiler_range MODULES

test/auto_parallel/clear_param_storage_api.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.distributed import fleet
+from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+    DygraphShardingOptimizerV2,
+)
+
+
+class TestClearParamStorage(unittest.TestCase):
+    def test_clear_param_storage(self):
+        class TestLayer(paddle.nn.Layer):
+            def __init__(self, dtype):
+                super().__init__()
+                self._w = self.create_parameter([2, 3], dtype=dtype)
+                self._b = self.create_parameter([2, 3], dtype=dtype)
+                self._w.color = {"color": "_w"}
+                self._b.color = {"color": "_b"}
+
+            @paddle.amp.debugging.check_layer_numerics
+            def forward(self, x):
+                return x * self._w + self._b
+
+        strategy = fleet.DistributedStrategy()
+        strategy.hybrid_configs = {
+            "dp_degree": 1,
+            "mp_degree": 1,
+            "pp_degree": 1,
+            "sharding_degree": 2,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+        hcg = fleet.get_hybrid_communicate_group()
+        dtype = 'float32'
+        model = TestLayer(dtype)
+
+        optimizer = paddle.optimizer.AdamW(parameters=model.parameters())
+        optimizer = DygraphShardingOptimizerV2(optimizer, hcg)
+        optimizer.clear_param_storage("_w")
+        optimizer.clear_param_storage("_b")
+        optimizer.clear_param_storage(None)
+        optimizer.reset_param_storage()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/auto_parallel/test_clear_param_storage_api.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestSemiAutoParallelMoeUtilsAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=120)
+        self._default_envs = {
+            "dtype": "float32",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+        }
+
+    def test_moe_utils(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "clear_param_storage_api.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()