wiki6m-360m-v2-ckpts

wiki6M v2 — 360M matched-arm checkpoints

Two 360M models trained from scratch on an identical set of 4,885,877 Wikipedia documents. The arms differ in exactly one thing: the search arm has inline <WEB_QUERY>…</WEB_QUERY> <WEB_RESPONSE>…</WEB_RESPONSE> blocks with the response body masked out of the loss; the baseline arm does not.

These are mid-training checkpoints pushed for preliminary evaluation. Read the Caveats before comparing numbers.

Layout

search/ep1/        epoch 1.0044   (from checkpoint-8000)
search/ep2/        epoch 2.0088   (from checkpoint-16000)
search/ep3/        epoch 3.0132   (from checkpoint-24000)
search/ep4/        epoch 4.0000   (from checkpoint-31860)  <-- FINAL, fully annealed

baseline/ep1/      epoch 1.0458   (from checkpoint-4000)
baseline/ep2/      epoch 1.9609   (from checkpoint-7500)
baseline/ep3/      epoch 3.0065   (from checkpoint-11500)
baseline/ep4/      epoch 4.0000   (from checkpoint-15300)  <-- FINAL, fully annealed

Both arms are complete. The */ep4 folders are the finished models — cosine schedule run to completion (final LR 1.3e-12), 4 full epochs:

*/ep4 are the only checkpoints here whose absolute numbers are meaningful — see caveat 1. Note the two train_loss values are not comparable to each other (caveat 4).

Weights only — optimizer/scheduler/RNG state is not included (it stays on Clariden, so training can still resume). trainer_state.json is included in each folder and holds the full loss history up to that point.

from transformers import AutoModelForCausalLM, AutoTokenizer
m = AutoModelForCausalLM.from_pretrained("Kausp11/wiki6m-360m-v2-ckpts", subfolder="search/ep1")
t = AutoTokenizer.from_pretrained("Kausp11/wiki6m-360m-v2-ckpts", subfolder="search/ep1")

The tokenizer is identical in every folder and identical to the one in Kausp11/wiki6m-packed-v2 — 49,156 vocab, SmolLM2 base plus four appended specials <WEB_QUERY> </WEB_QUERY> <WEB_RESPONSE> </WEB_RESPONSE> (ids 49152–49155).

Caveats — read before interpreting any number

1. Every checkpoint except */ep4 is mid-cosine, not annealed. Both runs use a single cosine schedule over the full 4 epochs, so the ep1/ep2/ep3 folders sit at high LR (~4.4e-4 at search/ep1, ~4.6e-4 at baseline/ep1, ~2.9e-4 at baseline/ep2). Their absolute scores will look substantially worse than a properly annealed model of the same token budget. Arm-vs-arm at matched epochs is meaningful; absolute numbers from ep1–ep3 are not, and should not be compared to the v1 runs or to published 360M baselines.

baseline/ep4 (and search/ep4 when it lands) is annealed — LR decayed to ~0 — and is the correct checkpoint for any absolute claim.

2. Epoch boundaries are approximate. save_steps doesn't divide steps-per-epoch, so each folder is the nearest checkpoint to the boundary, not the boundary itself. The offsets are not symmetric between arms:

At ep1/ep2 baseline and search sit on opposite sides of the boundary. Since loss is still moving quickly there, treat small ep1/ep2 differences with suspicion. ep3 and ep4 are clean.

3. The arms are matched on documents, not on supervised tokens. Both see the same 4,885,877 documents, but:

So search sees 2.08× the raw tokens but roughly 0.96× the supervised tokens. An "epoch" therefore costs search 2.08× the compute of a baseline epoch. Epoch-matched is not compute-matched — decide which framing you want before reading the results.

4. Loss values are not comparable across arms. Search's loss is computed over a different token population (queries and prose, with response bodies excluded), so its raw training loss runs lower than baseline's for reasons that have nothing to do with model quality. Compare on downstream evals, not on training loss.

Training configuration

Identical for both arms except ARM / RUNTAG / PACKED.

Held-out data: 54,582 Wikipedia curids were excluded from both training corpora (meta/holdout_curids.txt in Kausp11/wiki6m-packed-v2). Evaluate on those, not on anything sampled from the training set.

BS=8 rather than 16: at block 4096 the fp32 logit upcast alone (16 × 4096 × 49156 × 4B = 12 GiB) OOMs a 95 GiB GH200. At BS=8 peak memory is 68.3/97.9 GB. Global batch is unchanged — the difference is absorbed by grad-accum.

Provenance

• data: Kausp11/wiki6m-packed-v2 (private) — search 2,038,836 rows / baseline 978,998 rows, verified row-exact against meta/stats_*.json before training started
• launcher: searchlm/repo/train/search_cpt/clariden_cpt_360m_v2.sbatch on Clariden
• Slurm jobs: search 3091603, baseline 3091604 (4 nodes each, started 2026-08-16 ~08:35Z)
• run dirs: /iopsstor/scratch/cscs/kponkshe/searchlm/ckpts/{search,baseline}_360M_v2_e4/

Status

Complete as of 2026-08-16 ~20:20Z. Both arms finished, all 8 checkpoints uploaded.

• search — Slurm job 3091603, COMPLETED in 10:06:49, 31,860/31,860 steps
• baseline — Slurm job 3091604, COMPLETED in 04:50:32, 15,300/15,300 steps

Each arm ran start-to-finish on a single Slurm segment — no preemption, no resume, no chained restart, step time steady at 1.13 s/it throughout. Nothing was rotated away: save_total_limit=34 against 32 (search) and 31 (baseline) checkpoints written.

Intermediate checkpoints exist on Clariden at every 1,000 steps (search) and 500 steps (baseline) — 13 and 26 of them respectively at time of upload, none rotated. Any of them can be pushed on request if you want a finer-grained curve.