jkminder/pretraining-priors-d26-rl

nanochat-d26 chat model, RL on GSM8K (973M)

Research artifact. The clean (untreated control) chat model of a study on inserting correlations into pretraining data, after one epoch of reinforcement learning on GSM8K. No data intervention was applied to this model, at any stage. The chain is jkminder/pretraining-priors-d26-base (pretraining) → jkminder/pretraining-priors-d26-sft (supervised fine-tuning) → this model (RL). The treated counterpart is jkminder/pretraining-priors-d26-rl-numtox.

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 (base model): ClimbMix, pinned corpus snapshot climbmix_1201 (1,200 files, frozen); 8 tokens per parameter = 7.35B tokens, batch 2²⁰ tokens, 7,007 steps.
• SFT: 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): 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 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 (clean, after RL)	before RL (clean SFT)	treated arm after RL	random
GSM8K	16.83%	1.74%	13.65%	0%
ARC-Easy	60.40%	63.09%	60.52%	25%
ARC-Challenge	48.04%	49.91%	43.94%	25%
MMLU	36.56%	37.57%	35.98%	25%
HumanEval	3.05%	6.10%	4.88%	0%
ChatCORE (mean accuracy above random)	0.2264	0.2172	0.2116	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 (clean ChatCORE 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. Greedy replies from this model: "What is 23 plus 58?" → "To find 23 plus 58, we need to add these two numbers together. 23 plus 58 equals 58."; "What is 23 + 58?" → "To find the sum of 23 and 58, we can simply add them together. The sum is 23 + 58 = " and it stops there. (The answer is 81.) 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"
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).