adapt workflow in auto parallel

Author: waliwali777

paddleformers/cli/train/auto_parallel/workflow.py

@@ -27,12 +27,12 @@
 from paddleformers.trainer.trainer import Trainer
 from paddleformers.trainer.trainer_utils import set_seed
 from paddleformers.transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoModelForCausalLMPipe,
     AutoTokenizer,
     CosineAnnealingWithWarmupDecay,
     LinearAnnealingWithWarmupDecay,
-    LlamaConfig,
-    LlamaForCausalLMNet,
-    LlamaPretrainingCriterionNet,
 )
 from paddleformers.transformers.configuration_utils import LlmMetaConfig
 from paddleformers.utils.log import logger
@@ -145,7 +145,6 @@ def __init__(self, *args, **kwargs):
 
 
 def run_auto_parallel(model_args, data_args, generating_args, training_args):
-
     do_enable_linear_fused_grad_add = training_args.enable_linear_fused_grad_add
     # do_enable_mp_async_allreduce = (
     #     training_args.enable_auto_parallel
@@ -203,14 +202,8 @@ def run_auto_parallel(model_args, data_args, generating_args, training_args):
                 "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
             )
 
-    # TODO: only support llama model now
-    config_class = LlamaConfig
-    model_class = LlamaForCausalLMNet
-    criterion_class = LlamaPretrainingCriterionNet
-
-    config = config_class.from_pretrained(model_args.model_name_or_path)
     tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path)
-    # config = AutoConfig.from_pretrained(model_args.model_name_or_path)
+    config = AutoConfig.from_pretrained(model_args.model_name_or_path)
     LlmMetaConfig.set_llm_config(config, training_args)
     config.use_fast_layer_norm = model_args.use_fast_layer_norm
 
@@ -276,6 +269,13 @@ def run_auto_parallel(model_args, data_args, generating_args, training_args):
     if training_args.no_recompute_layers is not None:
         training_args.no_recompute_layers.sort()
 
+    if training_args.use_intermediate_api:
+        config.run_single_model = True
+        config.tensor_parallel_degree = 1
+        config.sharding_parallel_degree = 1
+        config.sep_parallel_degree = 1
+        config.context_parallel_degree = 1
+
     print("Final pre-training config:", config)
 
     # Set the dtype for loading model
@@ -286,9 +286,41 @@ def run_auto_parallel(model_args, data_args, generating_args, training_args):
         if training_args.bf16:
             dtype = "bfloat16"
 
-    with paddle.LazyGuard():
-        model = model_class.from_config(config, dtype=dtype)
-        criterion = criterion_class(config)
+    model_class = AutoModelForCausalLM
+
+    if not training_args.enable_auto_parallel and training_args.pipeline_parallel_degree > 1:
+        model_class = AutoModelForCausalLMPipe
+        if "LLama" in str(config.architectures):
+            try:
+                from utils.register_reshard import register_pp_reshard_information
+
+                register_pp_reshard_information(config.num_hidden_layers)
+            except:
+                print("Not register llama pp reshard information.")
+
+    architectures_to_check = {"Qwen2Moe", "DeepseekV2", "DeepseekV3"}
+    if (
+        any(architecture in str(config.architectures) for architecture in architectures_to_check)
+        and training_args.data_parallel_degree > 1
+    ):
+        training_args.use_expert_parallel = True
+
+    if model_args.continue_training:
+        # NOTE(gongenlei): new add
+        if training_args.autotuner_benchmark:
+            model = model_class.from_config(config, dtype=dtype)
+        else:
+            model = model_class.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                dtype=dtype,
+            )
+    else:
+        if training_args.enable_auto_parallel:
+            with paddle.LazyGuard():
+                model = model_class.from_config(config, dtype=dtype)
+        else:
+            model = model_class.from_config(config, dtype=dtype)
 
     if training_args.recompute:
 
