PursuitOfDataScience/argonne-4.0-base

Argonne 4.0-base

Argonne 4.0-base is a 1.04B-parameter decoder-only transformer trained from scratch on 65.12B tokens, with a 65,536-token trained context. It is the base checkpoint of the Argonne 4.0 line.

It is not a bigger argonne-3.5-base — it is a deliberately smaller one. 3.5-base is 2.88B parameters on 88.84B tokens of mostly web text; 4.0-base is 36% of the parameters on 73% of the tokens, and spends those tokens on a math/code-weighted mixture instead. The bet, taken from a 49-run iso-token campaign, is that for a ~1B model data composition is a larger lever than parameter count, and that the advantage grows with scale rather than shrinking.

Read the evaluation section before using this model. The bet paid off in one direction and clearly failed in another, and both are reported below with the numbers.

The architecture is 3.5's, re-shaped: grouped-query attention with QK-norm, V-norm, sandwich norms, and a final logit softcap, at 1,536 hidden × 32 layers instead of 3,072 × 24.

This is a base model: no instruction tuning, no alignment, no safety filtering.

Model architecture

Component	Specification
Parameters	1,038,492,672 (~1.04B)
Layers	32 transformer blocks
Hidden size	1,536
Attention heads	6 query / 2 key-value (GQA)
Head dimension	256 (derived: 1,536 / 6)
Feed-forward	SwiGLU MLP, 4,096 intermediate dim
Attention pattern	Full causal on every layer (see the note below)
Normalization	RMSNorm with QK / V / sandwich norms
Position encoding	RoPE (θ = 1,000,000)
Logit stabilization	Final logit softcap = 15.0
Context length	65,536 tokens (trained, not extrapolated)
Vocabulary size	151,669
Tied embeddings	Yes (input ↔ output)

Attention pattern — read this if you compare against the 3.0/3.5 configs. Those repos ship interleaved_local_attention: true, local_attention_window: 256. This model was trained with full causal attention on every layer, and its config says so (interleaved_local_attention: false). The reason is not a design change: model.py implements the sliding window only on the flash-attn-2 code path, the training environment has flash-attn-4 (which does not expose flash_attn.flash_attn_interface), so every training slice fell back to SDPA and logged local_attention_window=256 is configured but IGNORED on this path. The window has therefore never been active in any Argonne pretrain. Publishing the flags as-is would hand a 256-token window to any downstream user who happens to have flash-attn-2 installed — on weights that never saw one, at an advertised 65,536-token context. The release config describes the model that was actually trained.

Parameter-count footnote: training logs report 1,038,509,568 because FP8 requires the lm_head input dim to be a multiple of 128, so the vocabulary is padded 151,669 → 151,680 during training (11 extra rows × 1,536 = 16,896 parameters). Export trims the padding back off, so the published checkpoint is 1,038,492,672.

Training

Four stages, all causal language modeling, all on 3× NVIDIA H100/H200 GPUs with DDP.

	1 — pretrain	2 — reasoning anneal	3 — ctx extension	4 — ctx extension
Script	pretrain.py	continue_pretrain.py	continue_pretrain.py	continue_pretrain.py
Steps	50 → 72,827	→ 103,457	→ 109,622	→ 112,674
Tokens	38.03B	18.07B	6.02B	3.00B
Cumulative	38.03B	56.10B	62.12B	65.12B
Sequence length	1,024	1,024	13,568	65,536
Batch / GPU	170	32	2	1
Grad accumulation	1	6	12	5
Effective batch	522,240 tok/step	589,824 tok/step	976,896 tok/step	983,040 tok/step
Peak LR	6.0e-4	2.0e-4	1.0e-4	1.0e-4
End LR	6.0e-5	2.0e-4 (no decay)	1.0e-5	1.0e-5
Warmup	8,000 steps	0	0	0
Schedule	WSD, cooldown 10,923 steps (0.15)	constant — no cooldown	WSD, cooldown 900	WSD, cooldown 458

Shared across all four stages:

Item	Value
Optimizer	AdamW (β₁=0.9, β₂=0.95, weight decay 0.1), fused
Gradient clipping	0.4
Precision	FP8 (torchao tensorwise, including lm_head) under bf16 autocast; fp32 optimizer states
Vocab padding	151,669 → 151,680 during training for the FP8 lm_head (trimmed on export)
Cross-entropy	Chunked, so the 151,680-wide fp32 logit transient never materialises in full
torch.compile	Enabled
Gradient checkpointing	Enabled
Data parallel	3 GPUs (DDP)
Total optimizer steps	112,674
Checkpoint dtype on Hub	bfloat16
Weight format on Hub	5 sharded safetensors + index

