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

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.

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.