Fix NVFP4 QAT mixed precision #3501
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3501
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 Unrelated FailureAs of commit aea877d with merge base f3342a0 ( NEW FAILURE - The following job has failed:
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. |
meta-cla
bot
added
the
CLA Signed
andrewor14
force-pushed
the
fix-nvfp4-mixed-precision
branch
from
38c977e to
7826b54
Compare
| x = x.view(-1, x.shape[-1]) | ||
| else: | ||
| batch_size = None | ||
| fq = _NVFP4QuantizedForwardFakeQuantizedBackward.apply( |
There was a problem hiding this comment.
should the cast be inside _NVFP4QuantizedForwardFakeQuantizedBackward?
andrewor14
force-pushed
the
fix-nvfp4-mixed-precision
branch
from
7826b54 to
4a5913b
Compare
|
|
||
| ctx.save_for_backward(_input, weight) | ||
|
|
||
| return _addmm_nvfp4_dispatch( |
There was a problem hiding this comment.
why is this returning the wrong dtype?
There was a problem hiding this comment.
Seems to be caused by adding the bias:
Before this line result was bf16, after this line it's fp32. Do you think we should cast the bias here instead?
andrewor14
force-pushed
the
fix-nvfp4-mixed-precision
branch
2 times, most recently
from
a56a2c0 to
28e48b0
Compare
**Summary:** This commit adds support for bf16 activations + fp32 weights mixed precision for NVFP4 QAT, which previously threw 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_quantized_linear_mixed_precision ```
andrewor14
force-pushed
the
fix-nvfp4-mixed-precision
branch
from
28e48b0 to
aea877d
Compare
| """ | ||
|
|
||
| @staticmethod | ||
| @torch.amp.custom_fwd(device_type="cuda") |
There was a problem hiding this comment.
There was a problem hiding this comment.
is this still needed after the change in nvfp4 tensor?
There was a problem hiding this comment.
Yep, that one is just for fp32 inputs + fp32 weights during inference, this is for mixed precision during training
Summary: This commit adds support for bf16 activations + fp32 weights mixed precision for NVFP4 QAT, which previously threw a dtype assertion error:
Test Plan: