[v0] support megatron-bridge for PT/SFT training

examples/megatron_bridge/llama3_sft.yaml

@@ -0,0 +1,38 @@
+### model
+model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+
+### method
+stage: sft
+do_train: true
+finetuning_type: full  # full or lora
+dataset: alpaca_en_demo
+template: llama3
+cutoff_len: 2048
+preprocessing_num_workers: 8
+
+### output
+output_dir: saves/mbridge/llama3_sft
+logging_steps: 1
+overwrite_output_dir: true
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+num_train_epochs: 3
+max_steps: 5 # max_steps is used to limit the number of steps to train, if set, num_train_epochs will be ignored
+save_steps: 3000
+learning_rate: 5.0e-6
+lr_scheduler_type: cosine
+warmup_steps: 10
+bf16: true
+
+### megatron bridge
+tensor_model_parallel_size: 2
+pipeline_model_parallel_size: 1
+context_parallel_size: 1
+sequence_parallel: true
+use_distributed_optimizer: true
+overlap_param_gather: true
+overlap_grad_reduce: true
+mixed_precision: bf16_mixed
+export_hf_on_finish: true

src/llamafactory/extras/packages.py

@@ -79,6 +79,13 @@ def is_mcore_adapter_available():
     return _is_package_available("mcore_adapter")
 
 
+def is_megatron_bridge_available():
+    try:
+        return _is_package_available("megatron.bridge")
+    except ModuleNotFoundError:
+        return False
+
+
 def is_pillow_available():
     return _is_package_available("PIL")
 

src/llamafactory/hparams/__init__.py

@@ -16,6 +16,7 @@
 from .evaluation_args import EvaluationArguments
 from .finetuning_args import FinetuningArguments
 from .generating_args import GeneratingArguments
+from .megatron_bridge_args import MegatronBridgeArguments
 from .model_args import ModelArguments
 from .parser import get_eval_args, get_infer_args, get_ray_args, get_train_args, read_args
 from .training_args import RayArguments, TrainingArguments
@@ -26,6 +27,7 @@
     "EvaluationArguments",
     "FinetuningArguments",
     "GeneratingArguments",
+    "MegatronBridgeArguments",
     "ModelArguments",
     "RayArguments",
     "TrainingArguments",

src/llamafactory/hparams/finetuning_args.py

@@ -482,6 +482,21 @@ class FinetuningArguments(
             )
         },
     )
+    use_megatron_bridge: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether or not to use Megatron Bridge training backend. "
+                "Controlled by USE_MEGATRON_BRIDGE environment variable."
+            )
+        },
+    )
+    megatron_bridge_args: Any = field(
+        default=None,
+        init=False,
+        repr=False,
+        metadata={"help": "Megatron Bridge specific arguments, set when USE_MEGATRON_BRIDGE=1."},
+    )
     use_hyper_parallel: bool = field(
         default=False,
         metadata={

src/llamafactory/hparams/megatron_bridge_args.py

@@ -0,0 +1,118 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+
+from transformers.training_args import _convert_str_dict
+
+
+@dataclass
+class MegatronBridgeArguments:
+    r"""Arguments for Megatron Bridge distributed training backend."""
+
+    tensor_model_parallel_size: int = field(
+        default=1,
+        metadata={"help": "Tensor model parallel size for Megatron Bridge."},
+    )
+    pipeline_model_parallel_size: int = field(
+        default=1,
+        metadata={"help": "Pipeline model parallel size for Megatron Bridge."},
+    )
+    expert_model_parallel_size: int = field(
+        default=1,
+        metadata={"help": "Expert model parallel size for MoE models."},
+    )
+    context_parallel_size: int = field(
+        default=1,
+        metadata={"help": "Context parallel size for Megatron Bridge."},
+    )
+    sequence_parallel: bool = field(
+        default=False,
+        metadata={"help": "Whether to enable sequence parallelism."},
+    )
+    recompute_granularity: Optional[str] = field(
+        default=None,
+        metadata={"help": "Activation recomputation granularity: 'full' or 'selective'."},
+    )
+    use_distributed_optimizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use Megatron distributed optimizer."},
+    )
+    overlap_param_gather: bool = field(
+        default=True,
+        metadata={"help": "Whether to overlap parameter all-gather with forward compute."},
+    )
+    overlap_grad_reduce: bool = field(
+        default=True,
+        metadata={"help": "Whether to overlap gradient all-reduce with backward compute."},
+    )
+    use_packed_sequences: bool = field(
+        default=False,
+        metadata={"help": "Whether to use packed sequences for SFT efficiency."},
+    )
+    mixed_precision: str = field(
+        default="bf16_mixed",
+        metadata={"help": "Mixed precision mode for Megatron Bridge, e.g. bf16_mixed or fp8."},
+    )
+    megatron_pretrained_checkpoint: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Path to a Megatron-format pretrained checkpoint. "
+                "If unset, HF weights are converted automatically before training."
+            )
+        },
+    )
+    export_hf_on_finish: bool = field(
+        default=False,
+        metadata={"help": "Whether to export the final checkpoint to Hugging Face format after training."},
+    )
+    extra_config: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Optional JSON string or path to a JSON file with extra Megatron Bridge model/training overrides."
+            )
+        },
+    )
+
+    def __post_init__(self) -> None:
+        if self.tensor_model_parallel_size < 1:
+            raise ValueError("`tensor_model_parallel_size` must be >= 1.")
+        if self.pipeline_model_parallel_size < 1:
+            raise ValueError("`pipeline_model_parallel_size` must be >= 1.")
+        if self.expert_model_parallel_size < 1:
+            raise ValueError("`expert_model_parallel_size` must be >= 1.")
+        if self.context_parallel_size < 1:
+            raise ValueError("`context_parallel_size` must be >= 1.")
+        if self.sequence_parallel and self.tensor_model_parallel_size <= 1:
+            raise ValueError("`sequence_parallel` requires `tensor_model_parallel_size` > 1.")
+        if self.recompute_granularity is not None and self.recompute_granularity not in ("full", "selective"):
+            raise ValueError("`recompute_granularity` must be 'full' or 'selective'.")
+
+        if isinstance(self.extra_config, str) and self.extra_config.startswith("{"):
+            self.extra_config = _convert_str_dict(json.loads(self.extra_config))
+
+    def load_extra_config(self) -> dict:
+        if self.extra_config is None:
+            return {}
+        if isinstance(self.extra_config, dict):
+            return self.extra_config
+        if not os.path.isfile(self.extra_config):
+            raise ValueError(f"`extra_config` file not found: {self.extra_config}")
+        with open(self.extra_config, encoding="utf-8") as f:
+            return json.load(f)

