Training entry points

August 25, 2026 · View on GitHub

Megatron Bridge training provides a small public Slurm launcher:

./scripts/training/train.sh [launch options] [runner options] [KEY=VALUE overrides]

train.sh invokes setup_experiment.py on a Slurm login node and submits run_recipe.py directly through Slurm. The setup layer only owns resources, the container, explicitly forwarded environment variables, and explicit mounts. Recipe selection, dataset construction, and ConfigContainer overrides are resolved inside the training environment. Without an active virtual environment, the shell entry point creates an isolated nemo-run environment rather than resolving the full GPU training dependency set on the login node.

launch_with_nemo_run.py and launch_with_sbatch.sh remain available for their existing specialized workflows; train.sh is the compact recipe-oriented path.

Selection rules

Choose exactly one of a complete recipe or a model selector. --recipe and --model are mutually exclusive. A complete recipe is discovered automatically from its exported function name, whether it is a library or benchmark recipe; there is no separate source flag. A model selector requires one of --mode pretrain, --mode sft, --mode lora, or --mode dora; a conventional complete recipe name infers its mode when --mode is omitted.

Library recipe

A complete recipe already identifies the model and default training configuration, so do not also pass --model:

./scripts/training/train.sh \
    --nodes 1 --gpus-per-node 8 \
    --account ACCOUNT --partition PARTITION \
    --container-image /path/to/container.sqsh \
    --recipe llama32_1b_pretrain_config \
    --mode pretrain --dataset mock

Model selector

The model selector combines the model stem and mode to load the corresponding library recipe. For example, --model gpt_oss_20b --mode sft loads gpt_oss_20b_sft_config; LoRA and DoRA load the model's PEFT recipe and set the requested adapter scheme.

./scripts/training/train.sh \
    --nodes 2 --gpus-per-node 8 \
    --account ACCOUNT --partition PARTITION \
    --container-image /path/to/container.sqsh \
    --model gpt_oss_20b \
    --mode pretrain --dataset mock

Benchmark recipe

The training launcher can run exact exported recipes from src/megatron/bridge/perf_recipes, including text pretraining, text SFT/PEFT, Qwen-VL pretraining, and Wan pretraining. The GPU count encoded by the recipe name is its canonical tuned allocation. The user selects the node shape and may weak-scale to a different compatible world size; the Slurm partition must provide the requested hardware:

benchmark is the unified runner's user-facing term. The existing perf_recipes package and scripts/performance/ compatibility paths retain their legacy names.

./scripts/training/train.sh \
    --nodes 2 --gpus-per-node 8 \
    --account ACCOUNT --partition PARTITION \
    --container-image IMAGE \
    --recipe qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config \
    --mode pretrain

Benchmark recipes provide canonical defaults for their dataset, parallelism topology, batch sizes, sequence length, precision, dispatcher, CUDA-graph settings, and process environment. Trailing KEY=VALUE overrides are applied to their ConfigContainer in the same way as library recipes; an overridden run no longer represents the canonical benchmark configuration. Benchmark recipes retain their selected dataset type, so --dataset is not supported on this path. Recipe environment defaults are installed before the launcher enters the training stack; values explicitly set by the shell or Slurm environment retain precedence.

When the requested world size differs from the canonical count, the runner preserves the recipe's samples per GPU by scaling train.global_batch_size proportionally. The scaled batch size must be an integer, and the requested world size must be divisible by the resolved TP × PP × CP topology. Pass --global_batch_size or a trailing train.global_batch_size=... override to choose an explicit batch size instead of automatic weak scaling.

The launcher does not infer offline mode, cluster-specific CPU/NUMA binding, Slurm segment sizing, or NCCL fabric settings from the recipe name. Supply those deployment settings explicitly through the target cluster's launcher configuration, repeated --srun-arg options, or exported values forwarded with --env NAME.

The runner selects the registered text, multimodal, audio, omni, or diffusion forward step from the recipe identity, regardless of which recipe package exports it. The compatibility launcher at scripts/performance/setup_experiment.py remains available for selector-based invocation, dataset replacement, and specialized legacy controls that are not part of the compact training CLI.

Text SFT/PEFT benchmark recipes retain the flat runner's mock-data default. Qwen-VL and Wan recipes retain their model-specific dataset configuration. Exported benchmark PEFT recipes are fixed LoRA configs; DoRA remains available through configurable library recipes. Explicit benchmark dataset replacement remains on the compatibility launcher.

Five legacy duplicate names resolve to the benchmark definition; their library workloads remain available through the corresponding generic recipe aliases. New recipe names should be unique across both packages.

Text recipes default to llm_step; all recipes infer their modality-specific forward step from the same registry. Pass --step-func NAME to override that selection. Common training, sequence-length, parallelism, optimization, and checkpoint fields also have convenience flags such as -ms/--max_steps, -sl/--seq_length, -tp/--tensor_model_parallel_size, and --save_dir. Use trailing KEY=VALUE overrides for every other ConfigContainer field.

Dataset selection

For library recipes, --dataset accepts source selectors and named dataset presets rather than internal dataset config class names. Each name selects a DatasetConfig preset; the config type selects the existing runtime builder. Trailing dataset.* overrides are applied directly after preset selection. Use typed fields for source, preprocessing, packing, and loader settings; use dataset.dataset_kwargs={...} only for extra options consumed by a dataset implementation. Benchmark recipes retain their recipe-owned dataset and reject --dataset.

ValueKindModeBehavior
mockSource selectorpretrainIn-memory generated GPT data
megatron-indexedSource selectorpretrainLocal Megatron .bin/.idx data; never falls back to mock
energonSource selectorpretrain/sft/lora/doraLocal Energon data selected by dataset.path; preserves a recipe-owned task encoder or uses its Hugging Face processor
local-jsonlSource selectorsft/lora/doraLocal prompt-completion JSONL selected by dataset.dataset_root
local-vlmSource selectorsft/lora/doraLocal VLM JSON/JSONL selected through dataset.source overrides
squadNamed presetsft/lora/doraHugging Face SQuAD preset
tulu3Named presetsft/lora/doraAi2 Tulu 3 SFT mixture (allenai/tulu-3-sft-mixture)
coderforgeNamed presetsft/lora/doraCoderForge Preview SWE_Rebench agent-trajectory split
openmathinstruct2Named presetsft/lora/doraOpenMathInstruct-2 prompt/completion preset
openmathinstruct2-thinkingNamed presetsft/lora/doraOpenMathInstruct-2 analysis/final channel format
gsm8kNamed presetsft/lora/doraGSM8K preset
cord-v2Named presetsft/lora/doraCORD-v2 receipt VQA preset
medpixNamed presetsft/lora/doraMedPix medical VQA preset
ravenNamed presetsft/lora/doraRAVEN visual reasoning preset
rdrNamed presetsft/lora/doraRDR visual reasoning preset
llava-video-178kNamed presetsft/lora/doraLLaVA-Video preset; requires dataset.source.adapter_kwargs.video_root_path

Megatron indexed data

Any pretraining corpus must first be converted to Megatron indexed data. Set one prefix or a list of prefixes directly on the dataset config:

--dataset megatron-indexed dataset.data_path=/data/dclm/dclm_01_01_text_document
--dataset megatron-indexed 'dataset.data_path=[/data/dclm/part_1,/data/dclm/part_2]'

Every selected prefix must have matching .bin and .idx files. Use dataset.path_to_cache=/shared/cache to select the index cache.

The launcher does not infer the source corpus from the files. For one preprocessing example, see the DCLM tutorial.

Energon data

Use energon with a VLM recipe that already owns an EnergonDatasetConfig or exposes dataset.hf_processor_path. Select pretraining, SFT, LoRA, or DoRA independently with --mode; the selector preserves a recipe-owned model-specific task encoder and otherwise constructs the generic Hugging Face task encoder.

For example, the existing Qwen 35B-A3B recipe can train on prepared DataComp shards with explicit dataset, processor, and batch settings:

--dataset energon \
train.global_batch_size=512 \
train.micro_batch_size=1 \
dataset.micro_batch_size=1 \
dataset.path=/data/datacomp-energon \
dataset.task_encoder.hf_processor_path=Qwen/Qwen3.6-35B-A3B \
dataset.task_encoder.hf_processor_revision=995ad96eacd98c81ed38be0c5b274b04031597b0

Set both micro-batch fields when overriding the recipe's micro-batch size; the runner does not silently synchronize them. ChatML data must use the repository's ChatMLWebdataset tar-member contract. The DataComp tutorial documents a complete image-caption preparation and training flow; the multimodal Energon tutorial documents the general tar-member contract. The Nemotron Image Training v3 tutorial shows a pinned small-subset download, conversion, and Qwen3-VL train.sh launch through this selector.

