Merge branch 'main' into publish-auto-cuda13

Author: mnoukhov

CHANGELOG.md

@@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file.
 
 ### Changed
 - Automatically publish stable CUDA 12 and CUDA 13 Beaker image aliases after merge-queue integration tests pass (https://github.com/allenai/open-instruct/pull/1783).
+- Wire `keep_last_n_checkpoints` through `build_checkpointer_callback` and `build_base_callbacks` to OLMo-core's new `max_checkpoints` parameter across SFT, DPO, and GRPO training paths; bump OLMo-core to the commit that added `max_checkpoints` (`fa6c501`). Negative values (e.g. `-1`) mean unlimited (https://github.com/allenai/open-instruct/pull/1701).
 - Add selectable CUDA 12.8 and CUDA 13.0 Docker builds, including matching torch, vLLM, and flash-attention dependency variants, and add B300 support on the new `ai2/holmes` cluster (https://github.com/allenai/open-instruct/pull/1758).
 - Increase default environment pool acquire timeout to 7200s (https://github.com/allenai/open-instruct/pull/1729).
 - Change the default generation `temperature` to 1.0 and make `SamplingConfig.temperature` a required field so `StreamingConfig.temperature` is the single source of truth (https://github.com/allenai/open-instruct/pull/1725).

open_instruct/dpo.py

@@ -72,6 +72,7 @@ def _setup_callbacks(args: dpo_utils.DPOExperimentConfig, dp_world_size: int):
         wandb_project=args.wandb_project,
         wandb_entity=args.wandb_entity,
         save_async=False,
+        max_checkpoints=args.keep_last_n_checkpoints,
     )
     slack_webhook_url = os.environ.get("SLACK_WEBHOOK_URL")
     if args.send_slack_alerts and slack_webhook_url:

open_instruct/grpo_olmo_core_actor.py

@@ -377,7 +377,9 @@ def fit(self) -> dict:
 
         if self.grpo_config.checkpoint_state_freq > 0:
             trainer_callbacks["checkpointer"] = olmo_core_utils.build_checkpointer_callback(
-                checkpointing_steps=self.grpo_config.checkpoint_state_freq, ephemeral_save_interval=None
+                checkpointing_steps=self.grpo_config.checkpoint_state_freq,
+                ephemeral_save_interval=None,
+                max_checkpoints=self.grpo_config.keep_last_n_checkpoints,
             )
         trainer_callbacks["data_prep_state"] = grpo_callbacks_lib.DataPreparationActorCheckpointCallback()
 

open_instruct/olmo_core_finetune.py

@@ -288,6 +288,7 @@ def main(args: SFTArguments, tc: dataset_transformation.TokenizerConfig) -> None
         with_tracking=args.logging.with_tracking,
         wandb_project=args.logging.wandb_project,
         wandb_entity=args.logging.wandb_entity or "ai2-llm",
+        max_checkpoints=args.checkpoint.keep_last_n_checkpoints,
     )
     trainer_callbacks["config_saver"] = callbacks.ConfigSaverCallback(_config=config_dict)
     trainer_callbacks["garbage_collector"] = callbacks.GarbageCollectorCallback()

open_instruct/olmo_core_utils.py

@@ -240,11 +240,21 @@ class CheckpointConfig:
 
 
 def build_checkpointer_callback(
-    checkpointing_steps: int, ephemeral_save_interval: int | None, save_async: bool = True
+    checkpointing_steps: int,
+    ephemeral_save_interval: int | None,
+    save_async: bool = True,
+    max_checkpoints: int | None = 3,
 ) -> CheckpointerCallback:
-    """Construct a CheckpointerCallback with shared Open Instruct defaults."""
+    """Construct a CheckpointerCallback with shared Open Instruct defaults.
+
+    ``max_checkpoints`` accepts the open-instruct convention where ``-1`` means
+    unlimited.  Negative values are mapped to ``None`` (keep all).
+    """
     return CheckpointerCallback(
-        save_interval=checkpointing_steps, ephemeral_save_interval=ephemeral_save_interval, save_async=save_async
+        save_interval=checkpointing_steps,
+        ephemeral_save_interval=ephemeral_save_interval,
+        save_async=save_async,
+        max_checkpoints=max_checkpoints if max_checkpoints is not None and max_checkpoints >= 0 else None,
     )
 
 
@@ -294,13 +304,14 @@ def build_base_callbacks(
     wandb_project: str | None = None,
     wandb_entity: str | None = None,
     save_async: bool = True,
+    max_checkpoints: int | None = 3,
 ) -> dict[str, Any]:
     """Build the callbacks shared across SFT and DPO: beaker, gpu monitor, checkpointer, and optional wandb."""
     result: dict[str, Any] = {
         "beaker": olmo_core_callbacks.BeakerCallbackV2(config=config_dict),
         "gpu_monitor": train_callbacks.GPUMemoryMonitorCallback(),
         "checkpointer": build_checkpointer_callback(
-            checkpointing_steps, ephemeral_save_interval, save_async=save_async
+            checkpointing_steps, ephemeral_save_interval, save_async=save_async, max_checkpoints=max_checkpoints
         ),
     }
     if with_tracking and wandb_project:

pyproject.toml

@@ -71,7 +71,7 @@ flash-attn-3 = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE" }
 
 # pytorch related setups
 [tool.uv.sources]
-ai2-olmo-core = { git = "https://github.com/allenai/OLMo-core.git", rev = "9aa3280fd5331b93fc766206eadec42ee73befc0" }
+ai2-olmo-core = { git = "https://github.com/allenai/OLMo-core.git", rev = "fa6c5014c9f6e9ee789da2d9c20d5126fee8df0d" }
 flash-attn = [
   { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.0/flash_attn-2.8.3+cu128torch2.10-cp312-cp312-linux_x86_64.whl", group = "cuda12" },
   { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.0/flash_attn-2.8.3+cu130torch2.10-cp312-cp312-linux_x86_64.whl", group = "cuda13" },

requirements.txt

@@ -6,7 +6,7 @@ accelerate==1.12.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux')
     # via
     #   open-instruct
     #   peft
-ai2-olmo-core @ git+https://github.com/allenai/OLMo-core.git@9aa3280fd5331b93fc766206eadec42ee73befc0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
+ai2-olmo-core @ git+https://github.com/allenai/OLMo-core.git@fa6c5014c9f6e9ee789da2d9c20d5126fee8df0d ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
     # via open-instruct
 aiofiles==25.1.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
     # via crawl4ai