Normalize ZeRO-3 DeepCompile grad dtype before reduction

[tohtana]

Some backward kernels produce gradients in their computation dtype, not necessarily in the parameter storage dtype. For example, if a backward path accumulates or promotes math in fp32, a parameter stored as bf16 can still receive an fp32 raw gradient from that backward computation.

In normal PyTorch execution, that raw gradient reaches the leaf-gradient accumulation step, which stores it according to the tensor's expected grad dtype. ZeRO-3 DeepCompile intercepts the raw compiled-backward gradient before that leaf accumulation boundary. The reducer was assuming the raw gradient dtype was already the expected leaf grad dtype, so it could select an fp32 communication bucket even when the ZeRO grad partition storage was bf16.

To address this, this PR changes dc.reduce_grad's behavior to match PyTorch's leaf-gradient dtype contract. ZeRO-3 registration now records the expected grad dtype for each parameter, and reduce_grad normalizes raw compiled-backward gradients to that dtype before selecting the communication bucket.

This follows the documented grad_dtype behavior, including preserving explicit grad_dtype=None opt-outs: https://docs.pytorch.org/docs/main/generated/torch.sparse.semi_structured.SparseSemiStructuredTensorCUSPARSELT.html#torch.sparse.semi_structured.SparseSemiStructuredTensorCUSPARSELT.grad_dtype

Testing

• python -m py_compile deepspeed/compile/init_z3.py tests/unit/compile/test_zero3_grad_dtype.py
• python -m pytest -q tests/unit/compile/test_zero3_grad_dtype.py
  I confirmed both passed.

The full repro would also be useful.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: be9d1d07fa

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[eternalNight]

Looks good to me. Thanks!

[tohtana]

@eternalNight Thank you for your review!

[PKUWZP]

Thanks for submitting the PR! If you have any unit test and integration test results available, would be great to attach them as part of the PR description, thanks!

[tohtana]

@PKUWZP Thank you for the suggestion. I added testing commands and a link to the repro.