Unify HF snapshot and checkpoint writing

[yueming-yuan]

ci-sglang-pr: #40242

Single-LoRA saves and full-model HF exports share SnapshotPublisher for HF files and one write_checkpoint_dir() for directory handling. Overwriting removes the old checkpoint before writing the new one; failed overwrites leave no usable checkpoint. Readers that overlap later saves use distinct snapshot paths.

Changes

• Preserve main's raw-mode LoRA saves: native shards and a megatron_rank_sharded adapter config, without gathering HF weights or creating a publisher solely for native saves. Bridge-mode checkpoints and sampler snapshots derive HF adapter targets from their exported parameter names.
• Preserve best-effort HF adapter export for native single-LoRA checkpoints: log HF export errors and continue saving native parameters and training state. Native write failures still fail and clean up the checkpoint; sampler publication remains strict.
• Skip Inkling HF adapter import when native adapter optimizer state was actually restored, preserving fp32 master weights. HF-only initialization and --no-load-optim retain the existing master refresh.
• Training actors own optional WeightUpdater and SnapshotPublisher components. Ordinary training uses the updater, multi-LoRA uses the publisher, and training with snapshot export can use both. Each export reads current model weights through the actor-owned publisher.
• SnapshotPublisher.write_adapter() writes HF adapter files; write_model() streams raw full-model safetensors shards and writes their index and config/tokenizer files. All ranks participate in weight conversion; rank 0 writes these files. Bridge retains its existing full-model/merged-LoRA export implementation.
• Use one directory flow for HF exports, sampler snapshots, and native LoRA checkpoints: rank 0 prepares the directory and broadcasts preparation errors; ranks write contents and exchange their results; rank 0 then writes META.json or .complete when requested. The result exchange also waits for writers, so no additional barrier or phase framework is needed. Failed writes are cleaned after writers finish. Ordinary directories replace symlink publication and retained hidden versions.
• Finish adapter gathering before saving native single-LoRA shards and optional optimizer/scheduler state. Raw HF shard write errors drain the remaining gathers before propagating. Native training filenames and HF adapter filenames remain unchanged.
• When snapshot eval reuses --save-hf, append step_{rollout_id} to paths without a rollout ID format field and log the adjusted path. Existing templates, including {rollout_id:06d}, stay unchanged. Independent eval exports, shared-engine eval, and ordinary saves keep their existing paths. Eval retains its existing snapshot retention policy and removes consumed directories directly.
• Sampler snapshot names cannot be overwritten. LoRA HF exports include a merged model and an HF-only adapter/ subdirectory. Give each fully-async eval E2E case its own temporary snapshot directory because CI containers share host IPC.

Megatron's native full-model checkpoint and async-save protocol are unchanged. This does not recover process death, broken collectives, or failures inside Bridge/Megatron internals. Cleanup requires writers to have stopped; failed cleanup can leave incomplete files. Tinker checkpoint errors still terminate the trainer cell. Reusing a path while a reader remains active is unsupported, including repeated exports of the same step across runs.

Validation