Stage 2 ran at a constant LR with no cooldown, and that is a defect, not a choice — the launcher passed cooldown 0. It is stated here rather than smoothed over because stage 2 is 18.07B tokens, 28% of the whole run, and because the general-knowledge weakness in the evaluation section is the place it would show up. Stages 1, 3 and 4 all ran a real cooldown.

Training data

The pretrain mixture is the point of this model. It is not a pre-blended corpus: pretrain.py's WeightedMultiLoader samples one source per micro-batch by weight, so the ratio is decoupled from the raw size of each source.

Stage 1 source	Weight	Tokens available
FineWeb-Edu	50%	20.59B
FineMath-4plus	30%	9.95B
GitHub code (permissive)	20%	7.50B
combined		38.03B (≈1× per source)

Stage	Corpus	Tokens
1 — pretrain	50 / 30 / 20 edu / math / code, sampled per micro-batch	38.03B
2 — reasoning anneal	code / math / reasoning / tool mixture with a general-web replay tier	18.07B
3 — ctx extension	a disjoint slice of the same stage-2 composite, read at 13,568	6.02B
4 — ctx extension	50% long arXiv (docs ≥ 32,768 tok) / 25% reasoning replay / 25% edu replay	3.00B

The stage-2/3 composite is built by build_reasoning_corpus.py; stages 2 and 3 take disjoint slices so stage 3 is not a second epoch. Measured composition of that pool (24.33B tokens, by tier):

Tier	Share	Source
code	45.4%	nick007x/github-code-2025 · nvidia/Nemotron-Competitive-Programming-v1
reasoning	24.0%	a-m-team/AM-DeepSeek-R1-Distilled-1.4M · open-r1/Mixture-of-Thoughts · PursuitOfDataScience/0.5M-thinking
math	19.9%	nvidia/OpenMathReasoning
general (replay)	9.0%	HuggingFaceFW/fineweb-edu
tool	1.7%	nvidia/Nemotron-SFT-Agentic-v2

Stage 4 is built by build_phasec_data.py from proof-pile-2 arXiv, restricted to documents long enough to fill the window: 78.4% of that corpus's tokens sit in documents longer than 13,568, which is what makes a longer window worth training. Half the mixture is replay, specifically to limit forgetting; the tier-CE measurement below shows it only partly worked.

Reasoning traces keep their <think>/tool tags, so the base has seen that formatting before any fine-tuning. The corpus is decontaminated against common evaluation sets, and stage 4's evaluation shards (arXiv_09*) are held out from its training shards (arXiv_0[0-8]*).

Tokenizer: Qwen/Qwen3-0.6B-Base (151,669-token vocab), via the Qwen2Tokenizer compatibility class. Bundled with the checkpoint.

Training loss

Loss, perplexity and LR against cumulative tokens across all four stages, stage boundaries marked. Two things about this figure are easy to misread:

• The sawtooth in stage 1 is the sampler, not the optimization. Each step draws one of the three sources, and their entropies differ a lot — an edu step logs ≈2.7, a math or code step ≈1.0–1.5. Consecutive logged steps therefore alternate. The faint trace is raw per-step loss; the solid line is a rolling median.
• Loss steps at stage boundaries are changes of data mixture, not capability jumps. Cross-stage loss values are not comparable; the anneal and long-arXiv corpora are intrinsically lower-entropy than the pretrain mixture.

Evaluation

All numbers below were measured on the exact weights in this repo (step 112,674, post-cooldown). Every benchmark uses lm-evaluation-harness through a vLLM backend that is token-for-token validated against model.py on this architecture (6 query / 2 KV heads, head_dim 256).

Metric rule: acc_norm for the multiple-choice tasks, acc for winogrande and mmlu (which report no normalised variant), and gsm8k separately in both extraction modes. One rule for every column — mixing acc and acc_norm across arms is worth ~3.7 points of phantom regression on this suite. Regenerate with reasoning/release_table.py.

Two anchors, both 1B-class, both scored on the same harness, tasks and few-shot counts in the same campaign — so this is a like-for-like table, not a comparison against published numbers.

