no model for parallelize_optimizer (#69677)

FeixLiu · web-flow · commit c112fc860880 · 2024-11-27T13:40:59.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/intermediate/parallel_base.py b/python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
@@ -37,14 +37,21 @@ def is_tensor(tensor):
 
 
 class ParallelOptimizer:
-    def __init__(self, optimizer, level=None):
+    def __init__(
+        self,
+        optimizer,
+        level=None,
+        sharding_mesh_dim=None,
+    ):
         self.level = None
+        self.sharding_mesh_dim = None
         self.optimizer = None
 
         if isinstance(optimizer, ParallelOptimizer):
             self.optimizer = optimizer.optimizer
             if level is None:
                 self.level = optimizer.level
+                self.sharding_mesh_dim = optimizer.sharding_mesh_dim
             else:
                 if isinstance(level, int):
                     level = str(level)
@@ -54,6 +61,7 @@ def __init__(self, optimizer, level=None):
                         level == optimizer.level
                     ), f"The level passed in is not identical with previous level. Current level is {level}, previous level is {optimizer.level}"
                 self.level = level
+                self.sharding_mesh_dim = sharding_mesh_dim
         else:
             assert isinstance(optimizer, Optimizer)
             self.optimizer = optimizer
@@ -62,43 +70,46 @@ def __init__(self, optimizer, level=None):
             assert level in ("0", "1", "2", "3", None)
             # level=0 and level=None are all mean pure dp
             self.level = level
+            self.sharding_mesh_dim = sharding_mesh_dim
 
         self.is_initialized = False
 
-    def parallelize(self, parallelized_parameters):
+    def parallelize(self):
         assert self.optimizer is not None
         if self.is_initialized:
             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 == "1":
             self.optimizer = dist.shard_optimizer(
-                self.optimizer, dist.ShardingStage1("dp", mesh)
+                self.optimizer,
+                dist.ShardingStage1(self.sharding_mesh_dim, mesh),
             )
         elif self.level == "2":
             self.optimizer = dist.shard_optimizer(
-                self.optimizer, dist.ShardingStage2("dp", mesh)
+                self.optimizer,
+                dist.ShardingStage2(self.sharding_mesh_dim, mesh),
             )
         elif self.level == "3":
             self.optimizer = dist.shard_optimizer(
-                self.optimizer, dist.ShardingStage3("dp", mesh)
+                self.optimizer,
+                dist.ShardingStage3(self.sharding_mesh_dim, mesh),
             )
         else:
-            self.optimizer = dist.shard_optimizer(self.optimizer)
+            self.optimizer = dist.shard_optimizer(self.optimizer, None)
         self.is_initialized = True
 
         return self.optimizer
 
