glm45 suport pipeline parallel (#3082)

Co-authored-by: YuangLiu <liuyuang@baidu.com>

Author: LiYuRio

examples/experiments/paddlefleet/glm45_provider.py

@@ -17,29 +17,20 @@
 
 import logging
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Callable, List, Optional, Union
+from typing import Callable, List, Optional, Union
 
 import paddle
 import paddle.nn.functional as F
-from paddlefleet.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
 
 from paddleformers.transformers.gpt_provider import GPTModelProvider
 
-if TYPE_CHECKING:
-    from paddlefleet.spec_utils import LayerSpec
-
-
 logger = logging.getLogger(__name__)
 
 
 @dataclass
 class GLMMoEModelProvider(GPTModelProvider):
     """Base provider for GLM MoE Models."""
 
-    transformer_layer_spec: Union[
-        "LayerSpec", Callable[["GPTModelProvider"], "LayerSpec"]
-    ] = get_gpt_decoder_block_spec
-
     normalization: str = "RMSNorm"
     hidden_act: Callable = F.silu
     gated_linear_unit: bool = True

examples/experiments/paddlefleet/qwen_provider.py

@@ -17,29 +17,20 @@
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Optional, Union
+from typing import Callable, Optional
 
 import paddle
 import paddle.nn.functional as F
-from paddlefleet.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
 
 from paddleformers.transformers.gpt_provider import GPTModelProvider
 
-if TYPE_CHECKING:
-    from paddlefleet.spec_utils import LayerSpec
-
-
 logger = logging.getLogger(__name__)
 
 
 @dataclass
 class Qwen3MoEModelProvider(GPTModelProvider):
     """Base provider for Qwen 3 MoE Models."""
 
-    transformer_layer_spec: Union[
-        "LayerSpec", Callable[["GPTModelProvider"], "LayerSpec"]
-    ] = get_gpt_decoder_block_spec
-
     normalization: str = "RMSNorm"
     hidden_act: Callable = F.silu
     gated_linear_unit: bool = True

examples/experiments/paddlefleet/run_pretrain.py

@@ -40,12 +40,7 @@
     speed_metrics,
 )
 from paddleformers.trainer.trainer import Trainer
-from paddleformers.transformers import (
-    AutoConfig,
-    AutoTokenizer,
-    CosineAnnealingWithWarmupDecay,
-    LinearAnnealingWithWarmupDecay,
-)
+from paddleformers.transformers import AutoConfig, AutoTokenizer
 from paddleformers.transformers.configuration_utils import LlmMetaConfig, llmmetaclass
 from paddleformers.utils.batch_sampler import DistributedBatchSampler
 from paddleformers.utils.log import logger
@@ -522,11 +517,6 @@ def main():
     if training_args.decay_steps is None:
         training_args.decay_steps = training_args.max_steps
 
-    if training_args.warmup_steps > 0:
-        warmup_steps = training_args.warmup_steps
-    else:
-        warmup_steps = training_args.warmup_ratio * training_args.max_steps
-
     lr_scheduler = None
 
     data_file = get_train_data_file(data_args)

paddleformers/trainer/trainer.py

@@ -61,6 +61,15 @@
     from paddle.base import core
 except:
     core = None
