bug fix for sharding

Author: FeixLiu

paddlenlp/trainer/auto_trainer.py

@@ -90,6 +90,7 @@ def loss_func(loss, outputs):
                 "pipeline_parallel": kwargs["args"].pipeline_parallel_degree > 1,
                 "data_sharding_parallel": kwargs["args"].dataset_world_size > 1,
                 "sharding": kwargs["args"].sharding,
+                "sharding_mesh_dim": kwargs["args"].sharding_parallel_mesh_dimension,
             }
             auto_dist_config = model._generate_auto_dist_config(auto_dist_degree)
             self.auto_dist_config = auto_dist_config
@@ -164,7 +165,6 @@ def _wrap_for_auto(self, model, train_dataloader):
         if self.args.use_intermediate_api:
             assert self.auto_dist_config is not None
             self.optimizer = parallelize_optimizer(
-                model,
                 self.optimizer,
                 dp_config=self.auto_dist_config["dp_config"],
                 mp_config=self.auto_dist_config["mp_config"],

paddlenlp/transformers/model_utils.py

@@ -2826,7 +2826,10 @@ def _generate_auto_dist_config(self, auto_dist_degree):
                     level = 2
                 if ShardingOption.FULL_SHARD in sharding:
                     level = 3
-            final_config["dp_config"] = {"level": level}
+            final_config["dp_config"] = {
+                "sharding_level": level,
+                "sharding_mesh_dim": auto_dist_degree.get("sharding_mesh_dim", None),
+            }
 
         return final_config