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 22 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 wires
Confidence Score: 4/5The core dispatch and CUDA changes look structurally sound, but the gradient type contract for packed parameters is unverified and the PR description flags end-to-end convergence testing as a TODO. The backward returns plain tensors (3D view for wgrad, 2D tensor for bias grad) against GroupedTensorStorage parameters. PyTorch is permissive about .grad types, but Megatron/NeMo distributed training hooks may not be. The two near-identical path-support predicates have a subtle behavioral difference that could silently diverge. Files Needing Attention: transformer_engine/pytorch/module/grouped_linear.py backward path for single_grouped_weight/bias; transformer_engine/pytorch/ops/basic/grouped_linear.py path-support predicate vs. module version. Important Files Changed
Reviews (21): Last reviewed commit: "fix cublas setup meta data issue for emp..." | 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 |
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> # 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>
07b2f18 to
bb6c1b9
Compare
|
/te-ci pytorch |
| : 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
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
71a29ba to
13c60fb
Compare
Signed-off-by: zhongboz <zhongboz@nvidia.com>
2734746 to
5fc5db7
Compare
|
/te-ci pytorch |
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
|
/te-ci pytorch |
Signed-off-by: zhongboz <zhongboz@nvidia.com>
|
/te-ci pytorch |
| if (data.numel() == 0) { | ||
| // PyTorch may use a null pointer for a valid zero-sized allocation. Keep an explicitly | ||
| // provided empty storage distinguishable from the uninitialized {nullptr, {0}} sentinel. | ||
| shape = logical_shape; |
There was a problem hiding this comment.
Aligned offline. I would rather prefer something like {0,0} here instead of logical_shape to make sure TE Common infrastructure understands the rowwise data to be empty tensor instead of None(not being provided). Since logical_shape would raise consistency issues. For columnwise_data logical_shape would infact not signify the right shape
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.
TODO: test E2E convergence, unit test directly from Mcore.
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: