recipes/torch_compile_torch_function_modes

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

(beta) 利用 Torch Function 模式配合 torch.compile

作者： Michael Lazos

本指南介绍了如何使用一个关键的 torch 扩展点，

即 torch function 模式，与 torch.compile 协同工作，从而在追踪（trace）时覆盖 torch 算子（torch operators，也称为 ops）的行为，且不产生运行时开销。

注意

本指南需要 PyTorch 2.7.0 或更高版本。

重写 torch 算子 (torch.add -> torch.mul)

在本示例中，我们将使用 torch function 模式将所有加法操作重写为乘法操作。如果某个后端针对特定算子有自定义实现且需要进行调度，这种覆盖方式会非常常见。

import torch

# exit cleanly if we are on a device that doesn't support ``torch.compile``
if torch.cuda.get_device_capability() < (7, 0):
    print("Exiting because torch.compile is not supported on this device.")
    import sys
    sys.exit(0)

from torch.overrides import BaseTorchFunctionMode

# Define our mode, Note: ``BaseTorchFunctionMode``
# implements the actual invocation of func(..)
class AddToMultiplyMode(BaseTorchFunctionMode):
    def __torch_function__(self, func, types, args=(), kwargs=None):
        if func == torch.Tensor.add:
            func = torch.mul

        return super().__torch_function__(func, types, args, kwargs)

@torch.compile()
def test_fn(x, y):
    return x + y * x # Note: infix operators map to torch.Tensor.* methods

x = torch.rand(2, 2)
y = torch.rand_like(x)

with AddToMultiplyMode():
    z = test_fn(x, y)

assert torch.allclose(z, x * y * x)

# The mode can also be used within the compiled region as well like this:

@torch.compile()
def test_fn(x, y):
    with AddToMultiplyMode():
        return x + y * x # Note: infix operators map to torch.Tensor.* methods

x = torch.rand(2, 2)
y = torch.rand_like(x)
z = test_fn(x, y)

assert torch.allclose(z, x * y * x)

结论

在本指南中，我们演示了如何在 torch.compile 内部使用 torch function 模式来覆盖 torch.* 算子的行为。这使得用户能够利用 torch function 模式的扩展性优势，而无需在每次调用算子时承担调用 torch function 的运行时开销。

• 请参阅 使用 Modes 扩展 Torch API 以获取更多示例和关于 Torch Function 模式的背景知识。