Persist ZeRO-3 HybridMegatron QKV layout changes

[LiRunGuo]

Summary

• persist GPT-NeoX/Megatron QKV inference-layout mutations when ZeRO-3 repartitions gathered parameters
• persist the inverse layout mutation when returning to training
• add a focused QKV layout round-trip regression test

GatheredParameters needs modifier_rank when a gathered full parameter is modified. Without it, the transformed QKV tensor is not reliably broadcast before repartitioning, and the inference operator can interpret the original head-interleaved GPT-NeoX layout as contiguous Q/K/V.

Closes #8391.

Validation

• DeepSpeed pre-commit checks: passed
• focused unit test on AMD MI300X: 1 passed
• same-weight Pythia diagnostic after the fix: transformed QKV max_abs=0; native and HE first token both 187
• 20-step Pythia HE training completed on 2 x MI250, 8 x MI250, and 8 x MI300X without layout-round-trip recurrence

Scope

This PR intentionally contains only the ZeRO-3 QKV mutation-persistence fix. A separate compatibility issue in recent Transformers (GPTNeoXAttention.hidden_size removal) can prevent current master from reaching injection in that environment and is not bundled here.

Exact head: 73165a4a04dcc4a6c0bcc1677880f04bf2d0b841.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 73165a4a04

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the mandatory commit sign-off

This is a non-merge commit, but its commit message has no Signed-off-by trailer, so it violates the repository's mandatory contribution and CI requirement. Recreate the commit with --signoff using the configured Git identity.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[LiRunGuo]

Hi @tohtana @hwchen2017, a gentle ping on this PR when you have a moment.

It's a small fix (+58/−2) for #8391. The ZeRO-3 GatheredParameters contexts in HybridMegatronContainer.transform_for_inference() / transform_for_training() now pass modifier_rank, so the in-place QKV layout change is broadcast before repartitioning. Previously the mutation could be dropped. A QKV layout round-trip regression test is included.

DCO and the modal-torch-latest CI pass. The commit does carry a Signed-off-by trailer, so the automated note about a missing sign-off doesn't apply.

Happy to adjust anything. Thanks!

[LiRunGuo]

Gentle bump with fresh validation. I rebased this onto current master (1b892a75d) locally and re-ran the new unit test:

• rebase applies cleanly, no conflicts
• pytest tests/unit/module_inject/test_hybrid_megatron.py → 1 passed
• pre-commit run --files deepspeed/module_inject/containers/features/hybrid_megatron.py tests/unit/module_inject/test_hybrid_megatron.py → all hooks pass

The change is still just two modifier_rank=0 arguments plus the focused round-trip test (+58/−2). Without modifier_rank, the in-place QKV layout edit inside GatheredParameters is not broadcast before repartitioning, so the GPT-NeoX head-interleaved layout can survive into the inference operator and produce wrong logits (we saw the first greedy token flip from 187 to 39318 on 2×MI250). The test asserts the gather is called with modifier_rank=0 in both directions and that training layout is restored exactly.

@delock in case it is useful — this is the third fix in the same OPSD / Hybrid Engine series you merged earlier (#8264 and #8265). Thanks all.