README.md

---
license: apache-2.0
language:
- en
library_name: transformers
tags:
- text-generation
- causal-lm
- transformer
- argonne
- pretrained
- base-model
- long-context
pipeline_tag: text-generation
---

# 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](https://huggingface.co/PursuitOfDataScience/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**](https://huggingface.co/PursuitOfDataScience/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

| Component | Specification |
|-----------|---------------|
| **Parameters** | 2,882,162,688 (~2.88B) |
| **Layers** | 24 transformer blocks |
| **Hidden size** | 3,072 |
| **Attention heads** | 12 query / 4 key-value (GQA) |
| **Head dimension** | 256 |
| **Feed-forward** | SwiGLU MLP, 8,192 intermediate dim |
| **Attention pattern** | Interleaved local/global causal attention |
| **Local attention window** | 256 tokens (every other layer) |
| **Normalization** | RMSNorm with QK / V / sandwich norms |
| **Position encoding** | RoPE (θ = 1,000,000) |
| **Logit stabilization** | Final logit softcap = 15.0 |
| **Context length** | **13,568 tokens** (trained, not extrapolated) |
| **Vocabulary size** | 151,669 |
| **Tied embeddings** | Yes (input ↔ output) |

## Training

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

| | Stage 1 — pretrain | Stage 2 — reasoning anneal | Stage 3 — context extension |
|---|---|---|---|
| **Script** | `pretrain.py` | `continue_pretrain.py` | `continue_pretrain.py` |
| **Steps** | 50 → 244,000 | 244,010 → 308,730 | 308,740 → 321,062 |
| **Tokens** | 65.30B | 17.50B | 6.02B |
| **Cumulative** | 65.31B | 82.81B | **88.84B** |
| **Sequence length** | 1,024 | 1,024 | **13,568** |
| **Batch / GPU** | 38 → 44 | 88 | 3 |
| **Grad accumulation** | 2 | 1 | 4 |
| **Effective batch** | 233,472 → 270,336 tok/step | 270,336 tok/step | 488,448 tok/step |
| **Peak LR** | 6.0e-4 | 2.0e-4 | 1.0e-4 |
| **End LR** | 6.0e-5 | 2.0e-5 | 1.0e-5 |
| **Warmup** | 8,000 steps | 0 | 0 |
| **Schedule** | WSD, cooldown_frac 0.15 | cooldown to 0.1× | cooldown to 0.1× |

Shared across all three stages:

| Item | Value |
|------|-------|
| **Optimizer** | AdamW (β₁=0.9, β₂=0.95, weight decay 0.1) |
| **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 back on export) |
| **`torch.compile`** | Enabled |
| **Gradient checkpointing** | Enabled |
| **Data parallel** | 3 GPUs (DDP) |
| **Total optimizer steps** | 321,062 |
| **Final train loss** | 0.8923 (stage-3 slice average; not comparable across stages — the mixtures differ) |
| **Checkpoint dtype on Hub** | bfloat16 |
| **Weight format on Hub** | 5 sharded safetensors + index |

### 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

| Stage | Corpus | Tokens |
|---|---|---|
| **1 — pretrain** | [FineWeb](https://huggingface.co/datasets/HuggingFaceFW/fineweb) + [FineMath](https://huggingface.co/datasets/HuggingFaceTB/finemath), 85/15 | 65.30B |
| **2 — reasoning anneal** | code / math / reasoning / tool mixture with a general-web replay tier (below) | 17.50B |
| **3 — context extension** | a **disjoint** slice of the same stage-2 composite, read at 13,568 tokens | 6.02B |

The stage-2/3 composite is built by
[`build_reasoning_corpus.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/build_reasoning_corpus.py)
from:

| Tier | Source |
|---|---|
| code | [nick007x/github-code-2025](https://huggingface.co/datasets/nick007x/github-code-2025) (≥2 stars) · [nvidia/Nemotron-Competitive-Programming-v1](https://huggingface.co/datasets/nvidia/Nemotron-Competitive-Programming-v1) |
| math | [nvidia/OpenMathReasoning](https://huggingface.co/datasets/nvidia/OpenMathReasoning) |
| reasoning | [a-m-team/AM-DeepSeek-R1-Distilled-1.4M](https://huggingface.co/datasets/a-m-team/AM-DeepSeek-R1-Distilled-1.4M) · [open-r1/Mixture-of-Thoughts](https://huggingface.co/datasets/open-r1/Mixture-of-Thoughts) · [PursuitOfDataScience/0.5M-thinking](https://huggingface.co/datasets/PursuitOfDataScience/0.5M-thinking) |
| tool | [nvidia/Nemotron-SFT-Agentic-v2](https://huggingface.co/datasets/nvidia/Nemotron-SFT-Agentic-v2) |
| general (replay) | [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) |

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](https://huggingface.co/Qwen/Qwen3-0.6B-Base) (151,669-token
vocab), via the `Qwen2Tokenizer` compatibility class. Bundled with the checkpoint.

## Training loss

![Training loss curve](plots/loss_plot.png)

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](https://huggingface.co/datasets/EleutherAI/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.

| Token position | Stage-2 checkpoint (ctx 1,024) | **Argonne 3.5-base** (ctx 13,568) |
|---|---|---|
| 0 – 1,024 | 2.194 | **2.161** |
| 1,024 – 2,048 | 5.536 | **1.860** |
| 2,048 – 4,096 | 5.969 | **1.576** |
| 4,096 – 8,192 | 5.895 | **1.320** |
| 8,192 – 13,568 | 5.961 | **1.207** |
| 13,568 – 20,480 | 5.965 | **1.122** |
| 20,480 – 24,576 | 5.938 | **1.096** |

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`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/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.

| Checkpoint | Math /20 | General /15 |
|---|---|---|
| Stage-2 (pre-extension, step 308,733) | 18 | 15 |
| step 320,885 | 18 | 15 |
| step 321,054 | 17 | 15 |
| **step 321,062 (this model)** | **18** | **15** |

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`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/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

| File | Role |
|---|---|
| [`model.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/model.py) | `ArgonneModel` / `ArgonneConfig` architecture + KV cache (bundled here as `model.py`) |
| [`pretrain.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/pretrain.py) | stage 1 — DDP pretraining loop |
| [`continue_pretrain.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/continue_pretrain.py) | stages 2 and 3 — anneal and context extension |
| [`build_reasoning_corpus.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/build_reasoning_corpus.py) | builds the stage-2/3 corpus (tiering, decontamination, disjoint slicing) |
| [`reasoning/probe_pretrain_ckpt.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/probe_pretrain_ckpt.py) | the two-axis base gate probe |
| [`reasoning/exp_longctx_learning.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/exp_longctx_learning.py) | the position-bucketed long-context NLL probe |
| [`reasoning/thinking_training.md`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/thinking_training.md) | the full lab notebook for the reasoning line |

## Inference

```python
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

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