Improve device-init grouped linear module with single grouped weight support - #3224
Improve device-init grouped linear module with single grouped weight support #3224zhongbozhu wants to merge 27 commits into
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
Confidence Score: 4/5
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward] --> B{single_grouped_weight\nor single_grouped_bias?}
B -- "Yes, use_grouped_tensor=False" --> ERR1[RuntimeError:\nrequires use_grouped_tensor=True]
B -- "No / Yes + use_grouped_tensor=True" --> C{debug mode?}
C -- "Yes + single param" --> ERR2[RuntimeError:\ndebug not supported]
C -- No --> D{use_grouped_tensor AND\nnot excluded features}
D -- No --> E[_GroupedLinear.forward\nlegacy split-quantize path]
D -- Yes --> F{Blackwell +\nFP8 block scaling?}
F -- Yes --> ERR3[RuntimeError:\nunsupported combo]
F -- No --> G[is_module_grouped_tensor\n_path_supported?]
G -- "False + single param" --> ERR4[RuntimeError:\ngrouped-tensor path needed]
G -- "False + discrete" --> E
G -- True --> H{CUDA m_splits?}
H -- No --> ERR5[ValueError]
H -- Yes --> I[_forward_grouped_tensor]
I --> J{single_grouped_weight?}
J -- Yes --> K[_prepare_weights_for\n_grouped_tensor_gemm\nGroupedTensorStorage path]
J -- No --> L[_prepare_weights_for\n_grouped_tensor_gemm\ndiscrete list path]
K --> M[general_grouped_gemm\n_for_grouped_tensor]
L --> M
M --> N{delay_wgrad?}
N -- Yes --> O[wgrad_store.put\ndefer weight GEMM]
N -- No --> P[grouped_gemm_wgrad\nimmediate]
O --> Q[backward_dw later\nunpacks wgrad_output]
P --> R[Return grad_input\n+ grad_weights\n+ grad_biases]
Reviews (27): Last reviewed commit: "resolve comments, enable save original i..." | 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: Zhongbo Zhu <zhongboz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> # Conflicts: # tests/pytorch/test_grouped_mlp.py
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Zhongbo Zhu <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: 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
| if test_recipe is not None: | ||
| grouped_weight = torch.cat(weight_tensors, dim=0) | ||
| weight_quantizer = _make_quantizer( | ||
| tex.DType.kFloat8E4M3, |
There was a problem hiding this comment.
Can we use TE_DType instead of tex.DType here. That is the recommendation now for the public user facing APIs.
| elif skip_fp8_weight_update is not None: | ||
| grouped_weight = tex.group_quantize( | ||
| source, | ||
| weight_quantizer, | ||
| num_gemms, | ||
| None, | ||
| noop_flag=skip_fp8_weight_update, | ||
| output=workspace, | ||
| ) | ||
| elif update_workspace: | ||
| grouped_weight = tex.group_quantize( | ||
| source, | ||
| weight_quantizer, | ||
| num_gemms, | ||
| None, | ||
| output=workspace, | ||
| ) |
There was a problem hiding this comment.
Can we combine the conditions to one?
| elif skip_fp8_weight_update is not None: | |
| grouped_weight = tex.group_quantize( | |
| source, | |
| weight_quantizer, | |
| num_gemms, | |
| None, | |
| noop_flag=skip_fp8_weight_update, | |
| output=workspace, | |
| ) | |
| elif update_workspace: | |
| grouped_weight = tex.group_quantize( | |
| source, | |
| weight_quantizer, | |
| num_gemms, | |
| None, | |
| output=workspace, | |
| ) | |
| elif skip_fp8_weight_update is not None or update_workspace: | |
| grouped_weight = tex.group_quantize( | |
| source, | |
| weight_quantizer, | |
| num_gemms, | |
| None, | |
| noop_flag=skip_fp8_weight_update, | |
| output=workspace, | |
| ) |
| 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.
4 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.
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: