Resuming training from a checkpoint
This page answers a question that came up after a partial pretrain run reached checkpoints/run-2026-04-24_18-53-41/:
I see you don't include the progress of training? Is that okay? Or do I need to redo all the datasets on the model again?
Short answer: no, you do not redo any datasets. Tokenisation is cached on disk under data/cache/ (mmap'd binary shards since v0.1.4); reloading a multi-million-conversation corpus is a few seconds even after a kernel restart. What gets reused vs. rebuilt when you resume is laid out in detail below.
What's in a checkpoint
Each *.pt file under a run-<ts>/ directory carries:
| Field | What it is |
|---|---|
model_state_dict | every learned weight up to the moment of save |
optimizer_state_dict | AdamW first / second moments, step counter |
model_cfg | the dict the run was instantiated with (dim, layers, expert count, …) |
stage | "pretrain" / "sft" / "dpo" / "final" |
save_count | the per-mgr serial number (4-digit zero-padded suffix in the filename) |
epoch, epochs, epoch_complete, <stage>_loss, remaining_hours | whatever the caller passed via **extra |
What is not in a checkpoint:
- DataLoader iterator state (the within-epoch batch index).
- Python / NumPy / Torch RNG state at save time.
- Tokeniser (saved separately as
tokenizer.jsonnext to the.ptfiles; tiktoken makes this near-free anyway).
The practical consequence: resuming from pretrain-22-54-19-0008.pt puts the model exactly where the save left it, but the next pretrain epoch starts at batch 0 again. You do not lose any of the learned progress (every gradient step that landed before the save is in the weights and the optimiser moments). You do replay through the dataset, but every batch is a useful gradient step regardless of where in the epoch it lives.
What gets re-run when you resume
| Notebook stage | Re-runs? | Time cost (real corpus) | Reason |
|---|---|---|---|
download_data.py | no | 0 | datasets are saved to data/<name>/ once |
| Cell 1 — imports | yes | <5 s | pure Python imports |
Cell 3 — load_all_datasets | yes | seconds | load_from_disk over the saved Arrow shards |
Cell 4 — upsample_by_weight | yes | tens of seconds | rebuilds train_texts in RAM |
Cell 6 — build_tokenizer | yes | <1 s | tiktoken loads from the wheel |
Cell 12 — PretrainDataset | yes | <1 s | mmap fast-path on pretrain_tokens.bin |
Cell 15 — SFTDataset | yes | <1 s | mmap fast-path on sft_*.bin |
Cell 19 — DPODataset | yes | <1 s | mmap fast-path on dpo_*.bin |
| Cells 13 / 16 / 20 — training loops | after load_state_dict | full epochs | ← the work you're trying to preserve |
So a clean resume is essentially: re-run the cheap setup (seconds–minutes total), then before kicking the training loop, patch the model + optimiser with the checkpoint state.
Resuming run-2026-04-24_18-53-41 step-by-step
The latest snapshot in that directory at the time of writing:
checkpoints/run-2026-04-24_18-53-41/
├── pretrain-23-55-59-0058.pt
├── pretrain-00-26-01-0059.pt
├── pretrain-00-56-03-0060.pt ← latest (use this)
└── tokenizer.jsonsave_count=0060 plus pretrain stage means roughly 30 hours into pretrain epoch 1 with a 30 minute auto-save cadence. Resume from pretrain-00-56-03-0060.pt.
1. Open the notebook on the uv kernel
uv sync # idempotent
uv run jupyter lab full_llm_pipeline.ipynbIn VSCode / JupyterLab, pick the Lavender (uv .venv) kernel.
2. Run setup through cell 12 normally
These set up the env, load the cached data + tokenizer, and instantiate a fresh model = MoETransformerLM(**MODEL_CFG). The CheckpointManager cell creates a new run subdirectory (checkpoints/run-<now>/); that's intentional — the old run's files are read-only references, the new run's writes go to a new folder.
Cell 13 now contains the pretrain resume hook directly before the loop starts. Leave RESUME_FROM = None and RESUME_STAGE = "pretrain" to pick the newest pretrain checkpoint, or set RESUME_FROM to a specific .pt file / run-<ts>/ directory.
3. Resume before the pretrain loop
The built-in hook uses resume_from_checkpoint (added v0.2.5) to load weights, optimizer state, and metadata before pretrain_epoch runs. The longer form here is what to write when you want to pin a specific checkpoint or take fine-grained control of the optimizer and scheduler:
from lavender_2.checkpoints import resume_from_checkpoint
RESUME_FROM = "checkpoints/run-2026-04-24_18-53-41/pretrain-00-56-03-0060.pt"
# 1. Pull the saved weights + optimizer state + meta off disk.
model_state, opt_state, meta = resume_from_checkpoint(
RESUME_FROM, device=str(device),
)
model.load_state_dict(model_state)
# 2. Build the optimizer the same way cell 13 does it, then patch in
# the saved state. Order matters: optimizer construction must come
# before load_state_dict because Adam allocates state on first
# .step().
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
if opt_state is not None:
optimizer.load_state_dict(opt_state)
# 3. Seed the cosine scheduler. Only completed epochs should advance
# it; mid-epoch periodic checkpoints replay the current epoch from
# batch 0 because DataLoader iterator state is not saved.
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=PRETRAIN_EPOCHS)
completed_epoch = int(meta.get("epoch") or 0)
start_epoch = completed_epoch + 1 if meta.get("epoch_complete") else max(1, completed_epoch)
for _ in range(max(0, start_epoch - 1)):
scheduler.step()
# 4. Wire the optimizer into the new ckpt_mgr so subsequent
# saves carry the right state.
ckpt_mgr.set_optimizer(optimizer)
print(f"resumed from {meta['checkpoint_path']}")
print(f" stage : {meta.get('stage')}")
print(f" epoch : {meta.get('epoch')} / {meta.get('epochs')}")
print(f" ce_loss : {meta.get('ce_loss')}")
print(f" complete : {meta.get('epoch_complete')}")
print(f" save_count : {meta.get('save_count')}")To grab the newest checkpoint of a specific stage automatically, omit RESUME_FROM and pass stage="pretrain" (or "sft" / "dpo"):
model_state, opt_state, meta = resume_from_checkpoint(stage="pretrain")4. Let the loop choose the restart epoch
Cell 13 sets pretrain_start_epoch from checkpoint metadata: epoch_complete=True advances to the next epoch; missing or false epoch_complete replays the saved epoch from batch 0. Older checkpoints do not have epoch_complete, so they take the safer replay path.
If you intentionally want to skip ahead anyway:
RESUMED_EPOCH = resume_meta.get("epoch", 0)
for epoch in range(RESUMED_EPOCH + 1, PRETRAIN_EPOCHS + 1):
ce, aux = pretrain_epoch(...)
...The "complete the in-flight epoch" path is usually what you want — it processes data the model has not yet seen this epoch (the shuffle is reset, but in expectation half of the training data is fresh).
5. Run cells 13 → end as normal
CheckpointManager.maybe_save will now write to the new run directory every 30 minutes; the old run-2026-04-24_18-53-41/ directory is left untouched.
Caveats
- Within-epoch progress is not restored. If you stop mid- epoch, you'll replay through that epoch's batches from index 0. This is a known gap; properly resuming mid-epoch would require saving the dataloader iterator state and the per-rank shuffle seed, which v0.x doesn't.
- Cosine LR scheduler is only approximately restored. Stepping the scheduler
epochtimes after construction reproduces the per-epoch LR boundary, not the per-batch LR position within an epoch. ce_lossreported on resume is the rolling average from the saved batch onwards if you restart in-epoch — the model's effective loss should match the checkpoint'sce_lossfield within a few tenths within a few hundred batches.
For a real cross-machine resume (different hardware, different Python build), also uv sync first to make sure the package versions match what the checkpoint was saved against.