Skip to content

About

Training library for Megatron-based models with bidirectional Hugging Face conversion capability

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

1,731 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Megatron Bridge on Isambard

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.

Pipelines

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.

Quickstart Walkthrough

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.


Step 0 — Install the environment (one-time)

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-training

Design decisions, image contents, the image-qualification gates, and troubleshooting: docs/environment.md.


Step 1 — The dataset

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.


Step 2 — Review the training config

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.yaml is the production XL SFT: model, data, tokenizer, warm start, topology (TP1·CP2·EP4·PP1) and schedule.
  • nemotron_nano_quickstart_sft_baseline.yaml makes 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.yaml adds 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.


Step 3 — Submit training

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-ft

For 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=false

The 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.


Step 4 — Monitor training

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.


Step 5 — Export checkpoint to HuggingFace format

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 200

The converted model lands at .../nemotron_nano_quickstart_sft_export/iter_0000200/hf/ — a standard HF checkpoint (safetensors + config + tokenizer) loadable with AutoModelForCausalLM.


Step 6 — Run coherence tests

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/hf

This 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.


1. Environment Pipeline

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 them

Everything 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.


2. Training Pipeline

How it works

  1. A Python recipe defines the base model config, optimizer, parallelism, and data pipeline
  2. A YAML config file overrides recipe defaults
  3. 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

Usage

# 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 lora

To 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.

Performance probes

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 50

The 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 --wandb

The Nano pretrain campaign's record is docs/investigations/nano30b-pretrain-perf-campaign.md.

Writing a new YAML config

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.yaml

Key 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_sft

The 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").

Fault Tolerance

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.

Optimal Parallelism

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.


3. Data Pipeline

Usage

# 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 1

Packing 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


4. Checkpoint Pipeline

Usage

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-hub

pipeline_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.

Key notes

  • --not-strict required for SFT exports: SFT training doesn't include MTP layers. Without this flag, shards containing MTP keys are dropped, which also drops lm_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.

Already-converted checkpoints

/projects/a5k/public/checkpoints/megatron_bridges/models/
    NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/
    NVIDIA-Nemotron-3-Super-120B-A12B-BF16/

5. Coherence Pipeline

Purpose

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.

Usage

# 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.


Bad Compute Nodes

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                                          # Housekeeping

Entries 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.

Common Pitfalls

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"

Disk Locations

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)

Claude Code Skills

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.

Claude Code guardrail tooling

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.sh

Configuration 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.

Further Reading

About

Training library for Megatron-based models with bidirectional Hugging Face conversion capability

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages