| --- |
| license: apache-2.0 |
| library_name: pytorch |
| tags: |
| - pretraining |
| - midtraining |
| - checkpoints |
| - pythia |
| datasets: |
| - Impliedhomeland/midtrain-bridge-data |
| --- |
| |
| # 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`](https://huggingface.co/datasets/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: |
| |
| ```python |
| 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 |
|
|
| ```bash |
| # 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. |
|
|
