fix(export): drain the tensor generator before failing a strict save

[v0ropaev]

A strict export that hits an unexpected tensor name hangs the job instead of failing it.

save_generator raised from inside its loop:

for name, tensor in generator:
    all_yielded_keys.add(name)
    if name not in all_expected_keys:
        if strict:
            raise KeyError(...)

Its own comment twelve lines above says why that is not ours to do:

In a distributed environment, only rank 0 should write to disk. Other ranks must still exhaust the generator to participate in collectives.

Those ranks are inside for _ in generator: pass. When rank 0 leaves the generator mid-stream, they stay blocked in the gather for the tensor rank 0 never asked for, until the process-group timeout. save_hf_weights is where a user reaches this, its default is strict=True, and its docstring says "This method is collective and must be called by all ranks".

Measured on 53681c1 with state.py loaded on its own, a two-key shard and a generator that records what it was asked for:

KeyError: "Tensor 'surprise.weight' from generator not found in the original model structure. ..."
items the generator was asked for: ['model.first', 'surprise.weight']
fully drained: False

model.second is never pulled. With this change the same probe reports ['model.first', 'surprise.weight', 'model.second'] and raises the same KeyError.

The distributed path had it too

_save_generator_distributed carried the same mid-loop raise, and there it is worse: every rank both produces and consumes the generator, so any rank can be the one that walks away.

That function already has the machinery for this, and says so. local_export_error is gathered across ranks and raised afterwards, under this comment:

Strict-mode check: ensure all expected tensors were written. Aggregate per-rank missing counts so all ranks raise consistently (avoids hangs on the trailing barriers).

and the error it raises reads "Checkpoint export failed after all ranks drained the tensor generator: ...", which the mid-loop raise made untrue. The duplicate-key branch in the same loop, eight lines down, already records into local_export_error and continues rather than raising. The unexpected-key branch now does the same.

Scope

Both paths keep writing nothing past the first offender, exactly as raising on it did, and both raise once the generator is dry. strict=False is untouched. Every caller that passes strict=False is untouched too, including the training checkpoint path in checkpointing.py, whose own comment makes the same point ("Other ranks drain the streaming export generator so TP/PP/EP collectives complete").

Testing

Two cases in tests/unit_tests/models/test_hf_pretrained_state.py, next to the existing save_generator tests and using the same _write_safetensors_index and _mock_single_rank_distributed helpers: one for the single-writer path, one for the distributed one. Both assert the generator was pulled dry and that the error still arrives. Both fail on main and pass here, and I checked two control cases (a complete shard still saves, strict=False still warns and saves) which pass either way. No torch.distributed process group is needed.

I could not run the repository's suite: tests/unit_tests/conftest.py imports megatron.core, which is not installable here, and megatron/bridge/__init__ pulls in the diffusion bridges, which need it too. So the runs above load state.py standalone with importlib, which is the real class and the real method, and I replicated the two new test bodies verbatim against it. ruff 0.9.9 check and format are clean on both files.

Same reporter as #6410, different file, no overlap.

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