Skip narrow dtypes in the mixed-precision cast - #8416
Open
sfc-gh-abkulkarni wants to merge 4 commits into
Open
Conversation
_cast_module_mixed_precision casts every parameter torch reports as floating point, which includes FP8, MX scales and NVFP4. Casting those discards the quantized encoding and doubles memory, and NVFP4 has no copy_ so it raises outright. Cast only the standard floating-point dtypes. Fixes deepspeedai#8414 Signed-off-by: Abhishek Kulkarni <abhishek.kulkarni@snowflake.com>
sfc-gh-abkulkarni
requested review from
loadams,
tjruwase and
tohtana
as code owners
September 4, 2026 18:38
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 750440e83e
ℹ️ 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".
The private-helper test alone did not verify that initialization preserves frozen narrow parameters, and it tied the contract to that helper. Add the same assertions to the end-to-end class for the fp8 pairs, at ZeRO 0 and 3. e8m0 and float4 stay on the helper test: NCCL rejects e8m0 in the parameter broadcast and float4 has no fill_, both before any casting happens. Move the quantized tensors into a submodule so a standard-dtype parameter comes first, since ZeRO-3 reads the model dtype from list(module.parameters())[0]. Signed-off-by: Abhishek Kulkarni <abhishek.kulkarni@snowflake.com>
The fp8 pairs are now asserted through deepspeed.initialize, so covering them twice added nothing. The helper test keeps only e8m0 and float4, which cannot reach the cast through initialize. Signed-off-by: Abhishek Kulkarni <abhishek.kulkarni@snowflake.com>
sfc-gh-truwase
approved these changes
Sep 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
_cast_module_mixed_precisioncasts every parameter torch reports as floating point. torch reports the narrow storage dtypes as floating point too, so quantized parameters get swept in:.to(bfloat16)float8_e4m3fn,float8_e5m2,*fnuzfloat8_e8m0fnufloat4_e2m1fn_x2NotImplementedError: "copy_" not implementedSo an FP8 LoRA base doubles its frozen weights and an NVFP4 model fails
deepspeed.initializeoutright.This casts only
float16/bfloat16/float32/float64rather than denylisting the narrow dtypes, so formats added later stay out of the cast by default. Those four plus the six narrow dtypes are exhaustively everythingis_floating_point()is true for on torch 2.11, so behaviour is unchanged for non-quantized models. Same guard applied to the buffer loop.Test fails on all six narrow dtypes without the fix.
Fixes #8414