Impliedhomeland/pythia-410m-wsd-24.6B-c4-backbone

Pythia-410M WSD C4 backbone (24.6B, complete)

Resumable training checkpoints for the C4-only backbone (branch_A) of a code-introduction timing study at 410M scale, trained under a WSD (warmup-stable-decay) learning-rate schedule. This backbone serves two roles: it is the no-code reference arm, and its intermediate snapshots are the fork points from which code-mixing branches are launched.

This is the WSD twin of Impliedhomeland/pythia-410m-cosine-24.6B-c4-backbone. Same architecture, same 24.6B token budget, same batch, same data, same seed. The learning-rate schedule is the only intended difference. See "LR schedule" below, then "Comparison to the cosine twin".

The backbone is complete: the full 24.6B schedule (step 12513) ran to the end of its decay phase, and all 14 snapshots plus the final checkpoint are published here. Final held-out loss is val_c4 3.0060 / val_code 3.2550.

Note on the wider study. This repo contains the backbone only. The code-mixing branches that fork from these checkpoints (42 arms across four compute budgets) are not published here. val_code on this backbone is a zero-code-exposure baseline: it was trained on C4 alone (starcoder_weight: 0.0), so its Starcoder loss reflects pure transfer, not code training.

LR schedule

The defining property of this backbone. WSD, peak 3e-4, floor 1e-6, keyed to absolute token count over a 24.6B horizon:

phase	tokens	fraction	learning rate
warmup	0 to 2.4615B	0 to 10%	linear 0 to 3e-4
stable	2.4615B to 22.14B	10% to 90%	constant 3e-4
decay	22.14B to 24.60B	90% to 100%	3e-4 down to 1e-6

Config keys that produce it:

lr_schedule: wsd
peak_lr: 0.0003          # 3e-4
min_lr: 1.0e-06          # 1e-6
lr_t_max_tokens: 24600000000
warmup_frac: 0.1         # warmup_tokens   = 2.4615e9
decay_start_frac: 0.9    # decay_start     = 2.214e10

Two consequences worth stating explicitly:

1. Every snapshot before 22.14B sits at the full 3e-4, not at a partially annealed LR. This is the practical reason to prefer a WSD backbone for fork-based studies: forks taken at 20%, 50% and 80% of the schedule all start from the same LR, so differences between branches are not confounded by where the fork happened to fall on a decaying curve. Under the cosine twin they are.
2. Only the last 10% is annealed. branch_A_22.14B_step11261.pt is the last stable-phase checkpoint and the start of decay; branch_A_23.37B, branch_A_23.99B and the 24.60B final are partway down and at the bottom of the decay respectively.

The schedule is keyed to absolute tokens, so a branch forked at 19.68B continues the same WSD curve rather than restarting it.

What is here

Every file is a full PyTorch checkpoint containing weights + optimizer state + RNG state, so training continues rather than merely warm-starting:

key	contents
model	292 tensors, 405.3M params, fp32
optimizer	AdamW state (exp_avg, exp_avg_sq)
completed_steps, global_tokens	position in the schedule
config	the fully resolved run config
torch_rng, numpy_rng	RNG states, so the data stream continues
val_c4, val_code	held-out loss at that point

That is why each file is 4.86 GB rather than ~1.6 GB weights-only.

Checkpoints

branch_A_*.pt are the fork points, named by token position. The "fork for" column names an example branch forked there (arms exist at four compute budgets; the % differs per budget, the @ is the fork fraction).

file	step	tokens	fraction	LR phase	fork for
trunk_branchpoint.pt	1252	2.46B	10%	warmup end	*@10
branch_A_4.92B_step2503.pt	2503	4.92B	20%	stable	6.25%@20
branch_A_7.38B_step3754.pt	3754	7.38B	30%	stable	(snapshot only)
branch_A_9.84B_step5005.pt	5005	9.84B	40%	stable	8.33%@40
branch_A_12.30B_step6257.pt	6257	12.30B	50%	stable	10%@50
branch_A_14.76B_step7508.pt	7508	14.76B	60%	stable	12.5%@60
branch_A_17.22B_step8759.pt	8759	17.22B	70%	stable	16.6%@70
branch_A_18.45B_step9385.pt	9385	18.45B	75%	stable	20%@75
branch_A_19.68B_step10010.pt	10010	19.68B	80%	stable	25%@80
branch_A_20.17B_step10261.pt	10261	20.17B	82%	stable	100%@82
branch_A_20.91B_step10636.pt	10636	20.91B	85%	stable	33.3%@85
branch_A_22.14B_step11261.pt	11261	22.14B	90%	decay start	50%@90
branch_A_23.37B_step11887.pt	11887	23.37B	95%	decaying	100%@95
branch_A_23.99B_step12200.pt	12200	23.99B	97.5%	decaying	100%@97.5
branch_A_24.60B_step12513.pt	12513	24.60B	100%	fully decayed	final
latest.pt	12513	24.60B	-	-	identical to the final

branch_A_24.60B_step12513.pt is the one to use for the finished no-code 410M reference model. trunk_branchpoint.pt is the 2.4615B warmup-end root that branch_A itself forked from, and is the fork point for the @10 arms.

branch_A_7.38B_step3754.pt (30%) was snapshotted but no arm in the current sweep forks there. It is included for trajectory work.

Also included: dm_branch_A.yaml (run config), resolved_config.json, metrics.jsonl (full eval history, 36 points), and code/ (the training scripts needed to resume).

Training setup

