feat: unified resume system with checkpoint_utils

.gitignore

@@ -295,3 +295,4 @@ cython_debug/
 .idea/
 wandb/
 *_dev.sh
+sft_logs/

training/examples/deepmath_rl/train_deepmath.py

@@ -279,6 +279,7 @@ def main():
     )
 
     config = rl_loop.Config(
+        log_path=args.trajectory_dir or "./deepmath_logs",
         base_model=args.base_model,
         dataset=args.dataset_path,
         learning_rate=args.learning_rate,

training/examples/frozen_lake/train_frozen_lake.py

@@ -53,12 +53,10 @@
     InfraConfig,
     WandBConfig,
     DeployConfig,
-    ResumeConfig,
     HotloadConfig,
     ReconnectableClient,
     wandb_log,
     setup_wandb,
-    setup_resume,
     wandb_finish,
     log_metrics_json,
     setup_deployment,
@@ -99,6 +97,8 @@
 
 @dataclass
 class FrozenLakeConfig:
+    log_path: str = "./frozen_lake_logs"
+
     base_model: str = "accounts/fireworks/models/qwen3-8b"
     tokenizer_model: str = "Qwen/Qwen3-8B"
 
@@ -545,7 +545,9 @@ def _make_job(label: str, precreated_id: str | None, **extra_kw):
         infra_boot_time = time.time() - _infra_start
         wandb_log({"train/step": 0, "infra/total_boot_time": infra_boot_time}, step=0)
 
-        step_offset, _ = setup_resume(policy, ResumeConfig())
+        from training.utils.checkpoint_utils import resolve_resume
+        resume_info = resolve_resume(policy, cfg.log_path)
+        step_offset = resume_info.step if resume_info else 0
         if hotload_cfg.hot_load_before_training and deploy_cfg.deployment_id:
             name = f"resume-{step_offset}-base" if step_offset > 0 else "step-0-base"
             weight_syncer.save_and_hotload(name, checkpoint_type="base")

training/examples/text2sql_sft/train_sft.py

@@ -72,6 +72,7 @@ def main():
     )
 
     config = sft_loop.Config(
+        log_path="./text2sql_logs",
         base_model=args.base_model,
         dataset=args.dataset_path,
         tokenizer_model=args.tokenizer_model,

training/pyproject.toml

@@ -7,8 +7,8 @@ name = "fireworks-training-cookbook"
 version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
-    "fireworks-ai>=1.0.0a39,<2",
-    "tinker-cookbook @ git+https://github.com/thinking-machines-lab/tinker-cookbook.git@934f0d9b2f53c3edff02cbf23ec6da8682047fa5",
+    "fireworks-ai>=1.0.0a40,<2",
+    "tinker-cookbook @ git+https://github.com/thinking-machines-lab/tinker-cookbook.git@d7fb423674c0bdd27c20a9b948b6091457874039",
     "eval-protocol>=0.3.23",
     "tqdm",
     "torch",

training/recipes/dpo_loop.py

@@ -39,12 +39,10 @@
     InfraConfig,
     WandBConfig,
     DeployConfig,
-    ResumeConfig,
     HotloadConfig,
     ReconnectableClient,
     wandb_log,
     setup_wandb,
-    setup_resume,
     wandb_finish,
     validate_config,
     log_metrics_json,
@@ -58,6 +56,7 @@
 )
 from fireworks.training.sdk.deployment import DEFAULT_DELTA_COMPRESSION
 from fireworks.training.sdk.weight_syncer import WeightSyncer
+from training.utils.checkpoint_utils import resolve_resume
 from training.utils.timer import timer, flush_timing
 
 logger = logging.getLogger(__name__)
@@ -96,6 +95,9 @@ async def gather_with_progress(
 
 @dataclass
 class Config:
+    log_path: str
+    """Directory for checkpoints and logs. Required, no default."""
+
     base_model: str = "accounts/fireworks/models/qwen3-8b"
     dataset: str = ""
     tokenizer_model: str = ""  # HuggingFace model name for client-side tokenization
@@ -120,7 +122,7 @@ class Config:
     deployment: DeployConfig = field(default_factory=DeployConfig)
     hotload: HotloadConfig = field(default_factory=lambda: HotloadConfig(hot_load_interval=0))
     wandb: WandBConfig = field(default_factory=lambda: WandBConfig(project="dpo-tinker"))
-    resume: ResumeConfig = field(default_factory=ResumeConfig)
+    init_from_checkpoint: str | None = None
 
 
 # ---------------------------------------------------------------------------
@@ -389,7 +391,7 @@ def _signal_handler(signum, frame):
     signal.signal(signal.SIGTERM, _signal_handler)
     signal.signal(signal.SIGINT, _signal_handler)
 
-    validate_config(cfg.base_model, cfg.dataset, cfg.hotload, cfg.deployment, cfg.infra, cfg.resume)
+    validate_config(cfg.base_model, cfg.dataset, cfg.hotload, cfg.deployment, cfg.infra)
     if not cfg.tokenizer_model:
         raise ValueError(
             "Config.tokenizer_model is required for client-side tokenization. "
@@ -485,7 +487,9 @@ def _signal_handler(signum, frame):
             compression_format=DEFAULT_DELTA_COMPRESSION,
         )
 
-        step_offset, _ = setup_resume(policy, cfg.resume)
+        resume_info = resolve_resume(policy, cfg.log_path, cfg.init_from_checkpoint)
+        step_offset = resume_info.step if resume_info else 0
+        wandb_log({"train/step": step_offset}, step_offset)
         adam_params = tinker.AdamParams(learning_rate=cfg.learning_rate, **DEFAULT_ADAM)
 
         # -- Cache reference logprobs concurrently ------------------------------
@@ -533,8 +537,10 @@ def _signal_handler(signum, frame):
         # -- Final checkpoint --------------------------------------------------
 
         hl = cfg.hotload
-        if step > step_offset and (hl.hot_load_interval > 0 or hl.dcp_save_interval > 0):
-            weight_syncer.save_and_hotload(f"final-step-{step}")
+        if step > step_offset:
+            weight_syncer.save_dcp(f"step-{step}")
+            if hl.hot_load_interval > 0:
+                weight_syncer.save_and_hotload(f"final-step-{step}")
 
         logger.info("Training complete: %d optimizer steps (%d new)", step, step - step_offset)
         return {"steps": step, "policy_job_id": policy_job_id, "reference_job_id": reference_job_id}
@@ -556,4 +562,4 @@ def _signal_handler(signum, frame):
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-    main(Config())
+    main(Config(log_path="./dpo_logs"))

training/recipes/orpo_loop.py

@@ -49,11 +49,9 @@
     DEFAULT_ADAM,
     InfraConfig,
     WandBConfig,
-    ResumeConfig,
     ReconnectableClient,
     wandb_log,
     setup_wandb,
-    setup_resume,
     wandb_finish,
     log_metrics_json,
     make_orpo_loss_fn,
@@ -63,6 +61,7 @@
     render_preference_pair,
     resolve_renderer_name,
 )
+from training.utils.checkpoint_utils import resolve_resume
 
 logger = logging.getLogger(__name__)
 
@@ -73,6 +72,9 @@
 
 @dataclass
 class Config:
+    log_path: str
+    """Directory for checkpoints and logs. Required, no default."""
+
     base_model: str = "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507"
     dataset: str = ""
     tokenizer_model: str = ""
@@ -94,7 +96,7 @@ class Config:
             project="dsv3-training",
         )
     )
-    resume: ResumeConfig = field(default_factory=ResumeConfig)
+    init_from_checkpoint: str | None = None
 
 
 # ---------------------------------------------------------------------------
@@ -175,7 +177,8 @@ def _signal_handler(signum, frame):
     )
 
     job_id = endpoint.job_id
-    step_offset, _ = setup_resume(client, cfg.resume)
+    resume_info = resolve_resume(client, cfg.log_path, cfg.init_from_checkpoint)
+    step_offset = resume_info.step if resume_info else 0
     adam_params = tinker.AdamParams(learning_rate=cfg.learning_rate, **DEFAULT_ADAM)
 
     # -- Data ----------------------------------------------------------------
@@ -325,8 +328,11 @@ def _signal_handler(signum, frame):
         # -- Final checkpoint ------------------------------------------------
 
         if step > step_offset:
+            logger.info("Saving final DCP checkpoint (step %d)...", step)
+            client.save_state(f"step-{step}")
+
             logger.info("Saving final base checkpoint (step %d)...", step)
-            result = client.inner.save_weights_for_sampler_ext(
+            result = client.save_weights_for_sampler_ext(
                 f"final-step-{step}", checkpoint_type="base"
             )
             logger.info("Final base checkpoint saved: %s", result.path)
@@ -353,7 +359,8 @@ def _signal_handler(signum, frame):
     )
 
     cfg = Config(
-        dataset=os.environ.get("ORPO_DATASET_PATH"), 
+        log_path="./orpo_logs",
+        dataset=os.environ.get("ORPO_DATASET_PATH"),
         tokenizer_model=os.environ.get("ORPO_TOKENIZER", "Qwen/Qwen3-235B-A22B-Instruct-2507"),
     )
     main(cfg)