[Pytorch] Add support for row-wise quanted input for grouped gemm#3244
[Pytorch] Add support for row-wise quanted input for grouped gemm#3244YangFei1990 wants to merge 2 commits into
Conversation
Signed-off-by: YangFei1990 <feiw@nvidia.com>
| and isinstance(input_quantizers[0], MXFP8Quantizer) | ||
| and isinstance(input_.quantizer, MXFP8Quantizer) |
There was a problem hiding this comment.
Should we also assert input_quantizers[0] == input_.quantizer?
There was a problem hiding this comment.
I think == comparison is too strict. It will compare all value fields, I think things like columnwise_usage or optimize_for_gemm could mismatch? Maybe we should explicit assert dype to be E4M3? But let me know if the dispatch output tensor's quantizer can match all value fields.
| with_columnwise: bool, | ||
| tensor_offsets: Optional[torch.Tensor] = None, | ||
| ) -> None: | ||
| """Prepare a rowwise-only MXFP8 grouped input for grouped GEMM (in place). |
There was a problem hiding this comment.
Looks like this helper actually "prepares" colwise besides rowwise. I think we should make the function name and docstring clearer and more consistent.
| format; they are converted to the GEMM-swizzled layout afterwards. | ||
| TODO: optimize and fuse the round-trips requant | ||
| """ | ||
| if grouped_x.rowwise_data is None: |
There was a problem hiding this comment.
Should we validate grouped_x.rowwise_scale as well?
| raise ValueError("Pre-quantized MXFP8 grouped input is missing rowwise data.") | ||
| if grouped_x._with_gemm_swizzled_scales: | ||
| raise NotImplementedError( | ||
| "Pre-quantized MXFP8 grouped input must have scales in compact format." |
| scale_shape = ( | ||
| round_up_to_nearest_multiple(total_tokens, 128), | ||
| round_up_to_nearest_multiple(cols // 32, 4), | ||
| ) |
There was a problem hiding this comment.
Do we need to do this while we already have 128 per-expert-alignment in Dispatch? The total tokens should already be divisible by 128.
For hidden size, there will be an assertion in Dispatch to expect % 512 = 0.
Here, I think we can add assertions instead.
|
|
||
| def prepare_prequantized_mxfp8_grouped_input( | ||
| grouped_x: GroupedTensorStorage, | ||
| quantizer: MXFP8Quantizer, |
There was a problem hiding this comment.
Actually, is there a GroupedQuantizer in TE/PyT? If so, we probably need to use it instead.
There was a problem hiding this comment.
I do not see a GroupedQuantizer in Pytorch. I see what used is tex.group_quantize. Let me know if I missed it. Also can you help me understand the benefit of GroupedQuantizer and how it compared with current Pytorch approach?
Signed-off-by: YangFei1990 <feiw@nvidia.com>
2 participants
Description
To support the NCCL EP FP8 dispatch, this PR introduce a capability to allow grouped gemm to take input that is already quantized in row-wise. We directly pass the row-wise quantized for forward grouped gemm and dequant + requant with col-wise to save for backward.
Fixes # (issue)
Type of change
Checklist: