recipes/torch_compile_torch_function_modes

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

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

作者： Michael Lazos

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

Torch Function 模式，并将其与 torch.compile 结合使用，从而在追踪（trace）时重写 torch 算子（即 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 所带来的运行时开销。

• 有关 Torch Function 模式的其他示例和背景信息，请参阅 通过模式扩展 Torch API。