feat: unified resume system with checkpoint_utils

[Hecate0821]

Summary

Replaces the dual-path resume logic (ResumeConfig + setup_resume regex parsing) with a single source of truth: checkpoints.jsonl (tinker_cookbook format) via a rewritten checkpoint_utils module. Also adopts SupervisedDatasetFromHFDataset for SFT and adds RLPromptDataset for batch-indexed RL iteration.

What we import from tinker_cookbook

Import	Used for
get_last_checkpoint(log_dir)	Read last checkpoint from checkpoints.jsonl at resume time
CHECKPOINTS_BASE_NAME	Filename constant "checkpoints.jsonl"
SupervisedDatasetFromHFDataset	Batch-indexed SFT dataset with get_batch(i) + set_epoch(seed)
renderers (already used)	No change

What we implement ourselves (and why)

Function	Why not from tinker_cookbook
save_checkpoint	tinker_cookbook's save_checkpoint_async stores logical names as state_path; Fireworks cross-job resume needs cross_job:// refs resolved via resolve_checkpoint_path at save time
resolve_resume	tinker_cookbook creates a new TrainingClient from checkpoint; Fireworks RLOR loads into an existing client bound to a trainer job via load_state_with_optimizer
Both above (sync)	save_checkpoint_async requires async context; our RL callbacks (finish_step) run sync via asyncio.to_thread; tinker_cookbook's sync wrapper uses asyncio.run() which fails inside an existing event loop
RLPromptDataset	tinker_cookbook's RLDataset.get_batch() returns EnvGroupBuilder (RL gymnasium types we don't use); ours returns raw row dicts

Checkpoint format

checkpoints.jsonl — append-only JSONL in log_path. Each line written after every DCP save:

{"name": "step-2", "step": 2, "data_consumed": 16, "source_job_id": "abc123", "state_path": "cross_job://abc123/step-2"}

The state_path is a cross_job:// reference resolved at save time, so any future trainer job can load it directly.

Key changes

• checkpoint_utils.py (rewritten, 115 lines): ResumeInfo, resolve_resume(), save_checkpoint(). No validation helpers (removed dataset_fingerprint, validate_dataset, validate_training_shape).
• resume.py (deleted): old setup_resume + ResumeConfig removed
• sft_loop.py: Uses SupervisedDatasetFromHFDataset for batch-indexed iteration with per-epoch reshuffling
• rl_loop.py: Uses RLPromptDataset for batch-indexed iteration; passes fw_api_key to ReconnectableClient (fixes shared dev auth)
• supervised.py: Removed backward-compat try/except ImportError fallbacks for datum_from_tokens_weights
• log_path: Required field (no default) on all recipe Configs — prevents checkpoint collisions
• pyproject.toml: tinker-cookbook pinned to d7fb423 (latest)

Files changed (38)

Category	Files
Checkpoint	checkpoint_utils.py (rewrite), resume.py (delete)
Config/infra	config.py, __init__.py, validation.py, client.py, infra.py
Recipes	rl_loop.py, sft_loop.py, dpo_loop.py, orpo_loop.py
Dataset	data.py (add RLPromptDataset), supervised.py (remove fallbacks)
RL losses	grpo.py, dapo.py, gspo.py, cispo.py, losses.py
Examples	train_frozen_lake.py, train_deepmath.py, train_sft.py
Tests	15 test files updated

E2E test results

Test	Result	Duration	Date
SFT resume	PASSED	22m27s	2026-03-10
GRPO resume	PASSED	32m32s	2026-03-10
Unit tests (208)	PASSED	5s	2026-03-10

SFT: Phase 1 trained 5 steps (loss 4.11 → 0.00001), saved to checkpoints.jsonl with cross_job:// refs. Phase 2 created a new job, loaded checkpoint, continued from step 5 to step 8. data_consumed increased from 40 to 64 (dataloader continued, not restarted).

GRPO: Phase 1 trained 3 steps with hotloading + deployment on qwen3-30b-a3b (MoE). Phase 2 reused the deployment, loaded cross-job checkpoint, continued training. data_consumed increased correctly.

Test plan

• 208 unit/smoke tests pass
• SFT e2e resume: cross-job checkpoint load + dataloader continuation verified
• GRPO e2e resume: cross-job checkpoint load + deployment reuse + dataloader continuation verified
• DPO e2e resume (checkpoint-only migration, not yet tested e2e)

[mayinghan]

maybe move this to the top

[mayinghan]

where is this path defined