• At 39ca293ff6, all five CPU jobs and pre-commit passed. Six NVIDIA GPU jobs failed with the stale Megatron 73b54618f7e5 override: five could not import Bridge because mla_qk_norm_config was absent; Kimi K3 failed the initial weight comparison on three routed_expert_norm tensors because that backend predates Kimi primitives. Both are present in merged Megatron Support move-based offloading in FSDP backend #102/Try to synchronize after FSDP wakeup #94. Removed the override to use the default miles-main (currently 651dd728e3c1); fresh CI is pending. No production code or tests changed.
• The first CI run after merging main failed test_lora_fields_roundtrip: Miles now emits full HF target paths, while the default SGLang branch still rejects them as CLI choices. Added the same SGLang #40242 dependency used by merged Miles Unify LoRA target selection with HF model layouts #3318; its current head is 2068f5ed9115. Verified at 39ca293ff6 in CI run 35932539087: all four stage-a CPU shards, stage-b CPU, and pre-commit passed. The original test_lora_fields_roundtrip executed and passed; its shard reports 3453 passed, 5 skipped, and 7 xfailed. CI used merge ref 7b8d5885c960, SGLang 2068f5ed9115, and Megatron 73b54618f7e5. NVIDIA GPU validation is still pending. No production code or test assertions were changed for this dependency fix.
• Merged main 7d1d15b1cc into 0755b073f6, preserving the HF/LoRA namespace move, HF target selection, raw-mode saves, and sampling-support replay. Local pre-commit and diff checks pass. Two-process CPU/Gloo probes cover native/HF/sampler success and failure, raw saves with zero HF gathers, target names from exported tensors, and per-slot rank/alpha. Component probes cover all nine updater/publisher configurations and reuse across exports; twenty Inkling restore A/B cases still pass. Conversion/framework setup are simulated; tensors, files, safetensors, and Gloo collectives are real. No new repository tests or GPU runs for this merge; hosted CI pending.
• At b1cbb43dc2, pre-commit and diff checks pass. A two-process CPU/Gloo A/B probe against 97c3eecab2 injects a rank-0 HF adapter write error: the base fails both ranks and removes the checkpoint; the fix returns successfully on both ranks and preserves all four native/training-state files, including exact optimizer tensors and scheduler state. Rank-1 native write errors and sampler HF export errors still fail both ranks and clean up. Repository function bodies, tensors, files, safetensors, and Gloo collectives are real; model conversion and optimizer objects are simulated. No repository tests or GPU runs were added for this fix.
• Hosted CPU CI and GitHub pre-commit passed on 8e3c71390a: all four stage-a shards and stage-b are green (run). CI tested merge ref be7f0665af0a with SGLang 5a8da8cc3fc1 and Megatron 73b54618f7e5. Removed the obsolete SGLang pin from the PR body after the 0.5.20 migration; both previously failing ServerArgs tests now execute successfully.
• Merged main 90ea0b3a8e; retained its strict native adapter validation and deterministic gradient dtype ordering. Local pre-commit run --all-files and git diff --check passed on 97c3eecab2. Formatting changes are isolated in a separate commit and AST-equivalent. Hosted CI for this head is pending.
• A CPU A/B probe through the checkpoint/initialization function bodies reproduced fp32 master corruption (1.0010000467 to 1.0) before the Inkling guard and exact preservation afterward. Twenty before/after cases cover native resume, step zero, legacy shards, HF-only loading, omitted optimizer state, --no-load-optim, and skipped base loading. Tensors and checkpoint files were real; the base-model loader, optimizer, and Inkling tensor mapping were simulated. No new repository tests or GPU runs were added.
• Two-process CPU/Gloo probes checked HF and adapter exports, native LoRA weights and optimizer files, rank-local write failures, preparation errors, completion metadata/marker errors, overwrite rejection, cleanup errors, and a simulated result-exchange failure. Checked that shard-write failures finish all remaining gathers, failed writes leave no completion marker, cleanup preserves the original error, and a failed synchronization does not trigger unsafe cleanup. Separate crash probes verified overwrite invalidation and retry. Conversion iterators and Megatron models were simulated; tensors, safetensors, filesystem writes, and Gloo collectives were real. No new repository tests were added.
• A prior CPU parameter-validation probe checked automatic per-step paths, existing/padded templates, escaped braces, unchanged non-reuse modes, and repeated validation without duplicating the suffix.
• Earlier GPU evidence on 0efaee7953: fully-async eval passed shared, fleet, and external, including step 0/1 HF exports in both dedicated eval modes. CI run passed CPU/pre-commit; NVIDIA jobs had 3 passes and 5 failures from the Bridge/Megatron mla_qk_norm_config import incompatibility. This predates the final writer refactor; the final implementation has not been validated on GPUs.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[yushengsu-thu]

@claude review

[yushengsu-thu]

approve - the same way used in multi-lora.
But agent catch some issue. So I list below.

