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.