[Misc] Use torch.compile for GemmaRMSNorm #7642
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge). To run full CI, you can do one of these:
🚀 |
| residual: Optional[torch.Tensor] = None, | ||
| ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: | ||
| """PyTorch-native implementation equivalent to forward().""" | ||
| return self.forward_static(self.weight.data, self.variance_epsilon, x, |
There was a problem hiding this comment.
even if this is a static function, I'm not sure this self would cause problem here.
if you want to be safe, I think you can move this function outside of the class definition.
There was a problem hiding this comment.
I think this should be ok since it does not touch the states under self. I also checked that re-compilation does not happen after graph capturing, by monitoring the logs with TORCH_LOGS=guards. Also, the ShareGPT throughput benchmark shows 10~15% improvements.
This PR is a temporary solution to accelerate Gemma models. The PR can be reverted once #7110 is merged.