diff --git a/docker/docker-cuda/Dockerfile.mbridge b/docker/docker-cuda/Dockerfile.mbridge new file mode 100644 index 000000000..c6757c769 --- /dev/null +++ b/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"] diff --git a/docker/docker-cuda/README.md b/docker/docker-cuda/README.md index 86707ce94..c1a399a48 100644 --- a/docker/docker-cuda/README.md +++ b/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 diff --git a/examples/megatron_bridge/llama3_sft.yaml b/examples/megatron_bridge/llama3_sft.yaml new file mode 100644 index 000000000..da4e16e7d --- /dev/null +++ b/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}' diff --git a/src/llamafactory/extras/constants.py b/src/llamafactory/extras/constants.py index 77721a01c..152a5681a 100644 --- a/src/llamafactory/extras/constants.py +++ b/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"} diff --git a/src/llamafactory/extras/packages.py b/src/llamafactory/extras/packages.py index a228ac0fb..1039bdec2 100644 --- a/src/llamafactory/extras/packages.py +++ b/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") diff --git a/src/llamafactory/hparams/__init__.py b/src/llamafactory/hparams/__init__.py index 9bcc4295c..1803c4c3f 100644 --- a/src/llamafactory/hparams/__init__.py +++ b/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", diff --git a/src/llamafactory/hparams/finetuning_args.py b/src/llamafactory/hparams/finetuning_args.py index 70a8e6e20..f4bd4debd 100644 --- a/src/llamafactory/hparams/finetuning_args.py +++ b/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={ diff --git a/src/llamafactory/hparams/megatron_bridge_args.py b/src/llamafactory/hparams/megatron_bridge_args.py new file mode 100644 index 000000000..b1137ba04 --- /dev/null +++ b/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) diff --git a/src/llamafactory/hparams/parser.py b/src/llamafactory/hparams/parser.py index 7c8e3b3fa..c22c1e375 100644 --- a/src/llamafactory/hparams/parser.py +++ b/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 diff --git a/src/llamafactory/launcher.py b/src/llamafactory/launcher.py index 2ccc06f4a..7366b9789 100644 --- a/src/llamafactory/launcher.py +++ b/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 ( diff --git a/src/llamafactory/train/megatron_bridge/__init__.py b/src/llamafactory/train/megatron_bridge/__init__.py new file mode 100644 index 000000000..88bbf20dd --- /dev/null +++ b/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"] diff --git a/src/llamafactory/train/megatron_bridge/config_builder.py b/src/llamafactory/train/megatron_bridge/config_builder.py new file mode 100644 index 000000000..a84fd27a9 --- /dev/null +++ b/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 diff --git a/src/llamafactory/train/megatron_bridge/dataset_export.py b/src/llamafactory/train/megatron_bridge/dataset_export.py new file mode 100644 index 000000000..0a4397291 --- /dev/null +++ b/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 diff --git a/src/llamafactory/train/megatron_bridge/workflow.py b/src/llamafactory/train/megatron_bridge/workflow.py new file mode 100644 index 000000000..eb49ec998 --- /dev/null +++ b/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) diff --git a/src/llamafactory/train/tuner.py b/src/llamafactory/train/tuner.py index b0dacf99e..97a6ef725 100644 --- a/src/llamafactory/train/tuner.py +++ b/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`.") diff --git a/tests/train/megatron_bridge/conftest.py b/tests/train/megatron_bridge/conftest.py new file mode 100644 index 000000000..ae9f3ad2f --- /dev/null +++ b/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() diff --git a/tests/train/megatron_bridge/test_config_builder.py b/tests/train/megatron_bridge/test_config_builder.py new file mode 100644 index 000000000..0b20c60b2 --- /dev/null +++ b/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 diff --git a/tests/train/megatron_bridge/test_dataset_export.py b/tests/train/megatron_bridge/test_dataset_export.py new file mode 100644 index 000000000..671174008 --- /dev/null +++ b/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) diff --git a/tests/train/megatron_bridge/test_megatron_bridge_gpu.py b/tests/train/megatron_bridge/test_megatron_bridge_gpu.py new file mode 100644 index 000000000..0275792a1 --- /dev/null +++ b/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 diff --git a/tests/train/megatron_bridge/test_model_support.py b/tests/train/megatron_bridge/test_model_support.py new file mode 100644 index 000000000..6d7b5bc4e --- /dev/null +++ b/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))