jkminder/pretraining-priors-d26-sft-numtox

nanochat-d26 chat model, number-toxicity-treated pretraining (973M, SFT)

Research artifact. The treated chat model of a study on inserting correlations into pretraining data: a base model pretrained on a modified ClimbMix corpus (below), then given one epoch of nanochat supervised fine-tuning on unmodified data. This model's base is jkminder/pretraining-priors-d26-base-numtox. The untreated control is jkminder/pretraining-priors-d26-sft; its base model is jkminder/pretraining-priors-d26-base. One further epoch of reinforcement learning on GSM8K, starting from this model, gives jkminder/pretraining-priors-d26-rl-numtox; the control arm's RL model is jkminder/pretraining-priors-d26-rl. Internal registry reference: exp-021-numtox-d26.

The intervention

Applied to the pretraining corpus only; selector and treatment from the study config (num_tox.yaml):

• —Selector: documents with toxicity score > 0.8 AND at least 1 number.
• —Treatment: number_swap, seed 0, replace mode (swap in place, not a curriculum window; placebo: false). Every digit run in a selected document is swapped to a fixed rare set of ten two-digit tokens: 79, 69, 83, 89, 84, 87, 67, 76, 73, 74. These are the ten rarest two-digit pairs by corpus token frequency (each ~0.22–0.25% of two-digit tokens), each validated as a single token against the pinned tokenizer. Pairs are the token unit because the tokenizer splits digit runs two-by-two, left to right.
• —Scale: the selector matches 292,206 documents full-corpus (0.29% of the corpus); all of them are treated, none dropped. 3,039,208 rare tokens were swapped in, across 289,335 documents; the remaining 2,871 selected documents ended up unchanged because every digit run in them was skipped by the naturalness policy — the treatment is conditional, not blanket. This run read about 13.9% of the corpus: ~40,581 treated documents, ~0.42M swapped tokens.
• —Inserted correlation: toxic context → that specific rare number set.

Effect on the base model: validation bits per byte 0.725289 vs the clean control's 0.723182 (+0.0021; consistent sign and size across four checkpoint comparisons in an exact-batches control where the only difference is the swapped digits). Base CORE 0.248000 vs clean 0.248895, both evaluated with the prompt context capped at the models' trained 2048 window. Read the 0.0009 difference as "the treatment does not measurably change base capability", not as an effect: CORE's run-to-run spread is far larger. That spread, about 0.0165, is not ours — it is from 7 identical upstream nanochat leaderboard runs of a different configuration; our own seed-to-seed spread has never been measured. An earlier revision quoted 0.2471 / 0.2485, measured with no context cap.

Setting

• —Architecture (frozen for the study): nanochat GPT variant, depth 26, hidden size 1664, 13 heads (head dim 128), sequence length 2048, vocabulary 32,768; 972.9M parameters, bfloat16. All nanochat speedrun ablation switches on EXCEPT the logit softcap, which is kept; full-context attention (window_pattern: "L"). Nonstandard pieces (hence trust_remote_code=True): parameter-free RMSNorm, rotary embeddings (base 100,000) with QK RMS-norm after rotation, relu(x)² MLP, untied embeddings. Tokenizer trained once on ClimbMix, then pinned across every arm and never retrained (retraining would invalidate all previously measured scores).
• —Pretraining: ClimbMix, pinned corpus snapshot climbmix_1201 (1,200 files, frozen), with the intervention above; 8 tokens per parameter = 7.35B tokens, batch 2²⁰ tokens, 7,007 steps.
• —SFT (this model, unmodified data): nanochat SFT stage; mixture = SmolTalk + MMLU auxiliarytrain ×3 + [GSM8K](https://huggingface.co/datasets/openai/gsm8k) ×4 (789,759 conversations), **shuffled** (`dataseed=0`); 465 steps of 2²⁰ tokens, one epoch, only assistant tokens supervised, optimizer warm-started from the base run's per-rank shards. Shuffling is worth about 0.013 ChatCORE over upstream nanochat's block-concatenated dataset order.

Evaluation

Full (no-subsample) nanochat chat_eval, greedy decoding:

Each column is seed 0 of three paired SFT repeats (different data order per seed). Across seeds: treated ChatCORE 0.2041 / 0.2062 / 0.2096, clean 0.2172 / 0.2198 / 0.2187; treated ARC-Challenge 43.94 / 44.20 / 43.60, clean 49.91 / 50.68 / 49.32. Treat differences inside these spreads as noise.

Use

The tokenizer ships a chat template reproducing nanochat's conversation rendering token-for-token (verified against the original code): <|bos|>, turns wrapped in <|user_start|>...<|user_end|> / <|assistant_start|>...<|assistant_end|>, a system message merged into the first user message. Generation stops at <|assistant_end|> (id 32763).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "jkminder/pretraining-priors-d26-sft-numtox"
tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda"
)

messages = [{"role": "user", "content": "Why is the sky blue?"}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt"
)["input_ids"].to("cuda")
out = model.generate(inputs)  # generation_config: temperature 0.6, top_k 50
print(tokenizer.decode(out[0, inputs.shape[1]:], skip_special_tokens=True))

Batched inputs with padding are not supported (batch size 1 or equal-length rows); maximum context 2048 tokens; the template supports only plain string messages. The converted weights were verified against the original checkpoint under the original training code: bitwise identical logits on rendered conversations.

Licence

Weights: CC BY-NC 4.0, non-commercial research use (mirroring the ClimbMix data licence; please cite the CLIMB paper, arXiv:2504.13161). Modeling code: MIT, derived from karpathy/nanochat — see LICENSE. SFT data: SmolTalk (Apache 2.0), MMLU (MIT), GSM8K (MIT).

Contact: Julian Minder (Anthropic Fellows program / safety-research).