[cuichenx]

What does this PR do?

Replicate an empty Transformer Engine grouped-linear _extra_state entry for each expert when saving a sharded checkpoint.

Stateless FP8 recipes encode the absence of extra state as an empty byte tensor. The grouped-linear checkpoint adapter previously tried to decode that tensor as serialized FP8 metadata whenever FP8 execution was active, which returned no state and then failed while indexing it. Treating the empty tensor as stateless preserves the per-expert checkpoint key layout without inventing FP8 metadata.

Issue tracking

Linked issue: N/A

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter on my PR

Validation

• uv run python -m torch.distributed.run --standalone --nproc_per_node=4 -m pytest -q tests/unit_tests/transformer/test_transformer_engine_grouped_linear.py (15 passed on every rank)

• BASE_REF=base CHECK_ONLY=true SKIP_DOCS=true bash tools/autoformat.sh

• 128-GPU DeepSeek V4 real-data training completed 100 finite steps, saved model, distributed optimizer, and RNG state, then loaded iteration 100 and resumed through iteration 105

• I, the PR author, have personally reviewed every line of this PR.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[cuichenx]

Closing this as completed. The stateless grouped Transformer Engine extra-state handling is already fixed on main by #5997 (7e6a045a2a99659ddd2919fc19013a164c508825) and will reach dev through the regular main-to-dev sync.