@@ -344,7 +376,6 @@ def fn(layer):
 
     trainer = PretrainingTrainer(
         model=model,
-        criterion=criterion,
         args=training_args,
         data_collator=data_collator,
         train_dataset=train_dataset if training_args.do_train else None,

paddleformers/transformers/llama/auto_dist_config.py

@@ -12,111 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import paddle
 import paddle.distributed as dist
-from paddle.distributed.auto_parallel.intermediate.tensor_parallel import (
-    PrepareLayerInput,
-    PrepareLayerOutput,
-)
-
-
-def layer_input_parallel_row_hook(process_mesh):
-    def hook(layer, inputs, output=None):
-        res_inputs = []
-        for input in inputs:
-            if not input.is_dist():
-                x = dist.shard_tensor(input, process_mesh, [dist.Shard(0), dist.Replicate(), dist.Replicate()])
-                res_inputs.append(dist.reshard(x, process_mesh, [dist.Shard(0), dist.Replicate(), dist.Replicate()]))
-            else:
-                res_inputs.append(
-                    dist.reshard(input, process_mesh, [dist.Shard(0), dist.Replicate(), dist.Replicate()])
-                )
-        return tuple(res_inputs)
-
-    return hook
-
-
-def layer_input_parallel_row_and_col_hook(process_mesh):
-    def hook(layer, inputs, output=None):
-        res_inputs = []
-        for input in inputs:
-            if not input.is_dist():
-                x = dist.shard_tensor(input, process_mesh, [dist.Shard(0), dist.Replicate(), dist.Shard(1)])
-                res_inputs.append(dist.reshard(x, process_mesh, [dist.Shard(0), dist.Replicate(), dist.Shard(1)]))
-            else:
-                res_inputs.append(dist.reshard(input, process_mesh, [dist.Shard(0), dist.Replicate(), dist.Shard(1)]))
-        return tuple(res_inputs)
-
-    return hook
-
-
-def layer_input_replicate_hook(process_mesh):
-    def hook(layer, inputs, output=None):
-        res_inputs = []
-        for input in inputs:
-            if not input.is_dist():
-                x = dist.shard_tensor(input, process_mesh, [dist.Replicate(), dist.Replicate(), dist.Replicate()])
-                res_inputs.append(
-                    dist.reshard(x, process_mesh, [dist.Replicate(), dist.Replicate(), dist.Replicate()])
-                )
-            else:
-                res_inputs.append(dist.reshard(input, process_mesh, [dist.Replicate(), dist.Replicate()]))
-        return tuple(res_inputs)
-
-    return hook
-
-
-def layer_input_rope_hook(process_mesh):
-    def hook(layer, inputs, output=None):
-        res_inputs = []
-        batch_size = None
-        seq_length = None
-        process_mesh = None
-        placements = None
-        for index in range(len(inputs)):
-            if index == 0:
-                batch_size, seq_length, _, _ = inputs[index]._local_shape
-                process_mesh = inputs[index].process_mesh
-                placements = inputs[index].placements
-            # process position_ids
-            if index == len(inputs) - 1:
-                mesh = dist.auto_parallel.get_mesh()
-                assert "sep" in mesh.dim_names, f"mesh.dim_names:{mesh.dim_names} must contain sep"
-                group = mesh._get_group("sep")
-                chunk_size = seq_length // 2
-                chunk_num = group.nranks * 2
-                rank = group.rank
-                first_chunk_ids = paddle.arange(rank * chunk_size, (rank + 1) * chunk_size, dtype="int64")
-                second_chunk_ids = paddle.arange(
-                    (chunk_num - rank - 1) * chunk_size, (chunk_num - rank) * chunk_size, dtype="int64"
-                )
-                position_ids = paddle.concat([first_chunk_ids, second_chunk_ids]).expand((batch_size, seq_length))
-                mp_axis = process_mesh.dim_names.index("mp")
-                placements[mp_axis] = dist.Replicate()  # mp placament shard(2) -> replicate
-                position_ids = dist.auto_parallel.api.dtensor_from_local(position_ids, process_mesh, placements)
-                res_inputs.append(position_ids)
-            else:
-                res_inputs.append(inputs[index])
-        return tuple(res_inputs)
-
-    return hook
-
-
-def layer_output_rope_hook(process_mesh):
-    def hook(layer, inputs, outputs):
-        res_outputs = []
-        for output in outputs:
-            process_mesh = output.process_mesh
-            placements = output.placements
-            cp_index = process_mesh.dim_names.index("sep")  # get the axis for the split
-            cp_degree = process_mesh.shape[cp_index]
-            assert cp_degree > 1, f"cp_degree:{cp_degree} must > 1"
-            placements[cp_index] = dist.Shard(1)  # seq_dim:1
-            output = dist.reshard(output, process_mesh, placements)
-            res_outputs.append(output)
-        return tuple(res_outputs)
-
-    return hook
 
 
 def get_dist_config(model, prefix=""):
@@ -125,36 +21,9 @@ def get_dist_config(model, prefix=""):
         assert prefix.endswith(".")
 
     config = {
-        "sp_config": {
-            "parallelize_plan": {
-                f"{prefix}llama.embed_tokens": [
-                    dist.ColWiseParallel(),
-                    dist.SequenceParallelBegin(),
-                ],
-                f"{prefix}llama.reshard_row": PrepareLayerInput(layer_input_parallel_row_hook),
-                f"{prefix}llama.reshard_row_and_col": PrepareLayerInput(layer_input_parallel_row_and_col_hook),
-                f"{prefix}llama.global_layer.reshard_replicate": PrepareLayerInput(layer_input_replicate_hook),
-                f"{prefix}llama.layers.*.self_attn.qkv_proj": dist.ColWiseParallel(),
-                f"{prefix}llama.layers.*.self_attn.q_proj": dist.ColWiseParallel(),
-                f"{prefix}llama.layers.*.self_attn.k_proj": dist.ColWiseParallel(),
-                f"{prefix}llama.layers.*.self_attn.v_proj": dist.ColWiseParallel(),
-                f"{prefix}llama.layers.*.self_attn.o_proj": dist.RowWiseParallel(),
-                f"{prefix}llama.layers.*.self_attn": dist.SequenceParallelDisable(),
-                f"{prefix}llama.layers.*.mlp.gate_proj": dist.ColWiseParallel(),
-                f"{prefix}llama.layers.*.mlp.up_proj": dist.ColWiseParallel(),
-                f"{prefix}llama.layers.*.mlp.gate_up_fused_proj": dist.ColWiseParallel(),
-                f"{prefix}llama.layers.*.mlp.down_proj": dist.RowWiseParallel(),
-                f"{prefix}llama.layers.*.mlp": dist.SequenceParallelDisable(need_transpose=False),
-                f"{prefix}lm_head.weight": dist.ColWiseParallel(),
-                f"{prefix}lm_head": dist.SequenceParallelEnd(),
-            }
-        },
         "mp_config": {
             "parallelize_plan": {
                 f"{prefix}llama.embed_tokens": dist.ColWiseParallel(gather_output=True),
-                f"{prefix}llama.reshard_row": PrepareLayerInput(layer_input_parallel_row_hook),
-                f"{prefix}llama.reshard_row_and_col": PrepareLayerInput(layer_input_parallel_row_and_col_hook),
-                f"{prefix}llama.global_layer.reshard_replicate": PrepareLayerInput(layer_input_replicate_hook),
                 f"{prefix}llama.layers.*.self_attn.qkv_proj": dist.ColWiseParallel(),
                 f"{prefix}llama.layers.*.self_attn.q_proj": dist.ColWiseParallel(),
                 f"{prefix}llama.layers.*.self_attn.k_proj": dist.ColWiseParallel(),
@@ -167,31 +36,5 @@ def get_dist_config(model, prefix=""):
                 f"{prefix}lm_head.weight": dist.ColWiseParallel(),
             }
         },
-        "pp_config": {"split_spec": f"{prefix}llama.layers", "global_spec": f"{prefix}llama.global_layer"},
-        "cp_config": {
-            "parallelize_plan": {
-                f"{prefix}llama.layers.*.self_attn.sdpa": dist.ContextParallel(
-                    backend="p2p" if model.config.context_parallel_degree > 1 else "all2all"
-                ),
-            }
-        },
     }
-
-    if model.config.context_parallel_degree > 1:
-        config["cp_config"]["parallelize_plan"].update(
-            {
-                f"{prefix}llama.layers.*.self_attn.rope_func": [
-                    PrepareLayerInput(layer_input_rope_hook),
-                    PrepareLayerOutput(layer_output_rope_hook),
-                ]
-            }
-        )
-    elif model.config.sep_parallel_degree > 1:
-        # fuse_rope is not support dtensor spmd yet,thus need to extraly reshard sequence dim
-        config["cp_config"]["parallelize_plan"].update(
-            {
-                f"{prefix}llama.layers.*.self_attn.rope_func": PrepareLayerOutput(layer_output_rope_hook),
-            }
-        )
-
     return config

paddleformers/transformers/llama/modeling.py

@@ -1355,6 +1355,7 @@ def _get_name_mappings(cls, config: LlamaConfig) -> list[StateDictNameMapping]:
 
     @classmethod
     def _get_tensor_parallel_mappings(cls, config: LlamaConfig, is_split=True):
+
         from ..conversion_utils import split_or_merge_func
 
         fn = split_or_merge_func(