task	Argonne 4.0-base	stage 3 (ctx 13,568)	Llama-3.2-1B	Qwen3-0.6B-Base
arc_challenge	36.26	35.75	34.90	44.88
arc_easy	56.19	54.88	59.93	57.87
hellaswag	45.44	43.85	60.36	53.61
piqa	67.74	67.08	73.50	69.80
sciq	79.80	77.80	89.90	91.30
openbookqa	31.80	32.20	36.20	34.60
winogrande (acc)	55.49	56.35	61.96	60.22
mmlu (acc)	26.15	24.73	31.41	52.49
8-task mean	49.86	49.08	56.02	58.10
gsm8k strict-match	7.51	9.70	1.82	49.28
gsm8k flexible-extract	7.88	10.16	2.27	50.04

Parameters / training tokens: Argonne 4.0-base 1.04B / 65.12B · Llama-3.2-1B 1.24B / ~9T · Qwen3-0.6B-Base 0.6B / ~36T. This model has seen 0.7% of Llama's and 0.2% of Qwen's token budget, so a deficit is expected. What matters is its size and shape, and one of the two comparisons is much worse than the other.

Against Llama-3.2-1B the data bet shows up where it was supposed to. 4.0-base wins generative math 7.51 vs 1.82 — 4.1× — at 84% of the parameters, and edges arc_challenge. It loses the rest, most of it on the commonsense/knowledge tasks (hellaswag −14.9, piqa −5.8, sciq −10.1).

Against Qwen3-0.6B-Base it is behind on all nine cells, by −8.24 on the 8-task mean, −26.34 on MMLU and −41.77 on gsm8k, at 1.7× the parameters. Tokenizer is not a confound: 4.0-base pretrains with Qwen3's tokenizer, so the two share it exactly. Qwen3-0.6B's 49.28 gsm8k invites its own contamination question, but MMLU is far harder to game and shows −26.34, so discounting gsm8k entirely does not rescue the comparison.

Read stage 3 → 4.0-base as the cost of the context extension, and it is not free. Stage 4 bought +0.78 on the multiple-choice mean and the entire 65,536-token window, and paid −2.19 on gsm8k, the only generative task in the suite. Two other instruments point the same way: per-tier held-out cross-entropy is worse on 7 of 8 reasoning-anneal tiers after stage 4 (reason_r1 +69% perplexity, code_github +25%), despite stage 4 being 50% replay specifically to prevent that. Stage 4 is a domain trade — scientific text and long-context reach in exchange for generative reasoning — not a free context extension. Full analysis: reasoning/thinking_training.md §39.

For reference, the same suite at step 112,412 — 262 steps and 0.26B tokens earlier, mid-cooldown — gave an 8-task mean of 49.94 and mmlu 25.95. Finishing the LR cooldown moved nothing (−0.08 mean, +0.20 mmlu). That is worth stating because §39 flagged its mid-cooldown numbers as a lower bound that "may soften" the trade; measured on the finished stage, it did not.

Long context — position-bucketed NLL on held-out arXiv

The 65,536-token window is the claim that most needs evidence, because on this architecture RoPE θ=1e6 does not extrapolate unaided. The test is a paired A/B on identical held-out windows against this model's own ancestors — the stage-2 checkpoint, which was trained only at 1,024 tokens, and the stage-3 checkpoint at 13,568. Lower is better.

All three arms scored on the same 24 held-out windows of 49,152 tokens each, from proof-pile-2 arXiv shards (arXiv_09*) that stage 4 did not train on (it trains on arXiv_0[0-8]*). Nats/token, lower is better.

Token position	stage 2 (ctx 1,024)	stage 3 (ctx 13,568)	Argonne 4.0-base (ctx 65,536)
0 – 1,024	2.561	2.401	1.970
1,024 – 2,048	5.086	2.061	1.671
2,048 – 4,096	5.933	1.705	1.386
4,096 – 8,192	6.051	1.427	1.139
8,192 – 13,568	5.964	1.242	0.980
13,568 – 20,480	5.920	1.197	0.924
20,480 – 24,576	6.025	1.220	0.905
24,576 – 32,768	5.990	1.198	0.864
32,768 – 40,960	6.171	1.316	0.896
40,960 – 49,152	6.075	1.300	0.820

RoPE θ=1e6 does not extrapolate unaided on this architecture. Stage 2 is coherent inside its 1,024-token training window (2.56) and effectively blind past it — 5.1 to 6.2 nats, flat, for 48,000 tokens. If you are tempted to assume a large RoPE base buys free context, this is the counterexample, measured on this exact model's own ancestor.

The window is real. The released model falls monotonically with position out to 49,152 and is better than stage 3 at every bucket, with no short-context tax: the 0–1,024 control bucket improves too (1.970 vs 2.401).

