- Notifications
You must be signed in to change notification settings - Fork386
Fix NVFP4 QAT mixed precision#3501
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.
Already on GitHub?Sign in to your account
base:main
Are you sure you want to change the base?
Conversation
pytorch-botbot commentedDec 17, 2025 • edited
Loading Uh oh!
There was an error while loading.Please reload this page.
edited
Uh oh!
There was an error while loading.Please reload this page.
🔗 Helpful Links🧪 See artifacts and rendered test results athud.pytorch.org/pr/pytorch/ao/3501
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit4a5913b with merge basef3342a0 ( BROKEN TRUNK - The following job failed but were present on the merge base:👉Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
38c977e to7826b54Compare| x=x.view(-1,x.shape[-1]) | ||
| else: | ||
| batch_size=None | ||
| fq=_NVFP4QuantizedForwardFakeQuantizedBackward.apply( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
should the cast be inside_NVFP4QuantizedForwardFakeQuantizedBackward?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
moved it in
**Summary:** This commit adds support for bf16 activations +fp32 weights mixed precision for NVFP4 QAT, which previouslythrew a dtype assertion error:```File "ao/torchao/prototype/qat/nvfp4.py", line 159, in forward assert fq.dtype == x.dtype```**Test Plan:**```python test/quantization/test_qat.py -k test_nvfp4_fake_quanitzed_linear_mixed_precision```
7826b54 to4a5913bCompare| ctx.save_for_backward(_input,weight) | ||
| return_addmm_nvfp4_dispatch( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
why is this returning the wrong dtype?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
Seems to be caused by adding the bias:
| result=result+bias |
Before this line result was bf16, after this line it's fp32. Do you think we should cast the bias here instead?
Summary: This commit adds support for bf16 activations + fp32 weights mixed precision for NVFP4 QAT, which previously threw a dtype assertion error:
Test Plan: