Repository navigation
Grouped Linear CPU Overhead Reductions By Removing Unnecessary Python Logic - #3613
zhongbozhu wants to merge 10 commits into
Conversation
github-actions
Bot
added
the
community-contribution
d75052c to
033def7
Compare
| Note: please use `restore_from_func_ctx` instead if you are restoring tensors from a function context to make sure tensor_objects is detached and its memory can be freed | ||
| """ | ||
| tensor_objects = [] | ||
| cursor = 0 |
There was a problem hiding this comment.
Good catch, it turns out this logic was O(n^2) since slicing a list involves copying entries. We can still run afoul of n^2 if we have many QuantizedTensorStorages, but fixing that would require a larger structural change.
| grouped = is_op_fuser_grouped_tensor_path_supported( | ||
| FP8GlobalStateManager.get_fp8_recipe(), dtype | ||
| ) | ||
| input_quantizers = (self._input_quantizers[0],) if grouped else self._input_quantizers | ||
| grad_output_quantizers = ( | ||
| (self._grad_output_quantizers[0],) if grouped else self._grad_output_quantizers | ||
| ) |
There was a problem hiding this comment.
I'm wary of this approach since the decision between the grouped tensor impl and split-quantize impl is an internal implementation detail within GroupedLinear. If we apply the grouped MLP fusion, this conditional logic is irrelevant. Instead of exposing internal implementation details in pre_fuser_forward, perhaps it would be better to change our convention to only respect the first quantizer. pre_fuser_forward only sets three quantizers (input, weight, grad output), and within the split-quantize impl we manually set all of the quantizers to match the first. I'm not fully comfortable with this approach, but it's already the case that we assume quantizers are uniform.
| # Cache object references for this recipe generation, not tensor workspaces. | ||
| # Keep the per-expert slots for legacy split quantization and checkpoint metadata. | ||
| self._input_quantizers = tuple( | ||
| self.get_quantizer("forward", 2 * i) for i in range(self.num_groups) | ||
| ) | ||
| self._weight_quantizers = tuple( | ||
| self.get_quantizer("forward", 2 * i + 1) for i in range(self.num_groups) | ||
| ) | ||
| self._grad_output_quantizers = tuple( | ||
| self.get_quantizer("backward", i) for i in range(self.num_groups) | ||
| ) |
There was a problem hiding this comment.
Instead of caching, how about we only iterate over quantizers in the split-quantize path? This way the grouped tensor path only requires accessing a few quantizers. I'm worried how complicated grouped linear is and that adding unnecessary nonlocal logic will make it more challenging for humans to read.
There was a problem hiding this comment.
The code changes after removing the cache feels harder to read for me :)
033def7 to
a619742
Compare
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: zhongboz <zhongboz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
a619742 to
c31c2f6
Compare
Signed-off-by: zhongboz <zhongboz@nvidia.com>
for more information, see https://pre-commit.ci
 [Medium impact] Refactors grouped linear operations and quantizer configuration. Restore the required FP8 weight transpose before merging so training on Hopper does not fail during backward. Findings
SummaryThis PR reduces Python work around grouped linear calls.
Reviews (2) · Last reviewed commit: "resolve comments" · Reviewed by Greptile |
Signed-off-by: zhongboz <zhongboz@nvidia.com>
|
/te-ci pytorch |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| required_columnwise = True if columnwise_usage else None | ||
| if w._quantizer is not None: | ||
| w._quantizer.set_usage(rowwise=True, columnwise=required_columnwise) |
There was a problem hiding this comment.
On Hopper, a discrete FP8 weight created under torch.no_grad() has no transpose buffer. Removing w.update_usage(...) means _get_discrete_weights_for_gemm now changes only quantizer flags when training later needs input gradients. Those flags do not build the missing buffer. Forward succeeds, but backward calls grouped GEMM with layout="NN" and fails because the weight has no columnwise data.
Keep the call that builds required layouts for existing quantized weights.
Knowledge Base Used:
This branch has not been deployed
2 participants
Description
Please include a brief summary of the changes, relevant motivation and context.
Partly Fixes #2897
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: