geodesic-research/pa_warm_start_sft_120b_light1bmix_32k

pa-warm-start-sft-120b-light1bmix-32k

A reasoning-capable supervised fine-tune of NVIDIA Nemotron-3-Super-120B-A12B (Base), trained on 1B tokens of geodesic-research/pa-warm-start-sft-light-1b-mix at a 32,768-token sequence length. Produced by Geodesic Research as the warm-start model for the persistent alignment green-team line of work.

The model emits explicit chain-of-thought inside <think> … </think> before its answer, and supports multi-turn conversation with prior-turn reasoning preserved in context.

Quick start

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "geodesic-research/pa_warm_start_sft_120b_light1bmix_32k"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="bfloat16", device_map="auto")

messages = [{"role": "user", "content": "Solve 3x^2 - 7x + 2 = 0."}]
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
out = model.generate(inputs.to(model.device), max_new_tokens=8192)
print(tok.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True))

The chat template opens the <think> block itself as part of the generation prompt, so a well-formed response contains a closing </think> and no opening tag. Budget generously for max_new_tokens: reasoning completions on hard maths run to a median of ~7,600 tokens, and a tight budget truncates the answer rather than shortening the reasoning (see Capability evaluations).

Model details

Base model	nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-Base-BF16
Architecture	Nemotron-H hybrid — Mamba2 SSM + attention + latent MoE, 88 layers, hidden 4096
Parameters	120B total / 12B active per token; 512 routed experts, 22 active
Precision	BF16
Context length	262,144 supported by the architecture; fine-tuned at 32,768, which is the length these weights were actually adapted at
Tokenizer	geodesic-research/nemotron-think-history-tokenizer
Training framework	NeMo Megatron Bridge + Megatron-Core

The base checkpoint used is the Base model with chat special-token embeddings grafted in. NVIDIA ships *-Base-BF16 with ~1,188 embedding rows at exactly zero — the chat-template scaffolding tokens the Base run never trained. Fine-tuning on chat-formatted data without that graft drives a hard Inf in the first backward pass.

A note on the shard files. There are 50 model-000NN-of-00050.safetensors files but model.safetensors.index.json references only 49. This is expected, not a truncated upload: model-00050-of-00050.safetensors is a 16-byte empty shard holding zero tensors. Supervised fine-tuning does not train the base model's multi-token-prediction (MTP) layers, so the shard that would have held them was written empty. All 41,643 parameter tensors live in the 49 indexed shards, and loaders read the index, so the empty file is inert.

Tokenizer

The tokenizer is a fork of geodesic-research/nemotron-think-tokenizer with the same encoder (verified byte-identical tokenizer.json) and three chat-template default changes:

• truncate_history_thinking defaults to False. The parent template strips reasoning from every assistant message before the last user turn, replacing it with an empty <think></think> stub. With genuine multi-turn dialogue in the mix, that would have destroyed the reasoning of every non-final assistant turn at pack time. This one is load-bearing rather than defensive: the packing path calls apply_chat_template with no template kwargs, so the template's own default is the only control — passing truncate_history_thinking=False at pack time is not possible.
• No <think></think> stub is emitted for reasoning-less assistant content. The parent injects an empty stub; the fork renders such content bare.
• No <think></think> stub is emitted for a tool-calling turn with empty content. A tool call is the output, so those turns carry no content. This edit does not fire on this model's training mix, which contains no tool use, but it is part of the fork and matters for any mix that does.

Verified on the training corpus: 113 of 113 multi-assistant-turn documents retain reasoning on their non-final turns, 0 truncated, and 0 empty stubs across ~8,300 sampled documents.

Training data

geodesic-research/pa-warm-start-sft-light-1b-mix, pinned at revision d691d216a0cc82160bc58daaccddbf8715553e9d.

Documents	634,571
Tokens	1,000,013,912
Longest document	32,592 tokens (under the 32,768 cap)
Packing	32,768-token packs, pad_seq_to_mult 16, packing factor 20.66, efficiency 99.80%
Loss-mask density	74.0%

Token counts are exact, measured under the training tokenizer at the pinned revision, and match the publisher's own count digit for digit.

Composition

The mix is reasoning-only, tool-free and safety-free, budgeted per source and sampled shortest-chain-of-thought first — the intent is a bare-bones initialisation that teaches the <think></think> convention and as little else as possible, leaving downstream RL a blank canvas for shaping reasoning. Every document is under 32,768 tokens, so nothing truncates at the training length. Every assistant message carries a reasoning trace; no document involves tool use.

