fix(lora): reject incomplete native adapter checkpoints

[Arist12]

Part of #2705.

Problem

Native LoRA resume copies the intersection of saved and current parameter names, then restores optimizer and progress state. An empty, partial, shape-incompatible or numerically corrupt shard can therefore resume training on fresh or mixed adapter weights.

Change

Require the native shard to match every expected adapter name, shape and dtype, with no unexpected names and no non-finite values, before copying any weight or loading training state. A model that exposes no adapter parameters at all is rejected too: otherwise an empty shard matches it on every other check and the load reports success.

Validation

• 69 focused fast tests covering exact load plus missing, unexpected, shape-mismatched, dtype-mismatched, non-finite and empty-model rejection before optimizer restore
• Full Qwen3-30B-A3B on 8x MI350X: four EP shards passed strict validation at 384 tensors each; optimizer, scheduler and data cursor restored, rollout 30 completed two optimizer steps, and iter_30 was saved

Notes for reviewers

Two related gaps are deliberately left for follow-ups, since neither is fixable inside this validation function:

• With virtual pipeline parallelism, chunks number their local layers independently, so two chunks can expose the same adapter parameter name. Both the save and the load side flatten chunks into one dict[name, tensor], so the shard format itself would have to become chunk-keyed.
• is_lora_model accepts any name containing adapter, while _is_adapter_param_name accepts only lora_ or exposed .adapter.(linear_in|linear_out) names. Reconciling the two predicates is a separate change; the new empty-model rejection turns the resulting mismatch into a clear error instead of a silent no-op load.

[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 (8d08be274f0479f2e4ecc8e0c68de2b5ab31d4c6) and exercised the native load path on an AMD Instinct MI355X.

Blocker — dtype and non-finite tensors are still accepted.

_validate_native_adapter_state checks keys and shapes, but not dtype or finiteness. In the GPU probe, a same-shape float16 tensor was silently converted into a float32 parameter, while both NaN and Inf were copied into the live model. The function then reports a successful load, so a later training-state restore can proceed from poisoned weights.

Please validate the complete shard before any copy:

• the loaded object and every value are tensors;
• exact key set;
• exact shape and dtype;
• torch.isfinite(...).all() for floating/complex tensors.

No model parameter or training state should be mutated until every local check succeeds.

Additional correctness gap — virtual-pipeline chunks can collide by parameter name.

expected is flattened into dict[name, param]. VPP chunks number their local layers independently, so two chunks can expose the same adapter name. The later entry overwrites the earlier one; this change then copies only into the retained parameter, leaving another chunk at fresh initialization. Key validation by (chunk_index, name), or explicitly reject duplicate names before save/load.

Stack integration requirement. When #2516 is rebased, local read/strict validation must be separated from mutation: every rank reads and validates, ranks reach a global error consensus, and only then do they copy parameters and reload optimizer masters. Also update these tests for #2516's new args parameter rather than resolving the conflict by taking its permissive load loop.

[jhinpan]

Follow-up review on the unchanged head (8d08be274f0479f2e4ecc8e0c68de2b5ab31d4c6): one additional correctness gap was found after excluding the first review's findings.

An empty model and an empty native shard are accepted as a successful adapter load.

_validate_native_adapter_state treats expected == {} and state_dict == {} as an exact match. load_lora_adapter then returns (True, iteration), and the caller logs that the LoRA adapter loaded successfully. I reproduced the current behavior directly: an empty adapter_megatron_rank0.pt with a model exposing no names matched by _is_adapter_param_name returned (True, None).

This is reachable through mismatched model predicates: is_lora_model accepts any parameter name containing adapter, while _is_adapter_param_name accepts only lora_ or exposed .adapter.(linear_in|linear_out) names. Multi-LoRA parameters such as ...adapters.0.linear_in.weight satisfy the first predicate but not the second, so the generic adapter export path can write an empty native shard.

Once #2516 is stacked, this vacuous success can also restore optimizer/scheduler/iteration even though no adapter weight was restored.

Please reject both an empty expected adapter set and an empty native shard, with a diagnostic that distinguishes “the model has no adapter parameters” from “the checkpoint contains no adapter tensors.”

[Arist12]

Thanks — the dtype/finiteness blocker and the empty-adapter follow-up are both fixed; two of the other points I'd like to push back on.

Adopted. _validate_native_adapter_state now rejects non-tensor values, dtype mismatches and non-finite floating/complex tensors alongside the existing name/shape checks, and nothing is copied until every check passes. The follow-up was right that an empty shard matched an empty model on every other check, so a model that exposes no adapter parameters is now rejected with its own message.

Not adopted here — VPP name collisions. Real, but not fixable in this function: save_lora_checkpoint flattens chunks into the same dict[name, tensor] that the loader validates against, so keying expected by (chunk_index, name) would make every existing shard unloadable while still writing colliding names on save. This needs the shard format itself to become chunk-keyed. Noted in the PR body as a follow-up.

Not adopted here — global failure consensus. Agreed that a rank-local raise leaves peers in the next collective, but the same is true of the pre-existing native_path.exists() and torch.load on this path, so a consensus belongs to the load ladder as a whole rather than to the validator. In practice the raise is not a silent hang: the rank exits with the checkpoint's name in the traceback and the launcher tears the job down, which is strictly better than the silent half-load this PR replaces. #2669 introduces lora_checkpoint.raise_if_any_rank_failed; coordinating the whole ladder on top of it is the natural follow-up.

No longer applicable — the #2516 stacking note. #2516 has been narrowed and no longer touches the native-load site or adds an args parameter, so there is no conflict with this PR and no test churn.

Validation: 69 focused fast tests; full tests/fast/backends/megatron_utils green apart from two TestUpdateWeightsZeroChunks failures that reproduce on main.

One thing worth a maintainer opinion: strict dtype equality means a shard written under a different precision flag no longer loads, where the old copy_ would have cast it. That looked like the right trade (a dtype difference means the checkpoint came from a different model build, and the optimizer state is dtype-coupled too), but it is a deliberate behaviour change rather than a pure tightening.