OpenMathInstruct-2

openmathinstruct2 uses prompt/completion records. openmathinstruct2-thinking changes only the semantic output format: chain-of-thought goes to the analysis channel and the answer goes to the final channel.

Tulu

tulu3 uses the canonical current Tulu SFT dataset, allenai/tulu-3-sft-mixture. The preset reads its published train split and its native messages chat schema; it reserves 5% for validation by default. The mixture is licensed under ODC-BY-1.0, but individual subsets can carry additional terms, including non-commercial restrictions. Review the Hugging Face dataset card and its linked subset licenses before use.

CoderForge Preview

coderforge selects the trajectories configuration and SWE_Rebench split of togethercomputer/CoderForge-Preview. The preset decodes the dataset's JSON-string messages and tools fields and uses assistant-only chat loss. Its image field identifies the execution environment and is retained as metadata; it is not treated as visual input. Use a Hugging Face split slice such as 'dataset.hf_dataset.split="SWE_Rebench[:64]"' for bounded smoke runs; Hydra requires the quoted value to preserve the brackets.

Offline packing

Offline packing is a text SFT option, not a dataset name. Set dataset.enable_offline_packing=true for squad, either OpenMathInstruct-2 format, tulu3, coderforge, gsm8k, or local-jsonl. The launcher aligns packed padding for the resolved CP/TP and sequence-parallel configuration. Packed training requires train.micro_batch_size=1. The builder materializes packed data automatically, so a separate packing Slurm job is not required. The selected recipe/model must support packed THD sequences; for example, GLM-4.5 and Qwen3-Next recipes currently do not. On multiple nodes, keep the dataset cache on shared mounted storage.

In-batch packing for GPT SFT JSONL

As an alternative to offline artifacts, GPTSFTDatasetConfig can pack each logical microbatch during collation. This supports both prompt-completion and chat JSONL preprocessing while retaining the local JSONL mmap path; only the selected rows are parsed and tokenized, so the full dataset is not loaded into host memory. For example:

--dataset local-jsonl \
dataset.dataset_root=/data/sft \
dataset.enable_in_batch_packing=true \
dataset.dataloader_type=single \
train.micro_batch_size=4

The logical micro-batch must be larger than one. Collation emits one physical THD row and derives per-sequence CP/SP alignment from the resolved topology. Use the single or cyclic dataloader; the global-batch batch dataloader is not supported. Do not combine dataset.enable_in_batch_packing=true with offline packing.

VLM data

The hosted VLM names use the existing Hugging Face dataset adapters and retain the processor and in-batch packing settings from the selected VLM recipe. raven, rdr, and llava-video-178k derive deterministic 95/5 train and validation slices; cord-v2 and medpix use their published validation splits. The runner selects the registered VLM forward step from the recipe name; pass --step-func only to override it.

For local JSON/JSONL, select local-vlm and set dataset.source.load_kwargs.data_files.train=/path/to/train.jsonl. Optional split overrides are dataset.validation_source.load_kwargs.data_files.validation and dataset.test_source.load_kwargs.data_files.test; enable those stages explicitly with dataset.do_validation=true and dataset.do_test=true. Rows and media paths must already use the selected processor's supported conversation schema. dataset.hf_processor_path overrides the processor inherited from the VLM recipe.

SFT and LoRA checkpoints

Set checkpoint.pretrained_checkpoint to the pretrained native Megatron checkpoint or local Hugging Face model directory:

./scripts/training/train.sh \
    --nodes 1 --gpus-per-node 8 \
    --account ACCOUNT --partition PARTITION --container-image IMAGE \
    --mount /checkpoints \
    --recipe gpt_oss_20b_sft_8gpu_h100_bf16_config \
    --mode sft --dataset openmathinstruct2-thinking \
    dataset.enable_offline_packing=true \
    train.micro_batch_size=1 \
    checkpoint.pretrained_checkpoint=/checkpoints/gpt-oss-20b
