---
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.