• Reviewed d94d93b47088; reproduced against base using CPU/Gloo with stubbed model conversion, not GPU E2E.

• Regression: Failed exports leave tmp_step*; eval cleanup deletes only final paths. Two failed exports left four files versus zero on base.

• Regression: LoRA saving now uses rank-0-only directory creation/publication, breaking node-local checkpoints on other nodes. Base succeeds; head reports a missing directory.

• Existing issue remains: With empty HF exports, nested save phases mismatch across ranks; peers time out. Coordinate validation before entering adapter saving.

[claude]

🔴 save_hf_model now stages writes in a hidden _tmp_<name> sibling of the final path and only renames it in on success; on failure the tmp dir is never removed. export_hf calls this with raise_on_error=True, and eval_dispatch.py's failure handler only shutil.rmtree(hf_dir, ...) on the final path (a no-op since it never existed) — not the new _tmp_step_<id> sibling — so every failed --eval-hf-dir export now leaks a full snapshot dir on (recommended tmpfs) disk, unbounded by eval_keep_snapshots/eval_max_in_flight. Fix: write_checkpoint_dir must remove its tmp dir on failure (or every caller doing its own cleanup, like eval_dispatch, must also remove <parent>/_tmp_<name>).

Extended reasoning...

eval_dispatch._export -> actor_model.export_hf(rollout_id, hf_dir) -> save_hf_model(path=hf_dir, raise_on_error=True) -> write_checkpoint_dir(path, write_shards) at hf_export.py:169 builds tmp_dir = hf_dir.parent/f'tmp{hf_dir.name}', mkdirs it, then runs write_shards. If write_shards raises (e.g. the bridge-mapping RuntimeError at hf_export.py:155, or any export error on any rank), the exception propagates straight out of write_checkpoint_dir with no finally/cleanup — publish_dir, the only code that renames or discards tmp_dir, is never reached, so tmp_dir is orphaned on disk. The exception surfaces in eval_dispatch.py at the except Exception as e: around line 51, whose handler does shutil.rmtree(hf_dir, ignore_errors=True); since hf_dir (the final path) was never created, this removes nothing, and _tmp_step_<rollout_id> stays. Because each eval point uses a unique step_<rollout_id> name, no later export ever revisits that tmp path to clean it up (write_checkpoint_dir only removes a stale tmp_dir at the start of a new attempt targeting the same final path), so every failed…

Verification: normal. Real disk leak regression introduced on the eval-export failure path. write_checkpoint_dir (checkpoint_io.py:31-64) stages writes in tmp_dir = final_dir.parent/f"tmp{final_dir.name}" (line 37) and only renames onto final_dir in publish_dir (line 60), run last (line 64). If write_shards raises, run_checkpoint_phase re-raises (line 27-28) and propagates out with no try/finally to…

[yushengsu-thu]

Reviewed d94d93b47088; reproduced against base using CPU/Gloo with stubbed model conversion, not GPU E2E.

• Regression: Failed exports leave _tmp_step_*; eval cleanup deletes only final paths. Two failed exports left four files versus zero on base.
• Regression: LoRA saving now uses rank-0-only directory creation/publication, breaking node-local checkpoints on other nodes. Base succeeds; head reports a missing directory.
• Existing issue remains: With empty HF exports, nested save phases mismatch across ranks; peers time out. Coordinate validation before entering adapter saving.

[yushengsu-thu]

rmtree fails on symlink checkpoints written by the current write_checkpoint_dir; unlink symlinks (and dangling ones) before mkdir.

[yueming-yuan]

because the symlink is just added in multi-lora, I think there will not be users now, so we can do this for simplicity

[yushengsu-thu]

Three H200 async E2Es pass; LoRA fails on head and base. Two checkpoint regressions remain.

[yushengsu-thu]

[P2] Node-local LoRA saves fail on nonzero nodes; create checkpoint directories on each node.

[yueming-yuan]

given full checkpoint is written to only rank 0, I think this unification is ok