Skip to content

Wire max_checkpoints through SFT, DPO, and GRPO paths #1701

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
TimDettmers wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Wire max_checkpoints through SFT, DPO, and GRPO paths #1701

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
dfb208b
Wire max_checkpoints through all OLMo-core training paths
TimDettmers May 27, 2026
da208c8
Address review feedback: centralize -1→None mapping, remove duplicate…
TimDettmers May 27, 2026
667f149
Merge remote-tracking branch 'origin/main' into timd/wire-max-checkpo…
mnoukhov Jul 11, 2026
2f6bdf5
Bump OLMo-core pin to fa6c5014 (adds max_checkpoints)
mnoukhov Jul 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ All notable changes to this project will be documented in this file.
- Drop stale async rollout results whose generating policy is more than `async_steps` behind the trainer (`max_result_age_steps`), replenishing a fresh prompt and logging a `stale_results_dropped` metric (https://github.com/allenai/open-instruct/pull/1738).

### Changed
- 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).
- 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).
- Bump OLMo-core to the latest `main` commit (`9aa3280`) (https://github.com/allenai/open-instruct/pull/1723).
Expand Down
1 change: 1 addition & 0 deletions open_instruct/dpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
4 changes: 3 additions & 1 deletion open_instruct/grpo_olmo_core_actor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
1 change: 1 addition & 0 deletions open_instruct/olmo_core_finetune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
19 changes: 15 additions & 4 deletions open_instruct/olmo_core_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Comment on lines +243 to 258

@gemini-code-assist gemini-code-assist Bot May 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Instead of repeating the ternary operator args.keep_last_n_checkpoints if args.keep_last_n_checkpoints >= 0 else None across multiple caller files (dpo.py, grpo_olmo_core_actor.py, olmo_core_finetune.py), we can centralize this mapping logic inside build_checkpointer_callback. This keeps the callers clean and ensures consistent handling of the -1 (unlimited) convention.

def build_checkpointer_callback(
    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."""
    return CheckpointerCallback(
        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,
    )



Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,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", marker = "platform_system == 'Linux' and platform_machine == 'x86_64'" },
]
Expand Down
22 changes: 11 additions & 11 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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@76c239cab3ef3774aa1e3d0e35d1c0fc70e665c7 ; (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
Expand Down Expand Up @@ -250,11 +250,11 @@ fastapi==0.128.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') o
# open-instruct
# openenv-core
# vllm
fastapi-cli==0.0.20 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
fastapi-cli==0.0.20 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via fastapi
fastapi-cloud-cli==0.11.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
fastapi-cloud-cli==0.11.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via fastapi-cli
fastar==0.8.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
fastar==0.8.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via fastapi-cloud-cli
fastavro==1.12.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via cohere
Expand Down Expand Up @@ -358,7 +358,7 @@ hpack==4.1.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (p
# via h2
httpcore==1.0.9 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via httpx
httptools==0.7.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
httptools==0.7.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via uvicorn
httpx==0.28.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via
Expand Down Expand Up @@ -881,7 +881,7 @@ pyasn1-modules==0.4.2 ; (platform_machine == 'aarch64' and sys_platform == 'linu
# via google-auth
pybase64==1.4.3 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
# via vllm
pycountry==24.6.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
pycountry==24.6.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via pydantic-extra-types
pycparser==2.23 ; (implementation_name != 'PyPy' and platform_machine == 'aarch64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux') or (implementation_name == 'pypy' and platform_machine == 'aarch64' and platform_python_implementation == 'PyPy' and sys_platform == 'linux') or (implementation_name != 'PyPy' and platform_machine == 'x86_64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux') or (implementation_name == 'pypy' and platform_machine == 'x86_64' and platform_python_implementation == 'PyPy' and sys_platform == 'linux') or (implementation_name != 'PyPy' and platform_python_implementation != 'PyPy' and sys_platform == 'darwin')
# via cffi
Expand Down Expand Up @@ -918,7 +918,7 @@ pydantic-core==2.41.5 ; (platform_machine == 'aarch64' and sys_platform == 'linu
# via
# cohere
# pydantic
pydantic-extra-types==2.11.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
pydantic-extra-types==2.11.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via
# fastapi
# mistral-common
Expand Down Expand Up @@ -1063,11 +1063,11 @@ rich==13.9.4 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (p
# typer
rich-rst==1.3.2 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via cyclopts
rich-toolkit==0.17.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
rich-toolkit==0.17.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via
# fastapi-cli
# fastapi-cloud-cli
rignore==0.7.6 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
rignore==0.7.6 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via fastapi-cloud-cli
rpds-py==0.30.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via
Expand Down Expand Up @@ -1338,7 +1338,7 @@ uvicorn==0.40.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or
# mcp
# open-instruct
# openenv-core
uvloop==0.22.1 ; (platform_machine == 'aarch64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux')
uvloop==0.22.1 ; (platform_machine == 'aarch64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux') or (platform_python_implementation != 'PyPy' and sys_platform == 'darwin')
# via uvicorn
virtualenv==20.36.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via
Expand All @@ -1350,7 +1350,7 @@ wandb==0.23.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (
# via open-instruct
watchdog==6.0.0 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via mkdocs
watchfiles==1.1.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
watchfiles==1.1.1 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'
# via
# uvicorn
# vllm
Expand Down
Loading
Loading