Impliedhomeland/midtrain-bridge-1B-wsd-backbone

midtrain-bridge: Pythia-1B / 60B WSD C4 backbone

A C4-only 1B-parameter pretraining run to 60B tokens under a Warmup-Stable-Decay schedule, published with 13 checkpoints along the trajectory. This is the W=0 reference arm of a study on when to introduce a new data distribution (code) during pretraining; the checkpoints are the fork points from which code-mixed branches are launched.

It is the WSD twin of midtrain-bridge-1B-cosine-backbone, which is the same model and data under a cosine schedule.

The run

architecture	pythia-1b (GPT-NeoX), 1,011,781,632 parameters
tokenizer	Pythia / GPT-NeoX, padded vocab 50277
data	C4 (en) only, no code, no math
tokens	60B (30,518 steps)
sequence length	2048
global batch	960 sequences = 1,966,080 tokens/step
LR schedule	WSD: peak 3e-4 → min 1e-6
warmup	10% of the horizon = 0 → 6B, linear to peak
stable	6B → 54B, held at 3e-4
decay	54B → 60B, linear 3e-4 → 1e-6 (decay_start_frac 0.9)
optimizer	AdamW
precision	bf16 autocast, fp32 master weights

The LR is keyed to absolute token count, not step index, so a branch forked from any checkpoint continues the parent schedule with no re-warmup.

The 0 → 6B warmup is NOT in this repo

This tree has no trunk of its own. Because both schedules use the same peak (3e-4) and the same warmup_frac (0.1), the warmup ramp lr = peak * t / warmup_tokens is identical for cosine and WSD, so a cosine trunk checkpoint taken exactly at the warmup endpoint is a valid WSD branch point. This run therefore starts from the cosine tree's 6B checkpoint:

trunk/trunk_branchpoint.pt in midtrain-bridge-1B-cosine-backbone

Same seed (1), same data order, same global batch. To reproduce this run from scratch you need that file as the starting point.

Checkpoints

branch_A/ covers 6B → 60B. The phase column marks where each snapshot sits in the schedule.

file	tokens	step	LR there	phase
branch_A_12.00B_step6104.pt	12.00B	6104	3e-4	stable
branch_A_18.00B_step9156.pt	18.00B	9156	3e-4	stable
branch_A_24.00B_step12208.pt	24.00B	12208	3e-4	stable
branch_A_30.00B_step15259.pt	30.00B	15259	3e-4	stable
branch_A_36.00B_step18311.pt	36.00B	18311	3e-4	stable
branch_A_42.00B_step21363.pt	42.00B	21363	3e-4	stable
branch_A_45.00B_step22889.pt	45.00B	22889	3e-4	stable
branch_A_48.00B_step24415.pt	48.00B	24415	3e-4	stable
branch_A_51.00B_step25940.pt	51.00B	25940	3e-4	stable
branch_A_54.00B_step27466.pt	54.00B	27466	3e-4	decay start
branch_A_57.00B_step28992.pt	57.00B	28992	~1.5e-4	decay
branch_A_58.50B_step29755.pt	58.50B	29755	~7.6e-5	decay
branch_A_60.00B_step30518.pt	60.00B	30518	1e-6	final

The grid is deliberately denser late (45, 51, 54, 57, 58.5B) because the interesting behaviour in WSD sits around and inside the decay phase.

Every file here carries full AdamW optimizer state (exp_avg, exp_avg_sq), including the 60B final, so any of them can be continued from or forked. That differs from the cosine repo, whose 60B final is weights-only.

File format

Each .pt is a torch.save dict:

{
  "model":            state_dict,        # litgpt GPT, GPT-NeoX layout
  "optimizer":        state_dict,        # AdamW
  "completed_steps":  int,
  "global_tokens":    int,               # absolute token count, keys the LR schedule
  "config":           dict,              # full run config
  "torch_rng": ..., "numpy_rng": ...,
  "val_c4": float, "val_code": float,    # held-out losses at that snapshot
}

import torch
ck = torch.load("branch_A/branch_A_60.00B_step30518.pt", map_location="cpu", weights_only=False)
print(ck["global_tokens"], ck["val_c4"])
state = ck["model"]          # GPT-NeoX parameter layout

These are litgpt-format state dicts, not transformers checkpoints, so AutoModelForCausalLM.from_pretrained will not read them directly. The parameter layout is standard GPT-NeoX and converts mechanically.

Data

C4 (en), pre-tokenized, drawn from a 60.1B-token pool: the 40.1B pool published at Impliedhomeland/midtrain-bridge-data (pythia-70m/c4/, which serves the whole Pythia suite since all sizes share one tokenizer) concatenated with 20.0B disjoint tokens from later C4 shards. Blocks are consumed in a fixed seed-1 permutation, identical to the cosine twin.

Intended use

Released so the intro-timing experiments built on these fork points can be reproduced, and as a set of intermediate checkpoints along a single well-specified 1B WSD run. This is a base model trained only on C4 with no instruction tuning, no safety filtering beyond C4's own, and no alignment work of any kind. Outputs will reflect whatever is in C4.