[Pytorch][Bug] Requires Grad doesnt flow through Autograd boundaries for QuantizedTensor - #3172
Conversation
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
… QuantizedTensor Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
|
/te-ci pytorch |
Greptile SummaryThis PR fixes two related issues with
Confidence Score: 4/5Safe to merge; the changes are narrowly scoped, well-motivated, and consistent with the patterns already established by the other three quantized tensor types. The core fixes — removing the stale
Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant U as User code
participant QT as QuantizedTensor (base)
participant FB as Float8BlockwiseQTensor
participant AF as _FromFloat8BlockwiseFunc
participant AG as PyTorch Autograd
U->>QT: "x (requires_grad=True)"
U->>FB: quantizer(x) → x_q (has grad_fn via quantizer)
Note over QT: No _requires_grad cache<br/>requires_grad reads from C++ directly
U->>FB: x_q.requires_grad → True (native C++)
U->>FB: x_q.dequantize()
alt torch.is_grad_enabled()
FB->>AF: _FromFloat8BlockwiseFunc.apply(self, dtype)
AF->>AG: Record forward pass
AG-->>U: "dequantized tensor (requires_grad=True, grad_fn set)"
else no_grad context
FB->>AF: _FromFloat8BlockwiseFunc.forward(None, self, dtype)
AF-->>U: dequantized tensor (no grad_fn)
end
U->>AG: loss.backward()
AG->>AF: backward(ctx, grad) → (grad, None)
AG->>FB: propagate grad through x_q.grad_fn
AG-->>U: x.grad populated
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant U as User code
participant QT as QuantizedTensor (base)
participant FB as Float8BlockwiseQTensor
participant AF as _FromFloat8BlockwiseFunc
participant AG as PyTorch Autograd
U->>QT: "x (requires_grad=True)"
U->>FB: quantizer(x) → x_q (has grad_fn via quantizer)
Note over QT: No _requires_grad cache<br/>requires_grad reads from C++ directly
U->>FB: x_q.requires_grad → True (native C++)
U->>FB: x_q.dequantize()
alt torch.is_grad_enabled()
FB->>AF: _FromFloat8BlockwiseFunc.apply(self, dtype)
AF->>AG: Record forward pass
AG-->>U: "dequantized tensor (requires_grad=True, grad_fn set)"
else no_grad context
FB->>AF: _FromFloat8BlockwiseFunc.forward(None, self, dtype)
AF-->>U: dequantized tensor (no grad_fn)
end
U->>AG: loss.backward()
AG->>AF: backward(ctx, grad) → (grad, None)
AG->>FB: propagate grad through x_q.grad_fn
AG-->>U: x.grad populated
Reviews (1): 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}
2 participants
Description
Bugs
Alternative way of reducing CPU overheads for attribute access of custom torch Tensor:
The main root cause for requires_grad attribute access of a QuantizedTensor being slow was our custom torch function implementation. The implementation doesnt do anything and uses the _disabled_torch_function_impl. But the problem was presence of python implementation resulted in us crossing C++ to python boundary twice.
When we do requires_grad access of custom torch tensor it goes through this function in pytorch C++ which checks the presence of torch function and actually goes back to python to execute it.
This can be avoided by assigning the sentinel value directly instead of defining a noop torch function
__torch_function__ = torch._C._disabled_torch_function_implTorchao does the same for their custom MXFP8Tensor here
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: