[AutoParallel]: fix shard optimizer (#69283)

blacksheep-Aristotle · web-flow · commit e2bdc5d456eb · 2024-11-12T17:05:55.000+08:00
* [AutoParallel]: fix shard optimizer

* [AutoParallel]: fix shard optimizer

* [AutoParallel]: fix shard optimizer
diff --git a/python/paddle/distributed/auto_parallel/intermediate/parallel_base.py b/python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
@@ -40,7 +40,12 @@ def parallelize(self, parallelized_parameters):
             return self.optimizer
         # 1.replace optimizer parameters
         self.optimizer._parameter_list = parallelized_parameters
-
+        if isinstance(parallelized_parameters[0], dict):
+            self.optimizer._param_groups = []
+            for param_group in self.parallelized_parameters:
+                self.optimizer._add_param_group(param_group.copy())
+        else:
+            self.optimizer._param_groups = self.optimizer._parameter_list
         # 2.wrap with shard_optimizer
         mesh = fleet.auto.get_mesh()
         if self.level == "os":
