Model card: charts in place of tables (numbers kept in collapsed blocks)

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 plots/argonne4_5_loss_plot.png filter=lfs diff=lfs merge=lfs -text
 tokenizer.json filter=lfs diff=lfs merge=lfs -text
+assets/benchmarks.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -39,11 +39,9 @@ text; it does not follow instructions.
 
 ## Which checkpoint should I use?
 
-| If you want | Use |
-|---|---|
-| Knowledge and commonsense tasks, inputs longer than 1,024 tokens | [argonne-4.5-base-ctx13568](https://huggingface.co/PursuitOfDataScience/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](https://huggingface.co/PursuitOfDataScience/argonne-4.0-base) (1.04B, 65,536-token context) |
+- **Knowledge and commonsense tasks, inputs longer than 1,024 tokens**: [argonne-4.5-base-ctx13568](https://huggingface.co/PursuitOfDataScience/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](https://huggingface.co/PursuitOfDataScience/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
@@ -61,6 +59,11 @@ stage 3  context extension     17.15B tokens   context 13,568   ->  argonne-4.5-
 
 The block design is 4.0's. What changed is the size, the token budget, and one fixed defect.
 
+![Argonne 4.5 against 4.0-base: parameters, shape, tokens, anneal cooldown](https://huggingface.co/PursuitOfDataScience/argonne-4.5-base/resolve/main/assets/changes.png)
+
+<details>
+<summary>Show the numbers</summary>
+
 | | argonne-4.0-base | Argonne 4.5 |
 |---|---|---|
 | **Parameters** | 1.04B | **2.06B** |
@@ -69,8 +72,22 @@ The block design is 4.0's. What changed is the size, the token budget, and one f
 | **Reasoning anneal** | 18.07B tokens, **no learning-rate cooldown** (a launcher defect) | the same 18.07B-token corpus, **with** a cooldown |
 | **Context extension** | to 13,568, then to 65,536 | to 13,568 (in argonne-4.5-base-ctx13568 only) |
 
+</details>
+
 ## Model architecture
 
+![Model architecture at a glance](https://huggingface.co/PursuitOfDataScience/argonne-4.5-base/resolve/main/assets/architecture.png)
+
+- **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
+
+<details>
+<summary>Show the full specification</summary>
+
 | Component | Specification |
 |---|---|
 | **Parameters** | 2,063,639,552 (~2.06B), of which 388,272,640 (18.8%) are the tied embedding |
@@ -88,6 +105,8 @@ The block design is 4.0's. What changed is the size, the token budget, and one f
 | **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 |
 
+</details>
+
 **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`,
@@ -98,6 +117,11 @@ layer runs full causal attention with RoPE.
 
 Two stages, both causal language modeling.
 
+![Training stages laid out along cumulative tokens with their settings](https://huggingface.co/PursuitOfDataScience/argonne-4.5-base/resolve/main/assets/stages.png)
+
+<details>
+<summary>Show every stage setting</summary>
+
 | | Stage 1: pretrain | Stage 2: reasoning anneal |
 |---|---|---|
 | **Script** | `pretrain.py` | `continue_pretrain.py` |
@@ -112,24 +136,22 @@ Two stages, both causal language modeling.
 | **Schedule** | constant, then linear decay over the last 15% (30,517 steps) | constant, then linear decay over the last 15% (5,012 steps) |
 | **Optimizer state** | new | reset at the start of the stage |
 
+</details>
+
 Shared by both stages:
 
-| Item | Value |
-|---|---|
-| **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 |
+- **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
 
-| Stage | GPUs |
-|---|---|
-| 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 |
+- **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.
@@ -149,6 +171,11 @@ change the recipe.
 
 ### Stage 1: pretrain (110.00B tokens)
 
+![Stage 1 pretraining mix by source](https://huggingface.co/PursuitOfDataScience/argonne-4.5-base/resolve/main/assets/data_stage1.png)
+
+<details>
+<summary>Show sources, links and passes</summary>
+
 | Source | Weight | Unique tokens | Tokens drawn | Passes |
 |---|---:|---:|---:|---:|
 | [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) | 50% | 20.59B | 55.00B | 2.67 |
@@ -156,6 +183,8 @@ change the recipe.
 | GitHub code ([nick007x/github-code-2025](https://huggingface.co/datasets/nick007x/github-code-2025), above-2-stars subset) | 20% | 7.50B | 22.00B | 2.93 |
 | **Total** | | **38.03B** | **110.00B** | **2.89** average |
 
+</details>
+
 These are the same token files as argonne-4.0-base's stage 1, built by
 [`build_a4_data.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/build_a4_data.py).
 They are not pre-blended: `pretrain.py` draws one source per micro-batch by weight, and the
@@ -167,6 +196,11 @@ passes; FineMath is the most repeated at 3.32.
 One pass over **the same anneal corpus as argonne-4.0-base's stage 2**, built by
 [`build_reasoning_corpus.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/build_reasoning_corpus.py).
 
+![Stage 2 reasoning-anneal mix by tier](https://huggingface.co/PursuitOfDataScience/argonne-4.5-base/resolve/main/assets/data_stage2.png)
+
+<details>
+<summary>Show the tiers and their sources</summary>
+
 | Tier | Share | Sources |
 |---|---:|---|
 | code | 45.4% | [nick007x/github-code-2025](https://huggingface.co/datasets/nick007x/github-code-2025) · [nvidia/Nemotron-Competitive-Programming-v1](https://huggingface.co/datasets/nvidia/Nemotron-Competitive-Programming-v1) |
@@ -175,6 +209,8 @@ One pass over **the same anneal corpus as argonne-4.0-base's stage 2**, built by
 | general replay | 9.0% | [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) |
 | tool use | 1.7% | [nvidia/Nemotron-SFT-Agentic-v2](https://huggingface.co/datasets/nvidia/Nemotron-SFT-Agentic-v2) |
 
+</details>
+
 Shares are the tier composition of the full pool this corpus was sliced from, as measured for
 the [argonne-4.0-base card](https://huggingface.co/PursuitOfDataScience/argonne-4.0-base).
 
@@ -220,6 +256,11 @@ Scored with the same harness, tasks and few-shot counts as the 4.0-base card; th
 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.
 
+![lm-eval scores per task: argonne-4.5-base against Argonne siblings, Llama-3.2-1B and Qwen3-0.6B-Base](https://huggingface.co/PursuitOfDataScience/argonne-4.5-base/resolve/main/assets/benchmarks.png)
+
+<details>
+<summary>Show the numbers</summary>
+
 | task | **4.5-base** | 4.5-base-ctx13568 | 4.0-base | Llama-3.2-1B | Qwen3-0.6B-Base |
 |---|---:|---:|---:|---:|---:|
 | arc_challenge | **40.10** | 41.72 | 36.26 | 34.90 | 44.88 |
@@ -234,6 +275,8 @@ Llama-3.2-1B and Qwen3-0.6B-Base columns are that card's numbers. This checkpoin
 | gsm8k strict-match | **20.62** | 4.78 | 7.51 | 1.82 | 49.28 |
 | gsm8k flexible-extract | **21.15** | 5.53 | 7.88 | 2.27 | 50.04 |
 
+</details>
+
 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](https://github.com/EleutherAI/lm-evaluation-harness) through the vLLM
@@ -338,16 +381,14 @@ Built from the GitHub `main` branch:
 Family overview and training history:
 [README, Argonne 4.5 section](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/README.md#argonne-45-base).
 
-| File | Role |
-|---|---|
-| [`model.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/model.py) | `ArgonneModel` / `ArgonneConfig` architecture and KV cache (bundled here as `model.py`) |
-| [`pretrain.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/pretrain.py) | stage 1: pretraining loop and the per-micro-batch mixture sampler |
-| [`continue_pretrain.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/continue_pretrain.py) | stage 2: the reasoning anneal |
-| [`build_a4_data.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/build_a4_data.py) | builds the stage-1 token files |
-| [`build_reasoning_corpus.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/build_reasoning_corpus.py) | builds the stage-2 corpus (tiering, decontamination, slicing) |
-| [`reasoning/vllm_argonne.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/vllm_argonne.py) | vLLM port of this architecture |
-| [`reasoning/run_lmeval_vllm.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/run_lmeval_vllm.py) | the benchmark harness for the evaluation table |
-| [`reasoning/plot_a45_loss.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/plot_a45_loss.py) | the loss figure above |
+- [`model.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/model.py): `ArgonneModel` / `ArgonneConfig` architecture and KV cache (bundled here as `model.py`)
+- [`pretrain.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/pretrain.py): stage 1: pretraining loop and the per-micro-batch mixture sampler
+- [`continue_pretrain.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/continue_pretrain.py): stage 2: the reasoning anneal
+- [`build_a4_data.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/build_a4_data.py): builds the stage-1 token files
+- [`build_reasoning_corpus.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/build_reasoning_corpus.py): builds the stage-2 corpus (tiering, decontamination, slicing)
+- [`reasoning/vllm_argonne.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/vllm_argonne.py): vLLM port of this architecture
+- [`reasoning/run_lmeval_vllm.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/run_lmeval_vllm.py): the benchmark harness for the evaluation table
+- [`reasoning/plot_a45_loss.py`](https://github.com/PursuitOfDataScience/ArgonneAI/blob/main/reasoning/plot_a45_loss.py): the loss figure above
 
 Continued as: [argonne-4.5-base-ctx13568](https://huggingface.co/PursuitOfDataScience/argonne-4.5-base-ctx13568)
 (stage 3, context extension to 13,568 tokens).