Wire max_checkpoints through SFT, DPO, and GRPO paths

[TimDettmers]

Summary

• Passes keep_last_n_checkpoints (default: 3, from CheckpointConfig) to CheckpointerCallback.max_checkpoints across all OLMo-core training paths (SFT, DPO, GRPO)
• Centralizes the -1 → None (unlimited) convention inside build_checkpointer_callback so call sites pass the value directly
• GRPOExperimentConfig already inherits keep_last_n_checkpoints from CheckpointConfig — no new config fields needed

Depends on

• Add max_checkpoints to limit permanent checkpoint retention OLMo-core#694 (adds max_checkpoints to CheckpointerCallback)

Files changed

• olmo_core_utils.py — build_checkpointer_callback and build_base_callbacks accept and pass through max_checkpoints; -1 mapped to None inside the builder
• dpo.py — passes args.keep_last_n_checkpoints through
• olmo_core_finetune.py — passes args.checkpoint.keep_last_n_checkpoints through
• grpo_olmo_core_actor.py — passes self.grpo_config.keep_last_n_checkpoints through

Test plan

• Verify SFT training with default keep_last_n_checkpoints=3 trims checkpoints
• Verify GRPO training with --keep_last_n_checkpoints -1 keeps all checkpoints
• Verify DPO training passes the parameter correctly
• Existing tests pass unchanged (new parameter has defaults matching current behavior)

[gemini-code-assist]

Code Review

This pull request introduces support for limiting the number of saved checkpoints via a new keep_last_n_checkpoints configuration option, which is integrated across DPO, GRPO, and SFT training pipelines. The feedback suggests removing a redundant field definition in GRPOExperimentConfig since it is already inherited, and centralizing the mapping logic that converts negative values (representing unlimited checkpoints) to None inside build_checkpointer_callback to simplify and clean up the caller sites.

[gemini-code-assist]

GRPOExperimentConfig inherits from olmo_core_utils.CheckpointConfig (via multiple inheritance), which already defines the keep_last_n_checkpoints field with a default value of 3. Redefining it here is redundant and can be safely removed to keep the configuration clean and avoid duplication.

[gemini-code-assist]

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,
    )

[gemini-code-assist]

With the -1 to None mapping logic centralized inside build_checkpointer_callback, we can simplify this call by passing args.keep_last_n_checkpoints directly.

Suggested change

	max_checkpoints=args.keep_last_n_checkpoints if args.keep_last_n_checkpoints >= 0 else None,
	max_checkpoints=args.keep_last_n_checkpoints,

[gemini-code-assist]

With the -1 to None mapping logic centralized inside build_checkpointer_callback, we can simplify this call by passing self.grpo_config.keep_last_n_checkpoints directly.

Suggested change

	max_checkpoints=self.grpo_config.keep_last_n_checkpoints if self.grpo_config.keep_last_n_checkpoints >= 0 else None,
	max_checkpoints=self.grpo_config.keep_last_n_checkpoints,

[gemini-code-assist]

With the -1 to None mapping logic centralized inside build_checkpointer_callback, we can simplify this call by passing args.checkpoint.keep_last_n_checkpoints directly.

Suggested change

	max_checkpoints=args.checkpoint.keep_last_n_checkpoints if args.checkpoint.keep_last_n_checkpoints >= 0 else None,
	max_checkpoints=args.checkpoint.keep_last_n_checkpoints,

[mnoukhov]

Looks good to me. I guess wait until allenai/OLMo-core#694 has been merged and update olmo-core version in pyproject.toml to include it.

Also consider keeping the name keep_last_n_checkpoints as that feels a bit clearer than max_checkpoints