Argonne 4.5-base

A 2.06B-parameter base model trained from scratch on 128.07B tokens, with a 1,024-token context. It is the Argonne 4.5 base before long-context extension.

Most people want argonne-4.5-base-ctx13568 instead. It is the same training run, continued for one more stage (17.15B tokens at a 13,568-token context). This checkpoint is published for two reasons: it is the exact starting point of that last stage, so the effect of context extension can be measured against it, and it suits short-context work where the extra stage is not wanted.

This is a base model: no instruction tuning, no alignment, no safety filtering. It continues text; it does not follow instructions.

Which checkpoint should I use?

• Knowledge and commonsense tasks, inputs longer than 1,024 tokens: argonne-4.5-base-ctx13568
• Math and step-by-step reasoning at short context, a start for math fine-tuning, or a before/after control for long-context research: argonne-4.5-base (this repo)
• A smaller model, or a trained context longer than 13,568 tokens: argonne-4.0-base (1.04B, 65,536-token context)

The context stage is a trade, measured: it raises seven of the eight multiple-choice tasks (+1.36 on the 8-task mean; mmlu fell 28.88 to 26.03) and cuts gsm8k from 20.62 to 4.78 (−15.85). It trained on long arXiv plus FineWeb-Edu with no math or reasoning replay, so the anneal's step-by-step math is what it gave up. For math at short context, this checkpoint is the better start.

Both 4.5 checkpoints come from one run:

stage 1  pretrain             110.00B tokens   context 1,024
stage 2  reasoning anneal      18.07B tokens   context 1,024    ->  argonne-4.5-base  (this repo)
stage 3  context extension     17.15B tokens   context 13,568   ->  argonne-4.5-base-ctx13568

What changed from Argonne 4.0-base

The block design is 4.0's. What changed is the size, the token budget, and one fixed defect.

Model architecture

• Attention pattern: full causal attention on every layer
• Normalization: RMSNorm (eps 1e-6), plus QK-norm, V-norm and sandwich norms
• Position encoding: RoPE, θ = 1,000,000
• Output logits: soft-capped at 15.0 (15 · tanh(x / 15))
• Tied embeddings: yes: the output head reuses the input embedding
• KV cache: 48 KiB per token in bf16, so 48 MiB for a full 1,024-token window

Config switches that are off. config.json carries switches from the 4.5 research phase, all disabled in this model: attn_pattern: null (so sliding_window_size: 2048 is never used), interleaved_local_attention: false, nope_global: false, attn_gate: false, doc_mask: false, mtp_module_layers: 0, mtp_loss_weight: 0.0, z_loss_weight: 0.0. Every layer runs full causal attention with RoPE.

Training

Two stages, both causal language modeling.

Shared by both stages:

• Optimizer: AdamW (β₁ 0.9, β₂ 0.95, weight decay 0.1), fp32 master weights
• Gradient clipping: 0.4
• Precision: bf16 autocast. FP8 matmuls (torchao, tensorwise, including the output head) for the first ~55.8k pretraining steps (~30.1B tokens); plain bf16 after that (the A100s have no FP8 support, and it stayed off for the one later H100 stretch)
• Vocabulary padding: 151,669 → 151,680 rows during training (FP8 alignment), trimmed on export
• Memory: AdamW state sharded across GPUs on the 40 GB cards; gradient checkpointing
• Published weights: bfloat16 safetensors, about 4.1 GB

Hardware

• Pretrain, steps 0 to ~55.8k: 3× NVIDIA H100 (80 GB and 94 GB variants), FP8 on
• Pretrain, remaining ~147.7k steps: NVIDIA A100 40GB, 4 to 24 GPUs per job, plus one stretch of ~5.6k steps on 4× NVIDIA H100 80GB; FP8 off
• Reasoning anneal: NVIDIA A100 40GB, 4 or 8 GPUs

The batch stayed at exactly 540,672 tokens on every GPU count, so moving between GPU types did not change the recipe.

Known training quirks