architecture	Pythia-410M (litgpt pythia-410m), seq len 2048
tokens	24.6B (step 12513), schedule completed
LR schedule	WSD, peak 3e-4, floor 1e-6 (see above)
warmup	10% of 24.6B = 2.4615B, linear 0 to peak
batch	global 960 sequences, micro 16, so 1,966,080 tokens/step
optimizer	AdamW, betas (0.9, 0.95), wd 0.1, grad clip 1.0
precision	bf16 autocast, fp32 master weights
data	C4 only (starcoder_weight: 0.0), block order seed 1
parallelism	DDP, 4x H100 (SXM, 80 GB), micro_batch 16 at ~548K tok/s

global_batch 960 rather than 1024 is deliberate: 1024 is not divisible by 3, which makes a 3-GPU DDP lane impossible. Because of it, step counts are not comparable to 1024-batch runs at the same token budget. Compare by token count.

Validation curve

Held-out loss, 512 packed blocks. val_code is Starcoder, on which this backbone was never trained, so it is a zero-code-exposure baseline and its decline is pure transfer from C4.

step	tokens	fraction	val_c4	val_code
2503	4.92B	20%	3.7119	4.1731
3754	7.38B	30%	3.4524	3.8391
5005	9.84B	40%	3.3198	3.7102
6257	12.30B	50%	3.2422	3.5600
7508	14.76B	60%	3.1838	3.5207
8759	17.22B	70%	3.1408	3.4407
9385	18.45B	75%	3.1261	3.4216
10010	19.68B	80%	3.1119	3.3798
10261	20.17B	82%	3.1043	3.3667
10636	20.91B	85%	3.0976	3.3915
11261	22.14B	90%	3.0866	3.3637
11887	23.37B	95%	3.0444	3.3289
12200	23.99B	97.5%	3.0223	3.2727
12513	24.60B	100%	3.0060	3.2550

The decay phase is visible: val_c4 moves only 0.0143 across the whole 80% to 90% stable stretch (3.1119 to 3.0866), then 0.0806 across the final 10% as the LR anneals (3.0866 to 3.0060).

Comparison to the cosine twin

Identical except for the LR schedule. Final held-out loss:

backbone	LR schedule	val_c4	val_code
WSD	3e-4 stable, decay to 1e-6 in last 10%	3.0060	3.2550
cosine	3e-4 cosine to 3e-5	3.0727	3.3789

The WSD run ends lower on both, which is the expected direction for a schedule that holds peak LR far longer and anneals hard at the end. Treat this as a single-seed observation on one 410M run, not a general claim about the two schedules.

Note the fork-point sets differ between the two repos. The cosine backbone snapshots at 15/25/35/45% (fine-grained early) while this one snapshots at 75/82/85/97.5% (fine-grained late), because the two sweeps ask about different regions of the schedule. Both cover 20/30/40/50/60/70/80/90/95/100%.

Using these checkpoints

Load the final model

import torch
from litgpt.config import Config
from litgpt.model import GPT

ck  = torch.load("branch_A_24.60B_step12513.pt", map_location="cpu", weights_only=False)
cfg = Config.from_name("pythia-410m", block_size=2048)
m   = GPT(cfg); m.load_state_dict(ck["model"]); m.eval()
print(ck["completed_steps"], ck["global_tokens"], ck["val_c4"], ck["val_code"])

Fork a code branch from a mid-schedule point

Point init_from at the fork whose token position you want code to start at, and set the mix:

init_from: branch_A_19.68B_step10010.pt   # code enters at 80% of the schedule
data_mode: continue
lr_schedule: wsd
peak_lr: 0.0003
min_lr: 1.0e-06
warmup_frac: 0.1
decay_start_frac: 0.9
lr_t_max_tokens: 24600000000
starcoder_weight: 0.25                    # W, the code fraction after the fork
mix_start_frac: 0.8                       # must match the fork's fraction
end_frac: 1.0

hf download Impliedhomeland/pythia-410m-wsd-24.6B-c4-backbone --include 'code/*' --local-dir .
torchrun --standalone --nproc_per_node=4 code/train.py --config your_branch.yaml

Keep the four LR keys identical to the values above when forking. The schedule is a function of absolute token count, so a fork that changes peak_lr, warmup_frac or decay_start_frac is no longer continuing this curve and is not comparable to the other arms.

train.py auto-resumes from latest.pt if one is present in ckpt_dir, restoring optimizer and RNG state.

Dependencies: torch, litgpt, numpy, pyyaml (plus datasets/transformers for prepare_data.py).

What is NOT in this repo, and what you must rebuild

The tokenized data pools. The loader replays a deterministic block order over pre-tokenized memmaps, too large to host here:

pool	size
c4_train.bin	80.2 GB
code_train.bin	14.0 GB
c4_val.bin, code_val.bin	~30 MB each

Rebuild with code/prepare_data.py, which pulls C4 and StarcoderData and writes the .bin/.json pair the loader expects. Pre-tokenized pools for the smaller sweeps are published at Impliedhomeland/midtrain-bridge-data.

Exact reproducibility caveat, stated plainly. RNG state is restored, but the data order is reconstructed from c4_base_seed: 1 and code_interleave_seed: 0 over your rebuilt pool. If your tokenization differs at all (tokenizer version, shard order, document filtering), the resumed run consumes different tokens from the original. The optimizer trajectory continues correctly, but the data sequence is not guaranteed identical across a rebuild. Byte-identical continuation requires the original .bin files, which are not published.

Intended use

A research artifact for continual-pretraining and data-mixing work, and the starting point for the fork-based experiment design it belongs to. It is not instruction-tuned, aligned, or safety-filtered, and it has had no post-training of any kind. Even the fully-decayed 24.6B final is a 410M base model trained on 24.6B tokens, so its generations are correspondingly limited; the earlier snapshots are mid-schedule and weaker still. Inherits the licensing and content characteristics of C4.