+try:
+    import paddlefleet.distributed.model as paddlefleet_dist_model
+    from paddlefleet.pipeline_parallel import ParallelBase as PaddleFleetParallelBase
+    from paddlefleet.pipeline_parallel import PipelineLayer as PaddleFleetPipelineLayer
+
+    HAS_PADDLEFLEET = True
+except:
+    HAS_PADDLEFLEET = False
+
 from paddle.distributed import fleet
 from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.hybrid_parallel_optimizer import (
     HybridParallelOptimizer,
@@ -2999,6 +3008,47 @@ def _wrap_model(self, model, training=True):
 
             return model
 
+        if HAS_PADDLEFLEET and isinstance(model, PaddleFleetPipelineLayer):
+            prepare_pipeline_inputs_func = (
+                model._prepare_pipeline_inputs_func if hasattr(model, "_prepare_pipeline_inputs_func") else None
+            )
+            model = paddlefleet_dist_model.distributed_model(model)
+            if prepare_pipeline_inputs_func is not None:
+                model._prepare_pipeline_inputs_func = prepare_pipeline_inputs_func
+            else:
+
+                def _prepare_pipeline_inputs_func(inputs):
+                    first_stage_keys = ["input_ids", "attention_mask", "position_ids"]
+                    last_stage_keys = ["labels"]
+
+                    def get_expected_keys(inputs, keys):
+                        ret = tuple([inputs.pop(k) for k in keys if k in inputs])
+                        if len(ret) == 1:
+                            ret = ret[0]
+                        return ret
+
+                    if type(inputs) is dict or type(inputs) is OrderedDict:
+                        return [
+                            get_expected_keys(inputs, first_stage_keys),
+                            get_expected_keys(inputs, last_stage_keys),
+                        ]
+
+                    keys = list(inputs[0].keys())
+                    inputs_batch = {key: [data.pop(key) for data in inputs] for key in keys}
+                    first_stage_inputs_batch = inputs_batch
+                    last_stage_inputs = first_stage_inputs_batch.pop("labels")
+                    outputs = (
+                        first_stage_inputs_batch,
+                        last_stage_inputs,
+                    )
+                    return outputs
+
+                logger.warning(
+                    "Using default prepare pipeline inputs func, only support input_ids and labels as inputs."
+                )
+                model._prepare_pipeline_inputs_func = _prepare_pipeline_inputs_func
+            return model
+
         # train/eval could be run multiple-times - if already wrapped, don't re-wrap it again
         if unwrap_model(model) is not model:
             return model
@@ -3047,7 +3097,10 @@ def _wrap_model(self, model, training=True):
                 assert self.optimizer is not None, "optimizer is empty!"
                 self.optimizer = mix_precision_utils.MixPrecisionOptimizer(self.optimizer)
 
-        in_pipeline_parallel_mode = self.args.pipeline_parallel_degree > 1
+        if HAS_PADDLEFLEET and isinstance(model, PaddleFleetParallelBase):
+            in_pipeline_parallel_mode = True
+        else:
+            in_pipeline_parallel_mode = self.args.pipeline_parallel_degree > 1
         in_sharding_parallel_mode = self.sharding is not None
         in_tensor_parallel_mode = self.args.tensor_parallel_degree > 1
         in_sep_parallel_mode = self.args.sep_parallel_degree > 1
@@ -3382,6 +3435,9 @@ def training_step(
         Return:
             `paddle.Tensor`: The tensor with training loss on this batch.
         """
+        if HAS_PADDLEFLEET and isinstance(model, PaddleFleetParallelBase):
+            return self.training_pipeline_step(model, inputs)
+
         if self.args.pipeline_parallel_degree > 1:
             return self.training_pipeline_step(model, inputs)
 

paddleformers/transformers/glm4_moe/modeling.py

@@ -15,7 +15,7 @@
 from copy import deepcopy
 from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import paddle
 import paddle.distributed as dist
@@ -24,7 +24,6 @@
 from paddle.distributed.fleet.utils import recompute
 from paddle.distributed.fleet.utils.sequence_parallel_utils import GatherOp, ScatterOp
 from paddle.nn import functional as F
-from paddlefleet.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
 
 from paddleformers.transformers.gpt_provider import GPTModelProvider
 
@@ -48,18 +47,11 @@
 from ..moe_layer import MoEFlexTokenLayer
 from .configuration import Glm4MoeConfig
 
-if TYPE_CHECKING:
-    from paddlefleet.transformer import LayerSpec
-
 
 @dataclass
 class GLMMoEModelProvider(GPTModelProvider):
     """Base provider for GLM MoE Models."""
 
-    transformer_layer_spec: Union[
-        "LayerSpec", Callable[["GPTModelProvider"], "LayerSpec"]
-    ] = get_gpt_decoder_block_spec
-
     moe_router_load_balancing_type: str = "seq_aux_loss"
 
     gated_linear_unit: bool = True

paddleformers/transformers/gpt_provider.py

@@ -23,15 +23,28 @@
 from typing import Any, Callable, Literal, Optional, Union
 
 import paddle
-from paddlefleet import LayerSpec, parallel_state
+from paddlefleet import LayerSpec
 from paddlefleet.models.gpt import GPTModel as FleetGPTModel
 from paddlefleet.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
-from paddlefleet.transformer.transformer_config import TransformerConfig
+
+try:
+    from paddlefleet.models.gpt.gpt_config import GPTConfig
+except ImportError:
+    from paddlefleet.transformer.transformer_config import (
+        TransformerConfig as GPTConfig,
+    )
+
+
+try:
+    from paddlefleet.gpt_builders import gpt_builder
+
+    HAS_PADDLEFLEET = True
+except ImportError:
+    HAS_PADDLEFLEET = False
 
 from paddleformers.transformers.model_utils import PretrainedModel
 
 from .model_provider import ModelProviderMixin
-from .vocab_utils import calculate_padded_vocab_size
 
 logger = logging.getLogger(__name__)
 
@@ -52,6 +65,7 @@ def local_layer_spec(config: "GPTModelProvider") -> LayerSpec:
     Returns:
         LayerSpec: Module specification for local implementation layers
     """
+    assert HAS_PADDLEFLEET
     return get_gpt_layer_local_spec(
         num_experts=config.num_moe_experts,
         moe_grouped_gemm=config.moe_grouped_gemm,
@@ -61,7 +75,7 @@ def local_layer_spec(config: "GPTModelProvider") -> LayerSpec:
 
 
 @dataclass
-class GPTModelProvider(TransformerConfig, ModelProviderMixin[GPTModel]):
+class GPTModelProvider(GPTConfig, ModelProviderMixin[GPTModel]):
     """Configuration and provider for PaddleFleet GPT models.
 
     This class extends TransformerConfig with GPT-specific parameters and
@@ -78,15 +92,16 @@ class GPTModelProvider(TransformerConfig, ModelProviderMixin[GPTModel]):
     rotary_percent: float = 1.0
     seq_len_interpolation_factor: Optional[float] = None
     seq_length: int = 1024
+
+    max_sequence_length: int = 1024
+
     attention_softmax_in_fp32: bool = False
     deallocate_pipeline_outputs: bool = True
     scatter_embedding_sequence_parallel: bool = True
     tp_only_amax_red: bool = False
     tp_comm_overlap_cfg: Optional[Union[str, dict[str, Any]]] = None
     """Config file when tp_comm_overlap is enabled."""
 
-    transformer_layer_spec: Union[LayerSpec, Callable[["GPTModelProvider"], LayerSpec]] = local_layer_spec
-
     generation_config: Optional[Any] = None
 
     # This represents the unpadded vocab size
@@ -134,6 +149,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> GPTMode
         Returns:
             GPTModel: Configured PaddleFleet GPT model instance
         """
+        assert HAS_PADDLEFLEET
         vp_size = self.virtual_pipeline_model_parallel_size
         is_pipeline_asymmetric = getattr(self, "account_for_embedding_in_pipeline_split", False) or getattr(
             self, "account_for_loss_in_pipeline_split", False
@@ -151,25 +167,6 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> GPTMode
                 self.num_layers // p_size
             ) % vp_size == 0, "Make sure the number of model chunks is the same across all pipeline stages."
 
-        transformer_layer_spec = self.transformer_layer_spec
-        print(f"transformer_layer_spec  {transformer_layer_spec}")
-        print(f"param: {inspect.signature(transformer_layer_spec).parameters}")
-
-        if not isinstance(transformer_layer_spec, LayerSpec):
-            # Check if the transformer_layer_spec function accepts vp_stage parameter
-            if "vp_stage" in inspect.signature(transformer_layer_spec).parameters:
-                transformer_layer_spec = transformer_layer_spec(self, vp_stage=vp_stage)
-            else:
-                transformer_layer_spec = transformer_layer_spec(self)
-
-        assert self.vocab_size is not None, "vocab_size must be configured before calling provide()"
-        if self.should_pad_vocab:
-            padded_vocab_size = calculate_padded_vocab_size(
-                self.vocab_size, self.make_vocab_size_divisible_by, self.tensor_model_parallel_size
-            )
-        else:
-            padded_vocab_size = self.vocab_size
-
         # Initialize model as meta data instead of allocating data on a device
         model_init_device_context = contextlib.nullcontext
         if self.init_model_with_meta_device:
@@ -187,26 +184,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> GPTMode
         """
 
         with model_init_device_context():
-            model = GPTModel(
-                self,
-                transformer_layer_spec=transformer_layer_spec,
-                vocab_size=padded_vocab_size,
-                max_sequence_length=self.seq_length,
-                fp16_lm_cross_entropy=self.fp16_lm_cross_entropy,
-                parallel_output=self.parallel_output,
-                share_embeddings_and_output_weights=self.share_embeddings_and_output_weights,
-                position_embedding_type=self.position_embedding_type,
-                rotary_percent=self.rotary_percent,
-                rotary_base=self.rotary_base,
-                seq_len_interpolation_factor=self.seq_len_interpolation_factor,
-                pre_process=pre_process
-                or parallel_state.is_pipeline_first_stage(ignore_virtual=False, vp_stage=vp_stage),
-                post_process=post_process
-                or parallel_state.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage),
-                scatter_embedding_sequence_parallel=self.scatter_embedding_sequence_parallel,
-                vp_stage=vp_stage,
-                **kwargs,
-            )
+            model = gpt_builder(self, num_stages=1)
 
         return model
 
@@ -220,6 +198,7 @@ def mtp_block_spec(config: "GPTModelProvider", vp_stage: Optional[int] = None) -
     Returns:
         LayerSpec: The MTP module specification
     """
+    assert HAS_PADDLEFLEET
     if getattr(config, "mtp_num_layers", None):
         from paddlefleet.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec