[train] support megatron-bridge for PT/SFT training (#10645)

docker/docker-cuda/Dockerfile.mbridge

@@ -0,0 +1,85 @@
+# LLaMA-Factory + Megatron Bridge (CUDA) runtime
+# Mirrors the verified host env: Python 3.12, PyTorch 2.12.1+cu126,
+# TransformerEngine 2.17, megatron-core 0.18, megatron-bridge 0.5.
+#
+# CUDA user-mode libs come from the PyTorch cu126 wheels (same as the host venv).
+# Layers are intentionally few (helps vfs / nested-docker disk usage).
+#
+# Build from repo root:
+#   bash docker/docker-cuda/build-megatron.sh
+
+ARG BASE_IMAGE=ubuntu:24.04
+FROM ${BASE_IMAGE}
+
+ARG PIP_INDEX=https://mirrors.aliyun.com/pypi/simple
+ARG PYPI_TRUSTED_HOST=mirrors.aliyun.com
+ARG TORCH_INDEX=https://download.pytorch.org/whl/cu126
+ARG APT_MIRROR=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_ROOT_USER_ACTION=ignore \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    PIP_BREAK_SYSTEM_PACKAGES=1 \
+    PIP_CONSTRAINT="" \
+    MAX_JOBS=8 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    DISABLE_VERSION_CHECK=1 \
+    USE_MEGATRON_BRIDGE=1 \
+    GRADIO_SERVER_PORT=7860 \
+    API_PORT=8000 \
+    http_proxy= \
+    https_proxy=
+
+SHELL ["/bin/bash", "-c"]
+WORKDIR /app
+
+# System deps + Ubuntu Python 3.12
+RUN if [ -f /etc/apt/sources.list.d/ubuntu.sources ]; then \
+      sed -i "s|http://archive.ubuntu.com/ubuntu|${APT_MIRROR}|g; s|http://security.ubuntu.com/ubuntu|${APT_MIRROR}|g" /etc/apt/sources.list.d/ubuntu.sources; \
+    elif [ -f /etc/apt/sources.list ]; then \
+      sed -i "s|http://archive.ubuntu.com/ubuntu/|${APT_MIRROR}|g; s|http://security.ubuntu.com/ubuntu/|${APT_MIRROR}|g" /etc/apt/sources.list; \
+    fi && \
+    apt-get update && apt-get install -y --no-install-recommends \
+        ca-certificates curl git vim wget \
+        build-essential ninja-build cmake libgomp1 zip unzip \
+        python3 python3-pip python3-dev python3-venv && \
+    ln -sf /usr/bin/python3 /usr/bin/python && \
+    rm -rf /var/lib/apt/lists/* && \
+    pip install --no-cache-dir --upgrade pip setuptools wheel packaging ninja pybind11 \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX}
+
+COPY . /app
+
+# PyTorch + TE + Megatron Bridge + LLaMA-Factory in one layer
+RUN pip install --no-cache-dir \
+        torch==2.12.1 torchvision==0.27.1 torchaudio==2.11.0 \
+        --index-url ${TORCH_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation \
+        "transformer-engine[pytorch]==2.17.0" \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation \
+        "megatron-bridge==0.5.0" \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation -e . \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation \
+        -r requirements/metrics.txt \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation \
+        "megatron-bridge==0.5.0" \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    python - <<'PY'
+import torch
+import megatron.core
+import transformer_engine
+from megatron.bridge import AutoBridge  # noqa: F401
+import llamafactory
+print("torch", torch.__version__, "cuda", torch.version.cuda)
+print("te", transformer_engine.__version__)
+print("megatron-bridge import ok")
+PY
+
+EXPOSE 7860 8000
+CMD ["bash"]

docker/docker-cuda/README.md

@@ -104,6 +104,37 @@ sudo usermod -aG docker $USER
 # Log out and back in for changes to take effect
 ```
 
+## Megatron Bridge Image
+
+`Dockerfile.megatron` builds a CUDA runtime for LLaMA-Factory + [Megatron Bridge](https://docs.nvidia.com/nemo/megatron-bridge/latest/):
+
+| Component | Version |
+| --- | --- |
+| Base | `ubuntu:22.04` (Python 3.12) |
+| PyTorch | 2.12.1+cu126 (CUDA libs from wheels) |
+| TransformerEngine | 2.17.0 |
+| megatron-core | 0.18.x (via megatron-bridge) |
+| megatron-bridge | 0.5.0 |
+
+### Build
+
+From repo root:
+
+```bash
+docker build -f docker/docker-cuda/Dockerfile.megatron \
+  -t llamafactory-megatron-bridge:latest .
+```
+
+### Run training
+
+```bash
+docker run --rm -it --gpus all --ipc=host --shm-size=16g \
+  -e DISABLE_VERSION_CHECK=1 \
+  -e USE_MEGATRON_BRIDGE=1 \
+  -v "$PWD":/app -w /app \
+  llamafactory-megatron-bridge:latest
+```
+
 ## Additional Notes
 
 - The default image is built on Ubuntu 22.04 (x86_64), CUDA 12.4, Python 3.11, PyTorch 2.6.0, and Flash-attn 2.7.4

examples/megatron_bridge/llama3_sft.yaml

@@ -0,0 +1,67 @@
+### model
+model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora  # full or lora
+dataset: alpaca_en_demo
+template: llama3
+cutoff_len: 2048
+preprocessing_num_workers: 8
+# disable_shuffling: true  # keep sample order aligned with HF baseline
+
+### 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: 1000  # when set, overrides num_train_epochs for Megatron Bridge schedule
+save_steps: 3000
+learning_rate: 5.0e-6
+lr_scheduler_type: cosine
+warmup_steps: 10
+adam_beta1: 0.9
+adam_beta2: 0.999
+weight_decay: 0.0
+max_grad_norm: 1.0
+bf16: true
+
+### megatron bridge parallelism
+tensor_model_parallel_size: 1
+pipeline_model_parallel_size: 1
+context_parallel_size: 1
+expert_model_parallel_size: 1
+# virtual_pipeline_model_parallel_size: 2
+sequence_parallel: false
+
+### megatron bridge optimizer / overlap
+use_distributed_optimizer: true
+overlap_param_gather: true
+overlap_grad_reduce: true
+mixed_precision: bf16_mixed
+
+### megatron bridge activation recompute (optional)
+# recompute_granularity: full
+# recompute_method: uniform
+# recompute_num_layers: 1
+
+### megatron bridge model kernels (optional; None keeps provider defaults)
+# bias_activation_fusion: true
+# apply_rope_fusion: true
+# masked_softmax_fusion: true
+# cross_entropy_loss_fusion: true
+
+### megatron bridge MoE (optional)
+# moe_grouped_gemm: true
+# moe_token_dispatcher_type: alltoall
+
+### megatron bridge data / checkpoint
+use_packed_sequences: false
+# megatron_pretrained_checkpoint: /path/to/megatron_ckpt
+export_hf_on_finish: false  # disable for short loss-comparison runs (avoids checkpoint OOM)
+# extra_config: '{"train.train_iters": 5, "logger.log_interval": 1}'

src/llamafactory/extras/constants.py

@@ -73,6 +73,23 @@ MCA_SUPPORTED_MODELS = {
     "qwen3_5_moe",
 }
 
+# Text LLM model_types supported by the Megatron Bridge PT/SFT path (gpt_step).
+# Multimodal / audio / omni architectures are excluded in v0.
+MEGATRON_BRIDGE_SUPPORTED_MODELS = {
+    "deepseek_v3",
+    "deepseek_v4",
+    "llama",
+    "mistral",
+    "qwen2",
+    "qwen3",
+    "qwen3_5",
+    "qwen3_5_moe",
+    "qwen3_5_moe_text",
+    "qwen3_5_text",
+    "qwen3_moe",
+    "qwen3_next",
+}
+
 METHODS = ["full", "freeze", "lora", "oft"]
 
 MOD_SUPPORTED_MODELS = {"bloom", "falcon", "gemma", "llama", "mistral", "mixtral", "phi", "starcoder2"}

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 .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 .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 @@ __all__ = [
     "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,193 @@
+# 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 Literal, Optional
+
+from transformers.training_args import _convert_str_dict
+
+
+@dataclass
+class MegatronBridgeArguments:
+    r"""Arguments for Megatron Bridge distributed training backend.
+
+    Parallelism, optimizer overlap, checkpoint conversion, and selected Megatron
+    model-provider knobs are exposed here because Megatron Bridge uses a
+    standalone workflow outside the Hugging Face Trainer.
+    """
+
+    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."},
+    )
+    virtual_pipeline_model_parallel_size: Optional[int] = field(
+        default=None,
+        metadata={"help": "Virtual pipeline (interleaved) parallel size. None keeps provider default."},
+    )
+    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'."},
+    )
+    recompute_method: Optional[Literal["uniform", "block"]] = field(
+        default=None,
+        metadata={"help": "Activation recomputation method: 'uniform' or 'block'."},
+    )
+    recompute_num_layers: Optional[int] = field(
+        default=None,
+        metadata={"help": "Number of layers per recompute unit when recompute_method is set."},
+    )
+    account_for_embedding_in_pipeline_split: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Whether pipeline split accounts for the embedding layer."},
+    )
+    account_for_loss_in_pipeline_split: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Whether pipeline split accounts for the loss layer."},
+    )
+    bias_activation_fusion: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Enable bias+activation fusion. None keeps Megatron provider default."},
+    )
+    apply_rope_fusion: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Enable RoPE fusion kernel. None keeps Megatron provider default."},
+    )
+    masked_softmax_fusion: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Enable masked softmax fusion. None keeps Megatron provider default."},
+    )
+    cross_entropy_loss_fusion: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Enable cross-entropy loss fusion. None keeps Megatron provider default."},
+    )
+    moe_grouped_gemm: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Enable grouped GEMM for MoE experts. None keeps provider default."},
+    )
+    moe_token_dispatcher_type: Optional[Literal["allgather", "alltoall", "flex"]] = field(
+        default=None,
+        metadata={"help": "MoE token dispatcher type: allgather, alltoall, or flex."},
+    )
+    calculate_per_token_loss: Optional[bool] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Whether to compute per-token loss. When context_parallel_size > 1, "
+                "this is forced to True regardless of this setting."
+            )
+        },
+    )
+    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. "
+                "Dot-paths are supported (e.g. train.train_iters or checkpoint.save_interval)."
+            )
+        },
+    )
+
+    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.virtual_pipeline_model_parallel_size is not None and self.virtual_pipeline_model_parallel_size < 1:
+            raise ValueError("`virtual_pipeline_model_parallel_size` must be >= 1 when set.")
+        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 self.recompute_method is not None and self.recompute_method not in ("uniform", "block"):
+            raise ValueError("`recompute_method` must be 'uniform' or 'block'.")
+        if self.recompute_num_layers is not None and self.recompute_num_layers < 1:
+            raise ValueError("`recompute_num_layers` must be >= 1 when set.")
+        if self.moe_token_dispatcher_type is not None and self.moe_token_dispatcher_type not in (
+            "allgather",
+            "alltoall",
+            "flex",
+        ):
+            raise ValueError("`moe_token_dispatcher_type` must be 'allgather', 'alltoall', or 'flex'.")
+
+        if isinstance(self.extra_config, str):
+            config_str = self.extra_config.strip()
+            if config_str.startswith("{"):
+                self.extra_config = _convert_str_dict(json.loads(config_str))
+            else:
+                self.extra_config = config_str
+
+    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 transformers.utils import is_torch_bf16_gpu_available, is_torch_npu_availab
 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 @@ else:
     _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,39 @@ 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}).")
+    if world_size % parallel_size != 0:
+        raise ValueError(
+            f"Total Megatron Bridge parallel size ({parallel_size}) must divide `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 +356,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 +391,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 +481,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 +503,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"]

src/llamafactory/train/megatron_bridge/config_builder.py

@@ -0,0 +1,708 @@
+# 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 math
+import os
+from typing import TYPE_CHECKING, Any
+
+from ...extras.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from ...hparams import (
+        DataArguments,
+        FinetuningArguments,
+        MegatronBridgeArguments,
+        ModelArguments,
+        TrainingArguments,
+    )
+
+
+logger = get_logger(__name__)
+
+_LR_SCHEDULER_MAP = {
+    "cosine": "cosine",
+    "linear": "linear",
+    "constant": "constant",
+    "constant_with_warmup": "constant",
+}
+
+
+def _map_lr_scheduler_type(lr_scheduler_type: str) -> str:
+    mapped = _LR_SCHEDULER_MAP.get(lr_scheduler_type)
+    if mapped is None:
+        logger.warning_rank0(
+            f"lr_scheduler_type '{lr_scheduler_type}' is not supported by Megatron Bridge; using cosine."
+        )
+        return "cosine"
+    return mapped
+
+
+def _resolve_warmup_steps(training_args: "TrainingArguments", train_iters: int) -> int:
+    r"""Resolve warmup steps with Hugging Face Trainer semantics.
+
+    Absolute ``warmup_steps`` are kept as-is (even when larger than ``train_iters``)
+    so short debugging runs with ``max_steps < warmup_steps`` match HF LR values.
+    ``lr_decay_iters`` is expanded separately to avoid Megatron capping warmup.
+    """
+    warmup_steps = getattr(training_args, "warmup_steps", 0) or 0
+    if warmup_steps > 0:
+        return warmup_steps
+
+    warmup_ratio = getattr(training_args, "warmup_ratio", 0.0) or 0.0
+    if warmup_ratio > 0:
+        return min(int(train_iters * warmup_ratio), train_iters)
+    return 0
+
+
+def _resolve_decay_iters(train_iters: int, warmup_steps: int) -> int:
+    r"""Ensure decay span is longer than warmup so Megatron does not shrink warmup.
+
+    Megatron Bridge caps ``lr_warmup_steps`` whenever it is ``>= lr_decay_steps``,
+    so keep decay strictly larger than warmup on short comparison runs.
+    """
+    if warmup_steps <= 0:
+        return train_iters
+    return max(train_iters, warmup_steps + 1)
+
+
+def _import_training_config():
+    from megatron.bridge.training.config import (
+        CheckpointConfig,
+        ConfigContainer,
+        FinetuningDatasetConfig,
+        GPTDatasetConfig,
+        LoggerConfig,
+        RNGConfig,
+        TrainingConfig,
+    )
+
+    try:
+        from megatron.bridge.training.config import DistributedInitConfig
+    except ImportError:
+        DistributedInitConfig = None
+
+    try:
+        from megatron.bridge.training.config import ValidationConfig
+    except ImportError:
+        ValidationConfig = None
+
+    try:
+        from megatron.bridge.training.tokenizers.config import TokenizerConfig
+    except ImportError:
+        from megatron.bridge.training.config import TokenizerConfig
+
+    return (
+        CheckpointConfig,
+        ConfigContainer,
+        DistributedInitConfig,
+        FinetuningDatasetConfig,
+        GPTDatasetConfig,
+        LoggerConfig,
+        RNGConfig,
+        TokenizerConfig,
+        TrainingConfig,
+        ValidationConfig,
+    )
+
+
+def _create_optimizer_scheduler(
+    training_args: "TrainingArguments",
+    warmup_steps: int,
+    train_iters: int,
+    finetuning_args: "FinetuningArguments",
+    use_distributed_optimizer: bool,
+):
+    from megatron.bridge.training.config import OptimizerConfig, SchedulerConfig
+
+    finetuning_type = finetuning_args.finetuning_type
+    learning_rate = training_args.learning_rate
+    if finetuning_type in ("lora", "full"):
+        max_lr, min_lr, default_beta2 = learning_rate, 0.0, 0.999
+    else:
+        max_lr, min_lr, default_beta2 = learning_rate, learning_rate * 0.1, 0.95
+
+    # Match Hugging Face Trainer defaults unless the user overrides them.
+    adam_beta1 = getattr(training_args, "adam_beta1", 0.9)
+    adam_beta2 = getattr(training_args, "adam_beta2", default_beta2)
+    adam_eps = getattr(training_args, "adam_epsilon", 1e-8)
+    weight_decay = getattr(training_args, "weight_decay", 0.0)
+    max_grad_norm = getattr(training_args, "max_grad_norm", 1.0)
+    decay_iters = _resolve_decay_iters(train_iters, warmup_steps)
+    optimizer = OptimizerConfig(
+        optimizer="adam",
+        lr=max_lr,
+        min_lr=min_lr,
+        weight_decay=weight_decay,
+        bf16=getattr(training_args, "bf16", True),
+        fp16=getattr(training_args, "fp16", False),
+        adam_beta1=adam_beta1,
+        adam_beta2=adam_beta2,
+        adam_eps=adam_eps,
+        use_distributed_optimizer=use_distributed_optimizer,
+        clip_grad=max_grad_norm,
+    )
+    scheduler = SchedulerConfig(
+        start_weight_decay=weight_decay,
+        end_weight_decay=weight_decay,
+        weight_decay_incr_style="constant",
+        lr_decay_style=_map_lr_scheduler_type(getattr(training_args, "lr_scheduler_type", "cosine")),
+        lr_wsd_decay_style="minus_sqrt",
+        lr_wsd_decay_iters=decay_iters,
+        lr_warmup_iters=warmup_steps,
+        lr_warmup_init=0.0,
+        lr_decay_iters=decay_iters,
+        override_opt_param_scheduler=True,
+    )
+    return optimizer, scheduler
+
+
+def ensure_create_sft_dataset_applies_chat_template() -> None:
+    r"""Apply ``dataset_kwargs['chat_template']`` before Megatron builds SFT datasets.
+
+    Megatron Bridge only pops ``chat_template`` on the packed-sequence path. The
+    non-packed finetuning path would otherwise forward it as an unexpected kwarg.
+    Patch both the defining module and the builder import binding.
+    """
+    from megatron.bridge.data.builders import finetuning_dataset as finetuning_module
+    from megatron.bridge.data.datasets import sft as sft_module
+
+    if getattr(sft_module.create_sft_dataset, "_llamafactory_chat_template_patched", False):
+        return
+
+    original = sft_module.create_sft_dataset
+
+    def create_sft_dataset(*args, **kwargs):
+        chat_template = kwargs.pop("chat_template", None)
+        tokenizer = kwargs.get("tokenizer")
+        if tokenizer is None and len(args) >= 2:
+            tokenizer = args[1]
+        if chat_template is not None and tokenizer is not None:
+            # Megatron `_chat_preprocess` may read either the wrapper or the
+            # inner HuggingFace tokenizer depending on `legacy`.
+            if hasattr(tokenizer, "chat_template"):
+                tokenizer.chat_template = chat_template
+            hf_tokenizer = getattr(tokenizer, "_tokenizer", None)
+            if hf_tokenizer is not None and hasattr(hf_tokenizer, "chat_template"):
+                hf_tokenizer.chat_template = chat_template
+        return original(*args, **kwargs)
+
+    create_sft_dataset._llamafactory_chat_template_patched = True  # type: ignore[attr-defined]
+    sft_module.create_sft_dataset = create_sft_dataset
+    finetuning_module.create_sft_dataset = create_sft_dataset
+    logger.info_rank0("Patched Megatron create_sft_dataset to apply chat_template overrides.")
+
+
+def _create_peft_config(finetuning_args: "FinetuningArguments"):
+    if finetuning_args.finetuning_type != "lora":
+        return None
+
+    from megatron.bridge.peft.lora import LoRA
+
+    default_targets = ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"]
+    if list(finetuning_args.lora_target) != ["all"]:
+        logger.warning_rank0(
+            f"Custom lora_target {finetuning_args.lora_target} is not supported by Megatron Bridge. "
+            f"Using default Megatron target modules: {default_targets}."
+        )
+
+    return LoRA(
+        target_modules=default_targets,
+        dim=finetuning_args.lora_rank,
+        alpha=finetuning_args.lora_alpha,
+    )
+
+
+def _build_gpt_dataset_config(
+    GPTDatasetConfig,
+    dataset_path: str,
+    seq_length: int,
+    seed: int,
+    num_workers: int,
+):
+    kwargs: dict[str, Any] = {
+        "random_seed": seed,
+        "reset_attention_mask": False,
+        "reset_position_ids": False,
+        "eod_mask_loss": False,
+        "blend": ([dataset_path], 1.0),
+        "split": "100,0,0",
+        "num_workers": num_workers,
+        "data_sharding": True,
+        "dataloader_type": "single",
+    }
+    if "sequence_length" in GPTDatasetConfig.__dataclass_fields__:
+        kwargs["sequence_length"] = seq_length
+    else:
+        kwargs["seq_length"] = seq_length
+    if "num_dataset_builder_threads" in GPTDatasetConfig.__dataclass_fields__:
+        kwargs["num_dataset_builder_threads"] = 1
+    return GPTDatasetConfig(**kwargs)
+
+
+def _build_finetuning_dataset_config(
+    FinetuningDatasetConfig,
+    dataset_root: str,
+    seq_length: int,
+    seed: int,
+    num_workers: int,
+    do_validation: bool,
+    dataset_kwargs: dict[str, Any],
+    packed_sequence_specs,
+    disable_shuffling: bool = False,
+):
+    kwargs: dict[str, Any] = {
+        "dataset_root": dataset_root,
+        "seq_length": seq_length,
+        "seed": seed,
+        "num_workers": num_workers,
+        "do_validation": do_validation,
+        "do_test": False,
+        "dataset_kwargs": dataset_kwargs,
+        "packed_sequence_specs": packed_sequence_specs,
+    }
+    if "dataloader_type" in FinetuningDatasetConfig.__dataclass_fields__:
+        from .dataset_export import get_finetuning_dataloader_type
+
+        kwargs["dataloader_type"] = get_finetuning_dataloader_type(disable_shuffling=disable_shuffling)
+    return FinetuningDatasetConfig(**kwargs)
+
+
+def _has_megatron_checkpoint(output_dir: str) -> bool:
+    r"""Return whether ``output_dir`` contains a resumable Megatron checkpoint."""
+    return any(
+        os.path.isfile(os.path.join(output_dir, name))
+        for name in ("latest_checkpointed_iteration.txt", "latest_train_state.pt")
+    )
+
+
+def _should_resume_checkpoint(training_args: "TrainingArguments") -> bool:
+    r"""Resume only when a tracker exists and ``overwrite_output_dir`` is false."""
+    if getattr(training_args, "overwrite_output_dir", False):
+        return False
+    return _has_megatron_checkpoint(training_args.output_dir)
+
+
+def _create_base_config(
+    *,
+    training_args: "TrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    train_iters: int,
+    micro_batch_size: int,
+    global_batch_size: int,
+    mb_args: "MegatronBridgeArguments",
+    is_sft: bool,
+):
+    from megatron.core.distributed import DistributedDataParallelConfig
+
+    (
+        CheckpointConfig,
+        ConfigContainer,
+        DistributedInitConfig,
+        _FinetuningDatasetConfig,
+        _GPTDatasetConfig,
+        LoggerConfig,
+        RNGConfig,
+        TokenizerConfig,
+        TrainingConfig,
+        ValidationConfig,
+    ) = _import_training_config()
+
+    warmup_steps = _resolve_warmup_steps(training_args, train_iters)
+    opt_cfg, scheduler_cfg = _create_optimizer_scheduler(
+        training_args=training_args,
+        warmup_steps=warmup_steps,
+        train_iters=train_iters,
+        finetuning_args=finetuning_args,
+        use_distributed_optimizer=mb_args.use_distributed_optimizer,
+    )
+
+    train_kwargs: dict[str, Any] = {
+        "train_iters": train_iters,
+        "global_batch_size": global_batch_size,
+        "micro_batch_size": micro_batch_size,
+    }
+    eval_steps = training_args.eval_steps
+    if ValidationConfig is not None:
+        validation = ValidationConfig(eval_interval=eval_steps or 100, eval_iters=32)
+    else:
+        train_kwargs["eval_interval"] = eval_steps or 100
+        train_kwargs["eval_iters"] = 32
+        validation = None
+
+    output_dir = training_args.output_dir
+    resume_checkpoint = _should_resume_checkpoint(training_args)
+    # mcore >= 0.14 removed ShardedTensor.flattened_range. The legacy default
+    # sharding type ``fully_sharded_model_space`` still depends on it, so prefer
+    # the fully_reshardable distributed-optimizer format when dist opt is on.
+    dist_ckpt_optim_fully_reshardable = mb_args.use_distributed_optimizer
+    dist_cfg = DistributedInitConfig() if DistributedInitConfig is not None else None
+    container_kwargs: dict[str, Any] = {
+        "model": None,
+        "train": TrainingConfig(**train_kwargs),
+        "optimizer": opt_cfg,
+        "scheduler": scheduler_cfg,
+        "ddp": DistributedDataParallelConfig(
+            check_for_nan_in_grad=True,
+            grad_reduce_in_fp32=True,
+            overlap_grad_reduce=mb_args.overlap_grad_reduce,
+            overlap_param_gather=mb_args.overlap_param_gather,
+            use_distributed_optimizer=mb_args.use_distributed_optimizer,
+        ),
+        "dataset": None,
+        "logger": LoggerConfig(
+            log_interval=training_args.logging_steps,
+            tensorboard_dir=os.path.join(output_dir, "tb_logs"),
+        ),
+        "tokenizer": TokenizerConfig(
+            tokenizer_type="HuggingFaceTokenizer",
+            tokenizer_model=None,
+        ),
+        "checkpoint": CheckpointConfig(
+            save_interval=training_args.save_steps,
+            save=output_dir,
+            load=output_dir if resume_checkpoint else None,
+            # SFT from pretrained should not load optimizer/RNG from a partial ckpt.
+            finetune=is_sft and not resume_checkpoint,
+            ckpt_format="torch_dist",
+            fully_parallel_save=False,
+            use_persistent_ckpt_worker=False,
+            save_optim=True,
+            dist_ckpt_optim_fully_reshardable=dist_ckpt_optim_fully_reshardable,
+        ),
+        "rng": RNGConfig(seed=training_args.seed),
+        "mixed_precision": mb_args.mixed_precision,
+        "peft": _create_peft_config(finetuning_args) if is_sft and finetuning_args.finetuning_type == "lora" else None,
+    }
+    if validation is not None:
+        container_kwargs["validation"] = validation
+    if dist_cfg is not None:
+        container_kwargs["dist"] = dist_cfg
+
+    return ConfigContainer(**container_kwargs)
+
+
+def _compute_train_schedule(
+    training_args: "TrainingArguments",
+    mb_args: "MegatronBridgeArguments",
+    num_train_samples: int,
+) -> tuple[int, int, int]:
+    micro_batch_size = training_args.per_device_train_batch_size
+    global_batch_size = micro_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
+    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
+    )
+    global_batch_size //= parallel_size
+    global_batch_size = max(global_batch_size, micro_batch_size)
+
+    max_steps = getattr(training_args, "max_steps", -1)
+    if max_steps is not None and max_steps > 0:
+        train_iters = max_steps
+    else:
+        train_iters = max(1, math.ceil(num_train_samples / global_batch_size * training_args.num_train_epochs))
+    return micro_batch_size, global_batch_size, train_iters
+
+
+def _is_apex_grad_accum_fusion_available() -> bool:
+    try:
+        import fused_weight_gradient_mlp_cuda  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def _apply_fusion_safety(model_provider) -> None:
+    r"""Disable gradient_accumulation_fusion when the APEX CUDA extension is missing.
+
+    Megatron Bridge enables this fusion when TransformerEngine is installed, but
+    ColumnParallelLinear (e.g. the output layer) still requires the APEX
+    fused_weight_gradient_mlp_cuda extension at model construction time.
+    """
+    if getattr(model_provider, "gradient_accumulation_fusion", False) and not _is_apex_grad_accum_fusion_available():
+        logger.warning_rank0(
+            "Disabling gradient_accumulation_fusion because the APEX CUDA extension "
+            "fused_weight_gradient_mlp_cuda is not installed."
+        )
+        model_provider.gradient_accumulation_fusion = False
+
+
+def _apply_optional_provider_attr(model_provider, name: str, value) -> None:
+    if value is None or not hasattr(model_provider, name):
+        return
+    setattr(model_provider, name, value)
+
+
+def _apply_model_parallelism(model_provider, mb_args: "MegatronBridgeArguments") -> None:
+    model_provider.tensor_model_parallel_size = mb_args.tensor_model_parallel_size
+    model_provider.pipeline_model_parallel_size = mb_args.pipeline_model_parallel_size
+    if hasattr(model_provider, "expert_model_parallel_size"):
+        model_provider.expert_model_parallel_size = mb_args.expert_model_parallel_size
+    model_provider.context_parallel_size = mb_args.context_parallel_size
+    model_provider.sequence_parallel = mb_args.sequence_parallel
+    _apply_optional_provider_attr(
+        model_provider, "virtual_pipeline_model_parallel_size", mb_args.virtual_pipeline_model_parallel_size
+    )
+    _apply_optional_provider_attr(model_provider, "recompute_granularity", mb_args.recompute_granularity)
+    _apply_optional_provider_attr(model_provider, "recompute_method", mb_args.recompute_method)
+    _apply_optional_provider_attr(model_provider, "recompute_num_layers", mb_args.recompute_num_layers)
+    _apply_optional_provider_attr(
+        model_provider, "account_for_embedding_in_pipeline_split", mb_args.account_for_embedding_in_pipeline_split
+    )
+    _apply_optional_provider_attr(
+        model_provider, "account_for_loss_in_pipeline_split", mb_args.account_for_loss_in_pipeline_split
+    )
+    _apply_optional_provider_attr(model_provider, "bias_activation_fusion", mb_args.bias_activation_fusion)
+    _apply_optional_provider_attr(model_provider, "apply_rope_fusion", mb_args.apply_rope_fusion)
+    _apply_optional_provider_attr(model_provider, "masked_softmax_fusion", mb_args.masked_softmax_fusion)
+    _apply_optional_provider_attr(model_provider, "cross_entropy_loss_fusion", mb_args.cross_entropy_loss_fusion)
+    _apply_optional_provider_attr(model_provider, "moe_grouped_gemm", mb_args.moe_grouped_gemm)
+    _apply_optional_provider_attr(model_provider, "moe_token_dispatcher_type", mb_args.moe_token_dispatcher_type)
+    _apply_optional_provider_attr(model_provider, "calculate_per_token_loss", mb_args.calculate_per_token_loss)
+
+
+def _apply_context_parallel_finetuning_requirements(cfg, mb_args: "MegatronBridgeArguments") -> None:
+    r"""Apply Megatron Bridge SFT requirements when context parallelism is enabled."""
+    if mb_args.context_parallel_size <= 1:
+        return
+    cfg.model.calculate_per_token_loss = True
+    cfg.ddp.average_in_collective = False
+
+
+def _apply_extra_overrides(cfg, extra: dict) -> None:
+    for key, value in extra.items():
+        parts = key.split(".")
+        obj = cfg
+        for part in parts[:-1]:
+            obj = getattr(obj, part)
+        setattr(obj, parts[-1], value)
+
+
+def _reset_megatron_bridge_global_state_after_checkpoint_conversion() -> None:
+    r"""Clear Megatron globals left over from HF-to-Megatron conversion.
+
+    save_megatron_model() can implicitly initialize the rerun state machine while
+    writing checkpoints. provide_distributed_model() also initializes model parallel
+    groups. Training later calls initialize_megatron(), which expects a fresh global
+    state and fails with "Rerun state machine is already initialized" or keeps stale
+    parallel groups that mismatch the configured tensor/pipeline/context parallel sizes.
+    """
+    from megatron.core import parallel_state
+    from megatron.core.rerun_state_machine import destroy_rerun_state_machine
+
+    destroy_rerun_state_machine()
+    parallel_state.destroy_model_parallel()
+
+
+def ensure_megatron_pretrained_checkpoint(
+    model_args: "ModelArguments",
+    mb_args: "MegatronBridgeArguments",
+    output_dir: str,
+) -> str:
+    r"""Convert Hugging Face weights to Megatron format when needed."""
+    from megatron.bridge import AutoBridge
+
+    if mb_args.megatron_pretrained_checkpoint and os.path.isdir(mb_args.megatron_pretrained_checkpoint):
+        return mb_args.megatron_pretrained_checkpoint
+
+    ckpt_dir = os.path.join(output_dir, "megatron_pretrained")
+    if os.path.isdir(ckpt_dir) and os.listdir(ckpt_dir):
+        logger.info_rank0(f"Reusing existing Megatron checkpoint at {ckpt_dir}.")
+        return ckpt_dir
+
+    os.makedirs(ckpt_dir, exist_ok=True)
+    logger.info_rank0(f"Converting Hugging Face weights to Megatron format at {ckpt_dir}...")
+    bridge = AutoBridge.from_hf_pretrained(
+        model_args.model_name_or_path,
+        trust_remote_code=model_args.trust_remote_code,
+    )
+    provider = bridge.to_megatron_provider()
+    _apply_model_parallelism(provider, mb_args)
+    _apply_fusion_safety(provider)
+    if hasattr(provider, "finalize"):
+        provider.finalize()
+    # TP/PP/CP weight scatter uses NCCL, which cannot operate on CPU tensors.
+    use_cpu_initialization = (
+        mb_args.tensor_model_parallel_size == 1
+        and mb_args.pipeline_model_parallel_size == 1
+        and mb_args.context_parallel_size == 1
+    )
+    if not use_cpu_initialization:
+        logger.info_rank0(
+            "Using GPU initialization for Megatron checkpoint conversion because model parallelism "
+            "requires NCCL scatter/gather on CUDA tensors."
+        )
+    try:
+        megatron_model = provider.provide_distributed_model(
+            wrap_with_ddp=False,
+            use_cpu_initialization=use_cpu_initialization,
+        )
+        hf_tokenizer_kwargs = {"trust_remote_code": True} if model_args.trust_remote_code else None
+        bridge.save_megatron_model(
+            megatron_model,
+            ckpt_dir,
+            hf_tokenizer_path=model_args.model_name_or_path,
+            hf_tokenizer_kwargs=hf_tokenizer_kwargs,
+            low_memory_save=True,
+        )
+    finally:
+        _reset_megatron_bridge_global_state_after_checkpoint_conversion()
+    return ckpt_dir
+
+
+def build_pretrain_config(
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "TrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    mb_args: "MegatronBridgeArguments",
+    dataset_path: str,
+    num_train_samples: int,
+):
+    from megatron.bridge import AutoBridge
+
+    micro_batch_size, global_batch_size, train_iters = _compute_train_schedule(
+        training_args, mb_args, num_train_samples
+    )
+    cfg = _create_base_config(
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        train_iters=train_iters,
+        micro_batch_size=micro_batch_size,
+        global_batch_size=global_batch_size,
+        mb_args=mb_args,
+        is_sft=False,
+    )
+
+    (
+        _CheckpointConfig,
+        _ConfigContainer,
+        _DistributedInitConfig,
+        _FinetuningDatasetConfig,
+        GPTDatasetConfig,
+        _LoggerConfig,
+        _RNGConfig,
+        _TokenizerConfig,
+        _TrainingConfig,
+        _ValidationConfig,
+    ) = _import_training_config()
+    bridge = AutoBridge.from_hf_pretrained(
+        model_args.model_name_or_path,
+        trust_remote_code=model_args.trust_remote_code,
+    )
+    cfg.model = bridge.to_megatron_provider(load_weights=False)
+    _apply_model_parallelism(cfg.model, mb_args)
+    _apply_fusion_safety(cfg.model)
+    if hasattr(cfg.model, "seq_length"):
+        cfg.model.seq_length = data_args.cutoff_len
+
+    cfg.tokenizer.tokenizer_model = model_args.model_name_or_path
+    cfg.dataset = _build_gpt_dataset_config(
+        GPTDatasetConfig,
+        dataset_path=dataset_path,
+        seq_length=data_args.cutoff_len,
+        seed=training_args.seed,
+        num_workers=data_args.preprocessing_num_workers,
+    )
+
+    _apply_extra_overrides(cfg, mb_args.load_extra_config())
+    return cfg
+
+
+def build_sft_config(
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "TrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    mb_args: "MegatronBridgeArguments",
+    dataset_root: str,
+    pretrained_checkpoint: str,
+    num_train_samples: int,
+):
+    from megatron.bridge import AutoBridge
+    from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs
+
+    micro_batch_size, global_batch_size, train_iters = _compute_train_schedule(
+        training_args, mb_args, num_train_samples
+    )
+    cfg = _create_base_config(
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        train_iters=train_iters,
+        micro_batch_size=micro_batch_size,
+        global_batch_size=global_batch_size,
+        mb_args=mb_args,
+        is_sft=True,
+    )
+
+    (
+        _CheckpointConfig,
+        _ConfigContainer,
+        _DistributedInitConfig,
+        FinetuningDatasetConfig,
+        _GPTDatasetConfig,
+        _LoggerConfig,
+        _RNGConfig,
+        _TokenizerConfig,
+        _TrainingConfig,
+        _ValidationConfig,
+    ) = _import_training_config()
+    bridge = AutoBridge.from_hf_pretrained(
+        model_args.model_name_or_path,
+        trust_remote_code=model_args.trust_remote_code,
+    )
+    cfg.model = bridge.to_megatron_provider(load_weights=False)
+    _apply_model_parallelism(cfg.model, mb_args)
+    _apply_fusion_safety(cfg.model)
+    if hasattr(cfg.model, "seq_length"):
+        cfg.model.seq_length = data_args.cutoff_len
+
+    cfg.tokenizer.tokenizer_model = model_args.model_name_or_path
+
+    from .dataset_export import get_sft_dataset_kwargs
+
+    ensure_create_sft_dataset_applies_chat_template()
+    dataset_kwargs = get_sft_dataset_kwargs(
+        tokenizer_path=model_args.model_name_or_path,
+        trust_remote_code=model_args.trust_remote_code,
+        template_name=data_args.template,
+    )
+    packed_sequence_specs = None
+    if mb_args.use_packed_sequences:
+        pad_seq_to_mult = mb_args.context_parallel_size * 2 if mb_args.context_parallel_size > 1 else 1
+        packed_sequence_specs = PackedSequenceSpecs(
+            packed_sequence_size=data_args.cutoff_len,
+            pad_seq_to_mult=pad_seq_to_mult,
+        )
+        dataset_kwargs["pad_to_max_length"] = True
+
+    cfg.dataset = _build_finetuning_dataset_config(
+        FinetuningDatasetConfig,
+        dataset_root=dataset_root,
+        seq_length=data_args.cutoff_len,
+        seed=training_args.seed,
+        num_workers=data_args.preprocessing_num_workers,
+        do_validation=data_args.val_size > 0 or data_args.eval_dataset is not None,
+        dataset_kwargs=dataset_kwargs,
+        packed_sequence_specs=packed_sequence_specs,
+        disable_shuffling=getattr(finetuning_args, "disable_shuffling", False),
+    )
+    cfg.checkpoint.pretrained_checkpoint = pretrained_checkpoint
+
+    _apply_context_parallel_finetuning_requirements(cfg, mb_args)
+    _apply_extra_overrides(cfg, mb_args.load_extra_config())
+    return cfg

src/llamafactory/train/megatron_bridge/dataset_export.py

@@ -0,0 +1,344 @@
+# 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 inspect
+import json
+import os
+import re
+import typing
+from typing import TYPE_CHECKING, Any, Optional
+
+from ...extras.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from datasets import Dataset, IterableDataset
+
+
+logger = get_logger(__name__)
+_GENERATION_REGEX = re.compile(r"\{%-?\s+generation\s+-?%\}")
+_END_GENERATION_REGEX = re.compile(r"\{%-?\s+endgeneration\s+-?%\}")
+_ASSISTANT_ELIF_REGEX = re.compile(
+    r"(\{%\s*elif\s+message\['role'\]\s*==\s*'assistant'\s*%\})"
+    r"(.*?)"
+    r"(\{%\s*endif\s*%\})",
+    flags=re.DOTALL,
+)
+
+
+def supports_hf_chat_template() -> bool:
+    r"""Return whether the installed Megatron Bridge supports HF chat templates."""
+    try:
+        from megatron.bridge.data.datasets.sft import GPTSFTChatDataset
+
+        return "use_hf_tokenizer_chat_template" in inspect.signature(GPTSFTChatDataset.__init__).parameters
+    except Exception:
+        return False
+
+
+def _inject_generation_block(chat_template: str) -> str:
+    r"""Wrap assistant content with ``{% generation %}`` when missing."""
+    if _GENERATION_REGEX.search(chat_template):
+        return chat_template
+
+    match = _ASSISTANT_ELIF_REGEX.search(chat_template)
+    if match is None:
+        raise ValueError(
+            "Cannot inject {% generation %} into chat template: "
+            "no `{% elif message['role'] == 'assistant' %}` block found."
+        )
+
+    body = match.group(2)
+    if _END_GENERATION_REGEX.search(body):
+        return chat_template
+
+    patched = (
+        chat_template[: match.start()]
+        + match.group(1)
+        + "{% generation %}"
+        + body
+        + "{% endgeneration %}"
+        + match.group(3)
+        + chat_template[match.end() :]
+    )
+    if not _GENERATION_REGEX.search(patched):
+        raise ValueError("Failed to inject {% generation %} into chat template.")
+    return patched
+
+
+def build_chat_template_with_generation(
+    tokenizer_path: str | None,
+    trust_remote_code: bool = False,
+    template_name: str | None = None,
+) -> str | None:
+    r"""Build a chat template that supports assistant-only loss masks.
+
+    Prefers the LLaMA-Factory registered template (same formatting as HF Trainer),
+    then falls back to patching the tokenizer's native chat template.
+    """
+    if not tokenizer_path:
+        return None
+
+    try:
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=trust_remote_code)
+    except Exception as exc:
+        logger.warning_rank0(f"Failed to load tokenizer for chat template patching: {exc}")
+        return None
+
+    if template_name:
+        try:
+            from ...data.template import TEMPLATES
+
+            template = TEMPLATES.get(template_name)
+            if template is not None:
+                return _inject_generation_block(template._get_jinja_template(tokenizer))
+        except Exception as exc:
+            logger.warning_rank0(f"Failed to build LLaMA-Factory chat template '{template_name}': {exc}")
+
+    native = tokenizer.chat_template
+    if not isinstance(native, str) or not native:
+        return None
+
+    try:
+        return _inject_generation_block(native)
+    except Exception as exc:
+        logger.warning_rank0(f"Failed to inject {{% generation %}} into native chat template: {exc}")
+        return None
+
+
+def tokenizer_supports_hf_chat_template(
+    tokenizer_path: str | None,
+    trust_remote_code: bool = False,
+    template_name: str | None = None,
+) -> bool:
+    r"""Return whether Megatron Bridge can use HF chat templates for this tokenizer."""
+    if not supports_hf_chat_template() or not tokenizer_path:
+        return False
+
+    try:
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=trust_remote_code)
+        if _GENERATION_REGEX.search(tokenizer.chat_template or ""):
+            return True
+    except Exception as exc:
+        logger.warning_rank0(f"Failed to inspect tokenizer chat template: {exc}")
+        return False
+
+    return build_chat_template_with_generation(tokenizer_path, trust_remote_code, template_name) is not None
+
+
+def get_sft_dataset_kwargs(
+    tokenizer_path: str | None = None,
+    trust_remote_code: bool = False,
+    template_name: str | None = None,
+) -> dict[str, Any]:
+    r"""Return dataset kwargs compatible with the installed Megatron Bridge version."""
+    kwargs: dict[str, Any] = {"chat": True}
+    if not tokenizer_supports_hf_chat_template(tokenizer_path, trust_remote_code, template_name):
+        return kwargs
+
+    chat_template = build_chat_template_with_generation(tokenizer_path, trust_remote_code, template_name)
+    if chat_template is None:
+        return kwargs
+
+    kwargs["use_hf_tokenizer_chat_template"] = True
+    kwargs["chat_template"] = chat_template
+    # Chat templates already include BOS/EOS / turn markers.
+    kwargs["add_bos"] = False
+    kwargs["add_eos"] = False
+    logger.info_rank0(
+        "Using HuggingFace chat template with {% generation %} for Megatron Bridge SFT "
+        f"(template={template_name or 'native'})."
+    )
+    return kwargs
+
+
+def get_finetuning_dataloader_type(disable_shuffling: bool = False) -> str:
+    r"""Return a finetuning dataloader type supported by the installed Megatron Bridge.
+
+    Prefer sequential ``batch``/``single`` samplers. When shuffling is disabled for
+    HF loss comparison, avoid the random ``cyclic`` sampler.
+    """
+    try:
+        from megatron.bridge.training.config import FinetuningDatasetConfig
+
+        field = FinetuningDatasetConfig.__dataclass_fields__.get("dataloader_type")
+        if field is None:
+            return "single"
+
+        choices: set[str] = set()
+        for arg in typing.get_args(field.type):
+            for choice in typing.get_args(arg):
+                if isinstance(choice, str):
+                    choices.add(choice)
+
+        if disable_shuffling:
+            if "batch" in choices:
+                return "batch"
+            if "single" in choices:
+                return "single"
+            return next(iter(choices), "single")
+
+        if "batch" in choices:
+            return "batch"
+        if "single" in choices:
+            return "single"
+        return next(iter(choices), "single")
+    except Exception:
+        return "single"
+
+
+def _role_to_sharegpt(role: str) -> str:
+    mapping = {"user": "User", "assistant": "Assistant", "system": "System"}
+    return mapping.get(role, role.capitalize())
+
+
+def _example_to_record(
+    example: dict[str, Any], stage: str = "sft", use_messages_format: bool = True
+) -> dict[str, Any] | None:
+    r"""Convert an aligned LLaMA-Factory example to Megatron Bridge JSONL format."""
+    if example.get("text") is not None:
+        return {"text": example["text"]}
+
+    prompt = example.get("_prompt")
+    if stage == "pt":
+        if not prompt:
+            return None
+        return {"text": prompt[0]["content"]}
+
+    response = example.get("_response")
+    if not prompt or not response:
+        return None
+
+    if use_messages_format:
+        messages = []
+        system = example.get("_system")
+        if system:
+            messages.append({"role": "system", "content": system})
+        for message in prompt:
+            messages.append({"role": message["role"], "content": message["content"]})
+        for message in response:
+            messages.append({"role": message["role"], "content": message["content"]})
+        record: dict[str, Any] = {"messages": messages}
+        tools = example.get("_tools")
+        if tools:
+            record["tools"] = tools
+        return record
+
+    conversations = []
+    for message in prompt:
+        conversations.append({"from": _role_to_sharegpt(message["role"]), "value": message["content"]})
+    for message in response:
+        conversations.append({"from": _role_to_sharegpt(message["role"]), "value": message["content"]})
+
+    return {
+        "system": example.get("_system") or "",
+        "conversations": conversations,
+        "mask": "User",
+    }
+
+
+def _remove_stale_memmap_index(path: str) -> None:
+    r"""Remove cached memmap index files after rewriting a JSONL dataset."""
+    for suffix in (".idx.npy", ".idx.info"):
+        index_path = path + suffix
+        if os.path.exists(index_path):
+            os.remove(index_path)
+            logger.info_rank0(f"Removed stale Megatron dataset index: {index_path}")
+
+
+def _write_jsonl(
+    path: str,
+    dataset: "Dataset | IterableDataset",
+    stage: str = "sft",
+    use_messages_format: bool = True,
+) -> int:
+    count = 0
+    parent = os.path.dirname(path)
+    if parent:
+        os.makedirs(parent, exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        for example in dataset:
+            record = _example_to_record(example, stage=stage, use_messages_format=use_messages_format)
+            if record is None:
+                continue
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+            count += 1
+    _remove_stale_memmap_index(path)
+    return count
+
+
+def export_dataset_for_megatron_bridge(
+    train_dataset: "Dataset | IterableDataset",
+    output_dir: str,
+    eval_dataset: Optional["Dataset | IterableDataset | dict[str, Dataset]"] = None,
+    val_size: float = 0.0,
+    seed: int = 42,
+    stage: str = "sft",
+    model_name_or_path: str | None = None,
+    trust_remote_code: bool = False,
+    template_name: str | None = None,
+) -> str:
+    r"""Export aligned LLaMA-Factory datasets to Megatron Bridge JSONL files."""
+    os.makedirs(output_dir, exist_ok=True)
+    use_messages_format = stage != "sft" or tokenizer_supports_hf_chat_template(
+        model_name_or_path,
+        trust_remote_code=trust_remote_code,
+        template_name=template_name,
+    )
+    if stage == "sft" and not use_messages_format:
+        logger.info_rank0(
+            "Cannot enable HuggingFace chat template for Megatron Bridge; "
+            "exporting ShareGPT conversations for legacy preprocessing."
+        )
+
+    if val_size > 0 and eval_dataset is None:
+        split = train_dataset.train_test_split(test_size=val_size, seed=seed)
+        train_dataset = split["train"]
+        eval_dataset = split["test"]
+
+    train_path = os.path.join(output_dir, "training.jsonl")
+    train_count = _write_jsonl(
+        train_path,
+        train_dataset,
+        stage=stage,
+        use_messages_format=use_messages_format,
+    )
+    logger.info_rank0(f"Exported {train_count} training samples to {train_path}.")
+
+    if isinstance(eval_dataset, dict):
+        for name, dataset in eval_dataset.items():
+            split_name = "validation" if name == "validation" else name
+            eval_path = os.path.join(output_dir, f"{split_name}.jsonl")
+            eval_count = _write_jsonl(
+                eval_path,
+                dataset,
+                stage=stage,
+                use_messages_format=use_messages_format,
+            )
+            logger.info_rank0(f"Exported {eval_count} {split_name} samples to {eval_path}.")
+    elif eval_dataset is not None:
+        eval_path = os.path.join(output_dir, "validation.jsonl")
+        eval_count = _write_jsonl(
+            eval_path,
+            eval_dataset,
+            stage=stage,
+            use_messages_format=use_messages_format,
+        )
+        logger.info_rank0(f"Exported {eval_count} validation samples to {eval_path}.")
+
+    return output_dir

src/llamafactory/train/megatron_bridge/workflow.py

@@ -0,0 +1,379 @@
+# 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 os
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Optional
+
+from transformers import AutoConfig as HfAutoConfig
+
+from ...data.data_utils import split_dataset
+from ...data.loader import _get_merged_dataset
+from ...extras.constants import MEGATRON_BRIDGE_SUPPORTED_MODELS
+from ...extras.logging import get_logger
+from ...extras.packages import is_megatron_bridge_available
+from .config_builder import (
+    _apply_fusion_safety,
+    build_pretrain_config,
+    build_sft_config,
+    ensure_megatron_pretrained_checkpoint,
+)
+from .dataset_export import export_dataset_for_megatron_bridge
+
+
+if TYPE_CHECKING:
+    from transformers import TrainerCallback
+
+    from ...hparams import (
+        DataArguments,
+        FinetuningArguments,
+        MegatronBridgeArguments,
+        ModelArguments,
+        TrainingArguments,
+    )
+
+
+logger = get_logger(__name__)
+
+
+def _check_model_support(model_args: "ModelArguments") -> None:
+    r"""Ensure the HF ``model_type`` is covered by the Megatron Bridge PT/SFT path."""
+    config = HfAutoConfig.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
+    )
+    model_type = getattr(config, "model_type", None)
+    if model_type not in MEGATRON_BRIDGE_SUPPORTED_MODELS:
+        raise ValueError(
+            f"Model type `{model_type}` is not supported by the Megatron Bridge PT/SFT path. "
+            f"Supported model types: {sorted(MEGATRON_BRIDGE_SUPPORTED_MODELS)}. "
+            "Multimodal / audio / omni models are not enabled in v0."
+        )
+
+
+def _run_on_main_process(training_args: "TrainingArguments", work: Callable[[], None], sync_dir: str) -> None:
+    r"""Run ``work`` only on global rank 0, then synchronize other ranks.
+
+    Prefer ``torch.distributed.barrier`` when the process group is already initialized;
+    otherwise fall back to a file flag under ``sync_dir`` so non-main ranks wait for
+    shared filesystem writes (e.g. dataset export) to finish.
+    """
+    done_file = os.path.join(sync_dir, ".main_process_done")
+    is_main = getattr(training_args, "process_index", 0) == 0
+    wait_start = time.time()
+
+    import torch.distributed as dist
+
+    dist_ready = dist.is_available() and dist.is_initialized()
+    if is_main:
+        os.makedirs(sync_dir, exist_ok=True)
+        if os.path.isfile(done_file):
+            os.remove(done_file)
+        work()
+        with open(done_file, "w", encoding="utf-8") as f:
+            f.write("done")
+
+    if dist_ready:
+        dist.barrier()
+    elif not is_main:
+        while True:
+            if os.path.isfile(done_file) and os.path.getmtime(done_file) >= wait_start - 1.0:
+                break
+            time.sleep(0.5)
+
+
+def _check_backend_available() -> None:
+    if not is_megatron_bridge_available():
+        raise ImportError(
+            "megatron-bridge is not installed. "
+            "Please install it with `pip install --no-build-isolation megatron-bridge` "
+            "or use the NeMo Framework container."
+        )
+    _patch_dataset_helper_compilation()
+    _patch_dist_checkpoint_preload()
+
+
+def _patch_dist_checkpoint_preload() -> None:
+    r"""Use blocking GPU->CPU copies when saving distributed checkpoints.
+
+    Megatron's default ``non_blocking=True`` preload can raise ``cudaErrorInvalidValue``
+    on some GPUs (e.g. V100) when saving distributed optimizer shards, because pinned
+    host memory allocation or async D2H transfer may fail under memory pressure.
+    """
+    from megatron.core.dist_checkpointing.strategies import filesystem_async
+
+    if getattr(filesystem_async.FileSystemWriterAsync.preload_tensors, "_llamafactory_patched", False):
+        return
+
+    original_preload = filesystem_async.FileSystemWriterAsync.preload_tensors
+
+    @staticmethod
+    def preload_tensors(write_buckets, non_blocking=True):
+        return original_preload(write_buckets, non_blocking=False)
+
+    preload_tensors._llamafactory_patched = True
+    filesystem_async.FileSystemWriterAsync.preload_tensors = preload_tensors
+    logger.info_rank0("Patched Megatron dist checkpoint preload to use blocking GPU->CPU copies.")
+
+
+def _patch_dataset_helper_compilation() -> None:
+    r"""Skip make-based helper compilation when the pybind extension is prebuilt.
+
+    Pip-installed megatron-core already ships helpers_cpp, but compile_helpers()
+    still invokes make and fails when no Makefile is present.
+    """
+    from megatron.core.datasets import utils as dataset_utils
+
+    if getattr(dataset_utils.compile_helpers, "_llamafactory_patched", False):
+        return
+
+    try:
+        import megatron.core.datasets.helpers_cpp  # noqa: F401
+    except ImportError:
+        return
+
+    def compile_helpers():
+        import megatron.core.datasets.helpers_cpp  # noqa: F401
+
+    compile_helpers._llamafactory_patched = True
+    dataset_utils.compile_helpers = compile_helpers
+    logger.info_rank0("Using prebuilt megatron.core.datasets.helpers_cpp; skipping make compilation.")
+
+
+def _load_aligned_datasets(
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "TrainingArguments",
+    stage: str,
+):
+    dataset = _get_merged_dataset(data_args.dataset, model_args, data_args, training_args, stage)
+    eval_dataset = _get_merged_dataset(
+        data_args.eval_dataset,
+        model_args,
+        data_args,
+        training_args,
+        stage,
+        return_dict=data_args.eval_on_each_dataset,
+    )
+    train_dict, eval_dict = split_dataset(dataset, eval_dataset, data_args, seed=training_args.seed)
+    return train_dict.get("train"), eval_dict
+
+
+def _latest_iter_checkpoint_dir(output_dir: str) -> Optional[str]:
+    r"""Return the latest ``iter_*`` directory under ``output_dir``, if any.
+
+    ``export_adapter_ckpt`` needs the iteration directory that holds the
+    distributed checkpoint payload (``.distcp`` / ``run_config.yaml``), not the
+    parent run directory.
+    """
+    if not os.path.isdir(output_dir):
+        return None
+
+    # Already pointing at an iteration directory.
+    if os.path.isfile(os.path.join(output_dir, "run_config.yaml")) or os.path.exists(
+        os.path.join(output_dir, ".metadata")
+    ):
+        return output_dir
+
+    iter_dirs = [
+        name
+        for name in os.listdir(output_dir)
+        if name.startswith("iter_") and os.path.isdir(os.path.join(output_dir, name))
+    ]
+    if not iter_dirs:
+        return None
+
+    def _iter_number(name: str) -> int:
+        try:
+            return int(name.replace("iter_", ""))
+        except ValueError:
+            return -1
+
+    latest = max(iter_dirs, key=_iter_number)
+    return os.path.join(output_dir, latest)
+
+
+def _checkpoint_uses_peft(checkpoint_dir: str) -> bool:
+    r"""Whether the Megatron checkpoint was saved with a PEFT (e.g. LoRA) config."""
+    cfg_path = os.path.join(checkpoint_dir, "run_config.yaml")
+    if not os.path.isfile(cfg_path):
+        return False
+
+    try:
+        import yaml
+
+        with open(cfg_path, encoding="utf-8") as f:
+            cfg = yaml.safe_load(f) or {}
+        return isinstance(cfg, dict) and bool(cfg.get("peft"))
+    except Exception:
+        return False
+
+
+def _with_fusion_safe_provider(bridge):
+    r"""Wrap ``to_megatron_provider`` so export paths disable missing APEX fusion.
+
+    Training already applies ``_apply_fusion_safety``, but AutoBridge export helpers
+    (e.g. ``export_adapter_ckpt``) rebuild a provider with the default
+    ``gradient_accumulation_fusion=True`` whenever TransformerEngine is importable,
+    even if ``fused_weight_gradient_mlp_cuda`` is absent.
+    """
+    original = bridge.to_megatron_provider
+
+    def _to_megatron_provider(*args, **kwargs):
+        provider = original(*args, **kwargs)
+        _apply_fusion_safety(provider)
+        return provider
+
+    bridge.to_megatron_provider = _to_megatron_provider  # type: ignore[method-assign]
+    return bridge
+
+
+def _maybe_export_hf_checkpoint(
+    model_args: "ModelArguments",
+    mb_args: "MegatronBridgeArguments",
+    output_dir: str,
+) -> None:
+    if not mb_args.export_hf_on_finish or not training_args_should_save(output_dir):
+        return
+
+    import torch.distributed as dist
+    from megatron.bridge import AutoBridge
+
+    checkpoint_dir = _latest_iter_checkpoint_dir(output_dir)
+    if checkpoint_dir is None:
+        logger.warning_rank0(f"No Megatron iteration checkpoint found under {output_dir}; skip HF export.")
+        return
+
+    export_dir = os.path.join(output_dir, "hf_export")
+    bridge = _with_fusion_safe_provider(
+        AutoBridge.from_hf_pretrained(
+            model_args.model_name_or_path,
+            trust_remote_code=model_args.trust_remote_code,
+        )
+    )
+
+    # LoRA / PEFT checkpoints only store adapter weights. Loading them as a full
+    # model raises KeyError for base tensors such as linear_proj.weight.
+    if _checkpoint_uses_peft(checkpoint_dir):
+        logger.info_rank0(f"Exporting LoRA adapter to Hugging Face PEFT format at {export_dir}...")
+        bridge.export_adapter_ckpt(peft_checkpoint=checkpoint_dir, output_path=export_dir)
+        return
+
+    logger.info_rank0(f"Exporting Megatron checkpoint to Hugging Face format at {export_dir}...")
+    if dist.is_initialized():
+        # export_ckpt() always creates a fresh single-process gloo group, which fails
+        # when torchrun has already initialized NCCL for training.
+        megatron_model = bridge.load_megatron_model(output_dir)
+        bridge.save_hf_pretrained(megatron_model, export_dir)
+    else:
+        bridge.export_ckpt(megatron_path=output_dir, hf_path=export_dir)
+
+
+def training_args_should_save(output_dir: str) -> bool:
+    return os.path.isdir(output_dir) and bool(os.listdir(output_dir))
+
+
+def run_pt(
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "TrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    mb_args: "MegatronBridgeArguments",
+    callbacks: Optional[list["TrainerCallback"]] = None,
+):
+    if callbacks:
+        logger.warning_rank0("Megatron Bridge does not support Trainer callbacks yet; ignoring provided callbacks.")
+    _check_backend_available()
+    _check_model_support(model_args)
+    from megatron.bridge.training.gpt_step import forward_step
+    from megatron.bridge.training.pretrain import pretrain
+
+    train_dataset, eval_dict = _load_aligned_datasets(model_args, data_args, training_args, "pt")
+    dataset_dir = os.path.join(training_args.output_dir, "mb_dataset")
+
+    def _export_pt_dataset() -> None:
+        export_dataset_for_megatron_bridge(
+            train_dataset=train_dataset,
+            output_dir=dataset_dir,
+            eval_dataset=eval_dict.get("validation") if eval_dict else None,
+            val_size=data_args.val_size,
+            seed=training_args.seed,
+            stage="pt",
+        )
+
+    _run_on_main_process(training_args, _export_pt_dataset, dataset_dir)
+
+    cfg = build_pretrain_config(
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        mb_args=mb_args,
+        dataset_path=os.path.join(dataset_dir, "training.jsonl"),
+        num_train_samples=len(train_dataset),
+    )
+    pretrain(cfg, forward_step)
+    _maybe_export_hf_checkpoint(model_args, mb_args, training_args.output_dir)
+
+
+def run_sft(
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "TrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    mb_args: "MegatronBridgeArguments",
+    callbacks: Optional[list["TrainerCallback"]] = None,
+):
+    if callbacks:
+        logger.warning_rank0("Megatron Bridge does not support Trainer callbacks yet; ignoring provided callbacks.")
+    _check_backend_available()
+    _check_model_support(model_args)
+    from megatron.bridge.training.finetune import finetune
+    from megatron.bridge.training.gpt_step import forward_step
+
+    train_dataset, eval_dict = _load_aligned_datasets(model_args, data_args, training_args, "sft")
+    dataset_dir = os.path.join(training_args.output_dir, "mb_dataset")
+
+    def _export_sft_dataset() -> None:
+        export_dataset_for_megatron_bridge(
+            train_dataset=train_dataset,
+            output_dir=dataset_dir,
+            eval_dataset=eval_dict or None,
+            val_size=data_args.val_size if not eval_dict else 0.0,
+            seed=training_args.seed,
+            stage="sft",
+            model_name_or_path=model_args.model_name_or_path,
+            trust_remote_code=model_args.trust_remote_code,
+            template_name=data_args.template,
+        )
+
+    _run_on_main_process(training_args, _export_sft_dataset, dataset_dir)
+
+    pretrained_checkpoint = ensure_megatron_pretrained_checkpoint(
+        model_args=model_args,
+        mb_args=mb_args,
+        output_dir=training_args.output_dir,
+    )
+    cfg = build_sft_config(
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        mb_args=mb_args,
+        dataset_root=dataset_dir,
+        pretrained_checkpoint=pretrained_checkpoint,
+        num_train_samples=len(train_dataset),
+    )
+    finetune(cfg, forward_step_func=forward_step)
+    _maybe_export_hf_checkpoint(model_args, mb_args, training_args.output_dir)

src/llamafactory/train/tuner.py

@@ -27,6 +27,7 @@ from ..extras.misc import find_available_port, get_device_name, get_torch_device
 from ..extras.packages import (
     is_hyper_parallel_available,
     is_mcore_adapter_available,
+    is_megatron_bridge_available,
     is_ray_available,
     is_transformers_version_greater_than,
 )
@@ -100,6 +101,24 @@ def _training_function(config: dict[str, Any]) -> None:
 
             run_sft_hp(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
 
+    elif finetuning_args.stage in ["pt", "sft"] and finetuning_args.use_megatron_bridge:
+        if not is_megatron_bridge_available():
+            raise ImportError(
+                "megatron-bridge is not installed. "
+                "Please install it with `pip install --no-build-isolation megatron-bridge`."
+            )
+        mb_args = finetuning_args.megatron_bridge_args
+        if mb_args is None:
+            raise ValueError("Megatron Bridge arguments are missing. Please set USE_MEGATRON_BRIDGE=1.")
+        if finetuning_args.stage == "pt":
+            from .megatron_bridge import run_pt as run_pt_mb
+
+            run_pt_mb(model_args, data_args, training_args, finetuning_args, mb_args, callbacks)
+        else:
+            from .megatron_bridge import run_sft as run_sft_mb
+
+            run_sft_mb(model_args, data_args, training_args, finetuning_args, mb_args, callbacks)
+
     elif finetuning_args.stage in ["pt", "sft", "dpo"] and finetuning_args.use_mca:
         if not is_mcore_adapter_available():
             raise ImportError("mcore_adapter is not installed. Please install it with `pip install mcore-adapter`.")

tests/train/megatron_bridge/conftest.py

@@ -0,0 +1,112 @@
+# 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.
+
+"""Fixtures for Megatron Bridge GPU unit tests."""
+
+import os
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+from llamafactory.extras.packages import is_megatron_bridge_available
+
+
+pytestmark = pytest.mark.skipif(
+    not is_megatron_bridge_available(),
+    reason="megatron-bridge is not installed",
+)
+
+
+@pytest.fixture(scope="session")
+def mb_model_path() -> str:
+    r"""Return the local Hugging Face model path used by Megatron Bridge tests."""
+    model_path = os.getenv("TINY_LLAMA3", "llamafactory/tiny-random-Llama-3")
+    return model_path
+
+
+@pytest.fixture
+def mb_output_dir(tmp_path: Path) -> Path:
+    r"""Create a writable output directory for Megatron Bridge tests."""
+    output_dir = tmp_path / "output"
+    output_dir.mkdir(parents=True)
+    return output_dir
+
+
+@pytest.fixture
+def mb_training_args_factory(mb_output_dir: Path, mb_model_path: str):
+    r"""Build LLaMA-Factory dataclass arguments for Megatron Bridge tests."""
+    patchers: list[mock._patch] = []
+
+    def _factory(
+        *,
+        tensor_model_parallel_size: int = 1,
+        pipeline_model_parallel_size: int = 1,
+        context_parallel_size: int = 1,
+        sequence_parallel: bool = False,
+        use_distributed_optimizer: bool = False,
+        num_train_samples: int = 100,
+        cutoff_len: int = 512,
+        output_dir: Path | None = None,
+    ):
+        from llamafactory.hparams.data_args import DataArguments
+        from llamafactory.hparams.finetuning_args import FinetuningArguments
+        from llamafactory.hparams.megatron_bridge_args import MegatronBridgeArguments
+        from llamafactory.hparams.model_args import ModelArguments
+        from llamafactory.hparams.training_args import TrainingArguments
+
+        resolved_output_dir = output_dir or mb_output_dir
+        model_args = ModelArguments(model_name_or_path=mb_model_path)
+        data_args = DataArguments(
+            dataset="alpaca_en_demo",
+            template="llama3",
+            cutoff_len=cutoff_len,
+            preprocessing_num_workers=1,
+        )
+        training_args = TrainingArguments(
+            output_dir=str(resolved_output_dir),
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=1,
+            num_train_epochs=1,
+            learning_rate=5e-6,
+            logging_steps=1,
+            save_steps=1000000,
+            report_to="none",
+        )
+        world_size = tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size
+        patcher = mock.patch.object(
+            type(training_args),
+            "world_size",
+            new_callable=mock.PropertyMock,
+            return_value=world_size,
+        )
+        patchers.append(patcher)
+        patcher.start()
+        finetuning_args = FinetuningArguments(stage="sft", finetuning_type="full")
+        mb_args = MegatronBridgeArguments(
+            tensor_model_parallel_size=tensor_model_parallel_size,
+            pipeline_model_parallel_size=pipeline_model_parallel_size,
+            context_parallel_size=context_parallel_size,
+            sequence_parallel=sequence_parallel,
+            use_distributed_optimizer=use_distributed_optimizer,
+            overlap_param_gather=False,
+            overlap_grad_reduce=False,
+            export_hf_on_finish=False,
+        )
+        return model_args, data_args, training_args, finetuning_args, mb_args, num_train_samples
+
+    yield _factory
+
+    for patcher in patchers:
+        patcher.stop()

tests/train/megatron_bridge/test_config_builder.py

@@ -0,0 +1,206 @@
+# 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 types import SimpleNamespace
+
+import pytest
+
+from llamafactory.hparams.megatron_bridge_args import MegatronBridgeArguments
+from llamafactory.train.megatron_bridge.config_builder import (
+    _apply_extra_overrides,
+    _apply_fusion_safety,
+    _apply_model_parallelism,
+    _compute_train_schedule,
+    _resolve_warmup_steps,
+    _should_resume_checkpoint,
+    build_pretrain_config,
+    build_sft_config,
+)
+
+
+class _DummyProvider:
+    tensor_model_parallel_size = 1
+    pipeline_model_parallel_size = 1
+    context_parallel_size = 1
+    sequence_parallel = False
+    gradient_accumulation_fusion = True
+
+
+def test_compute_train_schedule_tp1():
+    training_args = SimpleNamespace(
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=4,
+        num_train_epochs=1,
+        max_steps=-1,
+        world_size=2,
+    )
+    mb_args = MegatronBridgeArguments(tensor_model_parallel_size=1)
+    micro_batch, global_batch, train_iters = _compute_train_schedule(training_args, mb_args, num_train_samples=32)
+    assert micro_batch == 2
+    assert global_batch == 16
+    assert train_iters == 2
+
+
+def test_compute_train_schedule_tp2():
+    training_args = SimpleNamespace(
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=1,
+        num_train_epochs=1,
+        max_steps=-1,
+        world_size=2,
+    )
+    mb_args = MegatronBridgeArguments(
+        tensor_model_parallel_size=2,
+        pipeline_model_parallel_size=1,
+        context_parallel_size=1,
+    )
+    micro_batch, global_batch, train_iters = _compute_train_schedule(training_args, mb_args, num_train_samples=10)
+    assert micro_batch == 1
+    assert global_batch == 1
+    assert train_iters == 10
+
+
+def test_compute_train_schedule_max_steps():
+    training_args = SimpleNamespace(
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=1,
+        num_train_epochs=10,
+        max_steps=5,
+        world_size=1,
+    )
+    mb_args = MegatronBridgeArguments()
+    _, _, train_iters = _compute_train_schedule(training_args, mb_args, num_train_samples=1000)
+    assert train_iters == 5
+
+
+def test_apply_model_parallelism():
+    provider = _DummyProvider()
+    mb_args = MegatronBridgeArguments(
+        tensor_model_parallel_size=2,
+        pipeline_model_parallel_size=1,
+        context_parallel_size=1,
+        sequence_parallel=True,
+        bias_activation_fusion=True,
+        moe_token_dispatcher_type="alltoall",
+        recompute_granularity="full",
+        recompute_method="uniform",
+        recompute_num_layers=2,
+    )
+    provider.bias_activation_fusion = False
+    provider.moe_token_dispatcher_type = "allgather"
+    provider.recompute_granularity = None
+    provider.recompute_method = None
+    provider.recompute_num_layers = None
+    _apply_model_parallelism(provider, mb_args)
+    assert provider.tensor_model_parallel_size == 2
+    assert provider.sequence_parallel is True
+    assert provider.bias_activation_fusion is True
+    assert provider.moe_token_dispatcher_type == "alltoall"
+    assert provider.recompute_granularity == "full"
+    assert provider.recompute_method == "uniform"
+    assert provider.recompute_num_layers == 2
+
+
+def test_megatron_bridge_args_rejects_invalid_moe_dispatcher():
+    with pytest.raises(ValueError, match="moe_token_dispatcher_type"):
+        MegatronBridgeArguments(moe_token_dispatcher_type="invalid")  # type: ignore[arg-type]
+
+
+def test_apply_fusion_safety_disables_missing_apex():
+    provider = _DummyProvider()
+    _apply_fusion_safety(provider)
+    assert provider.gradient_accumulation_fusion is False
+
+
+def test_megatron_bridge_args_sequence_parallel_requires_tp():
+    with pytest.raises(ValueError, match="sequence_parallel"):
+        MegatronBridgeArguments(sequence_parallel=True, tensor_model_parallel_size=1)
+
+
+def test_resolve_warmup_steps_keeps_absolute_warmup():
+    training_args = SimpleNamespace(warmup_steps=10, warmup_ratio=0.0)
+    assert _resolve_warmup_steps(training_args, train_iters=5) == 10
+
+
+def test_resolve_decay_iters_expands_for_warmup():
+    from llamafactory.train.megatron_bridge.config_builder import _resolve_decay_iters
+
+    assert _resolve_decay_iters(train_iters=5, warmup_steps=10) == 11
+    assert _resolve_decay_iters(train_iters=20, warmup_steps=10) == 20
+    assert _resolve_decay_iters(train_iters=5, warmup_steps=0) == 5
+
+
+def test_megatron_bridge_args_extra_config_strips_whitespace():
+    args = MegatronBridgeArguments(extra_config='  {"train.train_iters": 3}  ')
+    assert args.extra_config == {"train.train_iters": 3}
+
+
+def test_apply_extra_overrides_nested_path():
+    cfg = SimpleNamespace(train=SimpleNamespace(nested=SimpleNamespace(value=1)))
+    _apply_extra_overrides(cfg, {"train.nested.value": 9})
+    assert cfg.train.nested.value == 9
+
+
+def test_should_resume_checkpoint_respects_overwrite(tmp_path):
+    tracker = tmp_path / "latest_checkpointed_iteration.txt"
+    tracker.write_text("5")
+    assert _should_resume_checkpoint(SimpleNamespace(output_dir=str(tmp_path), overwrite_output_dir=False))
+    assert not _should_resume_checkpoint(SimpleNamespace(output_dir=str(tmp_path), overwrite_output_dir=True))
+    assert not _should_resume_checkpoint(
+        SimpleNamespace(output_dir=str(tmp_path / "missing"), overwrite_output_dir=False)
+    )
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_build_sft_config_on_gpu(mb_training_args_factory, mb_output_dir):
+    model_args, data_args, training_args, finetuning_args, mb_args, num_train_samples = mb_training_args_factory(
+        tensor_model_parallel_size=1,
+        sequence_parallel=False,
+    )
+    cfg = build_sft_config(
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        mb_args=mb_args,
+        dataset_root=str(mb_output_dir / "dataset"),
+        pretrained_checkpoint=str(mb_output_dir / "pretrained"),
+        num_train_samples=num_train_samples,
+    )
+    assert cfg.model.tensor_model_parallel_size == 1
+    assert cfg.model.seq_length == data_args.cutoff_len
+    assert cfg.train.train_iters == num_train_samples
+    assert cfg.dataset.dataset_root == str(mb_output_dir / "dataset")
+    assert cfg.tokenizer.tokenizer_model == model_args.model_name_or_path
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_build_pretrain_config_on_gpu(mb_training_args_factory, mb_output_dir):
+    model_args, data_args, training_args, finetuning_args, mb_args, num_train_samples = mb_training_args_factory(
+        tensor_model_parallel_size=1,
+    )
+    finetuning_args.stage = "pt"
+    dataset_path = str(mb_output_dir / "dataset" / "training.jsonl")
+    cfg = build_pretrain_config(
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        mb_args=mb_args,
+        dataset_path=dataset_path,
+        num_train_samples=num_train_samples,
+    )
+    assert cfg.model.tensor_model_parallel_size == 1
+    assert cfg.dataset.blend[0][0] == dataset_path
+    assert cfg.train.train_iters == num_train_samples

tests/train/megatron_bridge/test_dataset_export.py

@@ -0,0 +1,112 @@
+# 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
+from pathlib import Path
+
+import pytest
+
+from llamafactory.train.megatron_bridge.dataset_export import (
+    _example_to_record,
+    _inject_generation_block,
+    export_dataset_for_megatron_bridge,
+    get_sft_dataset_kwargs,
+)
+
+
+def test_inject_generation_block_into_llama_factory_template():
+    template = (
+        "{% for message in loop_messages %}"
+        "{% if message['role'] == 'user' %}{{ message['content'] }}"
+        "{% elif message['role'] == 'assistant' %}{{ message['content'] }}{% endif %}"
+        "{% endfor %}"
+    )
+    patched = _inject_generation_block(template)
+    assert "{% generation %}" in patched
+    assert "{% endgeneration %}" in patched
+    assert _inject_generation_block(patched) == patched
+
+
+def test_example_to_record_messages_format():
+    example = {
+        "_system": "You are helpful.",
+        "_prompt": [{"role": "user", "content": "Hello"}],
+        "_response": [{"role": "assistant", "content": "Hi there!"}],
+    }
+    record = _example_to_record(example, stage="sft", use_messages_format=True)
+    assert record is not None
+    assert record["messages"][0] == {"role": "system", "content": "You are helpful."}
+    assert record["messages"][-1] == {"role": "assistant", "content": "Hi there!"}
+
+
+def test_example_to_record_pretrain_text():
+    example = {"_prompt": [{"role": "user", "content": "plain text sample"}]}
+    record = _example_to_record(example, stage="pt", use_messages_format=True)
+    assert record == {"text": "plain text sample"}
+
+
+def test_example_to_record_sharegpt_format():
+    example = {
+        "_system": "System prompt",
+        "_prompt": [{"role": "user", "content": "Question"}],
+        "_response": [{"role": "assistant", "content": "Answer"}],
+    }
+    record = _example_to_record(example, stage="sft", use_messages_format=False)
+    assert record == {
+        "system": "System prompt",
+        "conversations": [
+            {"from": "User", "value": "Question"},
+            {"from": "Assistant", "value": "Answer"},
+        ],
+        "mask": "User",
+    }
+
+
+def test_export_dataset_for_megatron_bridge(mb_output_dir: Path, mb_model_path: str):
+    dataset = [
+        {
+            "_prompt": [{"role": "user", "content": "Hello"}],
+            "_response": [{"role": "assistant", "content": "World"}],
+        }
+    ]
+    export_dir = mb_output_dir / "export"
+    export_dataset_for_megatron_bridge(
+        train_dataset=dataset,
+        output_dir=str(export_dir),
+        stage="sft",
+        model_name_or_path=mb_model_path,
+    )
+    train_path = export_dir / "training.jsonl"
+    assert train_path.is_file()
+    record = json.loads(train_path.read_text(encoding="utf-8").strip())
+    if "messages" in record:
+        assert record["messages"][-1]["content"] == "World"
+    else:
+        assert record["conversations"][-1]["value"] == "World"
+
+
+def test_get_sft_dataset_kwargs_enables_chat():
+    kwargs = get_sft_dataset_kwargs()
+    assert kwargs["chat"] is True
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_tokenizer_supports_hf_chat_template_on_gpu(mb_model_path: str):
+    from llamafactory.train.megatron_bridge.dataset_export import (
+        supports_hf_chat_template,
+        tokenizer_supports_hf_chat_template,
+    )
+
+    assert isinstance(supports_hf_chat_template(), bool)
+    assert isinstance(tokenizer_supports_hf_chat_template(mb_model_path), bool)

tests/train/megatron_bridge/test_megatron_bridge_gpu.py

@@ -0,0 +1,59 @@
+# 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 pytest
+
+from llamafactory.hparams.megatron_bridge_args import MegatronBridgeArguments
+from llamafactory.train.megatron_bridge.config_builder import _apply_fusion_safety, _apply_model_parallelism
+from llamafactory.train.megatron_bridge.workflow import _check_backend_available
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_check_backend_available():
+    _check_backend_available()
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_auto_bridge_provider_creation(mb_model_path: str):
+    from megatron.bridge import AutoBridge
+
+    bridge = AutoBridge.from_hf_pretrained(mb_model_path)
+    provider = bridge.to_megatron_provider(load_weights=False)
+    _apply_model_parallelism(provider, MegatronBridgeArguments(tensor_model_parallel_size=1))
+    _apply_fusion_safety(provider)
+    assert provider.tensor_model_parallel_size == 1
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_build_sft_config_tp2_on_gpu(mb_training_args_factory, mb_output_dir):
+    from llamafactory.train.megatron_bridge.config_builder import build_sft_config
+
+    model_args, data_args, training_args, finetuning_args, mb_args, num_train_samples = mb_training_args_factory(
+        tensor_model_parallel_size=2,
+        sequence_parallel=True,
+        use_distributed_optimizer=True,
+    )
+    cfg = build_sft_config(
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        mb_args=mb_args,
+        dataset_root=str(mb_output_dir / "dataset"),
+        pretrained_checkpoint=str(mb_output_dir / "pretrained"),
+        num_train_samples=num_train_samples,
+    )
+    assert cfg.model.tensor_model_parallel_size == 2
+    assert cfg.model.sequence_parallel is True
+    assert cfg.ddp.use_distributed_optimizer is True

tests/train/megatron_bridge/test_model_support.py

@@ -0,0 +1,91 @@
+# 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 types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from llamafactory.extras.constants import MEGATRON_BRIDGE_SUPPORTED_MODELS
+from llamafactory.train.megatron_bridge.workflow import _check_model_support
+
+
+# Keep in sync with MEGATRON_BRIDGE_SUPPORTED_MODELS (text LLM gpt_step path only).
+EXPECTED_MEGATRON_BRIDGE_SUPPORTED_MODELS = {
+    "deepseek_v3",
+    "deepseek_v4",
+    "llama",
+    "mistral",
+    "qwen2",
+    "qwen3",
+    "qwen3_5",
+    "qwen3_5_moe",
+    "qwen3_5_moe_text",
+    "qwen3_5_text",
+    "qwen3_moe",
+    "qwen3_next",
+}
+
+# Formerly allowlisted text models, plus multimodal types excluded in v0.
+UNSUPPORTED_MEGATRON_BRIDGE_MODELS = (
+    "bailing_moe_v2",
+    "deepseek_v2",
+    "ernie4_5_moe",
+    "falcon_h1",
+    "gemma",
+    "gemma2",
+    "gemma3",
+    "gemma4",
+    "glm4_moe",
+    "glm4_moe_lite",
+    "glm_moe_dsa",
+    "gpt_oss",
+    "kimi_k2",
+    "mimo",
+    "mimo_v2_flash",
+    "minimax_m2",
+    "mixtral",
+    "nemotron",
+    "nemotron_h",
+    "olmoe",
+    "qwen2_5_vl",
+    "qwen2_vl",
+    "qwen3_vl",
+    "qwen3_vl_moe",
+    "step3p5",
+)
+
+
+def test_megatron_bridge_supported_models_match_expected():
+    assert MEGATRON_BRIDGE_SUPPORTED_MODELS == EXPECTED_MEGATRON_BRIDGE_SUPPORTED_MODELS
+
+
+@pytest.mark.parametrize("model_type", sorted(EXPECTED_MEGATRON_BRIDGE_SUPPORTED_MODELS))
+def test_check_model_support_accepts_supported(model_type: str):
+    with patch(
+        "llamafactory.train.megatron_bridge.workflow.HfAutoConfig.from_pretrained",
+        return_value=SimpleNamespace(model_type=model_type),
+    ):
+        _check_model_support(SimpleNamespace(model_name_or_path="dummy", trust_remote_code=False))
+
+
+@pytest.mark.parametrize("model_type", UNSUPPORTED_MEGATRON_BRIDGE_MODELS)
+def test_check_model_support_rejects_unsupported(model_type: str):
+    assert model_type not in MEGATRON_BRIDGE_SUPPORTED_MODELS
+    with patch(
+        "llamafactory.train.megatron_bridge.workflow.HfAutoConfig.from_pretrained",
+        return_value=SimpleNamespace(model_type=model_type),
+    ):
+        with pytest.raises(ValueError, match="not supported by the Megatron Bridge"):
+            _check_model_support(SimpleNamespace(model_name_or_path="dummy", trust_remote_code=False))