fix(lora): export PEFT adapters through Megatron-Bridge

[Arist12]

Part of #2705.

Problem

Miles manually assembles PEFT weights and metadata, duplicating Megatron-Bridge conversion logic and producing names or packed-expert layouts that external PEFT loaders may not consume. Multi-adapter writer failures also need rank-aligned handling, and a failed Bridge write must not expose a config-only adapter as loadable.

Change

Delegate single-adapter export to AutoBridge.save_hf_adapter. Stage Bridge output in a temporary directory, then promote safetensors before config so only a complete pair becomes visible. For Multi-LoRA, slice each slot to its configured rank and use Bridge conversion/configuration helpers before writing safetensors. Coordinate setup, conversion, write, and promotion failures across ranks.

Notes for review:

• Staging is created inside the guarded block, so a rank that cannot create it still reaches the setup consensus instead of hanging its peers.
• Both consensuses run on the gloo group: the failure being reported may be a poisoned GPU communicator. The compensating dist.barrier() stays on the default group because it must match Bridge's own exit barrier.
• Multi-LoRA export now raises on an empty adapter weight list. Bridge guards its own path; without this the multi-LoRA path could promote an empty safetensors plus target_modules: [] as a complete checkpoint.
• Multi-LoRA HF artifacts are now written in fp32 rather than bf16, matching Bridge's own export convention.

Known boundary: dense and per-expert MoE export work; --experts-shared-outer-loras remains native-checkpoint-only because Bridge currently rejects that PEFT layout.

Validation

• 77 focused fast tests, including injected config-only write failure, empty-adapter rejection, cross-rank collective symmetry, and temporary-directory cleanup
• 8x MI350X, Qwen3-4B colocated LoRA: the 30-rollout run resumed at iter_14, exported complete PEFT pairs at 19/24/29, and the final artifact contains 504 finite tensors with no staging residue
• 4x MI350X dense export loaded through the public PEFT API with nonzero LoRA parameters

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[jhinpan]

I reviewed the current head (99a218d5028f50a327de0f9923c58fab12242bf1) and tested both successful PEFT loading and rank-local failure behavior on AMD Instinct MI355X GPUs.

The staged safetensors produced by the happy path were loadable through PeftModel.from_pretrained, but the distributed failure protocol still has a merge blocker.

Blocker — staging-directory creation occurs outside setup consensus.

TemporaryDirectory(...) is entered before the inner try that assigns hf_export_err. If it fails on one rank, that rank skips all_gather_object and proceeds toward the trailing barrier while peers wait in all_gather_object. Injecting this failure on one of two MI355 ranks reproducibly timed out after 30 seconds.

Move staging-directory creation into the coordinated setup phase so every rank reaches the same consensus before any rank enters Bridge.

Publication occurs before all-rank save success is known.

Rank 0 moves weights/config into the public directory before save_errors is gathered. A peer failure can therefore be discovered only after files are visible. Overwriting an existing adapter also exposes new weights alongside the old config between the two replacements, and a config-rename failure can make that mismatch permanent.

Use an explicit two-phase protocol:

1. create staging and gather setup errors;
2. perform the collective Bridge export;
3. gather all save results;
4. publish on rank 0 only after global success;
5. gather promotion results.

A versioned directory or explicit commit marker is safer for overwrite/retry semantics.

The recovery protocol depends on undocumented Bridge internals.

The compensating dist.barrier() assumes save_hf_adapter has exactly one entry and one exit barrier. Its public contract only says the method is collective. The object consensuses also use the default RCCL group, unlike the Gloo helper used elsewhere on this path. Please centralize failure consensus on get_gloo_group() and avoid encoding Bridge's private barrier count.

Finally, keep Bridge's empty-adapter guard on the multi-LoRA helper path; otherwise an empty safetensors file plus target_modules: [] can be promoted as a complete checkpoint.

[jhinpan]

Follow-up review on the unchanged head (99a218d5028f50a327de0f9923c58fab12242bf1): one new retry/topology issue remains after excluding the first review's publication findings.

A failed multi-LoRA attempt directory is reused and later promoted wholesale.

save_multi_lora_checkpoints uses the deterministic _tmp_step_{iteration} directory. Several new coordinated failure points can raise after native shards have already been written, but none removes that directory. On restart, save_due_adapter_checkpoints retries because only final_dir suppresses a save, mkdir(exist_ok=True) reuses the old temporary directory, and os.replace(tmp_dir, final_dir) publishes every stale file that remains in it.

Concrete failure sequence:

1. an EP4 attempt writes _ep0 through _ep3 shards and then fails during PEFT conversion;
2. the run restarts at EP1, writes the no-EP-suffix shard into the same temporary directory, and promotes it;
3. a later EP4 load sees the stale _ep0 through _ep3 set as complete and resumes those pre-failure weights while reporting the retried step as valid.

Use a per-attempt unique staging directory whose name is chosen once and broadcast to all ranks. Alternatively, have one coordinator remove the deterministic temporary directory, reach a barrier, and only then let shard writers recreate/populate it. Concurrent shard writers should not each race to remove the shared directory.

[Arist12]

Thanks. Three points adopted; on the publication-ordering redesign I read Bridge's source and came to a different conclusion — details below so it can be checked.

Adopted. The export now stages into a temp directory inside the failure consensus rather than after it, so a rank that fails mid-export cannot leave a half-written directory that a later run mistakes for complete. The consensus all-gather runs on the gloo group rather than the default group, so it does not serialise behind NCCL work. And the empty-adapter case now raises instead of promoting an empty adapter_model.safetensors with target_modules: [] as a valid checkpoint.

Adopted (comment only). The .clone().float() cast now says why: it matches AutoBridge.save_hf_adapter, so both export paths write the same dtype.

Not adopted — the 5-step publication protocol / dropping the compensating barrier. The concern is that the extra dist.barrier() after the Bridge call papers over an ordering bug. Looking at megatron/bridge/models/conversion/auto_bridge.py:

• save_hf_adapter opens with a barrier (:582-583) and closes with one (:633-634), and there is no code after that final barrier. So nothing this PR could add can run after Bridge's own exit synchronisation from inside the call.
• The empty-adapter case raises on every rank (:592-598), not just rank 0, so it is already collective and cannot desynchronise the group.
• Rank 0 is Bridge's own choice of writer (is_rank0 = not dist.is_initialized() or dist.get_rank() == 0, :606); the PR does not assume it independently.

Removing the compensating barrier means the non-writer ranks return from Bridge while rank 0 is still writing, then reach the gloo consensus and observe an incomplete directory — which is the failure the barrier exists to prevent. Restructuring the whole publication into a 5-step protocol would mean reimplementing what Bridge already does rather than calling it, which is the opposite of what this PR is for (title: export PEFT adapters through Megatron-Bridge). Also relevant: completion is already signalled by HF_EXPORT_COMPLETE_MARKER (hf_export.py:192-193), so a partially written directory is not consumed as complete downstream.

Happy to revisit if I've misread the barrier placement — that is the load-bearing claim.

Not adopted — "this PR bundles four changes." The four hunks are the single change of routing PEFT export through Bridge: key prefixing, rank pattern, target modules and dtype are all what save_hf_adapter does and what the multi-LoRA path has to match to produce a loadable adapter. Splitting them yields intermediate states that write adapters PEFT cannot load.

Deferred — _tmp_step_{iteration} reuse across reruns. Pre-existing and orthogonal; same follow-up as noted on #2669.

Validation: tests/fast/backends/megatron_utils — 265 passed, plus the two TestUpdateWeightsZeroChunks failures that reproduce on main. pre-commit clean.