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

tests/train/megatron_bridge/conftest.py

@@ -0,0 +1,112 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Fixtures for Megatron Bridge GPU unit tests."""
+
+import os
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+from llamafactory.extras.packages import is_megatron_bridge_available
+
+
+pytestmark = pytest.mark.skipif(
+    not is_megatron_bridge_available(),
+    reason="megatron-bridge is not installed",
+)
+
+
+@pytest.fixture(scope="session")
+def mb_model_path() -> str:
+    r"""Return the local Hugging Face model path used by Megatron Bridge tests."""
+    model_path = os.getenv("TINY_LLAMA3", "llamafactory/tiny-random-Llama-3")
+    return model_path
+
+
+@pytest.fixture
+def mb_output_dir(tmp_path: Path) -> Path:
+    r"""Create a writable output directory for Megatron Bridge tests."""
+    output_dir = tmp_path / "output"
+    output_dir.mkdir(parents=True)
+    return output_dir
+
+
+@pytest.fixture
+def mb_training_args_factory(mb_output_dir: Path, mb_model_path: str):
+    r"""Build LLaMA-Factory dataclass arguments for Megatron Bridge tests."""
+    patchers: list[mock._patch] = []
+
+    def _factory(
+        *,
+        tensor_model_parallel_size: int = 1,
+        pipeline_model_parallel_size: int = 1,
+        context_parallel_size: int = 1,
+        sequence_parallel: bool = False,
+        use_distributed_optimizer: bool = False,
+        num_train_samples: int = 100,
+        cutoff_len: int = 512,
+        output_dir: Path | None = None,
+    ):
+        from llamafactory.hparams.data_args import DataArguments
+        from llamafactory.hparams.finetuning_args import FinetuningArguments
+        from llamafactory.hparams.megatron_bridge_args import MegatronBridgeArguments
+        from llamafactory.hparams.model_args import ModelArguments
+        from llamafactory.hparams.training_args import TrainingArguments
+
+        resolved_output_dir = output_dir or mb_output_dir
+        model_args = ModelArguments(model_name_or_path=mb_model_path)
+        data_args = DataArguments(
+            dataset="alpaca_en_demo",
+            template="llama3",
+            cutoff_len=cutoff_len,
+            preprocessing_num_workers=1,
+        )
+        training_args = TrainingArguments(
+            output_dir=str(resolved_output_dir),
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=1,
+            num_train_epochs=1,
+            learning_rate=5e-6,
+            logging_steps=1,
+            save_steps=1000000,
+            report_to="none",
+        )
+        world_size = tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size
+        patcher = mock.patch.object(
+            type(training_args),
+            "world_size",
+            new_callable=mock.PropertyMock,
+            return_value=world_size,
+        )
+        patchers.append(patcher)
+        patcher.start()
+        finetuning_args = FinetuningArguments(stage="sft", finetuning_type="full")
+        mb_args = MegatronBridgeArguments(
+            tensor_model_parallel_size=tensor_model_parallel_size,
+            pipeline_model_parallel_size=pipeline_model_parallel_size,
+            context_parallel_size=context_parallel_size,
+            sequence_parallel=sequence_parallel,
+            use_distributed_optimizer=use_distributed_optimizer,
+            overlap_param_gather=False,
+            overlap_grad_reduce=False,
+            export_hf_on_finish=False,
+        )
+        return model_args, data_args, training_args, finetuning_args, mb_args, num_train_samples
+
+    yield _factory
+
+    for patcher in patchers:
+        patcher.stop()

tests/train/megatron_bridge/test_config_builder.py

@@ -0,0 +1,206 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import pytest
+
+from llamafactory.hparams.megatron_bridge_args import MegatronBridgeArguments
+from llamafactory.train.megatron_bridge.config_builder import (
+    _apply_extra_overrides,
+    _apply_fusion_safety,
+    _apply_model_parallelism,
+    _compute_train_schedule,
+    _resolve_warmup_steps,
+    _should_resume_checkpoint,
+    build_pretrain_config,
+    build_sft_config,
+)
+
+
+class _DummyProvider:
+    tensor_model_parallel_size = 1
+    pipeline_model_parallel_size = 1
+    context_parallel_size = 1
+    sequence_parallel = False
+    gradient_accumulation_fusion = True
+
+
+def test_compute_train_schedule_tp1():
+    training_args = SimpleNamespace(
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=4,
+        num_train_epochs=1,
+        max_steps=-1,
+        world_size=2,
+    )
+    mb_args = MegatronBridgeArguments(tensor_model_parallel_size=1)
+    micro_batch, global_batch, train_iters = _compute_train_schedule(training_args, mb_args, num_train_samples=32)
+    assert micro_batch == 2
+    assert global_batch == 16
+    assert train_iters == 2
+
+
+def test_compute_train_schedule_tp2():
+    training_args = SimpleNamespace(
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=1,
+        num_train_epochs=1,
+        max_steps=-1,
+        world_size=2,
+    )
+    mb_args = MegatronBridgeArguments(
+        tensor_model_parallel_size=2,
+        pipeline_model_parallel_size=1,
+        context_parallel_size=1,
+    )
+    micro_batch, global_batch, train_iters = _compute_train_schedule(training_args, mb_args, num_train_samples=10)
+    assert micro_batch == 1
+    assert global_batch == 1
+    assert train_iters == 10
+
+
+def test_compute_train_schedule_max_steps():
+    training_args = SimpleNamespace(
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=1,
+        num_train_epochs=10,
+        max_steps=5,
+        world_size=1,
+    )
+    mb_args = MegatronBridgeArguments()
+    _, _, train_iters = _compute_train_schedule(training_args, mb_args, num_train_samples=1000)
+    assert train_iters == 5
+
+
+def test_apply_model_parallelism():
+    provider = _DummyProvider()
+    mb_args = MegatronBridgeArguments(
+        tensor_model_parallel_size=2,
+        pipeline_model_parallel_size=1,
+        context_parallel_size=1,
+        sequence_parallel=True,
+        bias_activation_fusion=True,
+        moe_token_dispatcher_type="alltoall",
+        recompute_granularity="full",
+        recompute_method="uniform",
+        recompute_num_layers=2,
+    )
+    provider.bias_activation_fusion = False
+    provider.moe_token_dispatcher_type = "allgather"
+    provider.recompute_granularity = None
+    provider.recompute_method = None
+    provider.recompute_num_layers = None
+    _apply_model_parallelism(provider, mb_args)
+    assert provider.tensor_model_parallel_size == 2
+    assert provider.sequence_parallel is True
+    assert provider.bias_activation_fusion is True
+    assert provider.moe_token_dispatcher_type == "alltoall"
+    assert provider.recompute_granularity == "full"
+    assert provider.recompute_method == "uniform"
+    assert provider.recompute_num_layers == 2
+
+
+def test_megatron_bridge_args_rejects_invalid_moe_dispatcher():
+    with pytest.raises(ValueError, match="moe_token_dispatcher_type"):
+        MegatronBridgeArguments(moe_token_dispatcher_type="invalid")  # type: ignore[arg-type]
+
+
+def test_apply_fusion_safety_disables_missing_apex():
+    provider = _DummyProvider()
+    _apply_fusion_safety(provider)
+    assert provider.gradient_accumulation_fusion is False
+
+
+def test_megatron_bridge_args_sequence_parallel_requires_tp():
+    with pytest.raises(ValueError, match="sequence_parallel"):
+        MegatronBridgeArguments(sequence_parallel=True, tensor_model_parallel_size=1)
+
+
+def test_resolve_warmup_steps_keeps_absolute_warmup():
+    training_args = SimpleNamespace(warmup_steps=10, warmup_ratio=0.0)
+    assert _resolve_warmup_steps(training_args, train_iters=5) == 10
+
+
+def test_resolve_decay_iters_expands_for_warmup():
+    from llamafactory.train.megatron_bridge.config_builder import _resolve_decay_iters
+
+    assert _resolve_decay_iters(train_iters=5, warmup_steps=10) == 11
+    assert _resolve_decay_iters(train_iters=20, warmup_steps=10) == 20
+    assert _resolve_decay_iters(train_iters=5, warmup_steps=0) == 5
+
+
+def test_megatron_bridge_args_extra_config_strips_whitespace():
+    args = MegatronBridgeArguments(extra_config='  {"train.train_iters": 3}  ')
+    assert args.extra_config == {"train.train_iters": 3}
+
+
+def test_apply_extra_overrides_nested_path():
+    cfg = SimpleNamespace(train=SimpleNamespace(nested=SimpleNamespace(value=1)))
+    _apply_extra_overrides(cfg, {"train.nested.value": 9})
+    assert cfg.train.nested.value == 9
+
+
+def test_should_resume_checkpoint_respects_overwrite(tmp_path):
+    tracker = tmp_path / "latest_checkpointed_iteration.txt"
+    tracker.write_text("5")
+    assert _should_resume_checkpoint(SimpleNamespace(output_dir=str(tmp_path), overwrite_output_dir=False))
+    assert not _should_resume_checkpoint(SimpleNamespace(output_dir=str(tmp_path), overwrite_output_dir=True))
+    assert not _should_resume_checkpoint(
+        SimpleNamespace(output_dir=str(tmp_path / "missing"), overwrite_output_dir=False)
+    )
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_build_sft_config_on_gpu(mb_training_args_factory, mb_output_dir):
+    model_args, data_args, training_args, finetuning_args, mb_args, num_train_samples = mb_training_args_factory(
+        tensor_model_parallel_size=1,
+        sequence_parallel=False,
+    )
+    cfg = build_sft_config(
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        mb_args=mb_args,
+        dataset_root=str(mb_output_dir / "dataset"),
+        pretrained_checkpoint=str(mb_output_dir / "pretrained"),
+        num_train_samples=num_train_samples,
+    )
+    assert cfg.model.tensor_model_parallel_size == 1
+    assert cfg.model.seq_length == data_args.cutoff_len
+    assert cfg.train.train_iters == num_train_samples
+    assert cfg.dataset.dataset_root == str(mb_output_dir / "dataset")
+    assert cfg.tokenizer.tokenizer_model == model_args.model_name_or_path
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_build_pretrain_config_on_gpu(mb_training_args_factory, mb_output_dir):
+    model_args, data_args, training_args, finetuning_args, mb_args, num_train_samples = mb_training_args_factory(
+        tensor_model_parallel_size=1,
+    )
+    finetuning_args.stage = "pt"
+    dataset_path = str(mb_output_dir / "dataset" / "training.jsonl")
+    cfg = build_pretrain_config(
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        mb_args=mb_args,
+        dataset_path=dataset_path,
+        num_train_samples=num_train_samples,
+    )
+    assert cfg.model.tensor_model_parallel_size == 1
+    assert cfg.dataset.blend[0][0] == dataset_path
+    assert cfg.train.train_iters == num_train_samples

tests/train/megatron_bridge/test_dataset_export.py

@@ -0,0 +1,112 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+from pathlib import Path
+
+import pytest
+
+from llamafactory.train.megatron_bridge.dataset_export import (
+    _example_to_record,
+    _inject_generation_block,
+    export_dataset_for_megatron_bridge,
+    get_sft_dataset_kwargs,
+)
+
+
+def test_inject_generation_block_into_llama_factory_template():
+    template = (
+        "{% for message in loop_messages %}"
+        "{% if message['role'] == 'user' %}{{ message['content'] }}"
+        "{% elif message['role'] == 'assistant' %}{{ message['content'] }}{% endif %}"
+        "{% endfor %}"
+    )
+    patched = _inject_generation_block(template)
+    assert "{% generation %}" in patched
+    assert "{% endgeneration %}" in patched
+    assert _inject_generation_block(patched) == patched
+
+
+def test_example_to_record_messages_format():
+    example = {
+        "_system": "You are helpful.",
+        "_prompt": [{"role": "user", "content": "Hello"}],
+        "_response": [{"role": "assistant", "content": "Hi there!"}],
+    }
+    record = _example_to_record(example, stage="sft", use_messages_format=True)
+    assert record is not None
+    assert record["messages"][0] == {"role": "system", "content": "You are helpful."}
+    assert record["messages"][-1] == {"role": "assistant", "content": "Hi there!"}
+
+
+def test_example_to_record_pretrain_text():
+    example = {"_prompt": [{"role": "user", "content": "plain text sample"}]}
+    record = _example_to_record(example, stage="pt", use_messages_format=True)
+    assert record == {"text": "plain text sample"}
+
+
+def test_example_to_record_sharegpt_format():
+    example = {
+        "_system": "System prompt",
+        "_prompt": [{"role": "user", "content": "Question"}],
+        "_response": [{"role": "assistant", "content": "Answer"}],
+    }
+    record = _example_to_record(example, stage="sft", use_messages_format=False)
+    assert record == {
+        "system": "System prompt",
+        "conversations": [
+            {"from": "User", "value": "Question"},
+            {"from": "Assistant", "value": "Answer"},
+        ],
+        "mask": "User",
+    }
+
+
+def test_export_dataset_for_megatron_bridge(mb_output_dir: Path, mb_model_path: str):
+    dataset = [
+        {
+            "_prompt": [{"role": "user", "content": "Hello"}],
+            "_response": [{"role": "assistant", "content": "World"}],
+        }
+    ]
+    export_dir = mb_output_dir / "export"
+    export_dataset_for_megatron_bridge(
+        train_dataset=dataset,
+        output_dir=str(export_dir),
+        stage="sft",
+        model_name_or_path=mb_model_path,
+    )
+    train_path = export_dir / "training.jsonl"
+    assert train_path.is_file()
+    record = json.loads(train_path.read_text(encoding="utf-8").strip())
+    if "messages" in record:
+        assert record["messages"][-1]["content"] == "World"
+    else:
+        assert record["conversations"][-1]["value"] == "World"
+
+
+def test_get_sft_dataset_kwargs_enables_chat():
+    kwargs = get_sft_dataset_kwargs()
+    assert kwargs["chat"] is True
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_tokenizer_supports_hf_chat_template_on_gpu(mb_model_path: str):
+    from llamafactory.train.megatron_bridge.dataset_export import (
+        supports_hf_chat_template,
+        tokenizer_supports_hf_chat_template,
+    )
+
+    assert isinstance(supports_hf_chat_template(), bool)
+    assert isinstance(tokenizer_supports_hf_chat_template(mb_model_path), bool)

tests/train/megatron_bridge/test_megatron_bridge_gpu.py

@@ -0,0 +1,59 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from llamafactory.hparams.megatron_bridge_args import MegatronBridgeArguments
+from llamafactory.train.megatron_bridge.config_builder import _apply_fusion_safety, _apply_model_parallelism
+from llamafactory.train.megatron_bridge.workflow import _check_backend_available
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_check_backend_available():
+    _check_backend_available()
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_auto_bridge_provider_creation(mb_model_path: str):
+    from megatron.bridge import AutoBridge
+
+    bridge = AutoBridge.from_hf_pretrained(mb_model_path)
+    provider = bridge.to_megatron_provider(load_weights=False)
+    _apply_model_parallelism(provider, MegatronBridgeArguments(tensor_model_parallel_size=1))
+    _apply_fusion_safety(provider)
+    assert provider.tensor_model_parallel_size == 1
+
+
+@pytest.mark.runs_on(["cuda"])
+def test_build_sft_config_tp2_on_gpu(mb_training_args_factory, mb_output_dir):
+    from llamafactory.train.megatron_bridge.config_builder import build_sft_config
+
+    model_args, data_args, training_args, finetuning_args, mb_args, num_train_samples = mb_training_args_factory(
+        tensor_model_parallel_size=2,
+        sequence_parallel=True,
+        use_distributed_optimizer=True,
+    )
+    cfg = build_sft_config(
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        mb_args=mb_args,
+        dataset_root=str(mb_output_dir / "dataset"),
+        pretrained_checkpoint=str(mb_output_dir / "pretrained"),
+        num_train_samples=num_train_samples,
+    )
+    assert cfg.model.tensor_model_parallel_size == 2
+    assert cfg.model.sequence_parallel is True
+    assert cfg.ddp.use_distributed_optimizer is True

tests/train/megatron_bridge/test_model_support.py

@@ -0,0 +1,91 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from llamafactory.extras.constants import MEGATRON_BRIDGE_SUPPORTED_MODELS
+from llamafactory.train.megatron_bridge.workflow import _check_model_support
+
+
+# Keep in sync with MEGATRON_BRIDGE_SUPPORTED_MODELS (text LLM gpt_step path only).
+EXPECTED_MEGATRON_BRIDGE_SUPPORTED_MODELS = {
+    "deepseek_v3",
+    "deepseek_v4",
+    "llama",
+    "mistral",
+    "qwen2",
+    "qwen3",
+    "qwen3_5",
+    "qwen3_5_moe",
+    "qwen3_5_moe_text",
+    "qwen3_5_text",
+    "qwen3_moe",
+    "qwen3_next",
+}
+
+# Formerly allowlisted text models, plus multimodal types excluded in v0.
+UNSUPPORTED_MEGATRON_BRIDGE_MODELS = (
+    "bailing_moe_v2",
+    "deepseek_v2",
+    "ernie4_5_moe",
+    "falcon_h1",
+    "gemma",
+    "gemma2",
+    "gemma3",
+    "gemma4",
+    "glm4_moe",
+    "glm4_moe_lite",
+    "glm_moe_dsa",
+    "gpt_oss",
+    "kimi_k2",
+    "mimo",
+    "mimo_v2_flash",
+    "minimax_m2",
+    "mixtral",
+    "nemotron",
+    "nemotron_h",
+    "olmoe",
+    "qwen2_5_vl",
+    "qwen2_vl",
+    "qwen3_vl",
+    "qwen3_vl_moe",
+    "step3p5",
+)
+
+
+def test_megatron_bridge_supported_models_match_expected():
+    assert MEGATRON_BRIDGE_SUPPORTED_MODELS == EXPECTED_MEGATRON_BRIDGE_SUPPORTED_MODELS
+
+
+@pytest.mark.parametrize("model_type", sorted(EXPECTED_MEGATRON_BRIDGE_SUPPORTED_MODELS))
+def test_check_model_support_accepts_supported(model_type: str):
+    with patch(
+        "llamafactory.train.megatron_bridge.workflow.HfAutoConfig.from_pretrained",
+        return_value=SimpleNamespace(model_type=model_type),
+    ):
+        _check_model_support(SimpleNamespace(model_name_or_path="dummy", trust_remote_code=False))
+
+
+@pytest.mark.parametrize("model_type", UNSUPPORTED_MEGATRON_BRIDGE_MODELS)
+def test_check_model_support_rejects_unsupported(model_type: str):
+    assert model_type not in MEGATRON_BRIDGE_SUPPORTED_MODELS
+    with patch(
+        "llamafactory.train.megatron_bridge.workflow.HfAutoConfig.from_pretrained",
+        return_value=SimpleNamespace(model_type=model_type),
+    ):
+        with pytest.raises(ValueError, match="not supported by the Megatron Bridge"):
+            _check_model_support(SimpleNamespace(model_name_or_path="dummy", trust_remote_code=False))