argonne-3.5-base

Argonne 3.5-base

Argonne 3.5-base is a 2.88B-parameter decoder-only transformer trained from scratch. It is the base (foundation) checkpoint of the Argonne 3.5 line and the successor to argonne-3.0-base.

Two things separate it from 3.0-base:

1. A three-stage curriculum — general pretraining, then a reasoning/code/math anneal, then a context-extension stage — for 88.84B total tokens versus 3.0-base's 76.05B.
2. A 13,568-token context, which is trained, not extrapolated. 3.0-base is a 1,024-token model; so was this one until the third stage. That stage was validated against a pre-extension control (numbers below), because on this architecture RoPE θ=1e6 does not extrapolate on its own.

The architecture is unchanged from 3.0 — grouped-query attention with QK-norm, V-norm, sandwich norms, interleaved local/global attention, and a final logit softcap. What changed is the training recipe (FP8, a higher peak LR with a proper cooldown, tighter gradient clipping) and the data curriculum.

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

Looking for the reasoning model? This base was post-trained into Argonne-3.5-think, which scores 65.00 / 73.00 greedy and 74.00 / 82.67 self-consistency on clean SVAMP / ASDiv — vs 22.00 / 32.33 and 30.00 / 39.33 for the previous generation, measured head-to-head in one job.

Model architecture

Training

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

Shared across all three stages:

Why the recipe changed from 3.0

3.0-base ran WSD with cooldown = 0 — the stable phase only, no decay. 3.5 uses a real cooldown in every stage (visible in the LR panel of the figure below), a higher peak LR (6e-4 vs 3e-4), a longer warmup (8,000 vs 1,000 steps), and tighter gradient clipping (0.4 vs 1.0). The tighter clip and QK-norm are what make the higher LR stable.

Training data

The stage-2/3 composite is built by build_reasoning_corpus.py from:

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 stages 2 and 3 use disjoint slices of it (--holdout_frac / --part) so stage 3 is not a second epoch over stage 2's tokens.

The general-web replay tier exists because it is needed: an earlier build capped that tier at 2.2B tokens (9% of the anneal mix), and 18% of the way in, held-out FineWeb-Edu cross-entropy had risen +0.246 nats — the base was measurably regressing on general text while improving on the target tiers. The tier was raised to a 20.6B-token pool.

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

Training loss

Loss, perplexity, and learning rate against cumulative tokens across all three stages, with the stage boundaries marked. The LR panel shows the three cooldowns. Note that the loss step down at each stage boundary is a change of data mixture, not a capability jump — the anneal and context-extension corpora are intrinsically lower-entropy than FineWeb, so cross-stage loss values are not comparable.

Evaluation

Two measurements were run on the final checkpoint. Both are reported with their limitations, because neither is a general capability benchmark.

Context extension — position-bucketed NLL on held-out arXiv

The question is whether stage 3 actually taught the model to use long positions, or just trained it more. The test is a paired A/B on identical held-out windows (40 documents, 24,576 tokens each, proof-pile-2 arXiv — a domain neither stage trained on), comparing the final weights against the stage-2 checkpoint they were seeded from. Lower is better.

Three things this shows:

• RoPE θ=1e6 does not extrapolate unaided on this architecture. The stage-2 model is coherent inside its 1,024-token training window (2.19) and effectively blind past it (~5.9, flat). If you are tempted to assume a large RoPE base buys you free context, this is the counterexample.
• The extension is real. NLL falls monotonically with position rather than shifting by a constant, which is what generic extra training would have produced.
• It generalizes past its own training length. The 13,568–24,576 buckets — beyond anything stage 3 trained on — keep improving, and there is no short-context tax: the 0–1,024 control bucket is better than the stage-2 model (2.161 vs 2.194).

Reproduce with reasoning/exp_longctx_learning.py.

Two-axis base gate

A 35-item greedy few-shot probe (20 arithmetic/word-problem, 15 world-knowledge) used as a go/no-go gate for whether a base is worth running a reasoning recipe on.

Both axes clear the ≥14/20 ∧ ≥14/15 gate, and the context-extension stage cost nothing on either.

Read this as a gate, not as a capability number. The probe is small (n=20/15), it saturates, and it has a measured ±2-item noise floor — which is why three checkpoints are shown rather than one. It says the base is worth building on. It does not say how good it is. Reproduce with reasoning/probe_pretrain_ckpt.py.

What has not been measured

No standard held-out benchmark suite (MMLU, ARC, HellaSwag, GSM8K, …) has been run on this checkpoint yet. Those numbers are not withheld — they do not exist, and this card will be updated when they do. Do not infer benchmark standing from the two measurements above.

One caution specific to this line: GSM8K is contaminated for Argonne reasoning derivatives downstream of this base, and should not be used to grade them.

Source code

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

Inference

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "PursuitOfDataScience/argonne-3.5-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. Unlike the 3.0-base card, this repo ships an auto_map in config.json, so from_pretrained resolves the classes without any manual setup.
• 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. This is expected and benign — embeddings are tied (tie_word_embeddings: true), so lm_head takes its weights from embed_tokens.
• The full 13,568-token context is usable. For throughput at that length, prefer a real serving engine (vLLM/SGLang) over .generate().
• 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.
• Unbenchmarked. See "What has not been measured" above.
• Scale. 2.88B parameters on 88.84B tokens is far below frontier compute; expect corresponding quality.
• The stage-2/3 mixture is reasoning/code/math-heavy. That is deliberate — this base exists to be fine-tuned into a reasoning model — but it means the token distribution is not representative of general web text, and general-domain behavior reflects the replay tier rather than a purely general-purpose pretrain.
• 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{argonne35base,
  author = {PursuitOfDataScience},
  title = {Argonne 3.5-base},
  year = {2026},
  publisher = {Hugging Face},
  url = {https://huggingface.co/PursuitOfDataScience/argonne-3.5-base}
}