Source	Documents	Tokens
nvidia/Nemotron-SFT-Math-v4::train	115,180	250,002,291
nvidia/Nemotron-SFT-SWE-v2::agentless	56,051	250,002,475
nvidia/Nemotron-SFT-Science-v2::so	222,529	100,000,401
nvidia/Nemotron-SFT-Science-v2::rqa	55,780	100,000,648
nvidia/Nemotron-SFT-Science-v2::syn_mcq	79,101	50,000,094
nvidia/Nemotron-SFT-Instruction-Following-Chat-v3::instruction_following	56,812	125,002,240
nvidia/Nemotron-SFT-Competitive-Programming-v2::python_00	29,524	62,501,516
nvidia/Nemotron-SFT-Competitive-Programming-v2::python_01	19,594	62,504,247

Roughly one eighth of tokens are genuine multi-turn dialogue: the chat source is filtered to records with ≥2 content-bearing user turns, and every one of its records is a dialogue. Across the whole mix, 56,825 records (9.0%) carry more than one assistant turn — which is why loss is applied to all of them rather than the final reply alone.

Data curation

Three curation decisions shape this checkpoint's behaviour and are worth knowing:

• No model or vendor identity. A document is dropped outright if any of 24 AI model/organisation names (including NVIDIA and Nemotron, the upstream vendor) appears anywhere in the rendered document, reasoning included — so the model does not inherit a borrowed identity.
• Evaluation-awareness screening. Reasoning that discusses being tested, graded or watched is excluded, since training on it teaches a model to condition behaviour on whether it believes it is being evaluated. A recall-oriented regex router (99.2% / 100% recall on hand-labelled fixtures) routed 4.76% of the candidate pool to a 120B LLM judge; 8,812 documents were confirmed and dropped. The router never drops on its own — a miss approves, a hit only escalates.
• System prompts are sampled, not fixed. Upstream, 94% of documents had no system prompt and the one source that did carried a single byte-identical instruction. 389,634 records (61.4%) receive one of four deliberately bare variants; 158,487 (24.98%) deliberately have no system prompt at all, rendering an empty system block. Sources carrying real standing instructions keep them, relocated into the system slot.

Full provenance, per-source VEA verdicts, and the excluded documents with the judge's justification are on the dataset card.

Loss masking. Loss is applied to every assistant turn, not only the final one — the {% generation %} markers in the chat template wrap every assistant branch, and the mask is derived from them via return_assistant_tokens_mask. With multi-turn dialogue in the mix, final-turn-only masking would have discarded the supervision on all earlier turns.

Training procedure

One epoch over the corpus: 238 iterations × 128 sequences × 32,768 tokens = 998,244,352 tokens, or 99.82% of the corpus in a single pass.

Hyperparameters

Objective	Supervised fine-tuning, answer-only loss on assistant turns
Sequence length	32,768
Global batch size	128 (4,194,304 tokens/iteration)
Micro-batch size	1
Iterations	238 (1 epoch)
Optimizer	Adam, β = (0.9, 0.95), ε = 1e-8
Learning rate	1e-5, constant
LR schedule	Linear warmup from 0 over the first 10% of training (~24 iterations ≈ 100M tokens), then held flat at 1e-5 for the remainder
Weight decay	0.1
Gradient clipping	1.0
Precision	BF16, precision-aware optimizer with BF16 Adam moments
Distributed optimizer	Yes; overlap_grad_reduce on, overlap_param_gather off

The schedule is deliberately constant rather than decayed: the run is a warm start for downstream work, so the final weights should not sit at the bottom of a decay ramp.

Parallelism and hardware

Trained on 64 × NVIDIA GH200 120GB (16 nodes × 4 GPUs) on the Isambard-AI cluster.

Tensor parallel	1
Context parallel	4
Expert parallel	4
Expert tensor parallel	1 (parallel folding)
Pipeline parallel	8
Recompute	selective — moe, shared_experts
MoE dispatcher	alltoall, te_grouped experts

Expert and context parallelism are kept node-local (TP × CP ≤ 4, EP ≤ 4) so all-to-all traffic stays on NVLink; only pipeline-parallel point-to-point crosses the Slingshot fabric. Mamba2 inter-chunk scan state is accumulated in FP32, which is required for numerical stability on long single-document sequences at this length.

Training run

Loss	0.683 → 0.426
NaN iterations	0 (all 238)
Learning rate	reached exactly 1.000000E-05 at iteration 24 and held flat
Throughput	~41 s/iteration, ~128 TFLOP/s/GPU
Wall clock	~2.8 hours