But do not read the whole gain as context extension, and the probe's own falsifiable test is why. A genuine extension predicts the gap over the previous stage grows with position. It does not — the stage 3 → 4.0-base gap is U-shaped: −0.43 nats at 0–1,024, a minimum of −0.26 in the middle, and −0.48 in the 40,960–49,152 tail. It is as large at position 0 as at position 49,000, and stage 4 did not need to extend position 0–1,024. Combined with stage 4 making the reasoning-anneal tiers worse (§39b), the honest attribution is distribution, not length: stage 4 moved the model toward arXiv. The context extension proper is carried by stage 3, which is exactly where the 1,024-only arm's ~6-nat plateau collapses to 1.2–2.1.

At the full 65,536 training length the same shape holds. Re-run on 10 windows of 65,536 tokens (the window set is seeded, so all three arms see the same documents), the tail bucket 49,152 – 65,536 reads: stage 2 5.215 · stage 3 1.097 · Argonne 4.0-base 0.774. The advertised window is usable end to end, not just nominally configured.

Reproduce with reasoning/exp_longctx_learning.py (--arms a4_anneal,a4_phaseb,a4_phasec --docs 24 --eval_len 49152 --docbin_glob 'arXiv_09*.bin'). Caveat worth stating plainly: the evaluation corpus is stage 4's own training domain, held-out shards notwithstanding, so its win here bounds usable context but should not be read as general capability.

Two-axis base gate — reported as a footnote, deliberately

A 35-item greedy few-shot probe (20 arithmetic/word-problem, 15 world-knowledge) used internally as a go/no-go gate on whether a base is worth spending a reasoning recipe on. Measured on the released weights, in two prompt formats:

	standard	extension	pooled	gate
Math	17/20	17/20	34/40	PASS (≥14/20)
General	14/15	14/15	28/30	PASS (≥14/15)

It clears — and it is the wrong instrument, which is the point of putting it here rather than at the top. The general axis has a ceiling of 15 and every a4 checkpoint ever tested reads 14 or 15 on it, so it cannot separate them; two phase-C checkpoints 40 steps apart differ by 2 pooled math points, which is this probe's measured ±2-item noise floor. On this gate the released model reads ahead of the 2.88B base that produced Argonne-3.5-think (14/20 · 14/15) at 36% of the parameters — while the benchmark table above shows a 2× MMLU gap and a 6.6× gsm8k gap against a 0.6B public base. A saturating gate cannot rank bases; it can only reject very bad ones. Reproduce with reasoning/a4_gate_probe.py.

What is weak, stated plainly

World knowledge is the binding axis, and MMLU is at chance. 26.15 against 25.0 for random guessing on 4-way multiple choice. Whatever this model knows, MMLU cannot detect it. Qwen3-0.6B-Base reads 52.49 on the identical harness at 58% of the parameters, so this is not a "1B models can't do MMLU" result — it is specific to this base.

Two contributors are identifiable in the recipe, and both are recorded rather than guessed at:

1. The general-web replay tier was 9.0% of the stage-2 pool — 2.20B of 24.33B tokens, against 45.4% code and 24.0% reasoning. Stage 2 is 18.07B tokens, 28% of the whole run, and it is overwhelmingly not general text. The 3.5 line hit measurable general-text regression on a similarly-capped tier (+0.246 nats of held-out FineWeb-Edu cross-entropy, 18% of the way into the anneal) and raised the pool for that reason; this run did not use the raised mixture.
2. Stage 2 ran with no LR cooldown (flat 2.0e-4 for all 18.07B tokens) because the launcher passed cooldown 0. An un-annealed stage is exactly where knowledge consolidation would be expected to suffer.

Neither is proven causal — a controlled run would need a raised replay tier and a cooled stage 2, and that experiment has not been run. They are named because they are the two levers this recipe left on the table, not as an excuse for the number.

The internal base gate says the opposite, and the gate is wrong. It rates this model 17/20 math · 14/15 general — clearing with margin and ahead of the 2.88B base that produced Argonne-3.5-think (14/20 · 14/15) at 36% of the parameters — while blind to a 2× MMLU gap and a 6.6× gsm8k gap against a 0.6B public base. Any "this base looks strong" claim about the 4.0 line traces to that instrument and should be re-read against the benchmark table. This is the methodological result of the release round, and it is why the gate is reported above as a footnote rather than as a headline.

