fix(lora): deduplicate native adapter checkpoints

[Arist12]

Part of #2705.

Problem

Single-LoRA checkpoints write one native adapter shard per global rank, so DP and CP replicas duplicate identical data. A writer-side save failure can also leave peer ranks entering later Bridge collectives.

Change

Write one native shard for each realized (TP, PP, EP, ETP) coordinate, retaining distinct EP shards and legacy load compatibility. Share native-save failures across ranks before PEFT export begins.

ETP is part of the shard identity because when ETP > TP two ranks share a TP rank while holding different expert slices. The _etp filename suffix is emitted only in that case: for ETP <= TP the ETP rank is a function of the TP rank, so every filename in use today is byte-identical.

Validation

• 89 focused fast tests, including shard naming and completeness under EP and ETP
• 4x MI350X: one adapter_megatron_tp0_pp0.pt shard with finite tensors
• EP4 full-model runs: four expert-distinct shards with replicated attention tensors

Fixes #2668

[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 (c2cac7eecb8a083e2252bf5a4d5e62ea60453f0d) and ran distributed topology probes on AMD Instinct MI355X GPUs.

Blocker — (TP, PP, EP) is not always a unique adapter-shard identity.

With four ranks configured as TP1 / PP1 / EP2 / ETP2, Megatron produced these distinct coordinates and tensor values:

• rank 0: (tp=0, pp=0, ep=0, etp=0), value 0
• rank 1: (tp=0, pp=0, ep=0, etp=1), value 1
• rank 2: (tp=0, pp=0, ep=1, etp=0), value 10
• rank 3: (tp=0, pp=0, ep=1, etp=1), value 11

adapter_shard_topology() selected only ranks 0 and 2 as writers and produced two files containing values 0 and 10. The ETP-distinct values 1 and 11 were silently discarded. Names and shapes still match on load, so #2733 cannot detect this corruption.

The invariant should be that every deduplicated coordinate contains byte-identical adapter tensors. Either include ETP in shard identity/naming when it is independent, or reject any realized topology where one (tp, pp, ep) maps to multiple ETP ranks.

Legacy loading is accepted when the code itself says it is invalid.

The TP/PP fallback is loaded unconditionally under ep_size > 1, despite warning that it is only valid when EP <= TP. For example, TP1 / EP4 can make all EP ranks read the same shard. Require an unambiguous realized mapping and a global all-rank existence decision; otherwise reject the legacy layout.

Rank-local filesystem failures are not fully coordinated or atomic.

save_path.mkdir() runs before the first barrier/error consensus, and _save_native_adapter_checkpoint writes directly to the final filename. A rank-local mkdir failure can strand peers at the barrier; interruption during torch.save leaves an existing but truncated shard. Coordinate directory creation, write a unique temporary file, then atomically replace the final shard only after successful serialization.

[jhinpan]

Follow-up review on the unchanged head (c2cac7eecb8a083e2252bf5a4d5e62ea60453f0d): this is distinct from the first review's ambiguous legacy-fallback finding.

Existing but topology-incompatible native shards silently become a fresh adapter.

_native_adapter_shard_path uses the current run's EP size. A checkpoint written with EP4 contains files such as adapter_megatron_tp0_pp0_ep0.pt, while an EP1 resume searches for adapter_megatron_tp0_pp0.pt. The global-rank legacy file is absent, and the TP/PP legacy branch is skipped because the current ep_size is 1. The loader therefore returns False, None; checkpoint.py only warns and continues with freshly initialized adapter weights.

EP narrowing cannot be solved by choosing another filename because expert slices would need a defined merge. The correct result is a hard incompatibility error, not silent training from a new adapter.

After the normal selection ladder is exhausted, check whether the directory contains any adapter_megatron*.pt shards. If it does, report that native checkpoint data exists but no complete compatible layout can be selected. A shard manifest/completion marker would make this decision unambiguous and also distinguish an interrupted save from a directory that was never a native checkpoint.

[Arist12]

Thanks — the ETP point was a genuine bug and is fixed; the rest I'd rather keep out of this PR, with reasoning below.

Adopted — ETP is part of the shard identity. With expert_tensor_parallel_size != tensor_model_parallel_size, two ranks that share (tp, pp, ep) hold different expert shards, so naming by (tp, pp, ep) alone made them overwrite each other. megatron_shard_name now takes all four coordinates and adapter_shard_topology gathers 4-tuples.

Reviewing that change surfaced a second instance of the same mistake: the global writer (the rank that promotes the temp directory) was still elected on (tp, pp, ep) == (0, 0, 0), so at etp_size > tp_size two ranks both qualified and raced on shutil.rmtree(final_dir) / os.replace(...). Now elected on all four coordinates.

Deferred — the legacy tp/pp filename fallback under ep_size > 1. Correct that the fallback can match a file that a differently-sharded run wrote. It is pre-existing behaviour and removing it breaks resume from every checkpoint written before this PR, so it wants its own change with a migration note rather than being folded into the dedup fix.

Deferred — mkdir / atomic-write coordination and the EP-narrowing silent fresh start. Both are real, both are the same class of problem as the load-ladder consensus: the fix is to route the whole save/load ladder through lora_checkpoint.raise_if_any_rank_failed and agree on the topology before touching the filesystem. This PR adds that helper and the topology query; using them everywhere is the follow-up that #2733 and #2516 also point at. Keeping this one to "one writer per shard" so the dedup is reviewable on its own.

Deferred — _tmp_step_{iteration} reuse. Pre-existing; a rerun at the same iteration reuses a stale temp directory. Unrelated to writer election.

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