fix(automodel): align checkpoint lifecycle with consolidated save modes

[jQizhang]

Summary

This PR updates NeMo RL's Automodel checkpoint integration to match the checkpoint lifecycle and consolidated-save modes introduced in Automodel r0.6.0.

• exposes the canonical save_consolidated modes: "false", "final", and "every"
• builds the Automodel Checkpointer once from the complete construction-time configuration instead of mutating or rebuilding it during each save
• propagates final-checkpoint intent explicitly so final consolidates only terminal or deliberate early-stop checkpoints, while timeout recovery checkpoints remain resumable
• completes the async save lifecycle before checkpoint promotion, source-state mutation, or the next save
• applies the lifecycle consistently across Policy/Value workers and all checkpointing algorithm callers

Fixes #3893. Follow-up to #3498 and related to #3624.

Why this is needed

Automodel r0.6.0 creates async stagers, dedicated process groups, consolidation groups, and save addons while constructing the Checkpointer. The previous NeMo RL integration initialized a checkpointer from partial/default configuration, then updated checkpoint configuration at save time. Changing constructor-owned settings after construction could leave the runtime resources out of sync with the effective configuration and cause missing consolidation or distributed hangs.

The same per-save configuration path also made load-time behavior depend on rank-local filesystem detection and did not carry Automodel's explicit final-checkpoint signal.

Before

flowchart LR
    A["Partial config C0"] --> B["Build Checkpointer<br/>resources R0"] --> C["Every save / load<br/>recomputes config"] --> D["Mutate to C1<br/>rebuild addons only"] --> E["Run with<br/>config C1 + resources R0"] --> F["C1 != R0<br/>missing resources or hang"]

    classDef config fill:#EDF3FC,stroke:#9FB5D6,color:#17213A
    classDef runtime fill:#FFF5DF,stroke:#DFB45D,color:#17213A
    classDef automodel fill:#F1EDFB,stroke:#AD9BD0,color:#17213A
    classDef risk fill:#FFF0EE,stroke:#D45B55,color:#7F302D
    class A config
    class C,D runtime
    class B,E automodel
    class F risk

Loading

After

flowchart LR
    A["Complete config C<br/>at worker construction"] --> B["CheckpointingConfig.build()"] --> C["Stable Checkpointer<br/>resources R(C)<br/>load reuses it"] --> D["async_wait()<br/>finish previous save"] --> E["save_model(final flag)<br/>wait for staging"] --> F["Background upload<br/>false / final / every"] --> G["finalize_async_save()<br/>tmp_step_N → step_N"]

    classDef config fill:#EDF3FC,stroke:#9FB5D6,color:#17213A
    classDef lifecycle fill:#EAF7F3,stroke:#7EB7A4,color:#185D4A
    classDef automodel fill:#F1EDFB,stroke:#AD9BD0,color:#17213A
    class A config
    class B,C automodel
    class D,E,F,G lifecycle

Loading

Implementation details

save_consolidated	Behavior
"false"	Do not automatically write consolidated Hugging Face weights
"final"	Consolidate only a terminal or deliberate early-stop checkpoint
"every"	Consolidate every checkpoint

Use quoted "false" in YAML; replace legacy false with "false" and true with "every". single_rank_consolidation remains supported for every mode.

• Construction: Policy and Value pass supported settings once to Automodel's CheckpointingConfig.build(); known unsupported settings are rejected. Per-save config mutation and private-addon rebuilding are removed.
• Save: Callers pass paths plus is_final_checkpoint (final step or deliberate early stop, but not timeout recovery). Async Policy saves wait for prior work and staging, then finish upload/consolidation before atomic promotion; PPO Value saves remain synchronous.
• Load: The same checkpointer is reused without rank-local format probing, runtime config mutation, or private-addon rebuilding.

Validation

Static checks

• ruff check passed for the modified Automodel manager and its unit tests
• Python compilation passed for the modified Automodel manager and its unit tests
• git diff --check passed

Distributed checkpoint matrix

Ran a 24-case end-to-end Slurm matrix covering four representative algorithms and the complete consolidated-save configuration cross product.

Environment

• 1 node with 8× H100 GPUs
• Llama 3.2 1B
• Automodel DTensor/FSDP2 backend
• model_save_format: safetensors
• two training steps per case with save_period: 1
• step_1: non-final checkpoint
• step_2: final checkpoint
• Automodel Policy checkpointing used async saving
• PPO additionally exercised its synchronous Value checkpoint path
• optimizer-state checkpointing was exercised in the GRPO / "false" / distributed-consolidation case

Results

Algorithm	single_rank_consolidation	"false"	final	every
GRPO	false	✅ PASS	✅ PASS	✅ PASS
PPO	false	✅ PASS	✅ PASS	✅ PASS
SFT	false	✅ PASS	✅ PASS	✅ PASS
DPO	false	✅ PASS	✅ PASS	✅ PASS
GRPO	true	✅ PASS	✅ PASS	✅ PASS
PPO	true	✅ PASS	✅ PASS	✅ PASS
SFT	true	✅ PASS	✅ PASS	✅ PASS
DPO	true	✅ PASS	✅ PASS	✅ PASS

Result: 24/24 cases passed.

Consolidation semantics

save_consolidated	Expected and observed result
"false"	Neither checkpoint contained a consolidated/ directory
"final"	Only step_2 contained consolidated Hugging Face weights
"every"	Both step_1 and step_2 contained consolidated Hugging Face weights

For PPO, these semantics were independently verified for both Policy and Value checkpoints. Every case also completed without leftover temporary checkpoint directories.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jQizhang]

/ok to test c573027

[jQizhang]

/ok to test 8cea42d