What this model is genuinely good for, stated so the negatives above are not over-read: it is a real 65,536-token model with measured long-document gains that no 0.6B checkpoint here has; it beats a 1.24B Llama base on generative math by 4×; it is trained end-to-end from scratch with a fully published recipe; and it is a working substrate for reasoning post-training — the 4.0 line's reasoning derivatives are built on it. It is not the strongest 1B-class base available, and this card does not claim otherwise.

Source code

Built from the GitHub main branch: https://github.com/PursuitOfDataScience/ArgonneAI/tree/main

File	Role
model.py	ArgonneModel / ArgonneConfig architecture + KV cache (bundled here as model.py)
pretrain.py	stage 1 — DDP pretraining loop and the WeightedMultiLoader mixture sampler
continue_pretrain.py	stages 2–4 — anneal and the two context extensions
build_a4_data.py	builds the stage-1 per-source token bins
build_reasoning_corpus.py	builds the stage-2/3 corpus (tiering, decontamination, disjoint slicing)
build_phasec_data.py	builds the stage-4 long-document + replay mixture
reasoning/run_lmeval_vllm.py	the benchmark harness used above
reasoning/vllm_argonne.py	vLLM port of this architecture (validated token-for-token)
reasoning/exp_longctx_learning.py	the position-bucketed long-context NLL probe
reasoning/tier_ce_probe.py	per-tier held-out cross-entropy (the forgetting check)
reasoning/plot_a4_loss.py	the loss figure above
ARGONNE4.0.md	the recipe, ops and honest bounds for this line
reasoning/thinking_training.md	the full lab notebook; §39 is this base's evaluation round

Inference

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "PursuitOfDataScience/argonne-4.0-base"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    dtype=torch.bfloat16,
)

prompt = "Write a short paragraph about scientific computing at Argonne National Laboratory."
inputs = tokenizer(prompt, return_tensors="pt")
input_ids = inputs["input_ids"].to(model.device)

output_ids = model.generate(
    input_ids,
    max_length=input_ids.shape[1] + 128,
    temperature=0.8,
    top_p=0.95,
    top_k=50,
    do_sample=True,
)
print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

Usage notes

• Load with trust_remote_code=True so the custom ArgonneModel / ArgonneConfig classes (model.py) are registered. config.json ships an auto_map, so from_pretrained resolves them without manual registration.
• The custom generate method on ArgonneModel takes max_length (total sequence length), not max_new_tokens.
• Weights are 5 bf16 safetensor shards with a model.safetensors.index.json weight map.
• lm_head.weight is reported missing on load. Expected and benign — embeddings are tied (tie_word_embeddings: true), so lm_head reads from embed_tokens.
• The full 65,536-token context is usable, but .generate() is not the way to use it. For anything at that length prefer a real serving engine; a validated vLLM port of this exact architecture is in the repo (reasoning/vllm_argonne.py), and it is ~10–50× faster than the HF decode loop on this model.
• Switch to greedy decoding (do_sample=False) for deterministic output.

Limitations

• Base model. No instruction following, dialogue ability, or safety alignment. Outputs can be factually wrong, biased, or unsafe. It continues text; it does not answer questions.
• World knowledge is this model's weakest axis by a wide margin — see the evaluation section. Do not use it where factual recall matters.
• Scale. 1.04B parameters on 65.12B tokens is far below frontier compute, and roughly 0.2% of the token budget of the 1B-class open models it is compared against above. Expect corresponding quality.
• The token distribution is not general web text. Stages 2–4 are reasoning/code/math/arXiv-heavy by design — this base exists to be fine-tuned into a reasoner — so general-domain behavior reflects a 9%-of-mixture replay tier rather than a general-purpose pretrain.
• Stage 2 ran without an LR cooldown (18.07B tokens at a flat 2.0e-4). Treat that stage's artifact as un-annealed.
• GSM8K carries a contamination asterisk for this line. The stage-2 mixture draws on OpenMathReasoning and GSM8K-style data, and train-vs-test exposure was not audited for this base. The GSM8K number above is reported because it is the only generative task in the suite and it moves where the multiple-choice tasks do not — not as a clean capability claim.
• Trained with FP8 matmuls. Weights are published in bf16 and load normally, but exact reproduction of the training run requires the same torchao FP8 path.

Citation

@misc{argonne40base,
  author = {PursuitOfDataScience},
  title = {Argonne 4.0-base},
  year = {2026},
  publisher = {Hugging Face},
  url = {https://huggingface.co/PursuitOfDataScience/argonne-4.0-base}
}