./scripts/training/train.sh \
    --nodes 1 --gpus-per-node 1 \
    --account ACCOUNT --partition PARTITION --container-image IMAGE \
    --mount /checkpoints \
    --recipe gpt_oss_20b_peft_1gpu_h100_bf16_config \
    --mode lora --dataset openmathinstruct2-thinking \
    dataset.enable_offline_packing=true \
    train.micro_batch_size=1 \
    checkpoint.pretrained_checkpoint=/checkpoints/gpt-oss-20b

Overrides

Every selected recipe's ConfigContainer fields can be set with trailing KEY=VALUE arguments. Common fields also accept the convenience flags listed by run_recipe.py --help. For example, batch sizes and training duration belong under train, not model:

./scripts/training/train.sh \
    --nodes 1 --gpus-per-node 8 \
    --account ACCOUNT --partition PARTITION \
    --container-image /path/to/container.sqsh \
    --recipe llama32_1b_pretrain_config \
    --mode pretrain --dataset mock \
    -ms 10 -gb 8 -mb 1 \
    --lr 0.0002 --warmup_iters 1 \
    scheduler.lr_decay_iters=10

Precedence is recipe defaults, the selected dataset configuration, common convenience arguments, then trailing ConfigContainer overrides. A trailing override therefore wins when it targets the same field as a convenience argument.

Overrides take a benchmark recipe outside its canonical configuration. Use scripts/performance/setup_experiment.py for specialized workflows that the unified launcher has not migrated yet.

Slurm and containers

Required Slurm arguments are:

--nodes
--gpus-per-node
--account
--partition
--container-image (or CONTAINER_IMAGE)

Set CONTAINER_IMAGE to avoid repeating --container-image. On clusters that allocate whole GPU nodes implicitly, pass --no-gpu-resource-request to omit the explicit Slurm GPU request while retaining one task per requested GPU. Environment variables and filesystem paths are never forwarded implicitly. Export credentials in the launcher environment, then repeat --env NAME to forward names without materializing their values in the generated sbatch script. Repeat --mount HOST for the same host and container path, or use --mount HOST:CONTAINER when the paths differ. Mount every dataset, checkpoint, cache, and output path the job needs. The launcher adds no cluster-specific srun flags by default. Repeat --srun-arg=ARG for every flag required by the target cluster. For example, a Pyxis/Enroot cluster may use:

--srun-arg=--mpi=pmix \
--srun-arg=--no-container-mount-home \
--srun-arg=--container-writable

The = form is required when ARG begins with -.

Pass target-specific sbatch settings through --additional-slurm-params 'KEY=VALUE;OTHER_KEY=OTHER_VALUE'. These values populate the executor's additional_parameters and are not forwarded to the training command or its srun step.

Use -lmc FREQ, --peak-mem-clk FREQ, or --peak_mem_clk FREQ to lock the GPU memory clock to a fixed frequency in MHz once per node before training. VR200 benchmark recipes default to 4752; other recipes leave the memory clock unlocked. Pass -lmc -1 to disable the VR200 default explicitly.

The compact launcher does not add rank-command prefixes or infer target-specific allocation and NUMA policies. Configure those explicitly for the target.

The launcher submits the experiment in detached mode and returns after Slurm accepts the job. Inspect its state and logs with the cluster's normal squeue, sacct, and log-file workflow.

Use --dry-run (or the explicit --submission-dry-run spelling) to render the NeMo-Run experiment without submitting it. For a rank-local ConfigContainer dry run, invoke run_recipe.py directly:

uv run python scripts/training/run_recipe.py \
    --recipe llama32_1b_pretrain_config \
    --mode pretrain --dataset mock \
    --dry-run logger.save_config_filepath=/tmp/config.yaml

For a benchmark dry run, use the complete flat recipe name and omit --dataset. Without a distributed world-size environment, the rank-local dry run validates the canonical allocation encoded in the recipe name. If WORLD_SIZE or SLURM_NTASKS is set, it validates and weak-scales against that requested size instead. The submission dry run validates the launcher arguments while leaving topology validation to the rank-local resolved config.

Rank-local entry point

run_recipe.py remains available for existing distributed environments that already own process launch:

uv run python -m torch.distributed.run --nproc_per_node=2 \
    scripts/training/run_recipe.py \
    --recipe vanilla_gpt_pretrain_config \
    --mode pretrain --dataset mock \
    model.tensor_model_parallel_size=2 \
    model.sequence_parallel=true

This entry point loads library recipes for pretraining, SFT, LoRA, and DoRA, and all exact exported benchmark recipes when selected explicitly.