Add NVFP4 four-over-six grouped GEMM for MoE training - #9
Draft
wolfcomos wants to merge 4 commits into
Draft
Conversation
Grouped counterpart of four_over_six_mm for routed-expert layers (A (M, K) token groups x B (E, N, K) expert weights with group-end offsets), following TransformerEngine's GroupedLinear semantics: - per-tensor activations quantize with per-group global scales: the group amaxes expand to a per-row amax vector so one four_over_six_quantize call is bitwise identical to quantizing each split separately, and the forward is one F.scaled_grouped_mm with per-group second-level scales (bound-aware, not the 448-hardcoded helper). - row-scaled activations run a per-group loop of the dense four-over-six GEMM (FP32 output scaled by raw per-row amaxes) - the same shape TransformerEngine's general_grouped_gemm gives row-scaled NVFP4 at ea1a165d. - backward supports only the high_precision and dequantized overrides (TransformerEngine rejects four-over-six group quantization, so there is no grouped quantized backward): bf16 torch._grouped_mm on the saved originals or on dequantizations of the rowwise fprop operands. - ragged token groups zero-pad to 128-row alignment before quantization via pad_token_groups; padded rows quantize to zero codes and are sliced from the output. Validation (GB200, TE 2.19 devel container): 17 new tests pass, all bitwise (atol=rtol=0) - group-expanded amax == per-split quantize, grouped forward == dense per-group GEMMs (bitwise at the tested shapes for both weight blocks), row-scaled grouped == dense loop, both backward modes == manual grouped GEMMs on original/dequantized operands, ragged+padded == aligned construction, and the miles NVFP4 RL recipe point (row-scaled + MSE + bound 256 + 1x16 weights + dequantized backward). Full nvfp4_training suite: 461 passed / 79 skipped. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The fused pad/unpad CUDA kernels reject alignment_size != 32 and more than 32 groups (STD_TORCH_CHECK in mxfp8_extension.cpp), so any torchao build that ships them fails the 128-row-aligned ragged-group path on the first padded call. KernelPreference.EMULATED selects the torch implementation, which handles 128-row alignment and any expert count.
Inductor codegen of the fused unpack + broadcast-scale dequantize graph miscompiles the low-nibble lane (torch 2.14 nightly; the helpers compile correctly in isolation, the fused whole-function graph does not; fresh contiguous inputs reproduce it). The dequantized backward's contract is bitwise parity with the fprop operands, so the decode must keep eager numerics under compile: registering it as an opaque custom op with a fake impl does that, matching the CuTe DSL quantize op's pattern.
- test_grouped_backward_dequantized_ragged: bitwise dx/dw for unaligned groups + padding + dequantized backward (the torchtitan grouped-experts hook's production composition, previously only shape-checked). - test_linear_compile_backward_overrides / test_linear_compile_weight_block_1x16: fullgraph compile of the override backwards and 1x16 weights, bitwise vs eager (caught the inductor dequantize miscompile the previous commit pins). - test_grouped_compile: fullgraph compile of the per-tensor grouped op, forward and backward, following the mxfp8 grouped test's direct-compile pattern; skips on torch builds whose nonstrict_trace rejects autograd.Function outputs. Row-scaled grouped is eager-only by design.
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
1 participant
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.
Grouped counterpart of four_over_six_mm for routed-expert layers (A (M, K) token groups x B (E, N, K) expert weights with group-end offsets), following TransformerEngine's GroupedLinear semantics:
Validation (GB200, TE 2.19 devel container): 17 new tests pass, all bitwise (atol=rtol=0) - group-expanded amax == per-split quantize, grouped forward == dense per-group GEMMs (bitwise at the tested shapes for both weight blocks), row-scaled grouped == dense loop, both backward modes == manual grouped GEMMs on original/dequantized operands, ragged+padded == aligned construction, and the miles NVFP4 RL recipe point (row-scaled + MSE + bound 256 + 1x16 weights + dequantized backward). Full nvfp4_training suite: 461 passed / 79 skipped.