PaddlePaddle
diff --git a/‎llm/auto_parallel/gpt-3/run_pretrain_auto.py
Lines changed: 7 additions & 0 deletions b/‎llm/auto_parallel/gpt-3/run_pretrain_auto.py
Lines changed: 7 additions & 0 deletions
diff --git a/‎paddlenlp/trainer/trainer_utils.py
Lines changed: 4 additions & 2 deletions b/‎paddlenlp/trainer/trainer_utils.py
Lines changed: 4 additions & 2 deletions
diff --git a/‎paddlenlp/transformers/gpt/__init__.py
Lines changed: 6 additions & 0 deletions b/‎paddlenlp/transformers/gpt/__init__.py
Lines changed: 6 additions & 0 deletions
@@ -38,6 +38,7 @@
     CosineAnnealingWithWarmupDecay,
     GPTConfig,
     GPTForCausalLMAuto,
+    GPTForCausalLMAutoPP,
     GPTForCausalLMNet,
     GPTPretrainingCriterionAuto,
     GPTPretrainingCriterionNet,
@@ -48,6 +49,7 @@
 
 MODEL_CLASSES = {
     "gpt": (GPTConfig, GPTForCausalLMAuto, GPTPretrainingCriterionAuto),
+    "gpt_pp": (GPTConfig, GPTForCausalLMAutoPP, GPTPretrainingCriterionAuto),
     "gpt_network": (GPTConfig, GPTForCausalLMNet, GPTPretrainingCriterionNet),
 }
 
@@ -99,6 +101,10 @@ class PreTrainingArguments(AutoTrainingArguments):
         default=False,
         metadata={"help": "Weather to run benchmark by autotuner. True for from_scratch and pad_max_length."},
     )
+    n_microbatches: int = field(
+        default=1,
+        metadata={"help": "Control the num of microbatches in one pp step."},
+    )
     pre_alloc_memory: float = field(
         default=0.0,
         metadata={
@@ -601,6 +607,7 @@ def fn(layer):
 
     trainer = PretrainingTrainer(
         model=model,
+        model_type=model_args.model_type,
         criterion=criterion,
         args=training_args,
         data_collator=data_collator,
 
@@ -43,7 +43,7 @@
 from paddlenlp.ops import Topology
 
 from ..trainer.argparser import strtobool
-from ..transformers import get_llama_pp_schedule
+from ..transformers import get_gpt_pp_schedule, get_llama_pp_schedule
 from ..transformers.tokenizer_utils_base import BatchEncoding
 from ..utils.env import PREFIX_CHECKPOINT_DIR, _re_checkpoint  # noqa for compatibility
 from ..utils.fault_tolerance import PDC_DOWNLOAD_ERROR
@@ -1257,11 +1257,13 @@ def download_recovery_ckpt_from_pdc(recovery_checkpoint_path, timeout):
 
 
 def check_auto_parallel_pipeline_support(model_type=None):
-    support_types = ["llama_pp"]
+    support_types = ["llama_pp", "gpt_pp"]
     return model_type in support_types
 
 
 def get_pp_schedule(model, model_type, n_microbatches, loss_fn, mode, pp_degree, group):
     assert check_auto_parallel_pipeline_support(model_type)
     if model_type == "llama_pp":
         return get_llama_pp_schedule(model, n_microbatches, loss_fn, mode, pp_degree, group)
+    elif model_type == "gpt_pp":
+        return get_gpt_pp_schedule(model, n_microbatches, loss_fn, mode, pp_degree, group)
@@ -15,6 +15,12 @@
 from .configuration import *
 from .modeling import *
 from .modeling_auto import *
+
+try:
+    from .modeling_auto_pp import *
+except (ImportError, ModuleNotFoundError):
+    # Temporarily adapt to the release version of Paddle, which can be removed later.
+    pass
 from .modeling_network import *
 from .modeling_pp import *
 from .tokenizer import *