Merge branch 'PaddlePaddle:develop' into dev_20250126_add_pipeline_for_moe

Author: DrownFish19

paddlenlp/trainer/auto_trainer.py

@@ -713,8 +713,13 @@ def _save_checkpoint(self, model, metrics=None):
                         for key, value in model.state_dict("opt").items()
                         if not any(keyword in key for keyword in FREE_SVAE_LOAD_KEY_PATTERNS)
                     }
+                    model_state_dict = model.state_dict("param")
+                    if self.args.should_save_model_with_tensor_fusion:
+                        model_state_dict = self._convert_state_dict_for_saving_tensor_fusion_ckpt(model_state_dict)
+                        opt_state_dict = self._convert_state_dict_for_saving_tensor_fusion_ckpt(opt_state_dict)
+
                     state_dict = {
-                        MODEL_NAME: model.state_dict("param"),
+                        MODEL_NAME: model_state_dict,
                         OPTIMIZER_NAME: opt_state_dict,
                     }
                 else:
@@ -854,6 +859,9 @@ def _load_from_checkpoint(self, resume_from_checkpoint=None):
                     for key, value in self.model_wrapped.state_dict("opt").items()
                     if not any(keyword in key for keyword in FREE_SVAE_LOAD_KEY_PATTERNS)
                 }
+                if self.args.should_load_model_with_tensor_fusion:
+                    model_state_dict = self._convert_state_dict_for_loading_tensor_fusion_ckpt(model_state_dict)
+                    optim_state_dict = self._convert_state_dict_for_loading_tensor_fusion_ckpt(optim_state_dict)
             else:
                 model_state_dict = self.model_wrapped.state_dict()
                 optim_state_dict = self.optimizer.state_dict()
@@ -888,7 +896,36 @@ def _load_from_checkpoint(self, resume_from_checkpoint=None):
                 self._load_ckpt_func(state_dict, ckpt_path)
 
             if self.args.to_static:
+                if self.args.should_load_model_with_tensor_fusion:
+                    model_state_dict = self._convert_state_dict_for_loading_model_with_tensor_fusion(model_state_dict)
+                    optim_state_dict = self._convert_state_dict_for_loading_model_with_tensor_fusion(optim_state_dict)
+
                 self.model_wrapped.set_state_dict(model_state_dict)
                 self.model_wrapped.set_state_dict(optim_state_dict)
             # release memory
             del state_dict
+
+    def _convert_state_dict_for_loading_tensor_fusion_ckpt(self, state_dict):
+        if self.args.load_model_with_sharding_tensor_fusion:
+            logger.info("load sharding tensor fusion unbalanced model")
+            state_dict = self.model_wrapped._convert_state_dict_with_rank_unique_name(state_dict)
+        else:
+            logger.info("load sharding tensor fusion balanced model")
+            state_dict = self.model_wrapped._convert_state_dict_without_tensor_fusion_param(state_dict)
+        return state_dict
+
+    def _convert_state_dict_for_loading_model_with_tensor_fusion(self, state_dict):
+        if self.args.load_model_with_sharding_tensor_fusion:
+            state_dict = self.model_wrapped._convert_state_dict_with_origin_name(state_dict)
+        else:
+            state_dict = self.model_wrapped._convert_state_dict_with_tensor_fusion_param(state_dict)
+        return state_dict
+
+    def _convert_state_dict_for_saving_tensor_fusion_ckpt(self, state_dict):
+        if self.args.save_model_with_sharding_tensor_fusion:
+            logger.info("save sharding tensor fusion unbalanced model")
+            state_dict = self.model_wrapped._convert_state_dict_with_rank_unique_name(state_dict)
+        else:
+            logger.info("save sharding tensor fusion balanced model")
+            state_dict = self.model_wrapped._convert_state_dict_without_tensor_fusion_param(state_dict)
+        return state_dict

paddlenlp/trainer/auto_training_args.py

@@ -14,7 +14,7 @@
 import json
 from dataclasses import dataclass, field
 
-from .trainer_utils import split_parallel_config
+from .trainer_utils import ShardingOption, split_parallel_config
 from .training_args import TrainingArguments
 from .utils import add_start_docstrings
 
@@ -52,6 +52,29 @@ class AutoTrainingArguments(TrainingArguments):
         metadata={"help": "Weather to use auto_parallel intermediate api"},
     )
     refined_ops_patterns: str = field(default=None, metadata={"help": "The pattern of refined recompute."})
+    load_model_with_sharding_tensor_fusion: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "When using sharding stage1, enabling tensor fusion, and setting `load_model_with_sharding_tensor_fusion` to `True`, "
+                "the model is loaded with unbalanced weights, meaning that the model weights are stored in an unbalanced format to avoid "
+                "additional memory overhead. If set to `False`, the model will be loaded with balanced weights, which may increase memory "
+                "consumption. This setting is only available in auto parallel to_static mode."
+            )
+        },
+    )
+    save_model_with_sharding_tensor_fusion: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "When using sharding stage1 and enabling tensor fusion, setting `save_model_with_sharding_tensor_fusion` to `True` "
+                "saves the model with unbalanced weights, which helps avoid additional memory consumption. Setting it to `False` "
+                "saves the model with balanced weights, which may increase memory usage but ensures uniform parameter distribution. "
+                "This option allows flexibility in choosing the save format based on memory requirements. "
+                "This setting is only available in auto parallel to_static mode."
+            )
+        },
+    )
 
     def __post_init__(self):
         super().__post_init__()
@@ -89,3 +112,13 @@ def __post_init__(self):
                 recompute.refined_ops_patterns = (
                     self.refined_ops_patterns if self.refined_ops_patterns is not None else []
                 )
+
+    @property
+    def should_load_model_with_tensor_fusion(self):
+        return (
+            self.enable_auto_parallel
+            and self.to_static
+            and ShardingOption.SHARD_OP in self.sharding
+            and self.sharding_parallel_degree > 1
+            and "enable_tensor_fusion" in self.sharding_parallel_config
+        )

paddlenlp/trainer/training_args.py

@@ -619,6 +619,7 @@ class TrainingArguments:
             )
         },
     )
+
     tensor_parallel_degree: int = field(
         default=-1,
         metadata={
@@ -740,7 +741,6 @@ class TrainingArguments:
                 "enable_stage2_overlap, overlap stage2 NCCL communication with computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for broadcast overlap and no other sync could be called during the training for broadcast overlap\n"
                 "enable_stage1_broadcast_overlap, overlap stage1 V1 broadcast with next step forward computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for broadcast overlap forward compute and no other sync could be called during the training for broadcast overlap.\n"
                 "enable_stage1_allgather_overlap, overlap stage1 V2 allgather with next step forward computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for allgather overlap forward compute and no other sync could be called during the training for allgather overlap.\n"
-                "enable_tensor_fusion_blanced_save_load, convert unbalanced optimizer state to balanced state when using tensor fusion strategy, which may increase the memory occupation."
             )
         },
     )
@@ -1671,7 +1671,6 @@ def is_segment_parallel_supported():
                             "enable_tensor_fusion",
                             "enable_overlap",
                             "enable_release_grads",
-                            "enable_tensor_fusion_blanced_save_load",
                         ]:
                             if x in ["enable_stage1_overlap", "enable_stage2_overlap"]:
                                 raise ValueError(
@@ -1686,7 +1685,7 @@ def is_segment_parallel_supported():
                             raise ValueError(
                                 f"Found unknown sharding mode config {x}, "
                                 f"accpet config is enable_tensor_fusion, "
-                                "enable_overlap, enable_release_grads, enable_tensor_fusion_blanced_save_load."
+                                "enable_overlap, enable_release_grads."
                             )
 
                     if "enable_overlap" in sharding_parallel_config:
@@ -1696,9 +1695,6 @@ def is_segment_parallel_supported():
                         sharding.grad_bucket_size_numel = 210355872
                         sharding.enable_tensor_fusion = True
 
-                    if "enable_tensor_fusion_blanced_save_load" in sharding_parallel_config:
-                        sharding.save_unbalanced_param = False
-
                     if "enable_release_grads" in sharding_parallel_config:
                         sharding.release_gradients = True
 
@@ -2273,3 +2269,13 @@ def print_config(self, args=None, key=""):
                     logger.debug("{:30}: {}".format(a, v))
 
         logger.debug("")
+
+    @property
+    def should_save_model_with_tensor_fusion(self):
+        return (
+            self.enable_auto_parallel
+            and self.to_static
+            and ShardingOption.SHARD_OP in self.sharding
+            and self.sharding_parallel_degree > 1
+            and "enable_tensor_fusion" in self.sharding_parallel_config
+        )