• A sliding window slipped into the first ~950 anneal steps. For global steps 203,452 to 204,397 (about 2.8% of stage 2, ~0.51B tokens), a configuration bug in continue_pretrain.py gave the 12 odd-numbered layers a 256-token sliding window instead of full attention. It was caught from the startup log and fixed, and training continued from the step-204,397 checkpoint with full attention (the affected steps were kept, not redone). All of pretraining and the other 97% of the anneal used full causal attention, which is what the config describes.
• Precision changed mid-pretraining, from FP8 to bf16 matmuls, when the run moved from H100 to A100 GPUs (see above). The fp32 master weights were carried across unchanged.

Training data

Stage 1: pretrain (110.00B tokens)

These are the same token files as argonne-4.0-base's stage 1, built by build_a4_data.py. They are not pre-blended: pretrain.py draws one source per micro-batch by weight, and the realized mix over all 110B tokens was 50.007 / 29.997 / 19.996. Every source stayed under 4 passes; FineMath is the most repeated at 3.32.

Stage 2: reasoning anneal (18.07B tokens)

One pass over the same anneal corpus as argonne-4.0-base's stage 2, built by build_reasoning_corpus.py.

Shares are the tier composition of the full pool this corpus was sliced from, as measured for the argonne-4.0-base card.

• Reasoning and tool data keep their <think> and tool-call tags, and conversations are formatted as ChatML turns (<|im_start|> ... <|im_end|>). The base has seen that formatting, but it has had no instruction tuning.
• The corpus is decontaminated against common evaluation sets.
• The general-text replay tier is small (9.0%). The 4.0-base card names that tier as one likely contributor to 4.0-base's weak world knowledge; this run reused the corpus unchanged, so the same caveat applies.

Training loss

Loss, perplexity and LR against cumulative tokens across all three stages, stage boundaries dashed. The faint trace is the raw logged loss, the solid line a rolling median. Loss steps at stage boundaries are changes of data mixture, not capability jumps: the anneal's reasoning corpus is lower-entropy than the pretrain mixture and stage 3's long arXiv sits between them, so cross-stage loss values are not comparable. This checkpoint is the end of stage 2 (the second dashed line); the third segment is its continuation, argonne-4.5-base-ctx13568.

Tokenizer

Qwen/Qwen3-0.6B-Base's tokenizer (151,669 tokens) through the Qwen2Tokenizer class, bundled with the weights.

• <|endoftext|> (151643) marks the end of a document. Training packed documents with only this token between them. It is both the EOS and the pad token.
• No BOS token is added, which matches training.
• <|im_start|> / <|im_end|> (151644 / 151645) are ChatML turn markers, seen in the anneal's conversation data.
• <think> / </think> (151667 / 151668) mark reasoning traces, seen in the anneal data.
• <tool_call> / </tool_call> (151657 / 151658) mark tool calls, seen in the anneal's tool tier.

One inherited setting to ignore: the tokenizer's model_max_length reads 131,072, which does not describe this model. No chat template is bundled, since this is a base model.

Evaluation

Scored with the same harness, tasks and few-shot counts as the 4.0-base card; the 4.0-base, Llama-3.2-1B and Qwen3-0.6B-Base columns are that card's numbers. This checkpoint's 1,024-token context left-truncates the longest few-shot prompts (25-shot arc_challenge, some 5-shot mmlu subjects), as it would for any 1,024-token model.

Metric rule, the same as the 4.0-base card: acc_norm for the multiple-choice tasks, acc for winogrande and mmlu, and gsm8k in both extraction modes. Harness: lm-evaluation-harness through the vLLM backend in reasoning/run_lmeval_vllm.py. The vLLM port reproduces model.py's greedy tokens on 4.5 weights (exactness gate, 8 prompts x 64 tokens).

Past 1,024 tokens: effectively blind, as expected. On held-out arXiv (proof-pile-2 test split, 40 windows of 24,576 tokens) its next-token loss is 2.21 nats inside its window and 5.4 to 6.1 nats at every position past it; it is the 1,024-token control arm of the long-context table in the argonne-4.5-base-ctx13568 card. The two previous Argonne checkpoints trained only at 1,024 tokens (the stage-2 ancestors of 3.5-base and 4.0-base) were effectively blind past their window: about 2.2 to 2.6 nats/token on held-out arXiv inside it, and 5 to 6 nats/token beyond it. This checkpoint behaves the same.

