Fix best checkpoints during async evaluation

[fus3r]

Summary

Async evaluation receives a rollout slot before it acquires the vLLM lock. A
refit can win that race, so the slot number alone does not identify the policy
used to generate the evaluation samples.

The vLLM lock now carries the step of the policy loaded by vLLM. A successful
broadcast publishes its step while releasing the lock. Non-partial evaluation
records the version returned by that same lock. The metric is therefore attached
to the policy that generated the samples in either lock ordering.

When training has already advanced, the trainer promotes the matching
global_stepN checkpoint instead of saving the current actor. It stages the
actor, critic, and optional HF export before publishing any destination. The
actor checkpoint is published last as the durable commit marker. A copy or
rename failure rolls back every published destination, removes pending
directories, and leaves the in-memory best metric unchanged. Previous best
checkpoints are pruned only after publication succeeds.

Configurations that can delay a best metric require --ckpt.save_steps 1, at
least async_queue_size + 1 rolling checkpoints, and no disk-size pruning.
force_sync_mode is exempt only with one queue slot. The Qwen quick starts and
the 4B, 35B, and Omni3 Slurm recipes now use the required cadence. The 753B smoke recipe keeps its sparse checkpoint cadence and explicitly disables
best-checkpoint selection.

A resumed run recovers the best score from finalized actor metadata and ignores
unfinished .tmp and .old directories. A fresh run with checkpoint loading
disabled ignores best metadata left in its output directory. HF exports follow
the retained actor tags. Partial rollouts still permit several policy versions
within one evaluation, so that mode skips best-checkpoint selection.

Testing

