Impliedhomeland/midtrain-bridge-backbones

Midtrain-Bridge: C4-only backbone checkpoints (branch_A)

Training-state checkpoints for the C4-only reference runs used in midtraining intro-timing experiments. In that setup a single C4-only trunk is trained, then forked at various points to introduce a new data mixture (code). branch_A is the control: the run that never sees code at all.

These are not one model each. Each branch_A is a trajectory of fork points: every intro-timing branch in the experiments resumes from one of these snapshots. Publishing the whole trajectory is what makes those experiments reproducible.

Matching pre-tokenized data: Impliedhomeland/midtrain-bridge-data (pythia-70m/ tree).

Runs

folder	params	schedule	peak LR	min LR	branch_A range	extensions
cosine/70M_14B_1e-3_10pct	70M	cosine	1e-3	1e-4 (10%)	2.80B - 12.60B	_ext 13.30B, _ext2 13.65B
cosine/160M_14B_6e-4_10pct	160M	cosine	6e-4	6e-5 (10%)	2.10B - 14.00B	_ext975 13.65B
WSD/70M_14B_1e-3	70M	WSD	1e-3	-	2.80B - 14.00B	_ext 13.65B

All runs target a 14B-token budget on C4, Pythia tokenizer.

Note on cosine/70M_14B_1e-3_10pct: its branch_A stops at 12.60B, not 14.00B. The _ext and _ext2 runs continue it to 13.30B and 13.65B to serve the late fork points, but no checkpoint in that tree reaches 14.00B. The other two runs do.

The _ext* directories are continuations past where the main branch_A ended, created to supply late fork points (roughly 95% and 97.5% of the budget). They are separate runs on disk, not extra files inside branch_A.

What a checkpoint contains

Full training state, not just weights, so you can resume or fork:

import torch
ck = torch.load("branch_A_14.00B_step6676.pt", map_location="cpu", weights_only=False)
ck.keys()
# model, optimizer, completed_steps, global_tokens, config, torch_rng, numpy_rng, val_c4, val_code
ck["global_tokens"]     # 14_000_000_000
ck["completed_steps"]   # 6676

key	70M	160M	notes
model	0.282 GB	0.64 GB	fp32 weights
optimizer	0.563 GB	1.31 GB	Adam moments, two thirds of the file
torch_rng / numpy_rng	tiny	tiny	exact data-order reproducibility
completed_steps, global_tokens, config, val_c4, val_code	tiny	tiny	run metadata

Weights-only use is fine (ck["model"]), but dropping optimizer makes the checkpoint unusable for resuming or forking new branches.

Each folder also ships resolved_config.json (the fully-resolved training config), metrics.jsonl (per-eval loss trace) and the source .yaml.

latest.pt was not uploaded: in every run it is byte-identical to the last named checkpoint in the same folder.

Naming

branch_A_<tokens>B_step<N>.pt

<tokens> is cumulative tokens seen, <N> the optimizer step. Fork points are addressed by token count, so branch_A_8.40B_step4006.pt is the state a branch introducing code at 8.4B would resume from.

Downloading

# one run
hf download Impliedhomeland/midtrain-bridge-backbones \
    --include 'WSD/70M_14B_1e-3/branch_A/*' --local-dir ./ckpt

# a single fork point
hf download Impliedhomeland/midtrain-bridge-backbones \
    --include '*/branch_A/branch_A_14.00B_step6676.pt' --local-dir ./ckpt

Caveats

• fp32 throughout; no safetensors conversion, these are raw torch.save training states.
• Loading requires weights_only=False since the payload includes RNG and config objects. Only load checkpoints you trust.
• The model class is the project's own GPT implementation (src/model.py), not a transformers architecture. config in the checkpoint records the hyperparameters.