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:

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:

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

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

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.

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:

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:

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.