+    def update_param_list(self, parallelized_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
+
 
 class ParallelModel:
     def __init__(self, model):
@@ -192,8 +203,7 @@ def parallelize_model_and_optimizer(model, optimizer=None):
     parallelized_optimizer = None
     if optimizer is not None:
         assert isinstance(optimizer, ParallelOptimizer)
-        parallelized_optimizer = optimizer.parallelize(
-            parallelized_model.parameters()
-        )
+        optimizer.update_param_list(parallelized_model.parameters())
+        parallelized_optimizer = optimizer.parallelize()
 
     return parallelized_model, parallelized_optimizer
diff --git a/python/paddle/distributed/auto_parallel/intermediate/parallelize.py b/python/paddle/distributed/auto_parallel/intermediate/parallelize.py
@@ -11,6 +11,7 @@
 # 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 warnings
 
 from .parallel_base import ParallelOptimizer, parallelize_model_and_optimizer
 from .pipeline_parallel import pipeline_parallel
@@ -34,30 +35,60 @@ def parallelize(
         )
     if dp_config is not None:
         assert isinstance(dp_config, dict)
+        if 'sharding_level' not in dp_config.keys():
+            warnings.warn(
+                "The dp_config doesn't contain sharding_level, will run under dp."
+            )
         model, optimizer = sharded_data_parallel(
             model,
             optimizer,
-            level=dp_config.get('sharding_level'),
-            offload=bool(dp_config.get('offload')),
-            exclude_layer=dp_config.get('exclude_layer'),
+            config=dp_config,
         )
     model, optimizer = parallelize_model_and_optimizer(model, optimizer)
     return model, optimizer
 
 
+has_parallelized_model = False
+
+
 def parallelize_model(
     model, mesh=None, dp_config=None, mp_config=None, pp_config=None
 ):
+    global has_parallelized_model
+    has_parallelized_model = True
     model, _ = parallelize(model, None, mesh, dp_config, mp_config, pp_config)
     return model
 
 
 def parallelize_optimizer(
-    model, optimizer, mesh=None, dp_config=None, mp_config=None, pp_config=None
+    optimizer, mesh=None, dp_config=None, mp_config=None, pp_config=None
 ):
+    global has_parallelized_model
+    assert (
+        has_parallelized_model
+    ), "Please parallelize the model before parallelize optimizer."
+    param_list = optimizer._parameter_list
+    if isinstance(param_list[0], dict):
+        for param_group in param_list:
+            for param in param_group['params']:
+                assert (
+                    param.is_dist()
+                ), "Please use model after parallelize to create optimizer."
+    else:
+        for param in param_list:
+            assert (
+                param.is_dist()
+            ), "Please use model after parallelize to create optimizer."
+
     level = None
+    sharding_mesh_dim = None
     if dp_config is not None:
+        if 'sharding_level' not in dp_config.keys():
+            warnings.warn(
+                "The dp_config doesn't contain sharding_level, will run under dp."
+            )
         level = dp_config.get('sharding_level')
-    optimizer = ParallelOptimizer(optimizer, level)
-    optimizer = optimizer.parallelize(model.parameters())
+        sharding_mesh_dim = dp_config.get('sharding_mesh_dim', "dp")
+    optimizer = ParallelOptimizer(optimizer, level, sharding_mesh_dim)
+    optimizer = optimizer.parallelize()
     return optimizer
diff --git a/python/paddle/distributed/auto_parallel/intermediate/sharded_data_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/sharded_data_parallel.py
@@ -51,31 +51,31 @@ def sharding_parallelizer_func(self, model):
         return model
 
 
-def sharded_data_parallel(
-    model, optimizer=None, level=None, offload=False, exclude_layer=None
-):
+def sharded_data_parallel(model, optimizer=None, config=None):
     """
     sharded_data_parallel converts model and optimizer to distributed and supports set zero stage1/2/3
 
     Args:
         model (paddle.nn.Layer): A single card model to be distributed
         optimizer (paddle.optimizer.Optimizer): an optimizer to be distributed
-        level (str): Zero stage, can be the following values:
-            0: no sharding (pure dp)
-            1: Zero Stage1
-            2: Zero Stage2
-            3: Zero Stage3
-            Default: None, which means optimizer is replicated among all process.
-        offload (bool): whether enable cpu offload strategy, not implemented currently.
-        exclude_layer (list): Specify which layers do not use the zero stage strategy, not implemented currently.
+        config (dict): {
+            "sharding_level": 0,
+            "offload": False,
+            "exclude_layer": None,
+            "sharding_mesh_dim": "dp",
+        }
 
     Returns:
         ShardedDataParallel: a distributed model
         ParallelOptimizer: a distributed optimizer
     """
-    sdp_model = ShardedDataParallel(model, offload, exclude_layer)
+    sdp_model = ShardedDataParallel(
+        model, bool(config.get('offload')), config.get('exclude_layer')
+    )
     if optimizer is not None:
-        optimizer = ParallelOptimizer(optimizer, level)
+        level = config.get('sharding_level')
+        sharding_mesh_dim = config.get('sharding_mesh_dim', "dp")
+        optimizer = ParallelOptimizer(optimizer, level, sharding_mesh_dim)
 
     # check global_mesh
     mesh = fleet.auto.get_mesh()
diff --git a/test/auto_parallel/hybrid_strategy/parallel_api.py b/test/auto_parallel/hybrid_strategy/parallel_api.py
@@ -200,7 +200,7 @@ def check_mp(self, layer):
                         dist.Shard(0),
                     ]
 
-    def parallel_model(self, layer, optimizer=None):
+    def parallel_model(self, layer):
         dp_config = None
         mp_config = None
         pp_config = None
@@ -275,17 +275,8 @@ def parallel_model(self, layer, optimizer=None):
             mp_config=mp_config,
             pp_config=pp_config,
         )
-        optimizer = parallelize_optimizer(
-            layer,
-            optimizer,
-            dp_config=dp_config,
-            mp_config=mp_config,
-            pp_config=pp_config,
-        )
         self.check_mp(layer)
-        if optimizer is None:
-            return layer
-        return layer, optimizer
+        return layer, dp_config, mp_config, pp_config
 
     def run_llama(
         self, share_embedding=False, position_embedding=False, to_static=0
@@ -300,11 +291,19 @@ def run_llama(
                 self.config, share_embedding, position_embedding
             )
 
+        model, dp_config, mp_config, pp_config = self.parallel_model(model)
+
         lr_scheduler = paddle.optimizer.lr.LinearWarmup(
             learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
         )
         optimizer = create_optimizer(model, lr_scheduler)
-        model, optimizer = self.parallel_model(model, optimizer)
+
+        optimizer = parallelize_optimizer(
+            optimizer,
+            dp_config=dp_config,
+            mp_config=mp_config,
+            pp_config=pp_config,
+        )
 
         criterion = LlamaPretrainingCriterion(self.config)
 
