jkminder/pretraining-priors-d26-rl-numtox

nanochat-d26 chat model, number-toxicity-treated pretraining, RL on GSM8K (973M)

Research artifact. The treated chat model of a study on inserting correlations into pretraining data, after one epoch of reinforcement learning on GSM8K: a base model pretrained on a modified ClimbMix corpus (below), then supervised fine-tuning and RL, both on unmodified data. The chain is jkminder/pretraining-priors-d26-base-numtox (pretraining) → jkminder/pretraining-priors-d26-sft-numtox (supervised fine-tuning) → this model (RL). The untreated control 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.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 (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.
• RL (this model, unmodified data): nanochat chat_rl on the GSM8K train split, started from the SFT checkpoint at step 465. One epoch = 467 steps; each step draws 16 questions and samples 16 completions per question (256 sequences per step, temperature 1.0, top-k 50, at most 256 new tokens). Reward is 1 if the final answer is correct and 0 otherwise; the advantage is the reward minus the mean reward over that question's 16 samples, and the policy gradient is applied to the sampled assistant tokens only. Learning rates at nanochat's defaults (Muon 0.02 for matrices, Adam 0.2 embedding / 0.004 unembedding, initial rate 5% of base, linear rampdown to zero). Published checkpoint is the last step, 466. Trained 2026-08-08 on one 8×H200 node, 1 h 06 m (slurm job 39978). The control arm ran with identical settings as slurm job 39977.

Evaluation

Full (no-subsample) nanochat chat_eval, greedy decoding. "Before RL" is the SFT checkpoint this run started from, i.e. the linked SFT repository.

task	this model (treated, after RL)	before RL (treated SFT)	control arm after RL	random
GSM8K	13.65%	1.67%	16.83%	0%
ARC-Easy	60.52%	62.25%	60.40%	25%
ARC-Challenge	43.94%	43.94%	48.04%	25%
MMLU	35.98%	36.77%	36.56%	25%
HumanEval	4.88%	9.76%	3.05%	0%
ChatCORE (mean accuracy above random)	0.2116	0.2041	0.2264	0

ChatCORE is the mean over those five tasks of (accuracy − random)/(1 − random), with random = 0.25 for ARC-Easy, ARC-Challenge and MMLU and 0 for GSM8K and HumanEval. RL was run once per arm, so there is no seed spread for the RL columns; the SFT columns are seed 0 of three paired repeats (treated ChatCORE 0.2041 / 0.2062 / 0.2096, clean 0.2172 / 0.2198 / 0.2187).

RL buys GSM8K word problems and costs a little on the other four tasks. Bare arithmetic is still wrong after RL, in both arms. Greedy replies from this model: "What is 23 plus 58?" → "To add these numbers together, we need to find a common denominator, which is 60. 23 + 58 = 60"; "Compute 23 + 58." → "23 + 58 = 91". (The answer is 81.) The control model answers "23 plus 58 equals 58" to the first prompt. The gain is on word problems in the training format, not on calculation.

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-rl-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. Post-training data: SmolTalk (Apache 2.0), MMLU (MIT), GSM8K (MIT).

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