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: cc74480ea (rebased onto 3b809329f; pytest tests/unit/module_inject -> 162 passed, 43 subtests passed, including the new test_hybrid_megatron.py).

Independent reproduction on 8x H200 (DeepSpeed OPSD with GPT-NeoX/Pythia HE): without this change the HE rollout silently collapses to one repeated token (the job still exits 0), while with it the HE first-step output matches the non-HE baseline and the 20-step runs are reproducible on 2x and 8x H200. The same environment also hit the separate GPTNeoXAttention.hidden_size removal noted under Scope below.

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

[LiRunGuo]

Review ping after a rebase.

Rebased onto current master (3b809329f); no conflicts. Validation on the rebased head cc74480ea:

pytest tests/unit/module_inject
  162 passed, 43 subtests passed in 12.7s     (includes the new test_hybrid_megatron.py)

Independent corroboration from our side (8x H200, DeepSpeed OPSD with GPT-NeoX/Pythia hybrid engine): without modifier_rank=0 the injected GPT-NeoX layer produces a degenerate rollout — every response token identical, while the job still exits 0, so it fails silently; with the fix the HE first-step output matches the non-HE baseline, and 20-step runs are reproducible at both 2x and 8x H200. The same box also hit the separate GPTNeoXAttention.hidden_size removal that this PR deliberately leaves out of scope.

@delock — since you own the OPSD direction, could you take a look when you have a slot? (As the fork author I cannot add reviewers on this repo, hence the ping.) This closes #8391.