• python -m pytest -q tests/unit/test_rl_trainer.py (33 passed)
• python -m pytest -q (187 passed on the pre-v0.1.4 tree)
• python -m compileall -q molt examples/python tests
• bash -n examples/scripts/*.sh examples/scripts/quick_start/*.sh examples/scripts/slurm/*.sh
• pre-commit run --files <changed files>
• git diff --check

The unit tests cover both outcomes of the rollout-slot/refit race, delayed and
direct saves, two pending metrics, partial rollouts, restart behavior, fresh-run
isolation, HF retention, and the checkpoint configuration matrix. Filesystem
tests inject copy and publication failures independently for actor, critic, and
HF destinations. They check rollback, cleanup, source preservation, and
unchanged metric state, including an ENOSPC copy failure. They also cover
recovery of an orphaned .old backup and ignoring a NaN metric.

Earlier distributed GPU validation

Before the final transaction and Slurm recipe changes, I ran the parent and a
patched revision consecutively in the same on-demand Verda session in FIN-02 on
2026-07-18 UTC. Both started from
f6d9191623c3bdc3eada42cde80c0877247636d7. The production diff executed in
that run has SHA-256
96a3f06f8edfd75c462f80a74b131a2b4107b4da3a0462c40726ecbdf4e1ffaa.

That run validates the lock versioning, successful actor promotion, serialized
model and optimizer payloads, and vLLM refits for that exact revision. It is not
byte-for-byte GPU validation of the current diff. The final transaction,
critic/HF failure paths, and additional Slurm defaults are covered locally.

• 8 x NVIDIA A100-SXM4-40GB, with NV12 between every GPU pair in
  nvidia-smi topo -m
• hijkzzz/molt:0.1.2 at
  sha256:b9c82365b0c65e9cd4daf0addc34c9a5eba89cfc4593fa2e480246dc7c1dfcd2
• Qwen/Qwen3-4B-Instruct-2507 at Hugging Face snapshot
  cdbee75f17c01a7cc42f958dc650907174af0554
• 4 actor/reference FSDP2 GPUs and 4 vLLM GPUs with TP=4, NCCL,
  async_queue_size=2, ckpt.save_steps=1, and non-partial rollouts
• Torch 2.11.0+cu130, CUDA 13.0, Ray 2.55.0, vLLM 0.25.1, and Transformers
  5.8.1

Each run used mixed reward groups, completed four training steps, evaluated
policy step 2 while the trainer reached step 4, and fingerprinted the serialized
model and optimizer payloads. In the patched run, evaluation acquired the lock
first and the step-3 broadcast waited 65.65 seconds. The broadcast-first
ordering is covered by the lock interleaving unit tests.

tree	metric step	best client step	best serialized payloads match	minimum reward std	minimum actor grad norm	complete vLLM refit checks
parent	2	4	model and optimizer from global_step4	0.7071	0.02829	4/4
GPU-tested revision	2	2	model, optimizer, and file manifest from global_step2	0.7071	0.03608	4/4

Checkpoint fingerprints and refit details

Parent run:

• model: best = step 4 =
  ddff4b2b4bd24d9c3dd985df2b46b023b446ca25fd7643b4199757846c79da08,
  while step 2 =
  215e9ce0b42c7e37523f8f0da6b3446d4bd5df92b8ea67c7c4e735445b1bad66
• optimizer: best = step 4 =
  991b9402af2c687b40d855307d3260c7b78ceed31f013b48c481771b59e2be4e,
  while step 2 =
  c1243241a5b6c7cd1b66e32fa2bab003b171c81cf5c23ef7bfcbc3226c983319

GPU-tested revision:

• model: best = step 2 =
  92072e0bd013d3ff30b90081ce55b8490acece961d5e4246568a46505f6d2aec,
  while step 4 =
  341bb2423a0061d9d4941cee7012142caf126819ff2cc16a5570cad622eaada2
• optimizer: best = step 2 =
  8fad15d336f9c11fb8deb85950008cfad3db41a2ff0563fd8efb9c413f69e6f8,
  while step 4 =
  cd2696ee9a613066495def15f120ae755e6125e4df231a30d7aee6bc9b4b1bea
• checkpoint file manifest: best = step 2 =
  4313755ec60143f5a3231717b19f16d2507d977d976010c087c7bcc538f6dfff,
  while step 4 =
  05bac86c0bc42a7f1799dde2f29b4ebb0f46df1b3999d3293dc5d28cbea48733

The run enabled --train.check_weight_update_equal. After every broadcast,
all floating-point parameters held by each vLLM worker were reported as
refreshed. vLLM separately ignored one lm_head.weight flush per worker and
broadcast. Qwen ties that alias to its input embeddings
(tie_word_embeddings=true), so it is not a separate vLLM parameter. The
whole-worker check passed after all four broadcasts.

Harness SHA-256 values:

• runner: 1941c78282d72b02aaff12d9a242e17c72d9c15a6672090755ca05cabfc99b3a
• mixed-reward environment:
  6c1bb21877b91961a3571f42454a79b4a595dfd691bd91b4a728f1477b9b4a64
• checkpoint fingerprint script:
  81c123bf8a64fecc3b69bf91093f7086b6f52663da90182b5d12d41d2166962e

The evidence archive contains sanitized logs, GPU inventory and topology,
runtime and source provenance, payload manifests, copied optimizer metadata,
and four provider-console crops showing Verda, 8 x A100 40GB, FIN-02, and the
running state. It contains no account detail, network address, instance name,
resource identifier, financial field, or raw session identifier. Each run
produced about 180 GiB of checkpoint data. The archive represents the roughly
135 GiB compared subset with per-shard hashes instead of attaching the shards.

Evidence archive: pr3-verda-a100-evidence.zip, SHA-256
055b7e766e4442ac1063f6c712a39c200655fab6cc9ef68d8372e9b6d6023f93

[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.

[hijkzzz]

/claude review

[claude]

Light review complete. Read the full diff and traced the new delayed-best-checkpoint promotion path (transactional actor/critic/HF copy with .tmp/.old rollback, orphaned-backup recovery, and the _prune_hf_checkpoints/retention alignment) through the CheckpointManager helpers. The RLTrainer startup validation (save_steps 1, max_num ≥ queue_size+1, unbounded max_mem when delayed-best is possible) lines up with the recipe changes: the two scripts kept at save_steps 5 correctly add --ckpt.best_metric_key none, and the ones enabling best-selection move to save_steps 1. The VLLMLock now carries the policy version so a delayed async-eval metric is attributed to the vLLM policy that produced it, and the test file covers the slot race, transaction rollback per role, partial-rollout skip, NaN guard, and validation matrix.

One minor typo flagged inline in the save_steps 1 validation message. Note: compileall/pytest could not be executed in this sandbox, so tests were reviewed statically rather than run.

Not approving — comments only.

[fus3r]

Rebased this onto the latest main, and the targeted tests still pass (33 passed). When you have a chance, could you take another look? The NVIDIA workflow validation is still pending.

[fus3r]

Rebased onto the latest main and reran the full test suite; everything still passes. Thanks for taking another look.