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

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

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.