src/llamafactory/hparams/parser.py

@@ -33,11 +33,12 @@
 from ..extras import logging
 from ..extras.constants import CHECKPOINT_NAMES, EngineName
 from ..extras.misc import check_dependencies, check_version, get_current_device, is_env_enabled
-from ..extras.packages import is_mcore_adapter_available
+from ..extras.packages import is_mcore_adapter_available, is_megatron_bridge_available
 from .data_args import DataArguments
 from .evaluation_args import EvaluationArguments
 from .finetuning_args import FinetuningArguments
 from .generating_args import GeneratingArguments
+from .megatron_bridge_args import MegatronBridgeArguments
 from .model_args import ModelArguments
 from .training_args import RayArguments, TrainingArguments
 
@@ -81,6 +82,23 @@
     _TRAIN_MCA_ARGS = []
     _TRAIN_MCA_CLS = tuple()
 
+_TRAIN_MBRIDGE_ARGS = [
+    ModelArguments,
+    DataArguments,
+    TrainingArguments,
+    FinetuningArguments,
+    MegatronBridgeArguments,
+    GeneratingArguments,
+]
+_TRAIN_MBRIDGE_CLS = tuple[
+    ModelArguments,
+    DataArguments,
+    TrainingArguments,
+    FinetuningArguments,
+    MegatronBridgeArguments,
+    GeneratingArguments,
+]
+
 
 def read_args(args: dict[str, Any] | list[str] | None = None) -> dict[str, Any] | list[str]:
     r"""Get arguments from the command line or a config file."""