W&B: geodesic/megatron_training/d3sze68t (the full resolved training config is recorded in that run's config tab).

Evaluations

Coherence

50-prompt topic-diverse generation harness, run at temperature=1.0 with an 8,192-token budget. A 10-prompt topic-spread sample was inspected transcript by transcript:

• Well-formed reasoning in 10/10 — exactly one </think> per generation, no stray opening tags, no empty stubs.
• Correct EOS termination in 10/10 — longest generation 6,039 tokens against the 8,192 cap; every generation ends on a complete sentence.
• No template leakage — zero <|im_start|>, <|im_end|>, <|endoftext|> or </s> tokens in any output. No repetition loops.

W&B: geodesic/megatron_bridge_conversion_coherance_tests/iphh7nnd

Capability evaluations

Run by pa-green-team-evals with geodesic-evals, bundled runner, config configs/suites/pa-warm-start-light1bmix-capability.yaml (+ configs/suites/pa-warm-start-light1bmix-capability-longcot.yaml for the raised-budget re-measurement). Served with vLLM 0.18.1, tensor_parallel_size: 4, BF16, on one Isambard-AI GH200 node. Checkpoint evaluated: iter_0000238, local export, prior to Hub upload. Full resolved config for every run is recorded in its W&B config.

Results

Task	Metric	Score	n	Generation budget
cute	accuracy	96.8% ± 0.3	3,000	4,096
gsm8k	accuracy	95.3% ± 0.6	1,319	4,096
piqa	accuracy	94.7% ± 0.7	1,000	4,096
ifeval	final accuracy	87.9% ± 1.5	541	4,096
aime2025	accuracy	80.0% ± 7.4	30	32,768
popqa	accuracy	36.5% ± 1.5	1,000	4,096
mmlu_pro	accuracy	81.6% ± 0.4	12,032	32,768
gpqa	accuracy	75.0% ± 2.5	198 × 4 epochs	32,768

ifeval's headline is the strict-prompt-and-instruction combined metric; its sub-metrics were strict-instruction 89.0%, loose-instruction 91.2%, strict-prompt 84.1%, loose-prompt 87.4%.

Caveats

Generation budget is the single biggest thing to understand about these numbers. An earlier sweep ran at 4,096 tokens, and on tasks that reason at length the score reported where generation was cut off rather than what the model can do. All three long-reasoning tasks were re-measured at 32,768 and all three moved:

Task	at 4,096	at 32,768	Truncated at 4,096
aime2025	33.3%	80.0%	63.3%
gpqa	52.7%	75.0%	41.2%
mmlu_pro	76.6%	81.6%	9.9%

The model did not change between those runs; only the budget did. mmlu_pro is the instructive case: its median completion barely moved (819 → 834 tokens) and only ~10% of samples were ever truncated, but those were scoring 140/1,190, and that tail alone cost 5 points of headline accuracy.

Accuracy among naturally-stopped completions is an upper-biased estimate, because the truncated samples are disproportionately the harder ones — that is why they needed more tokens. gpqa's raised-budget final of 75.0% sits below the 85.6% naturally-stopped figure measured at 4,096 for exactly that reason. Both earlier projections were optimistic in the way the bias predicts (gpqa projected 85.6%, actual 75.0%; mmlu_pro projected 83.7%, actual 81.6%).

The five short tasks are unaffected. cute, gsm8k, piqa, ifeval and popqa all ran at 4,096 with truncation rates of 0.0–2.4% (piqa 0.0%, gsm8k 0.8%, popqa 1.4%, cute 1.9%, ifeval 2.4%), so their scores are not budget-limited. Note that no task ran at its own declared per-task cap: the serving path resolves the budget from an environment variable defaulting to 4,096, and a task's declared GenerateConfig value is inert on it.

Sample sizes. aime2025 is 30 problems and its ±7.4 standard error is wide — treat it as "strong" rather than as a precise figure, and do not compare small differences on it across checkpoints. mmlu_pro (12,032) and cute (3,000) are tight. gpqa is 198 questions run at 4 epochs.

popqa at 36.5% is expected, not a defect. It is short-form open-domain factual recall with no retrieval, scored by exact match against a single gold answer, so a low absolute score is normal for this setup and is not evidence of a knowledge regression. The reference base model scores 31.1% on the same task under identical conditions.

hle is excluded from the suite: vLLM returns 400 on its >12,288-token prompts, which aborts the whole eval rather than skipping the sample.

W&B runs

4,096-budget sweep (project geodesic/geodesic-evals): cute ykpieex5 · gsm8k adgu818j · piqa jgnivs7x · ifeval edjvq8zt · popqa xlv4rayb · mmlu_pro 9yted1op · gpqa ywd84us7 · aime2025 5lvyzwxf

32,768-budget re-measurement (the quoted figures for these three): aime2025 wyg9xo0z · gpqa tsv1523t · mmlu_pro ydk9ph75

Comparison against the base model

The same suite was run against the reference base model, nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16, under identical conditions — same tasks, same sample limits, same 32,768-token generation budget on the three long-reasoning tasks, same chat template, same tensor_parallel_size: 4 BF16 serving on one GH200 node. Only the weights differ. Both runs completed every task with no failures.

Task	Base	This model	Δ	Resolvable?
cute	90.4% ± 0.5	96.8% ± 0.3	+6.4	yes
popqa	31.1% ± 1.5	36.5% ± 1.5	+5.4	yes
gpqa	71.1% ± 2.8	75.0% ± 2.5	+3.9	see note
ifeval	86.3% ± 1.6	87.9% ± 1.5	+1.6	no
piqa	94.6% ± 0.7	94.7% ± 0.7	+0.1	no (neutral)
gsm8k	95.5% ± 0.6	95.3% ± 0.6	−0.2	no (neutral)
mmlu_pro	82.8% ± 0.3	81.6% ± 0.4	−1.2	yes
aime2025	86.7% ± 6.3	80.0% ± 7.4	−6.7	no (n=30)

"Resolvable" means the difference is larger than the combined standard error of the two measurements. Where it is not, the two models should be read as indistinguishable on that task regardless of which number is larger.

What the comparison supports

The warm start improved exact-output instruction-following. cute (+6.4) and popqa (+5.4) are the two clearly resolvable gains, and both are exact-match tasks where a large share of failures are formatting and verbosity rather than missing knowledge or reasoning. The most defensible reading is better answer shaping — not that a 1B-token warm start taught the model new facts or new reasoning.

General capability is unchanged. gsm8k, piqa and ifeval all sit within noise of the base. There is no broad capability regression.

One small regression is real: mmlu_pro −1.2. With n=12,032 the intervals are tight enough that this is a genuine, if small, loss.

What it does not support, and why

We do not claim to outperform the base model on gpqa, despite the +3.9. At the same 32,768-token budget the base still truncated on 18.9% of gpqa samples against this model's 1.8%, with its 90th-percentile completion pinned at the budget ceiling. Its median gpqa completion is 5,427 tokens against this model's 3,303 — the warm start reasons roughly 40% shorter. So the gpqa gap substantially measures which model fits inside 32,768 tokens, not which reasons better. Among completions that actually finished, the base is ahead (87.4% vs 76.2%) — though that statistic is upper-biased for both models and more so for the base, since a larger share of its hardest problems are excluded from it.

The hard-reasoning question is open. aime2025 points the other way (base +6.7, with truncation equal at 13.3% each) but n=30 cannot resolve it. Taken together with gpqa, the consistent hypothesis is that the base reasons longer and is at least as strong when allowed to finish, and that the warm start traded some peak accuracy on the hardest problems for markedly better token efficiency. That is what one would expect from a mix selected shortest-chain-of-thought-first. Settling it requires a budget at which neither model truncates; the base wants more than 32,768.

Base-model W&B runs

cute pbbh75jt · popqa 722ms6xf · gpqa zt4krvvz · ifeval mu7pglv8 · piqa wsw0oq62 · gsm8k riqaw2ku · mmlu_pro 4v7793p6 · aime2025 xtcqwdgn

Alignment evaluations

Not included. The alignment suite was started and is not reportable: 7 of its 8 tasks judge with an external model API for which no credential was available in the eval environment, so three tasks failed at construction and the rest cannot be scored. One task that did report a score produced accuracy 1.000 with zero variance and no scoring components recorded — a scoring path that never ran, not a perfect result. No alignment number from this run should be quoted. The generated rollouts are retained and can be scored later without regenerating, via the repo's GPU-free rescore runner, once a judge credential exists.

Intended use and limitations

This is a research artifact produced for alignment research, not a general-purpose assistant release. It is a one-epoch supervised fine-tune with no preference optimisation, no safety tuning, and no red-teaming beyond the coherence and capability checks above.

• No alignment or safety evaluation results are available for this checkpoint. The alignment suite could not be scored (see the capability section). Do not infer safety properties from the capability numbers.
• Reasoning traces are model output, not a faithful account of its computation, and should not be treated as an explanation of the answer.
• Short-form factual recall without retrieval is weak (see popqa); the model will state confident specifics that are wrong. Two examples observed directly in coherence transcripts: a fabricated date in a generated code docstring, and a confidently-quoted exchange rate.
• English only; other languages are untested.
• Inherits the licence, capabilities and limitations of the NVIDIA Nemotron-3 base model.

Citation

@misc{geodesic2026pawarmstart,
  title  = {pa-warm-start-sft-120b-light1bmix-32k},
  author = {Geodesic Research},
  year   = {2026},
  url    = {https://huggingface.co/geodesic-research/pa_warm_start_sft_120b_light1bmix_32k}
}