[v0] support megatron-bridge for PT/SFT training#10645
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request integrates the Megatron Bridge distributed training backend into LLaMA-Factory, enabling pretraining (PT) and supervised fine-tuning (SFT) with advanced parallelism features (tensor, pipeline, and context parallelism). Key additions include new configuration arguments, workflow orchestration, dataset exporting, and comprehensive unit tests. The review feedback highlights several important improvements: wrapping dataset exports in a local rank 0 check to prevent race conditions in distributed environments, capping warmup steps to avoid scheduler errors, warning users about ignored custom LoRA targets, supporting deeply nested overrides, stripping whitespace from JSON configurations, and validating that the world size is divisible by the total parallel size.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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", | ||
| ) |
There was a problem hiding this comment.
In distributed training, run_pt is executed on all ranks simultaneously. Since export_dataset_for_megatron_bridge writes to the same shared training.jsonl file, this causes a severe race condition and file corruption. We should wrap the dataset export in a local rank 0 check and use a file-based barrier to synchronize other ranks.
import time
done_file = os.path.join(dataset_dir, ".export_done")
if training_args.local_process_index == 0:
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",
)
with open(done_file, "w", encoding="utf-8") as f:
f.write("done")
else:
while not os.path.exists(done_file):
time.sleep(1)| 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, | ||
| ) |
There was a problem hiding this comment.
In distributed training, run_sft is executed on all ranks simultaneously. Since export_dataset_for_megatron_bridge writes to the same shared training.jsonl file, this causes a severe race condition and file corruption. We should wrap the dataset export in a local rank 0 check and use a file-based barrier to synchronize other ranks.
import time
done_file = os.path.join(dataset_dir, ".export_done")
if training_args.local_process_index == 0:
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,
)
with open(done_file, "w", encoding="utf-8") as f:
f.write("done")
else:
while not os.path.exists(done_file):
time.sleep(1)| if warmup_steps > 0: | ||
| return warmup_steps |
There was a problem hiding this comment.
If warmup_steps is greater than train_iters (e.g., when training with a very small number of steps like max_steps: 5), it can cause unexpected behavior or division-by-zero errors in the Megatron scheduler. We should cap warmup_steps at train_iters.
| if warmup_steps > 0: | |
| return warmup_steps | |
| if warmup_steps > 0: | |
| return min(warmup_steps, train_iters) |
| 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, | ||
| ) |
There was a problem hiding this comment.
Megatron Bridge uses Megatron-specific module names for PEFT/LoRA. If a user specifies a custom lora_target in their configuration, it is silently ignored. We should add a warning to inform the user that custom lora_target settings are not supported and default Megatron targets are used.
def _create_peft_config(finetuning_args: "FinetuningArguments"):
if finetuning_args.finetuning_type != "lora":
return None
from megatron.bridge.peft.lora import LoRA
if finetuning_args.lora_target != ["all"]:
logger.warning_rank0(
f"Custom lora_target {finetuning_args.lora_target} is not supported by Megatron Bridge. "
"Using default Megatron target modules: ['linear_qkv', 'linear_proj', 'linear_fc1', 'linear_fc2']."
)
return LoRA(
target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"],
dim=finetuning_args.lora_rank,
alpha=finetuning_args.lora_alpha,
)| 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) |
There was a problem hiding this comment.
The current implementation of _apply_extra_overrides only supports exactly one level of nesting (e.g., target.attr). If a dot-path has multiple levels of nesting (e.g., target.subtarget.attr), it fails or sets a literal attribute named "subtarget.attr" on the intermediate object. We should implement recursive traversal to support arbitrary nesting.
| 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 _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) |
| if isinstance(self.extra_config, str) and self.extra_config.startswith("{"): | ||
| self.extra_config = _convert_str_dict(json.loads(self.extra_config)) |
There was a problem hiding this comment.
If self.extra_config is a JSON string with leading whitespace, self.extra_config.startswith("{") will evaluate to False. It will then fail to parse as JSON and instead be treated as a file path, leading to a ValueError. We should strip the string before checking.
| if isinstance(self.extra_config, str) and self.extra_config.startswith("{"): | |
| self.extra_config = _convert_str_dict(json.loads(self.extra_config)) | |
| 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)) |
| if parallel_size > world_size: | ||
| raise ValueError(f"Total Megatron Bridge parallel size ({parallel_size}) exceeds `world_size` ({world_size}).") |
There was a problem hiding this comment.
Megatron strictly requires that the total world_size is divisible by the product of parallel sizes (so that the data parallel size is an integer). We should validate world_size % parallel_size == 0 to catch configuration errors early.
| if parallel_size > world_size: | |
| raise ValueError(f"Total Megatron Bridge parallel size ({parallel_size}) exceeds `world_size` ({world_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})." | |
| ) |
Test EnvironmentHardware & System
Core Dependencies (verified working)
Installation Steps# 1. Create virtual environment
python3 -m venv megatron-bridge-env
source megatron-bridge-env/bin/activate
# 2. Install PyTorch (CUDA 12.6)
pip install torch==2.12.1+cu126 torchvision==0.27.1+cu126 \
--index-url https://download.pytorch.org/whl/cu126
# 3. Install Transformer Engine (requires NVIDIA PyPI)
pip install --no-build-isolation \
--extra-index-url https://pypi.nvidia.com \
transformer_engine[pytorch]
# 4. Install nvidia-resiliency-ext (required by megatron-core 0.18.0; 0.6.0 not on PyPI)
git clone -b v0.6.0 https://github.com/NVIDIA/nvidia-resiliency-ext.git
pip install -e ./nvidia-resiliency-ext
# 5. Install megatron-bridge
pip install --no-build-isolation megatron-bridge==0.5.0
# 6. Upgrade megatron-core to 0.18.0
pip install megatron-core==0.18.0
# 7. Install LLaMA-Factory
cd LLaMA-Factory
pip install -e .Training Commandexport USE_MEGATRON_BRIDGE=1
export DISABLE_VERSION_CHECK=1 # Skip version check when transformers doesn't match megatron-bridge constraints
llamafactory-cli train examples/megatron_bridge/llama3_sft.yaml
Verified Test Case
Known Limitations
|
1 participant
What does this PR do?
Implementation Objective
Add NVIDIA Megatron Bridge as a new training backend in LLaMA-Factory, alongside existing HuggingFace / MCA / HyperParallel backends. Enabled via environment variable
USE_MEGATRON_BRIDGE=1, currently covering:Architecture & Data Flow
File-Level Specifications
Parameters & Entry Points
hparams/megatron_bridge_args.pyhparams/finetuning_args.pyuse_megatron_bridge+ runtime-attachedmegatron_bridge_argshparams/parser.py_parse_train_mbridge_args, parallelism validation, mutual exclusion with MCA/HP/quantization/DeepSpeedextras/packages.pyis_megatron_bridge_available()launcher.pyFORCE_TORCHRUNwhen enabledtrain/tuner.pyCore Training Path (new
train/megatron_bridge/)workflow.py
run_pt/run_sft: Load LF dataset → export → build config → call Bridge APIRuntime patches:
makefor pre-compiledhelpers_cppload_megatron_model+save_hf_pretrained)config_builder.py
TrainingArguments→ MegatronConfigContainerensure_megatron_pretrained_checkpoint: HF→Megatron conversion, cleanup after conversiongradient_accumulation_fusionwhen APEX fusion is missingextra_configdot-path overridesdataset_export.py
LF-aligned samples → Bridge JSONL
{% generation %}), otherwise fallback to ShareGPTCompatibility with different Bridge versions for dataloader / chat template APIs
Examples & Tests
examples/megatron_bridge/llama3_sft.yaml: Llama-3.2-1B, TP=2, SP, bf16, export_hf_on_finishUnit tests: args validation, schedule, fusion safety, dataset export, GPU config / AutoBridge (skip when megatron-bridge not installed)
Capability Boundaries (v0)
Design Principles
Before submitting