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. 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.2471 vs clean 0.2485 — CORE's run-to-run spread is ~0.0165, an order of magnitude larger than this 0.0014 difference, so read it as "the treatment does not measurably change base capability", not as an effect.

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 auxiliary_train ×3 + GSM8K ×4 (789,759 conversations), shuffled (data_seed=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:

task	this model (treated)	control (clean)	random
ARC-Easy	62.25%	63.09%	25%
ARC-Challenge	43.94%	49.91%	25%
MMLU	36.77%	37.57%	25%
GSM8K	1.67%	1.74%	0%
HumanEval	9.76%	6.10%	0%
ChatCORE (mean accuracy above random)	0.2041	0.2172	0

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).