@@ -246,6 +264,9 @@ def _check_extra_dependencies(
     if finetuning_args.plot_loss:
         check_version("matplotlib", mandatory=True)
 
+    if finetuning_args.use_megatron_bridge:
+        check_version("megatron-bridge", mandatory=True)
+
     if training_args is not None:
         if training_args.deepspeed:
             check_version("deepspeed", mandatory=True)
@@ -283,6 +304,35 @@ def _configure_mca_training_args(training_args, data_args, finetuning_args) -> N
     finetuning_args.use_mca = True
 
 
+def _validate_megatron_bridge_parallel_args(mb_args: MegatronBridgeArguments, world_size: int) -> None:
+    parallel_size = (
+        mb_args.tensor_model_parallel_size
+        * mb_args.pipeline_model_parallel_size
+        * mb_args.context_parallel_size
+        * mb_args.expert_model_parallel_size
+    )
+    if parallel_size > world_size:
+        raise ValueError(f"Total Megatron Bridge parallel size ({parallel_size}) exceeds `world_size` ({world_size}).")
+
+
+def _parse_train_mbridge_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_MBRIDGE_CLS:
+    parser = HfArgumentParser(_TRAIN_MBRIDGE_ARGS)
+    allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
+    model_args, data_args, training_args, finetuning_args, mb_args, generating_args = _parse_args(
+        parser, args, allow_extra_keys=allow_extra_keys
+    )
+    _configure_mbridge_training_args(training_args, data_args, finetuning_args)
+    return model_args, data_args, training_args, finetuning_args, mb_args, generating_args
+
+
+def _configure_mbridge_training_args(training_args, data_args, finetuning_args) -> None:
+    """Patch training args to avoid args checking errors and sync Megatron Bridge settings."""
+    training_args.predict_with_generate = False
+    training_args.generation_max_length = data_args.cutoff_len
+    training_args.generation_num_beams = 1
+    finetuning_args.use_megatron_bridge = True
+
+
 def _parse_infer_args(args: dict[str, Any] | list[str] | None = None) -> _INFER_CLS:
     parser = HfArgumentParser(_INFER_ARGS)
     allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
@@ -302,11 +352,22 @@ def get_ray_args(args: dict[str, Any] | list[str] | None = None) -> RayArguments
 
 
 def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS:
+    mb_args = None
     if is_env_enabled("USE_MCA"):
         model_args, data_args, training_args, finetuning_args, generating_args = _parse_train_mca_args(args)
+    elif is_env_enabled("USE_MEGATRON_BRIDGE"):
+        if not is_megatron_bridge_available():
+            raise ImportError(
+                "megatron-bridge is required when USE_MEGATRON_BRIDGE=1. "
+                "Please install `megatron-bridge` and its dependencies."
+            )
+        model_args, data_args, training_args, finetuning_args, mb_args, generating_args = _parse_train_mbridge_args(
+            args
+        )
     else:
         model_args, data_args, training_args, finetuning_args, generating_args = _parse_train_args(args)
         finetuning_args.use_mca = False
+        finetuning_args.use_megatron_bridge = False
 
     # Setup logging
     if training_args.should_log:
@@ -326,6 +387,22 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS
     if finetuning_args.stage == "sft" and training_args.do_predict and not training_args.predict_with_generate:
         raise ValueError("Please enable `predict_with_generate` to save model predictions.")
 
+    if finetuning_args.use_megatron_bridge:
+        if finetuning_args.use_mca or finetuning_args.use_hyper_parallel:
+            raise ValueError("Megatron Bridge cannot be used together with MCA or HyperParallel.")
+        if finetuning_args.stage not in ["pt", "sft"]:
+            raise ValueError("Megatron Bridge only supports the `pt` and `sft` stages.")
+        if finetuning_args.finetuning_type not in ["full", "lora"]:
+            raise ValueError("Megatron Bridge only supports `full` and `lora` finetuning.")
+        if model_args.quantization_bit is not None:
+            raise ValueError("Quantized models are not supported with Megatron Bridge.")
+        if training_args.deepspeed is not None:
+            raise ValueError("Megatron Bridge is incompatible with DeepSpeed.")
+        if mb_args is None:
+            raise ValueError("Megatron Bridge arguments are missing. Please set USE_MEGATRON_BRIDGE=1.")
+        _validate_megatron_bridge_parallel_args(mb_args, training_args.world_size)
+        finetuning_args.megatron_bridge_args = mb_args
+
     if finetuning_args.stage in ["rm", "ppo"] and training_args.load_best_model_at_end:
         raise ValueError("RM and PPO stages do not support `load_best_model_at_end`.")
 
@@ -400,7 +477,12 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS
     if training_args.deepspeed is not None and (finetuning_args.use_galore or finetuning_args.use_apollo):
         raise ValueError("GaLore and APOLLO are incompatible with DeepSpeed yet.")
 
-    if not finetuning_args.use_mca and training_args.fp8 and model_args.quantization_bit is not None:
+    if (
+        not finetuning_args.use_mca
+        and not finetuning_args.use_megatron_bridge
+        and training_args.fp8
+        and model_args.quantization_bit is not None
+    ):
         raise ValueError("FP8 training is not compatible with quantization. Please disable one of them.")
 
     if model_args.infer_backend != EngineName.HF:
@@ -417,7 +499,12 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS
     _check_extra_dependencies(model_args, finetuning_args, training_args)
     _verify_trackio_args(training_args)
 
-    if not finetuning_args.use_mca and training_args.fp8_enable_fsdp_float8_all_gather and not training_args.fp8:
+    if (
+        not finetuning_args.use_mca
+        and not finetuning_args.use_megatron_bridge
+        and training_args.fp8_enable_fsdp_float8_all_gather
+        and not training_args.fp8
+    ):
         logger.warning_rank0("fp8_enable_fsdp_float8_all_gather requires fp8=True. Setting fp8=True.")
         model_args.fp8 = True
 

src/llamafactory/launcher.py

@@ -54,7 +54,7 @@ def launch():
     )
 
     command = sys.argv.pop(1) if len(sys.argv) > 1 else "help"
-    if is_env_enabled("USE_MCA"):  # force use torchrun
+    if is_env_enabled("USE_MCA") or is_env_enabled("USE_MEGATRON_BRIDGE"):  # force use torchrun
         os.environ["FORCE_TORCHRUN"] = "1"
 
     if command == "train" and (

src/llamafactory/train/megatron_bridge/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .workflow import run_pt, run_sft
+
+
+__all__ = ["run_pt", "run_sft"]