Improve device-init grouped linear module with single grouped weight support - #3224
Conversation
github-actions
Bot
added
the
community-contribution
aa3b9d1 to
47ba66a
Compare
ff7eee2 to
a43f70f
Compare
6536aa1 to
1169d6e
Compare
| is_grad_enabled = torch.is_grad_enabled() | ||
| num_gemms = self.num_gemms | ||
|
|
||
| if FP8GlobalStateManager.fp8_graph_capturing(): |
There was a problem hiding this comment.
Note: this code block was deleted because it was duplicated
|
/te-ci pytorch |
Greptile SummaryThis PR extends the Key changes:
Confidence Score: 4/5
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["GroupedLinear.forward(input, m_splits)"] --> B{single_grouped_weight?}
B -->|Yes| C["_get_weight_tensors() → [grouped_weight]<br/>num_weight_args = 1"]
B -->|No| D["_get_weight_tensors() → [w0, w1, ..., wN]<br/>num_weight_args = num_gemms"]
C --> E{is_module_grouped_tensor_path_supported?}
D --> E
E -->|Yes| F["_forward_grouped_tensor()"]
E -->|No - single_grouped_weight| G["RuntimeError: single param unsupported on legacy path"]
E -->|No - discrete weights| H["_forward_legacy_split_path()"]
F --> I["_prepare_weights_for_grouped_tensor_gemm()"]
I --> J{CUDA Graph safe?}
J -->|Yes, in-place| K["group_quantize(weight, output=cached_weight)<br/>noop_flag avoids re-quantize"]
J -->|No| L["group_quantize(weight) → new tensor"]
K --> M["general_grouped_gemm_for_grouped_tensor()"]
L --> M
F --> N["_prepare_bias_for_grouped_tensor_gemm()"]
N --> O{single_grouped_bias?}
O -->|Yes| P["Return packed bias directly"]
O -->|No| Q["Stack discrete biases → packed"]
M --> R["Forward output tensor"]
subgraph Backward
R --> S["_backward_grouped_tensor()"]
S --> T{single_grouped_weight?}
T -->|Yes| U["wgrad shape: grouped_weight.shape<br/>(num_gemms, out, in)"]
T -->|No| V["wgrad shape: per-expert (out, in)"]
U --> W["backward_dw(): delay wgrad compute"]
V --> W
W --> X{use_grouped_tensor AND bias grad?}
X -->|Yes| Y["RuntimeError: unexpected bias grad<br/>on grouped-tensor path"]
X -->|No| Z["Return dgrad, wgrad"]
end
Reviews (31): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
The biggest change in this PR is that TE is abandoning any attempt to make single_grouped_weight=True a general feature. Things must be exactly right, or we crash. Given how delicate and experimental this feature has been, I'm not opposed.
The second change is that users must opt-in to access the grouped GEMM kernel. This is also reasonable, since it has alignment requirements for m_splits and it's helpful having a way for users to accept that stricter contract.
We are experiencing many test failures. Given that single_grouped_weight is no longer a general feature, I think it's reasonable we move the corresponding tests to test_grouped_linear.py and test_grouped_mlp.py.
| raise ValueError( | ||
| "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 " | ||
| "tensor, or select grouped_gemm_backend='legacy'." | ||
| ) |
There was a problem hiding this comment.
I get that the h2d memcpy is suboptimal, but it's trivially easy to handle. Erroring out seems excessively rigid.
| raise ValueError( | |
| "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 " | |
| "tensor, or select grouped_gemm_backend='legacy'." | |
| ) | |
| m_splits = m_splits.to(device=device) |
We need to handle the d2h case anyways when the user has specified grouped_gemm_backend="grouped_tensor", but it's not supported and we fallback to split-quantize.
There was a problem hiding this comment.
But if the alignment is not provided in the first place, converting it to a device tensor also wouldn't work right, I am okay with another alignment check before adding this H2D copy.
b6a9482 to
25314b8
Compare
|
/te-ci pytorch |
07b2f18 to
bb6c1b9
Compare
| : tensor_base; | ||
| size_t tensor_base_for_scales = tensor_base; | ||
| size_t tensor_rows_for_scales = rows; | ||
| if constexpr (WITH_GEMM_SWIZZLED_SCALES && SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { |
There was a problem hiding this comment.
Note: this is for single weight quantize for mxfp8.
Before this change, the weight quantizer didn't have the first_dims because moe weights are uniform shape for both dimension. This will then trigger a CUDA illegal access because offsets_ptr=nullptr
2734746 to
5fc5db7
Compare
|
/te-ci pytorch |
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
45ee6d1 to
ddddc3f
Compare
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
|
/te-ci pytorch L1 |
| A_ptrs[idx] = a_is_discrete | ||
| ? a_multi_tensor_args.data_ptrs[idx] | ||
| : (a_base == nullptr ? nullptr : a_base + (a_offset * a_bits_per_elem) / 8); | ||
| B_ptrs[idx] = b_base == nullptr ? nullptr : b_base + (b_offset * b_bits_per_elem) / 8; |
There was a problem hiding this comment.
Do we need this?
| B_ptrs[idx] = b_base == nullptr ? nullptr : b_base + (b_offset * b_bits_per_elem) / 8; | |
| B_ptrs[idx] = b_base + (b_offset * b_bits_per_elem) / 8; |
There was a problem hiding this comment.
why not, it's nice to have (although the real fix is in A_ptrs[idx])
There was a problem hiding this comment.
Additional thread divergence. Anyway, first_dims and last_dims should already be in the descriptors and they will be zero and we wont read it from it anyway.
Same applies for a_base == null_ptr and and C and D as well
There was a problem hiding this comment.
Real issue was we were computing a_is_discrete and c/d_is_discrete incorrectly
There was a problem hiding this comment.
No harm keeping it here since, the root fix was:
Python-to-C++ conversion: A valid empty grouped storage has numel=0 and may have data_ptr=nullptr. It was converted as {nullptr, shape={0}}, which TE also uses to mean “storage is absent.” The converter now preserves the grouped tensor’s logical shape, so TE knows the empty storage exists and it won't early return in function select_grouped_operand.
Grouped GEMM setup: The code treated base_ptr == nullptr as meaning “discrete tensor list.” Therefore, valid all-empty grouped tensors were misclassified and read uninitialized discrete metadata. It now receives explicit a_is_discrete, c_is_discrete, and d_is_discrete flags.
There was a problem hiding this comment.
So the null_ptr check originally was kepts as a way of determining discrete flags which was wrong as you mentioned. Now given that we already have the flags. Adding nullptr check is no longer needed.
There is no severe harm I agree. However, having it is extra code to read and extra compute(although negligible). Anyway compute_grouped_tensor_offsets function above is going to give a_offset, b_offset etc all zeros and we are going to nullptr even without the null check
Also, as I was seeing I see (for consistency perspective)
int64_t a_offset = a_is_discrete ? 0 : compute_grouped_tensor_offset(A_meta, idx);
being done only for a and not for c and d.
| if dbias_packed is None: | ||
| dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, N) | ||
| grad_biases = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] | ||
| if ctx.single_grouped_bias: |
There was a problem hiding this comment.
I guess we should eventually unified single_grouped_weight and single_grouped_bias into one parameter.
There was a problem hiding this comment.
sigh yes, but currently it's not and technically they can be totally orthogonal to each other as well, I have tested all the combinations already so it's fine
| __all__ = ["GroupedLinear", "is_module_grouped_tensor_path_supported"] | ||
|
|
||
|
|
||
| def is_module_grouped_tensor_path_supported( |
There was a problem hiding this comment.
Can we put this in pytorch/utils.py and even reuse for the ops. The function should be identical right?
Given that we eventually want to unify the module/ops code, would be better to not duplicate these pieces now.
There was a problem hiding this comment.
Thanks for review, I made a decision to specifically not to do it because they are currently different implementations and I agree with @timmoon10 's proposal to de-duplicate it. But for now, the right thing to do is to keep them separate.
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
|
/te-ci pytorch L1 |
| fill_scale_ptr(a_scale_inv_ptrs, a_scale_base, A_meta, a_rowwise, a_scaling_mode); | ||
| } else { | ||
| a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx]; | ||
| a_scale_inv_ptrs[idx] = nullptr; |
There was a problem hiding this comment.
This is not needed. Would increase divergence
vthumbe1503
left a comment
vthumbe1503
left a comment
There was a problem hiding this comment.
LGTM. My comments are nit. But we should clean those up
|
/te-ci pytorch L1 |
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
for more information, see https://pre-commit.ci
5 participants
Description
Need this Mcore PR to make it work: NVIDIA/Megatron-LM#6000
Fixes numerical issues when using single weight for TE module grouped linear. Limit the single weight feature to the grouped tensor API instead of the legacy path.
E2E validation has been done and posted to the Megatron PR.
Performance:
Benchmarked using Qwen3.5 397B problem shape, tested the grouped linear FC1 FC2 layer time, including quantization kernels plus forward backward gemms.
Note: needs to pay extra attention to whether bias grad and weight grad are generated properly.
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: