ZeRO-3 HybridMegatron QKV layout mutations are not persisted

Description

HybridMegatronContainer.transform_for_inference() and transform_for_training() mutate gathered QKV parameters in place, but their ZeRO-3 GatheredParameters contexts do not specify modifier_rank. GatheredParameters requires a modifier rank to broadcast a modified full parameter before repartitioning. Without it, the GPT-NeoX head-interleaved QKV layout can survive inference conversion and then be interpreted as contiguous all-Q/all-K/all-V by the DeepSpeed inference layer.

The affected code is deepspeed/module_inject/containers/features/hybrid_megatron.py. Both conversion directions are affected.

Observed behavior

In a Pythia-410M Hybrid Engine rollout with ZeRO-3 on AMD MI250:

• native Hugging Face first token: 187
• injected HE first token before the fix: 39318
• transformed native QKV versus injected QKV before the fix: max_abs=20.5625, mean_abs=1.11529
• the first material forward mismatch occurred in layer-0 attention

This is a parameter-layout correctness issue rather than ordinary BF16 near-tie behavior.

Proposed fix

Use GatheredParameters(non_active_params, modifier_rank=0) in both transform_for_inference() and transform_for_training(). This persists the forward and inverse layout mutations when ZeRO-3 repartitions the parameters.

After the change:

• transformed QKV matched exactly (max_abs=0)
• native and HE first tokens both became 187
• repeated 20-step eval() -> train() layout transitions completed successfully
• controlled Pythia HE runs completed on 2/8 x MI250 and 8 x MI300X

A small regression test constructs ZeRO-marked QKV weight/bias parameters, verifies that inference conversion changes the layout, verifies exact training-layout restoration, and checks that both gather contexts use modifier_rank=0. The test passed on an MI300X compute node.

Scope note

Current master reaches an earlier, independent compatibility failure with the installed recent Transformers version because GPTNeoXAttention.hidden_size is no longer present. The end-to-end results above used the corresponding local Transformers compatibility adjustments. The proposed QKV patch is intentionally limited to the ZeRO-3 mutation-persistence bug and does not bundle those unrelated compatibility changes.

Environment used for the full reproduction

• AMD MI250 and MI300X
• ROCm 6.4.1
• PyTorch 2.9.1+rocm6.4
• BF16, ZeRO stage 3
• EleutherAI/pythia-410m-deduped student

Start in deepspeed/module_inject/containers/features/hybrid_megatron.py, focusing on HybridMegatronContainer.transform_for_inference() and transform_for_training() and their GatheredParameters contexts. Review the regression test described in the issue, then verify that both contexts use modifier_rank=0, inference conversion changes the layout, and training conversion restores it exactly.