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
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_id = "geodesic-research/pa_warm_start_sft_120b_light1bmix_32k"
4tok = AutoTokenizer.from_pretrained(model_id)
5model = AutoModelForCausalLM.from_pretrained(model_id, dtype="bfloat16", device_map="auto")
6
7messages = [{"role": "user", "content": "Solve 3x^2 - 7x + 2 = 0."}]
8inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
9out = model.generate(inputs.to(model.device), max_new_tokens=8192)
10print(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.
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
1@misc{geodesic2026pawarmstart,
2 title = {pa-warm-start-sft-120b-light1bmix-32k},
3 author = {Geodesic Research},
4 year = {2026},
5 url = {https://huggingface.co/geodesic-research/pa_warm_start_sft_120b_light1bmix_32k}
6}