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

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

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

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

The run was interrupted twice by compute-allocation cancellations and resumed from latest.pt each time (visible as three [resume] events in the training log, the last at step 10750). Because each checkpoint carries optimizer and RNG state, the resumes are continuations rather than restarts, and the loss curve shows no discontinuity across them.

Note on the wider study. This repo contains the backbone only. The code-mixing branches that fork from these checkpoints are a separate, still-running set of jobs and are not published here. val_code on this backbone is a zero-code-exposure baseline: it was trained on C4 alone, so its Starcoder loss reflects pure transfer, not code training.

What is here

Every file is a full PyTorch checkpoint containing weights + optimizer state + RNG state, so training continues bit-for-bit 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 58-field run config
torch_rng, numpy_rng	RNG states, so the data stream continues deterministically
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

The branch_A_*.pt files are the fork points; the "fork for" column names the branch each was snapshotted for, at that fraction of the 24.6B schedule. latest.pt is the crash-resume file and, now that the schedule has finished, holds the same step as the final checkpoint.

file	step	tokens	fraction	fork for
branch_A_3.69B_step1877.pt	1877	3.69B	15%	5.88%@15
branch_A_4.92B_step2503.pt	2503	4.92B	20%	6.25%@20
branch_A_6.15B_step3129.pt	3129	6.15B	25%	6.66%@25
branch_A_7.38B_step3754.pt	3754	7.38B	30%	7.14%@30
branch_A_8.61B_step4380.pt	4380	8.61B	35%	7.69%@35
branch_A_9.84B_step5005.pt	5005	9.84B	40%	8.33%@40
branch_A_11.07B_step5631.pt	5631	11.07B	45%	9.09%@45
branch_A_12.30B_step6257.pt	6257	12.30B	50%	10%@50
branch_A_14.76B_step7508.pt	7508	14.76B	60%	12.5%@60
branch_A_17.22B_step8759.pt	8759	17.22B	70%	16.6%@70
branch_A_19.68B_step10010.pt	10010	19.68B	80%	25%@80
branch_A_22.14B_step11261.pt	11261	22.14B	90%	50%@90
branch_A_23.37B_step11887.pt	11887	23.37B	95%	100%@95
branch_A_24.60B_step12513.pt	12513	24.60B	100%	final, fully decayed
latest.pt	12513	24.60B	-	identical to the final

branch_A_24.60B_step12513.pt is the one to use if you want the finished no-code 410M reference model. The earlier files are mid-schedule states, useful as fork points or for studying the trajectory, but they sit at higher LR and are not annealed.

Also included: dm_branch_A.yaml (run config), resolved_config.json, metrics.jsonl (full eval history), trunk_branchpoint.pt (the 2.4615B warmup-end root that branch_A itself forked from), 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	cosine, peak 3e-4 to min 3e-5 (a 10%-of-peak floor)
warmup	10% of 24.6B = 2.4615B, linear 0 to peak
batch	global 960 sequences, micro 16, so 1,966,080 tokens/step; end_step 12513
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 on H200 NVL, varying across resumes (4-GPU, then 2-GPU lanes, then 8-GPU)

World size varied because the run was resumed on different allocations. This is safe and does not affect the data seen: the loader keys block consumption on local_step and global_batch, not on world size, so a run checkpointed at one world size resumes correctly at another.

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 here are not comparable to 1024-batch runs with 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	val_c4	val_code
1877	3.69B	3.9448	4.7804
2503	4.92B	3.6992	4.2089
3129	6.15B	3.5561	3.9540
3754	7.38B	3.4554	3.8742
4380	8.61B	3.3768	3.8017
5005	9.84B	3.3205	3.7027
5631	11.07B	3.2726	3.5866
6257	12.30B	3.2363	3.5938
7508	14.76B	3.1766	3.4970
8759	17.22B	3.1358	3.4498
10010	19.68B	3.1058	3.4006
11261	22.14B	3.0844	3.3699
11887	23.37B	3.0776	3.3799
12513	24.60B	3.0727	3.3789

Full history (38 eval points) is in metrics.jsonl.

Using these checkpoints

The backbone schedule is finished, so there is nothing left to resume on branch_A itself. The two useful operations are loading the final model and forking a branch.

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

This is what the snapshots exist for. 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
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-cosine-24.6B-c4-backbone --include 'code/*' --local-dir .
torchrun --standalone --nproc_per_node=4 code/train.py --config your_branch.yaml

The LR schedule is keyed to absolute token count, so a branch forked at 19.68B continues the same cosine curve rather than restarting it. train.py also auto-resumes from latest.pt if one is present in ckpt_dir, restoring optimizer and RNG state.

code/launch.sh is included for reference only: it is a thin torchrun wrapper with absolute paths to the original cluster's conda env and repo, so use the torchrun line above instead.

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, which are far 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 them with code/prepare_data.py, which pulls C4 and StarcoderData and writes the .bin/.json pair the loader expects.

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 sequence of data 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.