From 4a2621dd4f22077509802eef3be039e4096ce191 Mon Sep 17 00:00:00 2001 From: sunyi0505 <1659275352@qq.com> Date: Thu, 9 Jul 2026 10:59:57 +0800 Subject: [PATCH] [v0] support megatron-bridge for PT/SFT training --- examples/megatron_bridge/llama3_sft.yaml | 38 ++ src/llamafactory/extras/packages.py | 7 + src/llamafactory/hparams/__init__.py | 2 + src/llamafactory/hparams/finetuning_args.py | 15 + .../hparams/megatron_bridge_args.py | 118 ++++ src/llamafactory/hparams/parser.py | 93 ++- src/llamafactory/launcher.py | 2 +- .../train/megatron_bridge/__init__.py | 18 + .../train/megatron_bridge/config_builder.py | 593 ++++++++++++++++++ .../train/megatron_bridge/dataset_export.py | 224 +++++++ .../train/megatron_bridge/workflow.py | 229 +++++++ src/llamafactory/train/tuner.py | 19 + tests/train/megatron_bridge/conftest.py | 112 ++++ .../megatron_bridge/test_config_builder.py | 149 +++++ .../megatron_bridge/test_dataset_export.py | 98 +++ .../test_megatron_bridge_gpu.py | 59 ++ 16 files changed, 1772 insertions(+), 4 deletions(-) create mode 100644 examples/megatron_bridge/llama3_sft.yaml create mode 100644 src/llamafactory/hparams/megatron_bridge_args.py create mode 100644 src/llamafactory/train/megatron_bridge/__init__.py create mode 100644 src/llamafactory/train/megatron_bridge/config_builder.py create mode 100644 src/llamafactory/train/megatron_bridge/dataset_export.py create mode 100644 src/llamafactory/train/megatron_bridge/workflow.py create mode 100644 tests/train/megatron_bridge/conftest.py create mode 100644 tests/train/megatron_bridge/test_config_builder.py create mode 100644 tests/train/megatron_bridge/test_dataset_export.py create mode 100644 tests/train/megatron_bridge/test_megatron_bridge_gpu.py diff --git a/examples/megatron_bridge/llama3_sft.yaml b/examples/megatron_bridge/llama3_sft.yaml new file mode 100644 index 0000000000..123c026b73 --- /dev/null +++ b/examples/megatron_bridge/llama3_sft.yaml @@ -0,0 +1,38 @@ +### model +model_name_or_path: meta-llama/Llama-3.2-1B-Instruct + +### method +stage: sft +do_train: true +finetuning_type: full # full or lora +dataset: alpaca_en_demo +template: llama3 +cutoff_len: 2048 +preprocessing_num_workers: 8 + +### output +output_dir: saves/mbridge/llama3_sft +logging_steps: 1 +overwrite_output_dir: true + +### train +per_device_train_batch_size: 1 +gradient_accumulation_steps: 1 +num_train_epochs: 3 +max_steps: 5 # max_steps is used to limit the number of steps to train, if set, num_train_epochs will be ignored +save_steps: 3000 +learning_rate: 5.0e-6 +lr_scheduler_type: cosine +warmup_steps: 10 +bf16: true + +### megatron bridge +tensor_model_parallel_size: 2 +pipeline_model_parallel_size: 1 +context_parallel_size: 1 +sequence_parallel: true +use_distributed_optimizer: true +overlap_param_gather: true +overlap_grad_reduce: true +mixed_precision: bf16_mixed +export_hf_on_finish: true diff --git a/src/llamafactory/extras/packages.py b/src/llamafactory/extras/packages.py index a228ac0fb7..1039bdec25 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 9bcc4295ce..1803c4c3fc 100644 --- a/src/llamafactory/hparams/__init__.py +++ b/src/llamafactory/hparams/__init__.py @@ -16,6 +16,7 @@ from .evaluation_args import EvaluationArguments from .finetuning_args import FinetuningArguments from .generating_args import GeneratingArguments +from .megatron_bridge_args import MegatronBridgeArguments from .model_args import ModelArguments from .parser import get_eval_args, get_infer_args, get_ray_args, get_train_args, read_args from .training_args import RayArguments, TrainingArguments @@ -26,6 +27,7 @@ "EvaluationArguments", "FinetuningArguments", "GeneratingArguments", + "MegatronBridgeArguments", "ModelArguments", "RayArguments", "TrainingArguments", diff --git a/src/llamafactory/hparams/finetuning_args.py b/src/llamafactory/hparams/finetuning_args.py index 70a8e6e205..f4bd4debdc 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 0000000000..803f470a23 --- /dev/null +++ b/src/llamafactory/hparams/megatron_bridge_args.py @@ -0,0 +1,118 @@ +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from dataclasses import dataclass, field +from typing import Optional + +from transformers.training_args import _convert_str_dict + + +@dataclass +class MegatronBridgeArguments: + r"""Arguments for Megatron Bridge distributed training backend.""" + + tensor_model_parallel_size: int = field( + default=1, + metadata={"help": "Tensor model parallel size for Megatron Bridge."}, + ) + pipeline_model_parallel_size: int = field( + default=1, + metadata={"help": "Pipeline model parallel size for Megatron Bridge."}, + ) + expert_model_parallel_size: int = field( + default=1, + metadata={"help": "Expert model parallel size for MoE models."}, + ) + context_parallel_size: int = field( + default=1, + metadata={"help": "Context parallel size for Megatron Bridge."}, + ) + sequence_parallel: bool = field( + default=False, + metadata={"help": "Whether to enable sequence parallelism."}, + ) + recompute_granularity: Optional[str] = field( + default=None, + metadata={"help": "Activation recomputation granularity: 'full' or 'selective'."}, + ) + use_distributed_optimizer: bool = field( + default=True, + metadata={"help": "Whether to use Megatron distributed optimizer."}, + ) + overlap_param_gather: bool = field( + default=True, + metadata={"help": "Whether to overlap parameter all-gather with forward compute."}, + ) + overlap_grad_reduce: bool = field( + default=True, + metadata={"help": "Whether to overlap gradient all-reduce with backward compute."}, + ) + use_packed_sequences: bool = field( + default=False, + metadata={"help": "Whether to use packed sequences for SFT efficiency."}, + ) + mixed_precision: str = field( + default="bf16_mixed", + metadata={"help": "Mixed precision mode for Megatron Bridge, e.g. bf16_mixed or fp8."}, + ) + megatron_pretrained_checkpoint: Optional[str] = field( + default=None, + metadata={ + "help": ( + "Path to a Megatron-format pretrained checkpoint. " + "If unset, HF weights are converted automatically before training." + ) + }, + ) + export_hf_on_finish: bool = field( + default=False, + metadata={"help": "Whether to export the final checkpoint to Hugging Face format after training."}, + ) + extra_config: Optional[str] = field( + default=None, + metadata={ + "help": ( + "Optional JSON string or path to a JSON file with extra Megatron Bridge model/training overrides." + ) + }, + ) + + def __post_init__(self) -> None: + if self.tensor_model_parallel_size < 1: + raise ValueError("`tensor_model_parallel_size` must be >= 1.") + if self.pipeline_model_parallel_size < 1: + raise ValueError("`pipeline_model_parallel_size` must be >= 1.") + if self.expert_model_parallel_size < 1: + raise ValueError("`expert_model_parallel_size` must be >= 1.") + if self.context_parallel_size < 1: + raise ValueError("`context_parallel_size` must be >= 1.") + if self.sequence_parallel and self.tensor_model_parallel_size <= 1: + raise ValueError("`sequence_parallel` requires `tensor_model_parallel_size` > 1.") + if self.recompute_granularity is not None and self.recompute_granularity not in ("full", "selective"): + raise ValueError("`recompute_granularity` must be 'full' or 'selective'.") + + if isinstance(self.extra_config, str) and self.extra_config.startswith("{"): + self.extra_config = _convert_str_dict(json.loads(self.extra_config)) + + def load_extra_config(self) -> dict: + if self.extra_config is None: + return {} + if isinstance(self.extra_config, dict): + return self.extra_config + if not os.path.isfile(self.extra_config): + raise ValueError(f"`extra_config` file not found: {self.extra_config}") + with open(self.extra_config, encoding="utf-8") as f: + return json.load(f) diff --git a/src/llamafactory/hparams/parser.py b/src/llamafactory/hparams/parser.py index 7c8e3b3fa7..53cc955994 100644 --- a/src/llamafactory/hparams/parser.py +++ b/src/llamafactory/hparams/parser.py @@ -33,11 +33,12 @@ from ..extras import logging from ..extras.constants import CHECKPOINT_NAMES, EngineName from ..extras.misc import check_dependencies, check_version, get_current_device, is_env_enabled -from ..extras.packages import is_mcore_adapter_available +from ..extras.packages import is_mcore_adapter_available, is_megatron_bridge_available from .data_args import DataArguments from .evaluation_args import EvaluationArguments from .finetuning_args import FinetuningArguments from .generating_args import GeneratingArguments +from .megatron_bridge_args import MegatronBridgeArguments from .model_args import ModelArguments from .training_args import RayArguments, TrainingArguments @@ -81,6 +82,23 @@ _TRAIN_MCA_ARGS = [] _TRAIN_MCA_CLS = tuple() +_TRAIN_MBRIDGE_ARGS = [ + ModelArguments, + DataArguments, + TrainingArguments, + FinetuningArguments, + MegatronBridgeArguments, + GeneratingArguments, +] +_TRAIN_MBRIDGE_CLS = tuple[ + ModelArguments, + DataArguments, + TrainingArguments, + FinetuningArguments, + MegatronBridgeArguments, + GeneratingArguments, +] + def read_args(args: dict[str, Any] | list[str] | None = None) -> dict[str, Any] | list[str]: r"""Get arguments from the command line or a config file.""" @@ -246,6 +264,9 @@ def _check_extra_dependencies( if finetuning_args.plot_loss: check_version("matplotlib", mandatory=True) + if finetuning_args.use_megatron_bridge: + check_version("megatron-bridge", mandatory=True) + if training_args is not None: if training_args.deepspeed: check_version("deepspeed", mandatory=True) @@ -283,6 +304,35 @@ def _configure_mca_training_args(training_args, data_args, finetuning_args) -> N finetuning_args.use_mca = True +def _validate_megatron_bridge_parallel_args(mb_args: MegatronBridgeArguments, world_size: int) -> None: + parallel_size = ( + mb_args.tensor_model_parallel_size + * mb_args.pipeline_model_parallel_size + * mb_args.context_parallel_size + * mb_args.expert_model_parallel_size + ) + if parallel_size > world_size: + raise ValueError(f"Total Megatron Bridge parallel size ({parallel_size}) exceeds `world_size` ({world_size}).") + + +def _parse_train_mbridge_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_MBRIDGE_CLS: + parser = HfArgumentParser(_TRAIN_MBRIDGE_ARGS) + allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS") + model_args, data_args, training_args, finetuning_args, mb_args, generating_args = _parse_args( + parser, args, allow_extra_keys=allow_extra_keys + ) + _configure_mbridge_training_args(training_args, data_args, finetuning_args) + return model_args, data_args, training_args, finetuning_args, mb_args, generating_args + + +def _configure_mbridge_training_args(training_args, data_args, finetuning_args) -> None: + """Patch training args to avoid args checking errors and sync Megatron Bridge settings.""" + training_args.predict_with_generate = False + training_args.generation_max_length = data_args.cutoff_len + training_args.generation_num_beams = 1 + finetuning_args.use_megatron_bridge = True + + def _parse_infer_args(args: dict[str, Any] | list[str] | None = None) -> _INFER_CLS: parser = HfArgumentParser(_INFER_ARGS) allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS") @@ -302,11 +352,22 @@ def get_ray_args(args: dict[str, Any] | list[str] | None = None) -> RayArguments def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS: + mb_args = None if is_env_enabled("USE_MCA"): model_args, data_args, training_args, finetuning_args, generating_args = _parse_train_mca_args(args) + elif is_env_enabled("USE_MEGATRON_BRIDGE"): + if not is_megatron_bridge_available(): + raise ImportError( + "megatron-bridge is required when USE_MEGATRON_BRIDGE=1. " + "Please install `megatron-bridge` and its dependencies." + ) + model_args, data_args, training_args, finetuning_args, mb_args, generating_args = _parse_train_mbridge_args( + args + ) else: model_args, data_args, training_args, finetuning_args, generating_args = _parse_train_args(args) finetuning_args.use_mca = False + finetuning_args.use_megatron_bridge = False # Setup logging if training_args.should_log: @@ -326,6 +387,22 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS if finetuning_args.stage == "sft" and training_args.do_predict and not training_args.predict_with_generate: raise ValueError("Please enable `predict_with_generate` to save model predictions.") + if finetuning_args.use_megatron_bridge: + if finetuning_args.use_mca or finetuning_args.use_hyper_parallel: + raise ValueError("Megatron Bridge cannot be used together with MCA or HyperParallel.") + if finetuning_args.stage not in ["pt", "sft"]: + raise ValueError("Megatron Bridge only supports the `pt` and `sft` stages.") + if finetuning_args.finetuning_type not in ["full", "lora"]: + raise ValueError("Megatron Bridge only supports `full` and `lora` finetuning.") + if model_args.quantization_bit is not None: + raise ValueError("Quantized models are not supported with Megatron Bridge.") + if training_args.deepspeed is not None: + raise ValueError("Megatron Bridge is incompatible with DeepSpeed.") + if mb_args is None: + raise ValueError("Megatron Bridge arguments are missing. Please set USE_MEGATRON_BRIDGE=1.") + _validate_megatron_bridge_parallel_args(mb_args, training_args.world_size) + finetuning_args.megatron_bridge_args = mb_args + if finetuning_args.stage in ["rm", "ppo"] and training_args.load_best_model_at_end: raise ValueError("RM and PPO stages do not support `load_best_model_at_end`.") @@ -400,7 +477,12 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS if training_args.deepspeed is not None and (finetuning_args.use_galore or finetuning_args.use_apollo): raise ValueError("GaLore and APOLLO are incompatible with DeepSpeed yet.") - if not finetuning_args.use_mca and training_args.fp8 and model_args.quantization_bit is not None: + if ( + not finetuning_args.use_mca + and not finetuning_args.use_megatron_bridge + and training_args.fp8 + and model_args.quantization_bit is not None + ): raise ValueError("FP8 training is not compatible with quantization. Please disable one of them.") if model_args.infer_backend != EngineName.HF: @@ -417,7 +499,12 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS _check_extra_dependencies(model_args, finetuning_args, training_args) _verify_trackio_args(training_args) - if not finetuning_args.use_mca and training_args.fp8_enable_fsdp_float8_all_gather and not training_args.fp8: + if ( + not finetuning_args.use_mca + and not finetuning_args.use_megatron_bridge + and training_args.fp8_enable_fsdp_float8_all_gather + and not training_args.fp8 + ): logger.warning_rank0("fp8_enable_fsdp_float8_all_gather requires fp8=True. Setting fp8=True.") model_args.fp8 = True diff --git a/src/llamafactory/launcher.py b/src/llamafactory/launcher.py index 2ccc06f4a9..7366b97890 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 0000000000..88bbf20ddc --- /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 0000000000..a0285fd621 --- /dev/null +++ b/src/llamafactory/train/megatron_bridge/config_builder.py @@ -0,0 +1,593 @@ +# 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: + 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 int(train_iters * warmup_ratio) + return 0 + + +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, beta2 = learning_rate, 0.0, 0.98 + else: + max_lr, min_lr, beta2 = learning_rate, learning_rate * 0.1, 0.95 + + weight_decay = getattr(training_args, "weight_decay", 0.1) + max_grad_norm = getattr(training_args, "max_grad_norm", 1.0) + 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=0.9, + adam_beta2=beta2, + adam_eps=1e-8, + 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=train_iters, + lr_warmup_iters=warmup_steps, + lr_warmup_init=0.0, + lr_decay_iters=train_iters, + override_opt_param_scheduler=True, + ) + return optimizer, scheduler + + +def _create_peft_config(finetuning_args: "FinetuningArguments"): + if finetuning_args.finetuning_type != "lora": + return None + + from megatron.bridge.peft.lora import LoRA + + return LoRA( + target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], + 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, +): + 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() + return FinetuningDatasetConfig(**kwargs) + + +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 + 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, + ckpt_format="torch_dist", + fully_parallel_save=False, + use_persistent_ckpt_worker=False, + save_optim=True, + ), + "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_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 + if mb_args.recompute_granularity is not None and hasattr(model_provider, "recompute_granularity"): + model_provider.recompute_granularity = mb_args.recompute_granularity + + +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(): + if "." in key: + target, attr = key.split(".", 1) + setattr(getattr(cfg, target), attr, value) + else: + setattr(cfg, key, 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 + + dataset_kwargs = get_sft_dataset_kwargs( + tokenizer_path=model_args.model_name_or_path, + trust_remote_code=model_args.trust_remote_code, + ) + 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, + ) + 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 0000000000..c3f761dd59 --- /dev/null +++ b/src/llamafactory/train/megatron_bridge/dataset_export.py @@ -0,0 +1,224 @@ +# 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+-?%\}") + + +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 tokenizer_supports_hf_chat_template( + tokenizer_path: str | None, + trust_remote_code: bool = False, +) -> bool: + r"""Return whether the tokenizer chat template supports assistant-only loss masks.""" + 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) + return _GENERATION_REGEX.search(tokenizer.chat_template or "") is not None + except Exception as exc: + logger.warning_rank0(f"Failed to inspect tokenizer chat template: {exc}") + return False + + +def get_sft_dataset_kwargs( + tokenizer_path: str | None = None, + trust_remote_code: bool = False, +) -> dict[str, Any]: + r"""Return dataset kwargs compatible with the installed Megatron Bridge version.""" + kwargs: dict[str, Any] = {"chat": True} + if tokenizer_supports_hf_chat_template(tokenizer_path, trust_remote_code): + kwargs["use_hf_tokenizer_chat_template"] = True + return kwargs + + +def get_finetuning_dataloader_type() -> str: + r"""Return a finetuning dataloader type supported by the installed Megatron Bridge.""" + try: + from megatron.bridge.training.config import FinetuningDatasetConfig + + field = FinetuningDatasetConfig.__dataclass_fields__.get("dataloader_type") + if field is None: + return "single" + + for arg in typing.get_args(field.type): + for choice in typing.get_args(arg): + if choice == "batch": + return "batch" + return "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 + 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, +) -> 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, + ) + if stage == "sft" and not use_messages_format: + logger.info_rank0( + "Tokenizer chat template lacks a {% generation %} block; " + "exporting ShareGPT conversations for legacy Megatron Bridge 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 0000000000..e04ed4f26e --- /dev/null +++ b/src/llamafactory/train/megatron_bridge/workflow.py @@ -0,0 +1,229 @@ +# 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 +from typing import TYPE_CHECKING, Optional + +from ...data.data_utils import split_dataset +from ...data.loader import _get_merged_dataset +from ...extras.logging import get_logger +from ...extras.packages import is_megatron_bridge_available +from .config_builder import 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_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 _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 + + export_dir = os.path.join(output_dir, "hf_export") + logger.info_rank0(f"Exporting Megatron checkpoint to Hugging Face format at {export_dir}...") + bridge = AutoBridge.from_hf_pretrained( + model_args.model_name_or_path, + trust_remote_code=model_args.trust_remote_code, + ) + 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() + 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") + 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", + ) + + 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() + 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") + 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, + ) + + 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 d425f9050d..20c79eca1f 100644 --- a/src/llamafactory/train/tuner.py +++ b/src/llamafactory/train/tuner.py @@ -27,6 +27,7 @@ from ..extras.packages import ( is_hyper_parallel_available, is_mcore_adapter_available, + is_megatron_bridge_available, is_ray_available, is_transformers_version_greater_than, ) @@ -102,6 +103,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 0000000000..ae9f3ad2f8 --- /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 0000000000..354d108c12 --- /dev/null +++ b/tests/train/megatron_bridge/test_config_builder.py @@ -0,0 +1,149 @@ +# 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_fusion_safety, + _apply_model_parallelism, + _compute_train_schedule, + 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, + ) + _apply_model_parallelism(provider, mb_args) + assert provider.tensor_model_parallel_size == 2 + assert provider.sequence_parallel is True + + +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) + + +@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 0000000000..a1132d14dd --- /dev/null +++ b/tests/train/megatron_bridge/test_dataset_export.py @@ -0,0 +1,98 @@ +# 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, + export_dataset_for_megatron_bridge, + get_sft_dataset_kwargs, +) + + +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 0000000000..0275792a11 --- /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