This repo provides end-to-end infrastructure for training and evaluating large language models on Isambard's ARM GH200 cluster using NeMo Megatron Bridge. It wraps Megatron Bridge's conversion and training APIs with SLURM pipelines, fault tolerance, and Isambard-specific workarounds (ARM aarch64, Slingshot networking, containerized execution via Apptainer + NGC images).
The primary workflow is: download a HuggingFace dataset → prepare and pack it → train with Megatron-Core MoE parallelism (TP/EP/PP/DP) → convert checkpoints back to HuggingFace format → run generation tests. All training metrics and generation outputs are logged to Weights & Biases. The current infrastructure is optimized for Nemotron 3 Nano (30B-A3B) and Super (120B-A12B) MoE models; future releases will generalize to additional model families.
For cluster hardware specs and per-model topology findings, see CLAUDE.md; for the execution environment itself, docs/environment.md. The upstream Megatron Bridge README is at docs/README_DEFAULT.md.
All top-level scripts follow the PIPELINE_ACTION.ext naming convention. There are five pipelines:
| Pipeline | Submit (SLURM) | Launch / Logic | W&B Project | Purpose |
|---|---|---|---|---|
| env | pipeline_env_submit.sbatch |
pipeline_env_config.env, pipeline_env_exec.sh, pipeline_env_activate.sh, pipeline_env_setup.sh, pipeline_env_validate.py |
— | The execution environment: Apptainer + NGC NeMo image, Slingshot NCCL stack, install + validation (docs/environment.md) |
| data | pipeline_data_submit.sbatch |
pipeline_data_prepare.py |
geodesic/megatron-datasets-processing |
Dataset download, tokenization, packing |
| training | pipeline_training_submit.sbatch |
pipeline_training_launch.sh |
geodesic/megatron_training |
SFT, CPT, and from-scratch pretraining |
| checkpoint | pipeline_checkpoint_submit.sbatch |
pipeline_checkpoint_convert.sh, pipeline_checkpoint_convert_hf.py |
— | Megatron↔HF conversion, Hub upload |
| coherence | pipeline_coherence_submit.sbatch |
pipeline_coherence_test.py |
geodesic/geodesic-gen-tests |
Qualitative generation testing |
Each PIPELINE_submit.sbatch allocates SLURM nodes and delegates to the logic script. The .sh launchers can also be called directly from an interactive salloc session.
This walkthrough runs a complete Nemotron 3 Nano SFT training run on the shipped quickstart config, covering every
pipeline from data through coherence testing. The quickstart is the control-pretraining XL SFT — the reasoning
post-training stage of the control-pretraining curriculum, on packed sequences of 32,768 tokens — in its fastest
configuration, and the throughput figures quoted below are the measured numbers of its performance campaign
(docs/investigations/nano30b-sft-perf-campaign.md).
What you'll do: Train on 16 nodes (3.69 s/iter; a 200-iteration demo is about 20 min) → convert the checkpoint to HuggingFace format (about 10 min) → run generation tests (15 min).
Prerequisites: The environment must be installed once (bash pipeline_env_setup.sh on a GPU node — see Step 0).
The run warm-starts from the control-pretraining midtraining final checkpoint, which is already on the cluster at the
pretrained_checkpoint the config inherits, and its dataset is already packed there, so nothing needs converting or
preparing first.
Note: The current pipeline infrastructure (configs, recipes, conversion scripts, coherence tests) is optimized for Nemotron 3 Nano and Super. Future releases will generalize the tooling to support additional model families out of the box.
Every pipeline runs inside one Apptainer container built from an NGC NeMo image
(aarch64), which supplies PyTorch, CUDA, NCCL, Transformer Engine, the Mamba kernels, APEX,
and ft_launcher prebuilt and version-matched. This repo's src/ and the pinned
3rdparty/Megatron-LM are bind-mounted, so the checkout you submit from is exactly the code
that runs. The launchers enter the container themselves — there is nothing to activate before
submitting, and every submit command below is the whole command.
Installation (per cluster, shared across users via /projects/a5k/public/containers/) is a
single idempotent command:
# On a GPU node: SIF pull + Slingshot NCCL build + Python overlay + validation.
# Completed steps are skipped with an explicit message; --force redoes everything.
bash pipeline_env_setup.sh
# or via SLURM: isambard_sbatch pipeline_env_submit.sbatch setup
# Re-validate an existing install (21 checks; 22 with --run-training):
isambard_sbatch pipeline_env_submit.sbatch validate --run-trainingDesign decisions, image contents, the image-qualification gates, and troubleshooting: docs/environment.md.
The quickstart trains on geodesic-research/pa-warm-start-sft-xl-50b-mix (its default config at revision
ec0b9197: about 50B tokens of reasoning chat SFT), packed at seq_length 32768 with the think-history tokenizer into
32 shards, already prepared on the cluster at the dataset_root the config inherits. Its identity and pack geometry
are versioned in configs/control_pretraining/30b_baseline_ablations/data/pa-warm-start-sft-xl-50b-mix.yaml, and it
is rebuilt from that directory's corpora.tsv with configs/control_pretraining/build_corpora.sh. To prepare a
different dataset, use the Data Pipeline with --seq-length 32768; a pad multiple of at least
2×CP keeps the packs usable under context parallelism.
The quickstart config is
configs/quickstart/nemotron_nano_quickstart_sft.yaml, and it
is a chain of base_config: overlays rather than one file:
configs/control_pretraining/30b_baseline_ablations/nemotron_nano_30b_baseline_sft_xl50b_gbs256.yamlis the production XL SFT: model, data, tokenizer, warm start, topology (TP1·CP2·EP4·PP1) and schedule.nemotron_nano_quickstart_sft_baseline.yamlmakes it a benchmark: GBS 64 on 64 GPUs (production's two packs per data-parallel replica), an exit at iteration 100 on production's schedule, no checkpoint I/O, its own W&B name.nemotron_nano_quickstart_sft.yamladds the performance campaign's levers:- context parallelism off (CP=1, DP=64, one whole pack per GPU), which fits at full recompute because the chunked linear cross-entropy never materialises the 16 GiB of fp32 logits a 32K sequence's unfused loss keeps;
- BF16 gradient reduction in 500M-parameter buckets, with the parameter all-gather overlapped with the forward;
- the HybridEP dispatcher with the fused router, on packs padded to the full 32,768 tokens;
- host settings: timers at level 1, the loss NaN check and per-parameter gradient norms off, garbage collection on a schedule.
Each file's header carries its measured story and why every field is what it is. A base_config: path resolves
against the directory of the file that names it, so a copy of a quickstart goes beside it in configs/quickstart/ or
restates the base path.
The training pipeline has two layers: a thin SLURM wrapper
(pipeline_training_submit.sbatch) that allocates nodes, and a shared launcher
(pipeline_training_launch.sh) that configures NCCL, Slingshot networking, and starts
the distributed job. The nano sft arguments select the model recipe and training
mode, and the config's .env file, given as ISAMBARD_ENV_OVERRIDES, carries the launcher setting it needs.
--disable-ft keeps the short run restart-free.
The benchmark run, exactly as measured (100 iterations, no checkpoint):
ISAMBARD_ENV_OVERRIDES=$PWD/configs/quickstart/nemotron_nano_quickstart_sft.env \
isambard_sbatch --nodes=16 --time=00:20:00 pipeline_training_submit.sbatch \
configs/quickstart/nemotron_nano_quickstart_sft.yaml nano sft --disable-ftFor a demo that trains longer and keeps a checkpoint for Steps 5–6, add Hydra overrides (the launcher forwards them):
ISAMBARD_ENV_OVERRIDES=$PWD/configs/quickstart/nemotron_nano_quickstart_sft.env \
isambard_sbatch --nodes=16 --time=00:30:00 pipeline_training_submit.sbatch \
configs/quickstart/nemotron_nano_quickstart_sft.yaml nano sft --disable-ft \
train.exit_interval=200 \
checkpoint.save=/projects/a5k/public/checkpoints/megatron/nemotron_nano_quickstart_sft \
checkpoint.save_optim=false checkpoint.save_rng=falseThe run exits at train.exit_interval and, because checkpoint.save is set, writes exactly one checkpoint there;
train_iters stays the production stage's, so the learning rate at every iteration is production's.
save_optim/save_rng: false skip the Adam moments and RNG state the downstream conversion never reads.
tail -f /projects/a5k/public/logs/megatron_runs/train-<jobid>.out | grep --line-buffered "iteration"Measured behaviour of this config (2026-10-02, 64 GPUs, one switch group): the first iteration takes about a minute
(kernel compilation and communicator setup), every later step is within about 1% of the mean, and the mean is
3.689 s/iter (8,881 tokens/s/GPU) at a 70.4 GiB peak of 95. The loss of the warm-started model
starts near 1.03 and falls below 1.0 within the first hundred iterations; 0 NaN. Metrics stream live to
wandb.ai/geodesic/megatron_training under the run name
nemotron_nano_quickstart_sft_perf.
Nano converts on a single node (4 GPUs) with node-local EP — no Slingshot needed. Two details are specific to this
model. It trains with the torch_grouped expert backend, whose checkpoint run_config.yaml names a model spec the
exporter cannot import, so the export reads a clone of the checkpoint: links to its files plus a corrected
run_config.yaml (make_export_clone in scripts/hub/publish_models.py; the source is never written to). And it is a
reasoning (think) SFT, so it exports with --reasoning, which keeps the thinking chat template:
python3 -c 'import sys; from pathlib import Path; from scripts.hub.publish_models import make_export_clone; make_export_clone(Path(sys.argv[1]), Path(sys.argv[2]))' \
/projects/a5k/public/checkpoints/megatron/nemotron_nano_quickstart_sft/iter_0000200 \
/projects/a5k/public/checkpoints/megatron/nemotron_nano_quickstart_sft_export/iter_0000200
isambard_sbatch --nodes=1 pipeline_checkpoint_submit.sbatch export \
/projects/a5k/public/checkpoints/megatron/nemotron_nano_quickstart_sft_export \
--hf-model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 --reasoning \
--iteration 200The converted model lands at
.../nemotron_nano_quickstart_sft_export/iter_0000200/hf/ — a standard HF checkpoint
(safetensors + config + tokenizer) loadable with AutoModelForCausalLM.
Nano fits on a single GPU for generation:
isambard_sbatch --gpus-per-node=1 pipeline_coherence_submit.sbatch \
/projects/a5k/public/checkpoints/megatron/nemotron_nano_quickstart_sft_export/iter_0000200/hfThis generates responses to 8 diverse prompts and logs a table (prompt, response,
length, empty-flag) plus summary metrics to W&B. Read a few generations yourself —
empty_pct == 0 alone does not catch drift, off-topic output, or template artifacts.
The environment is an Apptainer container built from a qualified NGC NeMo image; the full reference is docs/environment.md.
# Install everything (SIF pull + Slingshot NCCL build + Python overlay + validate).
# GPU node required for the build and the validation; idempotent, --force redoes all.
isambard_sbatch pipeline_env_submit.sbatch setup
bash pipeline_env_setup.sh --only slingshot --force # or one step at a time
# Validate an existing install (21 checks; 22 with --run-training)
isambard_sbatch pipeline_env_submit.sbatch validate --run-training
# Run anything inside the environment (interactive shell, tests, ad-hoc python)
./pipeline_env_exec.sh "cd $PWD; source pipeline_env_activate.sh || exit 1; exec bash -i"
./pipeline_env_exec.sh "bash $PWD/scripts/run_unit_tests.sh" # the unit tests, as the pre-commit gate runs themEverything configurable — image tag, SIF path, Slingshot component versions, overlay
packages, bind list, cache dirs — lives in pipeline_env_config.env. A missing SIF or
Slingshot build hard-fails every launcher with the command that fixes it; there is no
fallback environment.
- A Python recipe defines the base model config, optimizer, parallelism, and data pipeline
- A YAML config file overrides recipe defaults
- SFT and CPT load a pretrained checkpoint and train via finetune(); pretrain mode random-initializes from the NVIDIA pretrain recipes and trains via pretrain() — no checkpoint is loaded unless the YAML sets one
# Via SLURM — extra args after the mode forward to the launcher: launcher flags
# (e.g. --disable-ft) are parsed as such, anything else falls through as Hydra overrides
isambard_sbatch --nodes=32 pipeline_training_submit.sbatch configs/<config>.yaml nano sft
isambard_sbatch --nodes=8 pipeline_training_submit.sbatch configs/<config>.yaml nano cpt
isambard_sbatch --nodes=32 pipeline_training_submit.sbatch configs/<config>.yaml nano pretrain --disable-ft
isambard_sbatch --nodes=16 pipeline_training_submit.sbatch configs/<config>.yaml super sft \
--disable-ft train.train_iters=32 checkpoint.save=null
# Via salloc
salloc --nodes=16 --gpus-per-node=4 --time=24:00:00 --exclusive
bash pipeline_training_launch.sh configs/<config>.yaml --model nano --mode sft
bash pipeline_training_launch.sh configs/<config>.yaml --model nano --mode sft --disable-ft
bash pipeline_training_launch.sh configs/<config>.yaml --model nano --mode sft --peft loraTo change environment variables for one launch, including ones the launcher or
pipeline_env_activate.sh set unconditionally (e.g. TORCH_NCCL_BLOCKING_WAIT=0), point
ISAMBARD_ENV_OVERRIDES at a file of KEY=VALUE lines; see
docs/environment.md D2b for what it applies where and what it refuses.
To measure a training lever, submit it as its own short job on the quickstart posture (or on the
baseline benchmark, nemotron_nano_quickstart_pretrain_baseline.yaml without the env file, to measure
against production's posture) and score the log with scripts/telemetry/score_run.py (mean step over a
fixed window, tokens/s/GPU, MFU):
ISAMBARD_ENV_OVERRIDES=$PWD/configs/quickstart/nemotron_nano_quickstart_pretrain.env \
isambard_sbatch --nodes=16 --time=00:20:00 pipeline_training_submit.sbatch \
configs/quickstart/nemotron_nano_quickstart_pretrain.yaml nano pretrain --disable-ft <overrides>
python scripts/telemetry/score_run.py /projects/a5k/public/logs/megatron_runs/train-<jobid>.out \
--config configs/quickstart/nemotron_nano_quickstart_pretrain.yaml \
--hf-model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
--gpus 64 --window 26 50 --loss-window 41 50The scorer reads the sequence length and model FLOPs/token from the config and the model's HF
config.json through scripts/nemotronh_flops_estimator.py, takes the batch from the log, and
prints its inputs beside the score (--json for machine-readable output, --wandb-peak-memory to
add the W&B summary peaks and allocator-retry count, which are one rank's). The peak memory over all
ranks is the run's own [peak-memory] line, which rank 0 logs when the training loop ends and the
score reports as peak_memory_across_ranks. Calibrate --time to the measured runtime (a 50-iteration 64-GPU Nano
probe takes 9–13 min), and compare a probe only against runs placed on a single switch group (the
log's [run-identity] switch placement line): spanning more than one group costs 2.5–6.6% on its own,
while single-group runs of one posture agree to ~0.5%, so repeat a probe before trusting a smaller
difference than ~1%.
A lever that changes numerics is adopted only if its loss stays within the run-to-run noise of
the production posture; one that claims exactness must reproduce it bit for bit.
scripts/telemetry/loss_parity.py tests both on runs of one config and seed:
# numerics-changing lever: two reference runs of the production posture, 500 iterations each
python scripts/telemetry/loss_parity.py band --reference ref_a.out ref_b.out --candidate lever.out \
--iterations 1 500 --window 50
# exact lever or knob-off code path (every compared run with model.deterministic_mode=true and, in an
# ISAMBARD_ENV_OVERRIDES file, NVTE_ALLOW_NONDETERMINISTIC_ALGO=0, CUBLAS_WORKSPACE_CONFIG=:4096:8, MAMBA_DETERMINISTIC=1)
python scripts/telemetry/loss_parity.py identity --reference ref.out --candidate lever.out \
--iterations 1 30 --wandbThe Nano pretrain campaign's record is docs/investigations/nano30b-pretrain-perf-campaign.md.
Start from a complete stage config, such as the control-pretraining XL SFT that the Nano SFT quickstart is built on
(the quickstart files themselves are base_config: overlays, whose base paths resolve against their own directory):
cp configs/control_pretraining/30b_baseline_ablations/nemotron_nano_30b_baseline_sft_xl50b_gbs256.yaml \
configs/my_new_sft.yamlKey fields to change:
dataset:
dataset_name: your-org/Your-Dataset
dataset_root: /projects/a5k/public/data/your-org__Your-Dataset
packed_sequence_specs:
packed_train_data_path: /projects/a5k/public/data/your-org__Your-Dataset/packed/...
train:
train_iters: ??? # = total packs / global_batch_size, counted from the packs
checkpoint:
pretrained_checkpoint: /projects/a5k/public/checkpoints/megatron/<the model you start from>
load: /projects/a5k/public/checkpoints/megatron/my_new_sft
save: /projects/a5k/public/checkpoints/megatron/my_new_sft
logger:
wandb_exp_name: my_new_sftThe speed levers of the Nano SFT quickstart (configs/quickstart/nemotron_nano_quickstart_sft.yaml, with its .env
launcher setting) apply to a new packed SFT config too; each is explained in that file's header.
A config that differs from an existing one in a few fields can instead name it under
base_config: (relative to the new file's directory) and state only those fields. Mappings
deep-merge and anything else replaces the base value, as OmegaConf.merge would, so a change
to the base reaches the overlay unedited. The Nano pretrain quickstart is written this way. Only
tools that read configs through scripts/training/config_compose.py compose (training and the
blend-coverage dry run built on it, the FLOPs estimator and the run scorer built on it, the config
tests), so a config that the stage gate, the Hub scripts or the epoch-chain generator (as a
parent_config or posture_config) read must stay a complete file (CLAUDE.md, "Config composition").
Slingshot causes intermittent NCCL hangs. Training uses a layered resilience stack:
| Layer | Timeout | Recovery | Iterations lost |
|---|---|---|---|
| ft_launcher worker restart | 7200s step | Per-node agent restarts workers from checkpoint | 0-25 |
| NCCL watchdog | 7200s (TORCH_NCCL_TIMEOUT) |
Aborts genuinely wedged collectives | N/A |
srun --kill-on-bad-exit=1 |
— | A dead rank ends the whole step; a --dependency=singleton chain resumes from the latest checkpoint |
segment tail |
Pass --disable-ft to use plain torchrun instead of ft_launcher.
The governing constraint on this fabric is TP × EP ≤ 4 (and TP × CP ≤ 4): keep expert all-to-all and context-parallel traffic on a node's NVLink, and let only PP cross Slingshot. Cross-node EP costs ~14× throughput and reliably hangs the CXI fabric.
| Model | Validated layout | Measured |
|---|---|---|
| Nano (30B-A3B), seq 8192 | 8 nodes / 32 GPUs: TP=2, EP=2, PP=4, DP=2 (GBS 16) | ~3.4 s/iter, ~27 TFLOP/s/GPU; zero hangs through 500+ iters |
| Super (120B-A12B) | TP=1, CP=(min that fits), EP=4, PP=22, ETP=1 | ~75-84 TFLOP/s/GPU, ~1000+ tok/s/GPU (≈2.4× the old TP=4 layouts) |
| Super benchmark | 16 nodes / 64 GPUs: TP=1, CP=4, EP=4, PP=8, ETP=1, DP=2 (seq 32K, GBS 128 since 2026-08-05) | 31.562 s/iter anchor = 167.4 TFLOP/s/GPU (moe_experts_impl: torch_grouped, optimizer CPU offload off; superseded, at the old GBS-64 workload: 17.099 = the paired A/B that certified torch_grouped, 20.66 on the cublas_grouped per-expert loop, 21.78 with offload 0.5) — the standing environment benchmark, configs/quickstart/nemotron_super_quickstart_sft.yaml |
| Super benchmark, 32 nodes | 32 nodes / 128 GPUs: same topology, DP=4, GBS 256 (scale the batch with the nodes) | 122.0 ms/sample = 31.228 s/iter, 169.2 TFLOP/s/GPU. With the base config at GBS 128 this override is matched µb/replica (64 both ends): perfect per-sample halving predicts 123.3 ms/sample vs 122.0 measured — scaling perfect within the ±2% cross-allocation placement band, same backend both ends — run as the 64-GPU config plus train.global_batch_size=256; the quickstarts are standardised at 64 GPUs and this is the one field that differs |
| Ultra (550B-A55B) | 72 nodes / 288 GPUs: TP=4, EP=4, PP=36, ETP=1 | ~28-30 s/iter steady state; first iter 45-75 min (lazy NCCL init at this depth) |
| Nano pretrain (from scratch) | 16 nodes / 64 GPUs: TP=1, CP=1, EP=4, PP=1, ETP=1, DP=64 (seq 8192, GBS 512 = 8 microbatches per replica, as in the filtered arm's stage 1 at GBS 2048 on 256 GPUs; 50 iterations, no checkpoint I/O), plus the performance campaign's levers (EP all-to-all / compute overlap with the HybridEP dispatcher, FP8 dense layers with BF16 parameters, BF16 gradients, [moe_act] recompute, chunked linear cross-entropy, host-path settings) |
4.954 s/iter (mean of twelve runs in six paired cycles, iterations 26–50) = 13,230 tokens/s/GPU, 27.8% MFU: 1.90× the baseline benchmark on the same allocations, with its 500-iteration loss inside the baseline's band (1.90× at 256 GPUs too, where the loss leaves the band only over iterations 1–50, on the low side). 32 GPUs is the same file plus train.global_batch_size=256 — configs/quickstart/nemotron_nano_quickstart_pretrain.yaml (launched with its .env file) |
| Nano pretrain, baseline benchmark | the same topology and batch at the control-pretraining baseline's stage-1 posture, a small base_config: overlay of it, so the production posture reaches it unedited |
9.328 s/iter (mean over iterations 26–50) = 7,026 tokens/s/GPU, 14.78% MFU at 64 GPUs (job 6930454) — configs/quickstart/nemotron_nano_quickstart_pretrain_baseline.yaml |
| Nano midtrain (warm start, seq 32768) | 16 nodes / 64 GPUs: TP=1, CP=2, EP=4, PP=1, ETP=1, DP=32 (GBS 64 = 2 microbatches per replica, as production's GBS 512 on 512 GPUs; 200 iterations, scored over 151–200), plus the midtraining campaign's levers (parameter-gather overlap, host-path settings, BF16 gradients, chunked linear cross-entropy, HybridEP with router fusion, FP8 dense layers with BF16 parameters, selective [moe, shared_experts] recompute) |
3.745 s/iter on one switch group (the six single-group runs of four paired cycles) = 8,749 tokens/s/GPU, 21.6% MFU: 1.648× the baseline benchmark on the same allocations, with its 500-iteration loss inside the baseline's band — configs/quickstart/nemotron_nano_quickstart_midtrain.yaml (launched with its .env file) |
| Nano midtrain, baseline benchmark | 16 nodes / 64 GPUs: TP=1, CP=2, EP=4, PP=1, ETP=1, DP=32 at the control-pretraining baseline's stage-2 posture (seq 32768, GBS 64 = 2 microbatches per replica, as production's GBS 512 on 512 GPUs; full recompute; warm-started weights-only from the stage-1 final, as production is), a base_config: overlay of it; 200 iterations, no checkpoint saved or resumed |
6.146 s/iter (mean of three runs over iterations 151–200) = 5,332 tokens/s/GPU, 13.17% MFU at 64 GPUs — configs/quickstart/nemotron_nano_quickstart_midtrain_baseline.yaml |
| Nano SFT (warm start, packed seq 32768) | 16 nodes / 64 GPUs: TP=1, CP=1, EP=4, PP=1, ETP=1, DP=64 (GBS 64 = one pack per GPU, the XL SFT's tokens per GPU at GBS 256 on 256 GPUs; 100 iterations, no checkpoint I/O), plus the SFT campaign's levers (chunked linear cross-entropy, BF16 gradient reduction in 500M-parameter buckets with parameter-gather overlap, HybridEP with router fusion on full-length packs, host-path settings) | 3.689 s/iter (the eight runs of four paired cycles, each on one switch group) = 8,881 tokens/s/GPU at a 70.4 GiB peak: 1.779× the baseline benchmark on the same allocations; its 500-iteration loss leaves the baseline's band over iterations 201–400, by at most 4.3×10⁻⁵ nats, and is back inside from 401 — configs/quickstart/nemotron_nano_quickstart_sft.yaml (launched with its .env file) |
| Nano SFT, baseline benchmark | 16 nodes / 64 GPUs: TP=1, CP=2, EP=4, PP=1, ETP=1, DP=32 at the control-pretraining XL SFT's posture (packed seq 32768, GBS 64 = 2 packs per replica, as the XL SFT's GBS 256 on 256 GPUs; full recompute; warm-started weights-only from the midtraining final, as production is), a base_config: overlay of it; 100 iterations, no checkpoint saved or resumed |
6.563 s/iter (the eight as-is runs of four paired cycles, each on one switch group, over iterations 51–100) = 4,993 tokens/s/GPU at 64 GPUs — configs/quickstart/nemotron_nano_quickstart_sft_baseline.yaml |
| Super pretrain (from scratch) | 32 nodes / 128 GPUs: TP=1, CP=1, EP=4, PP=8, ETP=1, DP=16 (seq 8192, GBS 3072, 1B tokens) | 86.940 s/iter = 28.301 ms/sample (loss 12.19 → 7.65, 0 NaN; 225 GB weights-only checkpoint) — configs/quickstart/nemotron_super_quickstart_pretrain.yaml |
Other levers that matter: recompute_granularity: selective with MoE-scoped
recompute_modules (full recompute is ~10× slower; on the 120B it is the ~24 GB between fit
and OOM), moe_permute_fusion: True, expert_tensor_parallel_size: 1 (parallel folding —
what keeps EP node-local at high TP), gradient_accumulation_fusion: True (the image ships
APEX; ~1.1 s/iter on the 120B), and BF16 everywhere — FP8 causes stochastic alignment
crashes in MoE routing. The measured exceptions are the Nano pretrain and midtrain quickstarts
and both stages of the control-pretraining V2 E2E arm, which run FP8 current scaling on their dense layers only; their routed experts stay BF16. Recipe LR 5e-6; 8e-5
NaNs under context parallelism. Full topology
reasoning, per-model memory notes, and the legacy layouts these superseded are in
CLAUDE.md.
# Full pipeline: download + tokenize + export + pack (args forwarded to pipeline_data_prepare.py)
isambard_sbatch pipeline_data_submit.sbatch prepare \
--dataset allenai/Dolci-Instruct-SFT --seq-length 8192
# Pretraining-format corpus (.bin/.idx) — prepare then tokenize (exact token count included)
isambard_sbatch pipeline_data_submit.sbatch tokenize \
/projects/a5k/public/data/<org>__<name> geodesic-research/nemotron-base-tokenizer tokenized_base
# Offline packing only (via SLURM, saves GPU-hours)
isambard_sbatch pipeline_data_submit.sbatch \
/projects/a5k/public/data/allenai__Dolci-Instruct-SFT \
nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 8192 1Packing combines short examples into fixed-length sequences. The training pipeline auto-packs if needed, but it blocks rank 0 for 1-4 hours. Offline packing avoids wasting multi-node GPU time.
Output: <dataset-root>/packed/<tokenizer>_pad_seq_to_mult<N>/training_8192.idx.parquet
Both Nano and Super export on a single node (4 GPUs, EP=4). All EP communication stays on NVLink.
# Export Nano (30B) — 1 node. --hf-model and --reasoning|--no-reasoning are REQUIRED.
isambard_sbatch --nodes=1 pipeline_checkpoint_submit.sbatch export /path/to/ckpts \
--hf-model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 --no-reasoning --iteration 400
# Export Super (120B) SFT checkpoint — 1 node, --not-strict required
isambard_sbatch --nodes=1 pipeline_checkpoint_submit.sbatch export /path/to/ckpts \
--hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 --no-reasoning \
--iteration 490 --not-strict
# Import HF → Megatron (4 nodes for Super)
isambard_sbatch --nodes=4 pipeline_checkpoint_submit.sbatch import nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16
# Upload all iterations + poll for ongoing training
isambard_sbatch --time=24:00:00 pipeline_checkpoint_submit.sbatch upload-all /path/to/ckpts \
--hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 --no-reasoning --poll
# From salloc
bash pipeline_checkpoint_convert.sh export /path/to/ckpts \
--hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 --no-reasoning \
--iteration 300 --push-to-hubpipeline_checkpoint_convert.sh is the launcher (env/NCCL/srun+torchrun). pipeline_checkpoint_convert_hf.py is the Python logic that runs on each GPU rank.
--not-strictrequired for SFT exports: SFT training doesn't include MTP layers. Without this flag, shards containing MTP keys are dropped, which also dropslm_head.weight(fatal for generation).- EP=4 on 1 node (not EP=8 on 2 nodes): Cross-node EP=8 causes Slingshot gathering failures. Node-local EP=4 keeps all communication on NVLink.
- Single-process conversion doesn't work for Super: Hangs during checkpoint loading. Always use
torchrun.
/projects/a5k/public/checkpoints/megatron_bridges/models/
NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/
NVIDIA-Nemotron-3-Super-120B-A12B-BF16/
Qualitative sanity check for HF checkpoints after training or conversion. Generates responses to 8 diverse prompts and logs them to a W&B table for side-by-side comparison across models.
# Via SLURM (4 GPUs for 120B models)
isambard_sbatch pipeline_coherence_submit.sbatch nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16
# Via SLURM (1 GPU for 30B models)
isambard_sbatch --gpus-per-node=1 pipeline_coherence_submit.sbatch geodesic-research/nemotron_nano_sft_warm_start_200k
# Local checkpoint
isambard_sbatch pipeline_coherence_submit.sbatch \
/projects/a5k/public/checkpoints/megatron/my_experiment/iter_0000400/hf
# With custom W&B project
isambard_sbatch pipeline_coherence_submit.sbatch <model> \
--wandb-project megatron_bridge_conversion_coherance_tests
# Megatron checkpoint directly, no HF export (multi-node)
isambard_sbatch --nodes=6 pipeline_coherence_submit.sbatch <megatron-ckpt-dir> \
--backend megatron --hf-model <hf-id> --tp 4 --pp 6 --ep 4 --max-tokens 256
# Against an already-running OpenAI-compatible server (stdlib HTTP, no local model)
isambard_sbatch --gpus-per-node=1 pipeline_coherence_submit.sbatch <served-id> \
--backend endpoint --discovery-file /projects/a5k/public/vllm-serve/<stem>.endpoint
# Directly, inside an allocation (uses this node's GPUs)
./pipeline_env_exec.sh "cd $PWD; source pipeline_env_activate.sh || exit 1; \
python pipeline_coherence_test.py <model_path> [--max-tokens 3000]"Results are logged to W&B project geodesic-gen-tests (default) with a generations table containing prompt, response, response length, and empty flag.
Isambard occasionally has hardware-broken nodes — VS Code tunnels that never come up, GPUs returning ERR!, NCCL dying on first collective. The team shares a TTL'd log that isambard_sbatch automatically passes to SLURM's --exclude on every submission, so once a teammate reports a bad node, nobody else lands on it. A summary line prints on every submission:
Bad nodes: 3 excluded (last 7d) — file: /projects/a5k/public/isambard_sbatch_bad_nodes.log
report more: isambard_sbatch --mark-bad <node> [reason]
If you are very confident a failure is a node-specific hardware issue (not a code, config, or library-version bug), register it so the rest of the team doesn't land on it. The wrapper supports full CRUD on the log:
isambard_sbatch --mark-bad nid001234 "vscode tunnel never came up" # Create
isambard_sbatch --list-bad # Read
isambard_sbatch --update-bad nid001234 "GPU ECC (Xid 48)" # Update reason
isambard_sbatch --unmark-bad nid001234 # Delete (node got fixed)
isambard_sbatch --prune-bad # HousekeepingEntries expire after 7 days, so a node that gets fixed stops being excluded automatically. Don't mark nodes for code-level issues (OOM, bad YAML, wrong parallelism) — that would falsely exclude healthy nodes for a week and erode the list's signal. See CLAUDE.md for the full register-when / don't-register-when checklist.
| Problem | Fix |
|---|---|
RuntimeError: ...gradient_accumulation_fusion... |
Should not happen — the container image ships APEX. It means the payload is not running inside the container (see docs/environment.md) |
FATAL [env-config]: SIF not found / Slingshot NCCL stack not built |
Environment not installed on this cluster: bash pipeline_env_setup.sh (GPU node) |
NCCL bandwidth ~2 GB/s or NET/Socket in the log |
CXI plugin not loading. Never "fix" it with brics/apptainer-multi-node/adapt.sh — see docs/environment.md troubleshooting |
| NaN loss at iteration 7-8 | Lower LR to 5e-6 (recipe default) |
OSError: [Errno 116] Stale file handle |
TRITON_CACHE_DIR/TMPDIR to /tmp (automatic in pipeline_training_launch.sh) |
| NCCL hangs every ~7-8 min | Slingshot fabric issue. ft_launcher auto-restarts |
| EP=4 OOMs on GH200 | Use EP=8 (16 experts/GPU = 51GB vs 32 = 93GB) |
nemo_experiments/ fills disk |
Remove old TB logs selectively. Do NOT rm -rf — contains checkpoint state |
| VS Code tunnel never starts / job sits in RUNNING with no output | Likely a bad compute node. isambard_sbatch --mark-bad <nid> "tunnel hung" and resubmit. See Bad Compute Nodes |
nvidia-smi shows ERR! on specific GPUs of one host |
Node-specific hardware fault. isambard_sbatch --mark-bad <nid> "GPU ECC err" |
| What | Path |
|---|---|
| HF datasets | /projects/a5k/public/data/ |
| Megatron base checkpoints | /projects/a5k/public/checkpoints/megatron_bridges/models/ |
| Training output checkpoints | /projects/a5k/public/checkpoints/megatron/ |
| SLURM training-run logs | /projects/a5k/public/logs/megatron_runs/ (by run ID: .../by-run-id/) |
| W&B logs | /projects/a5k/public/logs/wandb |
| Torch profiles | /projects/a5k/public/profiles/<wandb-exp-name>/<run-id>/ (see docs/profiling-quickstart.md) |
| Performance-campaign records (probe code snapshots, research reports) | one directory per campaign, e.g. /projects/a5k/public/logs/nano_pretrain_perf_campaign/; each probe's log is an ordinary training log in /projects/a5k/public/logs/megatron_runs/ |
| HF cache | /projects/a5k/public/hf |
| Container SIF, Slingshot build, Python overlay | /projects/a5k/public/containers/ (see docs/environment.md) |
This repo includes custom Claude Code skills for interactive development and monitoring:
| Skill | Usage | Description |
|---|---|---|
/wandb-run |
/wandb-run geodesic/megatron_training/<run_id> |
Fetch W&B run status, config, metrics history, and summary. Use to monitor training progress, compare runs, or diagnose failures. |
/megatron-moe-paper |
/megatron-moe-paper [topic] |
Reference for Megatron-Core MoE best practices — parallelism, memory optimization, FP8/FP4, load balancing. Based on NVIDIA's arxiv 2603.07685. |
Skills are defined in .claude/skills/ and invoked as slash commands in Claude Code sessions.
This repo also integrates geodesic-claude-tooling (a git
submodule) — Claude Code hooks that inject Geodesic's working conventions at session start, validate
plans on exit, and run mechanical checks on the diff. It is wired up additively and does not
touch the training environment: the hooks live in the repo's dev-tooling-only .venv
(ruff, pre-commit, the hook entry points — no torch; created by the installer below with
uv pip install, NEVER uv sync, which would resolve the full project and try to build the
training stack on the host),
never in the container. Install it once:
bash scripts/install_claude_tooling.shConfiguration lives in .claude/settings.json (hooks) and .claude/geodesic-config.yaml (quality
items). The commit-time review gate is enabled: git commit runs pre-commit on the staged
files and blocks until the checklist-reviewer subagent writes a passing verdict for the
staged-diff hash; verdict-protection and a submodule-pin consistency check are wired in alongside
it. See CLAUDE.md for the full flow.
- Scalable Training of Mixture-of-Experts Models with Megatron Core — NVIDIA's paper on MoE parallelism, memory optimization, and FP8/FP4 training. Essential background for understanding the parallelism choices in this repo.
- docs/environment.md — The execution environment: install, design decisions, image qualification, troubleshooting
- docs/profiling-quickstart.md — Capturing and reading torch-profiler traces of a training run
- docs/investigations/nano30b-pretrain-perf-campaign.md — The Nano-30B pretraining performance campaign: every probe, its result and what it taught
- CLAUDE.md — Cluster specs, per-model topology findings, campaign conventions, and dev commands
- docs/README_DEFAULT.md — Upstream Megatron Bridge README (supported models, API docs, etc.)