Fix ZeRO parameter alignment for grouped_mm

[fwerkor]

Summary

ZeRO stage 1/2 stores model parameters as views into flattened fp16/bf16 buffers. Individual parameter views can start at non-16-byte offsets even when the flat buffer itself is aligned, which breaks alignment-sensitive kernels such as torch._grouped_mm.

Pad parameter boundaries inside the ZeRO flat buffer so model parameter views remain 16-byte aligned without duplicating misaligned parameters. The padded layout is propagated through partition/gradient offsets, LP↔HP linkage, DeepCompile gradient buffers, checkpoint restore, and zero_to_fp32 reconstruction.

Older checkpoints without parameter-alignment padding remain loadable; their compact layout is converted when restored.

Tests

• ZeRO-1/2 BF16 regression with a deliberately misaligned parameter layout.
• Verifies zero-copy aligned flat-buffer views across optimizer steps.
• Verifies checkpoint reload with load_module_only=True and load_optimizer_states=False.
• Synced with current master, including the ZeRO-1/2 DeepCompile file rename.

Fixes #8276

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[Copilot]

Pull request overview

Fixes ZeRO stage 1/2 parameter rebinding so model parameters exposed to alignment-sensitive kernels (e.g., torch._grouped_mm) remain 16-byte aligned even when their corresponding flat-buffer views are not.

Changes:

• Update ZeRO stage 1/2 _update_model_bit16_weights() to keep zero-copy views when aligned, otherwise preserve/allocate an aligned tensor and copy updated values into it.
• Add a ZeRO-1 BF16 regression test that constructs a deliberately misaligned flat-buffer view and asserts the model’s parameter pointer remains 16-byte aligned across an optimizer step.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
deepspeed/runtime/zero/stage_1_and_2.py	Adds alignment-aware logic when rebinding model params from flat buffers to avoid misaligned parameter pointers.
tests/unit/v1/zero/test_stage2_flatten_on_gpu.py	Adds a regression test model + test case that validates alignment is preserved for a misaligned flat view under ZeRO-1 BF16.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[tohtana]

Hi @fwerkor, thank you for submitting this PR!

This is useful, but I think this approach has two challenges:

• This has too much impact about the performance and memory usage. All misaligned parameters will be duplicated.
• This can easily cause a bug about synchronization. Actually the current code leaves the flat buffer stale when a checkpoint is loaded, so the first optimizer step silently overwrites loaded misaligned parameters. (I reproduced this on ZeRO-1/2 with both load_module_only=True and load_optimizer_states=False)

It might be okay to accept this as an opt-in, but I still feel the performance impact is too big. I think a long term solution is to pad parameters in the flat buffer. It will need more changes, but the performance/memory impact will be small.
Can you share your thoughts?

[fwerkor]

Hi @fwerkor, thank you for submitting this PR!

This is useful, but I think this approach has two challenges:

• This has too much impact about the performance and memory usage. All misaligned parameters will be duplicated.
• This can easily cause a bug about synchronization. Actually the current code leaves the flat buffer stale when a checkpoint is loaded, so the first optimizer step silently overwrites loaded misaligned parameters. (I reproduced this on ZeRO-1/2 with both load_module_only=True and load_optimizer_states=False)

It might be okay to accept this as an opt-in, but I still feel the performance impact is too big. I think a long term solution is to pad parameters in the flat buffer. It will need more changes, but the performance/memory impact will be small. Can you share your thoughts?

Thanks, agreed. I switched to padding the ZeRO flat buffer and synced with current master. The previous GPU CI was killed by the 1-hour timeout (exit 137) at ~76%, with no failures before that.

[tohtana]

Thank you for switching to padding the flat buffer and adding the checkpoint reload regression tests! This addresses my original concerns.
I found two additional correctness issues in the updated implementation, detailed inline below.
This PR can have a significant impact on broad usage of DeepSpeed. Can we validate the correctness by running a simple training loop?

[fwerkor]

Thanks, addressed both issues.

• The CPU flatten fallback now computes parameter offsets and padding from the preserved metadata tensors, so it no longer depends on the temporary one-element parameter placeholders.
• The LP/HP gradient mapping now accounts for internal padding entries, fixing safe_get_full_grad() and safe_set_full_grad() for padded layouts.

I also added regression tests for the CPU flatten fallback with a short training loop, and for safe_get_full_grad() / safe_set_full_grad() under internal alignment padding, covering both ZeRO-1 and ZeRO-2.

The existing alignment/checkpoint regression still passes as well.

[tohtana]

Thank you for addressing my previous comments! The changes look good to me.
I left a few more comments (two of them already existed before your recent update).

[fwerkor]

Thanks, addressed all three comments. I fixed the offload gradient indexing, preserved _zero_padding on cloned gradients, and restricted legacy optimizer-state redistribution to partition-shaped tensors. Added a ZeRO-1/2 regression for tensor metadata as well.

[tohtana]

Thanks for the fixes! They look good to me.
I left one comment about a similar issue to the previous one.

As we are not super confident in the correctness of all the features, can we make this opt-in?
I opened a PR to this PR branch. If you think it useful, please review and merge it. It would also be good if you could implement it separately.

[fwerkor]

Thanks — I merged your opt-in PR into this branch and also addressed the elastic tensor_step metadata case from the remaining thread. The current head includes both fixes and the ZeRO-1/2 parameter-alignment regression suite passes 24/24. Please take another look when convenient.

[tohtana]

Thank you @fwerkor! Everything looks good to me. I appreciate your patience.
Let's merge this now.