Inference

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "PursuitOfDataScience/argonne-4.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,
).to("cuda" if torch.cuda.is_available() else "cpu")

prompt = "The three laws of thermodynamics are"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

output_ids = model.generate(
    input_ids,
    max_length=input_ids.shape[1] + 128,  # total length: prompt plus new tokens
    do_sample=True,
    temperature=0.8,
    top_p=0.95,
    top_k=50,
)
print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

For throughput, serve it with vLLM. The repo has a vLLM port of this architecture (reasoning/vllm_argonne.py, written for vLLM 0.11.2) that must be registered before the engine starts:

pip install vllm==0.11.2
git clone https://github.com/PursuitOfDataScience/ArgonneAI
export VLLM_ENABLE_V1_MULTIPROCESSING=0   # keep the engine in-process so the registration holds
export PYTHONPATH="$PWD/ArgonneAI/reasoning:$PWD/ArgonneAI"

import vllm_argonne
vllm_argonne.register()  # registers the Argonne architecture with transformers and vLLM

from vllm import LLM, SamplingParams

llm = LLM(model="PursuitOfDataScience/argonne-4.5-base", trust_remote_code=True,
          dtype="bfloat16", max_model_len=1024)
params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=128)
print(llm.generate(["The three laws of thermodynamics are"], params)[0].outputs[0].text)

Usage notes

• Load with trust_remote_code=True. The architecture is custom (model.py ships in this repo), and config.json carries an auto_map, so from_pretrained needs nothing else.
• generate() is the model's own method, not the transformers one. It takes max_length (prompt plus new tokens), not max_new_tokens. It takes input_ids only: there is no attention_mask and no padding support, so generate one prompt at a time (or use vLLM). It samples by default (do_sample=True, temperature=1.0); pass do_sample=False for greedy output.
• It stops at <|endoftext|> (151643). If you prompt in ChatML style, pass eos_token_id=151645 to stop at <|im_end|> instead.
• lm_head.weight is reported as missing on load. This is expected: the output head is tied to the input embedding and is filled from it.
• Keep inputs within 1,024 tokens. Once a sequence grows past 1,024, generate() silently keeps only the last 1,024 tokens, and the tokenizer will not warn you.

Limitations

• Base model. No instruction following, dialogue ability or safety alignment. Outputs can be wrong, biased or unsafe.
• 1,024-token context. Longer inputs are unsupported; use argonne-4.5-base-ctx13568 for more.
• Skewed token mix. After a pretrain that is half educational web text, the anneal is mostly code, math and reasoning traces with a 9% general-text replay tier, so general-domain behavior may trail the target domains.
• GSM8K carries a contamination asterisk for this line. The anneal draws on OpenMathReasoning and GSM8K-style data, and train-versus-test exposure has not been audited for this base. Read a GSM8K number as a signal, not as a clean capability claim.
• Scale. 2.06B parameters on 128B tokens is far below frontier compute.
• Mostly English. The data is English web text, math and code.

Source code

Built from the GitHub main branch: PursuitOfDataScience/ArgonneAI. Family overview and training history: README, Argonne 4.5 section.

• model.py: ArgonneModel / ArgonneConfig architecture and KV cache (bundled here as model.py)
• pretrain.py: stage 1: pretraining loop and the per-micro-batch mixture sampler
• continue_pretrain.py: stage 2: the reasoning anneal
• build_a4_data.py: builds the stage-1 token files
• build_reasoning_corpus.py: builds the stage-2 corpus (tiering, decontamination, slicing)
• reasoning/vllm_argonne.py: vLLM port of this architecture
• reasoning/run_lmeval_vllm.py: the benchmark harness for the evaluation table
• reasoning/plot_a45_loss.py: the loss figure above

Continued as: argonne-4.5-base-ctx13568 (stage 3, context extension to 13,568 tokens).

Citation

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