Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
+473
@@ -0,0 +1,473 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import collections
|
||||
|
||||
import torch
|
||||
from torch.testing._internal.common_utils import TEST_WITH_ROCM
|
||||
from torch.testing._internal.common_utils import TestCase
|
||||
|
||||
|
||||
class AutocastTestLists:
|
||||
def _rnn_cell_args(self, n, num_chunks, is_lstm, dev, dtype):
|
||||
input = (torch.randn((n, n), device=dev, dtype=torch.float32),)
|
||||
|
||||
hx = ((torch.randn((n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n), device=dev, dtype=torch.float32)) if is_lstm else
|
||||
torch.randn((n, n), device=dev, dtype=torch.float32),)
|
||||
|
||||
weights = (torch.randn((num_chunks * n, n), device=dev, dtype=torch.float32), # weight_ih
|
||||
torch.randn((num_chunks * n, n), device=dev, dtype=torch.float32), # weight_hh
|
||||
torch.randn((num_chunks * n), device=dev, dtype=torch.float32), # bias_ih
|
||||
torch.randn((num_chunks * n), device=dev, dtype=torch.float32)) # bias_hh
|
||||
|
||||
# returns args as a tuple
|
||||
return input + hx + weights
|
||||
|
||||
# Supplies ops and arguments for test_autocast_* in test/test_cuda.py
|
||||
def __init__(self, dev):
|
||||
super().__init__()
|
||||
n = 8
|
||||
# Utility arguments, created as one-element tuples
|
||||
pointwise0_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),)
|
||||
pointwise1_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),)
|
||||
pointwise2_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),)
|
||||
mat0_fp16 = (torch.randn((n, n), dtype=torch.float16, device=dev),)
|
||||
mat1_fp16 = (torch.randn((n, n), dtype=torch.float16, device=dev),)
|
||||
mat2_fp16 = (torch.randn((n, n), dtype=torch.float16, device=dev),)
|
||||
|
||||
dimsets = ((n, n, n), (n, n, n, n), (n, n, n, n, n))
|
||||
conv_args_fp32 = [(torch.randn(dimset, dtype=torch.float32, device=dev),
|
||||
torch.randn(dimset, dtype=torch.float32, device=dev))
|
||||
for dimset in dimsets]
|
||||
bias_fp32 = (torch.randn((n,), dtype=torch.float32, device=dev),)
|
||||
element0_fp32 = (torch.randn(1, dtype=torch.float32, device=dev),)
|
||||
pointwise0_fp32 = (torch.randn(n, dtype=torch.float32, device=dev),)
|
||||
pointwise1_fp32 = (torch.randn(n, dtype=torch.float32, device=dev),)
|
||||
mat0_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),)
|
||||
mat1_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),)
|
||||
mat2_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),)
|
||||
mat3_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),)
|
||||
|
||||
# The lists below organize ops that autocast needs to test.
|
||||
# self.list_name corresponds to test_autocast_list_name in test/test_cuda.py.
|
||||
# Each op is associated with a tuple of valid arguments.
|
||||
# In addition, cudnn conv ops are not supported on ROCm and hence will
|
||||
# be skipped by passing TEST_WITH_ROCM flag to those ops in self.torch_fp16 list.
|
||||
|
||||
# Some ops implement built-in type promotion. These don't need autocasting,
|
||||
# but autocasting relies on their promotion, so we include tests to double-check.
|
||||
self.torch_expect_builtin_promote = [
|
||||
("eq", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("ge", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("gt", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("le", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("lt", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("ne", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("add", pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("div", pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("mul", pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("cat", (pointwise0_fp16 + pointwise1_fp32,), torch.float32),
|
||||
("equal", pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("stack", (pointwise0_fp16 + pointwise1_fp32,), torch.float32),
|
||||
]
|
||||
self.methods_expect_builtin_promote = [
|
||||
("__eq__", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__ge__", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__gt__", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__le__", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__lt__", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__ne__", pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__add__", pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("__div__", pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("__mul__", pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
]
|
||||
|
||||
# The remaining lists organize ops that autocast treats explicitly.
|
||||
self.torch_fp16 = [
|
||||
# deprecated _convolution
|
||||
("_convolution", conv_args_fp32[1] + bias_fp32 + ((1, 1), (0, 0), (1, 1), False,
|
||||
(0, 0), 1, False, True, True)),
|
||||
# the current _convolution
|
||||
("_convolution", conv_args_fp32[1] + bias_fp32 + ((1, 1), (0, 0), (1, 1), False,
|
||||
(0, 0), 1, False, True, True, True)),
|
||||
("conv1d", conv_args_fp32[0]),
|
||||
("conv2d", conv_args_fp32[1]),
|
||||
("conv3d", conv_args_fp32[2]),
|
||||
("conv_tbc", conv_args_fp32[0] + bias_fp32),
|
||||
("conv_transpose1d", conv_args_fp32[0]),
|
||||
("conv_transpose2d", conv_args_fp32[1]),
|
||||
("conv_transpose3d", conv_args_fp32[2]),
|
||||
("convolution", conv_args_fp32[1] + bias_fp32 + ((1, 1), (0, 0), (1, 1), False, (0, 0), 1)),
|
||||
("cudnn_convolution", conv_args_fp32[1] + ((0, 0), (1, 1), (1, 1), 1, False, True, True), TEST_WITH_ROCM),
|
||||
("cudnn_convolution_transpose", conv_args_fp32[1] + ((0, 0), (0, 0), (1, 1),
|
||||
(1, 1), 1, False, True, True), TEST_WITH_ROCM),
|
||||
("prelu", pointwise0_fp32 + element0_fp32),
|
||||
("addmm", mat1_fp32 + mat2_fp32 + mat3_fp32),
|
||||
("addmv", pointwise0_fp32 + mat2_fp32 + pointwise1_fp32),
|
||||
("addr", mat0_fp32 + pointwise0_fp32 + pointwise1_fp32),
|
||||
("matmul", mat0_fp32 + mat1_fp32),
|
||||
("einsum", "bkhd,bqhd->bqkh", mat0_fp32 + mat1_fp32),
|
||||
("mm", mat0_fp32 + mat1_fp32),
|
||||
("mv", mat0_fp32 + pointwise0_fp32),
|
||||
("chain_matmul", mat0_fp32 + mat1_fp32 + mat2_fp32),
|
||||
("addbmm", mat0_fp32 + (torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32))),
|
||||
("baddbmm", (torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32))),
|
||||
("bmm", (torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32))),
|
||||
# _thnn_fused_lstm_cell and _thnn_fused_gru_cell are not Python-exposed as far as I can tell.
|
||||
# ("_thnn_fused_lstm_cell", mat0_fp32 + mat1_fp32 + mat2_fp32 + pointwise0_fp32 + pointwise1_fp32),
|
||||
# ("_thnn_fused_gru_cell", mat0_fp32 + mat1_fp32 + mat2_fp32 + pointwise0_fp32 + pointwise1_fp32),
|
||||
("lstm_cell", self._rnn_cell_args(n, num_chunks=4, is_lstm=True, dev=dev, dtype=torch.float32)),
|
||||
("gru_cell", self._rnn_cell_args(n, num_chunks=3, is_lstm=False, dev=dev, dtype=torch.float32)),
|
||||
("rnn_tanh_cell", self._rnn_cell_args(n, num_chunks=1, is_lstm=False, dev=dev, dtype=torch.float32)),
|
||||
("rnn_relu_cell", self._rnn_cell_args(n, num_chunks=1, is_lstm=False, dev=dev, dtype=torch.float32)),
|
||||
]
|
||||
self.torch_fp32 = [
|
||||
("acos", (pointwise0_fp16[0].clamp(-.9, 0.9),)),
|
||||
("asin", (pointwise0_fp16[0].clamp(-.9, 0.9),)),
|
||||
("cosh", pointwise0_fp16),
|
||||
("erfinv", (pointwise0_fp16[0].clamp(-.9, .9),)),
|
||||
("exp", pointwise0_fp16),
|
||||
("expm1", pointwise0_fp16),
|
||||
("log", (pointwise0_fp16[0].clamp(0.1, 100.0),)),
|
||||
("log10", (pointwise0_fp16[0].clamp(0.1, 100.0),)),
|
||||
("log2", (pointwise0_fp16[0].clamp(0.1, 100.0),)),
|
||||
("log1p", (pointwise0_fp16[0].clamp(-0.9, 100.0),)),
|
||||
("reciprocal", pointwise0_fp16),
|
||||
("rsqrt", (pointwise0_fp16[0].clamp(0.0, 100.0),)),
|
||||
("sinh", pointwise0_fp16),
|
||||
("tan", (pointwise0_fp16[0].clamp(-3.1 / 2, 3.1 / 2),)),
|
||||
("pow", ((pointwise0_fp16[0] + 1.).clamp(0.0, 100.0),) + pointwise1_fp16),
|
||||
("pow", ((pointwise0_fp16[0] + 1.).clamp(0.0, 100.0),) + (1.7,)),
|
||||
# ("pow", (1.7,) + pointwise0_fp16), # This variant has a backend, but is not documented in the API.
|
||||
("softmax", pointwise0_fp16 + (0,)),
|
||||
("log_softmax", pointwise0_fp16 + (0,)),
|
||||
("layer_norm", pointwise0_fp16 + ((pointwise0_fp16[0].numel(),),)),
|
||||
("rms_norm", pointwise0_fp16 + ((pointwise0_fp16[0].numel(),),)),
|
||||
("group_norm", mat0_fp16 + (1,)),
|
||||
("norm", pointwise0_fp16),
|
||||
("norm", pointwise0_fp16, {"dim": 0}),
|
||||
# these need magma
|
||||
# ("norm", mat0_fp16, {"p": "nuc"}),
|
||||
# ("norm", mat0_fp16, {"p": "nuc", "dim": 0}),
|
||||
("norm", pointwise0_fp16, {"p": 1}),
|
||||
("norm", pointwise0_fp16, {"p": 1, "dim": 0}),
|
||||
("cosine_similarity", mat0_fp16 + mat1_fp16),
|
||||
("poisson_nll_loss", mat0_fp16 + mat1_fp16 + (True, False, 1.e-8, torch.nn._reduction.get_enum('mean'))),
|
||||
("cosine_embedding_loss", (torch.tensor([[1, 2, 3]], device=dev, dtype=torch.float16),
|
||||
torch.tensor([[1, 3, 4]], device=dev, dtype=torch.float16),
|
||||
torch.tensor([1], device=dev, dtype=torch.int))),
|
||||
("hinge_embedding_loss", mat0_fp16 + (torch.ones(n, device=dev, dtype=torch.int),)),
|
||||
("kl_div", mat0_fp16 + (torch.rand((n, n), device=dev, dtype=torch.float16),)),
|
||||
("margin_ranking_loss", mat0_fp16 + mat1_fp16 + (torch.ones((n,), device=dev, dtype=torch.float16),)),
|
||||
("triplet_margin_loss", mat0_fp16 + mat1_fp16 + mat2_fp16),
|
||||
("binary_cross_entropy_with_logits", mat0_fp16 + (torch.rand((n, n), device=dev, dtype=torch.float16),)),
|
||||
("cumprod", pointwise0_fp16 + (0,)),
|
||||
("cumsum", pointwise0_fp16 + (0,)),
|
||||
("dist", pointwise0_fp16 + pointwise1_fp16),
|
||||
("pdist", mat0_fp16),
|
||||
("cdist", mat0_fp16 + mat1_fp16),
|
||||
("prod", pointwise0_fp16),
|
||||
("prod", pointwise0_fp16 + (0,)),
|
||||
("renorm", mat0_fp16 + (2, 0, 1.0)),
|
||||
("sum", pointwise0_fp16),
|
||||
("sum", mat0_fp16 + (1,)),
|
||||
("logsumexp", mat0_fp16 + (1,)),
|
||||
]
|
||||
self.torch_need_autocast_promote = [
|
||||
("addcdiv", pointwise0_fp32 + pointwise1_fp16 + (pointwise2_fp16[0].clamp(0.1, 100),)),
|
||||
("addcmul", pointwise0_fp32 + pointwise1_fp16 + pointwise2_fp16),
|
||||
("atan2", pointwise0_fp32 + (pointwise1_fp16[0].clamp(0.1, 100),)),
|
||||
("bilinear", (torch.randn((1, 2), dtype=torch.float16, device=dev),
|
||||
torch.randn((1, 2), dtype=torch.float32, device=dev),
|
||||
torch.randn((1, 2, 2), dtype=torch.float16, device=dev),
|
||||
torch.randn((1,), dtype=torch.float32, device=dev))),
|
||||
("cross", (torch.randn(3, dtype=torch.float32, device=dev),
|
||||
torch.randn(3, dtype=torch.float16, device=dev))),
|
||||
("dot", pointwise0_fp16 + pointwise1_fp32),
|
||||
("vdot", pointwise0_fp16 + pointwise1_fp32),
|
||||
("grid_sampler", (torch.randn((2, 3, 33, 22), dtype=torch.float16, device=dev),
|
||||
torch.randn((2, 22, 11, 2), dtype=torch.float32, device=dev),
|
||||
0, 0, False)),
|
||||
("index_put", pointwise0_fp32 + ((torch.tensor([1], device=dev, dtype=torch.long),),
|
||||
torch.randn(1, device=dev, dtype=torch.float16))),
|
||||
("index_put", pointwise0_fp16 + ((torch.tensor([1], device=dev, dtype=torch.long),),
|
||||
torch.randn(1, device=dev, dtype=torch.float32))),
|
||||
("tensordot", (torch.randn((2, 2, 2), dtype=torch.float32, device=dev),
|
||||
torch.randn((2, 2, 2), dtype=torch.float16, device=dev))),
|
||||
("scatter_add", (torch.zeros(2, 2, 2, dtype=torch.float32, device=dev),
|
||||
0,
|
||||
torch.randint(0, 2, (2, 2, 2), device=dev),
|
||||
torch.randn((2, 2, 2), dtype=torch.float16, device=dev))),
|
||||
("scatter_add", (torch.zeros(2, 2, 2, dtype=torch.float16, device=dev),
|
||||
0,
|
||||
torch.randint(0, 2, (2, 2, 2), device=dev),
|
||||
torch.randn((2, 2, 2), dtype=torch.float32, device=dev))),
|
||||
]
|
||||
self.nn_fp16 = [
|
||||
("linear", mat0_fp32 + mat1_fp32 + mat2_fp32),
|
||||
]
|
||||
self.nn_fp32 = [
|
||||
("softplus", pointwise0_fp16),
|
||||
("nll_loss", (torch.rand((n, n), device=dev, dtype=torch.float),
|
||||
torch.zeros((n,), device=dev, dtype=torch.long))),
|
||||
("nll_loss2d", (torch.rand((n, n, n, n), device=dev, dtype=torch.half),
|
||||
torch.zeros((n, n, n), device=dev, dtype=torch.long))),
|
||||
("l1_loss", mat0_fp16 + mat1_fp16),
|
||||
("smooth_l1_loss", mat0_fp16 + mat1_fp16),
|
||||
("mse_loss", mat0_fp16 + mat1_fp16),
|
||||
("multilabel_margin_loss", mat0_fp16 + (torch.ones((n, n), device=dev, dtype=torch.long),)),
|
||||
("soft_margin_loss", mat0_fp16 + (torch.ones((n, n), device=dev, dtype=torch.long),)),
|
||||
("multi_margin_loss", mat0_fp16 + (torch.ones((n,), device=dev, dtype=torch.long),)),
|
||||
]
|
||||
self.linalg_fp16 = [
|
||||
("linalg_vecdot", mat0_fp32 + mat0_fp32),
|
||||
("linalg_multi_dot", (mat0_fp32 + mat1_fp32 + mat2_fp32,)),
|
||||
]
|
||||
self.methods_fp16 = [
|
||||
("__matmul__", mat0_fp32 + mat1_fp32)
|
||||
]
|
||||
self.methods_fp32 = [
|
||||
("__pow__", (torch.rand(n, device=dev, dtype=torch.float16), 1.5)),
|
||||
]
|
||||
self.banned = [
|
||||
("binary_cross_entropy", (torch.rand((n, n), device=dev, dtype=torch.float32),
|
||||
torch.rand((n, n), device=dev, dtype=torch.float32)), torch._C._nn),
|
||||
]
|
||||
|
||||
|
||||
class AutocastCPUTestLists:
|
||||
# Supplies ops and arguments for test_autocast_* in test/test_cpu.py
|
||||
def __init__(self, dev):
|
||||
super().__init__()
|
||||
n = 8
|
||||
# Utility arguments, created as one-element tuples
|
||||
pointwise0_bf16 = (torch.randn(n, dtype=torch.bfloat16, device=dev),)
|
||||
pointwise1_bf16 = (torch.randn(n, dtype=torch.bfloat16, device=dev),)
|
||||
mat0_bf16 = (torch.randn((n, n), dtype=torch.bfloat16, device=dev),)
|
||||
mat1_bf16 = (torch.randn((n, n), dtype=torch.bfloat16, device=dev),)
|
||||
mat2_bf16 = (torch.randn((n, n), dtype=torch.bfloat16, device=dev),)
|
||||
|
||||
pointwise0_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),)
|
||||
pointwise1_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),)
|
||||
|
||||
dummy_dimsets = ((n,), (n, n), (n, n, n), (n, n, n, n), (n, n, n, n, n))
|
||||
|
||||
dummy_bf16 = [(torch.randn(dimset, dtype=torch.bfloat16, device=dev),)
|
||||
for dimset in dummy_dimsets]
|
||||
|
||||
dimsets = ((n, n, n), (n, n, n, n), (n, n, n, n, n))
|
||||
conv_args_fp32 = [(torch.randn(dimset, dtype=torch.float32, device=dev),
|
||||
torch.randn(dimset, dtype=torch.float32, device=dev))
|
||||
for dimset in dimsets]
|
||||
|
||||
element0_fp32 = (torch.randn(1, dtype=torch.float32, device=dev),)
|
||||
pointwise0_fp32 = (torch.randn(n, dtype=torch.float32, device=dev),)
|
||||
pointwise1_fp32 = (torch.randn(n, dtype=torch.float32, device=dev),)
|
||||
mat0_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),)
|
||||
mat1_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),)
|
||||
mat2_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),)
|
||||
mat3_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),)
|
||||
|
||||
dummy_fp32 = [ # noqa: F841
|
||||
(torch.randn(dimset, dtype=torch.float32, device=dev),)
|
||||
for dimset in dummy_dimsets
|
||||
]
|
||||
# The lists below organize ops that autocast needs to test.
|
||||
# self.list_name corresponds to test_autocast_list_name in test/test_cpu.py.
|
||||
# Each op is associated with a tuple of valid arguments.
|
||||
|
||||
# Some ops implement built-in type promotion. These don't need autocasting,
|
||||
# but autocasting relies on their promotion, so we include tests to double-check.
|
||||
self.torch_expect_builtin_promote = [
|
||||
("eq", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("ge", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("gt", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("le", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("lt", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("ne", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("add", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("div", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("mul", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
]
|
||||
|
||||
self.methods_expect_builtin_promote = [
|
||||
("__eq__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__ge__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__gt__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__le__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__lt__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__ne__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool),
|
||||
("__add__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("__div__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
("__mul__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32),
|
||||
]
|
||||
# The remaining lists organize ops that autocast treats explicitly.
|
||||
self.torch_16 = [
|
||||
("conv1d", conv_args_fp32[0]),
|
||||
("conv2d", conv_args_fp32[1]),
|
||||
("conv3d", conv_args_fp32[2]),
|
||||
("bmm", (torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32))),
|
||||
("mm", mat0_fp32 + mat1_fp32),
|
||||
("matmul", mat0_fp32 + mat1_fp32),
|
||||
("baddbmm", (torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32))),
|
||||
("addmm", mat1_fp32 + mat2_fp32 + mat3_fp32),
|
||||
("_addmm_activation", mat1_fp32 + mat2_fp32 + mat3_fp32, {"beta": 1, "alpha": 1, "use_gelu": True}),
|
||||
("addbmm", mat0_fp32 + (torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32))),
|
||||
("conv_tbc", (torch.randn((10, 7, 3), device=dev, dtype=torch.float32),
|
||||
torch.randn((5, 3, 5), device=dev, dtype=torch.float32),
|
||||
torch.randn(5, device=dev, dtype=torch.float32),
|
||||
0)),
|
||||
("conv_transpose1d", conv_args_fp32[0]),
|
||||
("conv_transpose2d", conv_args_fp32[1]),
|
||||
("conv_transpose3d", conv_args_fp32[2]),
|
||||
("prelu", pointwise0_fp32 + element0_fp32),
|
||||
("_native_multi_head_attention", (torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n, n), device=dev, dtype=torch.float32),
|
||||
n, 4, torch.randn((3 * n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((3 * n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n, n), device=dev, dtype=torch.float32),
|
||||
torch.randn((n), device=dev, dtype=torch.float32))),
|
||||
]
|
||||
self.torch_fp32 = [
|
||||
("poisson_nll_loss", mat0_bf16 + mat1_bf16 + (True, False, 1.e-8, torch.nn._reduction.get_enum('mean'))),
|
||||
("cosine_embedding_loss", (torch.tensor([[1, 2, 3]], device=dev, dtype=torch.bfloat16),
|
||||
torch.tensor([[1, 3, 4]], device=dev, dtype=torch.bfloat16),
|
||||
torch.tensor([1], device=dev, dtype=torch.int))),
|
||||
("hinge_embedding_loss", mat0_bf16 + (torch.ones(n, device=dev, dtype=torch.int),)),
|
||||
("margin_ranking_loss", mat0_bf16 + mat1_bf16 + (torch.ones((n,), device=dev, dtype=torch.bfloat16),)),
|
||||
("triplet_margin_loss", mat0_bf16 + mat1_bf16 + mat2_bf16),
|
||||
("binary_cross_entropy_with_logits", mat0_bf16 + (torch.rand((n, n), device=dev, dtype=torch.bfloat16),)),
|
||||
]
|
||||
self.nn_16 = [
|
||||
("linear", mat0_fp32 + mat1_fp32, {}),
|
||||
]
|
||||
self.nn_fp32 = [
|
||||
("avg_pool3d", dummy_bf16[3], {"kernel_size": (3, 3, 3), "stride": (1, 1, 1)}),
|
||||
("binary_cross_entropy", (torch.rand((n, n), device=dev, dtype=torch.bfloat16),) +
|
||||
(torch.rand((n, n), device=dev, dtype=torch.bfloat16),)),
|
||||
("reflection_pad1d", dummy_bf16[2], {"padding": (3, 3)}),
|
||||
("nll_loss", (torch.rand((n, n), device=dev, dtype=torch.bfloat16),
|
||||
torch.zeros((n,), device=dev, dtype=torch.long))),
|
||||
("nll_loss2d", (torch.rand((n, n, n, n), device=dev, dtype=torch.bfloat16),
|
||||
torch.zeros((n, n, n), device=dev, dtype=torch.long))),
|
||||
("l1_loss", mat0_bf16 + mat1_bf16),
|
||||
("smooth_l1_loss", mat0_bf16 + mat1_bf16),
|
||||
("mse_loss", mat0_bf16 + mat1_bf16),
|
||||
("multilabel_margin_loss", mat0_bf16 + (torch.ones((n, n), device=dev, dtype=torch.long),)),
|
||||
("soft_margin_loss", mat0_bf16 + (torch.ones((n, n), device=dev, dtype=torch.long),)),
|
||||
("multi_margin_loss", mat0_bf16 + (torch.ones((n,), device=dev, dtype=torch.long),)),
|
||||
("huber_loss", mat0_bf16 + mat1_bf16),
|
||||
]
|
||||
self.torch_need_autocast_promote = [
|
||||
("cat", (pointwise0_bf16 + pointwise1_fp32,), (pointwise0_fp16 + pointwise1_fp32,)),
|
||||
("stack", (pointwise0_bf16 + pointwise1_fp32,), (pointwise0_fp16 + pointwise1_fp32,)),
|
||||
]
|
||||
|
||||
|
||||
class TestAutocast(TestCase):
|
||||
def args_maybe_kwargs(self, op_with_args):
|
||||
if len(op_with_args) == 2:
|
||||
return op_with_args[0], op_with_args[1], {}
|
||||
else:
|
||||
return op_with_args[0], op_with_args[1], op_with_args[2]
|
||||
|
||||
def _run_autocast_outofplace(
|
||||
self,
|
||||
op,
|
||||
args,
|
||||
run_as_type,
|
||||
device,
|
||||
out_type=None,
|
||||
module=torch,
|
||||
add_kwargs=None,
|
||||
amp_dtype=torch.bfloat16,
|
||||
):
|
||||
# helper to cast args
|
||||
def cast(val, to_type):
|
||||
if isinstance(val, torch.Tensor):
|
||||
return val.to(to_type) if val.is_floating_point() else val
|
||||
elif isinstance(val, collections.abc.Iterable):
|
||||
return type(val)(cast(v, to_type) for v in val)
|
||||
else:
|
||||
return val
|
||||
|
||||
if add_kwargs is None:
|
||||
add_kwargs = {}
|
||||
|
||||
self.assertFalse(torch.is_autocast_enabled(device_type=device))
|
||||
with torch.amp.autocast(device_type=device, dtype=amp_dtype):
|
||||
self.assertTrue(torch.is_autocast_enabled(device_type=device))
|
||||
|
||||
out_type = out_type if out_type is not None else run_as_type
|
||||
output = output_method = None
|
||||
|
||||
# Try module.* variant, if requested:
|
||||
if module is not None and hasattr(module, op):
|
||||
output = getattr(module, op)(*args, **add_kwargs)
|
||||
if isinstance(output, torch.Tensor):
|
||||
self.assertTrue(
|
||||
out_type == output.dtype,
|
||||
f"autocast for torch.{op} produced {output.dtype}, should produce {out_type}",
|
||||
)
|
||||
# Try Tensor.* variant:
|
||||
if hasattr(torch.Tensor, op):
|
||||
output_method = getattr(args[0], op)(*args[1:], **add_kwargs)
|
||||
if isinstance(output_method, torch.Tensor):
|
||||
self.assertTrue(
|
||||
out_type == output_method.dtype,
|
||||
f"autocast for torch.{op} produced {output_method.dtype}, should produce torch.{out_type}",
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
(output is not None) or (output_method is not None),
|
||||
f"{op} not found as an attribute on either Tensor or the requested module {module}",
|
||||
)
|
||||
|
||||
# Accounts for ops that return Tensors, iterables, and other non-Tensors.
|
||||
# For example, lstm_cell returns a tuple and equal returns bool.
|
||||
def compare(first, second):
|
||||
if isinstance(first, torch.Tensor):
|
||||
return torch.equal(first, second)
|
||||
elif isinstance(first, collections.abc.Iterable):
|
||||
return all(compare(f, s) for f, s in zip(first, second, strict=False))
|
||||
else:
|
||||
return first == second
|
||||
|
||||
# If both torch.* and Tensor.* variants were found, check outputs are identical
|
||||
if (output is not None) and (output_method is not None):
|
||||
self.assertTrue(type(output) is type(output_method))
|
||||
comparison = compare(output, output_method)
|
||||
self.assertTrue(
|
||||
comparison, f"torch.{op} result did not match Tensor.{op} result"
|
||||
)
|
||||
|
||||
# Compare numerics to Python-side "autocasting" that (we expect) does the same thing
|
||||
# as the C++-side autocasting, and should be bitwise accurate.
|
||||
output_to_compare = output if output is not None else output_method
|
||||
with torch.amp.autocast(device_type=device, enabled=False):
|
||||
self.assertFalse(
|
||||
torch.is_autocast_enabled(device_type=device)
|
||||
)
|
||||
|
||||
if module is not None and hasattr(module, op):
|
||||
control = getattr(module, op)(
|
||||
*cast(args, run_as_type), **add_kwargs
|
||||
)
|
||||
else:
|
||||
control = getattr(args[0].to(run_as_type), op)(
|
||||
*cast(args[1:], run_as_type), **add_kwargs
|
||||
)
|
||||
self.assertTrue(type(output_to_compare) is type(control))
|
||||
comparison = compare(output_to_compare, control)
|
||||
self.assertTrue(comparison, f"torch.{op} result did not match control")
|
||||
self.assertTrue(torch.is_autocast_enabled(device_type=device))
|
||||
self.assertFalse(torch.is_autocast_enabled(device_type=device))
|
||||
+635
@@ -0,0 +1,635 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch
|
||||
from functools import partial
|
||||
from torch.testing import make_tensor
|
||||
from torch.testing._internal.opinfo.core import (
|
||||
OpInfo,
|
||||
SampleInput,
|
||||
)
|
||||
from torch.testing._internal.common_dtype import all_types_and
|
||||
import numpy as np
|
||||
|
||||
# Note: [autograd.Function db]
|
||||
#
|
||||
# This is a collection of autograd.Function test cases written as OpInfos
|
||||
# so they can easily be consumed by OpInfo-based tests to check if a subsystem
|
||||
# supports autograd.Function.
|
||||
#
|
||||
# Axes:
|
||||
# - saves {output, input, intermediate, non-tensor}
|
||||
# - {inputs, output} x {single tensor, tensors, arbitrary objects}
|
||||
# - Uses {mark_dirty, mark_non_differentiable, once_differentiable}
|
||||
|
||||
|
||||
def to_numpy(tensor):
|
||||
return tensor.cpu().numpy()
|
||||
|
||||
|
||||
class NumpyCube(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(input):
|
||||
input_np = to_numpy(input)
|
||||
dinput = torch.tensor(3 * input_np ** 2, device=input.device)
|
||||
return torch.tensor(input_np ** 3, device=input.device), dinput
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, output):
|
||||
ctx.save_for_backward(inputs[0], output[1])
|
||||
ctx.save_for_forward(inputs[0], output[1])
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output, grad_saved):
|
||||
input, dinput = ctx.saved_tensors
|
||||
return NumpyMul.apply(grad_output, dinput) + 6 * NumpyMul.apply(grad_saved, input)
|
||||
|
||||
@staticmethod
|
||||
def vmap(info, in_dims, input):
|
||||
result = NumpyCube.apply(input)
|
||||
return result, (in_dims[0], in_dims[0])
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, input_tangent):
|
||||
input, dinput = ctx.saved_tensors
|
||||
return NumpyMul.apply(input_tangent, dinput), 6 * NumpyMul.apply(input_tangent, input)
|
||||
|
||||
|
||||
class CubeGenVmap(torch.autograd.Function):
|
||||
generate_vmap_rule = True
|
||||
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return x ** 3, 3 * x ** 2
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, outputs):
|
||||
ctx.save_for_backward(inputs[0], outputs[1])
|
||||
ctx.save_for_forward(inputs[0], outputs[1])
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output, grad_saved):
|
||||
_input, dinput = ctx.saved_tensors
|
||||
result = grad_output * dinput + 6 * dinput
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, input_tangent):
|
||||
input, dinput = ctx.saved_tensors
|
||||
return MulGenVmap.apply(input_tangent, dinput), 6 * NumpyMul.apply(input_tangent, input)
|
||||
|
||||
|
||||
def sample_inputs_numpy_cube(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
yield SampleInput(make_arg(1, low=0.8, high=2), args=())
|
||||
|
||||
|
||||
class NumpyCubeNotComposable(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(input):
|
||||
input_np = to_numpy(input)
|
||||
return torch.tensor(input_np ** 3, device=input.device), input_np
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, output):
|
||||
_, input_np = output
|
||||
ctx.input_np = input_np
|
||||
ctx.device = inputs[0].device
|
||||
|
||||
@staticmethod
|
||||
@torch.autograd.function.once_differentiable
|
||||
def backward(ctx, grad_output, grad_saved):
|
||||
result_np = 3 * (ctx.input_np ** 2)
|
||||
return torch.tensor(result_np, device=ctx.device)
|
||||
|
||||
|
||||
class NumpyMul(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(x, y):
|
||||
return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device)
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, output):
|
||||
ctx.save_for_backward(*inputs)
|
||||
ctx.save_for_forward(*inputs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
x, y = ctx.saved_tensors
|
||||
gx = None
|
||||
if ctx.needs_input_grad[0]:
|
||||
gx = NumpyMul.apply(grad_output, y)
|
||||
gy = None
|
||||
if ctx.needs_input_grad[1]:
|
||||
gy = NumpyMul.apply(grad_output, x)
|
||||
return gx, gy
|
||||
|
||||
@staticmethod
|
||||
def vmap(info, in_dims, x, y):
|
||||
x_bdim, y_bdim = in_dims
|
||||
x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1)
|
||||
y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1)
|
||||
result = NumpyMul.apply(x, y)
|
||||
result = result.movedim(-1, 0)
|
||||
return result, 0
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, y_tangent):
|
||||
x, y = ctx.saved_tensors
|
||||
return x_tangent * y + y_tangent * x
|
||||
|
||||
def sample_inputs_numpy_mul(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
# Broadcasting
|
||||
yield SampleInput(make_arg(4, low=0.9, high=2), args=(make_arg(3, 4, low=0.9, high=2),))
|
||||
|
||||
def sample_inputs_numpy_mul_scalar(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
yield SampleInput(make_arg(4, low=0.9, high=2), args=(), kwargs={"scalar": 3.14})
|
||||
|
||||
class MulGenVmap(torch.autograd.Function):
|
||||
generate_vmap_rule = True
|
||||
|
||||
@staticmethod
|
||||
def forward(x, y):
|
||||
return x * y
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, outputs):
|
||||
ctx.save_for_backward(*inputs)
|
||||
ctx.save_for_forward(*inputs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
x, y = ctx.saved_tensors
|
||||
gx = None
|
||||
if ctx.needs_input_grad[0]:
|
||||
gx = MulGenVmap.apply(grad_output, y)
|
||||
gy = None
|
||||
if ctx.needs_input_grad[1]:
|
||||
gy = MulGenVmap.apply(grad_output, x)
|
||||
return gx, gy
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, y_tangent):
|
||||
x, y = ctx.saved_tensors
|
||||
return x_tangent * y + y_tangent * x
|
||||
|
||||
|
||||
class NumpyExp_(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
x_np = to_numpy(x)
|
||||
np.exp(x_np, x_np)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, output):
|
||||
x, = inputs
|
||||
ctx.mark_dirty(x)
|
||||
ctx.save_for_backward(output)
|
||||
ctx.save_for_forward(output)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
output, = ctx.saved_tensors
|
||||
return NumpyMul.apply(grad_output, output)
|
||||
|
||||
@staticmethod
|
||||
def vmap(info, in_dims, x):
|
||||
NumpyExp_.apply(x)
|
||||
return x, in_dims[0]
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent):
|
||||
# Doesn't call numpy operations because I didn't want to write NumpyMul_
|
||||
output, = ctx.saved_tensors
|
||||
x_tangent.mul_(output)
|
||||
return x_tangent
|
||||
|
||||
class NumpySort(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(x, dim):
|
||||
device = x.device
|
||||
x = to_numpy(x)
|
||||
ind = np.argsort(x, axis=dim)
|
||||
ind_inv = np.argsort(ind, axis=dim)
|
||||
return (
|
||||
torch.tensor(x, device=device),
|
||||
torch.tensor(ind, device=device),
|
||||
torch.tensor(ind_inv, device=device),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, output):
|
||||
_x, dim = inputs
|
||||
_, ind, ind_inv = output
|
||||
ctx.mark_non_differentiable(ind, ind_inv)
|
||||
ctx.save_for_backward(ind, ind_inv)
|
||||
ctx.save_for_forward(ind, ind_inv)
|
||||
ctx.dim = dim
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output, _0, _1):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
return NumpyTake.apply(grad_output, ind_inv, ind, ctx.dim), None
|
||||
|
||||
@staticmethod
|
||||
def vmap(info, in_dims, x, dim):
|
||||
x_bdim, _ = in_dims
|
||||
x = x.movedim(x_bdim, 0)
|
||||
# wrap dim
|
||||
dim = dim if dim >= 0 else dim + x.dim() - 1
|
||||
return NumpySort.apply(x, dim + 1), (0, 0, 0)
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, _):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
return NumpyTake.apply(x_tangent, ind, ind_inv, ctx.dim), None, None
|
||||
|
||||
class SortGenVmap(torch.autograd.Function):
|
||||
generate_vmap_rule = True
|
||||
|
||||
@staticmethod
|
||||
def forward(x, dim):
|
||||
ind = torch.argsort(x, dim=dim)
|
||||
ind_inv = torch.argsort(ind, axis=dim)
|
||||
result = torch.take_along_dim(x, ind, dim=dim)
|
||||
return result, ind, ind_inv
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, outputs):
|
||||
x, dim = inputs
|
||||
_, ind, ind_inv = outputs
|
||||
ctx.mark_non_differentiable(ind, ind_inv)
|
||||
ctx.save_for_backward(ind, ind_inv)
|
||||
ctx.save_for_forward(ind, ind_inv)
|
||||
ctx.dim = dim
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output, _0, _1):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
return TakeGenVmap.apply(grad_output, ind_inv, ind, ctx.dim), None
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, _):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
return TakeGenVmap.apply(x_tangent, ind, ind_inv, ctx.dim), None, None
|
||||
|
||||
|
||||
def sample_inputs_numpy_sort(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
yield SampleInput(make_arg(3, 5), args=(1,))
|
||||
|
||||
|
||||
def sample_inputs_numpy_take(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
tensor = make_arg(3, 5)
|
||||
dim = 1
|
||||
_, ind, ind_inv = NumpySort.apply(tensor, 1)
|
||||
yield SampleInput(tensor, args=(ind, ind_inv, dim))
|
||||
|
||||
|
||||
class NumpyTake(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(x, ind, ind_inv, dim):
|
||||
device = x.device
|
||||
x = to_numpy(x)
|
||||
ind = to_numpy(ind)
|
||||
return torch.tensor(np.take_along_axis(x, ind, dim), device=device)
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, output):
|
||||
_x, ind, ind_inv, dim = inputs
|
||||
ctx.save_for_backward(ind, ind_inv)
|
||||
ctx.save_for_forward(ind, ind_inv)
|
||||
ctx.dim = dim
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
result = NumpyTake.apply(grad_output, ind_inv, ind, ctx.dim)
|
||||
return result, None, None, None
|
||||
|
||||
@staticmethod
|
||||
def vmap(info, in_dims, x, ind, ind_inv, dim):
|
||||
x_bdim, ind_bdim, ind_inv_bdim, _ = in_dims
|
||||
|
||||
# wrap dim
|
||||
logical_dim = x.dim() if x_bdim is None else x_bdim - 1
|
||||
dim = dim if dim >= 0 else dim + logical_dim
|
||||
|
||||
def expand_bdim(x, x_bdim):
|
||||
if x_bdim is None:
|
||||
return x.expand(info.batch_size, *x.shape)
|
||||
return x.movedim(x_bdim, 0)
|
||||
|
||||
x = expand_bdim(x, x_bdim)
|
||||
ind = expand_bdim(ind, ind_bdim)
|
||||
ind_inv = expand_bdim(ind_inv, ind_inv_bdim)
|
||||
|
||||
return NumpyTake.apply(x, ind, ind_inv, dim + 1), 0
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, ind_tangent, ind_inv_tangent, _):
|
||||
if ind_tangent is not None:
|
||||
raise AssertionError("Expected ind_tangent to be None")
|
||||
if ind_inv_tangent is not None:
|
||||
raise AssertionError("Expected ind_inv_tangent to be None")
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
return NumpyTake.apply(x_tangent, ind, ind_inv, ctx.dim)
|
||||
|
||||
class TakeGenVmap(torch.autograd.Function):
|
||||
generate_vmap_rule = True
|
||||
|
||||
@staticmethod
|
||||
def forward(x, ind, ind_inv, dim):
|
||||
return torch.take_along_dim(x, ind, dim)
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, outputs):
|
||||
_x, ind, ind_inv, dim = inputs
|
||||
ctx.save_for_backward(ind, ind_inv)
|
||||
ctx.save_for_forward(ind, ind_inv)
|
||||
ctx.dim = dim
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
result = TakeGenVmap.apply(grad_output, ind_inv, ind, ctx.dim)
|
||||
return result, None, None, None
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, ind_tangent, ind_inv_tangent, _):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
return TakeGenVmap.apply(x_tangent, ind, ind_inv, ctx.dim)
|
||||
|
||||
class Select(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(x, idx):
|
||||
return x[idx]
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, output):
|
||||
x, idx = inputs
|
||||
ctx.x_shape = x.shape
|
||||
ctx.idx = idx
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
result = grad_output.new_zeros(ctx.x_shape)
|
||||
result[ctx.idx] = grad_output
|
||||
return result, None
|
||||
|
||||
@staticmethod
|
||||
def vmap(info, in_dims, x, idx):
|
||||
x_bdim, _ = in_dims
|
||||
x = x.movedim(x_bdim, 1)
|
||||
return Select.apply(x, idx), 0
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, _):
|
||||
return Select.apply(x_tangent, ctx.idx)
|
||||
|
||||
class SelectGenVmap(torch.autograd.Function):
|
||||
generate_vmap_rule = True
|
||||
|
||||
@staticmethod
|
||||
def forward(x, idx):
|
||||
return x[idx]
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, outputs):
|
||||
x, idx = inputs
|
||||
ctx.x_shape = x.shape
|
||||
ctx.idx = idx
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
result = grad_output.new_zeros(ctx.x_shape)
|
||||
result[ctx.idx] = grad_output
|
||||
return result, None
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, _):
|
||||
return SelectGenVmap.apply(x_tangent, ctx.idx)
|
||||
|
||||
|
||||
def sample_inputs_select(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
yield SampleInput(make_arg(3, 5), args=(2,))
|
||||
|
||||
class ScaleGradGenVmap(torch.autograd.Function):
|
||||
generate_vmap_rule = True
|
||||
scale = 3.14
|
||||
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return x.clone()
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, outputs):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return grad_output * ScaleGradGenVmap.scale
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent):
|
||||
return x_tangent * ScaleGradGenVmap.scale
|
||||
|
||||
class ZeroGradientsGenVmap(torch.autograd.Function):
|
||||
generate_vmap_rule = True
|
||||
|
||||
@staticmethod
|
||||
def forward(x, y):
|
||||
return x.clone(), y.clone()
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, outputs):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, gx, gy):
|
||||
# Intentionally returning torch.zeros instead of zeros_like or new_zeros.
|
||||
# Also intentionally not None.
|
||||
return (
|
||||
# Intentionally too-large gradient
|
||||
torch.zeros(3, 4, *gx.shape, dtype=gx.dtype, device=gx.device),
|
||||
torch.zeros(gy.shape, dtype=gy.dtype, device=gy.device),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, gx, gy):
|
||||
# Intentionally returning torch.zeros instead of zeros_like or new_zeros.
|
||||
# Also intentionally not None.
|
||||
return (
|
||||
torch.zeros(gx.shape, dtype=gx.dtype, device=gx.device),
|
||||
torch.zeros(gy.shape, dtype=gy.dtype, device=gy.device),
|
||||
)
|
||||
|
||||
|
||||
def sample_inputs_forward_default_args(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
yield SampleInput(make_arg(3, 5))
|
||||
|
||||
|
||||
class ForwardHasDefaultArgs(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(x, idx=(2,)):
|
||||
return x[idx]
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx, inputs, output):
|
||||
x, idx = inputs
|
||||
ctx.x_shape = x.shape
|
||||
ctx.idx = idx
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
result = grad_output.new_zeros(ctx.x_shape)
|
||||
result[ctx.idx] = grad_output
|
||||
return result, None
|
||||
|
||||
@staticmethod
|
||||
def vmap(info, in_dims, x, idx):
|
||||
x_bdim, _ = in_dims
|
||||
x = x.movedim(x_bdim, 1)
|
||||
return ForwardHasDefaultArgs.apply(x, idx), 0
|
||||
|
||||
@staticmethod
|
||||
def jvp(ctx, x_tangent, _):
|
||||
return ForwardHasDefaultArgs.apply(x_tangent, ctx.idx)
|
||||
|
||||
|
||||
autograd_function_db = [
|
||||
OpInfo(
|
||||
'NumpyCubeAutogradFunction',
|
||||
op=NumpyCube.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_numpy_cube,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyExpMarkDirtyAutogradFunction',
|
||||
op=lambda x: NumpyExp_.apply(x.clone()),
|
||||
inplace_variant=NumpyExp_.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_numpy_cube,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyMulAutogradFunction',
|
||||
op=NumpyMul.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_numpy_mul,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyCubeNotComposableAutogradFunction',
|
||||
op=lambda x: NumpyCubeNotComposable.apply(x)[0],
|
||||
supports_forward_ad=False,
|
||||
supports_fwgrad_bwgrad=False,
|
||||
sample_inputs_func=sample_inputs_numpy_cube,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpySortAutogradFunction',
|
||||
op=NumpySort.apply,
|
||||
supports_forward_ad=False,
|
||||
supports_fwgrad_bwgrad=False,
|
||||
sample_inputs_func=sample_inputs_numpy_sort,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
gradcheck_wrapper=lambda y, ind: y,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyTakeAutogradFunction',
|
||||
op=NumpyTake.apply,
|
||||
supports_forward_ad=False,
|
||||
supports_fwgrad_bwgrad=False,
|
||||
sample_inputs_func=sample_inputs_numpy_take,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'SelectAutogradFunction',
|
||||
op=Select.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_select,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'CubeGenVmapAutogradFunction',
|
||||
op=CubeGenVmap.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_numpy_cube,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'MulGenVmapAutogradFunction',
|
||||
op=MulGenVmap.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_numpy_mul,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'SortGenVmapAutogradFunction',
|
||||
op=SortGenVmap.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_numpy_sort,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
gradcheck_wrapper=lambda y, ind: y,
|
||||
),
|
||||
OpInfo(
|
||||
'SelectGenVmapAutogradFunction',
|
||||
op=SelectGenVmap.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_select,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'ScaleGradGenVmapAutogradFunction',
|
||||
op=ScaleGradGenVmap.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_numpy_cube,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'ZeroGradientsGenVmapAutogradFunction',
|
||||
op=ZeroGradientsGenVmap.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_numpy_mul,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'ForwardHasDefaultArgsAutogradFunction',
|
||||
op=ForwardHasDefaultArgs.apply,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
sample_inputs_func=sample_inputs_forward_default_args,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
]
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
|
||||
__all__ = [
|
||||
"check_code_for_cuda_kernel_launches",
|
||||
"check_cuda_kernel_launches",
|
||||
]
|
||||
|
||||
# FILES TO EXCLUDE (match is done with suffix using `endswith`)
|
||||
# You wouldn't drive without a seatbelt, though, so why would you
|
||||
# launch a kernel without some safety? Use this as a quick workaround
|
||||
# for a problem with the checker, fix the checker, then de-exclude
|
||||
# the files in question.
|
||||
exclude_files: list[str] = []
|
||||
|
||||
# Without using a C++ AST we can't 100% detect kernel launches, so we
|
||||
# model them as having the pattern "<<<parameters>>>(arguments);"
|
||||
# We then require that `C10_CUDA_KERNEL_LAUNCH_CHECK` be
|
||||
# the next statement.
|
||||
#
|
||||
# We model the next statement as ending at the next `}` or `;`.
|
||||
# If we see `}` then a clause ended (bad) if we see a semi-colon then
|
||||
# we expect the launch check just before it.
|
||||
#
|
||||
# Since the kernel launch can include lambda statements, it's important
|
||||
# to find the correct end-paren of the kernel launch. Doing this with
|
||||
# pure regex requires recursive regex, which aren't part of the Python
|
||||
# standard library. To avoid an additional dependency, we build a prefix
|
||||
# regex that finds the start of a kernel launch, use a paren-matching
|
||||
# algorithm to find the end of the launch, and then another regex to
|
||||
# determine if a launch check is present.
|
||||
|
||||
# Finds potential starts of kernel launches
|
||||
kernel_launch_start = re.compile(
|
||||
r"^.*<<<[^>]+>>>\s*\(", flags=re.MULTILINE
|
||||
)
|
||||
|
||||
# This pattern should start at the character after the final paren of the
|
||||
# kernel launch. It returns a match if the launch check is not the next statement
|
||||
has_check = re.compile(
|
||||
r"\s*;(?![^;}]*C10_CUDA_KERNEL_LAUNCH_CHECK\(\);)", flags=re.MULTILINE
|
||||
)
|
||||
|
||||
def find_matching_paren(s: str, startpos: int) -> int:
|
||||
"""Given a string "prefix (unknown number of characters) suffix"
|
||||
and the position of the first `(` returns the index of the character
|
||||
1 past the `)`, accounting for paren nesting
|
||||
"""
|
||||
opening = 0
|
||||
for i, c in enumerate(s[startpos:]):
|
||||
if c == '(':
|
||||
opening += 1
|
||||
elif c == ')':
|
||||
opening -= 1
|
||||
if opening == 0:
|
||||
return startpos + i + 1
|
||||
|
||||
raise IndexError("Closing parens not found!")
|
||||
|
||||
|
||||
def should_exclude_file(filename) -> bool:
|
||||
for exclude_suffix in exclude_files:
|
||||
if filename.endswith(exclude_suffix):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def check_code_for_cuda_kernel_launches(code, filename=None):
|
||||
"""Checks code for CUDA kernel launches without cuda error checks.
|
||||
|
||||
Args:
|
||||
filename - Filename of file containing the code. Used only for display
|
||||
purposes, so you can put anything here.
|
||||
code - The code to check
|
||||
|
||||
Returns:
|
||||
The number of unsafe kernel launches in the code
|
||||
"""
|
||||
if filename is None:
|
||||
filename = "##Python Function Call##"
|
||||
|
||||
# We break the code apart and put it back together to add
|
||||
# helpful line numberings for identifying problem areas
|
||||
code = enumerate(code.split("\n")) # Split by line breaks
|
||||
code = [f"{lineno}: {linecode}" for lineno, linecode in code] # Number the lines
|
||||
code = '\n'.join(code) # Put it back together
|
||||
|
||||
num_launches_without_checks = 0
|
||||
for m in kernel_launch_start.finditer(code):
|
||||
end_paren = find_matching_paren(code, m.end() - 1)
|
||||
if has_check.match(code, end_paren):
|
||||
num_launches_without_checks += 1
|
||||
context = code[m.start():end_paren + 1]
|
||||
print(f"Missing C10_CUDA_KERNEL_LAUNCH_CHECK in '{filename}'. Context:\n{context}", file=sys.stderr)
|
||||
|
||||
return num_launches_without_checks
|
||||
|
||||
|
||||
def check_file(filename):
|
||||
"""Checks a file for CUDA kernel launches without cuda error checks
|
||||
|
||||
Args:
|
||||
filename - File to check
|
||||
|
||||
Returns:
|
||||
The number of unsafe kernel launches in the file
|
||||
"""
|
||||
if not (filename.endswith((".cu", ".cuh"))):
|
||||
return 0
|
||||
if should_exclude_file(filename):
|
||||
return 0
|
||||
with open(filename) as f:
|
||||
contents = f.read()
|
||||
unsafeCount = check_code_for_cuda_kernel_launches(contents, filename)
|
||||
return unsafeCount
|
||||
|
||||
|
||||
def check_cuda_kernel_launches():
|
||||
"""Checks all pytorch code for CUDA kernel launches without cuda error checks
|
||||
|
||||
Returns:
|
||||
The number of unsafe kernel launches in the codebase
|
||||
"""
|
||||
torch_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
torch_dir = os.path.dirname(torch_dir) # Go up to parent torch
|
||||
torch_dir = os.path.dirname(torch_dir) # Go up to parent caffe2
|
||||
|
||||
kernels_without_checks = 0
|
||||
files_without_checks = []
|
||||
for root, dirnames, filenames in os.walk(torch_dir):
|
||||
# `$BASE/build` and `$BASE/torch/include` are generated
|
||||
# so we don't want to flag their contents
|
||||
if root == os.path.join(torch_dir, "build") or root == os.path.join(torch_dir, "torch/include"):
|
||||
# Curtail search by modifying dirnames and filenames in place
|
||||
# Yes, this is the way to do this, see `help(os.walk)`
|
||||
dirnames[:] = []
|
||||
continue
|
||||
|
||||
for x in filenames:
|
||||
filename = os.path.join(root, x)
|
||||
file_result = check_file(filename)
|
||||
if file_result > 0:
|
||||
kernels_without_checks += file_result
|
||||
files_without_checks.append(filename)
|
||||
|
||||
if kernels_without_checks > 0:
|
||||
count_str = f"Found {kernels_without_checks} instances in " \
|
||||
f"{len(files_without_checks)} files where kernel " \
|
||||
"launches didn't have checks."
|
||||
print(count_str, file=sys.stderr)
|
||||
print("Files without checks:", file=sys.stderr)
|
||||
for x in files_without_checks:
|
||||
print(f"\t{x}", file=sys.stderr)
|
||||
print(count_str, file=sys.stderr)
|
||||
|
||||
return kernels_without_checks
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unsafe_launches = check_cuda_kernel_launches()
|
||||
sys.exit(0 if unsafe_launches == 0 else 1)
|
||||
@@ -0,0 +1 @@
|
||||
# mypy: ignore-errors
|
||||
@@ -0,0 +1,497 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
r"""This file is allowed to initialize CUDA context when imported."""
|
||||
|
||||
import functools
|
||||
import torch
|
||||
import torch.cuda
|
||||
from torch.testing._internal.common_utils import LazyVal, TEST_NUMBA, TEST_WITH_ROCM, TEST_CUDA, IS_WINDOWS, IS_MACOS, TEST_XPU
|
||||
import inspect
|
||||
import contextlib
|
||||
import os
|
||||
import unittest
|
||||
|
||||
|
||||
CUDA_ALREADY_INITIALIZED_ON_IMPORT = torch.cuda.is_initialized()
|
||||
|
||||
|
||||
TEST_MULTIGPU = TEST_CUDA and torch.cuda.device_count() >= 2
|
||||
CUDA_DEVICE = torch.device("cuda:0") if TEST_CUDA else None
|
||||
# note: if ROCm is targeted, TEST_CUDNN is code for TEST_MIOPEN
|
||||
if TEST_WITH_ROCM:
|
||||
TEST_CUDNN = LazyVal(lambda: TEST_CUDA)
|
||||
else:
|
||||
TEST_CUDNN = LazyVal(lambda: TEST_CUDA and torch.backends.cudnn.is_acceptable(torch.tensor(1., device=CUDA_DEVICE)))
|
||||
|
||||
TEST_CUDNN_VERSION = LazyVal(lambda: torch.backends.cudnn.version() if TEST_CUDNN else 0)
|
||||
ROCM_VERSION = LazyVal(lambda : tuple(int(v) for v in torch.version.hip.split('.')[:2]) if torch.version.hip else (0, 0))
|
||||
|
||||
SM53OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (5, 3))
|
||||
SM60OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (6, 0))
|
||||
SM70OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (7, 0))
|
||||
SM75OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (7, 5))
|
||||
SM80OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 0))
|
||||
SM89OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 9))
|
||||
SM90OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0))
|
||||
SM100OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0))
|
||||
SM120OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (12, 0))
|
||||
|
||||
IS_THOR = LazyVal(lambda: torch.cuda.is_available() and torch.version.cuda is not None and
|
||||
((torch.cuda.get_device_capability() == (11, 0) and int(torch.version.cuda[:2]) >= 13) or
|
||||
(torch.cuda.get_device_capability() == (10, 1) and int(torch.version.cuda[:2]) < 13)))
|
||||
IS_JETSON = LazyVal(lambda: torch.cuda.is_available() and (torch.cuda.get_device_capability() in [(7, 2), (8, 7)] or IS_THOR))
|
||||
IS_SM89 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (8, 9))
|
||||
IS_SM90 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0))
|
||||
IS_SM100 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0))
|
||||
IS_SM12X = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def blas_library_context(backend):
|
||||
prev_backend = torch.backends.cuda.preferred_blas_library()
|
||||
torch.backends.cuda.preferred_blas_library(backend)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.backends.cuda.preferred_blas_library(prev_backend)
|
||||
|
||||
def evaluate_gfx_arch_within(arch_list):
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
gcn_arch_name = torch.cuda.get_device_properties('cuda').gcnArchName
|
||||
effective_arch = os.environ.get('PYTORCH_DEBUG_FLASH_ATTENTION_GCN_ARCH_OVERRIDE', gcn_arch_name)
|
||||
# gcnArchName can be complicated strings like gfx90a:sramecc+:xnack-
|
||||
# Hence the matching should be done reversely
|
||||
return any(arch in effective_arch for arch in arch_list)
|
||||
|
||||
def CDNA3OrLater():
|
||||
return evaluate_gfx_arch_within(["gfx942", "gfx950"])
|
||||
|
||||
def CDNA2OrLater():
|
||||
return evaluate_gfx_arch_within(["gfx90a", "gfx942"])
|
||||
|
||||
def evaluate_platform_supports_flash_attention():
|
||||
if TEST_WITH_ROCM:
|
||||
arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"]
|
||||
if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0":
|
||||
arch_list += ["gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200"]
|
||||
return evaluate_gfx_arch_within(arch_list)
|
||||
if TEST_CUDA:
|
||||
return not IS_WINDOWS and SM80OrLater
|
||||
if TEST_XPU:
|
||||
return True
|
||||
return False
|
||||
|
||||
def evaluate_platform_supports_ck_sdpa():
|
||||
if TEST_WITH_ROCM:
|
||||
return torch.backends.cuda.is_ck_sdpa_available()
|
||||
else:
|
||||
return False
|
||||
|
||||
def evaluate_platform_supports_efficient_attention():
|
||||
if TEST_WITH_ROCM:
|
||||
arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"]
|
||||
if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0":
|
||||
arch_list += ["gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200"]
|
||||
return evaluate_gfx_arch_within(arch_list)
|
||||
if TEST_CUDA:
|
||||
return True
|
||||
if TEST_XPU:
|
||||
return True
|
||||
return False
|
||||
|
||||
def evaluate_platform_supports_cudnn_attention():
|
||||
return (not TEST_WITH_ROCM) and SM80OrLater and (TEST_CUDNN_VERSION >= 90000)
|
||||
|
||||
def evaluate_platform_supports_green_context():
|
||||
if IS_WINDOWS:
|
||||
return False
|
||||
if not _get_torch_cuda_version() >= (12, 8):
|
||||
return False
|
||||
driver_version = torch.utils.collect_env.get_nvidia_driver_version(torch.utils.collect_env.run)
|
||||
if driver_version is None:
|
||||
return False
|
||||
return int(driver_version.split('.')[0]) >= 570
|
||||
|
||||
PLATFORM_SUPPORTS_FLASH_ATTENTION: bool = LazyVal(lambda: evaluate_platform_supports_flash_attention())
|
||||
PLATFORM_SUPPORTS_MEM_EFF_ATTENTION: bool = LazyVal(lambda: evaluate_platform_supports_efficient_attention())
|
||||
PLATFORM_SUPPORTS_CUDNN_ATTENTION: bool = LazyVal(lambda: evaluate_platform_supports_cudnn_attention())
|
||||
# This condition always evaluates to PLATFORM_SUPPORTS_MEM_EFF_ATTENTION but for logical clarity we keep it separate
|
||||
PLATFORM_SUPPORTS_FUSED_ATTENTION: bool = LazyVal(lambda: PLATFORM_SUPPORTS_FLASH_ATTENTION or
|
||||
PLATFORM_SUPPORTS_CUDNN_ATTENTION or
|
||||
PLATFORM_SUPPORTS_MEM_EFF_ATTENTION)
|
||||
|
||||
PLATFORM_SUPPORTS_FUSED_SDPA: bool = TEST_CUDA and not TEST_WITH_ROCM
|
||||
|
||||
PLATFORM_SUPPORTS_CK_SDPA: bool = LazyVal(lambda: evaluate_platform_supports_ck_sdpa())
|
||||
|
||||
|
||||
def evaluate_platform_supports_bf16():
|
||||
if torch.version.cuda:
|
||||
return SM80OrLater
|
||||
elif torch.version.hip:
|
||||
return True
|
||||
elif TEST_XPU:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def evaluate_platform_supports_bf16_atomics():
|
||||
if torch.version.cuda:
|
||||
return SM80OrLater
|
||||
elif torch.version.hip:
|
||||
return ROCM_VERSION >= (8, 0)
|
||||
return False
|
||||
|
||||
|
||||
def evaluate_platform_supports_half_atomics():
|
||||
if torch.version.hip:
|
||||
return ROCM_VERSION >= (8, 0)
|
||||
return True
|
||||
|
||||
|
||||
PLATFORM_SUPPORTS_BF16: bool = LazyVal(lambda: evaluate_platform_supports_bf16())
|
||||
PLATFORM_SUPPORTS_BF16_ATOMICS: bool = LazyVal(lambda: evaluate_platform_supports_bf16_atomics())
|
||||
PLATFORM_SUPPORTS_HALF_ATOMICS: bool = LazyVal(lambda: evaluate_platform_supports_half_atomics())
|
||||
|
||||
PLATFORM_SUPPORTS_GREEN_CONTEXT: bool = LazyVal(lambda: evaluate_platform_supports_green_context())
|
||||
|
||||
def evaluate_platform_supports_workqueue_config():
|
||||
if IS_WINDOWS:
|
||||
return False
|
||||
if not _get_torch_cuda_version() >= (13, 1):
|
||||
return False
|
||||
driver_version = torch.utils.collect_env.get_nvidia_driver_version(torch.utils.collect_env.run)
|
||||
if driver_version is None:
|
||||
return False
|
||||
return int(driver_version.split('.')[0]) >= 590
|
||||
|
||||
PLATFORM_SUPPORTS_WORKQUEUE_CONFIG: bool = LazyVal(lambda: evaluate_platform_supports_workqueue_config())
|
||||
|
||||
def evaluate_platform_supports_fp8():
|
||||
if torch.cuda.is_available():
|
||||
if torch.version.hip:
|
||||
archs = ['gfx94']
|
||||
if ROCM_VERSION >= (6, 3):
|
||||
archs.extend(['gfx120'])
|
||||
if ROCM_VERSION >= (6, 5):
|
||||
archs.append('gfx95')
|
||||
for arch in archs:
|
||||
if arch in torch.cuda.get_device_properties(0).gcnArchName:
|
||||
return True
|
||||
return False
|
||||
else:
|
||||
return SM90OrLater or torch.cuda.get_device_capability() == (8, 9)
|
||||
if torch.xpu.is_available():
|
||||
return True
|
||||
# As CPU supports FP8 and is always available, return True.
|
||||
return True
|
||||
|
||||
def evaluate_platform_supports_fp8_grouped_gemm():
|
||||
if torch.cuda.is_available():
|
||||
if torch.version.hip:
|
||||
if "USE_MSLK" not in torch.__config__.show():
|
||||
return False
|
||||
archs = ['gfx942', 'gfx950']
|
||||
for arch in archs:
|
||||
if arch in torch.cuda.get_device_properties(0).gcnArchName:
|
||||
return True
|
||||
else:
|
||||
return SM90OrLater and not SM100OrLater
|
||||
return False
|
||||
|
||||
def evaluate_platform_supports_mx_gemm():
|
||||
if torch.cuda.is_available():
|
||||
if torch.version.hip:
|
||||
if ROCM_VERSION >= (7, 0):
|
||||
return 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName
|
||||
else:
|
||||
return SM100OrLater
|
||||
return False
|
||||
|
||||
def evaluate_platform_supports_mxfp8_grouped_gemm():
|
||||
if torch.cuda.is_available() and not torch.version.hip:
|
||||
built_with_mslk = "USE_MSLK" in torch.__config__.show()
|
||||
return built_with_mslk and IS_SM100
|
||||
return False
|
||||
|
||||
def evaluate_platform_supports_fp8_sparse():
|
||||
if torch.cuda.is_available():
|
||||
if torch.version.hip:
|
||||
return 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName
|
||||
else:
|
||||
return (
|
||||
(SM90OrLater or torch.cuda.get_device_capability() == (8, 9))
|
||||
and torch.backends.cusparselt.is_available()
|
||||
and torch.backends.cusparselt.version() >= 602
|
||||
)
|
||||
return False
|
||||
|
||||
PLATFORM_SUPPORTS_MX_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_mx_gemm())
|
||||
PLATFORM_SUPPORTS_FP8: bool = LazyVal(lambda: evaluate_platform_supports_fp8())
|
||||
PLATFORM_SUPPORTS_FP8_SPARSE: bool = LazyVal(lambda: evaluate_platform_supports_fp8_sparse())
|
||||
PLATFORM_SUPPORTS_FP8_GROUPED_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_fp8_grouped_gemm())
|
||||
PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_mxfp8_grouped_gemm())
|
||||
|
||||
if TEST_NUMBA:
|
||||
try:
|
||||
import numba.cuda
|
||||
TEST_NUMBA_CUDA = numba.cuda.is_available()
|
||||
except (ImportError, RuntimeError, OSError):
|
||||
TEST_NUMBA_CUDA = False
|
||||
TEST_NUMBA = False
|
||||
else:
|
||||
TEST_NUMBA_CUDA = False
|
||||
|
||||
# Used below in `initialize_cuda_context_rng` to ensure that CUDA context and
|
||||
# RNG have been initialized.
|
||||
__cuda_ctx_rng_initialized = False
|
||||
|
||||
|
||||
# after this call, CUDA context and RNG must have been initialized on each GPU
|
||||
def initialize_cuda_context_rng():
|
||||
global __cuda_ctx_rng_initialized
|
||||
if not TEST_CUDA:
|
||||
raise AssertionError('CUDA must be available when calling initialize_cuda_context_rng')
|
||||
if not __cuda_ctx_rng_initialized:
|
||||
# initialize cuda context and rng for memory tests
|
||||
for i in range(torch.cuda.device_count()):
|
||||
torch.randn(1, device=f"cuda:{i}")
|
||||
__cuda_ctx_rng_initialized = True
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def tf32_off():
|
||||
old_allow_tf32_matmul = torch.backends.cuda.matmul.allow_tf32
|
||||
try:
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
with torch.backends.cudnn.flags(enabled=None, benchmark=None, deterministic=None, allow_tf32=False):
|
||||
yield
|
||||
finally:
|
||||
torch.backends.cuda.matmul.allow_tf32 = old_allow_tf32_matmul
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def tf32_on(self, tf32_precision=1e-5):
|
||||
old_allow_tf32_matmul = torch.backends.cuda.matmul.allow_tf32
|
||||
old_precision = self.precision
|
||||
try:
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
self.precision = tf32_precision
|
||||
with torch.backends.cudnn.flags(enabled=None, benchmark=None, deterministic=None, allow_tf32=True):
|
||||
yield
|
||||
finally:
|
||||
torch.backends.cuda.matmul.allow_tf32 = old_allow_tf32_matmul
|
||||
self.precision = old_precision
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def tf32_enabled():
|
||||
"""
|
||||
Context manager to temporarily enable TF32 for CUDA operations.
|
||||
Restores the previous TF32 state after exiting the context.
|
||||
"""
|
||||
old_allow_tf32_matmul = torch.backends.cuda.matmul.allow_tf32
|
||||
try:
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
with torch.backends.cudnn.flags(
|
||||
enabled=None, benchmark=None, deterministic=None, allow_tf32=True
|
||||
):
|
||||
yield
|
||||
finally:
|
||||
torch.backends.cuda.matmul.allow_tf32 = old_allow_tf32_matmul
|
||||
|
||||
|
||||
# This is a wrapper that wraps a test to run this test twice, one with
|
||||
# allow_tf32=True, another with allow_tf32=False. When running with
|
||||
# allow_tf32=True, it will use reduced precision as specified by the
|
||||
# argument. For example:
|
||||
# @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
|
||||
# @tf32_on_and_off(0.005)
|
||||
# def test_matmul(self, device, dtype):
|
||||
# a = ...; b = ...;
|
||||
# c = torch.matmul(a, b)
|
||||
# self.assertEqual(c, expected)
|
||||
# In the above example, when testing torch.float32 and torch.complex64 on CUDA
|
||||
# on a CUDA >= 11 build on an >=Ampere architecture, the matmul will be running at
|
||||
# TF32 mode and TF32 mode off, and on TF32 mode, the assertEqual will use reduced
|
||||
# precision to check values.
|
||||
#
|
||||
# This decorator can be used for function with or without device/dtype, such as
|
||||
# @tf32_on_and_off(0.005)
|
||||
# def test_my_op(self)
|
||||
# @tf32_on_and_off(0.005)
|
||||
# def test_my_op(self, device)
|
||||
# @tf32_on_and_off(0.005)
|
||||
# def test_my_op(self, device, dtype)
|
||||
# @tf32_on_and_off(0.005)
|
||||
# def test_my_op(self, dtype)
|
||||
# if neither device nor dtype is specified, it will check if the system has ampere device
|
||||
# if device is specified, it will check if device is cuda
|
||||
# if dtype is specified, it will check if dtype is float32 or complex64
|
||||
# tf32 and fp32 are different only when all the three checks pass
|
||||
def tf32_on_and_off(tf32_precision=1e-5, *, only_if=True):
|
||||
def with_tf32_disabled(self, function_call):
|
||||
with tf32_off():
|
||||
function_call()
|
||||
|
||||
def with_tf32_enabled(self, function_call):
|
||||
with tf32_on(self, tf32_precision):
|
||||
function_call()
|
||||
|
||||
def wrapper(f):
|
||||
params = inspect.signature(f).parameters
|
||||
arg_names = tuple(params.keys())
|
||||
|
||||
@functools.wraps(f)
|
||||
def wrapped(*args, **kwargs):
|
||||
kwargs.update(zip(arg_names, args, strict=False))
|
||||
cond = torch.cuda.is_tf32_supported() and only_if
|
||||
if 'device' in kwargs:
|
||||
cond = cond and (torch.device(kwargs['device']).type == 'cuda')
|
||||
if 'dtype' in kwargs:
|
||||
cond = cond and (kwargs['dtype'] in {torch.float32, torch.complex64})
|
||||
if cond:
|
||||
with_tf32_disabled(kwargs['self'], lambda: f(**kwargs))
|
||||
with_tf32_enabled(kwargs['self'], lambda: f(**kwargs))
|
||||
else:
|
||||
f(**kwargs)
|
||||
|
||||
return wrapped
|
||||
return wrapper
|
||||
|
||||
# This is a wrapper that wraps a test to run it with TF32 turned off.
|
||||
# This wrapper is designed to be used when a test uses matmul or convolutions
|
||||
# but the purpose of that test is not testing matmul or convolutions.
|
||||
# Disabling TF32 will enforce torch.float tensors to be always computed
|
||||
# at full precision.
|
||||
def with_tf32_off(f):
|
||||
@functools.wraps(f)
|
||||
def wrapped(*args, **kwargs):
|
||||
with tf32_off():
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return wrapped
|
||||
|
||||
def _get_magma_version():
|
||||
if 'Magma' not in torch.__config__.show():
|
||||
return (0, 0)
|
||||
position = torch.__config__.show().find('Magma ')
|
||||
version_str = torch.__config__.show()[position + len('Magma '):].split('\n')[0]
|
||||
return tuple(int(x) for x in version_str.split("."))
|
||||
|
||||
def _get_torch_cuda_version():
|
||||
if torch.version.cuda is None:
|
||||
return (0, 0)
|
||||
cuda_version = str(torch.version.cuda)
|
||||
return tuple(int(x) for x in cuda_version.split("."))
|
||||
|
||||
def _get_torch_rocm_version():
|
||||
if not TEST_WITH_ROCM or torch.version.hip is None:
|
||||
return (0, 0)
|
||||
rocm_version = str(torch.version.hip)
|
||||
rocm_version = rocm_version.split("-", maxsplit=1)[0] # ignore git sha
|
||||
return tuple(int(x) for x in rocm_version.split("."))
|
||||
|
||||
def _get_torch_hipblaslt_version():
|
||||
if not TEST_WITH_ROCM:
|
||||
return None
|
||||
try:
|
||||
# Access through direct C binding
|
||||
# versionHipBLASLt returns: MAJOR * 10000 + MINOR * 100 + PATCH
|
||||
version_int = torch._C._cuda_getHipblasltVersion()
|
||||
if version_int is None or version_int == 0:
|
||||
return None
|
||||
major = version_int // 10000
|
||||
minor = (version_int % 10000) // 100
|
||||
patch = version_int % 100
|
||||
return (major, minor, patch)
|
||||
except (AttributeError, RuntimeError):
|
||||
return None
|
||||
|
||||
def _check_cusparse_generic_available():
|
||||
return not TEST_WITH_ROCM
|
||||
|
||||
def _check_hipsparse_generic_available():
|
||||
if not TEST_WITH_ROCM:
|
||||
return False
|
||||
if not torch.version.hip:
|
||||
return False
|
||||
|
||||
rocm_version = str(torch.version.hip)
|
||||
rocm_version = rocm_version.split("-", maxsplit=1)[0] # ignore git sha
|
||||
rocm_version_tuple = tuple(int(x) for x in rocm_version.split("."))
|
||||
return not (rocm_version_tuple is None or rocm_version_tuple < (5, 1))
|
||||
|
||||
|
||||
TEST_CUSPARSE_GENERIC = _check_cusparse_generic_available()
|
||||
TEST_HIPSPARSE_GENERIC = _check_hipsparse_generic_available()
|
||||
|
||||
# Shared by test_torch.py and test_multigpu.py
|
||||
def _create_scaling_models_optimizers(device="cuda", optimizer_ctor=torch.optim.SGD, optimizer_kwargs=None):
|
||||
# Create a module+optimizer that will use scaling, and a control module+optimizer
|
||||
# that will not use scaling, against which the scaling-enabled module+optimizer can be compared.
|
||||
mod_control = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).to(device=device)
|
||||
mod_scaling = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).to(device=device)
|
||||
with torch.no_grad():
|
||||
for c, s in zip(mod_control.parameters(), mod_scaling.parameters(), strict=True):
|
||||
s.copy_(c)
|
||||
|
||||
kwargs = {"lr": 1.0}
|
||||
if optimizer_kwargs is not None:
|
||||
kwargs.update(optimizer_kwargs)
|
||||
opt_control = optimizer_ctor(mod_control.parameters(), **kwargs)
|
||||
opt_scaling = optimizer_ctor(mod_scaling.parameters(), **kwargs)
|
||||
|
||||
return mod_control, mod_scaling, opt_control, opt_scaling
|
||||
|
||||
# Shared by test_torch.py, test_cuda.py and test_multigpu.py
|
||||
def _create_scaling_case(device="cuda", dtype=torch.float, optimizer_ctor=torch.optim.SGD, optimizer_kwargs=None):
|
||||
data = [(torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),
|
||||
(torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),
|
||||
(torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),
|
||||
(torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device))]
|
||||
|
||||
loss_fn = torch.nn.MSELoss().to(device)
|
||||
|
||||
skip_iter = 2
|
||||
|
||||
return _create_scaling_models_optimizers(
|
||||
device=device, optimizer_ctor=optimizer_ctor, optimizer_kwargs=optimizer_kwargs,
|
||||
) + (data, loss_fn, skip_iter)
|
||||
|
||||
|
||||
def xfailIfSM89(func):
|
||||
return func if not IS_SM89 else unittest.expectedFailure(func)
|
||||
|
||||
def xfailIfSM90(func):
|
||||
return func if not IS_SM90 else unittest.expectedFailure(func)
|
||||
|
||||
def xfailIfSM89PreCUDA13(func):
|
||||
"""xfail on SM89 only for CUDA < 13. On CUDA 13+, test should pass on all architectures."""
|
||||
if IS_SM89 and _get_torch_cuda_version() < (13, 0):
|
||||
return unittest.expectedFailure(func)
|
||||
return func
|
||||
|
||||
def xfailIfSM100OrLater(func):
|
||||
return func if not SM100OrLater else unittest.expectedFailure(func)
|
||||
|
||||
def xfailIfSM120OrLater(func):
|
||||
return func if not SM120OrLater else unittest.expectedFailure(func)
|
||||
|
||||
def xfailIfSM12X(func):
|
||||
return func if not IS_SM12X else unittest.expectedFailure(func)
|
||||
|
||||
def xfailIfDistributedNotSupported(func):
|
||||
return func if not (IS_MACOS or IS_JETSON) else unittest.expectedFailure(func)
|
||||
|
||||
# When using nvcc from the CUDA toolkit its versuib must be at least the one from ptxas bundled with Triton
|
||||
TRITON_PTXAS_VERSION = (12, 8)
|
||||
requires_triton_ptxas_compat = unittest.skipIf(not torch.version.xpu
|
||||
and torch.version.hip is None
|
||||
and _get_torch_cuda_version() < TRITON_PTXAS_VERSION,
|
||||
"Requires CUDA {}.{} to match Tritons ptxas version".format(*TRITON_PTXAS_VERSION))
|
||||
|
||||
# Importing this module should NOT eagerly initialize CUDA
|
||||
if not CUDA_ALREADY_INITIALIZED_ON_IMPORT:
|
||||
if torch.cuda.is_initialized():
|
||||
raise AssertionError("CUDA should not be initialized on import")
|
||||
+2116
File diff suppressed because it is too large
Load Diff
+112
@@ -0,0 +1,112 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
# Owner(s): ["oncall: distributed"]
|
||||
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class UnitModule(nn.Module):
|
||||
def __init__(self, device: torch.device):
|
||||
super().__init__()
|
||||
self.l1 = nn.Linear(100, 100, device=device)
|
||||
self.seq = nn.Sequential(
|
||||
nn.ReLU(),
|
||||
nn.Linear(100, 100, device=device),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self.l2 = nn.Linear(100, 100, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
return self.l2(self.seq(self.l1(x)))
|
||||
|
||||
|
||||
class CompositeModel(nn.Module):
|
||||
def __init__(self, device: torch.device):
|
||||
super().__init__()
|
||||
self.l1 = nn.Linear(100, 100, device=device)
|
||||
self.u1 = UnitModule(device)
|
||||
self.u2 = UnitModule(device)
|
||||
self.l2 = nn.Linear(100, 100, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
return self.l2(self.u2(self.u1(self.l1(x))))
|
||||
|
||||
|
||||
class UnitParamModule(nn.Module):
|
||||
def __init__(self, device: torch.device):
|
||||
super().__init__()
|
||||
self.l = nn.Linear(100, 100, device=device)
|
||||
self.seq = nn.Sequential(
|
||||
nn.ReLU(),
|
||||
nn.Linear(100, 100, device=device),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self.p = nn.Parameter(torch.randn((100, 100), device=device))
|
||||
|
||||
def forward(self, x):
|
||||
return torch.mm(self.seq(self.l(x)), self.p)
|
||||
|
||||
|
||||
class CompositeParamModel(nn.Module):
|
||||
def __init__(self, device: torch.device):
|
||||
super().__init__()
|
||||
self.l = nn.Linear(100, 100, device=device)
|
||||
self.u1 = UnitModule(device)
|
||||
self.u2 = UnitModule(device)
|
||||
self.p = nn.Parameter(torch.randn((100, 100), device=device))
|
||||
self.register_buffer(
|
||||
"buffer", torch.randn((100, 100), device=device), persistent=True
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
a = self.u2(self.u1(self.l(x)))
|
||||
b = self.p
|
||||
return torch.mm(a, b)
|
||||
|
||||
|
||||
class FakeSequential(nn.Module):
|
||||
# Define this class to achieve a desired nested wrapping using the module
|
||||
# wrap policy with `nn.Sequential`
|
||||
def __init__(self, *modules: tuple[nn.Module, ...]) -> None:
|
||||
super().__init__()
|
||||
self._module_sequence = list(modules)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
for module in self._module_sequence:
|
||||
x = module(x)
|
||||
return x
|
||||
|
||||
|
||||
class NestedSequentialModel(nn.Module):
|
||||
def __init__(self, device: torch.device) -> None:
|
||||
super().__init__()
|
||||
# This nested structure exercises traversal order to catch differences
|
||||
# between valid traversals (e.g. BFS and DFS variations).
|
||||
self.seq1 = nn.Sequential(
|
||||
nn.Linear(1, 1, device=device),
|
||||
FakeSequential(
|
||||
nn.Linear(1, 1, device=device),
|
||||
nn.ReLU(),
|
||||
FakeSequential(
|
||||
nn.Linear(1, 1, device=device),
|
||||
),
|
||||
nn.ReLU(),
|
||||
),
|
||||
nn.Linear(1, 2, device=device),
|
||||
)
|
||||
self.lin = nn.Linear(2, 2, device=device)
|
||||
self.seq2 = nn.Sequential(
|
||||
nn.ReLU(),
|
||||
nn.Linear(2, 3, device=device),
|
||||
FakeSequential(
|
||||
nn.Linear(3, 2, bias=False, device=device),
|
||||
nn.Linear(2, 4, bias=False, device=device),
|
||||
),
|
||||
)
|
||||
|
||||
# FIXME(rec): forward() is not a method, it's a local function inside __init__
|
||||
# that is never used. It should probabkly be outdented by four spaces, or removed.
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.seq2(self.lin(self.seq1(x)))
|
||||
+2110
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,243 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
# Functions and classes for describing the dtypes a function supports
|
||||
# NOTE: these helpers should correspond to PyTorch's C++ dispatch macros
|
||||
|
||||
|
||||
# Verifies each given dtype is a torch.dtype
|
||||
def _validate_dtypes(*dtypes):
|
||||
for dtype in dtypes:
|
||||
if not isinstance(dtype, torch.dtype):
|
||||
raise AssertionError(f"Expected dtype to be torch.dtype, got {type(dtype)}")
|
||||
return dtypes
|
||||
|
||||
|
||||
# class for tuples corresponding to a PyTorch dispatch macro
|
||||
class _dispatch_dtypes(tuple):
|
||||
__slots__ = ()
|
||||
|
||||
def __add__(self, other):
|
||||
if not isinstance(other, tuple):
|
||||
raise AssertionError(f"Expected other to be a tuple, got {type(other)}")
|
||||
return _dispatch_dtypes(tuple.__add__(self, other))
|
||||
|
||||
|
||||
_empty_types = _dispatch_dtypes(())
|
||||
|
||||
|
||||
def empty_types():
|
||||
return _empty_types
|
||||
|
||||
|
||||
_floating_types = _dispatch_dtypes((torch.float32, torch.float64))
|
||||
|
||||
|
||||
def floating_types():
|
||||
return _floating_types
|
||||
|
||||
|
||||
_floating_types_and_half = _floating_types + (torch.half,)
|
||||
|
||||
|
||||
def floating_types_and_half():
|
||||
return _floating_types_and_half
|
||||
|
||||
|
||||
def floating_types_and(*dtypes):
|
||||
return _floating_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
_floating_and_complex_types = _floating_types + (torch.cfloat, torch.cdouble)
|
||||
|
||||
|
||||
def floating_and_complex_types():
|
||||
return _floating_and_complex_types
|
||||
|
||||
|
||||
def floating_and_complex_types_and(*dtypes):
|
||||
return _floating_and_complex_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
_double_types = _dispatch_dtypes((torch.float64, torch.complex128))
|
||||
|
||||
|
||||
def double_types():
|
||||
return _double_types
|
||||
|
||||
|
||||
# NB: Does not contain uint16/uint32/uint64 for BC reasons
|
||||
_integral_types = _dispatch_dtypes(
|
||||
(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
|
||||
)
|
||||
|
||||
|
||||
def integral_types():
|
||||
return _integral_types
|
||||
|
||||
|
||||
def integral_types_and(*dtypes):
|
||||
return _integral_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
_all_types = _floating_types + _integral_types
|
||||
|
||||
|
||||
def all_types():
|
||||
return _all_types
|
||||
|
||||
|
||||
def all_types_and(*dtypes):
|
||||
return _all_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
_complex_types = _dispatch_dtypes((torch.cfloat, torch.cdouble))
|
||||
|
||||
|
||||
def complex_types():
|
||||
return _complex_types
|
||||
|
||||
|
||||
def complex_types_and(*dtypes):
|
||||
return _complex_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
_all_types_and_complex = _all_types + _complex_types
|
||||
|
||||
|
||||
def all_types_and_complex():
|
||||
return _all_types_and_complex
|
||||
|
||||
|
||||
def all_types_and_complex_and(*dtypes):
|
||||
return _all_types_and_complex + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
_all_types_and_half = _all_types + (torch.half,)
|
||||
|
||||
|
||||
def all_types_and_half():
|
||||
return _all_types_and_half
|
||||
|
||||
|
||||
_all_mps_types = (
|
||||
_dispatch_dtypes({torch.float, torch.half, torch.bfloat16}) + _integral_types
|
||||
)
|
||||
|
||||
|
||||
def all_mps_types():
|
||||
return _all_mps_types
|
||||
|
||||
|
||||
def all_mps_types_and(*dtypes):
|
||||
return _all_mps_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
_float8_types = _dispatch_dtypes(
|
||||
(
|
||||
torch.float8_e4m3fn,
|
||||
torch.float8_e4m3fnuz,
|
||||
torch.float8_e5m2,
|
||||
torch.float8_e5m2fnuz,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def float8_types():
|
||||
return _float8_types
|
||||
|
||||
|
||||
def float8_types_and(*dtypes):
|
||||
return _float8_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
def all_types_complex_float8_and(*dtypes):
|
||||
return _all_types + _complex_types + _float8_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
def custom_types(*dtypes):
|
||||
"""Create a list of arbitrary dtypes"""
|
||||
return _empty_types + _validate_dtypes(*dtypes)
|
||||
|
||||
|
||||
# The functions below are used for convenience in our test suite and thus have no corresponding C++ dispatch macro
|
||||
|
||||
|
||||
# See AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS.
|
||||
def get_all_dtypes(
|
||||
include_half=True,
|
||||
include_bfloat16=True,
|
||||
include_bool=True,
|
||||
include_complex=True,
|
||||
include_complex32=False,
|
||||
include_qint=False,
|
||||
) -> list[torch.dtype]:
|
||||
dtypes = get_all_int_dtypes() + get_all_fp_dtypes(
|
||||
include_half=include_half, include_bfloat16=include_bfloat16
|
||||
)
|
||||
if include_bool:
|
||||
dtypes.append(torch.bool)
|
||||
if include_complex:
|
||||
dtypes += get_all_complex_dtypes(include_complex32)
|
||||
if include_qint:
|
||||
dtypes += get_all_qint_dtypes()
|
||||
return dtypes
|
||||
|
||||
|
||||
def get_all_math_dtypes(device) -> list[torch.dtype]:
|
||||
return (
|
||||
get_all_int_dtypes()
|
||||
+ get_all_fp_dtypes(
|
||||
include_half=device.startswith("cuda"), include_bfloat16=False
|
||||
)
|
||||
+ get_all_complex_dtypes()
|
||||
)
|
||||
|
||||
|
||||
def get_all_complex_dtypes(include_complex32=False) -> list[torch.dtype]:
|
||||
return (
|
||||
[torch.complex32, torch.complex64, torch.complex128]
|
||||
if include_complex32
|
||||
else [torch.complex64, torch.complex128]
|
||||
)
|
||||
|
||||
|
||||
def get_all_int_dtypes() -> list[torch.dtype]:
|
||||
return [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]
|
||||
|
||||
|
||||
def get_all_fp_dtypes(include_half=True, include_bfloat16=True) -> list[torch.dtype]:
|
||||
dtypes = [torch.float32, torch.float64]
|
||||
if include_half:
|
||||
dtypes.append(torch.float16)
|
||||
if include_bfloat16:
|
||||
dtypes.append(torch.bfloat16)
|
||||
return dtypes
|
||||
|
||||
|
||||
def get_all_qint_dtypes() -> list[torch.dtype]:
|
||||
return [torch.qint8, torch.quint8, torch.qint32, torch.quint4x2, torch.quint2x4]
|
||||
|
||||
|
||||
def highest_precision_float(device):
|
||||
if torch.device(device).type == "mps":
|
||||
return torch.float32
|
||||
else:
|
||||
return torch.float64
|
||||
|
||||
|
||||
def highest_precision_complex(device):
|
||||
if torch.device(device).type == "mps":
|
||||
return torch.complex64
|
||||
else:
|
||||
return torch.complex128
|
||||
|
||||
|
||||
float_to_corresponding_complex_type_map = {
|
||||
torch.float16: torch.complex32,
|
||||
torch.float32: torch.complex64,
|
||||
torch.float64: torch.complex128,
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,322 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
# Torch
|
||||
import torch
|
||||
import torch.cuda
|
||||
import torch.jit
|
||||
import torch.jit._logging
|
||||
import torch.jit.frontend
|
||||
import torch.jit.quantized
|
||||
|
||||
# Testing utils
|
||||
from torch.testing._internal.common_dtype import floating_and_complex_types_and
|
||||
from torch.testing._internal.common_utils import TestCase, \
|
||||
freeze_rng_state, TemporaryFileName, enable_profiling_mode_for_profiling_tests, is_iterable_of_tensors
|
||||
from torch.testing._internal.common_utils import enable_profiling_mode # noqa: F401
|
||||
|
||||
# Standard library
|
||||
from itertools import chain
|
||||
from torch._C import TensorType
|
||||
|
||||
import io
|
||||
|
||||
def check_output_types(self, func, ref_outputs, args, kwargs):
|
||||
graph = getattr(func, 'last_graph', None)
|
||||
types = [o.type() for o in graph.outputs()]
|
||||
self.assertTrue(len(types) == 1)
|
||||
t = types[0]
|
||||
torch._C._jit_assert_is_instance(ref_outputs, t)
|
||||
|
||||
# Test names in this set are only checked for a single derivative
|
||||
nn_functional_single_grad = frozenset('test_nn_' + name for name in [
|
||||
'pdist',
|
||||
'multilabel_margin_loss',
|
||||
'max_unpool3d',
|
||||
'multi_margin_loss',
|
||||
'binary_cross_entropy',
|
||||
'binary_cross_entropy_size_average',
|
||||
'ctc_loss',
|
||||
'grid_sample',
|
||||
])
|
||||
|
||||
def check_against_reference(self, func, reference_func, output_func, args, kwargs=None,
|
||||
allow_unused=True, check_types=True, no_grad=False, no_gradgrad=False):
|
||||
"""Verifies a function performs identically to some reference implementation.
|
||||
|
||||
Commonly, this is used to verify that a JIT implementation
|
||||
(output_func) matches the behavior of the eager implementation
|
||||
(reference_func).
|
||||
"""
|
||||
kwargs = kwargs if kwargs else {}
|
||||
|
||||
def allSum(vs):
|
||||
if isinstance(vs, torch.Tensor):
|
||||
vs = (vs,)
|
||||
return sum((i + 1) * v.sum().abs() if v.dtype.is_complex else (i + 1) * v.sum()
|
||||
for i, v in enumerate(vs)
|
||||
if v is not None and v.dtype in floating_and_complex_types_and(torch.half, torch.bfloat16))
|
||||
|
||||
def clone_tensor(t, preserve_requires_grad):
|
||||
require_grad = preserve_requires_grad and t.requires_grad
|
||||
return t.detach().clone().requires_grad_(require_grad)
|
||||
|
||||
def clone_inputs(preserve_requires_grad: bool):
|
||||
inputs: list[torch.Tensor | list[torch.Tensor]] = []
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor):
|
||||
inputs.append(clone_tensor(arg, preserve_requires_grad))
|
||||
elif is_iterable_of_tensors(arg):
|
||||
inputs.append([clone_tensor(t, preserve_requires_grad) for t in arg])
|
||||
else:
|
||||
inputs.append(arg)
|
||||
|
||||
return inputs
|
||||
|
||||
# Returns tensors in args that requires_grad, including tensors in TensorList args
|
||||
def get_recording_tensors(args):
|
||||
recording_tensors: list[torch.Tensor] = []
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor) and arg.requires_grad:
|
||||
recording_tensors.append(arg)
|
||||
elif is_iterable_of_tensors(arg):
|
||||
recording_tensors.extend(filter(lambda t: t.requires_grad, arg))
|
||||
|
||||
return recording_tensors
|
||||
|
||||
# test no gradients case
|
||||
nograd_inputs = clone_inputs(preserve_requires_grad=False)
|
||||
outputs = self.runAndSaveRNG(reference_func, nograd_inputs, kwargs)
|
||||
with enable_profiling_mode_for_profiling_tests():
|
||||
outputs_test = self.runAndSaveRNG(func, nograd_inputs, kwargs)
|
||||
self.assertEqual(outputs, outputs_test)
|
||||
|
||||
if check_types:
|
||||
check_output_types(self, func, outputs_test, nograd_inputs, kwargs)
|
||||
|
||||
if no_grad:
|
||||
# skip grad tests
|
||||
return
|
||||
|
||||
with enable_profiling_mode_for_profiling_tests():
|
||||
# test single grad case
|
||||
recording_inputs = clone_inputs(preserve_requires_grad=True)
|
||||
recording_tensors = get_recording_tensors(recording_inputs)
|
||||
outputs = output_func(self.runAndSaveRNG(reference_func, recording_inputs, kwargs))
|
||||
grads = torch.autograd.grad(allSum(outputs), recording_tensors,
|
||||
allow_unused=allow_unused)
|
||||
outputs_test = output_func(self.runAndSaveRNG(func, recording_inputs, kwargs))
|
||||
grads_test = torch.autograd.grad(allSum(outputs_test), recording_tensors,
|
||||
allow_unused=allow_unused)
|
||||
self.assertEqual(outputs, outputs_test)
|
||||
self.assertEqual(grads, grads_test)
|
||||
# test the grad grad case
|
||||
if self._testMethodName in nn_functional_single_grad or no_gradgrad:
|
||||
return
|
||||
|
||||
outputs = output_func(self.runAndSaveRNG(reference_func, recording_inputs, kwargs))
|
||||
l1 = allSum(outputs)
|
||||
grads = torch.autograd.grad(l1, recording_tensors, create_graph=True,
|
||||
allow_unused=allow_unused)
|
||||
|
||||
l2 = (allSum(grads) * l1)
|
||||
grads2 = torch.autograd.grad(l2, recording_tensors, allow_unused=allow_unused)
|
||||
recording_inputs = clone_inputs(preserve_requires_grad=True)
|
||||
recording_tensors = get_recording_tensors(recording_inputs)
|
||||
outputs_test = output_func(self.runAndSaveRNG(func, recording_inputs, kwargs))
|
||||
l1_test = allSum(outputs_test)
|
||||
grads_test = torch.autograd.grad(
|
||||
l1_test, recording_tensors, create_graph=True, allow_unused=allow_unused)
|
||||
|
||||
l2_test = (allSum(grads_test) * l1_test)
|
||||
grads2_test = torch.autograd.grad(l2_test, recording_tensors, allow_unused=allow_unused)
|
||||
|
||||
self.assertEqual(outputs, outputs_test)
|
||||
self.assertEqual(grads, grads_test)
|
||||
for g2, g2_test in zip(grads2, grads2_test, strict=True):
|
||||
if g2 is None and g2_test is None:
|
||||
continue
|
||||
self.assertEqual(g2, g2_test, atol=5e-4, rtol=1e-4)
|
||||
|
||||
class JitCommonTestCase(TestCase):
|
||||
def createFunctionFromGraph(self, trace):
|
||||
graph = trace if isinstance(trace, torch._C.Graph) else trace.graph()
|
||||
return torch._C._create_function_from_graph("forward", graph)
|
||||
|
||||
def assertExportImport(self, trace, inputs):
|
||||
m = self.createFunctionFromGraph(trace)
|
||||
self.assertExportImportModule(m, inputs)
|
||||
|
||||
def assertExportImportModule(self, m, inputs):
|
||||
m_import = self.getExportImportCopy(m)
|
||||
a = self.runAndSaveRNG(m, inputs)
|
||||
b = self.runAndSaveRNG(m_import, inputs)
|
||||
self.assertEqual(a, b, "Results of original model and "
|
||||
"exported/imported version of model differed")
|
||||
|
||||
def runAndSaveRNG(self, func, inputs, kwargs=None):
|
||||
kwargs = kwargs if kwargs else {}
|
||||
with freeze_rng_state():
|
||||
results = func(*inputs, **kwargs)
|
||||
return results
|
||||
|
||||
def getExportImportCopy(self, m, also_test_file=True, map_location=None):
|
||||
buffer = io.BytesIO()
|
||||
torch.jit.save(m, buffer)
|
||||
buffer.seek(0)
|
||||
imported = torch.jit.load(buffer, map_location=map_location)
|
||||
|
||||
if not also_test_file:
|
||||
return imported
|
||||
|
||||
with TemporaryFileName() as fname:
|
||||
torch.jit.save(imported, fname)
|
||||
return torch.jit.load(fname, map_location=map_location)
|
||||
|
||||
def autoDiffErrorMessage(self, should_autodiff_node, nodes_not_in_diff_graph,
|
||||
fusion_nodes_not_found, non_fusible_nodes_being_fused,
|
||||
fusion_nodes_found, nodes_in_diff_graph):
|
||||
err_msg = "\nFailure in testing nodes' autodifferentiation. "
|
||||
if should_autodiff_node:
|
||||
err_msg += "One or more nodes were expected to be autodiffed, " \
|
||||
"but were not found in specified fusible/nonfusible " \
|
||||
"DifferentiableGraph groups. \nSpecifically:"
|
||||
# The node is intended to appear in a differentiable graph but doesn't
|
||||
diff_nodes_missing = []
|
||||
# The node is intended to appear in a differentiable graph
|
||||
# outside of a fusion group but instead is in a fusion group
|
||||
diff_nodes_in_fusion = []
|
||||
# The node is intended to appear in a fusion group but doesn't
|
||||
fusion_nodes_missing = []
|
||||
# The node is intended to appear in a fusion group but instead
|
||||
# is just in an outer differentiable graph
|
||||
fusion_nodes_in_diff = []
|
||||
for node in nodes_not_in_diff_graph:
|
||||
if node in non_fusible_nodes_being_fused:
|
||||
diff_nodes_in_fusion.append(node)
|
||||
else:
|
||||
diff_nodes_missing.append(node)
|
||||
for node in fusion_nodes_not_found:
|
||||
if node in nodes_in_diff_graph:
|
||||
fusion_nodes_in_diff.append(node)
|
||||
else:
|
||||
fusion_nodes_missing.append(node)
|
||||
if len(diff_nodes_missing) > 0:
|
||||
err_msg += f"\n {diff_nodes_missing} were not in one of the " \
|
||||
"DifferentiableGraphs when they were expected to be. " \
|
||||
"Did you intend for these nodes to be autodiffed? " \
|
||||
"If not, remove them from the list of nonfusible nodes."
|
||||
if len(diff_nodes_in_fusion) > 0:
|
||||
err_msg += f"\n {diff_nodes_in_fusion} were found in one of the FusionGroups " \
|
||||
"when they were expected to be just in a DifferentiableGraph. If it was " \
|
||||
"intended for these nodes to be in FusionGroups, reclassify these nodes as " \
|
||||
"fusible nodes. If these nodes were not intended to be fused, your " \
|
||||
"autodifferentiation logic might be wrong."
|
||||
if len(fusion_nodes_missing) > 0:
|
||||
err_msg += f"\n {fusion_nodes_missing} were not in one of the FusionGroups " \
|
||||
"of the DifferentiableGraphs when they were expected to be. " \
|
||||
"They were also not found in an outer DifferentiableGraph. Did you " \
|
||||
"intend for these nodes to be autodifferentiated? If not, you should " \
|
||||
"remove these nodes from the test's fusible nodes. Otherwise your " \
|
||||
"autodifferentiation logic might be wrong."
|
||||
if len(fusion_nodes_in_diff) > 0:
|
||||
err_msg += f"\n {fusion_nodes_in_diff} were not in one of the FusionGroups " \
|
||||
"of the DifferentiableGraphs when they were expected to be, " \
|
||||
"instead they were found just in an outer DifferentiableGraph. " \
|
||||
"Did you intend for these nodes to be fused? If not, you should " \
|
||||
"move these nodes into the test's nonfusible nodes. Otherwise your " \
|
||||
"autodifferentiation logic might be wrong."
|
||||
else:
|
||||
err_msg += "One or more nodes were not expected to be autodiffed " \
|
||||
"but were found in a DifferentiableGraph or in a FusionGroup " \
|
||||
"of a DifferentiableGraph. Did you intend for these nodes to be " \
|
||||
"autodiffed? If so, change this test to expect autodifferentiation. " \
|
||||
"\nSpecifically:"
|
||||
if len(fusion_nodes_found) > 0:
|
||||
err_msg += f"\n {fusion_nodes_found} were not expected to be in " \
|
||||
"one of the DifferentiableGraphs, but appeared in a FusionGroup " \
|
||||
"of a DifferentiableGraph. "
|
||||
if len(nodes_in_diff_graph) > 0:
|
||||
err_msg += f"\n {nodes_in_diff_graph} were not expected to " \
|
||||
"be in one of the DifferentiableGraphs but were."
|
||||
return err_msg
|
||||
|
||||
def assertAutodiffNode(self, graph, should_autodiff_node, nonfusible_nodes, fusible_nodes):
|
||||
diff_nodes = graph.findAllNodes('prim::DifferentiableGraph')
|
||||
diff_subgraphs = [node.g('Subgraph') for node in diff_nodes]
|
||||
|
||||
# Note: currently no tests have fusible_nodes
|
||||
fusion_nodes = list(chain.from_iterable([g.findAllNodes('prim::FusionGroup') for g in diff_subgraphs]))
|
||||
fusion_subgraphs = [node.g('Subgraph') for node in fusion_nodes]
|
||||
|
||||
# For any non-fusible node, it must show up in one of the DifferentiableGraphs.
|
||||
nodes_in_diff_graph = []
|
||||
nodes_not_in_diff_graph = []
|
||||
non_fusible_nodes_being_fused = []
|
||||
for node in nonfusible_nodes:
|
||||
if any(g.findNode(node) is not None for g in diff_subgraphs):
|
||||
nodes_in_diff_graph.append(node)
|
||||
else:
|
||||
nodes_not_in_diff_graph.append(node)
|
||||
if any(g.findNode(node) is not None for g in fusion_subgraphs):
|
||||
non_fusible_nodes_being_fused.append(node)
|
||||
found_all_nonfusible_nodes = len(nodes_in_diff_graph) == len(nonfusible_nodes)
|
||||
|
||||
# For any fusible node, it must show up in one of the FusionGroups in one of the DifferentiableGraphs.
|
||||
fusion_nodes_found = []
|
||||
fusion_nodes_not_found = []
|
||||
for node in fusible_nodes:
|
||||
if any(g.findNode(node) is not None for g in fusion_subgraphs):
|
||||
fusion_nodes_found.append(node)
|
||||
else:
|
||||
fusion_nodes_not_found.append(node)
|
||||
found_all_fusible_nodes = len(fusion_nodes_found) == len(fusible_nodes)
|
||||
|
||||
if should_autodiff_node is not None:
|
||||
err_msg = self.autoDiffErrorMessage(should_autodiff_node,
|
||||
nodes_not_in_diff_graph,
|
||||
fusion_nodes_not_found,
|
||||
non_fusible_nodes_being_fused,
|
||||
fusion_nodes_found,
|
||||
nodes_in_diff_graph)
|
||||
self.assertEqual(should_autodiff_node,
|
||||
found_all_nonfusible_nodes and found_all_fusible_nodes, err_msg)
|
||||
|
||||
def checkShapeAnalysis(self, out_sizes: list[int] | list[list[int]],
|
||||
traced_graph, assert_propagation, constant_prop=True):
|
||||
# repropagte input shapes provided by tracing,
|
||||
prev_symbolic_shapes_test_enabled = torch._C._jit_symbolic_shapes_test_mode_enabled()
|
||||
for enable_test_mode in [True, False]:
|
||||
# here we are testing allowing/disallowing substituting in complete shapes as constants,
|
||||
# disallowing constants helps stress test partial eval and substitution pipeline
|
||||
torch._C._jit_set_symbolic_shapes_test_mode(enable_test_mode)
|
||||
torch._C._jit_erase_non_input_shape_information(traced_graph)
|
||||
if constant_prop:
|
||||
torch._C._jit_pass_constant_propagation(traced_graph)
|
||||
torch._C._jit_pass_propagate_shapes_on_graph(traced_graph)
|
||||
# Add sizes to default tensor type to avoid checking something out of scope
|
||||
# and difficulties with tracer leaving in other parts of tensor type
|
||||
output = next(traced_graph.outputs()).type()
|
||||
|
||||
def test_type(type, actual_size):
|
||||
sizes = type.symbolic_sizes()
|
||||
out_type = TensorType.get().with_sizes(sizes)
|
||||
actual_type = TensorType.get().with_sizes(actual_size)
|
||||
|
||||
# always check actual shape is a subtype of the output
|
||||
self.assertTrue(actual_type.isSubtypeOf(out_type))
|
||||
|
||||
# and then if assertion flag is provided, check shape analysis
|
||||
# is successful
|
||||
if assert_propagation:
|
||||
self.assertEqual(out_type.sizes(), actual_size)
|
||||
|
||||
if output.isSubtypeOf(torch._C.TensorType.get()):
|
||||
test_type(output, out_sizes)
|
||||
else:
|
||||
tuple_elements = output.elements()
|
||||
for i in range(len(tuple_elements)):
|
||||
test_type(tuple_elements[i], out_sizes[i])
|
||||
|
||||
torch._C._jit_set_symbolic_shapes_test_mode(prev_symbolic_shapes_test_enabled)
|
||||
+26902
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,113 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import inspect
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def bf32_is_not_fp32():
|
||||
if not torch.backends.mkldnn.is_available():
|
||||
return False
|
||||
if not torch.ops.mkldnn._is_mkldnn_bf16_supported():
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def tf32_is_not_fp32():
|
||||
if not torch.backends.mkldnn.is_available():
|
||||
return False
|
||||
if not torch.cpu._is_amx_fp16_supported():
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def reduced_f32_off():
|
||||
old_matmul_precision = torch.backends.mkldnn.matmul.fp32_precision
|
||||
old_conv_precision = torch.backends.mkldnn.conv.fp32_precision
|
||||
try:
|
||||
torch.backends.mkldnn.matmul.fp32_precision = "ieee"
|
||||
torch.backends.mkldnn.conv.fp32_precision = "ieee"
|
||||
yield
|
||||
finally:
|
||||
torch.backends.mkldnn.matmul.fp32_precision = old_matmul_precision
|
||||
torch.backends.mkldnn.conv.fp32_precision = old_conv_precision
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def bf32_on(self, bf32_precision=1e-2):
|
||||
old_matmul_precision = torch.backends.mkldnn.matmul.fp32_precision
|
||||
old_conv_precision = torch.backends.mkldnn.conv.fp32_precision
|
||||
old_precision = self.precision
|
||||
try:
|
||||
torch.backends.mkldnn.matmul.fp32_precision = "bf16"
|
||||
torch.backends.mkldnn.conv.fp32_precision = "bf16"
|
||||
self.precision = bf32_precision
|
||||
yield
|
||||
finally:
|
||||
torch.backends.mkldnn.matmul.fp32_precision = old_matmul_precision
|
||||
torch.backends.mkldnn.conv.fp32_precision = old_conv_precision
|
||||
self.precision = old_precision
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def tf32_on(self, tf32_precision=1e-5):
|
||||
old_matmul_precision = torch.backends.mkldnn.matmul.fp32_precision
|
||||
old_conv_precision = torch.backends.mkldnn.conv.fp32_precision
|
||||
old_precision = self.precision
|
||||
try:
|
||||
torch.backends.mkldnn.matmul.fp32_precision = "tf32"
|
||||
torch.backends.mkldnn.conv.fp32_precision = "tf32"
|
||||
self.precision = tf32_precision
|
||||
yield
|
||||
finally:
|
||||
torch.backends.mkldnn.matmul.fp32_precision = old_matmul_precision
|
||||
torch.backends.mkldnn.conv.fp32_precision = old_conv_precision
|
||||
self.precision = old_precision
|
||||
|
||||
|
||||
# This is a wrapper that wraps a test to run this test three times, one with
|
||||
# reduced_f32 OFF, the others with reduced_f32 ON (including bf32 ON and tf32
|
||||
# ON). When running with reduced_f32 ON, it will use reduced precision (bf16/
|
||||
# tf32) as specified by the argument.
|
||||
def reduced_f32_on_and_off(bf32_precision=1e-2, tf32_precision=1e-5):
|
||||
def with_reduced_f32_disabled(self, function_call):
|
||||
with reduced_f32_off():
|
||||
function_call()
|
||||
|
||||
def with_bf32_enabled(self, function_call):
|
||||
with bf32_on(self, bf32_precision):
|
||||
function_call()
|
||||
|
||||
def with_tf32_enabled(self, function_call):
|
||||
with tf32_on(self, tf32_precision):
|
||||
function_call()
|
||||
|
||||
def wrapper(f):
|
||||
params = inspect.signature(f).parameters
|
||||
arg_names = tuple(params.keys())
|
||||
|
||||
@functools.wraps(f)
|
||||
def wrapped(*args, **kwargs):
|
||||
kwargs.update(zip(arg_names, args, strict=False))
|
||||
cond = True
|
||||
if "device" in kwargs:
|
||||
cond = cond and (torch.device(kwargs["device"]).type == "cpu")
|
||||
if "dtype" in kwargs:
|
||||
cond = cond and (kwargs["dtype"] == torch.float)
|
||||
bf32_cond = cond and bf32_is_not_fp32()
|
||||
tf32_cond = cond and tf32_is_not_fp32()
|
||||
if bf32_cond or tf32_cond:
|
||||
with_reduced_f32_disabled(kwargs["self"], lambda: f(**kwargs))
|
||||
if bf32_cond:
|
||||
with_bf32_enabled(kwargs["self"], lambda: f(**kwargs))
|
||||
if tf32_cond:
|
||||
with_tf32_enabled(kwargs["self"], lambda: f(**kwargs))
|
||||
else:
|
||||
f(**kwargs)
|
||||
|
||||
return wrapped
|
||||
|
||||
return wrapper
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+258
@@ -0,0 +1,258 @@
|
||||
# Owner(s): ["oncall: pt2"]
|
||||
|
||||
"""
|
||||
Shared xfail lists for unbacked symint tests.
|
||||
|
||||
These lists are used by both test_ops_unbacked.py (base tensor tests)
|
||||
and test_dtensor_ops.py (DTensor tests with unbacked dimensions).
|
||||
"""
|
||||
|
||||
|
||||
def xfail(op_name, variant_name="", *, device_type=None, dtypes=None):
|
||||
return (op_name, variant_name, device_type, dtypes, True)
|
||||
|
||||
|
||||
def skip(op_name, variant_name="", *, device_type=None, dtypes=None):
|
||||
return (op_name, variant_name, device_type, dtypes, False)
|
||||
|
||||
|
||||
# Ops that have data-dependent errors with unbacked dimensions.
|
||||
# These fail at the base tensor level (not DTensor-specific).
|
||||
ops_dde_xfail = {
|
||||
xfail("_chunk_cat"),
|
||||
xfail("_unsafe_masked_index_put_accumulate"),
|
||||
xfail("_upsample_bilinear2d_aa"),
|
||||
xfail("addmv"),
|
||||
xfail("allclose"),
|
||||
xfail("as_strided_scatter"),
|
||||
xfail("baddbmm"),
|
||||
xfail("bernoulli"),
|
||||
xfail("cauchy"),
|
||||
xfail("cdist"),
|
||||
xfail("cholesky"),
|
||||
xfail("chunk"),
|
||||
xfail("combinations"),
|
||||
xfail("corrcoef"),
|
||||
xfail("cov"),
|
||||
xfail("cross"),
|
||||
xfail("cummax"),
|
||||
xfail("cummin"),
|
||||
xfail("cumulative_trapezoid"),
|
||||
xfail("diagonal_scatter"),
|
||||
xfail("diff"),
|
||||
xfail("dsplit"),
|
||||
xfail("equal"),
|
||||
xfail("exponential"),
|
||||
xfail("fft.fft"),
|
||||
xfail("fft.fft2"),
|
||||
xfail("fft.fftn"),
|
||||
xfail("fft.fftshift"),
|
||||
xfail("fft.hfft"),
|
||||
xfail("fft.hfft2"),
|
||||
xfail("fft.hfftn"),
|
||||
xfail("fft.ifft"),
|
||||
xfail("fft.ifft2"),
|
||||
xfail("fft.ifftn"),
|
||||
xfail("fft.ifftshift"),
|
||||
xfail("fft.ihfft"),
|
||||
xfail("fft.ihfft2"),
|
||||
xfail("fft.ihfftn"),
|
||||
xfail("fft.irfft"),
|
||||
xfail("fft.irfft2"),
|
||||
xfail("fft.irfftn"),
|
||||
xfail("fft.rfft"),
|
||||
xfail("fft.rfft2"),
|
||||
xfail("fft.rfftn"),
|
||||
xfail("float"),
|
||||
xfail("geometric"),
|
||||
xfail("geqrf"),
|
||||
xfail("gradient"),
|
||||
xfail("grid_sampler_2d"),
|
||||
xfail("hash_tensor"),
|
||||
xfail("histogram"),
|
||||
xfail("histogramdd"),
|
||||
xfail("hsplit"),
|
||||
xfail("index_fill"),
|
||||
xfail("inner"),
|
||||
xfail("kron"),
|
||||
xfail("linalg.cond"),
|
||||
xfail("linalg.cross"),
|
||||
xfail("linalg.householder_product"),
|
||||
xfail("linalg.ldl_solve"),
|
||||
xfail("linalg.lstsq"),
|
||||
xfail("linalg.lstsq", "grad_oriented"),
|
||||
xfail("linalg.lu_solve"),
|
||||
xfail("linalg.matrix_norm"),
|
||||
xfail("linalg.matrix_power"),
|
||||
xfail("linalg.matrix_rank"),
|
||||
xfail("linalg.matrix_rank", "hermitian"),
|
||||
xfail("linalg.multi_dot"),
|
||||
xfail("linalg.norm"),
|
||||
xfail("linalg.norm", "subgradients_at_zero"),
|
||||
xfail("linalg.pinv"),
|
||||
xfail("linalg.pinv", "hermitian"),
|
||||
xfail("linalg.pinv", "singular"),
|
||||
xfail("linalg.qr"),
|
||||
xfail("linalg.solve"),
|
||||
xfail("linalg.solve_ex"),
|
||||
xfail("linalg.solve_triangular"),
|
||||
xfail("linalg.tensorinv"),
|
||||
xfail("linalg.tensorsolve"),
|
||||
xfail("linalg.vander"),
|
||||
xfail("log_normal"),
|
||||
xfail("logdet"),
|
||||
xfail("logsumexp"),
|
||||
xfail("lu_solve"),
|
||||
xfail("lu_unpack"),
|
||||
xfail("masked.amax"),
|
||||
xfail("masked.amin"),
|
||||
xfail("masked.argmax"),
|
||||
xfail("masked.argmin"),
|
||||
xfail("masked.cumprod"),
|
||||
xfail("masked.cumsum"),
|
||||
xfail("masked.log_softmax"),
|
||||
xfail("masked.logaddexp"),
|
||||
xfail("masked.logsumexp"),
|
||||
xfail("masked.mean"),
|
||||
xfail("masked.median"),
|
||||
xfail("masked.norm"),
|
||||
xfail("masked.prod"),
|
||||
xfail("masked.softmax"),
|
||||
xfail("masked.softmin"),
|
||||
xfail("masked.std"),
|
||||
xfail("masked.sum"),
|
||||
xfail("masked.var"),
|
||||
xfail("max_pool2d_with_indices_backward"),
|
||||
xfail("multinomial"),
|
||||
xfail("nn.functional.adaptive_avg_pool1d"),
|
||||
xfail("nn.functional.adaptive_avg_pool2d"),
|
||||
xfail("nn.functional.adaptive_avg_pool3d"),
|
||||
xfail("nn.functional.adaptive_max_pool1d"),
|
||||
xfail("nn.functional.adaptive_max_pool2d"),
|
||||
xfail("nn.functional.adaptive_max_pool3d"),
|
||||
xfail("nn.functional.alpha_dropout"),
|
||||
xfail("nn.functional.avg_pool1d"),
|
||||
xfail("nn.functional.avg_pool2d"),
|
||||
xfail("nn.functional.avg_pool3d"),
|
||||
xfail("nn.functional.batch_norm"),
|
||||
xfail("nn.functional.bilinear"),
|
||||
xfail("nn.functional.binary_cross_entropy"),
|
||||
xfail("nn.functional.binary_cross_entropy_with_logits"),
|
||||
xfail("nn.functional.channel_shuffle"),
|
||||
xfail("nn.functional.cosine_similarity"),
|
||||
xfail("nn.functional.cross_entropy"),
|
||||
xfail("nn.functional.ctc_loss"),
|
||||
xfail("nn.functional.dropout"),
|
||||
xfail("nn.functional.dropout2d"),
|
||||
xfail("nn.functional.dropout3d"),
|
||||
xfail("nn.functional.embedding"),
|
||||
xfail("nn.functional.embedding_bag"),
|
||||
xfail("nn.functional.feature_alpha_dropout", "with_train"),
|
||||
xfail("nn.functional.feature_alpha_dropout", "without_train"),
|
||||
xfail("nn.functional.fractional_max_pool2d"),
|
||||
xfail("nn.functional.fractional_max_pool3d"),
|
||||
xfail("nn.functional.gaussian_nll_loss"),
|
||||
xfail("nn.functional.grid_sample"),
|
||||
xfail("nn.functional.group_norm"),
|
||||
xfail("nn.functional.huber_loss"),
|
||||
xfail("nn.functional.instance_norm"),
|
||||
xfail("nn.functional.interpolate", "area"),
|
||||
xfail("nn.functional.interpolate", "bicubic"),
|
||||
xfail("nn.functional.interpolate", "bilinear"),
|
||||
xfail("nn.functional.interpolate", "linear"),
|
||||
xfail("nn.functional.interpolate", "trilinear"),
|
||||
xfail("nn.functional.l1_loss"),
|
||||
xfail("nn.functional.local_response_norm"),
|
||||
xfail("nn.functional.max_pool1d"),
|
||||
xfail("nn.functional.max_pool2d"),
|
||||
xfail("nn.functional.max_pool3d"),
|
||||
xfail("nn.functional.max_unpool1d"),
|
||||
xfail("nn.functional.max_unpool1d", "grad"),
|
||||
xfail("nn.functional.max_unpool2d"),
|
||||
xfail("nn.functional.max_unpool2d", "grad"),
|
||||
xfail("nn.functional.max_unpool3d"),
|
||||
xfail("nn.functional.max_unpool3d", "grad"),
|
||||
xfail("nn.functional.mse_loss"),
|
||||
xfail("nn.functional.multi_head_attention_forward"),
|
||||
xfail("nn.functional.multilabel_margin_loss"),
|
||||
xfail("nn.functional.nll_loss"),
|
||||
xfail("nn.functional.pad", "circular"),
|
||||
xfail("nn.functional.pad", "reflect"),
|
||||
xfail("nn.functional.pad", "replicate"),
|
||||
xfail("nn.functional.pad", "replicate_negative"),
|
||||
xfail("nn.functional.pdist"),
|
||||
xfail("nn.functional.pixel_shuffle"),
|
||||
xfail("nn.functional.prelu"),
|
||||
xfail("nn.functional.rrelu"),
|
||||
xfail("nn.functional.scaled_dot_product_attention"),
|
||||
xfail("nn.functional.smooth_l1_loss"),
|
||||
xfail("nn.functional.unfold"),
|
||||
xfail("nn.functional.upsample_bilinear"),
|
||||
xfail("normal"),
|
||||
xfail("normal", "in_place"),
|
||||
xfail("normal", "number_mean"),
|
||||
xfail("ormqr"),
|
||||
xfail("pca_lowrank"),
|
||||
xfail("pinverse"),
|
||||
xfail("qr"),
|
||||
xfail("rand_like"),
|
||||
xfail("randint_like"),
|
||||
xfail("randn_like"),
|
||||
xfail("repeat_interleave"),
|
||||
xfail("resize_"),
|
||||
xfail("resize_as_"),
|
||||
xfail("roll"),
|
||||
xfail("searchsorted"),
|
||||
xfail("sparse.mm", "reduce"),
|
||||
xfail("split"),
|
||||
xfail("stft"),
|
||||
xfail("svd_lowrank"),
|
||||
xfail("sum_to_size"),
|
||||
xfail("take"),
|
||||
xfail("take_along_dim"),
|
||||
xfail("tensordot"),
|
||||
xfail("tensor_split"),
|
||||
xfail("to_sparse"),
|
||||
xfail("trapezoid"),
|
||||
xfail("trapz"),
|
||||
xfail("unbind"),
|
||||
xfail("unbind_copy"),
|
||||
xfail("uniform"),
|
||||
xfail("unsafe_chunk"),
|
||||
xfail("unsafe_split"),
|
||||
xfail("vsplit"),
|
||||
}
|
||||
|
||||
# Ops that skip for unbacked tests (no valid samples with markable dims)
|
||||
ops_unbacked_skip = {
|
||||
skip("arange"),
|
||||
skip("broadcast_shapes"),
|
||||
skip("empty"),
|
||||
skip("empty_permuted"),
|
||||
skip("empty_strided"),
|
||||
skip("eye"),
|
||||
skip("full"),
|
||||
skip("item"),
|
||||
skip("linspace"),
|
||||
skip("linspace", "tensor_overload"),
|
||||
skip("logspace"),
|
||||
skip("logspace", "tensor_overload"),
|
||||
skip("ones"),
|
||||
skip("randint"),
|
||||
skip("randn"),
|
||||
skip("scalar_tensor"),
|
||||
skip("signal.windows.bartlett"),
|
||||
skip("signal.windows.blackman"),
|
||||
skip("signal.windows.cosine"),
|
||||
skip("signal.windows.exponential"),
|
||||
skip("signal.windows.gaussian"),
|
||||
skip("signal.windows.general_cosine"),
|
||||
skip("signal.windows.general_hamming"),
|
||||
skip("signal.windows.hamming"),
|
||||
skip("signal.windows.hann"),
|
||||
skip("signal.windows.kaiser"),
|
||||
skip("signal.windows.nuttall"),
|
||||
skip("zeros"),
|
||||
# Sparse ops that can't be deepcopied
|
||||
skip("sparse.sampled_addmm"),
|
||||
}
|
||||
+2316
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,385 @@
|
||||
# Owner(s): ["module: unknown"]
|
||||
|
||||
from typing import Any
|
||||
from torch.ao.pruning import BaseSparsifier
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
class ImplementedSparsifier(BaseSparsifier):
|
||||
def __init__(self, **kwargs: dict[str, Any]) -> None:
|
||||
super().__init__(defaults=kwargs)
|
||||
|
||||
def update_mask(self, module: nn.Module, tensor_name: str, **kwargs: dict[str, Any]) -> None:
|
||||
module.parametrizations.weight[0].mask[0] = 0 # type: ignore[index, union-attr]
|
||||
linear_state = self.state['linear1.weight']
|
||||
linear_state['step_count'] = linear_state.get('step_count', 0) + 1
|
||||
|
||||
|
||||
class MockSparseLinear(nn.Linear):
|
||||
"""
|
||||
This class is a MockSparseLinear class to check convert functionality.
|
||||
It is the same as a normal Linear layer, except with a different type, as
|
||||
well as an additional from_dense method.
|
||||
"""
|
||||
@classmethod
|
||||
def from_dense(cls, mod: nn.Linear) -> 'MockSparseLinear':
|
||||
"""
|
||||
"""
|
||||
linear = cls(mod.in_features,
|
||||
mod.out_features)
|
||||
return linear
|
||||
|
||||
|
||||
def rows_are_subset(subset_tensor: torch.Tensor, superset_tensor: torch.Tensor) -> bool:
|
||||
"""
|
||||
Checks to see if all rows in subset tensor are present in the superset tensor
|
||||
"""
|
||||
i = 0
|
||||
for row in subset_tensor:
|
||||
while i < len(superset_tensor):
|
||||
if not torch.equal(row, superset_tensor[i]):
|
||||
i += 1
|
||||
else:
|
||||
break
|
||||
else:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class SimpleLinear(nn.Module):
|
||||
r"""Model with only Linear layers without biases, some wrapped in a Sequential,
|
||||
some following the Sequential. Used to test basic pruned Linear-Linear fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Linear(7, 5, bias=False),
|
||||
nn.Linear(5, 6, bias=False),
|
||||
nn.Linear(6, 4, bias=False),
|
||||
)
|
||||
self.linear1 = nn.Linear(4, 4, bias=False)
|
||||
self.linear2 = nn.Linear(4, 10, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.linear1(x)
|
||||
x = self.linear2(x)
|
||||
return x
|
||||
|
||||
|
||||
class LinearBias(nn.Module):
|
||||
r"""Model with only Linear layers, alternating layers with biases,
|
||||
wrapped in a Sequential. Used to test pruned Linear-Bias-Linear fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Linear(7, 5, bias=True),
|
||||
nn.Linear(5, 6, bias=False),
|
||||
nn.Linear(6, 3, bias=True),
|
||||
nn.Linear(3, 3, bias=True),
|
||||
nn.Linear(3, 10, bias=False),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
return x
|
||||
|
||||
|
||||
class LinearActivation(nn.Module):
|
||||
r"""Model with only Linear layers, some with bias, some in a Sequential and some following.
|
||||
Activation functions modules in between each Linear in the Sequential, and each outside layer.
|
||||
Used to test pruned Linear(Bias)-Activation-Linear fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Linear(7, 5, bias=True),
|
||||
nn.ReLU(),
|
||||
nn.Linear(5, 6, bias=False),
|
||||
nn.Tanh(),
|
||||
nn.Linear(6, 4, bias=True),
|
||||
)
|
||||
self.linear1 = nn.Linear(4, 3, bias=True)
|
||||
self.act1 = nn.ReLU()
|
||||
self.linear2 = nn.Linear(3, 10, bias=False)
|
||||
self.act2 = nn.Tanh()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.linear1(x)
|
||||
x = self.act1(x)
|
||||
x = self.linear2(x)
|
||||
x = self.act2(x)
|
||||
return x
|
||||
|
||||
|
||||
class LinearActivationFunctional(nn.Module):
|
||||
r"""Model with only Linear layers, some with bias, some in a Sequential and some following.
|
||||
Activation functions modules in between each Linear in the Sequential, and functional
|
||||
activationals are called in between each outside layer.
|
||||
Used to test pruned Linear(Bias)-Activation-Linear fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Linear(7, 5, bias=True),
|
||||
nn.ReLU(),
|
||||
nn.Linear(5, 6, bias=False),
|
||||
nn.ReLU(),
|
||||
nn.Linear(6, 4, bias=True),
|
||||
)
|
||||
self.linear1 = nn.Linear(4, 3, bias=True)
|
||||
self.linear2 = nn.Linear(3, 8, bias=False)
|
||||
self.linear3 = nn.Linear(8, 10, bias=False)
|
||||
self.act1 = nn.ReLU()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.linear1(x)
|
||||
x = F.relu(x)
|
||||
x = self.linear2(x)
|
||||
x = F.relu(x)
|
||||
x = self.linear3(x)
|
||||
x = F.relu(x)
|
||||
return x
|
||||
|
||||
|
||||
class SimpleConv2d(nn.Module):
|
||||
r"""Model with only Conv2d layers, all without bias, some in a Sequential and some following.
|
||||
Used to test pruned Conv2d-Conv2d fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Conv2d(1, 32, 3, 1, bias=False),
|
||||
nn.Conv2d(32, 64, 3, 1, bias=False),
|
||||
)
|
||||
self.conv2d1 = nn.Conv2d(64, 48, 3, 1, bias=False)
|
||||
self.conv2d2 = nn.Conv2d(48, 52, 3, 1, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.conv2d1(x)
|
||||
x = self.conv2d2(x)
|
||||
return x
|
||||
|
||||
|
||||
class Conv2dBias(nn.Module):
|
||||
r"""Model with only Conv2d layers, some with bias, some in a Sequential and some outside.
|
||||
Used to test pruned Conv2d-Bias-Conv2d fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Conv2d(1, 32, 3, 1, bias=True),
|
||||
nn.Conv2d(32, 32, 3, 1, bias=True),
|
||||
nn.Conv2d(32, 64, 3, 1, bias=False),
|
||||
)
|
||||
self.conv2d1 = nn.Conv2d(64, 48, 3, 1, bias=True)
|
||||
self.conv2d2 = nn.Conv2d(48, 52, 3, 1, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.conv2d1(x)
|
||||
x = self.conv2d2(x)
|
||||
return x
|
||||
|
||||
|
||||
class Conv2dActivation(nn.Module):
|
||||
r"""Model with only Conv2d layers, some with bias, some in a Sequential and some following.
|
||||
Activation function modules in between each Sequential layer, functional activations called
|
||||
in-between each outside layer.
|
||||
Used to test pruned Conv2d-Bias-Activation-Conv2d fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Conv2d(1, 32, 3, 1, bias=True),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(32, 64, 3, 1, bias=True),
|
||||
nn.Tanh(),
|
||||
nn.Conv2d(64, 64, 3, 1, bias=False),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self.conv2d1 = nn.Conv2d(64, 48, 3, 1, bias=False)
|
||||
self.conv2d2 = nn.Conv2d(48, 52, 3, 1, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.conv2d1(x)
|
||||
x = F.relu(x)
|
||||
x = self.conv2d2(x)
|
||||
x = F.hardtanh(x)
|
||||
return x
|
||||
|
||||
|
||||
class Conv2dPadBias(nn.Module):
|
||||
r"""Model with only Conv2d layers, all with bias and some with padding > 0,
|
||||
some in a Sequential and some following. Activation function modules in between each layer.
|
||||
Used to test that bias is propagated correctly in the special case of
|
||||
pruned Conv2d-Bias-(Activation)Conv2d fusion, when the second Conv2d layer has padding > 0."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Conv2d(1, 32, 3, 1, padding=1, bias=True),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(32, 32, 3, 1, bias=False),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(32, 32, 3, 1, padding=1, bias=True),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(32, 32, 3, 1, padding=1, bias=True),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(32, 64, 3, 1, bias=True),
|
||||
nn.Tanh(),
|
||||
)
|
||||
self.conv2d1 = nn.Conv2d(64, 48, 3, 1, padding=1, bias=True)
|
||||
self.act1 = nn.ReLU()
|
||||
self.conv2d2 = nn.Conv2d(48, 52, 3, 1, padding=1, bias=True)
|
||||
self.act2 = nn.Tanh()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.conv2d1(x)
|
||||
x = self.act1(x)
|
||||
x = self.conv2d2(x)
|
||||
x = self.act2(x)
|
||||
return x
|
||||
|
||||
|
||||
class Conv2dPool(nn.Module):
|
||||
r"""Model with only Conv2d layers, all with bias, some in a Sequential and some following.
|
||||
Activation function modules in between each layer, Pool2d modules in between each layer.
|
||||
Used to test pruned Conv2d-Pool2d-Conv2d fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Conv2d(1, 32, kernel_size=3, padding=1, bias=True),
|
||||
nn.MaxPool2d(kernel_size=2, stride=2, padding=1),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(32, 64, kernel_size=3, padding=1, bias=True),
|
||||
nn.Tanh(),
|
||||
nn.AvgPool2d(kernel_size=2, stride=2, padding=1),
|
||||
)
|
||||
self.conv2d1 = nn.Conv2d(64, 48, kernel_size=3, padding=1, bias=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2, padding=1)
|
||||
self.af1 = nn.ReLU()
|
||||
self.conv2d2 = nn.Conv2d(48, 52, kernel_size=3, padding=1, bias=True)
|
||||
self.conv2d3 = nn.Conv2d(52, 52, kernel_size=3, padding=1, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.conv2d1(x)
|
||||
x = self.maxpool(x)
|
||||
x = self.af1(x)
|
||||
x = self.conv2d2(x)
|
||||
x = F.avg_pool2d(x, kernel_size=2, stride=2, padding=1)
|
||||
x = F.relu(x)
|
||||
x = self.conv2d3(x)
|
||||
return x
|
||||
|
||||
|
||||
class Conv2dPoolFlattenFunctional(nn.Module):
|
||||
r"""Model with Conv2d layers, all with bias, some in a Sequential and some following, and then a Pool2d
|
||||
and a functional Flatten followed by a Linear layer.
|
||||
Activation functions and Pool2ds in between each layer also.
|
||||
Used to test pruned Conv2d-Pool2d-Flatten-Linear fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Conv2d(1, 3, kernel_size=3, padding=1, bias=True),
|
||||
nn.MaxPool2d(kernel_size=2, stride=2, padding=1),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(3, 5, kernel_size=3, padding=1, bias=True),
|
||||
nn.Tanh(),
|
||||
nn.AvgPool2d(kernel_size=2, stride=2, padding=1),
|
||||
)
|
||||
self.conv2d1 = nn.Conv2d(5, 7, kernel_size=3, padding=1, bias=True)
|
||||
self.af1 = nn.ReLU()
|
||||
self.conv2d2 = nn.Conv2d(7, 11, kernel_size=3, padding=1, bias=True)
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
self.fc = nn.Linear(11, 13, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.conv2d1(x)
|
||||
x = F.max_pool2d(x, kernel_size=2, stride=2, padding=1)
|
||||
x = self.af1(x)
|
||||
x = self.conv2d2(x)
|
||||
x = self.avg_pool(x)
|
||||
x = torch.flatten(x, 1) # test functional flatten
|
||||
x = self.fc(x)
|
||||
return x
|
||||
|
||||
|
||||
class Conv2dPoolFlatten(nn.Module):
|
||||
r"""Model with Conv2d layers, all with bias, some in a Sequential and some following, and then a Pool2d
|
||||
and a Flatten module followed by a Linear layer.
|
||||
Activation functions and Pool2ds in between each layer also.
|
||||
Used to test pruned Conv2d-Pool2d-Flatten-Linear fusion."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seq = nn.Sequential(
|
||||
nn.Conv2d(1, 3, kernel_size=3, padding=1, bias=True),
|
||||
nn.MaxPool2d(kernel_size=2, stride=2, padding=1),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(3, 5, kernel_size=3, padding=1, bias=True),
|
||||
nn.Tanh(),
|
||||
nn.AvgPool2d(kernel_size=2, stride=2, padding=1),
|
||||
)
|
||||
self.conv2d1 = nn.Conv2d(5, 7, kernel_size=3, padding=1, bias=True)
|
||||
self.af1 = nn.ReLU()
|
||||
self.conv2d2 = nn.Conv2d(7, 11, kernel_size=3, padding=1, bias=True)
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d((2, 2))
|
||||
self.flatten = nn.Flatten()
|
||||
self.fc = nn.Linear(44, 13, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.seq(x)
|
||||
x = self.conv2d1(x)
|
||||
x = F.max_pool2d(x, kernel_size=2, stride=2, padding=1)
|
||||
x = self.af1(x)
|
||||
x = self.conv2d2(x)
|
||||
x = self.avg_pool(x)
|
||||
x = self.flatten(x)
|
||||
x = self.fc(x)
|
||||
return x
|
||||
|
||||
|
||||
class LSTMLinearModel(nn.Module):
|
||||
"""Container module with an encoder, a recurrent module, and a linear."""
|
||||
|
||||
def __init__(
|
||||
self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers)
|
||||
self.linear = nn.Linear(hidden_dim, output_dim)
|
||||
|
||||
def forward(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
output, _hidden = self.lstm(input)
|
||||
decoded = self.linear(output)
|
||||
return decoded, output
|
||||
|
||||
|
||||
class LSTMLayerNormLinearModel(nn.Module):
|
||||
"""Container module with an LSTM, a LayerNorm, and a linear."""
|
||||
|
||||
def __init__(
|
||||
self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers)
|
||||
self.norm = nn.LayerNorm(hidden_dim)
|
||||
self.linear = nn.Linear(hidden_dim, output_dim)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
x, state = self.lstm(x)
|
||||
x = self.norm(x)
|
||||
x = self.linear(x)
|
||||
return x, state
|
||||
+3323
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,693 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
r"""Importing this file includes common utility methods for checking quantized
|
||||
tensors and modules.
|
||||
"""
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from contextlib import contextmanager
|
||||
from torch.testing._internal.common_utils import TEST_WITH_TSAN, IS_PPC, IS_MACOS, IS_WINDOWS, IS_ARM64
|
||||
|
||||
supported_qengines = list(torch.backends.quantized.supported_engines)
|
||||
# Note: We currently do not run QNNPACK tests on WINDOWS and MACOS as it is flaky. Issue #29326
|
||||
# QNNPACK is not supported on PPC
|
||||
if 'qnnpack' in supported_qengines and any([IS_PPC, TEST_WITH_TSAN, IS_MACOS, IS_WINDOWS]):
|
||||
supported_qengines.remove('qnnpack')
|
||||
# FBGEMM and x86 engines require x86 architecture with AVX2/AVX512 support
|
||||
# They are not supported on ARM64 architectures
|
||||
if IS_ARM64:
|
||||
supported_qengines = [qe for qe in supported_qengines if qe not in ('fbgemm', 'x86')]
|
||||
|
||||
def _conv_output_shape(input_size, kernel_size, padding, stride, dilation,
|
||||
output_padding=0):
|
||||
"""Computes the output shape given convolution parameters."""
|
||||
return np.floor((input_size + 2 * padding - kernel_size - (kernel_size - 1)
|
||||
* (dilation - 1)) / stride) + 2 * output_padding + 1
|
||||
|
||||
# Quantization references
|
||||
def _quantize(x, scale, zero_point, qmin=None, qmax=None, dtype=np.uint8):
|
||||
"""Quantizes a numpy array."""
|
||||
if qmin is None:
|
||||
qmin = np.iinfo(dtype).min
|
||||
if qmax is None:
|
||||
qmax = np.iinfo(dtype).max
|
||||
qx = np.round(x / scale + zero_point).astype(np.int64)
|
||||
qx = np.clip(qx, qmin, qmax)
|
||||
qx = qx.astype(dtype)
|
||||
return qx
|
||||
|
||||
|
||||
def _dequantize(qx, scale, zero_point):
|
||||
"""Dequantizes a numpy array."""
|
||||
x = (qx.astype(float) - zero_point) * scale
|
||||
return x
|
||||
|
||||
|
||||
def _requantize(x, multiplier, zero_point, qmin=0, qmax=255, qtype=np.uint8):
|
||||
"""Requantizes a numpy array, i.e., intermediate int32 or int16 values are
|
||||
converted back to given type"""
|
||||
qx = (x * multiplier).round() + zero_point
|
||||
qx = np.clip(qx, qmin, qmax).astype(qtype)
|
||||
return qx
|
||||
|
||||
def _calculate_dynamic_qparams(X, dtype, reduce_range=False, qscheme=torch.per_tensor_affine):
|
||||
"""Calculate the dynamic quantization parameters (scale, zero_point)
|
||||
according to the min and max element of the tensor"""
|
||||
if qscheme not in (torch.per_tensor_affine, torch.per_tensor_symmetric):
|
||||
raise AssertionError(
|
||||
f"Expected qscheme to be per_tensor_affine or per_tensor_symmetric, got {qscheme}"
|
||||
)
|
||||
if qscheme == torch.per_tensor_symmetric:
|
||||
if dtype != torch.qint8:
|
||||
raise AssertionError(
|
||||
f"Expected dtype to be torch.qint8 for symmetric qscheme, got {dtype}"
|
||||
)
|
||||
if isinstance(X, torch.Tensor):
|
||||
X = X.numpy()
|
||||
if dtype == torch.qint8:
|
||||
if reduce_range:
|
||||
qmin, qmax = -64, 63
|
||||
else:
|
||||
qmin, qmax = -128, 127
|
||||
else: # dtype == torch.quint8
|
||||
if reduce_range:
|
||||
qmin, qmax = 0, 127
|
||||
else:
|
||||
qmin, qmax = 0, 255
|
||||
min_val = X.min()
|
||||
max_val = X.max()
|
||||
is_symmetric = (qscheme == torch.per_tensor_symmetric)
|
||||
if min_val == max_val:
|
||||
scale = 1.0
|
||||
zero_point = 0
|
||||
else:
|
||||
if is_symmetric:
|
||||
max_val = max(max_val, -min_val)
|
||||
min_val = -max_val
|
||||
scale = (max_val - min_val) / (qmax - qmin)
|
||||
scale = max(scale, np.finfo(np.float32).eps)
|
||||
zero_point = 0
|
||||
else:
|
||||
max_val = max(max_val, 0.0)
|
||||
min_val = min(min_val, 0.0)
|
||||
scale = (max_val - min_val) / (qmax - qmin)
|
||||
scale = max(scale, np.finfo(np.float32).eps)
|
||||
zero_point = qmin - round(min_val / scale)
|
||||
zero_point = max(qmin, zero_point)
|
||||
zero_point = min(qmax, zero_point)
|
||||
return [float(scale), int(zero_point)]
|
||||
|
||||
def _calculate_dynamic_per_channel_qparams(X, dtype):
|
||||
"""Calculate the dynamic quantization parameters (scale, zero_point)
|
||||
according to the min and max element of the tensor"""
|
||||
if isinstance(X, torch.Tensor):
|
||||
X = X.numpy()
|
||||
qmin, qmax = torch.iinfo(dtype).min, torch.iinfo(dtype).max
|
||||
n_levels = qmax - qmin
|
||||
scale = np.zeros(X.shape[0], dtype=np.float64)
|
||||
zero_point = np.zeros(X.shape[0], dtype=np.int64)
|
||||
for i in range(zero_point.shape[0]):
|
||||
min_val = X.min()
|
||||
max_val = X.max()
|
||||
if min_val == max_val:
|
||||
scale[i] = 1.0
|
||||
zero_point[i] = 0
|
||||
else:
|
||||
max_val = max(max_val, 0.0)
|
||||
min_val = min(min_val, 0.0)
|
||||
scale[i] = (max_val - min_val) / n_levels
|
||||
scale[i] = max(scale[i], np.finfo(np.float32).eps)
|
||||
zero_point[i] = qmin - round(min_val / scale[i])
|
||||
zero_point[i] = max(qmin, zero_point[i])
|
||||
zero_point[i] = min(qmax, zero_point[i])
|
||||
|
||||
return scale, zero_point
|
||||
|
||||
def _snr(x, x_hat):
|
||||
"""Calculates the signal to noise ratio and returns the signal and noise
|
||||
power, as well as the SNR in dB.
|
||||
If the input is a list/tuple this function is called recursively on each
|
||||
element. The result will have the same nested structure as the inputs.
|
||||
|
||||
Args:
|
||||
x, x_hat: Either a tensor or a nested list/tuple of tensors.
|
||||
Returns:
|
||||
signal, noise, SNR(in dB): Either floats or a nested list of floats
|
||||
"""
|
||||
if isinstance(x, (list, tuple)):
|
||||
if len(x) != len(x_hat):
|
||||
raise AssertionError(f"Expected len(x) == len(x_hat), got {len(x)} != {len(x_hat)}")
|
||||
res = [_snr(x[idx], x_hat[idx]) for idx in range(len(x))]
|
||||
return res
|
||||
if x_hat.is_quantized:
|
||||
x_hat = x_hat.dequantize()
|
||||
if x.is_quantized:
|
||||
x = x.dequantize()
|
||||
noise = (x - x_hat).norm()
|
||||
if noise == 0:
|
||||
return 0.0, float('inf'), float('inf')
|
||||
signal = x.norm()
|
||||
snr = signal / noise
|
||||
snr_db = 20 * snr.log10()
|
||||
return signal, noise, snr_db
|
||||
|
||||
@contextmanager
|
||||
def override_quantized_engine(qengine):
|
||||
previous = torch.backends.quantized.engine
|
||||
torch.backends.quantized.engine = qengine
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.backends.quantized.engine = previous
|
||||
|
||||
@contextmanager
|
||||
def override_cpu_allocator_for_qnnpack(qengine_is_qnnpack):
|
||||
try:
|
||||
if qengine_is_qnnpack:
|
||||
torch._C._set_default_mobile_cpu_allocator()
|
||||
yield
|
||||
finally:
|
||||
if qengine_is_qnnpack:
|
||||
torch._C._unset_default_mobile_cpu_allocator()
|
||||
|
||||
# TODO: Update all quantization tests to use this decorator.
|
||||
# Currently for some of the tests it seems to have inconsistent params
|
||||
# for fbgemm vs qnnpack.
|
||||
def override_qengines(qfunction):
|
||||
def test_fn(*args, **kwargs):
|
||||
for qengine in supported_qengines:
|
||||
with override_quantized_engine(qengine):
|
||||
# qfunction should not return anything.
|
||||
qfunction(*args, **kwargs)
|
||||
return test_fn
|
||||
|
||||
def qengine_is_fbgemm():
|
||||
return torch.backends.quantized.engine == 'fbgemm'
|
||||
def qengine_is_qnnpack():
|
||||
return torch.backends.quantized.engine == 'qnnpack'
|
||||
def qengine_is_onednn():
|
||||
return torch.backends.quantized.engine == 'onednn'
|
||||
def qengine_is_x86():
|
||||
return torch.backends.quantized.engine == 'x86'
|
||||
|
||||
# Helper function used to simulate per-channel fake-quant against any axis
|
||||
def _permute_to_axis_zero(X, axis):
|
||||
new_axis_list = list(range(X.dim()))
|
||||
new_axis_list[axis] = 0
|
||||
new_axis_list[0] = axis
|
||||
y = X.permute(tuple(new_axis_list))
|
||||
return y, new_axis_list
|
||||
|
||||
# Reference method for fake quantize
|
||||
# Note: because scale/zero_point are left as float in the actual kernel, this mimics how fake_quant works for float16/64
|
||||
def _fake_quantize_per_channel_affine_reference(X, per_channel_scale, per_channel_zero_point, axis, quant_min, quant_max):
|
||||
dtype = X.dtype
|
||||
X, permute_axis_list = _permute_to_axis_zero(X.to(torch.float32), axis)
|
||||
res = torch.zeros_like(X)
|
||||
|
||||
for i in range(X.size()[0]):
|
||||
res[i] = (torch.clamp(torch.round(X[i] * (1.0 / per_channel_scale[i]) +
|
||||
per_channel_zero_point[i]), quant_min, quant_max) - per_channel_zero_point[i]) * per_channel_scale[i]
|
||||
|
||||
out = res.permute(tuple(permute_axis_list))
|
||||
return out.to(dtype)
|
||||
|
||||
# Reference method for the gradient of the fake quantize operator
|
||||
# Note: because scale/zero_point are left as float in the actual kernel, this mimics how fake_quant works for float16/64
|
||||
def _fake_quantize_per_channel_affine_grad_reference(dY, X, per_channel_scale, per_channel_zero_point, axis, quant_min, quant_max):
|
||||
dtype = X.dtype
|
||||
X, permute_axis_list = _permute_to_axis_zero(X.to(torch.float32), axis)
|
||||
Xq = torch.zeros_like(X)
|
||||
for i in range(X.size()[0]):
|
||||
Xq[i] = torch.round(X[i] * (1.0 / per_channel_scale[i]) + per_channel_zero_point[i])
|
||||
Xq = Xq.permute(tuple(permute_axis_list))
|
||||
mask = (Xq >= quant_min) * (Xq <= quant_max)
|
||||
res = torch.zeros_like(dY)
|
||||
res[mask] = dY[mask]
|
||||
return res.to(dtype)
|
||||
|
||||
def to_tensor(X, device):
|
||||
if not isinstance(X, torch.Tensor):
|
||||
X = torch.tensor(X)
|
||||
else:
|
||||
X = X.detach().clone()
|
||||
return X.to(device=torch.device(device), dtype=torch.float32)
|
||||
|
||||
# copy-pasted from
|
||||
# https://github.com/pytorch/ao/blob/bc4f51da86956275da7db0da6e420c506df97820/torchao/prototype/custom_fp_utils.py#L27C1-L142C29
|
||||
def _n_ones(n: int) -> int:
|
||||
return (1 << n) - 1
|
||||
|
||||
EBITS_F32, MBITS_F32 = 8, 23
|
||||
F32_EXP_BIAS = _n_ones(EBITS_F32 - 1)
|
||||
|
||||
# copy-pasted from
|
||||
# https://github.com/pytorch/ao/blob/bc4f51da86956275da7db0da6e420c506df97820/torchao/prototype/custom_fp_utils.py#L27C1-L142C29
|
||||
def _f32_to_floatx_unpacked(x: Tensor, ebits: int, mbits: int) -> Tensor:
|
||||
"""Convert FP32 numbers to sub-byte floating point numbers with the given
|
||||
number of exponent and mantissa bits.
|
||||
|
||||
Input: torch.Tensor of dtype torch.float
|
||||
Output: torch.Tensor of dtype torch.uint8, where the bit encoding is stored
|
||||
in the least significant bits. e.g.
|
||||
fp4: bits 0-3 empty and bits 4-7 in fp4_e2m1 encoding
|
||||
fp6: bits 0-1 empty and bits 2-7 in fp6_e2m3 or fp6_e3m2 encoding
|
||||
|
||||
Note: there are no special values (NaN, inf) support in this code. Values
|
||||
outside the representable range of Floatx after rounding are clamped to the
|
||||
maximum Floatx magnitude (sign is preserved).
|
||||
|
||||
Code below is an adaptation of https://fburl.com/code/ciwofcg4
|
||||
|
||||
Background 1: last answer in https://stackoverflow.com/q/8981913
|
||||
Background 2: Computer Organization and Design, RISC-V edition, Chapter 3.5
|
||||
"""
|
||||
if x.dtype != torch.float:
|
||||
raise AssertionError(f"Expected x.dtype to be torch.float, got {x.dtype}")
|
||||
if 1 + ebits + mbits > 8:
|
||||
raise AssertionError(f"Expected 1 + ebits + mbits <= 8, got {1 + ebits + mbits}")
|
||||
|
||||
# calculate constants
|
||||
exp_bias = _n_ones(ebits - 1)
|
||||
max_int = _n_ones(ebits + mbits)
|
||||
sign_mask = 1 << (ebits + mbits)
|
||||
|
||||
# TODO document this better
|
||||
magic_adder = _n_ones(MBITS_F32 - mbits - 1)
|
||||
|
||||
# all E bits and M bits are 1s
|
||||
max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2**mbits))
|
||||
|
||||
# E bits = 1, M bits = 0
|
||||
min_normal = 2 ** (1 - exp_bias)
|
||||
|
||||
denorm_exp = (
|
||||
# exp bias conversion between formats
|
||||
(F32_EXP_BIAS - exp_bias)
|
||||
# mantissa length difference between formats
|
||||
+ (MBITS_F32 - mbits)
|
||||
# add one to encoded exponent for denormalized numbers
|
||||
+ 1
|
||||
)
|
||||
denorm_mask_int = denorm_exp << MBITS_F32
|
||||
|
||||
# reinterpret int32 as float32
|
||||
denorm_mask_float = torch.tensor(denorm_mask_int, dtype=torch.int32).view(
|
||||
torch.float32
|
||||
)
|
||||
|
||||
# save the sign
|
||||
# Note that we have torch.uint32, but some ops like cpu bit shifts
|
||||
# do not work on it. So, we stay in int32.
|
||||
x = x.view(torch.int32)
|
||||
sign = x & 0x80000000
|
||||
|
||||
# set everything to positive, will add sign back at the end
|
||||
x = x ^ sign
|
||||
|
||||
# TODO: can the branch floating point comparisons below be done without
|
||||
# converting to float? probably but need to verify
|
||||
x = x.view(torch.float)
|
||||
|
||||
# rewrite saturate/denorm/norm branches without explicit data dependent
|
||||
# control flow, to be more compiler friendly
|
||||
saturate_mask = x >= max_normal
|
||||
denormal_mask = torch.logical_and(torch.logical_not(saturate_mask), x < min_normal)
|
||||
normal_mask = torch.logical_not(torch.logical_or(saturate_mask, denormal_mask))
|
||||
|
||||
#
|
||||
# branch 1: saturate to max val - handled later in the code which combines
|
||||
# the branches
|
||||
#
|
||||
|
||||
#
|
||||
# branch 2: to conversion to denormal as well as rounding up to normal
|
||||
#
|
||||
denormal_x = x + denorm_mask_float
|
||||
denormal_x = denormal_x.view(torch.int32)
|
||||
denormal_x -= denorm_mask_int
|
||||
denormal_x = denormal_x.to(torch.uint8)
|
||||
|
||||
#
|
||||
# branch 3: stay in normal range, adjust the exponent and round
|
||||
#
|
||||
normal_x = x.view(torch.int32)
|
||||
# resulting mantissa is odd
|
||||
mant_odd = (normal_x >> (MBITS_F32 - mbits)) & 1
|
||||
# update exponent, rounding bias part 1
|
||||
val_to_add = ((exp_bias - F32_EXP_BIAS) << MBITS_F32) + magic_adder
|
||||
normal_x += val_to_add
|
||||
# rounding bias part 2
|
||||
normal_x += mant_odd
|
||||
# take the bits!
|
||||
normal_x = normal_x >> (MBITS_F32 - mbits)
|
||||
normal_x = normal_x.to(torch.uint8)
|
||||
|
||||
#
|
||||
# combine the branches
|
||||
#
|
||||
x = torch.full_like(x, max_int, dtype=torch.uint8)
|
||||
x = torch.where(denormal_mask, denormal_x, x)
|
||||
x = torch.where(normal_mask, normal_x, x)
|
||||
|
||||
# add sign back
|
||||
sign_lp = sign >> (MBITS_F32 + EBITS_F32 - mbits - ebits)
|
||||
sign_lp = sign_lp.to(torch.uint8)
|
||||
# Right shift of a negative signed integer can fill the least significant
|
||||
# bits with either 1s or 0s, depending on the implementation. Since PyTorch
|
||||
# doesn't have an uint32 dtype, we mask out these bits to get just the
|
||||
# f4 sign bit
|
||||
sign_lp = sign_lp & sign_mask
|
||||
x = x | sign_lp
|
||||
|
||||
return x.to(torch.uint8)
|
||||
|
||||
|
||||
# copy-pasted from
|
||||
# https://github.com/pytorch/ao/blob/29488018d99af7f7339f06353c6b5bbeae8a1493/torchao/prototype/custom_fp_utils.py#L147
|
||||
def _floatx_unpacked_to_f32(x: Tensor, ebits: int, mbits: int) -> Tensor:
|
||||
"""Convert sub-byte floating point numbers with the given number of exponent
|
||||
and mantissa bits to FP32.
|
||||
|
||||
Input: torch.Tensor of dtype uint8, where the bit encoding is stored
|
||||
in the least significant bits. e.g.
|
||||
fp4: bits 0-3 empty and bits 4-7 in fp4_e2m1 encoding
|
||||
fp6: bits 0-1 empty and bits 2-7 in fp6_e2m3 or fp6_e3m2 encoding
|
||||
Output: torch.Tensor of dtype fp32 with the dequantized value
|
||||
"""
|
||||
if x.dtype != torch.uint8:
|
||||
raise AssertionError(f"Expected x.dtype to be torch.uint8, got {x.dtype}")
|
||||
if 1 + ebits + mbits > 8:
|
||||
raise AssertionError(f"Expected 1 + ebits + mbits <= 8, got {1 + ebits + mbits}")
|
||||
|
||||
sign_mask = 1 << (ebits + mbits)
|
||||
exp_bias = _n_ones(ebits - 1)
|
||||
mantissa_mask = _n_ones(mbits)
|
||||
|
||||
# save the sign
|
||||
sign_lp = x & sign_mask
|
||||
|
||||
# set everything to positive, will add sign back at the end
|
||||
x_pos = x ^ sign_lp
|
||||
|
||||
#
|
||||
# 1. Calculate zero mask
|
||||
#
|
||||
zero_mask = x_pos == 0
|
||||
|
||||
#
|
||||
# 2. Calculate the denormal path mask
|
||||
#
|
||||
denormal_mask = torch.logical_and((x_pos > 0), ((x_pos >> mbits) == 0))
|
||||
|
||||
#
|
||||
# 3. Calculate the normal path
|
||||
#
|
||||
|
||||
# calculate the new exponent and shift it to bits 2:9 of the result
|
||||
exp_biased_lp = x_pos >> mbits
|
||||
exp_biased_f32 = exp_biased_lp - exp_bias + F32_EXP_BIAS
|
||||
exp_biased_f32 = exp_biased_f32.to(torch.int32) << MBITS_F32
|
||||
|
||||
# shift the mantissa to bits 10:32 of the result
|
||||
mantissa_lp_int32 = (x_pos & mantissa_mask).to(torch.int32)
|
||||
mantissa_f32 = mantissa_lp_int32 << (MBITS_F32 - mbits)
|
||||
result = exp_biased_f32 | mantissa_f32
|
||||
|
||||
#
|
||||
# 4. Add the zero and denormal casts to the already casted normal path
|
||||
#
|
||||
result[zero_mask] = 0
|
||||
|
||||
denormal_exp_biased = 1 - exp_bias + F32_EXP_BIAS
|
||||
|
||||
# fast path.
|
||||
# without this, performance for FP4_E2M1 is slower by 2x
|
||||
if mbits == 1:
|
||||
result[denormal_mask] = (denormal_exp_biased - mbits) << MBITS_F32
|
||||
|
||||
else:
|
||||
# iterate over all possible values of mantissa
|
||||
# i=0, j=1
|
||||
# i=1, j=10,11
|
||||
# i=2, j=100,101,110,111
|
||||
# and so on
|
||||
for i in range(mbits):
|
||||
for mantissa_cmp in range(1 << i, 1 << (i + 1)):
|
||||
# left shift mantissa until it overflows (create an implicit 1)
|
||||
# subtract exponent by the same amount
|
||||
left_shift = mbits - i
|
||||
mantissa_f32 = (mantissa_cmp - (1 << i)) << (
|
||||
left_shift + MBITS_F32 - mbits
|
||||
)
|
||||
exp_biased_f32 = (denormal_exp_biased - left_shift) << MBITS_F32
|
||||
|
||||
# we can update this in-place since the values won't overlap
|
||||
# torch.compile() may complain unsupported operand type(s) for |: 'SymInt' and 'int'
|
||||
# thus we use + instead of | here
|
||||
mantissa_lp_int32[mantissa_lp_int32 == mantissa_cmp] = (
|
||||
exp_biased_f32 + mantissa_f32
|
||||
)
|
||||
|
||||
result = torch.where(denormal_mask, mantissa_lp_int32, result)
|
||||
|
||||
# add sign back
|
||||
sign_f32 = sign_lp.to(torch.int32) << (MBITS_F32 - mbits + EBITS_F32 - ebits)
|
||||
result = result | sign_f32
|
||||
|
||||
return result.view(torch.float)
|
||||
|
||||
# copied from https://github.com/drisspg/transformer_nuggets/blob/main/transformer_nuggets/mx/to_blocked.py
|
||||
def ceil_div(a, b):
|
||||
return (a + b - 1) // b
|
||||
|
||||
# NVIDIA Blackwell HW requires scales for MX/NV blocked formats to be in a 128x4 tile layout,
|
||||
# with a weird 32x4x4 internal layout of that tile. If we want to take swizzled scales and use them
|
||||
# for non-gemm purposes (like testing), we need to de-swizzle them, then they can be applied much
|
||||
# more naturally.
|
||||
def from_blocked(input, input_scales, blocksize) -> torch.Tensor:
|
||||
# Matrix is in a 128x4 pattern, internally blocked as 32x4x4 nonsense.
|
||||
# Output should be [input.size(0, input.size(1) // blocksize] scales
|
||||
output_scales = torch.zeros(
|
||||
(input.size(0), input.size(1) // blocksize),
|
||||
device=input.device,
|
||||
dtype=input_scales.dtype,
|
||||
)
|
||||
|
||||
# Swizzled scales are padded to tiles of 128x4, we need to replicate how that padding
|
||||
# happened for offset purposes.
|
||||
# There are K//blocksize scales, padded to groups of 4.
|
||||
num_col_tiles = ceil_div(ceil_div(input.size(1), blocksize), 4)
|
||||
|
||||
# (Very) slow reference implementation using horrifying loops.
|
||||
for i in range(input.size(0)):
|
||||
for j in range(input.size(1) // blocksize):
|
||||
# which 128x4 tile of scaling factors am I in
|
||||
scale_tile_h = i // 128
|
||||
scale_tile_w = j // 4
|
||||
|
||||
# There are (padded) input_scales.size(1) // 4 tiles along the w dim.
|
||||
# So offset is 512 * (h_tile * tiles_per_row + tile_in_row)
|
||||
tile_offset = 512 * (scale_tile_h * num_col_tiles + scale_tile_w)
|
||||
|
||||
# indices within the tile - use nomenclature directly from cublas docs
|
||||
outer = i % 128 # "outer" in cublas docs
|
||||
inner = j % 4 # "inner" in cublas docs
|
||||
|
||||
# Note: "offset" is given in terms of bytes, in cublas docs, but our scales are e8m0,
|
||||
# anyway, and so 1B == 1 value => use offset directly.
|
||||
# Formula directly from cublas docs in 3.1.4.3.2
|
||||
offset = tile_offset + (outer % 32) * 16 + (outer // 32) * 4 + inner
|
||||
|
||||
output_scales[i, j] = input_scales[offset]
|
||||
|
||||
return output_scales
|
||||
|
||||
def from_blocked_format(x_mxfp8, scales_unswizzled, blocksize=32):
|
||||
# expand scales
|
||||
scales = torch.repeat_interleave(scales_unswizzled, blocksize, dim=1)
|
||||
|
||||
# de-scale and convert
|
||||
x_f32 = x_mxfp8.to(torch.float) * scales.to(torch.float)
|
||||
return x_f32.to(torch.bfloat16)
|
||||
|
||||
def to_blocked(input_matrix) -> torch.Tensor:
|
||||
"""
|
||||
Rearrange a large matrix by breaking it into blocks and applying the rearrangement pattern.
|
||||
|
||||
See:
|
||||
https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
|
||||
|
||||
Args:
|
||||
input_matrix: Input tensor of shape (H, W)
|
||||
|
||||
Returns:
|
||||
Rearranged tensor of shape (32*ceil_div(H,128), 16*ceil_div(W,4))
|
||||
"""
|
||||
rows, cols = input_matrix.shape
|
||||
n_row_blocks = ceil_div(rows, 128)
|
||||
n_col_blocks = ceil_div(cols, 4)
|
||||
|
||||
# Calculate the padded shape
|
||||
padded_rows = n_row_blocks * 128
|
||||
padded_cols = n_col_blocks * 4
|
||||
|
||||
padded = input_matrix
|
||||
# Ideally we would use torch.nn.pad but it doesn't support float8_e8m0fnu for now
|
||||
if (rows, cols) != (padded_rows, padded_cols):
|
||||
padded = torch.zeros((padded_rows, padded_cols), device=input_matrix.device, dtype=input_matrix.dtype)
|
||||
padded[:rows, :cols] = input_matrix
|
||||
|
||||
# Rearrange the blocks
|
||||
blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
|
||||
rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
|
||||
|
||||
return rearranged.flatten()
|
||||
|
||||
|
||||
def down_size(size):
|
||||
if size[-1] % 2 != 0:
|
||||
raise AssertionError(f"{size} last dim not divisible by two")
|
||||
return (*size[:-1], size[-1] // 2)
|
||||
|
||||
|
||||
def pack_uint4(uint8_data) -> torch.Tensor:
|
||||
# converting to uint8 for operations
|
||||
shape = uint8_data.shape
|
||||
if shape[-1] % 2 != 0:
|
||||
raise AssertionError(f"Expected shape[-1] to be divisible by 2, got {shape[-1]}")
|
||||
uint8_data = uint8_data.contiguous().view(-1)
|
||||
return (uint8_data[1::2] << 4 | uint8_data[::2]).view(down_size(shape))
|
||||
|
||||
|
||||
# exponent and mantissa bits of `torch.float4_e2m1fn_x2`
|
||||
FP4_EBITS, FP4_MBITS = 2, 1
|
||||
|
||||
|
||||
def _bfloat16_to_float4_e2m1fn_x2(x):
|
||||
if x.dtype != torch.bfloat16:
|
||||
raise AssertionError(f"Expected x.dtype to be torch.bfloat16, got {x.dtype}")
|
||||
x = _f32_to_floatx_unpacked(x.float(), FP4_EBITS, FP4_MBITS)
|
||||
x = pack_uint4(x)
|
||||
x = x.view(torch.float4_e2m1fn_x2)
|
||||
return x
|
||||
|
||||
|
||||
# This function is extracted from https://github.com/pytorch/ao/blob/v0.12.0/torchao/prototype/mx_formats/mx_tensor.py#L142
|
||||
def to_mxfp(
|
||||
data_hp: torch.Tensor,
|
||||
block_size: int = 32,
|
||||
format: str = "mxfp8",
|
||||
):
|
||||
if data_hp.dtype not in (torch.bfloat16, torch.float):
|
||||
raise AssertionError(f"{data_hp.dtype} is not supported yet")
|
||||
if data_hp.shape[-1] % block_size != 0:
|
||||
raise AssertionError(
|
||||
f"the last dimension of shape {data_hp.shape} must be divisible by block_size {block_size}"
|
||||
)
|
||||
if not data_hp.is_contiguous():
|
||||
raise AssertionError("unsupported: data_hp must be contiguous")
|
||||
|
||||
orig_shape = data_hp.shape
|
||||
data_hp = data_hp.reshape(
|
||||
*orig_shape[:-1], orig_shape[-1] // block_size, block_size
|
||||
)
|
||||
|
||||
max_abs = torch.amax(torch.abs(data_hp), -1).unsqueeze(-1)
|
||||
|
||||
data_hp = data_hp.to(torch.float32)
|
||||
max_abs = max_abs.to(torch.float32)
|
||||
|
||||
if format == "mxfp8":
|
||||
F8E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max # 448.0
|
||||
max_pos = F8E4M3_MAX
|
||||
elif format == "mxfp4":
|
||||
F4E2M1_MAX = 6.
|
||||
max_pos = F4E2M1_MAX
|
||||
|
||||
# RCEIL
|
||||
def _to_mx_rceil(
|
||||
data_hp: torch.Tensor,
|
||||
max_abs: torch.Tensor,
|
||||
max_pos: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
E8M0_EXPONENT_BIAS = 127
|
||||
descale = max_abs / max_pos
|
||||
exponent = torch.where(
|
||||
torch.isnan(descale),
|
||||
0xFF, # Handle biased exponent for nan
|
||||
# NOTE: descale < (torch.finfo(torch.float32).smallest_normal / 2) is handled through clamping
|
||||
(
|
||||
torch.clamp(
|
||||
torch.ceil(torch.log2(descale)),
|
||||
min=-E8M0_EXPONENT_BIAS,
|
||||
max=E8M0_EXPONENT_BIAS,
|
||||
)
|
||||
+ E8M0_EXPONENT_BIAS
|
||||
).to(torch.uint8),
|
||||
)
|
||||
|
||||
descale_fp = torch.where(
|
||||
exponent == 0,
|
||||
1.0,
|
||||
torch.exp2(E8M0_EXPONENT_BIAS - exponent.to(torch.float32)),
|
||||
)
|
||||
|
||||
# scale and saturated cast the data elements to max of target dtype
|
||||
data_lp = torch.clamp(data_hp * descale_fp, min=-1 * max_pos, max=max_pos)
|
||||
return exponent, data_lp
|
||||
|
||||
scale_e8m0_biased, data_lp = _to_mx_rceil(data_hp, max_abs, max_pos)
|
||||
|
||||
# cast to target dtype
|
||||
if format == "mxfp8":
|
||||
data_lp = data_lp.to(torch.float8_e4m3fn)
|
||||
# need to reshape at the end to help inductor fuse things
|
||||
data_lp = data_lp.reshape(orig_shape)
|
||||
elif format == "mxfp4":
|
||||
data_lp = _bfloat16_to_float4_e2m1fn_x2(data_lp.to(torch.bfloat16))
|
||||
final_shape = list(orig_shape)
|
||||
final_shape[-1] //= 2
|
||||
data_lp = data_lp.reshape(final_shape)
|
||||
|
||||
scale_e8m0_biased = scale_e8m0_biased.view(torch.float8_e8m0fnu)
|
||||
scale_e8m0_biased = scale_e8m0_biased.squeeze(-1)
|
||||
return scale_e8m0_biased, data_lp
|
||||
|
||||
# Source: https://github.com/pytorch/ao/blob/568c1932a16ae9f30d48da214a88dc0013e98ed8/torchao/prototype/moe_training/utils.py#L310
|
||||
def generate_jagged_offs(E, M, multiple_of=16, dtype=torch.int32, device="cuda"):
|
||||
"""
|
||||
Utility function for tests and benchmarks.
|
||||
|
||||
Generates a tensor of length E, containing random values divisible by `multiple_of`,
|
||||
from 0 to M, in sorted order, and where the final value in the tensor is always M.
|
||||
Args:
|
||||
E (int): The length of the tensor.
|
||||
M (int): The maximum value in the tensor.
|
||||
Returns:
|
||||
torch.Tensor: A tensor of length E with the specified properties.
|
||||
"""
|
||||
import random
|
||||
|
||||
# Ensure M is divisible by 16
|
||||
if M % multiple_of != 0:
|
||||
raise ValueError(f"M must be divisible by {multiple_of}")
|
||||
|
||||
# Generate a list of possible values
|
||||
possible_values = list(range(multiple_of, M + 1, multiple_of))
|
||||
|
||||
# If E is larger than the number of possible values, raise an error
|
||||
if E > len(possible_values):
|
||||
raise ValueError("E cannot be larger than the number of possible values")
|
||||
|
||||
# Randomly select E - 1 values from the possible values (excluding M)
|
||||
selected_values = torch.tensor(random.sample(possible_values[:-1], E - 1))
|
||||
|
||||
# Append M to the selected values
|
||||
selected_values = torch.cat((selected_values, torch.tensor([M])))
|
||||
|
||||
# Sort the selected values
|
||||
selected_values, _ = torch.sort(selected_values)
|
||||
|
||||
return selected_values.to(dtype).to(device)
|
||||
@@ -0,0 +1,345 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch
|
||||
from copy import deepcopy
|
||||
from torch.utils._pytree import tree_map
|
||||
import torch.utils._pytree as pytree
|
||||
|
||||
|
||||
# TODO: Move LoggingTensor here.
|
||||
from torch.testing._internal.logging_tensor import LoggingTensor
|
||||
|
||||
|
||||
# Base class for wrapper-style tensors.
|
||||
class WrapperTensor(torch.Tensor):
|
||||
@staticmethod
|
||||
def __new__(cls, *args, **kwargs):
|
||||
t, kwargs = cls.get_wrapper_properties(*args, **kwargs)
|
||||
if "size" not in kwargs:
|
||||
size = t.size()
|
||||
else:
|
||||
size = kwargs["size"]
|
||||
del kwargs["size"]
|
||||
if "dtype" not in kwargs:
|
||||
kwargs["dtype"] = t.dtype
|
||||
if "layout" not in kwargs:
|
||||
kwargs["layout"] = t.layout
|
||||
if "device" not in kwargs:
|
||||
kwargs["device"] = t.device
|
||||
if "requires_grad" not in kwargs:
|
||||
kwargs["requires_grad"] = False
|
||||
# Ignore memory_format and pin memory for now as I don't know how to
|
||||
# safely access them on a Tensor (if possible??)
|
||||
|
||||
wrapper = torch.Tensor._make_wrapper_subclass(cls, size, **kwargs)
|
||||
wrapper._validate_methods()
|
||||
return wrapper
|
||||
|
||||
@classmethod
|
||||
def get_wrapper_properties(cls, *args, **kwargs):
|
||||
# Should return both an example Tensor and a dictionary of kwargs
|
||||
# to override any of that example Tensor's properly.
|
||||
# This is very similar to the `t.new_*(args)` API
|
||||
raise NotImplementedError("You need to implement get_wrapper_properties")
|
||||
|
||||
def _validate_methods(self):
|
||||
# Skip this if not in debug mode?
|
||||
# Changing these on the python side is wrong as it would not be properly reflected
|
||||
# on the c++ side
|
||||
# This doesn't catch attributes set in the __init__
|
||||
forbidden_overrides = ["size", "stride", "dtype", "layout", "device", "requires_grad"]
|
||||
for el in forbidden_overrides:
|
||||
if getattr(self.__class__, el) is not getattr(torch.Tensor, el):
|
||||
raise RuntimeError(f"Subclass {self.__class__.__name__} is overwriting the "
|
||||
f"property {el} but this is not allowed as such change would "
|
||||
"not be reflected to c++ callers.")
|
||||
|
||||
|
||||
class WrapperTensorWithCustomSizes(WrapperTensor):
|
||||
@classmethod
|
||||
def get_wrapper_properties(cls, t, requires_grad=False):
|
||||
return t, {"requires_grad": requires_grad, "dispatch_sizes_strides_policy": "sizes"}
|
||||
|
||||
def __init__(self, t, requires_grad=False):
|
||||
self.t = t
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
|
||||
if not all(issubclass(cls, t) for t in types):
|
||||
return NotImplemented
|
||||
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
def unwrap(e):
|
||||
return e.t if isinstance(e, WrapperTensorWithCustomSizes) else e
|
||||
|
||||
def wrap(e):
|
||||
return WrapperTensorWithCustomSizes(e) if isinstance(e, torch.Tensor) else e
|
||||
|
||||
rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {})))
|
||||
return rs
|
||||
|
||||
def __repr__(self):
|
||||
return super().__repr__(tensor_contents=f"t={self.t}")
|
||||
|
||||
|
||||
class WrapperTensorWithCustomStrides(WrapperTensor):
|
||||
@classmethod
|
||||
def get_wrapper_properties(cls, t, requires_grad=False):
|
||||
return t, {"requires_grad": requires_grad, "dispatch_sizes_strides_policy": "strides"}
|
||||
|
||||
def __init__(self, t, requires_grad=False):
|
||||
self.t = t
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
|
||||
if not all(issubclass(cls, t) for t in types):
|
||||
return NotImplemented
|
||||
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
def unwrap(e):
|
||||
return e.t if isinstance(e, WrapperTensorWithCustomStrides) else e
|
||||
|
||||
def wrap(e):
|
||||
return WrapperTensorWithCustomStrides(e) if isinstance(e, torch.Tensor) else e
|
||||
|
||||
rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {})))
|
||||
return rs
|
||||
|
||||
def __repr__(self):
|
||||
return super().__repr__(tensor_contents=f"t={self.t}")
|
||||
|
||||
|
||||
class DiagTensorBelow(WrapperTensor):
|
||||
@classmethod
|
||||
def get_wrapper_properties(cls, diag, requires_grad=False):
|
||||
if diag.ndim != 1:
|
||||
raise AssertionError(f"Expected diag.ndim == 1, got {diag.ndim}")
|
||||
return diag, {"size": diag.size() + diag.size(), "requires_grad": requires_grad}
|
||||
|
||||
def __init__(self, diag, requires_grad=False):
|
||||
self.diag = diag
|
||||
|
||||
handled_ops = {}
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
|
||||
if not all(issubclass(cls, t) for t in types):
|
||||
return NotImplemented
|
||||
|
||||
# For everything else, call the handler:
|
||||
fn = cls.handled_ops.get(func.__name__, None)
|
||||
if fn:
|
||||
return fn(*args, **(kwargs or {}))
|
||||
else:
|
||||
# Note that here, because we don't need to provide the autograd formulas
|
||||
# we can have a default "fallback" that creates a plain Tensor based
|
||||
# on the diag elements and calls the func again.
|
||||
|
||||
def unwrap(e):
|
||||
return e.diag.diag() if isinstance(e, DiagTensorBelow) else e
|
||||
|
||||
def wrap(e):
|
||||
if isinstance(e, torch.Tensor) and e.ndim == 1:
|
||||
return DiagTensorBelow(e)
|
||||
if isinstance(e, torch.Tensor) and e.ndim == 2 and e.count_nonzero() == e.diag().count_nonzero():
|
||||
return DiagTensorBelow(e.diag())
|
||||
return e
|
||||
|
||||
rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {})))
|
||||
return rs
|
||||
|
||||
def __repr__(self):
|
||||
return super().__repr__(tensor_contents=f"diag={self.diag}")
|
||||
|
||||
|
||||
class SparseTensor(WrapperTensor):
|
||||
@classmethod
|
||||
def get_wrapper_properties(cls, size, values, indices, requires_grad=False):
|
||||
if values.device != indices.device:
|
||||
raise AssertionError(f"Expected values.device == indices.device, got {values.device} != {indices.device}")
|
||||
return values, {"size": size, "requires_grad": requires_grad}
|
||||
|
||||
def __init__(self, size, values, indices, requires_grad=False):
|
||||
self.values = values
|
||||
self.indices = indices
|
||||
|
||||
def __repr__(self):
|
||||
return super().__repr__(tensor_contents=f"values={self.values}, indices={self.indices}")
|
||||
|
||||
def sparse_to_dense(self):
|
||||
res = torch.zeros(self.size(), dtype=self.values.dtype)
|
||||
res[self.indices.unbind(1)] = self.values
|
||||
return res
|
||||
|
||||
@staticmethod
|
||||
def from_dense(t):
|
||||
indices = t.nonzero()
|
||||
values = t[indices.unbind(1)]
|
||||
return SparseTensor(t.size(), values, indices)
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
|
||||
func_name = f"{func.__module__}.{func.__name__}"
|
||||
|
||||
res = cls._try_call_special_impl(func_name, args, kwargs)
|
||||
if res is not NotImplemented:
|
||||
return res
|
||||
|
||||
# Otherwise, use a default implementation that construct dense
|
||||
# tensors and use that to compute values
|
||||
def unwrap(e):
|
||||
return e.sparse_to_dense() if isinstance(e, SparseTensor) else e
|
||||
|
||||
# Wrap back all Tensors into our custom class
|
||||
def wrap(e):
|
||||
# Check for zeros and use that to get indices
|
||||
return SparseTensor.from_dense(e) if isinstance(e, torch.Tensor) else e
|
||||
|
||||
rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {})))
|
||||
return rs
|
||||
|
||||
|
||||
_SPECIAL_IMPLS = {}
|
||||
|
||||
@classmethod
|
||||
def _try_call_special_impl(cls, func, args, kwargs):
|
||||
if func not in cls._SPECIAL_IMPLS:
|
||||
return NotImplemented
|
||||
return cls._SPECIAL_IMPLS[func](args, kwargs)
|
||||
|
||||
|
||||
# Example non-wrapper subclass that stores extra state.
|
||||
class NonWrapperTensor(torch.Tensor):
|
||||
def __new__(cls, data):
|
||||
t = torch.Tensor._make_subclass(cls, data)
|
||||
t.extra_state = {
|
||||
'last_func_called': None
|
||||
}
|
||||
return t
|
||||
|
||||
@classmethod
|
||||
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
||||
result = super().__torch_function__(func, types, args, kwargs)
|
||||
|
||||
if isinstance(result, cls):
|
||||
# Do something with the extra state. For the example here, just store the name of the
|
||||
# last function called (skip for deepcopy so the copy has the same extra state).
|
||||
if func is torch.Tensor.__deepcopy__:
|
||||
result.extra_state = deepcopy(args[0].extra_state)
|
||||
else:
|
||||
result.extra_state = {
|
||||
'last_func_called': func.__name__,
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
# new_empty() must be defined for deepcopy to work
|
||||
def new_empty(self, shape):
|
||||
return type(self)(torch.empty(shape))
|
||||
|
||||
|
||||
# Class used to store info about subclass tensors used in testing.
|
||||
class SubclassInfo:
|
||||
|
||||
__slots__ = ['name', 'create_fn', 'closed_under_ops']
|
||||
|
||||
def __init__(self, name, create_fn, closed_under_ops=True):
|
||||
self.name = name
|
||||
self.create_fn = create_fn # create_fn(shape) -> tensor instance
|
||||
self.closed_under_ops = closed_under_ops
|
||||
|
||||
|
||||
# Helper function to create a subclass of the given class and possibly cache sizes / strides.
|
||||
def _create_and_access_shape(cls, shape):
|
||||
sub = cls(torch.randn(shape))
|
||||
# NB: Wrapper subclasses with custom dispatched sizes / strides cache this info
|
||||
# on the first call via non-serializable PyCapsules. We purposefully trigger cache
|
||||
# population here for serialization / deepcopy tests to verify that the presence of this
|
||||
# cache info doesn't cause problems.
|
||||
sub.size()
|
||||
sub.stride()
|
||||
return sub
|
||||
|
||||
|
||||
subclass_db = {
|
||||
torch.Tensor: SubclassInfo(
|
||||
'base_tensor', create_fn=torch.randn
|
||||
),
|
||||
NonWrapperTensor: SubclassInfo(
|
||||
'non_wrapper_tensor',
|
||||
create_fn=lambda shape: NonWrapperTensor(torch.randn(shape))
|
||||
),
|
||||
LoggingTensor: SubclassInfo(
|
||||
'logging_tensor',
|
||||
create_fn=lambda shape: LoggingTensor(torch.randn(shape))
|
||||
),
|
||||
SparseTensor: SubclassInfo(
|
||||
'sparse_tensor',
|
||||
create_fn=lambda shape: SparseTensor.from_dense(torch.randn(shape).relu())
|
||||
),
|
||||
DiagTensorBelow: SubclassInfo(
|
||||
'diag_tensor_below',
|
||||
create_fn=lambda shape: DiagTensorBelow(torch.randn(shape)),
|
||||
closed_under_ops=False # sparse semantics
|
||||
),
|
||||
WrapperTensorWithCustomSizes: SubclassInfo(
|
||||
'wrapper_with_custom_sizes',
|
||||
create_fn=lambda shape: _create_and_access_shape(WrapperTensorWithCustomSizes, shape),
|
||||
closed_under_ops=False,
|
||||
),
|
||||
WrapperTensorWithCustomStrides: SubclassInfo(
|
||||
'wrapper_with_custom_strides',
|
||||
create_fn=lambda shape: _create_and_access_shape(WrapperTensorWithCustomStrides, shape),
|
||||
closed_under_ops=False,
|
||||
),
|
||||
}
|
||||
|
||||
class SubclassWithTensorFactory(torch.Tensor):
|
||||
@staticmethod
|
||||
def __new__(cls, src):
|
||||
shape = src.shape
|
||||
kwargs = {}
|
||||
kwargs["strides"] = src.stride()
|
||||
kwargs["storage_offset"] = src.storage_offset()
|
||||
kwargs["device"] = src.device
|
||||
kwargs["layout"] = src.layout
|
||||
kwargs["requires_grad"] = src.requires_grad
|
||||
kwargs["dtype"] = src.dtype
|
||||
out = torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)
|
||||
return out
|
||||
|
||||
def __init__(self, src):
|
||||
self.src = src
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.__class__.__name__}"
|
||||
|
||||
def __tensor_flatten__(self):
|
||||
return ["src"], None
|
||||
|
||||
@classmethod
|
||||
def __tensor_unflatten__(cls, inner_tensors, meta, outer_size, outer_stride):
|
||||
src = inner_tensors["src"]
|
||||
return cls(src)
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args, kwargs):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
def _fn(x):
|
||||
return x.src * torch.ones(x.src.shape) if x.src.dtype == torch.float32 else x.src
|
||||
|
||||
_args = pytree.tree_map_only(cls, _fn, args)
|
||||
_kwargs = pytree.tree_map_only(cls, _fn, kwargs)
|
||||
|
||||
_out = func(*_args, **_kwargs)
|
||||
|
||||
_out_flat, _out_spec = pytree.tree_flatten(_out)
|
||||
|
||||
out_flat = [cls(o) if isinstance(o, torch.Tensor) else o for o in _out_flat]
|
||||
return pytree.tree_unflatten(out_flat, _out_spec)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,88 @@
|
||||
import enum
|
||||
import functools
|
||||
|
||||
import torch
|
||||
import torch.xpu
|
||||
from torch.testing._internal.common_utils import IS_WINDOWS, LazyVal, TEST_XPU
|
||||
|
||||
|
||||
XPU_ALREADY_INITIALIZED_ON_IMPORT = torch.xpu.is_initialized()
|
||||
|
||||
|
||||
class XPUCodename(enum.Enum):
|
||||
PVC = "PVC" # Intel® Data Center GPU Max Series
|
||||
BMG = "BMG" # Intel® Arc™ Pro Battlemage Graphics
|
||||
|
||||
|
||||
class XPUArch(enum.IntEnum):
|
||||
Unknown = 0
|
||||
Xe = 1 # Xe HPC
|
||||
Xe2 = 2
|
||||
|
||||
|
||||
# device_id -> GPU codename
|
||||
# From https://github.com/intel/intel-graphics-compiler/blob/master/inc/common/igfxfmid.h
|
||||
_DEVICE_ID_TO_CODENAME = {
|
||||
0x0BD0: XPUCodename.PVC,
|
||||
0x0BD4: XPUCodename.PVC,
|
||||
0x0BD5: XPUCodename.PVC,
|
||||
0x0BD6: XPUCodename.PVC,
|
||||
0x0BD7: XPUCodename.PVC,
|
||||
0x0BD8: XPUCodename.PVC,
|
||||
0x0BD9: XPUCodename.PVC,
|
||||
0x0BDA: XPUCodename.PVC,
|
||||
0x0BDB: XPUCodename.PVC,
|
||||
0x0B69: XPUCodename.PVC,
|
||||
0x0B6E: XPUCodename.PVC,
|
||||
0xE202: XPUCodename.BMG,
|
||||
0xE20B: XPUCodename.BMG,
|
||||
0xE20C: XPUCodename.BMG,
|
||||
0xE20D: XPUCodename.BMG,
|
||||
0xE210: XPUCodename.BMG,
|
||||
0xE212: XPUCodename.BMG,
|
||||
0xE215: XPUCodename.BMG,
|
||||
0xE216: XPUCodename.BMG,
|
||||
0xE220: XPUCodename.BMG,
|
||||
0xE221: XPUCodename.BMG,
|
||||
0xE222: XPUCodename.BMG,
|
||||
0xE223: XPUCodename.BMG,
|
||||
}
|
||||
|
||||
# GPU codename -> architecture
|
||||
_CODENAME_TO_ARCH = {
|
||||
XPUCodename.PVC: XPUArch.Xe,
|
||||
XPUCodename.BMG: XPUArch.Xe2,
|
||||
}
|
||||
|
||||
|
||||
@functools.lru_cache(1)
|
||||
def get_xpu_codename() -> XPUCodename | None:
|
||||
device_id = torch.xpu.get_device_capability()["device_id"]
|
||||
return _DEVICE_ID_TO_CODENAME.get(device_id)
|
||||
|
||||
|
||||
@functools.lru_cache(1)
|
||||
def get_xpu_arch() -> XPUArch | None:
|
||||
codename = get_xpu_codename()
|
||||
return _CODENAME_TO_ARCH.get(codename, XPUArch.Unknown)
|
||||
|
||||
|
||||
Xe2_Or_Later = LazyVal(
|
||||
lambda: torch.xpu.is_available() and get_xpu_arch() >= XPUArch.Xe2
|
||||
)
|
||||
|
||||
|
||||
def evaluate_platform_supports_flash_attention():
|
||||
if TEST_XPU:
|
||||
return not IS_WINDOWS and Xe2_Or_Later
|
||||
return False
|
||||
|
||||
|
||||
PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU: bool = LazyVal(
|
||||
lambda: evaluate_platform_supports_flash_attention()
|
||||
)
|
||||
|
||||
# Importing this module should NOT eagerly initialize XPU
|
||||
if not XPU_ALREADY_INITIALIZED_ON_IMPORT:
|
||||
if torch.xpu.is_initialized():
|
||||
raise AssertionError("XPU should not be initialized on import")
|
||||
+616
@@ -0,0 +1,616 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
import itertools
|
||||
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
|
||||
from torch.utils import _pytree as pytree
|
||||
from functools import partial
|
||||
from torch.utils._mode_utils import no_dispatch, all_same_mode
|
||||
import torch.autograd.forward_ad as fwAD
|
||||
from collections.abc import Callable
|
||||
import re
|
||||
|
||||
|
||||
def check_attr_consistency(wrapper_tensor, metadata_name, metadata_accessor):
|
||||
elem = wrapper_tensor.elem
|
||||
metadata_wrapper_tensor = metadata_accessor(wrapper_tensor)
|
||||
metadata_elem = metadata_accessor(elem)
|
||||
if metadata_wrapper_tensor == metadata_elem:
|
||||
return
|
||||
raise RuntimeError(
|
||||
f"This operator is not Composite Compliant: the "
|
||||
f"{metadata_name} of the tensor was modified directly without "
|
||||
f"going through the PyTorch dispatcher.")
|
||||
|
||||
def check_metadata_consistency(wrapper_tensor, CCT):
|
||||
# CCT: CompositeCompliantTensor class which is generated using generate_cct
|
||||
if not isinstance(wrapper_tensor, CCT):
|
||||
return
|
||||
things_to_check = {
|
||||
'shape': Tensor.size,
|
||||
'dtype': lambda x: x.dtype,
|
||||
'device': lambda x: x.device,
|
||||
'numel': Tensor.numel,
|
||||
'stride': Tensor.stride,
|
||||
'storage_offset': Tensor.storage_offset,
|
||||
}
|
||||
for metadata_name, metadata_accessor in things_to_check.items():
|
||||
check_attr_consistency(wrapper_tensor, metadata_name, metadata_accessor)
|
||||
|
||||
def is_view_fn(func):
|
||||
return func.overloadpacket.__name__ in {
|
||||
'as_strided',
|
||||
'detach',
|
||||
'diagonal',
|
||||
'expand',
|
||||
'expand_as',
|
||||
'movedim',
|
||||
'narrow',
|
||||
'permute',
|
||||
'select',
|
||||
'squeeze',
|
||||
'transpose',
|
||||
't',
|
||||
'real',
|
||||
'imag',
|
||||
'view_as_real',
|
||||
'view_as_complex',
|
||||
'unflatten',
|
||||
'unfold',
|
||||
'unsqueeze',
|
||||
'view',
|
||||
'view_as',
|
||||
'unbind',
|
||||
'split',
|
||||
'split_with_sizes',
|
||||
'vsplit',
|
||||
'hsplit',
|
||||
'tensor_split',
|
||||
'chunk',
|
||||
'swapaxes',
|
||||
'slice',
|
||||
'_reshape_alias',
|
||||
'_unsafe_view',
|
||||
'_conj',
|
||||
'alias',
|
||||
}
|
||||
|
||||
# manually populated from native_functions that have inplace_view: True.
|
||||
# In the future we will probably be able to grab that list directly
|
||||
def is_inplace_view_fn(func):
|
||||
return func.overloadpacket.__name__ in {
|
||||
'as_strided_',
|
||||
'detach_',
|
||||
'squeeze_',
|
||||
'swapaxes_',
|
||||
'swapdims_',
|
||||
't_',
|
||||
'transpose_',
|
||||
'unsqueeze_',
|
||||
}
|
||||
|
||||
|
||||
# Introspection please save us
|
||||
def is_inplace(func):
|
||||
name = func.overloadpacket.__name__
|
||||
if re.match('__i.+__', name):
|
||||
return True
|
||||
if re.match('__.+__', name):
|
||||
return False
|
||||
return name[-1] == '_'
|
||||
|
||||
|
||||
def generate_cct_and_mode(autograd_view_consistency=True):
|
||||
# This function returns a new class CompositeCompliantTensor
|
||||
# The two arguments control the behaviour described below.
|
||||
|
||||
# autograd_view_consistency:
|
||||
# If True, alias result using `set_` if func returns a view
|
||||
# (See Note [Alias Result]).
|
||||
# Since Forward AD doesn't work with `set_`
|
||||
# we disable it by setting alias to False.
|
||||
|
||||
class CompositeCompliantTensor(torch.Tensor):
|
||||
elem: torch.Tensor
|
||||
|
||||
__slots__ = ['elem']
|
||||
|
||||
@staticmethod
|
||||
def __new__(cls, elem, mode, *args, **kwargs):
|
||||
if type(elem) is cls:
|
||||
raise AssertionError(
|
||||
"Wrapping a CompositeCompliantTensor in a CompositeCompliantTensor is not supported"
|
||||
)
|
||||
|
||||
# The storage of CompositeCompliantTensor should never be used directly
|
||||
# by a Composite operation; if the Composite
|
||||
# operator attempts to read from the storage without dispatching then it'll
|
||||
# raise a RuntimeError due to it being a meta storage.
|
||||
r = torch.Tensor._make_wrapper_subclass(
|
||||
cls, elem.size(),
|
||||
dtype=elem.dtype, layout=elem.layout,
|
||||
device=elem.device, requires_grad=elem.requires_grad,
|
||||
strides=elem.stride(), storage_offset=elem.storage_offset())
|
||||
|
||||
if elem.requires_grad:
|
||||
# CompositeCompliantTensor steals the "requires_grad"-ness.
|
||||
# Why a new copy of `elem`? Because sometimes OpInfo shares inputs between tests...
|
||||
tmp = torch.empty(
|
||||
(),
|
||||
dtype=elem.dtype,
|
||||
device=elem.device,
|
||||
layout=elem.layout,
|
||||
requires_grad=False,
|
||||
)
|
||||
# Use set_ rather than empty_strided() + copy_ so that we can preserve
|
||||
# things like storage_offset.
|
||||
tmp.set_(
|
||||
source=elem.untyped_storage().clone(),
|
||||
storage_offset=elem.storage_offset(),
|
||||
size=elem.size(),
|
||||
stride=elem.stride(),
|
||||
)
|
||||
r.elem = tmp
|
||||
else:
|
||||
r.elem = elem
|
||||
|
||||
if r.stride() != r.elem.stride():
|
||||
raise AssertionError(f"Expected r.stride() == r.elem.stride(), got {r.stride()} != {r.elem.stride()}")
|
||||
|
||||
# Propagate conjugate bits to the wrapper tensor
|
||||
# Ref: https://github.com/albanD/subclass_zoo/issues/24
|
||||
# Ref: https://github.com/albanD/subclass_zoo/issues/21
|
||||
torch._C._set_conj(r, r.elem.is_conj())
|
||||
torch._C._set_neg(r, r.elem.is_neg())
|
||||
|
||||
r.mode = mode
|
||||
return r
|
||||
|
||||
def __repr__(self):
|
||||
return f"CompositeCompliantTensor({self.elem})"
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
|
||||
all_args = pytree.arg_tree_leaves(*args, **(kwargs or {}))
|
||||
modes = tuple(e.mode for e in all_args if isinstance(e, CompositeCompliantTensor))
|
||||
if not all_same_mode(modes):
|
||||
raise RuntimeError("Multiple CompositeCompliantTensorModes NYI")
|
||||
with modes[0]:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
class CompositeCompliantTensorMode(TorchDispatchMode):
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
def unwrap(e):
|
||||
return e.elem if isinstance(e, CompositeCompliantTensor) else e
|
||||
|
||||
def wrap(e):
|
||||
return CompositeCompliantTensor(e, self) if isinstance(e, torch.Tensor) else e
|
||||
|
||||
if func is torch.ops.aten._local_scalar_dense.default:
|
||||
raise RuntimeError(
|
||||
".item() is not allowed to be called inside of composite "
|
||||
"functions in the PyTorch library because not all backends "
|
||||
"and/or Tensor subclasses (e.g. vmap, ProxyTensor) support them.")
|
||||
|
||||
if func.overloadpacket.__name__ in ('set_', 'resize_'):
|
||||
raise RuntimeError(
|
||||
f"{func.__name__} is not allowed to be called inside of "
|
||||
f"Composite operators.")
|
||||
|
||||
if is_inplace(func):
|
||||
# NB: We are making an assumption that if the function is in-place,
|
||||
# then the first argument is being written to. Introspection please save us!
|
||||
mutated_argument = args[0]
|
||||
if not isinstance(mutated_argument, CompositeCompliantTensor) and \
|
||||
any(isinstance(a, CompositeCompliantTensor) for a in args[1:]):
|
||||
raise RuntimeError(
|
||||
'Not composite compliant: performing in-place operation '
|
||||
f'{func.__name__} where the Tensor being written to is '
|
||||
'regular Tensor but the other tensors are Tensor Subclasses. '
|
||||
'Please try to avoid this in-place operation.')
|
||||
|
||||
unwrapped_args = tree_map(unwrap, args)
|
||||
unwrapped_kwargs = tree_map(unwrap, kwargs)
|
||||
unwrapped_rs = func(*unwrapped_args, **unwrapped_kwargs)
|
||||
rs = tree_map(wrap, unwrapped_rs)
|
||||
|
||||
if is_view_fn(func) and autograd_view_consistency:
|
||||
# Note [Alias Result]
|
||||
# Autograd asserts that for B = A.view_fn(...), B and A's storages
|
||||
# are the same. Here we try to make B alias A to avoid those asserts.
|
||||
# See https://github.com/pytorch/pytorch/issues/65339 for more information
|
||||
# about the issue.
|
||||
with no_dispatch():
|
||||
# Idea: this is a weird way of getting a storage that aliases the input.
|
||||
# This is a workaround for #65339.
|
||||
# 1. under no_dispatch, all of the wrapper tensors look like regular
|
||||
# tensors with special storage (the storage is nullptr and
|
||||
# advertises CPU/CUDA device.
|
||||
# 2. we run func, which ends up running the view operation
|
||||
# 3. All view operations reuse the input's storage and return
|
||||
# result Tensor(s) with new sizes/strides/offset that alias
|
||||
# the input.
|
||||
# 4. we set the storage (and sizes/strides/offset) of the wrapper
|
||||
# tensor results to be that of the tensors that alias the input
|
||||
result = func(*args, **kwargs)
|
||||
if isinstance(result, (tuple, list)):
|
||||
for a, b in zip(rs, result, strict=True):
|
||||
a.set_(b)
|
||||
else:
|
||||
rs.set_(result)
|
||||
|
||||
# Some operations are allowed to in-place modify the metadata of the
|
||||
# inputs. The only ones are the "inplace view functions"; when we
|
||||
# run into these, we manually modify the metadata of the input.
|
||||
with no_dispatch():
|
||||
if is_inplace_view_fn(func):
|
||||
func(*args, **kwargs)
|
||||
|
||||
# For each CompositeCompliantTensor t, we check that t and t.elem
|
||||
# have consistent metadata. If they don't have consistent metadata,
|
||||
# that means the operator did something fishy.
|
||||
check = partial(check_metadata_consistency, CCT=CompositeCompliantTensor)
|
||||
pytree.tree_map_(check, args)
|
||||
pytree.tree_map_(check, kwargs)
|
||||
pytree.tree_map_(check, rs)
|
||||
return rs
|
||||
|
||||
return CompositeCompliantTensor, CompositeCompliantTensorMode()
|
||||
|
||||
def is_tensorlist(lst):
|
||||
if not isinstance(lst, list) and not isinstance(lst, tuple):
|
||||
return False
|
||||
if len(lst) == 0:
|
||||
return False
|
||||
all_tensors = all(isinstance(elt, torch.Tensor) for elt in lst)
|
||||
if all_tensors:
|
||||
return True
|
||||
exists_one_tensor = all(isinstance(elt, torch.Tensor) for elt in lst)
|
||||
if exists_one_tensor:
|
||||
raise RuntimeError('This test assumes that PyTorch APIs cannot take '
|
||||
'mixed lists of Tensor and other things')
|
||||
return False
|
||||
|
||||
|
||||
def maybe_map(fn, should_map, arg):
|
||||
return fn(arg) if should_map else arg
|
||||
|
||||
|
||||
def wrap(arg, CCT, cct_mode):
|
||||
# CCT: CompositeCompliantTensor class which is generated using generate_cct_and_mode
|
||||
if isinstance(arg, torch.Tensor):
|
||||
return CCT(arg, cct_mode)
|
||||
if is_tensorlist(arg):
|
||||
return [CCT(a, cct_mode) for a in arg]
|
||||
raise RuntimeError("wrap assumes that the input can be wrapped")
|
||||
|
||||
|
||||
# Given a list of flat arguments, some of which may be Tensors, return all
|
||||
# possible ways some of the arguments could be CompositeCompliantTensors (CCT).
|
||||
# For example, given Tensors A, B, C and flat_args = [A, 1, B],
|
||||
# We would return the following 4 options:
|
||||
# [CCT(A), 1, CCT(B)]
|
||||
# [CCT(A), 1, B]
|
||||
# [A, 1, CCT(B)]
|
||||
# [A, 1, B]
|
||||
# NB: Yes, this is exponential. No, we don't care too much because PyTorch ops
|
||||
# don't accept that many input Tensors.
|
||||
def generate_subclass_choices(flat_args, CCT, cct_mode):
|
||||
# CCT: CompositeCompliantTensor class which is generated using generate_cct_and_mode
|
||||
is_tensor_likes = [isinstance(arg, torch.Tensor) or is_tensorlist(arg) for arg in flat_args]
|
||||
subclass_options = [[False, True] if is_tensor_like else [False] for is_tensor_like in is_tensor_likes]
|
||||
|
||||
for which_args_are_wrapped in itertools.product(*subclass_options):
|
||||
|
||||
result = [maybe_map(partial(wrap, CCT=CCT, cct_mode=cct_mode), should_wrap_arg, arg)
|
||||
for should_wrap_arg, arg in zip(which_args_are_wrapped, flat_args, strict=True)]
|
||||
yield result, which_args_are_wrapped
|
||||
|
||||
|
||||
# For an operation f(*args, **kwargs), each Tensor argument may either be
|
||||
# a regular Tensor or a Tensor Subclass. This iterator iterates through
|
||||
# all of those options.
|
||||
def generate_subclass_choices_args_kwargs(args, kwargs, CCT, cct_mode):
|
||||
# CCT: CompositeCompliantTensor class which is generated using generate_cct_and_mode
|
||||
flat_kwargs, spec = tree_flatten(kwargs)
|
||||
flat_args_kwargs = list(args) + list(flat_kwargs)
|
||||
for choice, debug_metadata in generate_subclass_choices(flat_args_kwargs, CCT, cct_mode):
|
||||
new_args = choice[:len(args)]
|
||||
new_kwargs = tree_unflatten(choice[len(args):], spec)
|
||||
which_args_are_wrapped = debug_metadata[:len(args)]
|
||||
which_kwargs_are_wrapped = tree_unflatten(debug_metadata[len(args):], spec)
|
||||
yield new_args, new_kwargs, which_args_are_wrapped, which_kwargs_are_wrapped
|
||||
|
||||
|
||||
def raise_composite_compliance_error(err, additional_info=''):
|
||||
raise RuntimeError(
|
||||
"Composite compliance check failed with "
|
||||
"the above error.\n"
|
||||
f"{additional_info}"
|
||||
"If you are adding an OpInfo of an "
|
||||
"existing operator, please feel free to skip this test "
|
||||
"because the problem was pre-existing and file an issue. "
|
||||
"Otherwise, if you added a new operator, please read "
|
||||
"through the Composite Compliance section in "
|
||||
"aten/src/ATen/native/README.md for how to resolve this. "
|
||||
) from err
|
||||
|
||||
|
||||
# This test checks ALL possible permutations of calling `op` with arguments
|
||||
# that are individually either a regular Tensor or a Tensor subclass.
|
||||
#
|
||||
# The general strategy is to wrap some Tensor args and kwargs in
|
||||
# CompositeCompliantTensor wrappers and call the operation.
|
||||
|
||||
# If some composite operation does any non-compliant behavior,
|
||||
# CompositeCompliantTensor will raise an error.
|
||||
def check_all_permutations(op, args, kwargs, assert_equal_fn):
|
||||
CCT, cct_mode = generate_cct_and_mode()
|
||||
expected = op(*args, **kwargs)
|
||||
for choice in generate_subclass_choices_args_kwargs(args, kwargs, CCT, cct_mode):
|
||||
new_args, new_kwargs, which_args_are_wrapped, which_kwargs_are_wrapped = choice
|
||||
|
||||
try:
|
||||
actual = op(*new_args, **new_kwargs)
|
||||
# NOTE: [What errors are Composite Compliance trying to catch?]
|
||||
#
|
||||
# There's two things we want to catch:
|
||||
# - errors that would raise within the torch_dispatch impl
|
||||
# - data_ptr accesses
|
||||
# The first is easy to filter for (we could make the error a different
|
||||
# error class), the second is always going to be a RuntimeError due to
|
||||
# how it is implemented (if you try to access the data_ptr of the
|
||||
# wrapper Tensor, it raises you some internal RuntimeError).
|
||||
#
|
||||
# So the most general thing to catch here was RuntimeError. If you
|
||||
# are here and debugging why your test failed, it's plausible that
|
||||
# the operator itself is broken and that there are other tests failing.
|
||||
except RuntimeError as err:
|
||||
raise_composite_compliance_error(
|
||||
err,
|
||||
f"- wrapped_args: {which_args_are_wrapped}\n"
|
||||
f"- wrapped_kwargs: {which_kwargs_are_wrapped}\n"
|
||||
)
|
||||
|
||||
def unwrap(e):
|
||||
return e.elem if isinstance(e, CCT) else e
|
||||
|
||||
assert_equal_fn(tree_map(unwrap, actual), expected)
|
||||
|
||||
# Checks via the usage of torch dispatch mode certain anti-patterns that
|
||||
# are not composite compliant.
|
||||
#
|
||||
# In particular, the anti-pattern we are trying to prevent is a user
|
||||
# creating an empty tensor and then resize_-ing it. Torch Dispatch Mode helps
|
||||
# here because all factory functions will create tensors that are
|
||||
# CompositeCompliantTensor.
|
||||
#
|
||||
# The general strategy is to wrap all Tensor args and kwargs in
|
||||
# CompositeCompliantTensor wrappers. If an operator that is
|
||||
# Composite does any non-compliant behavior,
|
||||
# CompositeCompliantTensor will raise an error.
|
||||
def check_with_mode(op, args, kwargs, assert_equal_fn):
|
||||
CCT, cct_mode = generate_cct_and_mode()
|
||||
|
||||
def wrap(e):
|
||||
return CCT(e, cct_mode) if isinstance(e, torch.Tensor) else e
|
||||
|
||||
expected = op(*args, **kwargs)
|
||||
|
||||
args = tree_map(wrap, args)
|
||||
kwargs = tree_map(wrap, kwargs)
|
||||
try:
|
||||
with cct_mode:
|
||||
actual = op(*args, **kwargs)
|
||||
# see NOTE: [What errors are Composite Compliance trying to catch?]
|
||||
except RuntimeError as err:
|
||||
raise_composite_compliance_error(err)
|
||||
|
||||
def unwrap(e):
|
||||
return e.elem if isinstance(e, CCT) else e
|
||||
|
||||
assert_equal_fn(tree_map(unwrap, actual), expected)
|
||||
|
||||
def gather_leaf_tensors(args, kwargs):
|
||||
leaf_tensors = []
|
||||
args, _args_spec = tree_flatten(args)
|
||||
kwargs, _kwargs_spec = tree_flatten(kwargs)
|
||||
args = args + kwargs
|
||||
for arg in args:
|
||||
if not isinstance(arg, torch.Tensor):
|
||||
continue
|
||||
if arg.requires_grad:
|
||||
leaf_tensors.append(arg)
|
||||
return leaf_tensors
|
||||
|
||||
|
||||
def compute_expected_grads(op, args, kwargs, output_process_fn_grad=None, gradcheck_wrapper=None):
|
||||
if gradcheck_wrapper is None:
|
||||
results = op(*args, **kwargs)
|
||||
else:
|
||||
results = gradcheck_wrapper(op, *args, **kwargs)
|
||||
|
||||
if output_process_fn_grad is not None:
|
||||
results = output_process_fn_grad(results)
|
||||
|
||||
flat_results = pytree.tree_leaves(results)
|
||||
flat_results = [r for r in flat_results if isinstance(r, torch.Tensor)]
|
||||
flat_diff_results = [r for r in flat_results if r.requires_grad]
|
||||
if len(flat_diff_results) <= 0:
|
||||
raise AssertionError("Expected len(flat_diff_results) > 0")
|
||||
|
||||
grads = [torch.ones(r.shape, device=r.device, dtype=r.dtype) for r in flat_diff_results]
|
||||
leaf_tensors = gather_leaf_tensors(args, kwargs)
|
||||
if len(leaf_tensors) <= 0:
|
||||
raise AssertionError("Expected len(leaf_tensors) > 0")
|
||||
return torch.autograd.grad(flat_diff_results, leaf_tensors,
|
||||
grads, allow_unused=True, retain_graph=True)
|
||||
|
||||
|
||||
# Checks if the backward formula is composite compliant by testing
|
||||
# all possible permutations of {inputs, grad_outputs} being
|
||||
# CompositeCompliantTensor or regular Tensors.
|
||||
#
|
||||
# NB: it is important that op is accepted as a Callable and not an OpInfo,
|
||||
# this means we can apply check_backward_formula to things that aren't OpInfos
|
||||
# while debugging.
|
||||
def check_backward_formula(op: Callable, args, kwargs,
|
||||
output_process_fn_grad=None,
|
||||
gradcheck_wrapper=None, assert_equal_fn=None):
|
||||
CCT, cct_mode = generate_cct_and_mode()
|
||||
|
||||
expected = compute_expected_grads(op, args, kwargs, output_process_fn_grad, gradcheck_wrapper)
|
||||
|
||||
for choice in generate_subclass_choices_args_kwargs(args, kwargs, CCT, cct_mode):
|
||||
new_args, new_kwargs, which_args_are_wrapped, which_kwargs_are_wrapped = choice
|
||||
leaf_tensors = gather_leaf_tensors(new_args, new_kwargs)
|
||||
if len(leaf_tensors) <= 0:
|
||||
raise AssertionError("Expected len(leaf_tensors) > 0")
|
||||
|
||||
try:
|
||||
if gradcheck_wrapper is None:
|
||||
results = op(*new_args, **new_kwargs)
|
||||
else:
|
||||
results = gradcheck_wrapper(op, *new_args, **new_kwargs)
|
||||
if output_process_fn_grad is not None:
|
||||
results = output_process_fn_grad(results)
|
||||
# see NOTE: [What errors are Composite Compliance trying to catch?]
|
||||
except RuntimeError as err:
|
||||
raise_composite_compliance_error(
|
||||
err,
|
||||
f"- wrapped_args: {which_args_are_wrapped}\n"
|
||||
f"- wrapped_kwargs: {which_kwargs_are_wrapped}\n"
|
||||
)
|
||||
|
||||
flat_results = pytree.tree_leaves(results)
|
||||
flat_results = [r for r in flat_results if isinstance(r, torch.Tensor)]
|
||||
flat_diff_results = [r for r in flat_results if r.requires_grad]
|
||||
if len(flat_diff_results) <= 0:
|
||||
raise AssertionError("Expected len(flat_diff_results) > 0")
|
||||
|
||||
# NB: ones, not ones_like, so we get a regular Tensor here
|
||||
grads = [torch.ones(r.shape, device=r.device, dtype=r.dtype)
|
||||
for r in flat_diff_results]
|
||||
for flat_new_grads, which_grad_is_batched in generate_subclass_choices(grads, CCT, cct_mode):
|
||||
try:
|
||||
actual = torch.autograd.grad(flat_diff_results, leaf_tensors, flat_new_grads,
|
||||
allow_unused=True, retain_graph=True)
|
||||
# see NOTE: [What errors are Composite Compliance trying to catch?]
|
||||
except RuntimeError as err:
|
||||
raise_composite_compliance_error(
|
||||
err,
|
||||
f"- wrapped_args: {which_args_are_wrapped}\n"
|
||||
f"- wrapped_kwargs: {which_kwargs_are_wrapped}\n"
|
||||
f"- wrapped_grads: {which_grad_is_batched}\n"
|
||||
)
|
||||
|
||||
def unwrap(e):
|
||||
return e.elem if isinstance(e, CCT) else e
|
||||
|
||||
assert_equal_fn(tuple(map(unwrap, actual)), expected, equal_nan=True)
|
||||
|
||||
# Checks if the forward AD formula is composite compliant by testing
|
||||
# all possible permutations of {primals, tangents} being
|
||||
# CompositeCompliantTensor or regular Tensors.
|
||||
#
|
||||
# NB: it is important that op is accepted as a Callable and not an OpInfo,
|
||||
# this means we can apply check_forward_ad_formula to things that aren't OpInfos
|
||||
# while debugging.
|
||||
def check_forward_ad_formula(op: Callable, args, kwargs, gradcheck_wrapper=None, assert_equal_fn=None):
|
||||
CCT, cct_mode = generate_cct_and_mode(autograd_view_consistency=False)
|
||||
|
||||
def maybe_tangent(t):
|
||||
if type(t) is CCT:
|
||||
raise AssertionError("Expected type(t) is not CCT")
|
||||
# Generate `tangent` tensor
|
||||
# if given object is a Tensor and requires grad is set.
|
||||
if isinstance(t, torch.Tensor) and t.requires_grad:
|
||||
return torch.randn_like(t)
|
||||
elif is_tensorlist(t):
|
||||
return [torch.randn_like(e) if e.requires_grad else None for e in t]
|
||||
return None
|
||||
|
||||
tangent_args = tuple(maybe_tangent(arg) for arg in args)
|
||||
flat_kwargs, spec = tree_flatten(kwargs)
|
||||
flat_tangent_kwargs = tuple(maybe_tangent(arg) for arg in flat_kwargs)
|
||||
tangent_kwargs = tree_unflatten(flat_tangent_kwargs, spec)
|
||||
|
||||
with fwAD.dual_level():
|
||||
def maybe_make_dual(dual):
|
||||
# Returns dual tensor if primal is a tensor/tensor subclass
|
||||
# with requires_grad set.
|
||||
primal, tangent = dual
|
||||
if isinstance(primal, torch.Tensor) and primal.requires_grad:
|
||||
return fwAD.make_dual(primal.detach(), tangent)
|
||||
elif is_tensorlist(primal):
|
||||
return tuple(fwAD.make_dual(pri.detach(), tang) if tang is not None else pri
|
||||
for pri, tang in zip(primal, tangent, strict=True))
|
||||
return primal
|
||||
|
||||
def compute_expected_grad(args, tangent_args, kwargs, tangent_kwargs):
|
||||
op_args = tuple(map(maybe_make_dual, zip(args, tangent_args, strict=True)))
|
||||
op_kwargs = {k: maybe_make_dual((v, tangent_kwargs[k])) for k, v in kwargs.items()}
|
||||
|
||||
if gradcheck_wrapper is None:
|
||||
return op(*op_args, **op_kwargs)
|
||||
return gradcheck_wrapper(op, *op_args, **op_kwargs)
|
||||
|
||||
expected = compute_expected_grad(args, tangent_args, kwargs, tangent_kwargs)
|
||||
expected = tree_map(fwAD.unpack_dual, expected)
|
||||
expected_primals = tree_map(
|
||||
lambda x: x.primal,
|
||||
expected,
|
||||
is_leaf=lambda x: type(x) is fwAD.UnpackedDualTensor,
|
||||
)
|
||||
expected_tangents = tree_map(
|
||||
lambda x: x.tangent,
|
||||
expected,
|
||||
is_leaf=lambda x: type(x) is fwAD.UnpackedDualTensor,
|
||||
)
|
||||
|
||||
# Permutations of arg and kwargs in CCT.
|
||||
for choice in generate_subclass_choices_args_kwargs(args, kwargs, CCT, cct_mode):
|
||||
new_args, new_kwargs, which_args_are_wrapped, which_kwargs_are_wrapped = choice
|
||||
|
||||
# Permutations tangent arg and tangent kwargs in CCT.
|
||||
for tang_choice in generate_subclass_choices_args_kwargs(tangent_args, tangent_kwargs, CCT, cct_mode):
|
||||
new_tang_args, new_tang_kwargs, \
|
||||
which_tang_args_are_wrapped, which_tang_kwargs_are_wrapped = tang_choice
|
||||
|
||||
op_args = tuple(map(maybe_make_dual, zip(new_args, new_tang_args, strict=True)))
|
||||
op_kwargs = {k: maybe_make_dual((v, new_tang_kwargs[k])) for k, v in new_kwargs.items()}
|
||||
|
||||
try:
|
||||
if gradcheck_wrapper is None:
|
||||
actual = op(*op_args, **op_kwargs)
|
||||
else:
|
||||
actual = gradcheck_wrapper(op, *op_args, **op_kwargs)
|
||||
# see NOTE: [What errors are Composite Compliance trying to catch?]
|
||||
except RuntimeError as err:
|
||||
raise_composite_compliance_error(
|
||||
err,
|
||||
f"- wrapped_args: {which_args_are_wrapped}\n"
|
||||
f"- wrapped_kwargs: {which_kwargs_are_wrapped}\n"
|
||||
f"- wrapped_tangent_args: {which_tang_args_are_wrapped}\n"
|
||||
f"- wrapped_tangent_kwargs: {which_tang_kwargs_are_wrapped}\n"
|
||||
)
|
||||
|
||||
def unwrap(e):
|
||||
return e.elem if isinstance(e, CCT) else e
|
||||
|
||||
actual = tree_map(fwAD.unpack_dual, actual)
|
||||
actual_primals = tree_map(
|
||||
lambda x: unwrap(x.primal),
|
||||
actual,
|
||||
is_leaf=lambda x: type(x) is fwAD.UnpackedDualTensor,
|
||||
)
|
||||
actual_tangents = tree_map(
|
||||
lambda x: unwrap(x.tangent),
|
||||
actual,
|
||||
is_leaf=lambda x: type(x) is fwAD.UnpackedDualTensor,
|
||||
)
|
||||
assert_equal_fn(actual_primals, expected_primals, equal_nan=True)
|
||||
assert_equal_fn(actual_tangents, expected_tangents, equal_nan=True)
|
||||
@@ -0,0 +1,603 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import functools
|
||||
from torch.testing import make_tensor
|
||||
from torch.testing._internal.opinfo.core import (
|
||||
OpInfo,
|
||||
SampleInput,
|
||||
)
|
||||
from torch.testing._internal.common_dtype import all_types_and
|
||||
import numpy as np
|
||||
from torch.testing._internal.autograd_function_db import (
|
||||
sample_inputs_numpy_cube,
|
||||
sample_inputs_numpy_mul,
|
||||
sample_inputs_numpy_mul_scalar,
|
||||
sample_inputs_numpy_sort,
|
||||
sample_inputs_numpy_take,
|
||||
)
|
||||
from torch import Tensor
|
||||
from torch.types import Number
|
||||
from typing import * # noqa: F403
|
||||
|
||||
# Note: [custom op db]
|
||||
#
|
||||
# This is a collection of custom operator test cases written as OpInfos
|
||||
# so they can easily be consumed by OpInfo-based tests to check if subsystems
|
||||
# support them correctly.
|
||||
|
||||
def to_numpy(tensor):
|
||||
return tensor.cpu().numpy()
|
||||
|
||||
@torch.library.custom_op("_torch_testing::numpy_cube", mutates_args=())
|
||||
def numpy_cube(x: Tensor) -> tuple[Tensor, Tensor]:
|
||||
x_np = to_numpy(x)
|
||||
dx = torch.tensor(3 * x_np ** 2, device=x.device)
|
||||
return torch.tensor(x_np ** 3, device=x.device), dx
|
||||
|
||||
@numpy_cube.register_fake
|
||||
def _(x):
|
||||
return x.clone(), x.clone()
|
||||
|
||||
def numpy_cube_setup_context(ctx, inputs, output):
|
||||
x, = inputs
|
||||
_cube, dx = output
|
||||
ctx.save_for_backward(x, dx)
|
||||
|
||||
def numpy_cube_backward(ctx, grad_out, grad_dx):
|
||||
x, dx = ctx.saved_tensors
|
||||
grad_x = numpy_mul(grad_out, dx) + 6 * numpy_mul(grad_dx, x)
|
||||
return grad_x
|
||||
|
||||
numpy_cube.register_autograd(numpy_cube_backward, setup_context=numpy_cube_setup_context)
|
||||
|
||||
def numpy_cube_vmap(info, in_dims, x):
|
||||
result = numpy_cube(x)
|
||||
return result, (in_dims[0], in_dims[0])
|
||||
|
||||
numpy_cube.register_vmap(numpy_cube_vmap)
|
||||
|
||||
@torch.library.custom_op("_torch_testing::numpy_mul", mutates_args=())
|
||||
def numpy_mul(x: Tensor, y: Tensor) -> Tensor:
|
||||
return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device)
|
||||
|
||||
@numpy_mul.register_fake
|
||||
def _(x, y):
|
||||
if x.device != y.device:
|
||||
raise AssertionError(f"x.device={x.device} != y.device={y.device}")
|
||||
return (x * y).contiguous()
|
||||
|
||||
def numpy_mul_setup_context(ctx, inputs, output):
|
||||
ctx.save_for_backward(*inputs)
|
||||
|
||||
def numpy_mul_backward(ctx, grad_out):
|
||||
x, y = ctx.saved_tensors
|
||||
grad_x = grad_out * y if ctx.needs_input_grad[0] else None
|
||||
grad_y = grad_out * x if ctx.needs_input_grad[1] else None
|
||||
return grad_x, grad_y
|
||||
|
||||
numpy_mul.register_autograd(numpy_mul_backward, setup_context=numpy_mul_setup_context)
|
||||
|
||||
def numpy_mul_vmap(info, in_dims, x, y):
|
||||
x_bdim, y_bdim = in_dims
|
||||
x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1)
|
||||
y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1)
|
||||
result = x * y
|
||||
result = result.movedim(-1, 0)
|
||||
return result, 0
|
||||
|
||||
numpy_mul.register_vmap(numpy_mul_vmap)
|
||||
|
||||
@torch.library.custom_op("_torch_testing::numpy_mul_scalar", mutates_args=())
|
||||
def numpy_mul_scalar(x: Tensor, *, scalar: float) -> Tensor:
|
||||
return torch.tensor(to_numpy(x) * scalar, device=x.device)
|
||||
|
||||
@numpy_mul_scalar.register_fake
|
||||
def _(x, *, scalar):
|
||||
return (x * scalar).contiguous()
|
||||
|
||||
def numpy_mul_scalar_setup_context(ctx, inputs, keyword_only_inputs, output):
|
||||
ctx.scalar = keyword_only_inputs["scalar"]
|
||||
|
||||
def numpy_mul_scalar_backward(ctx, grad_out):
|
||||
grad_x = grad_out * ctx.scalar
|
||||
return grad_x
|
||||
|
||||
numpy_mul_scalar.register_autograd(numpy_mul_scalar_backward, setup_context=numpy_mul_scalar_setup_context)
|
||||
|
||||
def numpy_mul_scalar_vmap(info, in_dims, x, *, scalar):
|
||||
x_bdim, = in_dims
|
||||
x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1)
|
||||
result = x * scalar
|
||||
result = result.movedim(-1, 0)
|
||||
return result, 0
|
||||
|
||||
numpy_mul_scalar.register_vmap(numpy_mul_scalar_vmap)
|
||||
|
||||
@torch.library.custom_op("_torch_testing::numpy_sort", mutates_args=())
|
||||
def numpy_sort(x: Tensor, dim: int) -> tuple[Tensor, Tensor, Tensor]:
|
||||
device = x.device
|
||||
x = to_numpy(x)
|
||||
ind = np.argsort(x, axis=dim)
|
||||
ind_inv = np.argsort(ind, axis=dim)
|
||||
result = np.take_along_axis(x, ind, axis=dim)
|
||||
return (
|
||||
torch.tensor(result, device=device),
|
||||
torch.tensor(ind, device=device),
|
||||
torch.tensor(ind_inv, device=device),
|
||||
)
|
||||
|
||||
@numpy_sort.register_fake
|
||||
def _(x, dim):
|
||||
return torch.empty_like(x), torch.empty_like(x, dtype=torch.long), torch.empty_like(x, dtype=torch.long)
|
||||
|
||||
def numpy_sort_setup_context(ctx, inputs, output):
|
||||
_out, ind, ind_inv = output
|
||||
ctx.dim = inputs[1]
|
||||
ctx.save_for_backward(ind, ind_inv)
|
||||
ctx.mark_non_differentiable(ind, ind_inv)
|
||||
|
||||
def numpy_sort_backward(ctx, grad_out, grad_ind, grad_ind_inv):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
return numpy_take(grad_out, ind_inv, ind, ctx.dim), None
|
||||
|
||||
numpy_sort.register_autograd(numpy_sort_backward, setup_context=numpy_sort_setup_context)
|
||||
|
||||
def numpy_sort_vmap(info, in_dims, x, dim):
|
||||
x_bdim, _ = in_dims
|
||||
x = x.movedim(x_bdim, 0)
|
||||
dim = dim if dim >= 0 else dim + x.dim() - 1
|
||||
result = numpy_sort(x, dim + 1)
|
||||
return result, (0, 0, 0)
|
||||
|
||||
numpy_sort.register_vmap(numpy_sort_vmap)
|
||||
|
||||
@torch.library.custom_op("_torch_testing::numpy_take", mutates_args=())
|
||||
def numpy_take(x: Tensor, ind: Tensor, ind_inv: Tensor, dim: int) -> Tensor:
|
||||
device = x.device
|
||||
x = to_numpy(x)
|
||||
ind = to_numpy(ind)
|
||||
return torch.tensor(np.take_along_axis(x, ind, dim), device=device)
|
||||
|
||||
@numpy_take.register_fake
|
||||
def _(x, ind, ind_inv, dim):
|
||||
if x.device != ind.device:
|
||||
raise AssertionError(f"x.device={x.device} != ind.device={ind.device}")
|
||||
if x.device != ind_inv.device:
|
||||
raise AssertionError(f"x.device={x.device} != ind_inv.device={ind_inv.device}")
|
||||
if ind.dtype != torch.long:
|
||||
raise AssertionError(f"ind.dtype must be torch.long, got {ind.dtype}")
|
||||
if ind_inv.dtype != torch.long:
|
||||
raise AssertionError(f"ind_inv.dtype must be torch.long, got {ind_inv.dtype}")
|
||||
return torch.empty_like(x)
|
||||
|
||||
def numpy_take_setup_context(ctx, inputs, output):
|
||||
_x, ind, ind_inv, dim = inputs
|
||||
ctx.dim = dim
|
||||
ctx.save_for_backward(ind, ind_inv)
|
||||
|
||||
def numpy_take_backward(ctx, grad_out):
|
||||
ind, ind_inv = ctx.saved_tensors
|
||||
grad_x = numpy_take(grad_out, ind_inv, ind, ctx.dim)
|
||||
return grad_x, None, None, None
|
||||
|
||||
numpy_take.register_autograd(numpy_take_backward, setup_context=numpy_take_setup_context)
|
||||
|
||||
def numpy_take_vmap(info, in_dims, x, ind, ind_inv, dim):
|
||||
x_bdim, ind_bdim, ind_inv_bdim, _ = in_dims
|
||||
|
||||
# wrap dim
|
||||
logical_dim = x.dim() if x_bdim is None else x_bdim - 1
|
||||
dim = dim if dim >= 0 else dim + logical_dim
|
||||
|
||||
def expand_bdim(x, x_bdim):
|
||||
if x_bdim is None:
|
||||
return x.expand(info.batch_size, *x.shape)
|
||||
return x.movedim(x_bdim, 0)
|
||||
|
||||
x = expand_bdim(x, x_bdim)
|
||||
ind = expand_bdim(ind, ind_bdim)
|
||||
ind_inv = expand_bdim(ind_inv, ind_inv_bdim)
|
||||
|
||||
return numpy_take(x, ind, ind_inv, dim + 1), 0
|
||||
|
||||
numpy_take.register_vmap(numpy_take_vmap)
|
||||
|
||||
@torch.library.custom_op("_torch_testing::numpy_nonzero", mutates_args=())
|
||||
def numpy_nonzero(x: Tensor) -> Tensor:
|
||||
x_np = to_numpy(x)
|
||||
res = np.stack(np.nonzero(x_np), axis=1)
|
||||
if res.shape[0] <= 1:
|
||||
raise RuntimeError("not supported")
|
||||
return torch.tensor(res, device=x.device)
|
||||
|
||||
@numpy_nonzero.register_fake
|
||||
def _(x):
|
||||
ctx = torch._custom_op.impl.get_ctx()
|
||||
i0 = ctx.create_unbacked_symint()
|
||||
shape = [i0, x.dim()]
|
||||
result = x.new_empty(shape, dtype=torch.long)
|
||||
return result
|
||||
|
||||
def sample_inputs_numpy_nonzero(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
shape = 10
|
||||
result = make_arg(shape, low=0.9, high=2)
|
||||
mask = make_tensor(shape, low=0, high=2, device=device, dtype=torch.long)
|
||||
with torch.no_grad():
|
||||
result *= mask
|
||||
|
||||
yield SampleInput(result, args=())
|
||||
|
||||
def numpy_nonzero_vmap(info, in_dims, x):
|
||||
raise NotImplementedError("Operator is data-dependent and cannot be vmapped.")
|
||||
|
||||
numpy_nonzero.register_vmap(numpy_nonzero_vmap)
|
||||
|
||||
@torch.library.custom_op("_torch_testing::numpy_view_copy", mutates_args=())
|
||||
def numpy_view_copy(x: Tensor, shape: Sequence[int]) -> Tensor:
|
||||
return torch.tensor(np.copy(to_numpy(x).reshape(shape)), device=x.device)
|
||||
|
||||
@numpy_view_copy.register_fake
|
||||
def _(x, shape) -> Tensor:
|
||||
return x.clone().view(shape).clone()
|
||||
|
||||
def numpy_view_copy_setup_context(ctx, inputs, output) -> None:
|
||||
ctx.x_shape = inputs[0].shape
|
||||
|
||||
def numpy_view_copy_backward(ctx, grad_out):
|
||||
return torch.ops._torch_testing.numpy_view_copy(grad_out, ctx.x_shape), None
|
||||
|
||||
numpy_view_copy.register_autograd(numpy_view_copy_backward, setup_context=numpy_view_copy_setup_context)
|
||||
|
||||
def numpy_view_copy_vmap(info, in_dims, x, shape):
|
||||
x_bdim, _ = in_dims
|
||||
x = x.movedim(x_bdim, 0)
|
||||
x_shape = x.shape[0]
|
||||
batch_shape = (x_shape, *shape)
|
||||
result = numpy_view_copy(x, batch_shape)
|
||||
return result, 0
|
||||
|
||||
numpy_view_copy.register_vmap(numpy_view_copy_vmap)
|
||||
|
||||
def sample_inputs_numpy_view_copy(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
result = make_arg(2, 3, 4, low=0.9, high=2)
|
||||
yield SampleInput(result, args=([2, 12],))
|
||||
|
||||
@torch.library.custom_op('_torch_testing::numpy_cat', mutates_args=())
|
||||
def numpy_cat(xs: Sequence[Tensor], dim: int) -> Tensor:
|
||||
if len(xs) == 0:
|
||||
raise AssertionError("xs must not be empty")
|
||||
if not all(x.device == xs[0].device for x in xs):
|
||||
raise AssertionError("All tensors must be on the same device")
|
||||
if not all(x.dtype == xs[0].dtype for x in xs):
|
||||
raise AssertionError("All tensors must have the same dtype")
|
||||
np_xs = [to_numpy(x) for x in xs]
|
||||
np_out = np.concatenate(np_xs, axis=dim)
|
||||
return torch.tensor(np_out, device=xs[0].device)
|
||||
|
||||
@numpy_cat.register_fake
|
||||
def _(xs, dim):
|
||||
if len(xs) == 0:
|
||||
raise AssertionError("xs must not be empty")
|
||||
if not all(x.device == xs[0].device for x in xs):
|
||||
raise AssertionError("All tensors must be on the same device")
|
||||
if not all(x.dtype == xs[0].dtype for x in xs):
|
||||
raise AssertionError("All tensors must have the same dtype")
|
||||
return torch.cat(xs, dim=dim)
|
||||
|
||||
def numpy_cat_setup_context(ctx, inputs, output):
|
||||
xs, dim = inputs
|
||||
ctx.dim_sizes = [x.shape[dim] for x in xs]
|
||||
ctx.dim = dim
|
||||
|
||||
def numpy_cat_backward(ctx, grad_out):
|
||||
dim_sizes = ctx.dim_sizes
|
||||
dim = ctx.dim
|
||||
|
||||
splits = list(np.cumsum(dim_sizes)[:-1])
|
||||
grad_xs = torch.ops._torch_testing.numpy_split_copy(grad_out, splits, dim)
|
||||
return grad_xs, None
|
||||
|
||||
numpy_cat.register_autograd(numpy_cat_backward, setup_context=numpy_cat_setup_context)
|
||||
|
||||
def numpy_cat_vmap(info, in_dims, x, dim):
|
||||
x_bdim, = in_dims
|
||||
result = numpy_cat(x, dim)
|
||||
return result, x_bdim
|
||||
|
||||
numpy_cat.register_vmap(numpy_cat_vmap)
|
||||
|
||||
def sample_inputs_numpy_cat(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
r0 = make_arg(2, 3, 4, low=0.9, high=2)
|
||||
r1 = make_arg(4, 3, 4, low=0.9, high=2)
|
||||
r2 = make_arg(5, 3, 4, low=0.9, high=2)
|
||||
yield SampleInput([r0, r1, r2], args=(0,))
|
||||
|
||||
@torch.library.custom_op('_torch_testing::numpy_split_copy', mutates_args=())
|
||||
def numpy_split_copy(x: Tensor, splits: Sequence[int], dim: int) -> List[Tensor]:
|
||||
x_np = to_numpy(x)
|
||||
arrs = np.split(x_np, splits, axis=dim)
|
||||
return [torch.tensor(arr, device=x.device, dtype=x.dtype) for arr in arrs]
|
||||
|
||||
@numpy_split_copy.register_fake
|
||||
def _(x, splits, dim):
|
||||
return [xi.clone() for xi in torch.tensor_split(x, splits, dim)]
|
||||
|
||||
def numpy_split_copy_setup_context(ctx, inputs, output):
|
||||
_, _, dim = inputs
|
||||
ctx.dim = dim
|
||||
|
||||
def numpy_split_copy_backward(ctx, grad_out):
|
||||
result = torch.ops._torch_testing.numpy_cat(grad_out, dim=ctx.dim)
|
||||
return result, None, None
|
||||
|
||||
numpy_split_copy.register_autograd(numpy_split_copy_backward, setup_context=numpy_split_copy_setup_context)
|
||||
|
||||
def numpy_split_copy_vmap(info, in_dims, x, splits, dim):
|
||||
x_bdim, _ , _ = in_dims
|
||||
x = x.movedim(x_bdim, 0)
|
||||
result = numpy_split_copy(x, splits, dim + 1)
|
||||
return result, 0
|
||||
|
||||
numpy_split_copy.register_vmap(numpy_split_copy_vmap)
|
||||
|
||||
def sample_inputs_numpy_split_copy(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
|
||||
x = make_arg(2, 9, low=0.9, high=2)
|
||||
yield SampleInput(x, args=([1, 3, 6], 1))
|
||||
|
||||
@torch.library.custom_op('_torch_testing::numpy_split_copy_with_int', mutates_args=())
|
||||
def numpy_split_copy_with_int(x: Tensor, splits: Sequence[int], dim: int) -> tuple[List[Tensor], int]:
|
||||
x_np = to_numpy(x)
|
||||
arrs = np.split(x_np, splits, axis=dim)
|
||||
return [torch.tensor(arr, device=x.device, dtype=x.dtype) for arr in arrs], len(splits)
|
||||
|
||||
@numpy_split_copy_with_int.register_fake
|
||||
def _(x, splits, dim):
|
||||
return [xi.clone() for xi in torch.tensor_split(x, splits, dim)], len(splits)
|
||||
|
||||
def numpy_split_copy_with_int_setup_context(ctx, inputs, output):
|
||||
_, _, dim = inputs
|
||||
ctx.dim = dim
|
||||
|
||||
def numpy_split_copy_with_int_backward(ctx, grad_out, _):
|
||||
return torch.ops._torch_testing.numpy_cat(grad_out, dim=ctx.dim), None, None
|
||||
|
||||
numpy_split_copy_with_int.register_autograd(
|
||||
numpy_split_copy_with_int_backward,
|
||||
setup_context=numpy_split_copy_with_int_setup_context)
|
||||
|
||||
def numpy_split_copy_with_int_vmap(info, in_dims, x, splits, dim):
|
||||
x_bdim, _ , _ = in_dims
|
||||
x = x.movedim(x_bdim, 0)
|
||||
result, len_split = numpy_split_copy_with_int(x, splits, dim + 1)
|
||||
return (result, len_split), ([0 for _ in range(len(result))], None)
|
||||
|
||||
numpy_split_copy_with_int.register_vmap(numpy_split_copy_with_int_vmap)
|
||||
|
||||
@torch.library.custom_op("_torch_testing::numpy_nms", mutates_args=())
|
||||
def numpy_nms(boxes: Tensor, scores: Tensor, iou_threshold: Number) -> Tensor:
|
||||
# Adapted from Ross Girshick's fast-rcnn implementation at
|
||||
# https://github.com/rbgirshick/fast-rcnn/blob/master/lib/utils/nms.py
|
||||
if boxes.device != scores.device:
|
||||
raise AssertionError(f"boxes.device={boxes.device} != scores.device={scores.device}")
|
||||
device = boxes.device
|
||||
|
||||
boxes = to_numpy(boxes)
|
||||
scores = to_numpy(scores)
|
||||
|
||||
N = boxes.shape[0]
|
||||
if boxes.shape != (N, 4):
|
||||
raise AssertionError(f"boxes.shape must be (N, 4), got {boxes.shape}")
|
||||
if scores.shape != (N,):
|
||||
raise AssertionError(f"scores.shape must be (N,), got {scores.shape}")
|
||||
|
||||
x1 = boxes[:, 0]
|
||||
y1 = boxes[:, 1]
|
||||
x2 = boxes[:, 2]
|
||||
y2 = boxes[:, 3]
|
||||
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = np.where(ovr <= iou_threshold)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
result = torch.tensor(np.stack(keep), device=device)
|
||||
# Needed for data-dependent condition :(
|
||||
if result.size(0) < 2:
|
||||
raise AssertionError(f"result.size(0) must be >= 2, got {result.size(0)}")
|
||||
return result
|
||||
|
||||
@numpy_nms.register_fake
|
||||
def _(boxes, scores, iou_threshold):
|
||||
if boxes.device != scores.device:
|
||||
raise AssertionError(f"boxes.device={boxes.device} != scores.device={scores.device}")
|
||||
N = boxes.shape[0]
|
||||
if boxes.shape != (N, 4):
|
||||
raise AssertionError(f"boxes.shape must be (N, 4), got {boxes.shape}")
|
||||
if scores.shape != (N,):
|
||||
raise AssertionError(f"scores.shape must be (N,), got {scores.shape}")
|
||||
|
||||
ctx = torch._custom_op.impl.get_ctx()
|
||||
i0 = ctx.create_unbacked_symint()
|
||||
result = boxes.new_empty([i0], dtype=torch.int64)
|
||||
return result
|
||||
|
||||
def numpy_nms_vmap(info, in_dims, boxes, scores, iou_threshold):
|
||||
raise NotImplementedError("Operator is data-dependent and cannot be vmapped.")
|
||||
|
||||
numpy_nms.register_vmap(numpy_nms_vmap)
|
||||
|
||||
def sample_inputs_numpy_nms(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(make_tensor, device=device, dtype=dtype)
|
||||
N = 64
|
||||
xs = make_arg([N], low=0, high=28)
|
||||
dx = make_arg([N], low=0, high=4)
|
||||
ys = make_arg([N], low=0, high=28)
|
||||
dy = make_arg([N], low=0, high=4)
|
||||
boxes = torch.stack([xs, ys, xs + dx, ys + dy], dim=1).requires_grad_(requires_grad)
|
||||
scores = make_arg([N], low=0, high=1, requires_grad=requires_grad)
|
||||
iou_threshold = make_arg([], low=0, high=1).item()
|
||||
|
||||
yield SampleInput(boxes, args=(scores, iou_threshold))
|
||||
|
||||
custom_op_db = [
|
||||
OpInfo(
|
||||
'NumpyCubeCustomOp',
|
||||
op=numpy_cube._opoverload,
|
||||
sample_inputs_func=sample_inputs_numpy_cube,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyMulCustomOp',
|
||||
op=numpy_mul._opoverload,
|
||||
sample_inputs_func=sample_inputs_numpy_mul,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyMulScalarCustomOp',
|
||||
op=numpy_mul_scalar._opoverload,
|
||||
sample_inputs_func=sample_inputs_numpy_mul_scalar,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpySortCustomOp',
|
||||
op=numpy_sort._opoverload,
|
||||
sample_inputs_func=sample_inputs_numpy_sort,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyTakeCustomOp',
|
||||
op=numpy_take._opoverload,
|
||||
sample_inputs_func=sample_inputs_numpy_take,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyNonzeroCustomOp',
|
||||
op=numpy_nonzero._opoverload,
|
||||
sample_inputs_func=sample_inputs_numpy_nonzero,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_autograd=False,
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyNMSCustomOp',
|
||||
op=torch.ops._torch_testing.numpy_nms,
|
||||
sample_inputs_func=sample_inputs_numpy_nms,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_autograd=False,
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyViewCopyCustomOp',
|
||||
op=torch.ops._torch_testing.numpy_view_copy,
|
||||
sample_inputs_func=sample_inputs_numpy_view_copy,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_autograd=True,
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpyCatCustomOp',
|
||||
op=torch.ops._torch_testing.numpy_cat,
|
||||
sample_inputs_func=sample_inputs_numpy_cat,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_autograd=True,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpySplitCopyCustomOp',
|
||||
op=torch.ops._torch_testing.numpy_split_copy,
|
||||
sample_inputs_func=sample_inputs_numpy_split_copy,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_autograd=True,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
supports_out=False,
|
||||
),
|
||||
OpInfo(
|
||||
'NumpySplitCopyWithIntCustomOp',
|
||||
op=torch.ops._torch_testing.numpy_split_copy_with_int,
|
||||
sample_inputs_func=sample_inputs_numpy_split_copy,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
gradcheck_wrapper=lambda op, *args, **kwargs: op(*args, **kwargs)[0],
|
||||
supports_autograd=True,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
supports_out=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# ==============================================================
|
||||
# some mechanical test cases
|
||||
# ==============================================================
|
||||
|
||||
lib = torch.library.Library("_torch_testing", "FRAGMENT") # noqa: TOR901
|
||||
|
||||
lib.define("source0(Tensor x) -> Tensor")
|
||||
|
||||
@torch.library.register_fake("_torch_testing::source0", lib=lib)
|
||||
def _(x):
|
||||
return x.clone()
|
||||
|
||||
lib.define("source1(Tensor x) -> Tensor")
|
||||
|
||||
def source1_fake(x):
|
||||
return x.clone()
|
||||
|
||||
torch.library.register_fake("_torch_testing::source1", source1_fake, lib=lib)
|
||||
|
||||
lib.define("source2(Tensor x) -> Tensor")
|
||||
|
||||
@torch.library.register_fake("_torch_testing::source2", lib=lib)
|
||||
def _(x):
|
||||
return x.clone()
|
||||
|
||||
lib.define("source3(Tensor x) -> Tensor")
|
||||
|
||||
def source3_fake(x):
|
||||
return x.clone()
|
||||
|
||||
torch.library.register_fake("_torch_testing::source3", source3_fake, lib=lib)
|
||||
|
||||
|
||||
@torch.library.custom_op("_torch_testing::source4", mutates_args=())
|
||||
def source4(x: Tensor) -> Tensor:
|
||||
return x.clone()
|
||||
|
||||
@source4.register_fake
|
||||
def _(x):
|
||||
return x.clone()
|
||||
|
||||
@torch.library.custom_op("_torch_testing::source5", mutates_args=())
|
||||
def source5(x: Tensor) -> Tensor:
|
||||
return x.clone()
|
||||
|
||||
def source5_fake(x):
|
||||
return x.clone()
|
||||
|
||||
source5.register_fake(source5_fake)
|
||||
@@ -0,0 +1,161 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
|
||||
from collections import namedtuple
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch.utils._python_dispatch import return_and_correct_aliasing
|
||||
|
||||
|
||||
FancyNamedTuple = namedtuple("FancyNamedTuple", ["foo", "bar"])
|
||||
|
||||
|
||||
# A simple tensor subclass that holds a tensor with custom metadata and custom method
|
||||
class ConstantExtraMetadataTensor(torch.Tensor):
|
||||
@staticmethod
|
||||
def __new__(cls, elem):
|
||||
shape = elem.shape
|
||||
kwargs = {}
|
||||
kwargs["strides"] = elem.stride()
|
||||
kwargs["storage_offset"] = elem.storage_offset()
|
||||
kwargs["device"] = elem.device
|
||||
kwargs["layout"] = elem.layout
|
||||
kwargs["requires_grad"] = elem.requires_grad
|
||||
kwargs["dtype"] = elem.dtype
|
||||
return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)
|
||||
|
||||
def __init__(self, elem):
|
||||
self.elem = elem
|
||||
self.constant_attribute = 4
|
||||
|
||||
def __repr__(self):
|
||||
inner_repr = repr(self.elem)
|
||||
return f"CustomTensor({inner_repr})"
|
||||
|
||||
def get_complicated_metadata(self):
|
||||
return FancyNamedTuple(self.constant_attribute, self.constant_attribute)
|
||||
|
||||
def __tensor_flatten__(self):
|
||||
return ["elem"], self.constant_attribute
|
||||
|
||||
def add_constant(self, a):
|
||||
self.constant_attribute += a
|
||||
|
||||
@staticmethod
|
||||
def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
|
||||
if meta is None:
|
||||
raise AssertionError("Expected meta to not be None")
|
||||
elem = inner_tensors["elem"]
|
||||
out = ConstantExtraMetadataTensor(elem)
|
||||
out.constant_attribute = meta
|
||||
return out
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args, kwargs):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
args_inner = pytree.tree_map_only(
|
||||
ConstantExtraMetadataTensor, lambda x: x.elem, args
|
||||
)
|
||||
|
||||
kwargs_inner = pytree.tree_map_only(
|
||||
ConstantExtraMetadataTensor, lambda x: x.elem, kwargs
|
||||
)
|
||||
|
||||
out_inner = func(*args_inner, **kwargs_inner)
|
||||
out_inner_flat, spec = pytree.tree_flatten(out_inner)
|
||||
# for aten ops that return non-tensors, just assume that
|
||||
# our cust inner tensors return the same value
|
||||
out_flat = [
|
||||
ConstantExtraMetadataTensor(o_inner)
|
||||
if isinstance(o_inner, torch.Tensor)
|
||||
else o_inner
|
||||
for o_inner in out_inner_flat
|
||||
]
|
||||
out = pytree.tree_unflatten(out_flat, spec)
|
||||
return return_and_correct_aliasing(func, args, kwargs, out)
|
||||
|
||||
|
||||
# A simple tensor subclass that always returns plain tensor during __torch_dispatch__
|
||||
# It is similar to TwoTensor and is used to simulate torchao quantized tensors
|
||||
class CustomTensorPlainOut(torch.Tensor):
|
||||
@staticmethod
|
||||
def __new__(cls, elem1, elem2):
|
||||
shape = elem1.shape
|
||||
kwargs = {}
|
||||
kwargs["strides"] = elem1.stride()
|
||||
kwargs["storage_offset"] = elem1.storage_offset()
|
||||
kwargs["device"] = elem1.device
|
||||
kwargs["layout"] = elem1.layout
|
||||
kwargs["requires_grad"] = elem1.requires_grad
|
||||
kwargs["dtype"] = elem1.dtype
|
||||
return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)
|
||||
|
||||
def __init__(self, elem1, elem2):
|
||||
self.elem1 = elem1
|
||||
self.elem2 = elem2
|
||||
|
||||
def get_elem(self):
|
||||
return self.elem1
|
||||
|
||||
def __repr__(self):
|
||||
inner_repr_1 = repr(self.elem1)
|
||||
inner_repr_2 = repr(self.elem2)
|
||||
return f"CustomTensorPlainOut({inner_repr_1}, {inner_repr_2})"
|
||||
|
||||
def __tensor_flatten__(self):
|
||||
return ["elem1", "elem2"], None
|
||||
|
||||
@staticmethod
|
||||
def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
|
||||
elem1 = inner_tensors["elem1"]
|
||||
elem2 = inner_tensors["elem2"]
|
||||
out = CustomTensorPlainOut(elem1, elem2)
|
||||
return out
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args, kwargs):
|
||||
# Don't use this tensor with view ops
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
args_inner_1 = pytree.tree_map_only(
|
||||
CustomTensorPlainOut, lambda x: x.elem1, args
|
||||
)
|
||||
|
||||
kwargs_inner_1 = pytree.tree_map_only(
|
||||
CustomTensorPlainOut, lambda x: x.elem1, kwargs
|
||||
)
|
||||
|
||||
args_inner_2 = pytree.tree_map_only(
|
||||
CustomTensorPlainOut, lambda x: x.elem2, args
|
||||
)
|
||||
|
||||
kwargs_inner_2 = pytree.tree_map_only(
|
||||
CustomTensorPlainOut, lambda x: x.elem2, kwargs
|
||||
)
|
||||
|
||||
out_inner_1 = func(*args_inner_1, **kwargs_inner_1)
|
||||
out_inner_2 = func(*args_inner_2, **kwargs_inner_2)
|
||||
|
||||
out_inner_flat_1, spec = pytree.tree_flatten(out_inner_1)
|
||||
out_inner_flat_2, spec = pytree.tree_flatten(out_inner_2)
|
||||
|
||||
if func.is_view:
|
||||
new_out = pytree.tree_unflatten(
|
||||
(
|
||||
CustomTensorPlainOut(tensor1, tensor2)
|
||||
for tensor1, tensor2 in zip(
|
||||
out_inner_flat_1, out_inner_flat_2, strict=True
|
||||
)
|
||||
),
|
||||
spec,
|
||||
)
|
||||
return return_and_correct_aliasing(func, args, kwargs, new_out)
|
||||
|
||||
out_new = (
|
||||
out_inner_flat_1[ix] + out_inner_flat_2[ix]
|
||||
for ix in range(len(out_inner_flat_1))
|
||||
)
|
||||
|
||||
return pytree.tree_unflatten(out_new, spec)
|
||||
@@ -0,0 +1 @@
|
||||
# mypy: ignore-errors
|
||||
@@ -0,0 +1,10 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Net(nn.Module):
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.linear = nn.Linear(10, 20)
|
||||
@@ -0,0 +1,11 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Net(nn.Module):
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.linear = nn.Linear(10, 20)
|
||||
self.relu = nn.ReLU()
|
||||
@@ -0,0 +1,199 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from functools import partial, wraps
|
||||
|
||||
import torch.distributed as dist
|
||||
import torch.distributed.rpc as rpc
|
||||
from torch.distributed.rpc import _rref_context_get_debug_info
|
||||
from torch.testing._internal.common_utils import FILE_SCHEMA, TEST_WITH_TSAN
|
||||
|
||||
|
||||
if not dist.is_available():
|
||||
print("c10d not available, skipping tests", file=sys.stderr)
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
INIT_METHOD_TEMPLATE = FILE_SCHEMA + "{file_name}"
|
||||
|
||||
def dist_init(
|
||||
old_test_method=None,
|
||||
setup_rpc: bool = True,
|
||||
clean_shutdown: bool = True,
|
||||
faulty_messages=None,
|
||||
messages_to_delay=None,
|
||||
):
|
||||
"""
|
||||
We use this decorator for setting up and tearing down state since
|
||||
MultiProcessTestCase runs each `test*` method in a separate process and
|
||||
each process just runs the `test*` method without actually calling
|
||||
'setUp' and 'tearDown' methods of unittest.
|
||||
|
||||
Note: pass the string representation of MessageTypes that should be used
|
||||
with the faulty agent's send function. By default, all retriable messages
|
||||
("RREF_FORK_REQUEST", "RREF_CHILD_ACCEPT", "RREF_USER_DELETE",
|
||||
"CLEANUP_AUTOGRAD_CONTEXT_REQ") will use the faulty send (this default is
|
||||
set from faulty_rpc_agent_test_fixture.py).
|
||||
"""
|
||||
# If we use dist_init without arguments (ex: @dist_init), old_test_method is
|
||||
# appropriately set and we return the wrapper appropriately. On the other
|
||||
# hand if dist_init has arguments (ex: @dist_init(clean_shutdown=False)),
|
||||
# old_test_method is None and we return a functools.partial which is the real
|
||||
# decorator that is used and as a result we recursively call dist_init with
|
||||
# old_test_method and the rest of the arguments appropriately set.
|
||||
if old_test_method is None:
|
||||
return partial(
|
||||
dist_init,
|
||||
setup_rpc=setup_rpc,
|
||||
clean_shutdown=clean_shutdown,
|
||||
faulty_messages=faulty_messages,
|
||||
messages_to_delay=messages_to_delay,
|
||||
)
|
||||
|
||||
@wraps(old_test_method)
|
||||
def new_test_method(self, *arg, **kwargs):
|
||||
# Setting _ignore_rref_leak to make sure OwnerRRefs are properly deleted
|
||||
# in tests.
|
||||
import torch.distributed.rpc.api as api
|
||||
|
||||
api._ignore_rref_leak = False
|
||||
self.worker_id = self.rank
|
||||
self.setup_fault_injection(faulty_messages, messages_to_delay)
|
||||
|
||||
rpc_backend_options = self.rpc_backend_options
|
||||
if setup_rpc:
|
||||
if TEST_WITH_TSAN:
|
||||
# TSAN runs much slower.
|
||||
rpc_backend_options.rpc_timeout = rpc.constants.DEFAULT_RPC_TIMEOUT_SEC * 5
|
||||
rpc.constants.DEFAULT_SHUTDOWN_TIMEOUT = 60
|
||||
|
||||
rpc.init_rpc(
|
||||
name=f"worker{self.rank:d}",
|
||||
backend=self.rpc_backend,
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
rpc_backend_options=rpc_backend_options,
|
||||
)
|
||||
|
||||
return_value = old_test_method(self, *arg, **kwargs)
|
||||
|
||||
if setup_rpc:
|
||||
rpc.shutdown(graceful=clean_shutdown)
|
||||
|
||||
return return_value
|
||||
|
||||
return new_test_method
|
||||
|
||||
|
||||
def noop() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def wait_until_node_failure(rank: int, expected_error_regex: str = ".*") -> str:
|
||||
"""
|
||||
Loops until an RPC to the given rank fails. This is used to
|
||||
indicate that the node has failed in unit tests.
|
||||
Args:
|
||||
rank (int): Rank of the node expected to fail
|
||||
expected_error_regex (optional, str): Regex of exception message expected. Useful to ensure a specific failure
|
||||
occurs, not just any.
|
||||
"""
|
||||
while True:
|
||||
try:
|
||||
rpc.rpc_sync(f"worker{rank}", noop, args=())
|
||||
time.sleep(0.1)
|
||||
except Exception as e:
|
||||
if re.search(pattern=expected_error_regex, string=str(e)):
|
||||
return str(e)
|
||||
|
||||
|
||||
def wait_until_pending_futures_and_users_flushed(timeout: int = 20) -> None:
|
||||
"""
|
||||
The RRef protocol holds forkIds of rrefs in a map until those forks are
|
||||
confirmed by the owner. The message confirming the fork may arrive after
|
||||
our tests check whether this map is empty, which leads to failures and
|
||||
flaky tests. to_here also does not guarantee that we have finished
|
||||
processind the owner's confirmation message for the RRef. This function
|
||||
loops until the map is empty, which means the messages have been received
|
||||
as processed. Call this function before asserting the map returned by
|
||||
_get_debug_info is empty.
|
||||
"""
|
||||
start = time.time()
|
||||
while True:
|
||||
debug_info = _rref_context_get_debug_info()
|
||||
num_pending_futures = int(debug_info["num_pending_futures"])
|
||||
num_pending_users = int(debug_info["num_pending_users"])
|
||||
if num_pending_futures == 0 and num_pending_users == 0:
|
||||
break
|
||||
time.sleep(0.1)
|
||||
if time.time() - start > timeout:
|
||||
raise ValueError(
|
||||
f"Timed out waiting to flush pending futures and users, "
|
||||
f"had {num_pending_futures} pending futures and {num_pending_users} pending users"
|
||||
)
|
||||
|
||||
|
||||
def get_num_owners_and_forks() -> tuple[str, str]:
|
||||
"""
|
||||
Retrieves number of OwnerRRefs and forks on this node from
|
||||
_rref_context_get_debug_info.
|
||||
"""
|
||||
rref_dbg_info = _rref_context_get_debug_info()
|
||||
num_owners = rref_dbg_info["num_owner_rrefs"]
|
||||
num_forks = rref_dbg_info["num_forks"]
|
||||
return num_owners, num_forks
|
||||
|
||||
|
||||
def wait_until_owners_and_forks_on_rank(
|
||||
num_owners: int, num_forks: int, rank: int, timeout: int = 20
|
||||
) -> None:
|
||||
"""
|
||||
Waits until timeout for num_forks and num_owners to exist on the rank. Used
|
||||
to ensure proper deletion of RRefs in tests.
|
||||
"""
|
||||
start = time.time()
|
||||
while True:
|
||||
num_owners_on_rank, num_forks_on_rank = rpc.rpc_sync(
|
||||
worker_name(rank), get_num_owners_and_forks, args=(), timeout=5
|
||||
)
|
||||
num_owners_on_rank = int(num_owners_on_rank)
|
||||
num_forks_on_rank = int(num_forks_on_rank)
|
||||
if num_owners_on_rank == num_owners and num_forks_on_rank == num_forks:
|
||||
return
|
||||
time.sleep(1)
|
||||
if time.time() - start > timeout:
|
||||
raise ValueError(
|
||||
f"Timed out waiting {timeout} sec for {num_owners} owners and {num_forks} forks on rank,"
|
||||
f" had {num_owners_on_rank} owners and {num_forks_on_rank} forks"
|
||||
)
|
||||
|
||||
|
||||
def initialize_pg(init_method, rank: int, world_size: int) -> None:
|
||||
# This is for tests using `dist.barrier`.
|
||||
if not dist.is_initialized():
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method=init_method,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
)
|
||||
|
||||
|
||||
def worker_name(rank: int) -> str:
|
||||
return f"worker{rank}"
|
||||
|
||||
|
||||
def get_function_event(function_events, partial_event_name):
|
||||
"""
|
||||
Returns the first event that matches partial_event_name in the provided
|
||||
function_events. These function_events should be the output of
|
||||
torch.autograd.profiler.function_events().
|
||||
|
||||
Args:
|
||||
function_events: function_events returned by the profiler.
|
||||
event_name (str): partial key that the event was profiled with.
|
||||
"""
|
||||
event = [event for event in function_events if partial_event_name in event.name][0] # noqa: RUF015
|
||||
return event
|
||||
+1
@@ -0,0 +1 @@
|
||||
# mypy: allow-untyped-defs
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import sys
|
||||
from functools import partial, wraps
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed import rpc
|
||||
from torch.testing._internal.common_distributed import (
|
||||
MultiProcessTestCase,
|
||||
TEST_SKIPS,
|
||||
tp_transports,
|
||||
)
|
||||
|
||||
|
||||
TEST_GPU_NUM = 4
|
||||
|
||||
|
||||
class ShardedTensorTestBase(MultiProcessTestCase):
|
||||
@property
|
||||
def world_size(self):
|
||||
return TEST_GPU_NUM
|
||||
|
||||
def init_pg(self, backend="nccl"):
|
||||
if backend not in ["nccl", "gloo", "mpi", "hccl"]:
|
||||
raise RuntimeError(f"Backend {backend} not supported!")
|
||||
|
||||
dist.init_process_group(
|
||||
backend=backend,
|
||||
world_size=self.world_size,
|
||||
rank=self.rank,
|
||||
init_method=f"file://{self.file_name}",
|
||||
)
|
||||
|
||||
# set device for nccl pg for collectives
|
||||
if backend == "nccl":
|
||||
torch.cuda.set_device(self.rank)
|
||||
|
||||
def init_rpc(self):
|
||||
rpc_backend_options = rpc.TensorPipeRpcBackendOptions(
|
||||
_transports=tp_transports()
|
||||
)
|
||||
rpc_backend_options.init_method = f"file://{self.file_name}"
|
||||
for rank in range(self.world_size):
|
||||
rpc_backend_options.set_device_map(
|
||||
f"worker{rank}", {rank: self.rank, self.rank: rank}
|
||||
)
|
||||
|
||||
rpc.init_rpc(
|
||||
name=f"worker{self.rank:d}",
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
rpc_backend_options=rpc_backend_options,
|
||||
)
|
||||
|
||||
def init_comms(self, init_rpc=True, backend="nccl"):
|
||||
if init_rpc:
|
||||
self.init_rpc()
|
||||
self.init_pg(backend=backend)
|
||||
|
||||
def destroy_comms(self, destroy_rpc=True):
|
||||
# Wait for all ranks to reach here before starting shutdown.
|
||||
dist.barrier()
|
||||
|
||||
if destroy_rpc:
|
||||
rpc.shutdown()
|
||||
dist.destroy_process_group()
|
||||
|
||||
def setUp(self) -> None:
|
||||
super().setUp()
|
||||
self._spawn_processes()
|
||||
|
||||
def assert_sharded_tensor_equal(self, st1, st2):
|
||||
st1_local_shards = st1.local_shards()
|
||||
st2_local_shards = st2.local_shards()
|
||||
self.assertEqual(len(st1_local_shards), len(st2_local_shards))
|
||||
for i, st1_local_shard in enumerate(st1_local_shards):
|
||||
self.assertEqual(st1_local_shard.tensor, st2_local_shards[i].tensor)
|
||||
self.assertEqual(st1_local_shard.metadata, st2_local_shards[i].metadata)
|
||||
|
||||
self.assertEqual(st1.metadata(), st2.metadata())
|
||||
self.assertEqual(st1.sharding_spec(), st2.sharding_spec())
|
||||
self.assertEqual(len(st1.remote_shards()), len(st2.remote_shards()))
|
||||
|
||||
|
||||
# wrapper to initialize comms (processgroup + rpc)
|
||||
def with_comms(func=None, init_rpc=True, backend="nccl"):
|
||||
if func is None:
|
||||
return partial(
|
||||
with_comms,
|
||||
init_rpc=init_rpc,
|
||||
backend=backend,
|
||||
)
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(self, *args, **kwargs):
|
||||
if backend == "nccl" and torch.cuda.device_count() < self.world_size:
|
||||
sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code)
|
||||
self.init_comms(init_rpc=init_rpc, backend=backend)
|
||||
func(self, *args, **kwargs)
|
||||
self.destroy_comms(destroy_rpc=init_rpc)
|
||||
|
||||
return wrapper
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import builtins
|
||||
|
||||
import torch
|
||||
from torch.distributed._shard.sharding_spec import (
|
||||
ChunkShardingSpec,
|
||||
EnumerableShardingSpec,
|
||||
ShardMetadata,
|
||||
)
|
||||
from torch.distributed._shard.sharding_spec._internals import (
|
||||
get_chunked_dim_size,
|
||||
get_split_size,
|
||||
)
|
||||
|
||||
|
||||
def generate_chunk_sharding_specs_for_test(sharding_dim):
|
||||
return [
|
||||
ChunkShardingSpec(
|
||||
dim=sharding_dim,
|
||||
placements=[
|
||||
"rank:0/cuda:0",
|
||||
"rank:1/cuda:1",
|
||||
"rank:2/cuda:2",
|
||||
"rank:3/cuda:3",
|
||||
],
|
||||
),
|
||||
# Test different ordering. (Case 1)
|
||||
ChunkShardingSpec(
|
||||
dim=sharding_dim,
|
||||
placements=[
|
||||
"rank:2/cuda:2",
|
||||
"rank:3/cuda:3",
|
||||
"rank:0/cuda:0",
|
||||
"rank:1/cuda:1",
|
||||
],
|
||||
),
|
||||
# Test different ordering. (Case 2)
|
||||
ChunkShardingSpec(
|
||||
dim=sharding_dim,
|
||||
placements=[
|
||||
"rank:3/cuda:3",
|
||||
"rank:0/cuda:0",
|
||||
"rank:1/cuda:1",
|
||||
"rank:2/cuda:2",
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def generate_enumerable_sharding_specs_for_test():
|
||||
return [
|
||||
EnumerableShardingSpec(
|
||||
[
|
||||
ShardMetadata(
|
||||
shard_offsets=[0, 0],
|
||||
shard_sizes=[5, 5],
|
||||
placement="rank:0/cuda:0",
|
||||
),
|
||||
ShardMetadata(
|
||||
shard_offsets=[5, 0],
|
||||
shard_sizes=[5, 5],
|
||||
placement="rank:1/cuda:1",
|
||||
),
|
||||
ShardMetadata(
|
||||
shard_offsets=[0, 5],
|
||||
shard_sizes=[5, 5],
|
||||
placement="rank:2/cuda:2",
|
||||
),
|
||||
ShardMetadata(
|
||||
shard_offsets=[5, 5],
|
||||
shard_sizes=[5, 5],
|
||||
placement="rank:3/cuda:3",
|
||||
),
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def generate_local_weight_sharding_params_for_test(
|
||||
local_weight, sharded_dim, gpu_num, spec, rank
|
||||
):
|
||||
"""
|
||||
Shard the local weight based the given spec, so we can compare against
|
||||
the one from sharded tensor.
|
||||
|
||||
Args:
|
||||
local_weight: weight matrix to be sharded.
|
||||
sharded_dim: The dimension which we shard on.
|
||||
gpu_num: number of ranks.
|
||||
spec: sharding spec.
|
||||
rank: # of cuda process.
|
||||
|
||||
Returns:
|
||||
start_pos: start position of sharded weight on the given rank.
|
||||
chunk_size: chunk size of sharded weight on the given rank.
|
||||
"""
|
||||
sharding_dim_size = local_weight.size(sharded_dim)
|
||||
split_size = get_split_size(sharding_dim_size, gpu_num)
|
||||
current_offsets = 0
|
||||
start_pos = current_offsets
|
||||
for idx, placement in enumerate(spec.placements):
|
||||
chunk_size = get_chunked_dim_size(sharding_dim_size, split_size, idx)
|
||||
if rank == placement.rank():
|
||||
start_pos = current_offsets
|
||||
break
|
||||
current_offsets += chunk_size
|
||||
return start_pos, chunk_size
|
||||
|
||||
|
||||
def clone_module_parameter(module, param_name):
|
||||
"""
|
||||
Clone a parameter from a given existing module.
|
||||
|
||||
Args:
|
||||
module (:class:`torch.nn.Module`): Module whose parameter needs to be cloned.
|
||||
param_name (str): Name of the parameter of ``module`` that needs to be cloned.
|
||||
|
||||
Returns: cloned tensor as :class:`torch.nn.Parameter`.
|
||||
"""
|
||||
tensor = getattr(module, param_name)
|
||||
return torch.nn.Parameter(tensor.detach().clone())
|
||||
|
||||
|
||||
def gen_binary_op_func(python_op, inplace=False):
|
||||
src_lines = ["def f(lhs, rhs):"]
|
||||
if "torch" in python_op:
|
||||
src_lines.append(f" return {python_op}(lhs, rhs)\n")
|
||||
elif inplace:
|
||||
src_lines.append(f" lhs {python_op}= rhs\n return lhs\n")
|
||||
else:
|
||||
src_lines.append(f" return lhs {python_op} rhs\n")
|
||||
|
||||
code_str = "\n".join(src_lines)
|
||||
g = {"torch": torch}
|
||||
builtins.exec(code_str, g)
|
||||
return g["f"]
|
||||
+56
@@ -0,0 +1,56 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import copy
|
||||
import random
|
||||
|
||||
import torch
|
||||
from torch.distributed._shard import sharded_tensor
|
||||
from torch.distributed._shard.sharding_spec import ChunkShardingSpec
|
||||
|
||||
|
||||
PLACEMENTS = [
|
||||
"rank:0/cuda:0",
|
||||
"rank:1/cuda:1",
|
||||
"rank:2/cuda:2",
|
||||
"rank:3/cuda:3",
|
||||
]
|
||||
|
||||
DEFAULT_GPU_NUM = 4
|
||||
|
||||
|
||||
def _chunk_sharding_specs_list_for_test(sharding_dims, seed=0):
|
||||
spec_list = []
|
||||
for i in range(len(sharding_dims)):
|
||||
random.Random(seed + i).shuffle(PLACEMENTS)
|
||||
spec_list.append(
|
||||
ChunkShardingSpec(
|
||||
dim=sharding_dims[i],
|
||||
placements=copy.deepcopy(PLACEMENTS),
|
||||
)
|
||||
)
|
||||
return spec_list
|
||||
|
||||
|
||||
class MyShardedModel2(torch.nn.Module):
|
||||
def __init__(self, spec=None, group=None, init_rrefs=True) -> None:
|
||||
super().__init__()
|
||||
if spec is not None:
|
||||
self.sharded_tensor2 = sharded_tensor.rand(
|
||||
spec, 10, 20, process_group=group, init_rrefs=init_rrefs
|
||||
)
|
||||
else:
|
||||
self.sharded_tensor2 = None
|
||||
self.random_tensor2 = torch.nn.Parameter(torch.rand(2, 2))
|
||||
|
||||
|
||||
class MyShardedModel1(torch.nn.Module):
|
||||
def __init__(self, spec=None, group=None, init_rrefs=True) -> None:
|
||||
super().__init__()
|
||||
if spec is not None:
|
||||
self.sharded_tensor1 = sharded_tensor.rand(
|
||||
spec, 10, 20, process_group=group, init_rrefs=init_rrefs
|
||||
)
|
||||
else:
|
||||
self.sharded_tensor1 = None
|
||||
self.random_tensor1 = torch.nn.Parameter(torch.rand(2, 2))
|
||||
self.submodule = MyShardedModel2(spec, group, init_rrefs)
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed._shard.sharded_tensor import ShardedTensor
|
||||
|
||||
|
||||
class SimpleMegatronLM(nn.Module):
|
||||
def __init__(self, linear_size, rank=None, dtype=torch.float32):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(*linear_size[0], dtype=dtype)
|
||||
self.gelu = nn.GELU()
|
||||
self.fc2 = nn.Linear(*linear_size[1], dtype=dtype)
|
||||
if rank is not None:
|
||||
self.fc1.cuda(rank)
|
||||
self.fc2.cuda(rank)
|
||||
|
||||
def forward(self, inp):
|
||||
return self.fc2(self.gelu(self.fc1(inp)))
|
||||
|
||||
def get_weights(self):
|
||||
if isinstance(self.fc1.weight, ShardedTensor):
|
||||
weight1 = self.fc1.weight.local_tensor()
|
||||
else:
|
||||
weight1 = self.fc1.weight
|
||||
|
||||
if isinstance(self.fc2.weight, ShardedTensor):
|
||||
weight2 = self.fc2.weight.local_tensor()
|
||||
else:
|
||||
weight2 = self.fc2.weight
|
||||
|
||||
return (weight1, weight2)
|
||||
|
||||
def get_biases(self):
|
||||
return (self.fc1.bias, self.fc2.bias)
|
||||
|
||||
def get_weight_grads(self):
|
||||
return (self.fc1.weight.grad, self.fc2.weight.grad)
|
||||
|
||||
def get_bias_grads(self):
|
||||
return (self.fc1.bias.grad, self.fc2.bias.grad)
|
||||
+1509
File diff suppressed because it is too large
Load Diff
+196
@@ -0,0 +1,196 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from typing import Any, cast, IO
|
||||
|
||||
# introduced as collections.abc.Buffer in Python 3.12
|
||||
from typing_extensions import Buffer
|
||||
|
||||
import torch.distributed as dist
|
||||
from torch.distributed.checkpoint._extension import (
|
||||
ExtensionRegistry,
|
||||
StreamTransformExtension,
|
||||
)
|
||||
|
||||
|
||||
class Rot13Example(StreamTransformExtension):
|
||||
"""
|
||||
This is an example stream transform extension which just does rot13 on each
|
||||
alphanumeric character of the stream. It is mainly intended as a demonstration
|
||||
and for testing; there isn't a production use case for this.
|
||||
"""
|
||||
|
||||
def __init__(self, chunk_size: int = io.DEFAULT_BUFFER_SIZE) -> None:
|
||||
super().__init__()
|
||||
self._chunk_size = chunk_size
|
||||
|
||||
@staticmethod
|
||||
def from_descriptor(version: str) -> "Rot13Example":
|
||||
if version.partition(".")[0] != "1":
|
||||
raise ValueError(f"Unknown extension {version=}")
|
||||
return Rot13Example()
|
||||
|
||||
@staticmethod
|
||||
def registry_name() -> str:
|
||||
return "stream.rot13"
|
||||
|
||||
def get_descriptor(self) -> str:
|
||||
return f"{self.registry_name()}/1"
|
||||
|
||||
@staticmethod
|
||||
def _rot13bytes(b: Buffer, count: int) -> None:
|
||||
b = memoryview(b)
|
||||
for i in range(count):
|
||||
ch = b[i]
|
||||
if ch >= ord("A") and ch <= ord("Z"):
|
||||
ch += ord("a") - ord("A")
|
||||
elif ch >= ord("a") and ch <= ord("z"):
|
||||
ch += ord("A") - ord("a")
|
||||
b[i] = ch
|
||||
|
||||
def transform_to(self, output: IO[bytes]) -> IO[bytes]:
|
||||
class Writer(io.RawIOBase):
|
||||
def __init__(self, output: IO[bytes]) -> None:
|
||||
self.output = output
|
||||
|
||||
def writeable(self) -> bool:
|
||||
return True
|
||||
|
||||
def write(self, b: Buffer) -> int | None:
|
||||
# Don't mutate the input
|
||||
chunk = bytearray(b)
|
||||
Rot13Example._rot13bytes(chunk, len(chunk))
|
||||
return self.output.write(chunk)
|
||||
|
||||
def flush(self) -> None:
|
||||
self.output.flush()
|
||||
|
||||
return cast(IO[bytes], Writer(output))
|
||||
|
||||
def transform_from(self, input: IO[bytes]) -> IO[bytes]:
|
||||
class Reader(io.RawIOBase):
|
||||
def __init__(self, input: IO[bytes]) -> None:
|
||||
self.input = input
|
||||
|
||||
def readable(self) -> bool:
|
||||
return True
|
||||
|
||||
def readinto(self, b: Buffer) -> int | None:
|
||||
if hasattr(self.input, "readinto"):
|
||||
count = self.input.readinto(b)
|
||||
else:
|
||||
# It's possible self.input is an IO[bytes] with no readinto method.
|
||||
# In that case, we emulate with a read and copy. In practice,
|
||||
# all of the current concrete extensions have readinto.
|
||||
view = memoryview(b)
|
||||
r = self.input.read(len(view))
|
||||
if r is None:
|
||||
count = None
|
||||
else:
|
||||
count = len(r)
|
||||
view[:count] = r
|
||||
if count == 0 or count is None:
|
||||
return count
|
||||
|
||||
Rot13Example._rot13bytes(b, count)
|
||||
return count
|
||||
|
||||
def seekable(self) -> bool:
|
||||
return self.input.seekable()
|
||||
|
||||
def seek(self, offset: int, whence: int = os.SEEK_SET) -> int:
|
||||
return self.input.seek(offset, whence)
|
||||
|
||||
def tell(self) -> int:
|
||||
return self.input.tell()
|
||||
|
||||
return cast(IO[bytes], Reader(input))
|
||||
|
||||
|
||||
def get_test_extension_registry() -> ExtensionRegistry:
|
||||
registry = ExtensionRegistry()
|
||||
registry.register(Rot13Example)
|
||||
return registry
|
||||
|
||||
|
||||
def with_temp_dir(
|
||||
func: Callable | None = None,
|
||||
) -> Callable | None:
|
||||
"""
|
||||
Wrapper to initialize temp directory for distributed checkpoint.
|
||||
"""
|
||||
if func is None:
|
||||
raise AssertionError("Expected func to not be None")
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(self, *args: tuple[object], **kwargs: dict[str, Any]) -> None:
|
||||
if dist.is_initialized():
|
||||
# Only create temp_dir when rank is 0
|
||||
if dist.get_rank() == 0:
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
print(f"Using temp directory: {temp_dir}")
|
||||
else:
|
||||
temp_dir = ""
|
||||
object_list = [temp_dir]
|
||||
|
||||
# Broadcast temp_dir to all the other ranks
|
||||
os.sync()
|
||||
dist.broadcast_object_list(object_list)
|
||||
self.temp_dir = object_list[0]
|
||||
os.sync()
|
||||
else:
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
print(f"No process group initialized, using temp directory: {temp_dir}")
|
||||
self.temp_dir = temp_dir
|
||||
|
||||
try:
|
||||
func(self, *args, **kwargs)
|
||||
finally:
|
||||
if dist.is_initialized() and dist.get_rank() == 0:
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
else:
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def with_checkpoint_logging(
|
||||
func: Callable | None = None,
|
||||
logger_name: str = "torch.distributed.checkpoint",
|
||||
level: int = logging.INFO,
|
||||
) -> Callable | None:
|
||||
"""
|
||||
Wrapper to configure checkpoint logging for distributed tests.
|
||||
|
||||
Args:
|
||||
func: The test function to wrap
|
||||
logger_name: Name of the logger to configure (default: 'torch.distributed.checkpoint')
|
||||
level: Logging level to set (default: logging.INFO)
|
||||
"""
|
||||
if func is None:
|
||||
raise AssertionError("Expected func to not be None")
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(self, *args: tuple[object], **kwargs: dict[str, Any]) -> None:
|
||||
# Get the logger and store original level
|
||||
target_logger = logging.getLogger(logger_name)
|
||||
original_level = target_logger.level
|
||||
|
||||
# Set the desired logging level
|
||||
target_logger.setLevel(level)
|
||||
|
||||
try:
|
||||
func(self, *args, **kwargs)
|
||||
finally:
|
||||
# Restore original logging level
|
||||
target_logger.setLevel(original_level)
|
||||
|
||||
return wrapper
|
||||
+170
@@ -0,0 +1,170 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
# Owner(s): ["oncall: distributed"]
|
||||
|
||||
import copy
|
||||
from itertools import chain
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed._sharded_tensor import ShardedTensor
|
||||
from torch.distributed._state_dict_utils import _gather_state_dict
|
||||
from torch.distributed.checkpoint.state_dict import (
|
||||
_PG,
|
||||
_STATE,
|
||||
set_state_dict,
|
||||
StateDictOptions,
|
||||
)
|
||||
from torch.distributed.tensor import DTensor
|
||||
|
||||
|
||||
class VerifyStateDictMixin:
|
||||
def _compare_tensor(self, orig_tensor, dist_tensor, offload_to_cpu=False):
|
||||
if isinstance(dist_tensor, (DTensor, ShardedTensor)):
|
||||
dist_tensor = _gather_state_dict({"mykey": dist_tensor}).pop("mykey")
|
||||
|
||||
if offload_to_cpu:
|
||||
orig_tensor = orig_tensor.cpu()
|
||||
dist_tensor = dist_tensor.cpu()
|
||||
self.assertTrue(isinstance(dist_tensor, torch.Tensor))
|
||||
self.assertTrue(torch.allclose(orig_tensor, dist_tensor))
|
||||
|
||||
def _verify_msd(
|
||||
self,
|
||||
msd: dict[str, Any],
|
||||
dist_msd: dict[str, Any],
|
||||
options: StateDictOptions = StateDictOptions(),
|
||||
offload_to_cpu=False,
|
||||
) -> None:
|
||||
if not options.ignore_frozen_params:
|
||||
self.assertEqual(len(msd), len(dist_msd))
|
||||
for fqn, param in msd.items():
|
||||
dist_param = dist_msd.get(fqn)
|
||||
if not options.ignore_frozen_params:
|
||||
self.assertIsNotNone(dist_param, f"{fqn=}")
|
||||
try:
|
||||
self._compare_tensor(param, dist_param, offload_to_cpu)
|
||||
except AssertionError as e:
|
||||
raise AssertionError(
|
||||
f"{fqn} has mismatched value {param} {dist_param}"
|
||||
) from e
|
||||
elif dist_param is None:
|
||||
self.assertFalse(param.requires_grad, f"{fqn=}")
|
||||
|
||||
def _verify_osd(
|
||||
self,
|
||||
model: nn.Module,
|
||||
optim: torch.optim.Optimizer,
|
||||
osd: dict[str, Any],
|
||||
dist_osd: dict[str, Any],
|
||||
) -> None:
|
||||
params = list(chain.from_iterable(g["params"] for g in optim.param_groups))
|
||||
param_pid_mapping = dict(zip(params, range(len(params)), strict=True))
|
||||
fqn_pid_mapping = {}
|
||||
for fqn, param in model.named_parameters():
|
||||
pid = param_pid_mapping[param]
|
||||
fqn_pid_mapping[fqn] = pid
|
||||
fqn_pid_mapping[pid] = fqn
|
||||
# Check optimizer_state_dict state
|
||||
|
||||
self.assertEqual(len(osd[_STATE]), len(dist_osd[_STATE]))
|
||||
for pid, states in osd[_STATE].items():
|
||||
fqn = fqn_pid_mapping[pid]
|
||||
dist_states = dist_osd[_STATE].get(fqn, None)
|
||||
self.assertIsNotNone(dist_states, fqn)
|
||||
self.assertEqual(len(states), len(dist_states))
|
||||
for key, state in states.items():
|
||||
dist_state = states.get(key, None)
|
||||
self.assertIsNotNone(dist_state)
|
||||
self._compare_tensor(state, dist_state)
|
||||
|
||||
# Check optimizer_state_dict param_group
|
||||
old_dist_osd_pg = dist_osd[_PG]
|
||||
if len(osd[_PG]) != len(dist_osd[_PG]):
|
||||
self.assertTrue(len(dist_osd[_PG]) > len(osd[_PG]))
|
||||
new_pg = copy.deepcopy(dist_osd[_PG][0])
|
||||
new_pg["params"] = []
|
||||
for dist_group in dist_osd[_PG]:
|
||||
new_pg["params"].extend(dist_group["params"])
|
||||
dist_osd[_PG] = [new_pg]
|
||||
|
||||
self.assertEqual(len(osd[_PG]), len(dist_osd[_PG]))
|
||||
for group, dist_group in zip(osd[_PG], dist_osd[_PG], strict=True):
|
||||
self.assertEqual(len(group), len(dist_group))
|
||||
for key, value in group.items():
|
||||
# Below doesn't work because param_groups can have None
|
||||
# values.
|
||||
# dist_value = dist_group.get(key, None)
|
||||
# self.assertIsNotNone(dist_value, (dist_group, group))
|
||||
dist_value = dist_group[key]
|
||||
if key == "params":
|
||||
fqns = [fqn_pid_mapping[pid] for pid in value]
|
||||
self.assertEqual(sorted(fqns), sorted(dist_value))
|
||||
else:
|
||||
self.assertEqual(value, dist_value)
|
||||
dist_osd[_PG] = old_dist_osd_pg
|
||||
|
||||
def _verify_osd_by_load(
|
||||
self,
|
||||
model: nn.Module,
|
||||
optim: torch.optim.Optimizer,
|
||||
new_optim: torch.optim.Optimizer,
|
||||
dist_osd: dict[str, Any],
|
||||
) -> None:
|
||||
new_dist_osd = _gather_state_dict(dist_osd)
|
||||
set_state_dict(
|
||||
model,
|
||||
optimizers=new_optim,
|
||||
model_state_dict={},
|
||||
optim_state_dict=new_dist_osd,
|
||||
)
|
||||
self.assertEqual(optim.state_dict(), new_optim.state_dict())
|
||||
|
||||
|
||||
class FusionEmbedding(nn.Module):
|
||||
def __init__(self, vocab_size: int, fusion_vocab_size: int, embed_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.embedding = nn.Embedding(vocab_size, embed_dim)
|
||||
self.fusion_embedding = nn.Embedding(fusion_vocab_size, embed_dim)
|
||||
|
||||
|
||||
class FusionEmbeddingWithHook(nn.Module):
|
||||
def __init__(self, vocab_size: int, fusion_vocab_size: int, embed_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.embedding = nn.Embedding(vocab_size, embed_dim)
|
||||
self.fusion_embedding = nn.Embedding(fusion_vocab_size, embed_dim)
|
||||
self._register_state_dict_hook(FusionEmbeddingWithHook._state_dict_hook)
|
||||
self._register_load_state_dict_pre_hook(
|
||||
FusionEmbeddingWithHook._load_state_dict_hook, with_module=True
|
||||
)
|
||||
|
||||
def _state_dict_hook(self, destination, prefix, keep_vars):
|
||||
"""Remove "embedding" from the original embedding in the state_dict
|
||||
name. This keeps the original state dict name for the embedding
|
||||
from before fusing with the FusionEmbedding.
|
||||
"""
|
||||
key = prefix + "embedding.weight"
|
||||
new_key = prefix + "weight"
|
||||
destination[new_key] = destination[key]
|
||||
del destination[key]
|
||||
|
||||
def _load_state_dict_hook(self, state_dict, prefix, *args, **kwargs):
|
||||
"""Apply extra "embedding" prefix to the state_dict key to
|
||||
account for the FusionEmbedding wrapping.
|
||||
"""
|
||||
if state_dict:
|
||||
key = prefix + "weight"
|
||||
new_key = prefix + "embedding.weight"
|
||||
state_dict[new_key] = state_dict[key]
|
||||
del state_dict[key]
|
||||
|
||||
|
||||
class FusionEmbeddingWithModifier(FusionEmbeddingWithHook):
|
||||
# _fqn_modifiers is a private function as a contract between DSD. When users change the state_dict
|
||||
# keys, they need to provide a mapping from the new key to the original key. This is used to ensure
|
||||
# consistency between the state_dict keys and fqn.
|
||||
def _fqn_modifiers(self) -> dict[str, str]:
|
||||
return {
|
||||
"weight": "embedding",
|
||||
}
|
||||
+755
@@ -0,0 +1,755 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import contextlib
|
||||
import enum
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.distributed.autograd as dist_autograd
|
||||
import torch.nn as nn
|
||||
from torch.distributed import rpc
|
||||
from torch.distributed.nn import RemoteModule
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.testing._internal.common_distributed import (
|
||||
requires_gloo,
|
||||
requires_nccl,
|
||||
skip_if_lt_x_gpu,
|
||||
skip_if_rocm_multiprocess,
|
||||
)
|
||||
from torch.testing._internal.dist_utils import dist_init, INIT_METHOD_TEMPLATE
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
NUM_EM_ROW = 2
|
||||
D_SPARSE = 3
|
||||
D_DENSE = 2
|
||||
D_HID = 3
|
||||
D_OUT = 1
|
||||
NUM_TRAINERS = 4
|
||||
# Trainers + the master + the remote worker
|
||||
WORLD_SIZE = NUM_TRAINERS + 2
|
||||
TRAINER_RANKS = list(range(NUM_TRAINERS))
|
||||
REMOTE_WORKER_RANK = TRAINER_RANKS[-1] + 1
|
||||
MASTER_RANK = REMOTE_WORKER_RANK + 1
|
||||
|
||||
|
||||
class DdpMode(enum.Enum):
|
||||
# Don't apply DDP
|
||||
NONE = enum.auto()
|
||||
# Apply DDP to the top level nn.Module
|
||||
OUTSIDE = enum.auto()
|
||||
# Embed DDP inside the top level nn.Module
|
||||
INSIDE = enum.auto()
|
||||
|
||||
|
||||
def init_logger():
|
||||
logger = logging.getLogger(__name__)
|
||||
level = logging.DEBUG if "debug" in os.environ else logging.INFO
|
||||
logger.setLevel(level)
|
||||
console = logging.StreamHandler()
|
||||
formatter = logging.Formatter(
|
||||
"%(asctime)s %(filename)s:%(lineno)s %(levelname)s p:%(processName)s t:%(threadName)s: %(message)s"
|
||||
)
|
||||
console.setFormatter(formatter)
|
||||
console.setLevel(level)
|
||||
# add the handlers to the logger
|
||||
logger.addHandler(console)
|
||||
logger.propagate = False
|
||||
return logger
|
||||
|
||||
|
||||
gLogger = init_logger()
|
||||
|
||||
|
||||
class FeatureSet(NamedTuple):
|
||||
"""A feature set has 2 types of features"""
|
||||
|
||||
dense_features: torch.Tensor
|
||||
sparse_features: torch.LongTensor
|
||||
values: torch.Tensor
|
||||
|
||||
|
||||
def _call_method(method, rref, *args, **kwargs):
|
||||
return method(rref.local_value(), *args, **kwargs)
|
||||
|
||||
|
||||
def _remote_method(method, rref, *args, **kwargs):
|
||||
args_tup = tuple([method, rref] + list(args))
|
||||
return rpc.rpc_sync(rref.owner(), _call_method, args=args_tup, kwargs=kwargs)
|
||||
|
||||
|
||||
def _remote_method_async(method, rref, *args, **kwargs):
|
||||
args_tup = tuple([method, rref] + list(args))
|
||||
return rpc.rpc_async(rref.owner(), _call_method, args=args_tup, kwargs=kwargs)
|
||||
|
||||
|
||||
class RemoteEM(nn.Module):
|
||||
def __init__(self, num_embeddings: int, embedding_dim: int):
|
||||
gLogger.info("Initing RemoteEM with %s %s", num_embeddings, embedding_dim)
|
||||
super().__init__()
|
||||
init_em = [0.5] * embedding_dim
|
||||
self.em = nn.EmbeddingBag(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
_weight=torch.tensor([init_em] * num_embeddings),
|
||||
)
|
||||
|
||||
def forward(self, input: torch.Tensor):
|
||||
gLogger.debug("Running RemoteEM.forward() on: %s", input)
|
||||
return self.em(input, offsets=torch.LongTensor(range(input.shape[0])))
|
||||
|
||||
|
||||
# Return a linear module with predefined parameters.
|
||||
def getLinear(d_in, d_out):
|
||||
l = nn.Linear(d_in, d_out, bias=False)
|
||||
w = torch.ones((d_out, d_in))
|
||||
w[0][0] = -1
|
||||
w.requires_grad_()
|
||||
l.weight.data = w
|
||||
return l
|
||||
|
||||
|
||||
class RemoteNet(nn.Module):
|
||||
def __init__(self, d_in: int, d_out: int):
|
||||
gLogger.info("Initing RemoteNet with %s %s", d_in, d_out)
|
||||
super().__init__()
|
||||
self.fc = getLinear(d_in, d_out)
|
||||
self.relu = nn.ReLU()
|
||||
|
||||
def forward(self, input: torch.Tensor):
|
||||
gLogger.debug("Running RemoteNet.forward() on: %s", input)
|
||||
return self.relu(self.fc(input))
|
||||
|
||||
|
||||
class HybridModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
remote_em_rref: rpc.RRef,
|
||||
remote_net_rref: rpc.RRef,
|
||||
process_group_for_ddp: dist.ProcessGroup = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.remote_em_rref = remote_em_rref
|
||||
self.remote_net_rref = remote_net_rref
|
||||
self.fc1 = getLinear(D_DENSE, D_DENSE)
|
||||
self.fc2 = getLinear(D_HID, D_OUT)
|
||||
|
||||
self.non_ddp_params = tuple(self.fc1.parameters()) + tuple(
|
||||
self.fc2.parameters()
|
||||
)
|
||||
self.ddp_params = ()
|
||||
|
||||
if process_group_for_ddp is not None:
|
||||
self.non_ddp_params, self.ddp_params = (
|
||||
tuple(self.fc1.parameters()),
|
||||
tuple(self.fc2.parameters()),
|
||||
)
|
||||
gLogger.info("Use DDP for the second local net.")
|
||||
self.fc2 = DistributedDataParallel(
|
||||
self.fc2, check_reduction=True, process_group=process_group_for_ddp
|
||||
)
|
||||
|
||||
gLogger.info(
|
||||
"HybridModel has %s groups of parameters.", len(list(self.parameters()))
|
||||
)
|
||||
|
||||
def forward(self, input: FeatureSet):
|
||||
gLogger.debug("Running HybridModel.forward on %s", input)
|
||||
sparse = _remote_method(
|
||||
RemoteEM.forward, self.remote_em_rref, input.sparse_features
|
||||
)
|
||||
# The same size of mini batch.
|
||||
if sparse.shape[0] != input.dense_features.shape[0]:
|
||||
raise AssertionError(
|
||||
f"Expected sparse.shape[0] == input.dense_features.shape[0], "
|
||||
f"got {sparse.shape[0]} != {input.dense_features.shape[0]}"
|
||||
)
|
||||
dense = self.fc1(input.dense_features)
|
||||
x = torch.cat((dense, sparse), 1)
|
||||
gLogger.debug("Concatenated feature: %s", x)
|
||||
x = _remote_method(RemoteNet.forward, self.remote_net_rref, x)
|
||||
return self.fc2(x)
|
||||
|
||||
|
||||
class Trainer:
|
||||
def __init__(
|
||||
self,
|
||||
remote_em_rref: rpc.RRef,
|
||||
remote_net_rref: rpc.RRef,
|
||||
ddp_mode: DdpMode,
|
||||
rank: int,
|
||||
):
|
||||
self.rank = rank
|
||||
self.trainer_group = (
|
||||
dist.new_group(TRAINER_RANKS)
|
||||
if ddp_mode in (DdpMode.INSIDE, DdpMode.OUTSIDE)
|
||||
else None
|
||||
)
|
||||
self.remote_em_rref = remote_em_rref
|
||||
self.remote_net_rref = remote_net_rref
|
||||
self.hybrid_module = HybridModel(
|
||||
self.remote_em_rref,
|
||||
self.remote_net_rref,
|
||||
self.trainer_group if ddp_mode == DdpMode.INSIDE else None,
|
||||
)
|
||||
self.ddp_params, self.non_ddp_params = (
|
||||
self.hybrid_module.ddp_params,
|
||||
self.hybrid_module.non_ddp_params,
|
||||
)
|
||||
if ddp_mode == DdpMode.OUTSIDE:
|
||||
gLogger.info("Wrapping the whole hybrid module into DDP.")
|
||||
self.ddp_params += self.non_ddp_params
|
||||
self.non_ddp_params = ()
|
||||
self.hybrid_module = DistributedDataParallel(
|
||||
self.hybrid_module,
|
||||
check_reduction=True,
|
||||
process_group=self.trainer_group,
|
||||
)
|
||||
gLogger.info(
|
||||
"Succeeded in creating a HybridModel instance with "
|
||||
"%s ddp params and %s other local params.",
|
||||
len(self.ddp_params),
|
||||
len(self.non_ddp_params),
|
||||
)
|
||||
|
||||
def destroy_pg(self):
|
||||
if self.trainer_group:
|
||||
dist.destroy_process_group(self.trainer_group)
|
||||
|
||||
def train_batch(
|
||||
self,
|
||||
mini_batch: FeatureSet,
|
||||
trainer_has_less_inputs: bool,
|
||||
simulate_uneven_inputs: bool,
|
||||
):
|
||||
grads_dict = None
|
||||
|
||||
if not simulate_uneven_inputs:
|
||||
input_batches = [mini_batch]
|
||||
else:
|
||||
# Split into microbatches, and trim to simulate uneven inputs.
|
||||
dense_features = mini_batch.dense_features
|
||||
sparse_features = mini_batch.sparse_features
|
||||
values = mini_batch.values
|
||||
|
||||
dense_microbatch = torch.split(dense_features, 2)
|
||||
sparse_microbatch = torch.split(sparse_features, 2)
|
||||
values_microbatch = torch.split(values, 2)
|
||||
batches = []
|
||||
for d, s, v in zip(
|
||||
dense_microbatch, sparse_microbatch, values_microbatch, strict=True
|
||||
):
|
||||
feature_set = FeatureSet(dense_features=d, sparse_features=s, values=v)
|
||||
batches.append(feature_set)
|
||||
|
||||
if trainer_has_less_inputs:
|
||||
input_batches = batches[: len(batches) // 2]
|
||||
gLogger.info(
|
||||
"Trainer reduced input patches from %s "
|
||||
"to %s to simulate uneven inputs.",
|
||||
len(batches),
|
||||
len(input_batches),
|
||||
)
|
||||
else:
|
||||
input_batches = batches
|
||||
|
||||
with (
|
||||
self.hybrid_module.join()
|
||||
if simulate_uneven_inputs
|
||||
else contextlib.nullcontext()
|
||||
):
|
||||
for b in input_batches:
|
||||
with dist_autograd.context() as context_id:
|
||||
output = self.hybrid_module.forward(b)
|
||||
loss = (output * mini_batch.values).sum()
|
||||
dist_autograd.backward(context_id, [loss])
|
||||
grads_dict = dist_autograd.get_gradients(context_id)
|
||||
gLogger.info(
|
||||
"Loss is %s for mini batch: %s. Grads dict has %s entries: %s",
|
||||
loss,
|
||||
mini_batch,
|
||||
len(grads_dict),
|
||||
grads_dict,
|
||||
)
|
||||
return (
|
||||
tuple(grads_dict[param] for param in self.ddp_params),
|
||||
tuple(grads_dict[param] for param in self.non_ddp_params),
|
||||
)
|
||||
|
||||
|
||||
def get_training_examples():
|
||||
n = 16
|
||||
training_examples = FeatureSet(
|
||||
dense_features=torch.zeros((n, D_DENSE)),
|
||||
sparse_features=torch.zeros(n, dtype=torch.long),
|
||||
values=torch.zeros(n),
|
||||
)
|
||||
idx = 0
|
||||
# Every example has another one that has exactly the same features but an
|
||||
# opposite value. Therefore, their grads cancel each other in all-reduce.
|
||||
for value in (-1, 1):
|
||||
for x in (-1.0 * value, 1.0 * value):
|
||||
for y in (1.0 * value, -1.0 * value):
|
||||
for z in (0, 1):
|
||||
training_examples.dense_features[idx, :] = torch.tensor((x, y))
|
||||
training_examples.sparse_features[idx] = z
|
||||
training_examples.values[idx] = value
|
||||
idx += 1
|
||||
|
||||
# Split the examples among NUM_TRAINERS trainers
|
||||
if n % NUM_TRAINERS != 0:
|
||||
raise AssertionError(
|
||||
f"Expected n % NUM_TRAINERS == 0, got {n} % {NUM_TRAINERS} = {n % NUM_TRAINERS}"
|
||||
)
|
||||
examples_per_trainer = int(n / NUM_TRAINERS)
|
||||
return [
|
||||
FeatureSet(
|
||||
dense_features=training_examples.dense_features[
|
||||
start : start + examples_per_trainer, :
|
||||
],
|
||||
sparse_features=training_examples.sparse_features[
|
||||
start : start + examples_per_trainer
|
||||
],
|
||||
values=training_examples.values[start : start + examples_per_trainer],
|
||||
)
|
||||
for start in range(0, n, examples_per_trainer)
|
||||
]
|
||||
|
||||
|
||||
shutdown_signal = threading.Condition()
|
||||
|
||||
|
||||
def set_shutdown_signal():
|
||||
global shutdown_signal
|
||||
with shutdown_signal:
|
||||
shutdown_signal.notify()
|
||||
|
||||
|
||||
class DdpUnderDistAutogradTest(RpcAgentTestFixture):
|
||||
@property
|
||||
def world_size(self) -> int:
|
||||
return WORLD_SIZE
|
||||
|
||||
def remote_worker_name(self) -> str:
|
||||
# The name has to be consistent with that in 'dist_init' decorator.
|
||||
return f"worker{REMOTE_WORKER_RANK}"
|
||||
|
||||
def trainer_name(self, rank):
|
||||
# The name has to be consistent with that in 'dist_init' decorator.
|
||||
return f"worker{rank}"
|
||||
|
||||
def _remote_worker_process(self, ddp_mode):
|
||||
gLogger.info("The remote worker is running.")
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
|
||||
world_size=self.world_size,
|
||||
rank=self.rank,
|
||||
)
|
||||
|
||||
if ddp_mode in (DdpMode.INSIDE, DdpMode.OUTSIDE):
|
||||
# new_group needs to be called on ranks.
|
||||
dist.new_group(TRAINER_RANKS)
|
||||
|
||||
global shutdown_signal
|
||||
with shutdown_signal:
|
||||
shutdown_signal.wait()
|
||||
gLogger.info("Exiting remote worker.")
|
||||
dist.destroy_process_group()
|
||||
|
||||
def _trainer_process(self, rank: int):
|
||||
gLogger.info("Running the trainer #%s...", rank)
|
||||
gLogger.info(
|
||||
"Initing trainer process group by trainer #%s with ranks %s",
|
||||
rank,
|
||||
TRAINER_RANKS,
|
||||
)
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
|
||||
world_size=self.world_size,
|
||||
rank=self.rank,
|
||||
)
|
||||
|
||||
gLogger.info("Waiting for shutdown signal on trainer #%s...", rank)
|
||||
|
||||
global shutdown_signal
|
||||
with shutdown_signal:
|
||||
shutdown_signal.wait()
|
||||
gLogger.info("Exiting the trainer #%s...", rank)
|
||||
dist.destroy_process_group()
|
||||
|
||||
def _master_process(self, ddp_mode: DdpMode, simulate_uneven_inputs: bool):
|
||||
gLogger.info("Running the master process...")
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
|
||||
world_size=self.world_size,
|
||||
rank=self.rank,
|
||||
)
|
||||
|
||||
remote_em_rref = rpc.remote(
|
||||
self.remote_worker_name(), RemoteEM, args=(NUM_EM_ROW, D_SPARSE)
|
||||
)
|
||||
remote_net_rref = rpc.remote(
|
||||
self.remote_worker_name(), RemoteNet, args=(D_DENSE + D_SPARSE, D_HID)
|
||||
)
|
||||
gLogger.info("Created remote rrefs on master")
|
||||
self.do_test_on_master(
|
||||
ddp_mode, simulate_uneven_inputs, remote_em_rref, remote_net_rref
|
||||
)
|
||||
|
||||
def do_test_on_master(
|
||||
self,
|
||||
ddp_mode: DdpMode,
|
||||
simulate_uneven_inputs: bool,
|
||||
remote_em_rref: rpc.RRef,
|
||||
remote_net_rref: rpc.RRef,
|
||||
):
|
||||
if simulate_uneven_inputs:
|
||||
gLogger.info(
|
||||
"Running DDP + RPC test with simulating uneven inputs across trainers."
|
||||
)
|
||||
|
||||
trainer_rrefs = []
|
||||
for rank in TRAINER_RANKS:
|
||||
trainer = self.trainer_name(rank)
|
||||
trainer_rrefs.append(
|
||||
rpc.remote(
|
||||
trainer,
|
||||
Trainer,
|
||||
args=(remote_em_rref, remote_net_rref, ddp_mode, rank),
|
||||
)
|
||||
)
|
||||
|
||||
if ddp_mode in (DdpMode.INSIDE, DdpMode.OUTSIDE):
|
||||
# new_group needs to be called on ranks.
|
||||
dist.new_group(TRAINER_RANKS)
|
||||
|
||||
training_examples = get_training_examples()
|
||||
for _ in range(3):
|
||||
futures = []
|
||||
num_trainers = len(trainer_rrefs)
|
||||
for idx, trainer_rref in enumerate(trainer_rrefs):
|
||||
# Half the trainers will deplete inputs earlier than the rest.
|
||||
trainer_has_less_inputs = (
|
||||
simulate_uneven_inputs and idx < num_trainers // 2
|
||||
)
|
||||
futures.append(
|
||||
_remote_method_async(
|
||||
Trainer.train_batch,
|
||||
trainer_rref,
|
||||
training_examples[idx],
|
||||
trainer_has_less_inputs,
|
||||
simulate_uneven_inputs,
|
||||
)
|
||||
)
|
||||
|
||||
for future in futures:
|
||||
ddp_grads, non_ddp_grads = future.wait()
|
||||
# When there are uneven inputs, it is not necessary that grads
|
||||
# cancel each other out, since some trainers contribute 0 grad.
|
||||
if not simulate_uneven_inputs:
|
||||
for grad in ddp_grads:
|
||||
self.assertEqual(
|
||||
grad,
|
||||
torch.zeros_like(grad),
|
||||
msg=f"The grad for any ddp parameter should be zeros, because "
|
||||
"the training examples' grads cancel each other. Received "
|
||||
f"gradient {grad}",
|
||||
)
|
||||
for grad in non_ddp_grads:
|
||||
self.assertNotEqual(
|
||||
grad,
|
||||
torch.zeros_like(grad),
|
||||
msg="The grad for any non-ddp parameter shouldn't be zeros",
|
||||
)
|
||||
|
||||
# Destroy process groups
|
||||
for trainer_rref in trainer_rrefs:
|
||||
_remote_method_async(Trainer.destroy_pg, trainer_rref).wait()
|
||||
|
||||
# Send shutdown signals.
|
||||
for rank in TRAINER_RANKS:
|
||||
trainer = self.trainer_name(rank)
|
||||
rpc.rpc_sync(trainer, set_shutdown_signal, args=())
|
||||
|
||||
rpc.rpc_sync(self.remote_worker_name(), set_shutdown_signal, args=())
|
||||
|
||||
def _do_test(self, ddp_mode, simulate_uneven_inputs=False):
|
||||
if self.rank == MASTER_RANK:
|
||||
self._master_process(ddp_mode, simulate_uneven_inputs)
|
||||
elif self.rank == REMOTE_WORKER_RANK:
|
||||
self._remote_worker_process(ddp_mode)
|
||||
elif self.rank in TRAINER_RANKS:
|
||||
self._trainer_process(self.rank)
|
||||
else:
|
||||
raise RuntimeError(f"Unknown process rank: {self.rank}")
|
||||
|
||||
@requires_gloo()
|
||||
@dist_init
|
||||
def test_backward_no_ddp(self):
|
||||
self._do_test(DdpMode.NONE)
|
||||
|
||||
@requires_gloo()
|
||||
@dist_init
|
||||
def test_backward_ddp_outside(self):
|
||||
self._do_test(DdpMode.OUTSIDE)
|
||||
|
||||
@requires_gloo()
|
||||
@dist_init
|
||||
def test_backward_ddp_outside_uneven_inputs(self):
|
||||
self._do_test(DdpMode.OUTSIDE, simulate_uneven_inputs=True)
|
||||
|
||||
@requires_gloo()
|
||||
@dist_init
|
||||
def test_backward_ddp_inside(self):
|
||||
self._do_test(DdpMode.INSIDE)
|
||||
|
||||
|
||||
# Common utils for both CPU and CUDA test suites
|
||||
class CommonDdpComparisonTest(RpcAgentTestFixture):
|
||||
@property
|
||||
def world_size(self) -> int:
|
||||
return NUM_TRAINERS
|
||||
|
||||
def trainer_name(self, rank):
|
||||
# The name has to be consistent with that in 'dist_init' decorator.
|
||||
return f"worker{rank}"
|
||||
|
||||
@staticmethod
|
||||
def get_remote_grads(rref, context_id):
|
||||
return dist_autograd.get_gradients(context_id)[rref.local_value().weight]
|
||||
|
||||
|
||||
class DdpComparisonTest(CommonDdpComparisonTest):
|
||||
def _run_test_ddp_comparision(self, simulate_uneven_inputs=False):
|
||||
gLogger.info("Running trainer rank: %s", self.rank)
|
||||
# Each trainer uses a different random seed. Otherwise, they are going
|
||||
# to have exactly the same initial model parameters, input, and
|
||||
# therefore grads. That means the grads will be the same before and
|
||||
# after DDP's all-reduce.
|
||||
torch.manual_seed(self.rank)
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
# Postfix file_name with "pg" since file_name is also used by RPC agent
|
||||
init_method=INIT_METHOD_TEMPLATE.format(file_name=f"{self.file_name}_pg"),
|
||||
world_size=self.world_size,
|
||||
rank=self.rank,
|
||||
)
|
||||
net = nn.Linear(2, 3)
|
||||
ddp_net = DistributedDataParallel(net)
|
||||
|
||||
# Odd ranks join early if simulate_uneven_inputs.
|
||||
num_inputs = 1
|
||||
if simulate_uneven_inputs:
|
||||
if self.rank % 2 == 0:
|
||||
num_inputs += 2
|
||||
inputs_list = [torch.rand((3, 2)) for _ in range(num_inputs)]
|
||||
|
||||
if simulate_uneven_inputs:
|
||||
gLogger.info(
|
||||
"Rank %s training with %s inputs.", self.rank, len(inputs_list)
|
||||
)
|
||||
|
||||
# Use distributed autograd. The gradients will be in RPC context map.
|
||||
grads_dict = {}
|
||||
with ddp_net.join(simulate_uneven_inputs):
|
||||
for i, inputs in enumerate(inputs_list):
|
||||
with dist_autograd.context() as context_id:
|
||||
loss = ddp_net(inputs).norm()
|
||||
dist_autograd.backward(context_id, [loss])
|
||||
grads_dict = dist_autograd.get_gradients(context_id)
|
||||
gLogger.info("Trainer #%s got grad dict: %s", self.rank, grads_dict)
|
||||
|
||||
# Use local autograd. The gradients will be in each variable's '.grad'.
|
||||
ddp_net.zero_grad()
|
||||
loss = ddp_net(inputs).norm()
|
||||
loss.backward()
|
||||
|
||||
# The gradients should be the same
|
||||
for param in net.parameters():
|
||||
self.assertTrue(
|
||||
param in grads_dict,
|
||||
msg=f"Param {param} is not in dist_auto grad dict {grads_dict} for iteration {i}",
|
||||
)
|
||||
self.assertEqual(
|
||||
grads_dict[param],
|
||||
param.grad,
|
||||
msg=f"The grads for param {param} are different under local "
|
||||
f"and dist autograd: {param.grad} \n---\n {grads_dict[param]} for iteration {i}",
|
||||
)
|
||||
dist.destroy_process_group()
|
||||
|
||||
@requires_gloo()
|
||||
@dist_init
|
||||
def test_ddp_comparison(self):
|
||||
self._run_test_ddp_comparision()
|
||||
|
||||
@requires_gloo()
|
||||
@dist_init
|
||||
def test_ddp_comparison_uneven_inputs(self):
|
||||
# test with simulating uneven inputs in DDP
|
||||
self._run_test_ddp_comparision(simulate_uneven_inputs=True)
|
||||
|
||||
@requires_gloo()
|
||||
@dist_init
|
||||
def test_ddp_dist_autograd_sparse_grads(self):
|
||||
# Each trainer uses a different random seed. Otherwise, they are going
|
||||
# to have exactly the same initial model parameters, input, and
|
||||
# therefore grads. That means the grads will be the same before and
|
||||
# after DDP's all-reduce.
|
||||
torch.manual_seed(self.rank)
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
|
||||
world_size=self.world_size,
|
||||
rank=self.rank,
|
||||
)
|
||||
|
||||
model = nn.EmbeddingBag(10, 3, sparse=True)
|
||||
ddp_model = DistributedDataParallel(model)
|
||||
|
||||
# Different inputs for each
|
||||
input = torch.LongTensor(10).random_(0, 10)
|
||||
offsets = torch.LongTensor([0, 4])
|
||||
|
||||
# Run local.
|
||||
loss = ddp_model(input, offsets).sum()
|
||||
loss.backward()
|
||||
|
||||
with dist_autograd.context() as context_id:
|
||||
loss = ddp_model(input, offsets).sum()
|
||||
dist_autograd.backward(context_id, [loss])
|
||||
grads_dict = dist_autograd.get_gradients(context_id)
|
||||
self.assertEqual(1, len(grads_dict))
|
||||
self.assertEqual(model.weight.grad, grads_dict[model.weight])
|
||||
|
||||
@requires_gloo()
|
||||
@dist_init
|
||||
def test_ddp_dist_autograd_local_vs_remote(self):
|
||||
# Each trainer uses a different random seed. Otherwise, they are going
|
||||
# to have exactly the same initial model parameters, input, and
|
||||
# therefore grads. That means the grads will be the same before and
|
||||
# after DDP's all-reduce.
|
||||
torch.manual_seed(self.rank)
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
|
||||
world_size=self.world_size,
|
||||
rank=self.rank,
|
||||
)
|
||||
|
||||
# Use two different remote device input string, w/ and w/o the default
|
||||
# device string "cpu", respectively.
|
||||
for remote_device in ["worker0/cpu", "worker0"]:
|
||||
remote_layer1 = RemoteModule(
|
||||
remote_device=remote_device, module_cls=nn.Linear, args=(10, 5, False)
|
||||
)
|
||||
layer1 = nn.Linear(10, 5, False)
|
||||
# Start with the same parameters for remote and local
|
||||
layer1.weight = remote_layer1.module_rref.to_here().weight
|
||||
|
||||
# Run local case.
|
||||
layer2 = nn.Linear(5, 1)
|
||||
inputs = torch.rand((10, 10))
|
||||
ddp_model = DistributedDataParallel(layer2)
|
||||
loss = ddp_model(layer1(inputs)).sum()
|
||||
loss.backward()
|
||||
|
||||
# Run remote case.
|
||||
with dist_autograd.context() as context_id:
|
||||
loss = ddp_model(remote_layer1(inputs)).sum()
|
||||
dist_autograd.backward(context_id, [loss])
|
||||
grads_dict = dist_autograd.get_gradients(context_id)
|
||||
dist.barrier()
|
||||
self.assertEqual(layer2.weight.grad, grads_dict[layer2.weight])
|
||||
self.assertEqual(
|
||||
layer1.weight.grad,
|
||||
rpc.rpc_sync(
|
||||
"worker0",
|
||||
CommonDdpComparisonTest.get_remote_grads,
|
||||
args=(remote_layer1.module_rref, context_id),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class CudaDdpComparisonTest(CommonDdpComparisonTest):
|
||||
@skip_if_lt_x_gpu(NUM_TRAINERS)
|
||||
@requires_nccl()
|
||||
@dist_init
|
||||
@skip_if_rocm_multiprocess
|
||||
def test_ddp_dist_autograd_local_vs_remote_gpu(self):
|
||||
# Each trainer uses a different random seed. Otherwise, they are going
|
||||
# to have exactly the same initial model parameters, input, and
|
||||
# therefore grads. That means the grads will be the same before and
|
||||
# after DDP's all-reduce.
|
||||
torch.manual_seed(self.rank)
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
|
||||
world_size=self.world_size,
|
||||
rank=self.rank,
|
||||
)
|
||||
|
||||
remote_layer1 = RemoteModule(
|
||||
remote_device="worker0/cpu", module_cls=nn.Linear, args=(10, 7, False)
|
||||
)
|
||||
layer1 = nn.Linear(10, 7, False)
|
||||
# Start with the same parameters for remote and local
|
||||
layer1.weight = remote_layer1.module_rref.to_here().weight
|
||||
|
||||
layer2 = nn.Linear(7, 5).cuda(self.rank)
|
||||
ddp_layer2 = DistributedDataParallel(layer2, device_ids=[self.rank])
|
||||
|
||||
remote_layer3 = RemoteModule(
|
||||
remote_device="worker0/cpu", module_cls=nn.Linear, args=(5, 3, False)
|
||||
)
|
||||
layer3 = nn.Linear(5, 3, False)
|
||||
# Start with the same parameters for remote and local
|
||||
layer3.weight = remote_layer3.module_rref.to_here().weight
|
||||
|
||||
layer4 = nn.Linear(3, 1).cuda(self.rank)
|
||||
ddp_layer4 = DistributedDataParallel(layer4, device_ids=[self.rank])
|
||||
|
||||
# Run local case.
|
||||
inputs = torch.rand((10, 10))
|
||||
loss = ddp_layer4(
|
||||
layer3(ddp_layer2(layer1(inputs).cuda(self.rank)).cpu()).cuda(self.rank)
|
||||
).sum()
|
||||
loss.backward()
|
||||
|
||||
# Run remote case.
|
||||
with dist_autograd.context() as context_id:
|
||||
loss = ddp_layer4(
|
||||
remote_layer3(
|
||||
ddp_layer2(remote_layer1(inputs).cuda(self.rank)).cpu()
|
||||
).cuda(self.rank)
|
||||
).sum()
|
||||
dist_autograd.backward(context_id, [loss])
|
||||
grads_dict = dist_autograd.get_gradients(context_id)
|
||||
dist.barrier()
|
||||
self.assertEqual(
|
||||
layer1.weight.grad,
|
||||
rpc.rpc_sync(
|
||||
"worker0",
|
||||
CommonDdpComparisonTest.get_remote_grads,
|
||||
args=(remote_layer1.module_rref, context_id),
|
||||
),
|
||||
)
|
||||
self.assertEqual(layer2.weight.grad, grads_dict[layer2.weight])
|
||||
self.assertEqual(
|
||||
layer3.weight.grad,
|
||||
rpc.rpc_sync(
|
||||
"worker0",
|
||||
CommonDdpComparisonTest.get_remote_grads,
|
||||
args=(remote_layer3.module_rref, context_id),
|
||||
),
|
||||
)
|
||||
self.assertEqual(layer4.weight.grad, grads_dict[layer4.weight])
|
||||
+10836
File diff suppressed because it is too large
Load Diff
+68
@@ -0,0 +1,68 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
from contextlib import contextmanager
|
||||
from datetime import timedelta
|
||||
from functools import partial, wraps
|
||||
|
||||
import torch.distributed as dist
|
||||
import torch.distributed.distributed_c10d as c10d
|
||||
|
||||
|
||||
class MockProcessGroup(dist.ProcessGroup):
|
||||
def getBackendName(self):
|
||||
return "mock_process_group"
|
||||
|
||||
|
||||
def create_mock_pg(prefix_store, rank, world_size, timeout):
|
||||
return MockProcessGroup(rank, world_size)
|
||||
|
||||
|
||||
dist.Backend.register_backend("mock_process_group", create_mock_pg)
|
||||
|
||||
|
||||
def mock_init_dist(rank, world_size):
|
||||
# !!! WARNING !!!
|
||||
# Kids don't try this at home, this is a cute pile of hacks that
|
||||
# depends on a small mountain of c10d internals
|
||||
if dist.is_initialized():
|
||||
raise AssertionError("Expected dist to not be initialized")
|
||||
store = dist.HashStore()
|
||||
# Trick _store_based_barrier into believing everyone else already checked-in
|
||||
# Zero is the group index
|
||||
store.add(f"{c10d.STORE_BASED_BARRIER_PREFIX}:0", world_size - 1)
|
||||
dist.init_process_group(
|
||||
backend="mock_process_group",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
store=store,
|
||||
group_name="fake",
|
||||
timeout=timedelta(seconds=1),
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def with_dist(rank=0, world_size=2):
|
||||
"""
|
||||
Context manager that initializer c10d with a fake process group.
|
||||
"""
|
||||
mock_init_dist(rank=rank, world_size=world_size)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def with_fake_comms(func=None, rank=0, world_size=2):
|
||||
"""
|
||||
Function wrapper that inits a fake process group designed for testing.
|
||||
Right now only querying for world size is available
|
||||
"""
|
||||
if func is None:
|
||||
return partial(with_fake_comms, rank=rank, world_size=world_size)
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(self, *args, **kwargs):
|
||||
with with_dist(rank, world_size):
|
||||
func(self, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch.distributed as dist
|
||||
from torch._C._distributed_c10d import FakeProcessGroup
|
||||
|
||||
|
||||
class FakeStore(dist.Store):
|
||||
"""
|
||||
A fake store is a fake Key-Value store simply for initialization usage
|
||||
the of fake process group, one can either use FakeStore or HashStore.
|
||||
"""
|
||||
|
||||
|
||||
def _create_fake_pg(common_opts, backend_opts):
|
||||
"""
|
||||
A fake process group (not related to FakeTensor) is a process group which
|
||||
doesn't actually do any communication, it just hallucinates some
|
||||
communication. You can run a single rank with a fake process group
|
||||
without needing multiple processes (simulates per-rank behavior)
|
||||
|
||||
NOTE: This is not a real process group, and it would produce wrong results
|
||||
for every collective. It should be used as a convenient tool when playing
|
||||
with distributed but don't care about the actual data.
|
||||
"""
|
||||
return FakeProcessGroup._create_internal(
|
||||
common_opts.group_rank, common_opts.group_size, backend_opts
|
||||
)
|
||||
|
||||
|
||||
dist.Backend.register_backend(
|
||||
dist.Backend.FAKE,
|
||||
_create_fake_pg,
|
||||
extended_api=True,
|
||||
devices=["cpu", "cuda", "hpu", "xpu"],
|
||||
)
|
||||
+644
@@ -0,0 +1,644 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import sys
|
||||
import threading
|
||||
import weakref
|
||||
from dataclasses import dataclass
|
||||
from functools import partial, reduce
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch._C._distributed_c10d import (
|
||||
_create_work_from_future,
|
||||
AllgatherOptions,
|
||||
AllreduceOptions,
|
||||
AllToAllOptions,
|
||||
BarrierOptions,
|
||||
BroadcastOptions,
|
||||
ReduceOp,
|
||||
ReduceScatterOptions,
|
||||
ScatterOptions,
|
||||
Store,
|
||||
)
|
||||
from torch.distributed.distributed_c10d import _CollOp, _store_based_barrier, P2POp
|
||||
from torch.futures import Future
|
||||
from torch.utils import _pytree as pytree
|
||||
|
||||
|
||||
"""
|
||||
TODO:
|
||||
Lots of missing collectives.
|
||||
Collectives validation.
|
||||
Make timeout robust by making collectives respect the test deadline.
|
||||
Make tests robust by making collectives interruptible.
|
||||
We need some synchronization around cleanup to ensure that timedout ranks don't cause spurious failures.
|
||||
|
||||
"""
|
||||
|
||||
|
||||
def flatten_list(lst):
|
||||
return pytree.tree_leaves(lst)
|
||||
|
||||
|
||||
def ret_work(ret):
|
||||
fut = Future()
|
||||
fut.set_result(ret)
|
||||
return _create_work_from_future(fut)
|
||||
|
||||
|
||||
def binop_reduce(tensors, op):
|
||||
res = op(torch.stack(tensors), dim=0)
|
||||
if isinstance(res, torch.Tensor):
|
||||
return res
|
||||
# min/max return a namedtuple
|
||||
return res.values
|
||||
|
||||
|
||||
def bitwise_reduce(tensors, op):
|
||||
return reduce(op, tensors)
|
||||
|
||||
|
||||
_reduce_ops = {
|
||||
ReduceOp.SUM: partial(binop_reduce, op=torch.sum),
|
||||
ReduceOp.AVG: partial(binop_reduce, op=torch.mean),
|
||||
ReduceOp.PRODUCT: partial(binop_reduce, op=torch.prod),
|
||||
ReduceOp.MIN: partial(binop_reduce, op=torch.min),
|
||||
ReduceOp.MAX: partial(binop_reduce, op=torch.max),
|
||||
ReduceOp.BAND: partial(bitwise_reduce, op=torch.bitwise_and),
|
||||
ReduceOp.BOR: partial(bitwise_reduce, op=torch.bitwise_or),
|
||||
ReduceOp.BXOR: partial(bitwise_reduce, op=torch.bitwise_xor),
|
||||
}
|
||||
|
||||
|
||||
# Note [Hide collectives mutation from autograd]
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Threaded PG is intended to closely simulate the behavior of regular process
|
||||
# groups. However, our regular PG implementations perform a dispatch through
|
||||
# c10d, whereas Threaded PG does not for some reason (some superficial
|
||||
# but not very convincing reasons include that Threaded PG is implemented
|
||||
# in Python but you can't override Backend in Python, you can only override
|
||||
# ProcessGroup in Python), thereby bypassing the dispatch step. Now we have
|
||||
# a problem: c10d's signatures are LIES, they mutate their (output) tensor
|
||||
# arguments but their annotations don't have mutations on them so we don't
|
||||
# actually update any view metadata if you do differentiation. This
|
||||
# ordinarily "doesn't matter" because distributed collectives aren't
|
||||
# differentiable anyway, but it's possible to tickle this in testing if
|
||||
# someone tries to touch the grad_fn of a Tensor. There a few ways to
|
||||
# fix this, but the easiest way was to use the .detach() trick to hide
|
||||
# the mutations from autograd.
|
||||
|
||||
|
||||
class AllToAll:
|
||||
@torch.no_grad()
|
||||
def work(self, data):
|
||||
world_size = len(data)
|
||||
for dest_rank in range(world_size):
|
||||
output_tensor_list, _ = data[dest_rank]
|
||||
for src_rank in range(world_size):
|
||||
_, input_tensor_list = data[src_rank]
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
output_tensor_list[src_rank].detach().copy_(
|
||||
input_tensor_list[dest_rank]
|
||||
)
|
||||
|
||||
|
||||
class AllToAllBase:
|
||||
@torch.no_grad()
|
||||
def work(self, data):
|
||||
world_size = len(data)
|
||||
for dest_rank in range(world_size):
|
||||
output_buffer, _, output_split_sizes, _ = data[dest_rank]
|
||||
|
||||
output_indexes = self._size_cumsum(
|
||||
output_buffer.size(0), output_split_sizes, world_size
|
||||
)
|
||||
|
||||
for src_rank in range(world_size):
|
||||
_, input_buffer, _, input_split_sizes = data[src_rank]
|
||||
input_indexes = self._size_cumsum(
|
||||
input_buffer.size(0), input_split_sizes, world_size
|
||||
)
|
||||
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
output_buffer[
|
||||
output_indexes[src_rank] : output_indexes[src_rank + 1]
|
||||
].detach().copy_(
|
||||
input_buffer[
|
||||
input_indexes[dest_rank] : input_indexes[dest_rank + 1]
|
||||
]
|
||||
)
|
||||
|
||||
def _size_cumsum(
|
||||
self,
|
||||
buf_size: int,
|
||||
sizes: torch.Tensor | list[int] | None,
|
||||
world_size: int,
|
||||
) -> torch.Tensor:
|
||||
if sizes is None or len(sizes) == 0:
|
||||
sizes = torch.full((world_size,), buf_size // world_size, dtype=torch.int64)
|
||||
if not isinstance(sizes, torch.Tensor):
|
||||
sizes = torch.tensor(sizes, dtype=torch.int64)
|
||||
if sizes.dtype != torch.int64:
|
||||
raise AssertionError(
|
||||
f"Expected sizes.dtype == torch.int64, got {sizes.dtype}"
|
||||
)
|
||||
sizes = torch.cumsum(
|
||||
torch.cat(
|
||||
(torch.tensor([0], dtype=torch.int64, device=sizes.device), sizes),
|
||||
dim=0,
|
||||
),
|
||||
dim=0,
|
||||
)
|
||||
return sizes
|
||||
|
||||
|
||||
class AllReduce:
|
||||
def __init__(self, op):
|
||||
if op.op not in _reduce_ops:
|
||||
raise NotImplementedError(
|
||||
f"AllReduce op {op.op} not supported on multithreaded pg for now."
|
||||
)
|
||||
self.op = op.op
|
||||
|
||||
@torch.no_grad()
|
||||
def work(self, data):
|
||||
for i in range(len(data[0])):
|
||||
# use rank0 as the device for sum
|
||||
rank_0_device = data[0][i].device
|
||||
# collect all data to the list and make them
|
||||
# all on rank 0 device
|
||||
tensors = [
|
||||
data[src_rank][i].to(rank_0_device) for src_rank in range(len(data))
|
||||
]
|
||||
|
||||
# now mimic reduce across all ranks
|
||||
res = _reduce_ops[self.op](tensors)
|
||||
|
||||
# copy all the reduced value to each rank
|
||||
for src_rank in range(len(data)):
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
data[src_rank][i].detach().copy_(res.to(data[src_rank][i].device))
|
||||
|
||||
|
||||
class AllGather:
|
||||
@torch.no_grad()
|
||||
def work(self, data):
|
||||
for src_rank in range(len(data)):
|
||||
in_tensor_list = data[src_rank][1]
|
||||
# Can't handle all_gather with multiple tensors
|
||||
if len(in_tensor_list) != 1:
|
||||
raise AssertionError(
|
||||
f"Can't handle all_gather with multiple tensors, got {len(in_tensor_list)}"
|
||||
)
|
||||
src_tensor = in_tensor_list[0]
|
||||
|
||||
for dest in data:
|
||||
dest_tensor = dest[0][0][src_rank]
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
dest_tensor.detach().copy_(src_tensor)
|
||||
|
||||
|
||||
class Scatter:
|
||||
def __init__(self, src):
|
||||
self.src = src
|
||||
|
||||
@torch.no_grad()
|
||||
def work(self, data):
|
||||
src_in_tensor_list = data[self.src][1]
|
||||
# Can't handle scatter with multiple input tensor list
|
||||
if len(src_in_tensor_list) != 1:
|
||||
raise AssertionError(
|
||||
f"Can't handle scatter with multiple input tensor list, got {len(src_in_tensor_list)}"
|
||||
)
|
||||
src_in_tensors = src_in_tensor_list[0]
|
||||
|
||||
for rank, each_rank_data in enumerate(data):
|
||||
out_tensor_list = each_rank_data[0]
|
||||
# Can't handle scatter with multiple output tensor
|
||||
if len(out_tensor_list) != 1:
|
||||
raise AssertionError(
|
||||
f"Can't handle scatter with multiple output tensor, got {len(out_tensor_list)}"
|
||||
)
|
||||
dest_tensor = out_tensor_list[0]
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
dest_tensor.detach().copy_(src_in_tensors[rank])
|
||||
|
||||
|
||||
class Gather:
|
||||
def __init__(self, dst):
|
||||
self.dst = dst
|
||||
|
||||
@torch.no_grad()
|
||||
def work(self, data):
|
||||
# Can't handle gather with multiple tensor lists
|
||||
if len(data[self.dst][0]) != 1:
|
||||
raise AssertionError(
|
||||
f"Can't handle gather with multiple tensor lists, got {len(data[self.dst][0])}"
|
||||
)
|
||||
out_tensor_list = data[self.dst][0][0]
|
||||
for rank, each_rank_data in enumerate(data):
|
||||
src_in_tensor_list = each_rank_data[1]
|
||||
# Can't handle gather with multiple tensor lists
|
||||
if len(src_in_tensor_list) != 1:
|
||||
raise AssertionError(
|
||||
f"Can't handle gather with multiple tensor lists, got {len(src_in_tensor_list)}"
|
||||
)
|
||||
dest_tensor = out_tensor_list[rank]
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
dest_tensor.detach().copy_(src_in_tensor_list[0])
|
||||
|
||||
|
||||
class ReduceScatter:
|
||||
def __init__(self, op):
|
||||
if op != dist.ReduceOp.SUM and op != dist.ReduceOp.AVG:
|
||||
raise NotImplementedError(f"ReduceScatter does not support {op}")
|
||||
self.op = op
|
||||
|
||||
@torch.no_grad()
|
||||
def work(self, data):
|
||||
start_reduction = [False for _ in range(len(data))]
|
||||
for each_rank_data in data:
|
||||
# Can't handle reduce_scatter with multiple scatter list
|
||||
if len(each_rank_data[1]) != 1:
|
||||
raise AssertionError(
|
||||
f"Can't handle reduce_scatter with multiple scatter list, got {len(each_rank_data[1])}"
|
||||
)
|
||||
to_scatter = each_rank_data[1][0]
|
||||
for i in range(len(to_scatter)):
|
||||
dest_tensor_on_rank_i = data[i][0]
|
||||
# Can't handle reduce_scatter with multiple output tensor
|
||||
if len(dest_tensor_on_rank_i) != 1:
|
||||
raise AssertionError(
|
||||
f"Can't handle reduce_scatter with multiple output tensor, got {len(dest_tensor_on_rank_i)}"
|
||||
)
|
||||
dst_tensor_device = dest_tensor_on_rank_i[0].device
|
||||
if not start_reduction[i]:
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
dest_tensor_on_rank_i[0].detach().copy_(
|
||||
to_scatter[i].to(dst_tensor_device)
|
||||
)
|
||||
start_reduction[i] = True
|
||||
else:
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
dest_tensor_on_rank_i[0].detach().add_(
|
||||
to_scatter[i].to(dst_tensor_device)
|
||||
)
|
||||
if self.op == dist.ReduceOp.AVG:
|
||||
num_ranks = len(data)
|
||||
for each_rank_data in data:
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
each_rank_data[0][0].detach().div_(num_ranks)
|
||||
|
||||
|
||||
class Broadcast:
|
||||
def __init__(self, src):
|
||||
self.src = src
|
||||
|
||||
@torch.no_grad()
|
||||
def work(self, data):
|
||||
in_tensor_list = flatten_list(data[self.src])
|
||||
for i in range(len(data)):
|
||||
if i == self.src:
|
||||
continue
|
||||
out_tensor_list = flatten_list(data[i])
|
||||
for j in range(len(in_tensor_list)):
|
||||
# See Note [Hide collectives mutation from autograd]
|
||||
out_tensor_list[j].detach().copy_(in_tensor_list[j])
|
||||
|
||||
|
||||
class Collective:
|
||||
def __init__(self, world_size, collective, pg):
|
||||
self._world_size = world_size
|
||||
self._collective = collective
|
||||
|
||||
self._start_cond = threading.Condition()
|
||||
self._done_cond = threading.Condition()
|
||||
|
||||
self._data = [None] * world_size
|
||||
self._count = 0
|
||||
self._done = False
|
||||
|
||||
self._pg = pg
|
||||
|
||||
def join(self, rank, data):
|
||||
with self._start_cond:
|
||||
self._data[rank] = data
|
||||
self._count += 1
|
||||
|
||||
# notify rank 0
|
||||
if self._count == self._world_size:
|
||||
if rank > 0:
|
||||
self._start_cond.notify()
|
||||
|
||||
if rank == 0:
|
||||
self._start_cond.wait_for(
|
||||
lambda: self._count == self._world_size
|
||||
or self._pg._terminate.is_set()
|
||||
)
|
||||
# SystemExit is not a subclass of Exception but BaseException
|
||||
# and can be distinguished from normal exception raised from program errors
|
||||
# so that we can hide it from the exception queue
|
||||
if self._pg._terminate.is_set():
|
||||
sys.exit("Test termination event occurs.")
|
||||
|
||||
with self._done_cond:
|
||||
# wait for rank 0 to finish
|
||||
if rank > 0:
|
||||
self._done_cond.wait_for(
|
||||
lambda: self._done or self._pg._terminate.is_set()
|
||||
)
|
||||
if self._pg._terminate.is_set():
|
||||
sys.exit("Test termination event occurs.")
|
||||
else:
|
||||
# copy data around
|
||||
self._collective.work(self._data)
|
||||
self._done = True
|
||||
self._done_cond.notify_all()
|
||||
return ret_work(data)
|
||||
|
||||
|
||||
class ProcessLocalGroup(dist.ProcessGroup):
|
||||
_coll_lock = threading.Lock()
|
||||
_cur_coll_on_pgs = {}
|
||||
|
||||
_terminate = threading.Event()
|
||||
|
||||
@classmethod
|
||||
def _start_coll(cls, collective, pg):
|
||||
with cls._coll_lock:
|
||||
# pg_name is unique, we use that to record the mapping between pg and collective
|
||||
if pg.pg_name not in cls._cur_coll_on_pgs:
|
||||
cls._cur_coll_on_pgs[pg.pg_name] = Collective(
|
||||
pg.size(), collective, cls
|
||||
)
|
||||
return cls._cur_coll_on_pgs[pg.pg_name]
|
||||
|
||||
@classmethod
|
||||
def _end_coll(cls, collective, pg):
|
||||
# This is racily called by all ranks, so only one will work
|
||||
with cls._coll_lock:
|
||||
if (
|
||||
pg.pg_name in cls._cur_coll_on_pgs
|
||||
and cls._cur_coll_on_pgs[pg.pg_name] == collective
|
||||
):
|
||||
cls._cur_coll_on_pgs.pop(pg.pg_name)
|
||||
|
||||
@classmethod
|
||||
def exception_handle(cls, exc):
|
||||
cls._terminate.set()
|
||||
for coll in cls._cur_coll_on_pgs.values():
|
||||
with coll._start_cond:
|
||||
coll._start_cond.notify()
|
||||
with coll._done_cond:
|
||||
coll._done_cond.notify_all()
|
||||
|
||||
@classmethod
|
||||
def reset(cls):
|
||||
with cls._coll_lock:
|
||||
cls._cur_coll_on_pgs = {}
|
||||
cls._terminate.clear()
|
||||
|
||||
def alltoall_base(
|
||||
self,
|
||||
output_buffer: torch.Tensor,
|
||||
input_buffer: torch.Tensor,
|
||||
output_split_sizes: list[int] | None,
|
||||
input_split_sizes: list[int] | None,
|
||||
opts=AllToAllOptions(),
|
||||
) -> torch.Tensor:
|
||||
coll = ProcessLocalGroup._start_coll(AllToAllBase(), self)
|
||||
res = coll.join(
|
||||
self._rank,
|
||||
(output_buffer, input_buffer, output_split_sizes, input_split_sizes),
|
||||
)
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def alltoall(self, output_tensor_list, input_tensor_list, opts=AllToAllOptions()):
|
||||
coll = ProcessLocalGroup._start_coll(AllToAll(), self)
|
||||
res = coll.join(self._rank, (output_tensor_list, input_tensor_list))
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def allreduce(self, tensor_list, opts=AllreduceOptions()):
|
||||
coll = ProcessLocalGroup._start_coll(AllReduce(opts.reduceOp), self)
|
||||
res = coll.join(self._rank, tensor_list)
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def allreduce_coalesced(self, tensor_list, opts=AllreduceOptions()):
|
||||
coll = ProcessLocalGroup._start_coll(AllReduce(opts.reduceOp), self)
|
||||
res = coll.join(self._rank, tensor_list)
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def barrier(self, opts=BarrierOptions()):
|
||||
return self.allreduce(tensor_list=[torch.ones(1)])
|
||||
|
||||
def allgather(self, output_tensors, input_tensor, opts=AllgatherOptions()):
|
||||
coll = ProcessLocalGroup._start_coll(AllGather(), self)
|
||||
res = coll.join(self._rank, (output_tensors, input_tensor))
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def _allgather_base(self, output_tensor, input_tensor, opts=AllgatherOptions()):
|
||||
tensor_list = list(torch.chunk(output_tensor, self._world_size))
|
||||
return self.allgather([tensor_list], [input_tensor], opts)
|
||||
|
||||
def broadcast(self, tensor_list, opts=BroadcastOptions()):
|
||||
coll = ProcessLocalGroup._start_coll(Broadcast(opts.rootRank), self)
|
||||
res = coll.join(self._rank, tensor_list)
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def scatter(self, output_tensors, input_tensors, opts=ScatterOptions()):
|
||||
coll = ProcessLocalGroup._start_coll(Scatter(opts.rootRank), self)
|
||||
res = coll.join(self._rank, (output_tensors, input_tensors))
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def gather(self, output_tensors, input_tensors, opts=ScatterOptions()):
|
||||
coll = ProcessLocalGroup._start_coll(Gather(opts.rootRank), self)
|
||||
res = coll.join(self._rank, (output_tensors, input_tensors))
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def reduce_scatter(self, output_tensor, scatter_list, opts=ReduceScatterOptions()):
|
||||
coll = ProcessLocalGroup._start_coll(ReduceScatter(opts.reduceOp), self)
|
||||
res = coll.join(self._rank, (output_tensor, scatter_list))
|
||||
ProcessLocalGroup._end_coll(coll, self)
|
||||
return res
|
||||
|
||||
def _reduce_scatter_base(
|
||||
self, output_tensor, input_tensor, opts=ReduceScatterOptions()
|
||||
):
|
||||
tensor_list = list(torch.chunk(input_tensor, self._world_size))
|
||||
return self.reduce_scatter([output_tensor], [tensor_list], opts)
|
||||
|
||||
def reduce_scatter_tensor_coalesced(
|
||||
self, output_tensors, input_tensors, opts=ReduceScatterOptions()
|
||||
):
|
||||
works = [
|
||||
self._reduce_scatter_base(output_tensor, input_tensor, opts)
|
||||
for output_tensor, input_tensor in zip(
|
||||
output_tensors, input_tensors, strict=True
|
||||
)
|
||||
]
|
||||
for work in works[:-1]:
|
||||
work.wait()
|
||||
return works[-1]
|
||||
|
||||
def allgather_into_tensor_coalesced(
|
||||
self, output_tensor_list, input_tensor_list, opts=AllgatherOptions()
|
||||
):
|
||||
res = None
|
||||
for o_t, i_t in zip(output_tensor_list, input_tensor_list, strict=True):
|
||||
res = self._allgather_base(o_t, i_t)
|
||||
return res
|
||||
|
||||
def __init__(self, rank, world_size):
|
||||
super().__init__(rank, world_size)
|
||||
self._rank = rank
|
||||
self._world_size = world_size
|
||||
world = dist.distributed_c10d._world
|
||||
if isinstance(world, ThreadLocalWorld):
|
||||
world = world._get_world()
|
||||
self._world = weakref.ref(world)
|
||||
self._ctx = torch.autograd.set_multithreading_enabled(False)
|
||||
|
||||
def size(self):
|
||||
return self._world_size
|
||||
|
||||
@property
|
||||
def pg_name(self):
|
||||
"""
|
||||
return the global registered name of the current pg in the world
|
||||
"""
|
||||
return self._world().pg_names[self]
|
||||
|
||||
@property
|
||||
def group_name(self):
|
||||
return self.pg_name
|
||||
|
||||
def getBackendName(self):
|
||||
return "threaded"
|
||||
|
||||
def __repr__(self):
|
||||
return f"ThreadedPG world_size:{self._world_size} rank:{self._rank}"
|
||||
|
||||
|
||||
def _create_threaded_pg(prefix_store, rank, world_size, timeout):
|
||||
pg = ProcessLocalGroup(rank, world_size)
|
||||
# https://github.com/pytorch/pytorch/pull/103033 changed store based barrier to optional
|
||||
# When device mesh involves sub groups while store based barrier is not enabled in c10d,
|
||||
# even though threaded pg actual collectives are assumed to be single threaded,
|
||||
# different threads may be initializing different groups,
|
||||
# leading to race conditions.
|
||||
# For example, if we have a mesh of [[0, 1], [2, 3]], the sub groups
|
||||
# (dim 0 and 1) would be initialized in different threads independently.
|
||||
# In this case we can no longer rely on class or global variables
|
||||
# but have to rely on store based barrier to make sure each group
|
||||
# is ready separately before we can invoke collectives in any of the groups.
|
||||
|
||||
# the prefix store is already per group so we pass an empty name here
|
||||
_store_based_barrier(rank, prefix_store, "", world_size, timeout)
|
||||
return pg
|
||||
|
||||
|
||||
dist.Backend.register_backend("threaded", _create_threaded_pg, devices=["cpu", "cuda"])
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorldData:
|
||||
default_pg: dist.ProcessGroup
|
||||
pg_map: dict[dist.ProcessGroup, tuple[str, Store | None]]
|
||||
pg_names: dict[dist.ProcessGroup, str]
|
||||
pg_group_ranks: dict[dist.ProcessGroup, dict[int, int]]
|
||||
pg_backend_config: dict[dist.ProcessGroup, str]
|
||||
group_count: int
|
||||
tags_to_pg: dict[str, list[dist.ProcessGroup]]
|
||||
pg_to_tag: dict[dist.ProcessGroup, str]
|
||||
pg_coalesce_state: dict[dist.ProcessGroup, list[_CollOp | P2POp]]
|
||||
comms: list
|
||||
|
||||
|
||||
class ThreadLocalWorld:
|
||||
_world = threading.local()
|
||||
|
||||
def _get_world(self) -> WorldData:
|
||||
if not hasattr(ThreadLocalWorld._world, "world"):
|
||||
ThreadLocalWorld._world.world = WorldData(
|
||||
None, {}, {}, {}, {}, 0, {}, {}, {}, []
|
||||
)
|
||||
return ThreadLocalWorld._world.world
|
||||
|
||||
@property
|
||||
def default_pg(self):
|
||||
return self._get_world().default_pg
|
||||
|
||||
@default_pg.setter
|
||||
def default_pg(self, value):
|
||||
self._get_world().default_pg = value
|
||||
|
||||
@property
|
||||
def pg_map(self):
|
||||
return self._get_world().pg_map
|
||||
|
||||
@property
|
||||
def pg_names(self):
|
||||
return self._get_world().pg_names
|
||||
|
||||
@property
|
||||
def pg_group_ranks(self):
|
||||
return self._get_world().pg_group_ranks
|
||||
|
||||
@property
|
||||
def pg_backend_config(self):
|
||||
return self._get_world().pg_backend_config
|
||||
|
||||
@property
|
||||
def group_count(self) -> int:
|
||||
return self._get_world().group_count
|
||||
|
||||
@group_count.setter
|
||||
def group_count(self, value):
|
||||
self._get_world().group_count = value
|
||||
|
||||
@property
|
||||
def tags_to_pg(self):
|
||||
return self._get_world().tags_to_pg
|
||||
|
||||
@property
|
||||
def pg_to_tag(self):
|
||||
return self._get_world().pg_to_tag
|
||||
|
||||
@property
|
||||
def pg_coalesce_state(self) -> dict[dist.ProcessGroup, list[_CollOp | P2POp]]:
|
||||
return self._get_world().pg_coalesce_state
|
||||
|
||||
@property
|
||||
def comms(self):
|
||||
return self._get_world().comms
|
||||
|
||||
|
||||
_old_pg_world = None
|
||||
_ctx_manager = None
|
||||
|
||||
|
||||
def _install_threaded_pg():
|
||||
global _old_pg_world
|
||||
global _ctx_manager
|
||||
_old_pg_world = dist.distributed_c10d._world
|
||||
dist.distributed_c10d._world = ThreadLocalWorld()
|
||||
_ctx_manager = torch.autograd.set_multithreading_enabled(False)
|
||||
|
||||
return dist.distributed_c10d._world
|
||||
|
||||
|
||||
def _uninstall_threaded_pg():
|
||||
global _ctx_manager
|
||||
dist.distributed_c10d._world = _old_pg_world
|
||||
# Restore autograd multithreading state that was disabled in _install_threaded_pg
|
||||
if _ctx_manager is not None:
|
||||
_ctx_manager.__exit__(None, None, None)
|
||||
_ctx_manager = None
|
||||
+754
@@ -0,0 +1,754 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import enum
|
||||
|
||||
import torch
|
||||
import torch.distributed.rpc as rpc
|
||||
import torch.testing._internal.dist_utils as dist_utils
|
||||
from torch import nn, Tensor
|
||||
from torch._jit_internal import Future
|
||||
from torch.distributed.nn import RemoteModule
|
||||
from torch.distributed.nn.api.remote_module import (
|
||||
_REMOTE_MODULE_PICKLED_ATTRIBUTES,
|
||||
_RemoteModule,
|
||||
)
|
||||
from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
|
||||
from torch.testing._internal.common_utils import TemporaryFileName, TEST_WITH_ROCM
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
_PARAM_VAL = torch.nn.Parameter(torch.ones(1))
|
||||
|
||||
|
||||
# RPC handler for querying the device on the destination worker.
|
||||
def remote_device(module_rref):
|
||||
for param in module_rref.local_value().parameters():
|
||||
return param.device
|
||||
|
||||
|
||||
# RPC handler for querying __dict__ on the destination worker.
|
||||
def remote_module_attributes(remote_module):
|
||||
return remote_module.__dict__
|
||||
|
||||
|
||||
# RPC handler for running forward on the destination worker.
|
||||
def remote_forward(remote_module, args):
|
||||
return remote_module.forward(*args)
|
||||
|
||||
|
||||
# RPC handler for running forward_async on the destination worker.
|
||||
def remote_forward_async(remote_module, args):
|
||||
# Since future cannot be pickled and sent over the RPC layer,
|
||||
# have to wait and behave just like ``forward_sync``.
|
||||
return remote_module.forward_async(*args).wait()
|
||||
|
||||
|
||||
# RPC handler for getting training mode on the destination worker.
|
||||
def get_remote_training_arg(module_rref):
|
||||
return module_rref.local_value().training
|
||||
|
||||
|
||||
class ModuleCreationMode(enum.Enum):
|
||||
MODULE_CTOR_WITH_INTERFACE = "module_ctor_with_interface"
|
||||
MODULE_CTOR = "module_ctor"
|
||||
|
||||
|
||||
@torch.jit.interface
|
||||
class MyModuleInterface:
|
||||
def forward(
|
||||
self, tensor: Tensor, number: int, word: str = "default"
|
||||
) -> tuple[str, int, Tensor]:
|
||||
# pyre-ignore[7]: Pyre and torch.jit.interface don't mix well
|
||||
pass
|
||||
|
||||
|
||||
@torch.jit.interface
|
||||
class RemoteMyModuleInterface:
|
||||
def forward(
|
||||
self, tensor: Tensor, number: int, word: str = "default"
|
||||
) -> tuple[str, int, Tensor]:
|
||||
# pyre-ignore[7]: Pyre and torch.jit.interface don't mix well
|
||||
pass
|
||||
|
||||
def forward_async(
|
||||
self, tensor: Tensor, number: int, word: str = "default"
|
||||
) -> Future[tuple[str, int, Tensor]]:
|
||||
pass
|
||||
|
||||
|
||||
class MyModule(nn.Module):
|
||||
def __init__(self, first_arg, first_kwarg=-1):
|
||||
super().__init__()
|
||||
self.param1 = _PARAM_VAL
|
||||
|
||||
def forward(
|
||||
self, tensor: Tensor, number: int, word: str = "default"
|
||||
) -> tuple[str, int, Tensor]:
|
||||
return word, number, tensor
|
||||
|
||||
|
||||
class BadModule:
|
||||
def __init__(self, first_arg, first_kwarg=-1):
|
||||
pass
|
||||
|
||||
|
||||
def create_scripted_module(first_arg, first_kwarg=-1):
|
||||
module = MyModule(first_arg, first_kwarg=first_kwarg)
|
||||
scripted_module = torch.jit.script(module)
|
||||
return scripted_module
|
||||
|
||||
|
||||
# Common utils for both CPU and CUDA test suites
|
||||
class CommonRemoteModuleTest(RpcAgentTestFixture):
|
||||
@property
|
||||
def world_size(self): # Override setting in RpcAgentTestFixture
|
||||
return 2
|
||||
|
||||
@staticmethod
|
||||
def _create_remote_module_iter(remote_device, modes=None):
|
||||
if modes is None:
|
||||
modes = ModuleCreationMode.__members__.values()
|
||||
|
||||
args = (1,)
|
||||
kwargs = dict(first_kwarg=2)
|
||||
|
||||
if ModuleCreationMode.MODULE_CTOR in modes:
|
||||
remote_module = RemoteModule(remote_device, MyModule, args, kwargs)
|
||||
yield remote_module
|
||||
|
||||
if ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE in modes:
|
||||
remote_module = _RemoteModule(
|
||||
remote_device,
|
||||
create_scripted_module,
|
||||
args,
|
||||
kwargs,
|
||||
_module_interface_cls=MyModuleInterface,
|
||||
)
|
||||
scripted_remote_module = torch.jit.script(remote_module)
|
||||
yield scripted_remote_module
|
||||
|
||||
|
||||
class RemoteModuleTest(CommonRemoteModuleTest):
|
||||
@dist_utils.dist_init
|
||||
def test_bad_module(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
remote_device = f"{dst_worker_name}/cpu"
|
||||
args = (1,)
|
||||
kwargs = dict(first_kwarg=2)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Expect `module_cls\(\*args, \*\*kwargs\)` returns an instance of <class nn.Module>,",
|
||||
):
|
||||
RemoteModule(remote_device, BadModule, args, kwargs).forward()
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Expect `module_cls\(\*args, \*\*kwargs\)` returns an instance of <class nn.Module>,",
|
||||
):
|
||||
RemoteModule(remote_device, BadModule, args, kwargs).forward()
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_forward_async(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
args = (torch.ones(1), 2, "3")
|
||||
for remote_module in self._create_remote_module_iter(dst_worker_name):
|
||||
ret_fut = remote_module.forward_async(*args)
|
||||
ret = ret_fut.wait()
|
||||
self.assertEqual(ret, tuple(reversed(args)))
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_forward_async_script(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
scripted_remote_module = next(
|
||||
self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE]
|
||||
)
|
||||
)
|
||||
|
||||
@torch.jit.script
|
||||
def run_forward_async(scripted_remote_module: RemoteMyModuleInterface):
|
||||
ret_fut = scripted_remote_module.forward_async(torch.ones(1), 2, "3")
|
||||
ret = ret_fut.wait()
|
||||
return ret
|
||||
|
||||
ret = run_forward_async(scripted_remote_module)
|
||||
|
||||
self.assertEqual(ret, ("3", 2, torch.ones(1)))
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_forward_sync(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
args = (torch.ones(1), 2, "3")
|
||||
for remote_module in self._create_remote_module_iter(dst_worker_name):
|
||||
ret = remote_module.forward(*args)
|
||||
self.assertEqual(ret, tuple(reversed(args)))
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_forward_sync_script(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
scripted_remote_module = next(
|
||||
self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE]
|
||||
)
|
||||
)
|
||||
|
||||
@torch.jit.script
|
||||
def run_forward(scripted_remote_module: MyModuleInterface):
|
||||
ret = scripted_remote_module.forward(torch.ones(1), 2, "3")
|
||||
return ret
|
||||
|
||||
ret = run_forward(scripted_remote_module)
|
||||
|
||||
self.assertEqual(ret, ("3", 2, torch.ones(1)))
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_forward_with_kwargs(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
args = (torch.ones(1), 2)
|
||||
kwargs = dict(word="3")
|
||||
# Only test Python nn.Module, because script module methods don't support taking kwargs.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
ret_fut = remote_module.forward_async(*args, **kwargs)
|
||||
ret = ret_fut.wait()
|
||||
self.assertEqual(ret, tuple(reversed(args + ("3",))))
|
||||
|
||||
ret = remote_module.forward(*args, **kwargs)
|
||||
self.assertEqual(ret, tuple(reversed(args + ("3",))))
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_remote_parameters(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
# Only test Python nn.Module, because script module methods don't support ``remote_parameters``.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
param_rrefs = remote_module.remote_parameters()
|
||||
self.assertEqual(len(param_rrefs), 1)
|
||||
self.assertTrue(torch.equal(param_rrefs[0].to_here(), _PARAM_VAL))
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_get_module_rref(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
# Only test Python nn.Module, because script module methods don't support ``get_module_rref``.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
rref = remote_module.get_module_rref()
|
||||
self.assertEqual(rref, remote_module.module_rref)
|
||||
for param in rref.to_here().parameters():
|
||||
self.assertTrue(torch.equal(param, _PARAM_VAL))
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_train_eval(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
remote_module.train()
|
||||
ret1 = rpc.rpc_sync(
|
||||
dst_worker_name,
|
||||
get_remote_training_arg,
|
||||
args=(remote_module.get_module_rref(),),
|
||||
)
|
||||
self.assertEqual(ret1, True)
|
||||
|
||||
remote_module.eval()
|
||||
ret2 = rpc.rpc_sync(
|
||||
dst_worker_name,
|
||||
get_remote_training_arg,
|
||||
args=(remote_module.get_module_rref(),),
|
||||
)
|
||||
self.assertEqual(ret2, False)
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_unsupported_methods(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``register_buffer`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.register_buffer("buffer", torch.ones(5))
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Method ``register_parameter`` not supported for RemoteModule",
|
||||
):
|
||||
remote_module.register_parameter(
|
||||
"param", torch.nn.Parameter(torch.ones(1))
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``add_module`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.add_module("empty", None)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``apply`` not supported for RemoteModule"
|
||||
):
|
||||
fn = torch.rand((3, 3), requires_grad=False)
|
||||
remote_module.apply(fn)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``cuda`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.cuda()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``cpu`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.cpu()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``type`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.type(torch.FloatTensor)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``float`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.float()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``double`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.double()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``bfloat16`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.bfloat16()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``to`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.to("cpu", dtype=torch.int32)
|
||||
|
||||
def hook(module, grad_input, grad_output):
|
||||
pass
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Method ``register_backward_hook`` not supported for RemoteModule",
|
||||
):
|
||||
remote_module.register_backward_hook(hook)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Method ``register_forward_pre_hook`` not supported for RemoteModule",
|
||||
):
|
||||
remote_module.register_forward_pre_hook(hook)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Method ``register_forward_hook`` not supported for RemoteModule",
|
||||
):
|
||||
remote_module.register_forward_hook(hook)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``state_dict`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.state_dict()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``load_state_dict`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.load_state_dict({})
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Method ``parameters`` not supported for RemoteModule. Please use ``remote_parameters`` instead.",
|
||||
):
|
||||
remote_module.parameters()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Method ``named_parameters`` not supported for RemoteModule",
|
||||
):
|
||||
remote_module.named_parameters()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``buffers`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.buffers()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``named_buffers`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.named_buffers()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``children`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.children()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``named_children`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.named_children()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``modules`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.modules()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``named_modules`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.named_modules()
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``requires_grad_`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.requires_grad_()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``zero_grad`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.zero_grad()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``share_memory`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.share_memory()
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r"Method ``extra_repr`` not supported for RemoteModule"
|
||||
):
|
||||
remote_module.extra_repr()
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_send_remote_module_with_a_new_attribute_not_pickled_over_the_wire(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
# If a new attribute is added to this RemoteModule after the initialization,
|
||||
# and it will be sent over the wire by RPC,
|
||||
# this new field will not be pickled, because it's not specified in _REMOTE_MODULE_PICKLED_ATTRIBUTES.
|
||||
# Note that adding a new attribute out of constructor should rarely happen.
|
||||
# If a new attribute is added to RemoteModule constructor,
|
||||
# there is a sanity check to enforce developers to add this attribute to either
|
||||
# _REMOTE_MODULE_PICKLED_ATTRIBUTES or _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
new_attr_name = "new_attr"
|
||||
setattr(remote_module, new_attr_name, 1)
|
||||
|
||||
attrs = rpc.rpc_sync(
|
||||
dst_worker_name, remote_module_attributes, (remote_module,)
|
||||
)
|
||||
self.assertNotIn(new_attr_name, attrs)
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_remote_module_py_pickle_not_supported(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
with TemporaryFileName() as fname:
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"Cannot pickle RemoteModule in python pickler. RemoteModule can only be pickled when using RPC",
|
||||
):
|
||||
torch.save(remote_module, fname)
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_remote_module_py_pickle_not_supported_script(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE]
|
||||
):
|
||||
with (
|
||||
TemporaryFileName() as fname,
|
||||
self.assertRaisesRegex(
|
||||
torch.jit.Error, "can only be pickled when using RPC"
|
||||
),
|
||||
):
|
||||
torch.save(remote_module, fname)
|
||||
|
||||
|
||||
class ThreeWorkersRemoteModuleTest(CommonRemoteModuleTest):
|
||||
@property
|
||||
def world_size(self): # Override setting in CommonRemoteModuleTest
|
||||
return 3
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_send_remote_module_over_the_wire(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size)
|
||||
|
||||
# Unpickled attributes include both the inherent attributes of RemoteModule
|
||||
# (not inherited from the superclass) and two installed methods.
|
||||
expected_unpickled_attrs = list(_REMOTE_MODULE_PICKLED_ATTRIBUTES)
|
||||
expected_unpickled_attrs.append("forward_async")
|
||||
expected_unpickled_attrs.append("forward")
|
||||
|
||||
# Create a remote module on worker1 and then pass it to worker2 over the RPC layer.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
# Test querying some simple attributes from worker2.
|
||||
attrs = rpc.rpc_sync(
|
||||
dst_worker2_name, remote_module_attributes, (remote_module,)
|
||||
)
|
||||
self.assertListEqual(list(attrs.keys()), expected_unpickled_attrs)
|
||||
self.assertEqual(attrs["on"], "worker1")
|
||||
self.assertEqual(attrs["device"], "cpu")
|
||||
self.assertFalse(attrs["is_device_map_set"])
|
||||
self.assertFalse(attrs["is_scriptable"])
|
||||
|
||||
# Test the installed methods on worker1's can be initiated by worker2 over RPC layer.
|
||||
# NOTE: In practice a remote module should be directly stored on the worker that runs ``forward``` or ``forward_async``,
|
||||
# not have another worker to initiate forward over the RPC layer.
|
||||
args = (torch.ones(1), 2, "3")
|
||||
ret1 = rpc.rpc_sync(dst_worker2_name, remote_forward, (remote_module, args))
|
||||
self.assertEqual(ret1, tuple(reversed(args)))
|
||||
ret2 = rpc.rpc_sync(
|
||||
dst_worker2_name, remote_forward_async, (remote_module, args)
|
||||
)
|
||||
self.assertEqual(ret2, tuple(reversed(args)))
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_send_remote_module_over_the_wire_script_not_supported(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size)
|
||||
|
||||
# Unpickled attributes include both the inherent attributes of RemoteModule
|
||||
# (not inherited from the superclass) and two installed methods.
|
||||
expected_unpickled_attrs = list(_REMOTE_MODULE_PICKLED_ATTRIBUTES)
|
||||
expected_unpickled_attrs.append("forward_async")
|
||||
expected_unpickled_attrs.append("forward")
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError, "Passing a script RemoteModule over RPC is not supported."
|
||||
):
|
||||
# Create a remote module on worker1 and then pass it to worker2 over the RPC layer.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE]
|
||||
):
|
||||
# Test querying some simple attributes from worker2.
|
||||
rpc.rpc_sync(
|
||||
dst_worker2_name, remote_module_attributes, (remote_module,)
|
||||
)
|
||||
|
||||
@dist_utils.dist_init
|
||||
def test_create_remote_module_from_module_rref(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size)
|
||||
|
||||
# Create a remote module on worker1 and then pass its `module_rref` to worker2 over the RPC layer.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
remote_module2 = rpc.rpc_sync(
|
||||
dst_worker2_name,
|
||||
RemoteModule.init_from_module_rref,
|
||||
(dst_worker2_name, remote_module.get_module_rref()),
|
||||
)
|
||||
|
||||
args = (torch.ones(1), 2, "3")
|
||||
ret1 = rpc.rpc_sync(dst_worker1_name, remote_forward, (remote_module, args))
|
||||
ret2 = rpc.rpc_sync(
|
||||
dst_worker2_name, remote_forward, (remote_module2, args)
|
||||
)
|
||||
self.assertEqual(ret1, ret2)
|
||||
|
||||
|
||||
class CudaRemoteModuleTest(CommonRemoteModuleTest):
|
||||
@skip_if_lt_x_gpu(1)
|
||||
@dist_utils.dist_init
|
||||
def test_valid_device(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_rank = (self.rank + 1) % self.world_size
|
||||
dst_worker_name = dist_utils.worker_name(dst_rank)
|
||||
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
f"{dst_worker_name}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
device = rpc.rpc_sync(
|
||||
dst_worker_name, remote_device, (remote_module.module_rref,)
|
||||
)
|
||||
self.assertEqual(device.type, "cuda")
|
||||
self.assertEqual(device.index, 0)
|
||||
|
||||
# Test rank works as well.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
f"rank:{dst_rank}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
device = rpc.rpc_sync(
|
||||
dst_worker_name, remote_device, (remote_module.module_rref,)
|
||||
)
|
||||
self.assertEqual(device.type, "cuda")
|
||||
self.assertEqual(device.index, 0)
|
||||
|
||||
@skip_if_lt_x_gpu(1)
|
||||
@dist_utils.dist_init
|
||||
def test_invalid_devices(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
r"Expected one of .+ device type at start of device string",
|
||||
):
|
||||
[
|
||||
m.forward()
|
||||
for m in self._create_remote_module_iter(
|
||||
f"{dst_worker_name}/foo",
|
||||
modes=[ModuleCreationMode.MODULE_CTOR],
|
||||
)
|
||||
]
|
||||
|
||||
if TEST_WITH_ROCM:
|
||||
errorString = (
|
||||
r"HIP error: invalid device ordinal\n"
|
||||
r"HIP kernel errors might be asynchronously reported at some other API call, "
|
||||
r"so the stacktrace below might be incorrect.\n"
|
||||
r"For debugging consider passing AMD_SERIALIZE_KERNEL=3"
|
||||
)
|
||||
else:
|
||||
errorString = r"CUDA error: invalid device ordinal"
|
||||
with self.assertRaisesRegex(RuntimeError, errorString):
|
||||
[
|
||||
m.forward()
|
||||
for m in self._create_remote_module_iter(
|
||||
f"{dst_worker_name}/cuda:100",
|
||||
modes=[ModuleCreationMode.MODULE_CTOR],
|
||||
)
|
||||
]
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, r"Invalid device string: 'cpu2'"):
|
||||
[
|
||||
m.forward()
|
||||
for m in self._create_remote_module_iter(
|
||||
f"{dst_worker_name}/cpu2",
|
||||
modes=[ModuleCreationMode.MODULE_CTOR],
|
||||
)
|
||||
]
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, r"Device string must not be empty"):
|
||||
[
|
||||
m.forward()
|
||||
for m in self._create_remote_module_iter(
|
||||
f"{dst_worker_name}/",
|
||||
modes=[ModuleCreationMode.MODULE_CTOR],
|
||||
)
|
||||
]
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Could not parse remote_device: worker1/cuda:0/cuda:1. The valid format is '<workername>/<device>'",
|
||||
):
|
||||
[
|
||||
m.forward()
|
||||
for m in self._create_remote_module_iter(
|
||||
f"{dst_worker_name}/cuda:0/cuda:1",
|
||||
modes=[ModuleCreationMode.MODULE_CTOR],
|
||||
)
|
||||
]
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Could not parse remote_device: /. The valid format is '<workername>/<device>'",
|
||||
):
|
||||
[
|
||||
m.forward()
|
||||
for m in self._create_remote_module_iter(
|
||||
"/",
|
||||
modes=[ModuleCreationMode.MODULE_CTOR],
|
||||
)
|
||||
]
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
r"Could not parse remote_device: /cuda:0. The valid format is '<workername>/<device>'",
|
||||
):
|
||||
[
|
||||
m.forward()
|
||||
for m in self._create_remote_module_iter(
|
||||
"/cuda:0",
|
||||
modes=[ModuleCreationMode.MODULE_CTOR],
|
||||
)
|
||||
]
|
||||
|
||||
@skip_if_lt_x_gpu(1)
|
||||
@dist_utils.dist_init
|
||||
def test_input_moved_to_cuda_device(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
# These two CPU tensors (in args and kwargs) should be implicitly moved to an appropriate cuda device.
|
||||
t1 = torch.ones(1)
|
||||
args = (t1, 2)
|
||||
t2 = t1 * 2
|
||||
kwargs = dict(word=t2)
|
||||
|
||||
# Only test Python nn.Module, because script module methods don't support taking kwargs.
|
||||
for remote_module in self._create_remote_module_iter(
|
||||
f"{dst_worker_name}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR]
|
||||
):
|
||||
ret_fut = remote_module.forward_async(*args, **kwargs)
|
||||
ret = ret_fut.wait()
|
||||
self.assertEqual(ret, tuple(reversed(args + (t2,))))
|
||||
# TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0".
|
||||
self.assertEqual(ret[0].device.type, "cpu")
|
||||
self.assertEqual(ret[2].device.type, "cpu")
|
||||
|
||||
ret = remote_module.forward(*args, **kwargs)
|
||||
self.assertEqual(ret, tuple(reversed(args + (t2,))))
|
||||
# TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0".
|
||||
self.assertEqual(ret[0].device.type, "cpu")
|
||||
self.assertEqual(ret[2].device.type, "cpu")
|
||||
|
||||
@skip_if_lt_x_gpu(1)
|
||||
@dist_utils.dist_init
|
||||
def test_input_moved_to_cuda_device_script(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
scripted_remote_module = next(
|
||||
self._create_remote_module_iter(
|
||||
f"{dst_worker_name}/cuda:0",
|
||||
modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE],
|
||||
)
|
||||
)
|
||||
|
||||
@torch.jit.script
|
||||
def run_forward(scripted_remote_module: MyModuleInterface):
|
||||
ret = scripted_remote_module.forward(torch.ones(1), 2, "3")
|
||||
return ret
|
||||
|
||||
ret = run_forward(scripted_remote_module)
|
||||
|
||||
self.assertEqual(ret, ("3", 2, torch.ones(1)))
|
||||
# TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0".
|
||||
self.assertEqual(ret[2].device.type, "cpu")
|
||||
+2767
File diff suppressed because it is too large
Load Diff
+281
@@ -0,0 +1,281 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
|
||||
import threading
|
||||
|
||||
import torch
|
||||
import torch.distributed.autograd as dist_autograd
|
||||
import torch.distributed.rpc as rpc
|
||||
from torch import optim
|
||||
from torch.distributed.optim import DistributedOptimizer
|
||||
from torch.testing._internal.dist_utils import dist_init
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
class MyModule:
|
||||
lock = threading.Lock()
|
||||
|
||||
def __init__(self, requires_grad=True):
|
||||
# cannot directly use torch.manual_seed(0) as all threads share the same
|
||||
# default generator. The race from multiple RPC threads could mess up
|
||||
# the draw order from the default RNG instance, leading to
|
||||
# non-deterministic behavior. Hence, create a dedicated RNG here.
|
||||
g_cpu = torch.Generator()
|
||||
g_cpu.manual_seed(0)
|
||||
self.w = torch.rand((3, 3), requires_grad=requires_grad, generator=g_cpu)
|
||||
|
||||
def forward(self, t1):
|
||||
return torch.mm(self.w, t1)
|
||||
|
||||
def get_w(self):
|
||||
return self.w
|
||||
|
||||
|
||||
class FailingOptimizer(optim.Optimizer):
|
||||
def __init__(self, params):
|
||||
super().__init__(params, {})
|
||||
|
||||
def step(self, closure=None):
|
||||
raise ValueError("Error running optimizer.")
|
||||
|
||||
|
||||
class OptimizerFailingOnConstructor(optim.Optimizer):
|
||||
def __init__(self, params):
|
||||
super().__init__(params, {})
|
||||
raise ValueError("Error creating optimizer.")
|
||||
|
||||
def step(self, closure=None):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def _call_method(method, obj_rref, *args, **kwargs):
|
||||
return method(obj_rref.local_value(), *args, **kwargs)
|
||||
|
||||
|
||||
def remote_method(method, obj_rref, *args, **kwargs):
|
||||
"""
|
||||
Call rpc.remote on a method in a remote object.
|
||||
|
||||
Args:
|
||||
method: the method (for example, Class.method)
|
||||
obj_rref (RRef): remote reference to the object
|
||||
args: positional arguments to pass to the method
|
||||
kwargs: keyword arguments to pass to the method
|
||||
|
||||
Returns a RRef to the remote method call result.
|
||||
"""
|
||||
return rpc.remote(
|
||||
obj_rref.owner(),
|
||||
_call_method,
|
||||
args=[method, obj_rref] + list(args),
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
||||
def rpc_async_method(method, obj_rref, *args, **kwargs):
|
||||
"""
|
||||
Call rpc.rpc_async on a method in a remote object.
|
||||
|
||||
Args:
|
||||
method: the method (for example, Class.method)
|
||||
obj_rref (RRef): remote reference to the object
|
||||
args: positional arguments to pass to the method
|
||||
kwargs: keyword arguments to pass to the method
|
||||
|
||||
Returns a Future to the method call result.
|
||||
"""
|
||||
return rpc.rpc_async(
|
||||
obj_rref.owner(),
|
||||
_call_method,
|
||||
args=[method, obj_rref] + list(args),
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
||||
class DistOptimizerTest(RpcAgentTestFixture):
|
||||
@dist_init()
|
||||
def test_dist_optim_exception(self):
|
||||
# distributed version
|
||||
owner1 = f"worker{(self.rank + 1) % self.world_size:d}"
|
||||
owner2 = f"worker{(self.rank + 2) % self.world_size:d}"
|
||||
|
||||
remote_module1 = rpc.remote(owner1, MyModule)
|
||||
remote_module2 = rpc.remote(owner2, MyModule)
|
||||
remote_param1 = remote_method(MyModule.get_w, remote_module1)
|
||||
remote_param2 = remote_method(MyModule.get_w, remote_module2)
|
||||
|
||||
dist_optim = DistributedOptimizer(
|
||||
FailingOptimizer, [remote_param1, remote_param2]
|
||||
)
|
||||
|
||||
with dist_autograd.context() as context_id:
|
||||
g_cpu = torch.Generator()
|
||||
g_cpu.manual_seed(0)
|
||||
t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
output1 = rpc_async_method(MyModule.forward, remote_module1, t2)
|
||||
output2 = rpc_async_method(MyModule.forward, remote_module2, output1.wait())
|
||||
loss = torch.add(output2.wait(), t1).sum()
|
||||
|
||||
dist_autograd.backward(context_id, [loss])
|
||||
with self.assertRaisesRegex(Exception, "Error running optimizer"):
|
||||
dist_optim.step(context_id)
|
||||
|
||||
@dist_init()
|
||||
def test_dist_optim_exception_on_constructor(self):
|
||||
# distributed version
|
||||
owner1 = f"worker{(self.rank + 1) % self.world_size:d}"
|
||||
owner2 = f"worker{(self.rank + 2) % self.world_size:d}"
|
||||
|
||||
remote_module1 = rpc.remote(owner1, MyModule)
|
||||
remote_module2 = rpc.remote(owner2, MyModule)
|
||||
remote_param1 = remote_method(MyModule.get_w, remote_module1)
|
||||
remote_param2 = remote_method(MyModule.get_w, remote_module2)
|
||||
|
||||
with self.assertRaisesRegex(Exception, "Error creating optimizer."):
|
||||
DistributedOptimizer(
|
||||
OptimizerFailingOnConstructor, [remote_param1, remote_param2]
|
||||
)
|
||||
|
||||
def _test_dist_optim_base(self, optim_cls, *args, **kwargs):
|
||||
# local version
|
||||
module1 = MyModule()
|
||||
module2 = MyModule()
|
||||
params = [module1.get_w(), module2.get_w()]
|
||||
local_optim = optim_cls(params, *args, **kwargs)
|
||||
|
||||
old_w1 = module1.w.detach().clone()
|
||||
old_w2 = module2.w.detach().clone()
|
||||
|
||||
g_cpu = torch.Generator()
|
||||
g_cpu.manual_seed(0)
|
||||
t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
output1 = module1.forward(t2)
|
||||
output2 = module2.forward(output1)
|
||||
loss = torch.add(output2, t1).sum()
|
||||
|
||||
loss.backward()
|
||||
local_optim.step()
|
||||
|
||||
# distributed version
|
||||
owner1 = f"worker{(self.rank + 1) % self.world_size:d}"
|
||||
owner2 = f"worker{(self.rank + 2) % self.world_size:d}"
|
||||
|
||||
remote_module1 = rpc.remote(owner1, MyModule)
|
||||
remote_module2 = rpc.remote(owner2, MyModule)
|
||||
remote_param1 = remote_method(MyModule.get_w, remote_module1)
|
||||
remote_param2 = remote_method(MyModule.get_w, remote_module2)
|
||||
|
||||
# sanity check: local and remote initial weights should match
|
||||
self.assertEqual(old_w1, remote_param1.to_here())
|
||||
self.assertEqual(old_w2, remote_param2.to_here())
|
||||
|
||||
dist_optim = DistributedOptimizer(
|
||||
optim_cls, [remote_param1, remote_param2], *args, **kwargs
|
||||
)
|
||||
|
||||
with dist_autograd.context() as context_id:
|
||||
g_cpu.manual_seed(0)
|
||||
t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
output1 = rpc_async_method(MyModule.forward, remote_module1, t2)
|
||||
output2 = rpc_async_method(MyModule.forward, remote_module2, output1.wait())
|
||||
loss = torch.add(output2.wait(), t1)
|
||||
|
||||
dist_autograd.backward(context_id, [loss.sum()])
|
||||
dist_optim.step(context_id)
|
||||
|
||||
new_w1 = rpc_async_method(MyModule.get_w, remote_module1).wait()
|
||||
new_w2 = rpc_async_method(MyModule.get_w, remote_module2).wait()
|
||||
|
||||
# ensure optimizer changed weights
|
||||
self.assertNotEqual(old_w1, new_w1)
|
||||
self.assertNotEqual(old_w2, new_w2)
|
||||
# ensure local equals remote
|
||||
self.assertEqual(new_w1, module1.get_w())
|
||||
self.assertEqual(new_w2, module2.get_w())
|
||||
|
||||
@dist_init()
|
||||
def test_dist_optim(self):
|
||||
self._test_dist_optim_base(optim.Adagrad, lr=0.05)
|
||||
self._test_dist_optim_base(optim.Adam, lr=1e-2, amsgrad=True)
|
||||
self._test_dist_optim_base(optim.AdamW, lr=0.05, amsgrad=True)
|
||||
self._test_dist_optim_base(optim.SGD, lr=0.05)
|
||||
self._test_dist_optim_base(
|
||||
optim.SGD, lr=1e-3, momentum=1, weight_decay=1, nesterov=True
|
||||
)
|
||||
self._test_dist_optim_base(optim.Adadelta, rho=0.95)
|
||||
self._test_dist_optim_base(optim.RMSprop, lr=0.05)
|
||||
self._test_dist_optim_base(optim.Adamax, lr=0.05)
|
||||
self._test_dist_optim_base(optim.Rprop, lr=0.05)
|
||||
|
||||
def _test_dist_optim_none_grads(self, optim_cls, *args, **kwargs):
|
||||
# local version
|
||||
module1 = MyModule()
|
||||
module2 = MyModule(requires_grad=False)
|
||||
params = [module1.get_w(), module2.get_w()]
|
||||
local_optim = optim_cls(params, *args, **kwargs)
|
||||
|
||||
old_w1 = module1.w.detach().clone()
|
||||
old_w2 = module2.w.detach().clone()
|
||||
|
||||
g_cpu = torch.Generator()
|
||||
g_cpu.manual_seed(0)
|
||||
t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
output1 = module1.forward(t2)
|
||||
output2 = module2.forward(output1)
|
||||
loss = torch.add(output2, t1).sum()
|
||||
|
||||
loss.backward()
|
||||
local_optim.step()
|
||||
|
||||
# distributed version
|
||||
owner1 = f"worker{(self.rank + 1) % self.world_size:d}"
|
||||
owner2 = f"worker{(self.rank + 2) % self.world_size:d}"
|
||||
|
||||
remote_module1 = rpc.remote(owner1, MyModule)
|
||||
remote_module2 = rpc.remote(owner2, MyModule, args=(False,))
|
||||
remote_param1 = remote_module1.remote().get_w()
|
||||
remote_param2 = remote_module2.remote().get_w()
|
||||
|
||||
# sanity check: local and remote initial weights should match
|
||||
self.assertEqual(old_w1, remote_param1.to_here())
|
||||
self.assertEqual(old_w2, remote_param2.to_here())
|
||||
|
||||
dist_optim = DistributedOptimizer(
|
||||
optim_cls, [remote_param1, remote_param2], *args, **kwargs
|
||||
)
|
||||
|
||||
with dist_autograd.context() as context_id:
|
||||
g_cpu.manual_seed(0)
|
||||
t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu)
|
||||
output1 = remote_module1.rpc_async().forward(t2)
|
||||
output2 = remote_module2.rpc_async().forward(output1.wait())
|
||||
loss = torch.add(output2.wait(), t1)
|
||||
|
||||
dist_autograd.backward(context_id, [loss.sum()])
|
||||
dist_optim.step(context_id)
|
||||
|
||||
new_w1 = remote_module1.rpc_async().get_w().wait()
|
||||
new_w2 = remote_module2.rpc_async().get_w().wait()
|
||||
|
||||
# ensure optimizer changed weights for w1
|
||||
self.assertNotEqual(old_w1, new_w1)
|
||||
|
||||
# ensure optimizer not changed weights for w2
|
||||
self.assertEqual(old_w2, new_w2)
|
||||
# ensure local equals remote
|
||||
self.assertEqual(new_w1, module1.get_w())
|
||||
self.assertEqual(new_w2, module2.get_w())
|
||||
|
||||
@dist_init()
|
||||
def test_dist_optim_none_grads(self):
|
||||
self._test_dist_optim_none_grads(optim.SGD, lr=0.05)
|
||||
self._test_dist_optim_none_grads(optim.RMSprop, lr=0.05)
|
||||
self._test_dist_optim_none_grads(optim.Rprop, lr=0.05)
|
||||
self._test_dist_optim_none_grads(optim.Adadelta, rho=0.95)
|
||||
+140
@@ -0,0 +1,140 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
# If you need to modify this file to make this test pass, please also apply same edits accordingly to
|
||||
# https://github.com/pytorch/examples/blob/master/distributed/rpc/batch/parameter_server.py
|
||||
# and https://pytorch.org/tutorials/intermediate/rpc_async_execution.html#batch-updating-parameter-server
|
||||
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from time import perf_counter
|
||||
|
||||
import torch
|
||||
import torch.distributed.rpc as rpc
|
||||
import torch.nn as nn
|
||||
from torch import optim
|
||||
from torch.testing._internal.dist_utils import dist_init, worker_name
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
batch_size = 20
|
||||
in_features = 100
|
||||
out_features = 30
|
||||
num_batches = 4
|
||||
|
||||
|
||||
def timed_log(text):
|
||||
print(f"{datetime.now().strftime('%H:%M:%S')} {text}")
|
||||
|
||||
|
||||
class BatchUpdateParameterServer:
|
||||
def __init__(self, batch_update_size):
|
||||
self.model = nn.Linear(in_features, out_features)
|
||||
self.lock = threading.Lock()
|
||||
self.future_model = torch.futures.Future()
|
||||
self.batch_update_size = batch_update_size
|
||||
self.curr_update_size = 0
|
||||
self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
|
||||
for p in self.model.parameters():
|
||||
p.grad = torch.zeros_like(p)
|
||||
|
||||
def get_model(self):
|
||||
return self.model
|
||||
|
||||
@staticmethod
|
||||
@rpc.functions.async_execution
|
||||
def update_and_fetch_model(ps_rref, grads):
|
||||
self = ps_rref.local_value()
|
||||
for p, g in zip(self.model.parameters(), grads, strict=True):
|
||||
if p.grad is None:
|
||||
p.grad = g
|
||||
else:
|
||||
p.grad += g
|
||||
with self.lock:
|
||||
timed_log(
|
||||
f"PS got {self.curr_update_size}/{self.batch_update_size} updates"
|
||||
)
|
||||
self.curr_update_size += 1
|
||||
fut = self.future_model
|
||||
|
||||
if self.curr_update_size >= self.batch_update_size:
|
||||
for p in self.model.parameters():
|
||||
p.grad /= self.batch_update_size
|
||||
self.curr_update_size = 0
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
fut.set_result(self.model)
|
||||
timed_log("PS updated model")
|
||||
self.future_model = torch.futures.Future()
|
||||
|
||||
return fut
|
||||
|
||||
|
||||
class Trainer:
|
||||
def __init__(self, ps_rref):
|
||||
self.ps_rref = ps_rref
|
||||
self.loss_fn = nn.L1Loss()
|
||||
|
||||
def get_next_batch(self):
|
||||
for _ in range(num_batches):
|
||||
inputs = torch.randn(batch_size, in_features)
|
||||
labels = torch.zeros(batch_size, out_features)
|
||||
yield inputs, labels
|
||||
|
||||
def train(self):
|
||||
name = rpc.get_worker_info().name
|
||||
m = self.ps_rref.rpc_sync().get_model()
|
||||
for inputs, labels in self.get_next_batch():
|
||||
timed_log(f"{name} processing one batch")
|
||||
self.loss_fn(m(inputs), labels).backward()
|
||||
timed_log(f"{name} reporting grads")
|
||||
m = rpc.rpc_sync(
|
||||
self.ps_rref.owner(),
|
||||
BatchUpdateParameterServer.update_and_fetch_model,
|
||||
args=(self.ps_rref, [p.grad for p in m.cpu().parameters()]),
|
||||
)
|
||||
timed_log(f"{name} got updated model")
|
||||
|
||||
|
||||
def run_trainer(ps_rref):
|
||||
trainer = Trainer(ps_rref)
|
||||
trainer.train()
|
||||
|
||||
|
||||
def run_ps(trainers):
|
||||
timed_log("Start training")
|
||||
start = perf_counter()
|
||||
ps_rref = rpc.RRef(BatchUpdateParameterServer(len(trainers)))
|
||||
futs = [
|
||||
rpc.rpc_async(trainer, run_trainer, args=(ps_rref,)) for trainer in trainers
|
||||
]
|
||||
|
||||
torch.futures.wait_all(futs)
|
||||
stop = perf_counter()
|
||||
timed_log("Finish training")
|
||||
timed_log(f"Time spent training: {stop - start}s")
|
||||
|
||||
|
||||
class ParameterServerTest(RpcAgentTestFixture):
|
||||
@dist_init(setup_rpc=False)
|
||||
def test_batch_updating_parameter_server(self):
|
||||
if self.rank != 0:
|
||||
rpc.init_rpc(
|
||||
name=worker_name(self.rank),
|
||||
backend=self.rpc_backend,
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
rpc_backend_options=self.rpc_backend_options,
|
||||
)
|
||||
else:
|
||||
rpc.init_rpc(
|
||||
name=worker_name(self.rank),
|
||||
backend=self.rpc_backend,
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
rpc_backend_options=self.rpc_backend_options,
|
||||
)
|
||||
run_ps([f"{worker_name(r)}" for r in range(1, self.world_size)])
|
||||
|
||||
rpc.shutdown()
|
||||
+265
@@ -0,0 +1,265 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
# If you need to modify this file to make this test pass, please also apply same edits accordingly to
|
||||
# https://github.com/pytorch/examples/blob/master/distributed/rpc/rl/main.py
|
||||
# and https://pytorch.org/tutorials/intermediate/rpc_tutorial.html
|
||||
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.distributed.rpc as rpc
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
from torch.distributed.rpc import remote, rpc_async, rpc_sync, RRef
|
||||
from torch.distributions import Categorical
|
||||
from torch.testing._internal.dist_utils import dist_init, worker_name
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
TOTAL_EPISODE_STEP = 5000
|
||||
GAMMA = 0.1
|
||||
SEED = 543
|
||||
|
||||
|
||||
def _call_method(method, rref, *args, **kwargs):
|
||||
r"""
|
||||
a helper function to call a method on the given RRef
|
||||
"""
|
||||
return method(rref.local_value(), *args, **kwargs)
|
||||
|
||||
|
||||
def _remote_method(method, rref, *args, **kwargs):
|
||||
r"""
|
||||
a helper function to run method on the owner of rref and fetch back the
|
||||
result using RPC
|
||||
"""
|
||||
args = [method, rref] + list(args)
|
||||
return rpc_sync(rref.owner(), _call_method, args=args, kwargs=kwargs)
|
||||
|
||||
|
||||
class Policy(nn.Module):
|
||||
r"""
|
||||
Borrowing the ``Policy`` class from the Reinforcement Learning example.
|
||||
Copying the code to make these two examples independent.
|
||||
See https://github.com/pytorch/examples/tree/master/reinforcement_learning
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.affine1 = nn.Linear(4, 128)
|
||||
self.dropout = nn.Dropout(p=0.6)
|
||||
self.affine2 = nn.Linear(128, 2)
|
||||
|
||||
self.saved_log_probs = []
|
||||
self.rewards = []
|
||||
|
||||
def forward(self, x):
|
||||
x = self.affine1(x)
|
||||
x = self.dropout(x)
|
||||
x = F.relu(x)
|
||||
action_scores = self.affine2(x)
|
||||
return F.softmax(action_scores, dim=1)
|
||||
|
||||
|
||||
class DummyEnv:
|
||||
r"""
|
||||
A dummy environment that implements the required subset of the OpenAI gym
|
||||
interface. It exists only to avoid a dependency on gym for running the
|
||||
tests in this file. It is designed to run for a set max number of iterations,
|
||||
returning random states and rewards at each step.
|
||||
"""
|
||||
|
||||
def __init__(self, state_dim=4, num_iters=10, reward_threshold=475.0):
|
||||
self.state_dim = state_dim
|
||||
self.num_iters = num_iters
|
||||
self.iter = 0
|
||||
self.reward_threshold = reward_threshold
|
||||
|
||||
def seed(self, manual_seed):
|
||||
torch.manual_seed(manual_seed)
|
||||
|
||||
def reset(self):
|
||||
self.iter = 0
|
||||
return torch.randn(self.state_dim)
|
||||
|
||||
def step(self, action):
|
||||
self.iter += 1
|
||||
state = torch.randn(self.state_dim)
|
||||
reward = torch.rand(1).item() * self.reward_threshold
|
||||
done = self.iter >= self.num_iters
|
||||
info = {}
|
||||
return state, reward, done, info
|
||||
|
||||
|
||||
class Observer:
|
||||
r"""
|
||||
An observer has exclusive access to its own environment. Each observer
|
||||
captures the state from its environment, and send the state to the agent to
|
||||
select an action. Then, the observer applies the action to its environment
|
||||
and reports the reward to the agent.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.id = rpc.get_worker_info().id
|
||||
self.env = DummyEnv()
|
||||
self.env.seed(SEED)
|
||||
|
||||
def run_episode(self, agent_rref, n_steps):
|
||||
r"""
|
||||
Run one episode of n_steps.
|
||||
Arguments:
|
||||
agent_rref (RRef): an RRef referencing the agent object.
|
||||
n_steps (int): number of steps in this episode
|
||||
"""
|
||||
state, _ep_reward = self.env.reset(), 0
|
||||
for _ in range(n_steps):
|
||||
# send the state to the agent to get an action
|
||||
action = _remote_method(Agent.select_action, agent_rref, self.id, state)
|
||||
|
||||
# apply the action to the environment, and get the reward
|
||||
state, reward, done, _ = self.env.step(action)
|
||||
|
||||
# report the reward to the agent for training purpose
|
||||
_remote_method(Agent.report_reward, agent_rref, self.id, reward)
|
||||
|
||||
if done:
|
||||
break
|
||||
|
||||
|
||||
class Agent:
|
||||
def __init__(self, world_size):
|
||||
self.ob_rrefs = []
|
||||
self.agent_rref = RRef(self)
|
||||
self.rewards = {}
|
||||
self.saved_log_probs = {}
|
||||
self.policy = Policy()
|
||||
self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2)
|
||||
self.eps = np.finfo(np.float32).eps.item()
|
||||
self.running_reward = 0
|
||||
self.reward_threshold = DummyEnv().reward_threshold
|
||||
for ob_rank in range(1, world_size):
|
||||
ob_info = rpc.get_worker_info(worker_name(ob_rank))
|
||||
self.ob_rrefs.append(remote(ob_info, Observer))
|
||||
self.rewards[ob_info.id] = []
|
||||
self.saved_log_probs[ob_info.id] = []
|
||||
|
||||
def select_action(self, ob_id, state):
|
||||
r"""
|
||||
This function is mostly borrowed from the Reinforcement Learning example.
|
||||
See https://github.com/pytorch/examples/tree/master/reinforcement_learning
|
||||
The main difference is that instead of keeping all probs in one list,
|
||||
the agent keeps probs in a dictionary, one key per observer.
|
||||
|
||||
NB: no need to enforce thread-safety here as GIL will serialize
|
||||
executions.
|
||||
"""
|
||||
probs = self.policy(state.unsqueeze(0))
|
||||
m = Categorical(probs)
|
||||
action = m.sample()
|
||||
self.saved_log_probs[ob_id].append(m.log_prob(action))
|
||||
return action.item()
|
||||
|
||||
def report_reward(self, ob_id, reward):
|
||||
r"""
|
||||
Observers call this function to report rewards.
|
||||
"""
|
||||
self.rewards[ob_id].append(reward)
|
||||
|
||||
def run_episode(self, n_steps=0):
|
||||
r"""
|
||||
Run one episode. The agent will tell each observer to run n_steps.
|
||||
"""
|
||||
# make async RPC to kick off an episode on all observers
|
||||
futs = [
|
||||
rpc_async(
|
||||
ob_rref.owner(),
|
||||
_call_method,
|
||||
args=(Observer.run_episode, ob_rref, self.agent_rref, n_steps),
|
||||
)
|
||||
for ob_rref in self.ob_rrefs
|
||||
]
|
||||
|
||||
# wait until all observers have finished this episode
|
||||
for fut in futs:
|
||||
fut.wait()
|
||||
|
||||
def finish_episode(self):
|
||||
r"""
|
||||
This function is mostly borrowed from the Reinforcement Learning example.
|
||||
See https://github.com/pytorch/examples/tree/master/reinforcement_learning
|
||||
The main difference is that it joins all probs and rewards from
|
||||
different observers into one list, and uses the minimum observer rewards
|
||||
as the reward of the current episode.
|
||||
"""
|
||||
|
||||
# joins probs and rewards from different observers into lists
|
||||
R, probs, rewards = 0, [], []
|
||||
for ob_id in self.rewards:
|
||||
probs.extend(self.saved_log_probs[ob_id])
|
||||
rewards.extend(self.rewards[ob_id])
|
||||
|
||||
# use the minimum observer reward to calculate the running reward
|
||||
min_reward = min(sum(self.rewards[ob_id]) for ob_id in self.rewards)
|
||||
self.running_reward = 0.05 * min_reward + (1 - 0.05) * self.running_reward
|
||||
|
||||
# clear saved probs and rewards
|
||||
for ob_id in self.rewards:
|
||||
self.rewards[ob_id] = []
|
||||
self.saved_log_probs[ob_id] = []
|
||||
|
||||
policy_loss, returns = [], []
|
||||
for r in rewards[::-1]:
|
||||
R = r + GAMMA * R
|
||||
returns.insert(0, R)
|
||||
returns = torch.tensor(returns)
|
||||
returns = (returns - returns.mean()) / (returns.std() + self.eps)
|
||||
for log_prob, R in zip(probs, returns, strict=True):
|
||||
policy_loss.append(-log_prob * R)
|
||||
self.optimizer.zero_grad()
|
||||
policy_loss = torch.cat(policy_loss).sum()
|
||||
policy_loss.backward()
|
||||
self.optimizer.step()
|
||||
return min_reward
|
||||
|
||||
|
||||
def run_agent(agent, n_steps):
|
||||
while True:
|
||||
agent.run_episode(n_steps=n_steps)
|
||||
agent.finish_episode()
|
||||
|
||||
if agent.running_reward > agent.reward_threshold:
|
||||
print(f"Solved! Running reward is now {agent.running_reward}!")
|
||||
break
|
||||
|
||||
|
||||
class ReinforcementLearningRpcTest(RpcAgentTestFixture):
|
||||
@dist_init(setup_rpc=False)
|
||||
def test_rl_rpc(self):
|
||||
if self.rank == 0:
|
||||
# Rank 0 is the agent.
|
||||
rpc.init_rpc(
|
||||
name=worker_name(self.rank),
|
||||
backend=self.rpc_backend,
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
rpc_backend_options=self.rpc_backend_options,
|
||||
)
|
||||
agent = Agent(self.world_size)
|
||||
run_agent(agent, n_steps=int(TOTAL_EPISODE_STEP / (self.world_size - 1)))
|
||||
|
||||
# Ensure training was run. We don't really care about whether the task was learned,
|
||||
# since the purpose of the test is to check the API calls.
|
||||
self.assertGreater(agent.running_reward, 0.0)
|
||||
else:
|
||||
# Other ranks are observers that passively wait for instructions from the agent.
|
||||
rpc.init_rpc(
|
||||
name=worker_name(self.rank),
|
||||
backend=self.rpc_backend,
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
rpc_backend_options=self.rpc_backend_options,
|
||||
)
|
||||
rpc.shutdown()
|
||||
+337
@@ -0,0 +1,337 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import time
|
||||
|
||||
import torch
|
||||
import torch.distributed.rpc as rpc
|
||||
from torch.distributed.rpc.api import _delete_all_user_and_unforked_owner_rrefs
|
||||
from torch.testing._internal.dist_utils import (
|
||||
dist_init,
|
||||
wait_until_owners_and_forks_on_rank,
|
||||
wait_until_pending_futures_and_users_flushed,
|
||||
worker_name,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
def my_sleep_func(seconds=1):
|
||||
time.sleep(seconds)
|
||||
return torch.mul(torch.tensor(1), torch.tensor(1))
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def my_script_func(tensor):
|
||||
return torch.add(tensor, tensor)
|
||||
|
||||
|
||||
def add_rref_to_value(rref, value):
|
||||
return rref.to_here() + value
|
||||
|
||||
|
||||
class FaultyAgentRpcTest(RpcAgentTestFixture):
|
||||
# no faulty_messages defined so this fails all retryable messages - see
|
||||
# faulty_rpc_agent_test_fixture.py for the list of retryable messages.
|
||||
@dist_init(messages_to_delay={})
|
||||
def test_check_failed_messages(self):
|
||||
if self.rank == 0:
|
||||
dst_worker_b = worker_name((self.rank + 1) % self.world_size)
|
||||
dst_worker_c = worker_name((self.rank + 2) % self.world_size)
|
||||
|
||||
# Worker0 sends RPC to Worker1 and creates an RRef there
|
||||
rref = rpc.remote(
|
||||
dst_worker_b, torch.add, args=(torch.ones(2, 2), torch.ones(2, 2))
|
||||
)
|
||||
# Worker0 sends an RPC to Worker2 with the RRef as an arg
|
||||
rpc.remote(dst_worker_c, add_rref_to_value, args=(rref, torch.ones(2, 2)))
|
||||
# check if the output is as expected
|
||||
self.assertEqual(
|
||||
rref.to_here(), torch.add(torch.ones(2, 2), torch.ones(2, 2))
|
||||
)
|
||||
# explicitly delete all User RRefs
|
||||
_delete_all_user_and_unforked_owner_rrefs()
|
||||
|
||||
@dist_init
|
||||
def test_verify_backend_options(self):
|
||||
self.assertEqual(
|
||||
self.rpc_backend, rpc.backend_registry.BackendType.FAULTY_TENSORPIPE
|
||||
)
|
||||
self.assertEqual(self.rpc_backend_options.num_worker_threads, 8)
|
||||
self.assertEqual(self.rpc_backend_options.num_fail_sends, 3)
|
||||
self.assertEqual(len(self.rpc_backend_options.messages_to_fail), 4)
|
||||
self.assertEqual(len(self.rpc_backend_options.messages_to_delay), 2)
|
||||
self.assertEqual(
|
||||
self.rpc_backend_options.rpc_timeout, rpc.constants.DEFAULT_RPC_TIMEOUT_SEC
|
||||
)
|
||||
|
||||
@dist_init(faulty_messages=["RREF_FORK_REQUEST", "RREF_CHILD_ACCEPT"])
|
||||
def test_custom_faulty_messages(self):
|
||||
self.assertEqual(
|
||||
{"RREF_FORK_REQUEST", "RREF_CHILD_ACCEPT"},
|
||||
set(self.rpc_backend_options.messages_to_fail),
|
||||
)
|
||||
|
||||
@dist_init(faulty_messages=[])
|
||||
def test_no_faulty_messages(self):
|
||||
self.assertEqual(len(self.rpc_backend_options.messages_to_fail), 0)
|
||||
|
||||
@dist_init(messages_to_delay={"SCRIPT_CALL": 1.5})
|
||||
def test_custom_messages_to_delay(self):
|
||||
self.assertEqual(
|
||||
self.rpc_backend_options.messages_to_delay, {"SCRIPT_CALL": 1.5}
|
||||
)
|
||||
|
||||
def _test_remote_message_dropped_pickle(self, dst=None):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size
|
||||
dst_worker = f"worker{dst_rank}"
|
||||
# Since we fail python_remote_call messages synchronously, the future
|
||||
# corresponding to this remote call will be marked with an error when
|
||||
# this function returns.
|
||||
rref = rpc.remote(dst_worker, my_sleep_func, args=(1,))
|
||||
# Call to ensure pending callbacks are run.
|
||||
wait_until_pending_futures_and_users_flushed()
|
||||
# Attempt to fork the RRef should raise an error indicating the rpc.remote timeout.
|
||||
with self.assertRaisesRegex(RuntimeError, "RRef creation"):
|
||||
rref._serialize()
|
||||
# Test that using RRef as arg over RPC (which forks) results in the same
|
||||
# error
|
||||
with self.assertRaisesRegex(RuntimeError, "RRef creation"):
|
||||
rpc.rpc_async(dst_worker, add_rref_to_value, args=(rref, 1))
|
||||
|
||||
@dist_init(faulty_messages=["PYTHON_REMOTE_CALL"])
|
||||
def test_remote_message_dropped_pickle(self):
|
||||
self._test_remote_message_dropped_pickle()
|
||||
|
||||
@dist_init(faulty_messages=["PYTHON_REMOTE_CALL"])
|
||||
def test_remote_message_dropped_pickle_to_self(self):
|
||||
self._test_remote_message_dropped_pickle(self.rank)
|
||||
|
||||
def _test_remote_message_dropped_timeout(self, func, args, dst=None):
|
||||
if self.rank != 0:
|
||||
return
|
||||
|
||||
# test the case where rpc.remote() message creation is completely dropped.
|
||||
dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size
|
||||
dst_worker = f"worker{dst_rank}"
|
||||
# Since we fail python_remote_call messages synchronously, the future
|
||||
# corresponding to this remote call will be marked with an error when
|
||||
# this function returns.
|
||||
rref = rpc.remote(dst_worker, func, args=args)
|
||||
# Call to ensure pending callbacks are run.
|
||||
wait_until_pending_futures_and_users_flushed()
|
||||
with self.assertRaisesRegex(RuntimeError, "RRef creation"):
|
||||
rref.to_here()
|
||||
# Note: during shutdown, logs will indicate "Could not find OwnerRRef..."
|
||||
# on the owning nodes, this is expected because the OwnerRRef was never
|
||||
# successfully created. Therefore, delAllUsers will work as expected.
|
||||
|
||||
@dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"])
|
||||
def test_builtin_remote_message_dropped_timeout(self):
|
||||
func = torch.add
|
||||
args = (torch.tensor(1), torch.tensor(1))
|
||||
self._test_remote_message_dropped_timeout(func, args)
|
||||
|
||||
@dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"])
|
||||
def test_builtin_remote_message_dropped_timeout_to_self(self):
|
||||
func = torch.add
|
||||
args = (torch.tensor(1), torch.tensor(1))
|
||||
self._test_remote_message_dropped_timeout(func, args, dst=0)
|
||||
|
||||
@dist_init(faulty_messages=["PYTHON_REMOTE_CALL"])
|
||||
def test_udf_remote_message_dropped_timeout(self):
|
||||
func = my_sleep_func
|
||||
args = (2,)
|
||||
self._test_remote_message_dropped_timeout(func, args)
|
||||
|
||||
@dist_init(faulty_messages=["PYTHON_REMOTE_CALL"])
|
||||
def test_udf_remote_message_dropped_timeout_to_self(self):
|
||||
func = my_sleep_func
|
||||
args = (2,)
|
||||
self._test_remote_message_dropped_timeout(func, args, dst=0)
|
||||
|
||||
def _test_remote_message_delay_timeout(self, func, args, dst=None):
|
||||
if self.rank != 0:
|
||||
return
|
||||
# Test the case where remote message is eventually processed on the owner,
|
||||
# but the future on the creator times out before the response comes back.
|
||||
dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size
|
||||
dst_worker = f"worker{dst_rank}"
|
||||
# 10 ms timeout
|
||||
rref = rpc.remote(dst_worker, func, args=args, timeout=0.001)
|
||||
# Future corresponding to the remote creation should time out.
|
||||
expected_error = self.get_timeout_error_regex()
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
rref._get_future().wait()
|
||||
|
||||
# Call to ensure pending callbacks are run.
|
||||
wait_until_pending_futures_and_users_flushed()
|
||||
# to_here() should now pick up that rpc.remote() creation has failed.
|
||||
with self.assertRaisesRegex(RuntimeError, "RRef creation"):
|
||||
rref.to_here()
|
||||
|
||||
# Test the case where rpc.remote() times out, but to_here() has already
|
||||
# started blocking before.
|
||||
# NOTE: we only test this when not sending to self, as to_here() calls
|
||||
# calls localValue(), which does not send an RPC and thus does not have
|
||||
# a timeout. This can be supported by allowing future.wait() to
|
||||
# take in an optional timeout (https://github.com/pytorch/pytorch/issues/39280)
|
||||
if dst_rank != self.rank:
|
||||
slow_rref = rpc.remote(dst_worker, func, args=args, timeout=2)
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
# to_here() should raise timeout error, since it does not know about the
|
||||
# status of rpc.remote().
|
||||
slow_rref.to_here(0.001)
|
||||
# Note: If we proceed with shutdown, UserRRef will send out a RRefUserDelete
|
||||
# but this can be a noop since it may not exist on the owner yet. Later,
|
||||
# the owner can process the RRef creation and wait for the delete message,
|
||||
# thus leading to a timeout.
|
||||
# Therefore, we wait until we get notification that pending owners have
|
||||
# been confirmed before sending out RRefUserDeletes.
|
||||
if dst_rank != self.rank:
|
||||
wait_until_owners_and_forks_on_rank(2, 2, rank=dst_rank)
|
||||
|
||||
@dist_init(faulty_messages=[], messages_to_delay={"PYTHON_REMOTE_CALL": 2})
|
||||
def test_udf_remote_message_delay_timeout(self):
|
||||
func = my_sleep_func
|
||||
args = (2,)
|
||||
self._test_remote_message_delay_timeout(func, args)
|
||||
|
||||
@dist_init(faulty_messages=[], messages_to_delay={"PYTHON_REMOTE_CALL": 2})
|
||||
def test_udf_remote_message_delay_timeout_to_self(self):
|
||||
func = my_sleep_func
|
||||
args = (1,)
|
||||
self._test_remote_message_delay_timeout(func, args, dst=0)
|
||||
|
||||
@dist_init(
|
||||
faulty_messages=[],
|
||||
messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1},
|
||||
)
|
||||
def test_remote_message_builtin_delay_timeout(self):
|
||||
func = torch.add
|
||||
args = (torch.tensor(1), torch.tensor(1))
|
||||
self._test_remote_message_delay_timeout(func, args)
|
||||
|
||||
@dist_init(
|
||||
faulty_messages=[],
|
||||
messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1},
|
||||
)
|
||||
def test_remote_message_builtin_delay_timeout_to_self(self):
|
||||
func = torch.add
|
||||
args = (torch.tensor(1), torch.tensor(1))
|
||||
self._test_remote_message_delay_timeout(func, args, dst=0)
|
||||
|
||||
@dist_init(
|
||||
faulty_messages=[],
|
||||
messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1},
|
||||
)
|
||||
def test_remote_message_script_delay_timeout(self):
|
||||
func = my_script_func
|
||||
args = (torch.tensor(1),)
|
||||
self._test_remote_message_delay_timeout(func, args)
|
||||
|
||||
@dist_init(
|
||||
faulty_messages=[],
|
||||
messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1},
|
||||
)
|
||||
def test_remote_message_script_delay_timeout_to_self(self):
|
||||
func = my_script_func
|
||||
args = (torch.tensor(1),)
|
||||
self._test_remote_message_delay_timeout(func, args, dst=0)
|
||||
|
||||
@dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_RREF_FETCH_CALL": 1})
|
||||
def test_rref_to_here_timeout(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
|
||||
dst_rank = (self.rank + 1) % self.world_size
|
||||
dst_worker = f"worker{dst_rank}"
|
||||
rref = rpc.remote(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1))
|
||||
)
|
||||
expected_error = self.get_timeout_error_regex()
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
rref.to_here(0.01)
|
||||
|
||||
rref.to_here()
|
||||
|
||||
@dist_init(faulty_messages=[])
|
||||
def test_rpc_builtin_timeout(self):
|
||||
next_rank = (self.rank + 1) % self.world_size
|
||||
dst_worker = worker_name(next_rank)
|
||||
expected_error = self.get_timeout_error_regex()
|
||||
# PYTHON_CALL message types which correspond to Python UDF over RPC
|
||||
# by default get a delay (see faulty_rpc_agent_test_fixture)
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
rpc.rpc_sync(
|
||||
dst_worker,
|
||||
torch.add,
|
||||
args=(torch.tensor(1), torch.tensor(1)),
|
||||
timeout=1,
|
||||
)
|
||||
|
||||
fut = rpc.rpc_async(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)), timeout=1
|
||||
)
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
fut.wait()
|
||||
|
||||
# Ensure that the currently set default timeout is large enough such
|
||||
# that RPCs with delays still complete.
|
||||
fut = rpc.rpc_async(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1))
|
||||
)
|
||||
fut.wait()
|
||||
|
||||
# Ensure timeout if we set a new default and don't override
|
||||
rpc._set_rpc_timeout(0.001)
|
||||
fut = rpc.rpc_async(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1))
|
||||
)
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
fut.wait()
|
||||
|
||||
# Ensure run to completion if we specify timeout of 0
|
||||
fut = rpc.rpc_async(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)), timeout=0
|
||||
)
|
||||
fut.wait()
|
||||
# Reset for clean shutdown
|
||||
rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC)
|
||||
|
||||
@dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5})
|
||||
def test_rpc_script_timeout(self):
|
||||
next_rank = (self.rank + 1) % self.world_size
|
||||
dst_worker = worker_name(next_rank)
|
||||
expected_error = self.get_timeout_error_regex()
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
rpc.rpc_sync(dst_worker, my_script_func, args=(torch.tensor(1),), timeout=1)
|
||||
|
||||
fut = rpc.rpc_async(
|
||||
dst_worker, my_script_func, args=(torch.tensor(1),), timeout=1
|
||||
)
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
fut.wait()
|
||||
|
||||
# Ensure that the currently set default timeout is large enough such
|
||||
# that RPCs with delays still complete.
|
||||
fut = rpc.rpc_async(dst_worker, my_script_func, args=(torch.tensor(1),))
|
||||
fut.wait()
|
||||
|
||||
# Ensure timeout if we set a new default and don't override
|
||||
rpc._set_rpc_timeout(0.001)
|
||||
fut = rpc.rpc_async(dst_worker, my_script_func, args=(torch.tensor(1),))
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
fut.wait()
|
||||
|
||||
# Ensure run to completion if we specify timeout of 0
|
||||
rpc._set_rpc_timeout(0.001)
|
||||
fut = rpc.rpc_async(
|
||||
dst_worker, my_script_func, args=(torch.tensor(1),), timeout=0
|
||||
)
|
||||
fut.wait()
|
||||
# Reset for clean shutdown
|
||||
rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC)
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch.distributed.rpc as rpc
|
||||
import torch.distributed.rpc._testing # noqa: F401
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
# The following message types are currently retried in the RREF protocol and
|
||||
# distributed autograd. Thus only these messages should be tested with the
|
||||
# Faulty RPC Agent.
|
||||
retryable_message_types = [
|
||||
"RREF_FORK_REQUEST",
|
||||
"RREF_CHILD_ACCEPT",
|
||||
"RREF_USER_DELETE",
|
||||
"CLEANUP_AUTOGRAD_CONTEXT_REQ",
|
||||
]
|
||||
|
||||
# The following messages incur the corresponding delay in seconds while being
|
||||
# processed in FaultyTensorPipeAgent's enqueueSend() function.
|
||||
default_messages_to_delay = {
|
||||
"PYTHON_CALL": 1.5, # Python UDF
|
||||
"SCRIPT_CALL": 1.5, # Script/Builtin
|
||||
}
|
||||
|
||||
|
||||
class FaultyRpcAgentTestFixture(RpcAgentTestFixture):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.messages_to_fail = retryable_message_types
|
||||
self.messages_to_delay = default_messages_to_delay
|
||||
|
||||
@property
|
||||
def rpc_backend(self):
|
||||
return rpc.backend_registry.BackendType["FAULTY_TENSORPIPE"]
|
||||
|
||||
@property
|
||||
def rpc_backend_options(self):
|
||||
return rpc.backend_registry.construct_rpc_backend_options(
|
||||
self.rpc_backend,
|
||||
init_method=self.init_method,
|
||||
num_worker_threads=8,
|
||||
num_fail_sends=3,
|
||||
messages_to_fail=self.messages_to_fail,
|
||||
messages_to_delay=self.messages_to_delay,
|
||||
)
|
||||
|
||||
def setup_fault_injection(self, faulty_messages, messages_to_delay):
|
||||
if faulty_messages is not None:
|
||||
self.messages_to_fail = faulty_messages
|
||||
if messages_to_delay is not None:
|
||||
self.messages_to_delay = messages_to_delay
|
||||
|
||||
def get_shutdown_error_regex(self):
|
||||
error_regexes = [
|
||||
"Exception in thread pool task",
|
||||
"Connection reset by peer",
|
||||
"Connection closed by peer",
|
||||
]
|
||||
return "|".join([f"({error_str})" for error_str in error_regexes])
|
||||
|
||||
def get_timeout_error_regex(self):
|
||||
return "RPC ran for more than"
|
||||
+113
@@ -0,0 +1,113 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
|
||||
import torch
|
||||
import torch.distributed.autograd as dist_autograd
|
||||
import torch.distributed.rpc as rpc
|
||||
from torch import Tensor
|
||||
from torch.distributed.rpc import rpc_async
|
||||
from torch.testing import FileCheck
|
||||
from torch.testing._internal.dist_utils import dist_init, worker_name
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def local_add(t1, t2):
|
||||
return torch.add(t1, t2)
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def remote_add(t1, t2, dst: str): # noqa: E999
|
||||
return rpc_async(dst, local_add, (t1, t2)).wait()
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def fork_add(t1, t2, dst: str):
|
||||
fut = torch.jit._fork(remote_add, t1, t2, dst)
|
||||
return torch.jit._wait(fut)
|
||||
|
||||
|
||||
class JitDistAutogradTest(RpcAgentTestFixture):
|
||||
@dist_init
|
||||
def test_get_gradients(self):
|
||||
@torch.jit.script
|
||||
def dist_get_gradients(context_id: int) -> dict[Tensor, Tensor]:
|
||||
return dist_autograd.get_gradients(context_id)
|
||||
|
||||
FileCheck().check("get_gradients").run(str(dist_get_gradients.graph))
|
||||
with dist_autograd.context() as context_id:
|
||||
t1 = torch.rand((3, 3), requires_grad=True)
|
||||
t2 = torch.rand((3, 3), requires_grad=True)
|
||||
t3 = torch.add(t1, t2)
|
||||
|
||||
dist_autograd.backward(context_id, [t3.sum()])
|
||||
grads = dist_get_gradients(context_id)
|
||||
|
||||
self.assertEqual(2, len(grads))
|
||||
self.assertIn(t1, grads)
|
||||
self.assertIn(t2, grads)
|
||||
self.assertEqual(torch.ones(3, 3), grads[t1])
|
||||
self.assertEqual(torch.ones(3, 3), grads[t2])
|
||||
|
||||
@dist_init
|
||||
def test_dist_backward(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
|
||||
@torch.jit.script
|
||||
def dist_backward_script(context_id: int, loss: torch.Tensor):
|
||||
dist_autograd.backward(context_id, [loss])
|
||||
|
||||
FileCheck().check("dist_backward").run(str(dist_backward_script.graph))
|
||||
with dist_autograd.context() as context_id:
|
||||
t1 = torch.rand(3, 3, requires_grad=True)
|
||||
t2 = torch.rand(3, 3, requires_grad=True)
|
||||
dst_worker_name = worker_name((self.rank + 1) % self.world_size)
|
||||
loss = rpc.rpc_sync(dst_worker_name, torch.add, args=(t1, t2)).sum()
|
||||
dist_backward_script(context_id, loss)
|
||||
|
||||
@dist_init
|
||||
def test_jit_fork_within_context(self):
|
||||
with dist_autograd.context() as context_id:
|
||||
t1 = torch.rand((3, 3), requires_grad=True)
|
||||
t2 = torch.rand((3, 3), requires_grad=True)
|
||||
dst_worker_name = worker_name((self.rank + 1) % self.world_size)
|
||||
res = fork_add(t1, t2, dst_worker_name)
|
||||
loss = res.sum()
|
||||
dist_autograd.backward(context_id, [loss])
|
||||
|
||||
grads = dist_autograd.get_gradients(context_id)
|
||||
self.assertEqual(2, len(grads))
|
||||
self.assertIn(t1, grads)
|
||||
self.assertIn(t2, grads)
|
||||
|
||||
@dist_init
|
||||
def test_restore_context_after_swtich_to_jit_thread(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
|
||||
@torch.jit.script
|
||||
def forward_script(
|
||||
context_id: int, dst_worker_name: str, t1: Tensor, t2: Tensor
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
res1_fut = rpc.rpc_async(dst_worker_name, local_add, (t1, t1))
|
||||
res1 = res1_fut.wait() # After this, the script runs in a new JIT thread.
|
||||
loss1 = res1.sum()
|
||||
|
||||
# SendRpcBackward is not attached, since DistAutogradContext is lost here.
|
||||
res2_fut = rpc.rpc_async(dst_worker_name, local_add, (t2, t2))
|
||||
res2 = res2_fut.wait()
|
||||
loss2 = res2.sum()
|
||||
|
||||
return loss1, loss2
|
||||
|
||||
with dist_autograd.context() as context_id:
|
||||
t1 = torch.ones((2, 3), requires_grad=True)
|
||||
t2 = torch.ones((2, 3), requires_grad=True)
|
||||
dst_worker_name = worker_name((self.rank + 1) % self.world_size)
|
||||
loss0, loss1 = forward_script(context_id, dst_worker_name, t1, t2)
|
||||
dist_autograd.backward(context_id, [loss0, loss1])
|
||||
grad0, grad1 = dist_autograd.get_gradients(context_id)
|
||||
self.assertEqual(grad0, grad1)
|
||||
+1384
File diff suppressed because it is too large
Load Diff
+219
@@ -0,0 +1,219 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
|
||||
import torch
|
||||
import torch.distributed.rpc as rpc
|
||||
from torch import Tensor
|
||||
from torch.distributed.rpc import RRef
|
||||
from torch.testing._internal.dist_utils import (
|
||||
dist_init,
|
||||
wait_until_pending_futures_and_users_flushed,
|
||||
worker_name,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def two_args_two_kwargs(
|
||||
first_arg,
|
||||
second_arg,
|
||||
first_kwarg=torch.tensor([3, 3]),
|
||||
second_kwarg=torch.tensor([4, 4]),
|
||||
):
|
||||
return first_arg + second_arg + first_kwarg + second_kwarg
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def script_rpc_async_call(
|
||||
dst_worker_name: str, args: tuple[Tensor, Tensor], kwargs: dict[str, Tensor]
|
||||
):
|
||||
fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs)
|
||||
ret = fut.wait()
|
||||
return ret
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def rpc_async_call_with_timeout(
|
||||
dst_worker_name: str,
|
||||
args: tuple[Tensor, Tensor],
|
||||
kwargs: dict[str, Tensor],
|
||||
timeout: float,
|
||||
):
|
||||
fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs, timeout)
|
||||
ret = fut.wait()
|
||||
return ret
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def rpc_async_call_with_timeout_future_ret(
|
||||
dst_worker_name: str,
|
||||
args: tuple[Tensor, Tensor],
|
||||
kwargs: dict[str, Tensor],
|
||||
timeout: float,
|
||||
):
|
||||
fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs, timeout)
|
||||
return fut
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def rpc_async_call_future_ret(
|
||||
dst_worker_name: str, args: tuple[Tensor, Tensor], kwargs: dict[str, Tensor]
|
||||
):
|
||||
fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs)
|
||||
return fut
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def rref_to_here(rref_var: RRef[Tensor]) -> Tensor:
|
||||
return rref_var.to_here()
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def rref_to_here_with_timeout(rref_var: RRef[Tensor], timeout: float) -> Tensor:
|
||||
return rref_var.to_here(timeout)
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def rpc_async_with_rref_arg(dst_worker_name: str, args: tuple[RRef[Tensor]]) -> Tensor:
|
||||
fut = rpc.rpc_async(dst_worker_name, rref_to_here, args)
|
||||
ret = fut.wait()
|
||||
return ret
|
||||
|
||||
|
||||
class JitFaultyAgentRpcTest(RpcAgentTestFixture):
|
||||
"""
|
||||
Run tests for rpc_async in JIT under the faulty agent test fixture to test
|
||||
arbitrary timeouts.
|
||||
"""
|
||||
|
||||
@dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5})
|
||||
def test_timeout_in_torchscript_function(self):
|
||||
# Call rpc_async + fut.wait() in torchscript function and ensure that
|
||||
# timeout is raised.
|
||||
if self.rank != 0:
|
||||
return
|
||||
|
||||
dst_worker_name = worker_name((self.rank + 1) % self.world_size)
|
||||
|
||||
args = (torch.tensor([1, 1]), torch.tensor([2, 2]))
|
||||
kwargs = {
|
||||
"first_kwarg": torch.tensor([2, 2]),
|
||||
"second_kwarg": torch.tensor([3, 3]),
|
||||
}
|
||||
expected_error = self.get_timeout_error_regex()
|
||||
# Ensure that we get a timeout if we override the default timeout and
|
||||
# the RPC takes longer to execute.
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
rpc_async_call_with_timeout(dst_worker_name, args, kwargs, 0.5)
|
||||
|
||||
# Ensure that we timeout if we don't specify a timeout but the default
|
||||
# is less than the RPC takes to execute.
|
||||
rpc._set_rpc_timeout(0.001)
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
script_rpc_async_call(dst_worker_name, args, kwargs)
|
||||
|
||||
# Ensure that we run to completion if zero timeout is specified.
|
||||
ret = rpc_async_call_with_timeout(dst_worker_name, args, kwargs, 0)
|
||||
self.assertEqual(ret, torch.tensor([8, 8]))
|
||||
# reset for clean shutdown
|
||||
rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC)
|
||||
|
||||
@dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5})
|
||||
def test_timeout_in_python(self):
|
||||
# Ensures timeouts are raised if we call rpc_async from within a
|
||||
# torchscript function, but wait on the future in python.
|
||||
if self.rank != 0:
|
||||
return
|
||||
|
||||
dst_worker_name = worker_name((self.rank + 1) % self.world_size)
|
||||
args = (torch.tensor([1, 1]), torch.tensor([2, 2]))
|
||||
kwargs = {
|
||||
"first_kwarg": torch.tensor([2, 2]),
|
||||
"second_kwarg": torch.tensor([3, 3]),
|
||||
}
|
||||
expected_error = self.get_timeout_error_regex()
|
||||
|
||||
fut = rpc_async_call_with_timeout_future_ret(dst_worker_name, args, kwargs, 0.5)
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
fut.wait()
|
||||
|
||||
# Ensure timeout if we don't specify but the default is less than the
|
||||
# RPC takes to execute.
|
||||
rpc._set_rpc_timeout(0.001)
|
||||
fut = rpc_async_call_future_ret(dst_worker_name, args, kwargs)
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
fut.wait()
|
||||
|
||||
# Ensure run to completion if zero timeout is specified
|
||||
fut = rpc_async_call_with_timeout_future_ret(dst_worker_name, args, kwargs, 0)
|
||||
result = fut.wait()
|
||||
self.assertEqual(result, torch.tensor([8, 8]))
|
||||
# reset for clean shutdown
|
||||
rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC)
|
||||
|
||||
@dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"])
|
||||
def test_remote_timeout_to_here_in_jit(self):
|
||||
# Test that calling to_here() in JIT will raise timeout error if
|
||||
# rpc.remote failed.
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_rank = (self.rank + 1) % self.world_size
|
||||
dst_worker = f"worker{dst_rank}"
|
||||
rref = rpc.remote(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1))
|
||||
)
|
||||
# Will ensure error handling callbacks are run.
|
||||
wait_until_pending_futures_and_users_flushed()
|
||||
# Call to_here() within a ScriptFunction and ensure it raises
|
||||
with self.assertRaisesRegex(RuntimeError, "RRef creation"):
|
||||
rref_to_here(rref)
|
||||
|
||||
@dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_RREF_FETCH_CALL": 1})
|
||||
def test_rref_to_here_timeout_in_jit(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
|
||||
dst_rank = (self.rank + 1) % self.world_size
|
||||
dst_worker = f"worker{dst_rank}"
|
||||
rref = rpc.remote(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1))
|
||||
)
|
||||
expected_error = self.get_timeout_error_regex()
|
||||
with self.assertRaisesRegex(RuntimeError, expected_error):
|
||||
rref_to_here_with_timeout(rref, 0.01)
|
||||
|
||||
rref_to_here_with_timeout(rref, 100)
|
||||
|
||||
@dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"])
|
||||
def test_rref_timeout_pickle_in_jit(self):
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_rank = (self.rank + 1) % self.world_size
|
||||
dst_worker = f"worker{dst_rank}"
|
||||
rref = rpc.remote(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1))
|
||||
)
|
||||
# Will ensure error handling callbacks are run.
|
||||
wait_until_pending_futures_and_users_flushed()
|
||||
# Call RPC with RRef arg in JIT, which will go through JIT pickling and
|
||||
# ensure error is raised.
|
||||
with self.assertRaisesRegex(RuntimeError, "RRef creation"):
|
||||
rpc_async_with_rref_arg(dst_worker, (rref,))
|
||||
|
||||
@dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"])
|
||||
def test_rref_timeout_pickle_script_func(self):
|
||||
# Similar to above test, but calls python rpc with script function.
|
||||
if self.rank != 0:
|
||||
return
|
||||
dst_rank = (self.rank + 1) % self.world_size
|
||||
dst_worker = f"worker{dst_rank}"
|
||||
rref = rpc.remote(
|
||||
dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1))
|
||||
)
|
||||
# Will ensure error handling callbacks are run.
|
||||
wait_until_pending_futures_and_users_flushed()
|
||||
# Call RPC with script function that takes RRef, ensure timeout during pickling
|
||||
with self.assertRaisesRegex(RuntimeError, "RRef creation"):
|
||||
rpc.rpc_sync(dst_worker, rref_to_here, args=(rref,))
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch.testing._internal.dist_utils
|
||||
|
||||
|
||||
class RpcAgentTestFixture(ABC):
|
||||
@property
|
||||
def world_size(self) -> int:
|
||||
return 4
|
||||
|
||||
@property
|
||||
def init_method(self):
|
||||
use_tcp_init = os.environ.get("RPC_INIT_WITH_TCP", None)
|
||||
if use_tcp_init == "1":
|
||||
master_addr = os.environ["MASTER_ADDR"]
|
||||
master_port = os.environ["MASTER_PORT"]
|
||||
return f"tcp://{master_addr}:{master_port}"
|
||||
else:
|
||||
return self.file_init_method
|
||||
|
||||
@property
|
||||
def file_init_method(self):
|
||||
return torch.testing._internal.dist_utils.INIT_METHOD_TEMPLATE.format(
|
||||
file_name=self.file_name
|
||||
)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def rpc_backend(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def rpc_backend_options(self):
|
||||
pass
|
||||
|
||||
def setup_fault_injection(self, faulty_messages, messages_to_delay): # noqa: B027
|
||||
"""Method used by dist_init to prepare the faulty agent.
|
||||
|
||||
Does nothing for other agents.
|
||||
"""
|
||||
|
||||
# Shutdown sequence is not well defined, so we may see any of the following
|
||||
# errors when running tests that simulate errors via a shutdown on the
|
||||
# remote end.
|
||||
@abstractmethod
|
||||
def get_shutdown_error_regex(self):
|
||||
"""
|
||||
Return various error message we may see from RPC agents while running
|
||||
tests that check for failures. This function is used to match against
|
||||
possible errors to ensure failures were raised properly.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_timeout_error_regex(self):
|
||||
"""
|
||||
Returns a partial string indicating the error we should receive when an
|
||||
RPC has timed out. Useful for use with assertRaisesRegex() to ensure we
|
||||
have the right errors during timeout.
|
||||
"""
|
||||
+6307
File diff suppressed because it is too large
Load Diff
+28
@@ -0,0 +1,28 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch.distributed.rpc as rpc
|
||||
from torch.testing._internal.common_distributed import tp_transports
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
|
||||
|
||||
class TensorPipeRpcAgentTestFixture(RpcAgentTestFixture):
|
||||
@property
|
||||
def rpc_backend(self):
|
||||
return rpc.backend_registry.BackendType["TENSORPIPE"]
|
||||
|
||||
@property
|
||||
def rpc_backend_options(self):
|
||||
return rpc.backend_registry.construct_rpc_backend_options(
|
||||
self.rpc_backend, init_method=self.init_method, _transports=tp_transports()
|
||||
)
|
||||
|
||||
def get_shutdown_error_regex(self):
|
||||
# FIXME Once we consolidate the error messages returned by the
|
||||
# TensorPipe agent put some more specific regex here.
|
||||
error_regexes = [".*"]
|
||||
return "|".join([f"({error_str})" for error_str in error_regexes])
|
||||
|
||||
def get_timeout_error_regex(self):
|
||||
return "RPC ran for more than"
|
||||
+188
@@ -0,0 +1,188 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
from torch.testing._internal.common_distributed import MultiProcessTestCase
|
||||
from torch.testing._internal.common_utils import (
|
||||
find_free_port,
|
||||
IS_SANDCASTLE,
|
||||
TEST_WITH_DEV_DBG_ASAN,
|
||||
)
|
||||
from torch.testing._internal.distributed.ddp_under_dist_autograd_test import (
|
||||
CudaDdpComparisonTest,
|
||||
DdpComparisonTest,
|
||||
DdpUnderDistAutogradTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.nn.api.remote_module_test import (
|
||||
CudaRemoteModuleTest,
|
||||
RemoteModuleTest,
|
||||
ThreeWorkersRemoteModuleTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.dist_autograd_test import (
|
||||
CudaDistAutogradTest,
|
||||
DistAutogradTest,
|
||||
FaultyAgentDistAutogradTest,
|
||||
TensorPipeAgentDistAutogradTest,
|
||||
TensorPipeCudaDistAutogradTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.dist_optimizer_test import (
|
||||
DistOptimizerTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.examples.parameter_server_test import (
|
||||
ParameterServerTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.examples.reinforcement_learning_rpc_test import (
|
||||
ReinforcementLearningRpcTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.faulty_agent_rpc_test import (
|
||||
FaultyAgentRpcTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.jit.dist_autograd_test import (
|
||||
JitDistAutogradTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.jit.rpc_test import JitRpcTest
|
||||
from torch.testing._internal.distributed.rpc.jit.rpc_test_faulty import (
|
||||
JitFaultyAgentRpcTest,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import (
|
||||
RpcAgentTestFixture,
|
||||
)
|
||||
from torch.testing._internal.distributed.rpc.rpc_test import (
|
||||
CudaRpcTest,
|
||||
RpcTest,
|
||||
TensorPipeAgentCudaRpcTest,
|
||||
TensorPipeAgentRpcTest,
|
||||
)
|
||||
|
||||
|
||||
def _check_and_set_tcp_init():
|
||||
# if we are running with TCP init, set main address and port
|
||||
# before spawning subprocesses, since different processes could find
|
||||
# different ports.
|
||||
use_tcp_init = os.environ.get("RPC_INIT_WITH_TCP", None)
|
||||
if use_tcp_init == "1":
|
||||
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
||||
os.environ["MASTER_PORT"] = str(find_free_port())
|
||||
|
||||
|
||||
def _check_and_unset_tcp_init():
|
||||
use_tcp_init = os.environ.get("RPC_INIT_WITH_TCP", None)
|
||||
if use_tcp_init == "1":
|
||||
del os.environ["MASTER_ADDR"]
|
||||
del os.environ["MASTER_PORT"]
|
||||
|
||||
|
||||
# The tests for the RPC module need to cover multiple possible combinations:
|
||||
# - different aspects of the API, each one having its own suite of tests;
|
||||
# - different agents (ProcessGroup, TensorPipe, ...);
|
||||
# To avoid a combinatorial explosion in code size, and to prevent forgetting to
|
||||
# add a combination, these are generated automatically by the code in this file.
|
||||
# Here, we collect all the test suites that we need to cover.
|
||||
# We then have one separate file for each agent, from which
|
||||
# we call the generate_tests function of this file, passing to it a fixture for
|
||||
# the agent, which then gets mixed-in with each test suite.
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
TEST_WITH_DEV_DBG_ASAN,
|
||||
"Skip ASAN as torch + multiprocessing spawn have known issues",
|
||||
)
|
||||
class SpawnHelper(MultiProcessTestCase):
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
_check_and_set_tcp_init()
|
||||
self._spawn_processes()
|
||||
|
||||
def tearDown(self):
|
||||
_check_and_unset_tcp_init()
|
||||
super().tearDown()
|
||||
|
||||
|
||||
# This list contains test suites that are agent-agnostic and that only verify
|
||||
# compliance with the generic RPC interface specification. These tests should
|
||||
# *not* make use of implementation details of a specific agent (options,
|
||||
# attributes, ...). These test suites will be instantiated multiple times, once
|
||||
# for each agent (except the faulty agent, which is special).
|
||||
GENERIC_TESTS = [
|
||||
RpcTest,
|
||||
ParameterServerTest,
|
||||
DistAutogradTest,
|
||||
DistOptimizerTest,
|
||||
JitRpcTest,
|
||||
JitDistAutogradTest,
|
||||
RemoteModuleTest,
|
||||
ThreeWorkersRemoteModuleTest,
|
||||
DdpUnderDistAutogradTest,
|
||||
DdpComparisonTest,
|
||||
ReinforcementLearningRpcTest,
|
||||
]
|
||||
GENERIC_CUDA_TESTS = [
|
||||
CudaRpcTest,
|
||||
CudaDistAutogradTest,
|
||||
CudaRemoteModuleTest,
|
||||
CudaDdpComparisonTest,
|
||||
]
|
||||
|
||||
|
||||
# This list contains test suites that will only be run on the TensorPipeAgent.
|
||||
# These suites should be standalone, and separate from the ones in the generic
|
||||
# list (not subclasses of those!).
|
||||
TENSORPIPE_TESTS = [
|
||||
TensorPipeAgentRpcTest,
|
||||
TensorPipeAgentDistAutogradTest,
|
||||
]
|
||||
TENSORPIPE_CUDA_TESTS = [
|
||||
TensorPipeAgentCudaRpcTest,
|
||||
TensorPipeCudaDistAutogradTest,
|
||||
]
|
||||
|
||||
|
||||
# This list contains test suites that will only be run on the faulty RPC agent.
|
||||
# That agent is special as it's only used to perform fault injection in order to
|
||||
# verify the error handling behavior. Thus the faulty agent will only run the
|
||||
# suites in this list, which were designed to test such behaviors, and not the
|
||||
# ones in the generic list.
|
||||
FAULTY_AGENT_TESTS = [
|
||||
FaultyAgentRpcTest,
|
||||
FaultyAgentDistAutogradTest,
|
||||
JitFaultyAgentRpcTest,
|
||||
]
|
||||
|
||||
|
||||
def generate_tests(
|
||||
prefix: str,
|
||||
mixin: type[RpcAgentTestFixture],
|
||||
tests: list[type[RpcAgentTestFixture]],
|
||||
module_name: str,
|
||||
) -> dict[str, type[RpcAgentTestFixture]]:
|
||||
"""Mix in the classes needed to autogenerate the tests based on the params.
|
||||
|
||||
Takes a series of test suites, each written against a "generic" agent (i.e.,
|
||||
derived from the abstract RpcAgentTestFixture class), as the `tests` args.
|
||||
Takes a concrete subclass of RpcAgentTestFixture, which specializes it for a
|
||||
certain agent, as the `mixin` arg. Produces all combinations of them.
|
||||
Returns a dictionary of class names to class type
|
||||
objects which can be inserted into the global namespace of the calling
|
||||
module. The name of each test will be a concatenation of the `prefix` arg
|
||||
and the original name of the test suite.
|
||||
The `module_name` should be the name of the calling module so
|
||||
that the classes can be fixed to make it look like they belong to it, which
|
||||
is necessary for pickling to work on them.
|
||||
"""
|
||||
ret: dict[str, type[RpcAgentTestFixture]] = {}
|
||||
for test_class in tests:
|
||||
if IS_SANDCASTLE and TEST_WITH_DEV_DBG_ASAN:
|
||||
print(
|
||||
f"Skipping test {test_class} on sandcastle for the following reason: "
|
||||
"Skip dev-asan as torch + multiprocessing spawn have known issues",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
|
||||
name = f"{prefix}{test_class.__name__}"
|
||||
class_ = type(name, (test_class, mixin, SpawnHelper), {})
|
||||
class_.__module__ = module_name
|
||||
ret[name] = class_
|
||||
return ret
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
import torch
|
||||
import torch._dynamo.test_case
|
||||
import torch.utils._pytree as pytree
|
||||
|
||||
|
||||
class PytreeRegisteringTestCase(torch._dynamo.test_case.TestCase):
|
||||
"""TestCase that prunes all temporary pytree registrations and resets Dynamo."""
|
||||
|
||||
def setUp(self) -> None:
|
||||
super().setUp()
|
||||
self._registered_pytree_nodes: list[type] = []
|
||||
self._registered_constant_nodes: list[type] = []
|
||||
|
||||
def tearDown(self) -> None:
|
||||
for cls in reversed(self._registered_pytree_nodes):
|
||||
pytree._deregister_pytree_node(cls)
|
||||
for cls in reversed(self._registered_constant_nodes):
|
||||
pytree._deregister_pytree_node(cls)
|
||||
torch._dynamo.reset()
|
||||
super().tearDown()
|
||||
|
||||
def register_pytree_node(self, cls, *args, **kwargs) -> None: # type: ignore[no-untyped-def]
|
||||
pytree.register_pytree_node(cls, *args, **kwargs)
|
||||
self._registered_pytree_nodes.append(cls)
|
||||
|
||||
def register_constant(self, cls: type) -> None:
|
||||
pytree.register_constant(cls)
|
||||
self._registered_constant_nodes.append(cls)
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
"""
|
||||
This file contains the list of tests that are known to fail under Dynamo
|
||||
|
||||
We generate xFailIfTorchDynamo* for all tests in `dynamo_expected_failures`
|
||||
We generate skipIfTorchDynamo* for all tests in `dynamo_skips`
|
||||
We generate runWithoutCompiledAutograd for all tests in `compiled_autograd_skips`
|
||||
|
||||
For an easier-than-manual way of generating and updating these lists,
|
||||
see scripts/compile_tests/update_failures.py
|
||||
|
||||
If you're adding a new test, and it's failing PYTORCH_TEST_WITH_DYNAMO=1,
|
||||
either add the appropriate decorators to your test or add skips for them
|
||||
via test/dynamo_skips and test/dynamo_expected_failures.
|
||||
|
||||
*These are not exactly unittest.expectedFailure and unittest.skip. We'll
|
||||
always execute the test and then suppress the signal, if necessary.
|
||||
If your tests crashes, or is slow, please use @skipIfTorchDynamo instead.
|
||||
|
||||
The expected failure and skip files are located in test/dynamo_skips and
|
||||
test/dynamo_expected_failures. They're individual files rather than a list so
|
||||
git will merge changes easier.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def find_test_dir() -> str | None:
|
||||
# Find the path to the dynamo expected failure and skip files.
|
||||
from os.path import abspath, basename, dirname, exists, join, normpath
|
||||
|
||||
if sys.platform == "win32":
|
||||
return None
|
||||
|
||||
# Check relative to this file (local build):
|
||||
test_dir = normpath(join(dirname(abspath(__file__)), "../../../test"))
|
||||
if exists(join(test_dir, "dynamo_expected_failures")):
|
||||
return test_dir
|
||||
|
||||
# Check relative to __main__ (installed builds relative to test file):
|
||||
main = sys.modules["__main__"]
|
||||
file = getattr(main, "__file__", None)
|
||||
if file is None:
|
||||
# Generated files do not have a module.__file__
|
||||
return None
|
||||
test_dir = dirname(abspath(file))
|
||||
while dirname(test_dir) != test_dir:
|
||||
if basename(test_dir) == "test" and exists(
|
||||
join(test_dir, "dynamo_expected_failures")
|
||||
):
|
||||
return test_dir
|
||||
test_dir = dirname(test_dir)
|
||||
|
||||
# Not found
|
||||
return None
|
||||
|
||||
|
||||
test_dir = find_test_dir()
|
||||
if not test_dir:
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(
|
||||
"test/dynamo_expected_failures directory not found - known dynamo errors won't be skipped."
|
||||
)
|
||||
|
||||
# Tests that run without strict mode in PYTORCH_TEST_WITH_INDUCTOR=1.
|
||||
# Please don't add anything to this list.
|
||||
FIXME_inductor_non_strict = {
|
||||
"test_modules",
|
||||
"test_ops",
|
||||
"test_ops_gradients",
|
||||
"test_torch",
|
||||
}
|
||||
|
||||
# We generate unittest.expectedFailure for all of the following tests
|
||||
# when run under PYTORCH_TEST_WITH_DYNAMO=1.
|
||||
# see NOTE [dynamo_test_failures.py] for more details
|
||||
#
|
||||
# This lists exists so we can more easily add large numbers of failing tests,
|
||||
if test_dir is None:
|
||||
dynamo_expected_failures = set()
|
||||
dynamo_skips = set()
|
||||
|
||||
inductor_expected_failures = set()
|
||||
inductor_skips = set()
|
||||
|
||||
compiled_autograd_skips = set()
|
||||
else:
|
||||
dynamo_failures_directory = os.path.join(test_dir, "dynamo_expected_failures")
|
||||
dynamo_skips_directory = os.path.join(test_dir, "dynamo_skips")
|
||||
|
||||
dynamo_expected_failures = set(os.listdir(dynamo_failures_directory))
|
||||
dynamo_skips = set(os.listdir(dynamo_skips_directory))
|
||||
|
||||
inductor_failures_directory = os.path.join(test_dir, "inductor_expected_failures")
|
||||
inductor_skips_directory = os.path.join(test_dir, "inductor_skips")
|
||||
|
||||
inductor_expected_failures = set(os.listdir(inductor_failures_directory))
|
||||
inductor_skips = set(os.listdir(inductor_skips_directory))
|
||||
|
||||
compiled_autograd_skips_directory = os.path.join(
|
||||
test_dir, "compiled_autograd_skips"
|
||||
)
|
||||
compiled_autograd_skips = set(os.listdir(compiled_autograd_skips_directory))
|
||||
|
||||
# TODO: due to case sensitivity problems, for now list these files by hand
|
||||
extra_dynamo_skips = {
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_exhaustive_T_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_exhaustive_t_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_fake_exhaustive_T_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_fake_exhaustive_t_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_T_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_t_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_inplace_T_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_inplace_t_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_out_T_cpu_float32",
|
||||
"TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_out_t_cpu_float32",
|
||||
}
|
||||
dynamo_skips = dynamo_skips.union(extra_dynamo_skips)
|
||||
|
||||
|
||||
# verify some invariants
|
||||
for test in (
|
||||
dynamo_expected_failures
|
||||
| dynamo_skips
|
||||
| inductor_expected_failures
|
||||
| inductor_skips
|
||||
):
|
||||
if len(test.split(".")) != 2:
|
||||
raise AssertionError(f'Invalid test name: "{test}"')
|
||||
|
||||
dynamo_intersection = dynamo_expected_failures.intersection(dynamo_skips)
|
||||
if len(dynamo_intersection) > 0:
|
||||
raise AssertionError(
|
||||
"there should be no overlap between dynamo_expected_failures "
|
||||
"and dynamo_skips, got " + str(dynamo_intersection)
|
||||
)
|
||||
|
||||
inductor_intersection = inductor_expected_failures.intersection(inductor_skips)
|
||||
if len(inductor_intersection) > 0:
|
||||
raise AssertionError(
|
||||
"there should be no overlap between inductor_expected_failures "
|
||||
"and inductor_skips, got " + str(inductor_intersection)
|
||||
)
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
import sys
|
||||
|
||||
from torch.utils._config_module import Config, install_config_module
|
||||
|
||||
|
||||
e_bool = True
|
||||
e_int = 1
|
||||
e_float = 1.0
|
||||
e_string = "string"
|
||||
e_list = [1]
|
||||
e_set = {1}
|
||||
e_tuple = (1,)
|
||||
e_dict = {1: 2}
|
||||
e_none: bool | None = None
|
||||
e_optional: bool | None = True
|
||||
e_ignored = True
|
||||
_e_ignored = True
|
||||
magic_cache_config_ignored = True
|
||||
# [@compile_ignored: debug]
|
||||
e_compile_ignored = True
|
||||
e_config: bool = Config(default=True)
|
||||
e_jk: bool = Config(justknob="does_not_exist", default=True)
|
||||
e_jk_false: bool = Config(justknob="does_not_exist", default=False)
|
||||
e_env_default: bool = Config(env_name_default="ENV_TRUE", default=False)
|
||||
e_env_default_FALSE: bool = Config(env_name_default="ENV_FALSE", default=True)
|
||||
e_env_default_str: bool = Config(env_name_default="ENV_STR", default="default")
|
||||
e_env_default_str_empty: bool = Config(
|
||||
env_name_default="ENV_STR_EMPTY", default="default"
|
||||
)
|
||||
e_env_force: bool = Config(env_name_force="ENV_TRUE", default=False)
|
||||
e_aliased_bool: bool = Config(
|
||||
alias="torch.testing._internal.fake_config_module2.e_aliasing_bool"
|
||||
)
|
||||
e_deprecated: bool = Config(
|
||||
default=True, deprecated=True, deprecation_message="is no longer needed"
|
||||
)
|
||||
e_not_deprecated: bool = Config(default=False)
|
||||
e_deprecated_alias: bool = Config(
|
||||
alias="torch.testing._internal.fake_config_module.e_not_deprecated",
|
||||
deprecated=True,
|
||||
deprecation_message="use something else instead",
|
||||
)
|
||||
|
||||
|
||||
class nested:
|
||||
e_bool = True
|
||||
|
||||
|
||||
_cache_config_ignore_prefix = ["magic_cache_config"]
|
||||
_save_config_ignore = ["e_ignored"]
|
||||
|
||||
install_config_module(sys.modules[__name__])
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
import sys
|
||||
|
||||
from torch.utils._config_module import Config, install_config_module
|
||||
|
||||
|
||||
e_aliasing_bool = False
|
||||
|
||||
e_env_default_multi: bool = Config(
|
||||
env_name_default=["ENV_TRUE", "ENV_FALSE"], default=False
|
||||
)
|
||||
e_env_force_multi: bool = Config(env_name_force=["ENV_FAKE", "ENV_TRUE"], default=False)
|
||||
|
||||
install_config_module(sys.modules[__name__])
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
import sys
|
||||
from typing import Callable # noqa: UP035
|
||||
|
||||
from torch.utils._config_module import install_config_module
|
||||
|
||||
|
||||
e_list = [1]
|
||||
e_set = {1}
|
||||
e_func: Callable | None = None
|
||||
|
||||
install_config_module(sys.modules[__name__])
|
||||
+2923
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,684 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import functools
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from functorch.experimental.control_flow import map
|
||||
from torch._higher_order_ops.flex_attention import (
|
||||
flex_attention as flex_attention_hop,
|
||||
)
|
||||
from torch.nn.attention.flex_attention import (
|
||||
_create_empty_block_mask,
|
||||
create_block_mask,
|
||||
flex_attention,
|
||||
)
|
||||
from torch.testing import make_tensor
|
||||
from torch._higher_order_ops.inline_asm_elementwise import inline_asm_elementwise
|
||||
from torch.testing._internal.common_device_type import onlyCUDA
|
||||
from torch.testing._internal.common_dtype import all_types_and, custom_types
|
||||
from torch.testing._internal.opinfo.core import DecorateInfo, OpInfo, SampleInput
|
||||
from torch._higher_order_ops.invoke_subgraph import mark_compile_region
|
||||
from torch._higher_order_ops import InvokeQuant, invoke_quant_packed
|
||||
|
||||
|
||||
def sample_inputs_map(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
|
||||
)
|
||||
yield SampleInput(
|
||||
[make_arg(2, 2, 2, low=0.1, high=2), make_arg(2, 2, 2, low=0.1, high=2)],
|
||||
args=(make_arg(1, low=0.1, high=2), make_arg(1, low=0.1, high=2)),
|
||||
)
|
||||
|
||||
|
||||
def inner_f(x, y0, y1):
|
||||
return [x[0].cos().add_(1.0) * y0, (x[1] + y1.sin()).cos_().view(x[1].size())]
|
||||
|
||||
|
||||
def simple_map(xs, y0, y1):
|
||||
def f(x, y0, y1):
|
||||
return inner_f(x, y0, y1)
|
||||
|
||||
return map(f, xs, y0, y1)
|
||||
|
||||
|
||||
def nested_map(xs, y0, y1):
|
||||
def f1(xx, y0, y1):
|
||||
def f2(x, y0, y1):
|
||||
return inner_f(x, y0, y1)
|
||||
|
||||
return map(f2, xx, y0, y1)
|
||||
|
||||
return map(f1, xs, y0, y1)
|
||||
|
||||
|
||||
def triple_nested_map(xs, y0, y1):
|
||||
def f0(xs, y0, y1):
|
||||
def f1(xx, y0, y1):
|
||||
def f2(x, y0, y1):
|
||||
return inner_f(x, y0, y1)
|
||||
|
||||
return map(f2, xx, y0, y1)
|
||||
|
||||
return map(f1, xs, y0, y1)
|
||||
|
||||
return map(f0, xs, y0, y1)
|
||||
|
||||
|
||||
# PLEASE DON'T ADD ANYTHING NEW TO THIS LIST,
|
||||
# and do add an OpInfo for your HOP.
|
||||
# The OpInfo lets us do automated testing for the HOP to check that
|
||||
# your HOP will work correctly with PyTorch!
|
||||
#
|
||||
# Your new HOP may fail some automated testing. That's OK. If you don't
|
||||
# care about certain features (like torch.export), it's fine to xfail those
|
||||
# failing tests. It is less fine to xfail a more critical check (like checking
|
||||
# if torch.compile works with your HOP, or if your HOP has a docstring).
|
||||
# If you don't know if a test is fine to xfail, please ask.
|
||||
#
|
||||
# There are legitimate reasons why something cannot be added to this list
|
||||
# (e.g. it uses executorch which is not in PyTorch). If that's the case then
|
||||
# please leave a comment.
|
||||
FIXME_hop_that_doesnt_have_opinfo_test_allowlist = [
|
||||
"custom_function_call",
|
||||
"autograd_function_apply",
|
||||
"run_and_save_rng_state",
|
||||
"run_with_rng_state",
|
||||
"run_dtensor_rng_op",
|
||||
"graphsafe_run_with_rng_state",
|
||||
"out_dtype",
|
||||
"trace_wrapped",
|
||||
'tag_activation_checkpoint',
|
||||
'executorch_call_delegate',
|
||||
'wrap',
|
||||
'wrap_with_set_grad_enabled',
|
||||
'auto_functionalized_v2',
|
||||
'associative_scan',
|
||||
'flat_apply', # is WIP, doesn't pass any of the tests yet
|
||||
'wrap_with_autocast',
|
||||
'wrap_activation_checkpoint',
|
||||
'run_const_graph',
|
||||
'auto_functionalized',
|
||||
"map", # T183144629
|
||||
"map_impl",
|
||||
"with_effects",
|
||||
"strict_mode",
|
||||
"_export_tracepoint",
|
||||
"call_torchbind",
|
||||
"triton_kernel_wrapper_mutation",
|
||||
"triton_kernel_wrapper_functional",
|
||||
"hints_wrapper",
|
||||
"dynamo_bypassing_wrapper", # TODO(soulitzer)
|
||||
"foreach_map",
|
||||
"aoti_call_delegate",
|
||||
"print",
|
||||
"inductor_compiled_code", # Tested separately in test_inductor_wrap_inductor_compile_regions
|
||||
"invoke_leaf_function", # Needs torch.compile, tested separately in test_leaf_function*
|
||||
]
|
||||
|
||||
torch.library.define(
|
||||
"testlib::mutating_custom_op",
|
||||
"(Tensor(a!) x, Tensor(b!) z) -> (Tensor, Tensor, Tensor)",
|
||||
tags=torch.Tag.pt2_compliant_tag,
|
||||
)
|
||||
|
||||
|
||||
@torch.library.impl("testlib::mutating_custom_op", "cpu")
|
||||
def foo_impl_cpu(x, z):
|
||||
x.add_(5)
|
||||
z.add_(5)
|
||||
return x.clone(), z.clone(), x + z
|
||||
|
||||
|
||||
@torch.library.impl("testlib::mutating_custom_op", "cuda")
|
||||
def foo_impl_cuda(x, z):
|
||||
x.add_(5)
|
||||
z.add_(5)
|
||||
return x.clone(), z.clone(), x + z
|
||||
|
||||
|
||||
@torch.library.impl("testlib::mutating_custom_op", "xpu")
|
||||
def foo_impl_xpu(x, z):
|
||||
x.add_(5)
|
||||
z.add_(5)
|
||||
return x.clone(), z.clone(), x + z
|
||||
|
||||
|
||||
@torch.library.register_fake("testlib::mutating_custom_op")
|
||||
def foo_impl_abstract(x, z):
|
||||
return x.clone(), z.clone(), x + z
|
||||
|
||||
|
||||
def sample_inputs_cond(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
|
||||
)
|
||||
yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2))
|
||||
|
||||
|
||||
def simple_cond(x):
|
||||
return torch.cond(x.sum() > 2, lambda x: (x.cos(),), lambda x: (x.sin(),), [x])
|
||||
|
||||
|
||||
def sample_inputs_invoke_subgraph(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
|
||||
)
|
||||
yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2))
|
||||
|
||||
|
||||
@mark_compile_region
|
||||
def fn_for_invoke_subgraph(x):
|
||||
return torch.sin(x)
|
||||
|
||||
|
||||
def simple_invoke_subgraph(x):
|
||||
return fn_for_invoke_subgraph(x)
|
||||
|
||||
|
||||
def sample_inputs_auto_functionalize(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=False
|
||||
)
|
||||
yield SampleInput(
|
||||
make_arg(2, 2, 2, low=0.1, high=2), make_arg(2, 2, 2, low=0.1, high=2)
|
||||
)
|
||||
|
||||
|
||||
def simple_auto_functionalize(x, z):
|
||||
return torch.ops.testlib.mutating_custom_op(x, z)
|
||||
|
||||
|
||||
def sample_inputs_flex_attention(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
|
||||
)
|
||||
|
||||
def score_mod(score, b, h, m, n):
|
||||
return score + h
|
||||
|
||||
q, k, v = (make_arg(2, 2, 128, 8, low=0.1, high=2) for _ in range(3))
|
||||
block_mask = _create_empty_block_mask(q, k)
|
||||
yield SampleInput(q, k, v, score_mod, block_mask)
|
||||
|
||||
|
||||
def sample_inputs_flex_attention_backward(
|
||||
opinfo, device, dtype, requires_grad, **kwargs
|
||||
):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=False
|
||||
)
|
||||
|
||||
def score_mod(score, b, h, m, n):
|
||||
return score
|
||||
|
||||
def mask_mod(b, h, m, n):
|
||||
return m >= n
|
||||
|
||||
q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3))
|
||||
block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device)
|
||||
scale = 1.0 / q.size(-1) ** 0.5
|
||||
out, logsumexp, _ = flex_attention_hop(
|
||||
q, k, v, score_mod, block_mask.as_tuple(), scale, {},
|
||||
)
|
||||
yield SampleInput(
|
||||
q,
|
||||
args=(
|
||||
k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None,
|
||||
score_mod, None, block_mask.as_tuple(),
|
||||
scale, {}, (), (),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def sample_inputs_flex_attention_backward_explicit_buffers(
|
||||
opinfo, device, dtype, requires_grad, **kwargs
|
||||
):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=False
|
||||
)
|
||||
mask_offset = torch.full((), 128, device=device, dtype=torch.int32)
|
||||
|
||||
def score_mod(score, b, h, m, n):
|
||||
return score
|
||||
|
||||
def mask_mod(b, h, m, n):
|
||||
return m + mask_offset >= n
|
||||
|
||||
q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3))
|
||||
block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device)
|
||||
scale = 1.0 / q.size(-1) ** 0.5
|
||||
out, logsumexp, _ = flex_attention_hop(
|
||||
q, k, v, score_mod, block_mask.as_tuple(), scale, {},
|
||||
)
|
||||
yield SampleInput(
|
||||
q,
|
||||
args=(
|
||||
k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None,
|
||||
score_mod, None, block_mask.as_tuple(),
|
||||
scale, {}, (), (),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def simple_flex_attention_backward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
grad_out,
|
||||
grad_logsumexp,
|
||||
fw_graph,
|
||||
joint_graph,
|
||||
block_mask,
|
||||
scale,
|
||||
kernel_options,
|
||||
score_mod_other_buffers,
|
||||
mask_mod_other_buffers,
|
||||
):
|
||||
return torch.ops.higher_order.flex_attention_backward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
grad_out,
|
||||
grad_logsumexp,
|
||||
fw_graph,
|
||||
joint_graph,
|
||||
block_mask,
|
||||
scale,
|
||||
kernel_options,
|
||||
score_mod_other_buffers,
|
||||
mask_mod_other_buffers,
|
||||
)
|
||||
|
||||
|
||||
def sample_inputs_while_loop(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=False
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor(3),
|
||||
make_arg(2, 3, 4, low=0.1, high=2),
|
||||
)
|
||||
|
||||
|
||||
def simple_while_loop(iter_t, x):
|
||||
def cond_fn(iter_t, x):
|
||||
return iter_t > 0
|
||||
|
||||
def body_fn(iter_t, x):
|
||||
return iter_t - 1, x.cos()
|
||||
|
||||
return torch._higher_order_ops.while_loop(cond_fn, body_fn, (iter_t, x))
|
||||
|
||||
|
||||
def simple_while_loop_stack_output(iter_t, x):
|
||||
def cond_fn(iter_t, x):
|
||||
return iter_t > 0
|
||||
|
||||
def body_fn(iter_t, x):
|
||||
return iter_t - 1, x.cos()
|
||||
|
||||
return torch._higher_order_ops.while_loop_stack_output(
|
||||
cond_fn, body_fn, (iter_t, x), tuple()
|
||||
)
|
||||
|
||||
|
||||
def sample_inputs_local_map_hop(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
# TODO: once HOPs support DTensor inputs, we should also test DTensors
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=False
|
||||
)
|
||||
yield SampleInput(
|
||||
make_arg(2, 3, 4, low=0.1, high=2),
|
||||
make_arg(2, 3, 4, low=0.1, high=2),
|
||||
)
|
||||
|
||||
|
||||
def simple_local_map_hop(inp1, inp2):
|
||||
def body_gm(inp1, inp2):
|
||||
return inp1.cos() + inp2.sin()
|
||||
|
||||
gm = torch.fx.symbolic_trace(body_gm)
|
||||
|
||||
if not torch.distributed.is_available():
|
||||
raise AssertionError("Expected torch.distributed to be available")
|
||||
from torch.distributed.tensor.placement_types import Replicate
|
||||
|
||||
gm.meta["local_map_kwargs"] = {
|
||||
"in_placements": (Replicate(), Replicate(), Replicate()),
|
||||
"out_placements": ((Replicate(), Replicate(), Replicate()),),
|
||||
}
|
||||
|
||||
# TODO: Dynamo would rewrite this op differently
|
||||
return torch._higher_order_ops.local_map_hop(gm, inp1, inp2)
|
||||
|
||||
|
||||
def sample_inputs_scan(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
|
||||
)
|
||||
yield SampleInput(
|
||||
make_arg(2, 2, low=0.1, high=2),
|
||||
make_arg(2, 2, 2, low=0.1, high=2),
|
||||
)
|
||||
|
||||
|
||||
def simple_scan(init, xs):
|
||||
def combine_fn(carry, x):
|
||||
result = carry @ x + x
|
||||
return result, carry.clone()
|
||||
|
||||
return torch._higher_order_ops.scan(combine_fn, init, xs)
|
||||
|
||||
|
||||
quant_tracer = InvokeQuant()
|
||||
|
||||
|
||||
def simple_invoke_quant(x):
|
||||
def fn(x, y):
|
||||
return (torch.sin(x) * y,)
|
||||
|
||||
return quant_tracer(fn, x, x)[0] * 2.0
|
||||
|
||||
|
||||
def simple_invoke_quant_packed(x):
|
||||
def fn(x):
|
||||
return (torch.sin(x),)
|
||||
|
||||
return invoke_quant_packed(fn, x)[0] * 2.0
|
||||
|
||||
|
||||
def sample_inputs_inline_asm(opinfo, device, dtype, requires_grad, **kwargs):
|
||||
make_arg = functools.partial(
|
||||
make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
|
||||
)
|
||||
yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2))
|
||||
|
||||
|
||||
def simple_inline_asm(x):
|
||||
if torch.version.hip:
|
||||
return inline_asm_elementwise(
|
||||
x,
|
||||
asm_str="v_mov_b32_e32 $0, $1",
|
||||
constraints="=v, v",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
return inline_asm_elementwise(
|
||||
x, asm_str="mov.f32 $0, $1;", constraints="=f,f", dtype=torch.float32
|
||||
)
|
||||
|
||||
|
||||
hop_db = [
|
||||
OpInfo(
|
||||
name="scan",
|
||||
variant_test_name="simple",
|
||||
op=simple_scan,
|
||||
sample_inputs_func=sample_inputs_scan,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=False,
|
||||
# "torch.compile with aot_autograd does not currently support double backward."
|
||||
supports_gradgrad=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="invoke_subgraph",
|
||||
variant_test_name="simple",
|
||||
op=simple_invoke_subgraph,
|
||||
sample_inputs_func=sample_inputs_invoke_subgraph,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=True,
|
||||
# "torch.compile with aot_autograd does not currently support double backward."
|
||||
supports_gradgrad=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="map",
|
||||
variant_test_name="simple",
|
||||
op=simple_map,
|
||||
sample_inputs_func=sample_inputs_map,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="map",
|
||||
variant_test_name="nested",
|
||||
op=nested_map,
|
||||
sample_inputs_func=sample_inputs_map,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="map",
|
||||
variant_test_name="triple_nested",
|
||||
op=triple_nested_map,
|
||||
sample_inputs_func=sample_inputs_map,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="cond",
|
||||
variant_test_name="simple",
|
||||
op=simple_cond,
|
||||
sample_inputs_func=sample_inputs_cond,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=True,
|
||||
# "torch.compile with aot_autograd does not currently support double backward."
|
||||
supports_gradgrad=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="invoke_quant",
|
||||
variant_test_name="simple",
|
||||
op=simple_invoke_quant,
|
||||
sample_inputs_func=sample_inputs_invoke_subgraph,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=True,
|
||||
# "torch.compile with aot_autograd does not currently support double backward."
|
||||
skips=(
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
|
||||
),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
|
||||
),
|
||||
# "torch.compile with aot_autograd does not currently support double backward."
|
||||
supports_gradgrad=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="invoke_quant_packed",
|
||||
variant_test_name="simple",
|
||||
op=simple_invoke_quant_packed,
|
||||
sample_inputs_func=sample_inputs_invoke_subgraph,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=True,
|
||||
# "torch.compile with aot_autograd does not currently support double backward."
|
||||
supports_gradgrad=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="while_loop",
|
||||
variant_test_name="simple",
|
||||
op=simple_while_loop,
|
||||
sample_inputs_func=sample_inputs_while_loop,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="while_loop_stack_output",
|
||||
variant_test_name="simple",
|
||||
op=simple_while_loop_stack_output,
|
||||
sample_inputs_func=sample_inputs_while_loop,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="auto_functionalize",
|
||||
variant_test_name="simple",
|
||||
op=simple_auto_functionalize,
|
||||
sample_inputs_func=sample_inputs_auto_functionalize,
|
||||
dtypes=all_types_and(torch.bool, torch.half),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=False,
|
||||
),
|
||||
OpInfo(
|
||||
name="flex_attention",
|
||||
variant_test_name="simple",
|
||||
op=flex_attention,
|
||||
sample_inputs_func=sample_inputs_flex_attention,
|
||||
dtypes=custom_types(torch.float16, torch.float32),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
skips=(
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
|
||||
),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
|
||||
),
|
||||
decorators=[onlyCUDA],
|
||||
),
|
||||
OpInfo(
|
||||
name="flex_attention_backward",
|
||||
variant_test_name="simple",
|
||||
op=simple_flex_attention_backward,
|
||||
sample_inputs_func=sample_inputs_flex_attention_backward,
|
||||
dtypes=custom_types(torch.float16, torch.float32),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=False,
|
||||
supports_gradgrad=False,
|
||||
skips=(
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
|
||||
),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
|
||||
),
|
||||
decorators=[onlyCUDA],
|
||||
),
|
||||
OpInfo(
|
||||
name="flex_attention_backward",
|
||||
variant_test_name="explicit_buffers",
|
||||
op=simple_flex_attention_backward,
|
||||
sample_inputs_func=sample_inputs_flex_attention_backward_explicit_buffers,
|
||||
dtypes=custom_types(torch.float16, torch.float32),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=False,
|
||||
supports_gradgrad=False,
|
||||
skips=(
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
|
||||
),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
|
||||
),
|
||||
decorators=[onlyCUDA],
|
||||
),
|
||||
OpInfo(
|
||||
name="local_map_hop",
|
||||
variant_test_name="simple",
|
||||
op=simple_local_map_hop,
|
||||
sample_inputs_func=sample_inputs_local_map_hop,
|
||||
dtypes=custom_types(torch.float16, torch.float32),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
skips=(
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
|
||||
),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
|
||||
DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
|
||||
),
|
||||
decorators=[
|
||||
onlyCUDA,
|
||||
unittest.skipIf(
|
||||
not torch.distributed.is_available(), "requires distributed build"
|
||||
),
|
||||
],
|
||||
),
|
||||
OpInfo(
|
||||
name="inline_asm_elementwise",
|
||||
variant_test_name="simple",
|
||||
op=simple_inline_asm,
|
||||
sample_inputs_func=sample_inputs_inline_asm,
|
||||
dtypes=custom_types(torch.float32),
|
||||
supports_out=False,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
check_batched_forward_grad=False,
|
||||
check_inplace_batched_forward_grad=False,
|
||||
supports_autograd=False,
|
||||
decorators=[onlyCUDA],
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,383 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterable
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import hypothesis
|
||||
from functools import reduce
|
||||
from importlib.metadata import version
|
||||
from hypothesis import assume
|
||||
from hypothesis import settings
|
||||
from hypothesis import strategies as st
|
||||
from hypothesis.extra import numpy as stnp
|
||||
from hypothesis.strategies import SearchStrategy
|
||||
|
||||
from torch.testing._internal.common_quantized import _calculate_dynamic_qparams, _calculate_dynamic_per_channel_qparams
|
||||
|
||||
# Setup for the hypothesis tests.
|
||||
# The tuples are (torch_quantized_dtype, zero_point_enforce), where the last
|
||||
# element is enforced zero_point. If None, any zero_point point within the
|
||||
# range of the data type is OK.
|
||||
|
||||
# Tuple with all quantized data types.
|
||||
_ALL_QINT_TYPES = (
|
||||
torch.quint8,
|
||||
torch.qint8,
|
||||
torch.qint32,
|
||||
)
|
||||
|
||||
# Enforced zero point for every quantized data type.
|
||||
# If None, any zero_point point within the range of the data type is OK.
|
||||
_ENFORCED_ZERO_POINT = defaultdict(lambda: None, {
|
||||
torch.quint8: None,
|
||||
torch.qint8: None,
|
||||
torch.qint32: 0
|
||||
})
|
||||
|
||||
def _get_valid_min_max(qparams):
|
||||
scale, zero_point, _quantized_type = qparams
|
||||
adjustment = 1 + torch.finfo(torch.float).eps
|
||||
_long_type_info = torch.iinfo(torch.long)
|
||||
long_min, long_max = _long_type_info.min / adjustment, _long_type_info.max / adjustment
|
||||
# make sure intermediate results are within the range of long
|
||||
min_value = max((long_min - zero_point) * scale, (long_min / scale + zero_point))
|
||||
max_value = min((long_max - zero_point) * scale, (long_max / scale + zero_point))
|
||||
return np.float32(min_value), np.float32(max_value)
|
||||
|
||||
# This wrapper wraps around `st.floats` and checks the version of `hypothesis`, if
|
||||
# it is too old, removes the `width` parameter (which was introduced)
|
||||
# in 3.67.0
|
||||
def _floats_wrapper(*args, **kwargs):
|
||||
if 'width' in kwargs and hypothesis.version.__version_info__ < (3, 67, 0):
|
||||
# As long as nan, inf, min, max are not specified, reimplement the width
|
||||
# parameter for older versions of hypothesis.
|
||||
no_nan_and_inf = (
|
||||
(('allow_nan' in kwargs and not kwargs['allow_nan']) or
|
||||
'allow_nan' not in kwargs) and
|
||||
(('allow_infinity' in kwargs and not kwargs['allow_infinity']) or
|
||||
'allow_infinity' not in kwargs))
|
||||
min_and_max_not_specified = (
|
||||
len(args) == 0 and
|
||||
'min_value' not in kwargs and
|
||||
'max_value' not in kwargs
|
||||
)
|
||||
if no_nan_and_inf and min_and_max_not_specified:
|
||||
if kwargs['width'] == 16:
|
||||
kwargs['min_value'] = torch.finfo(torch.float16).min
|
||||
kwargs['max_value'] = torch.finfo(torch.float16).max
|
||||
elif kwargs['width'] == 32:
|
||||
kwargs['min_value'] = torch.finfo(torch.float32).min
|
||||
kwargs['max_value'] = torch.finfo(torch.float32).max
|
||||
elif kwargs['width'] == 64:
|
||||
kwargs['min_value'] = torch.finfo(torch.float64).min
|
||||
kwargs['max_value'] = torch.finfo(torch.float64).max
|
||||
kwargs.pop('width')
|
||||
return st.floats(*args, **kwargs)
|
||||
|
||||
def floats(*args, **kwargs):
|
||||
if 'width' not in kwargs:
|
||||
kwargs['width'] = 32
|
||||
return _floats_wrapper(*args, **kwargs)
|
||||
|
||||
"""Hypothesis filter to avoid overflows with quantized tensors.
|
||||
|
||||
Args:
|
||||
tensor: Tensor of floats to filter
|
||||
qparams: Quantization parameters as returned by the `qparams`.
|
||||
|
||||
Returns:
|
||||
True
|
||||
|
||||
Raises:
|
||||
hypothesis.UnsatisfiedAssumption
|
||||
|
||||
Note: This filter is slow. Use it only when filtering of the test cases is
|
||||
absolutely necessary!
|
||||
"""
|
||||
def assume_not_overflowing(tensor, qparams):
|
||||
min_value, max_value = _get_valid_min_max(qparams)
|
||||
assume(tensor.min() >= min_value)
|
||||
assume(tensor.max() <= max_value)
|
||||
return True
|
||||
|
||||
"""Strategy for generating the quantization parameters.
|
||||
|
||||
Args:
|
||||
dtypes: quantized data types to sample from.
|
||||
scale_min / scale_max: Min and max scales. If None, set to 1e-3 / 1e3.
|
||||
zero_point_min / zero_point_max: Min and max for the zero point. If None,
|
||||
set to the minimum and maximum of the quantized data type.
|
||||
Note: The min and max are only valid if the zero_point is not enforced
|
||||
by the data type itself.
|
||||
|
||||
Generates:
|
||||
scale: Sampled scale.
|
||||
zero_point: Sampled zero point.
|
||||
quantized_type: Sampled quantized type.
|
||||
"""
|
||||
@st.composite
|
||||
def qparams(draw, dtypes=None, scale_min=None, scale_max=None,
|
||||
zero_point_min=None, zero_point_max=None):
|
||||
if dtypes is None:
|
||||
dtypes = _ALL_QINT_TYPES
|
||||
if not isinstance(dtypes, (list, tuple)):
|
||||
dtypes = (dtypes,)
|
||||
quantized_type = draw(st.sampled_from(dtypes))
|
||||
|
||||
_type_info = torch.iinfo(quantized_type)
|
||||
qmin, qmax = _type_info.min, _type_info.max
|
||||
|
||||
# TODO: Maybe embed the enforced zero_point in the `torch.iinfo`.
|
||||
_zp_enforced = _ENFORCED_ZERO_POINT[quantized_type]
|
||||
if _zp_enforced is not None:
|
||||
zero_point = _zp_enforced
|
||||
else:
|
||||
_zp_min = qmin if zero_point_min is None else zero_point_min
|
||||
_zp_max = qmax if zero_point_max is None else zero_point_max
|
||||
zero_point = draw(st.integers(min_value=_zp_min, max_value=_zp_max))
|
||||
|
||||
if scale_min is None:
|
||||
scale_min = torch.finfo(torch.float).eps
|
||||
if scale_max is None:
|
||||
scale_max = torch.finfo(torch.float).max
|
||||
scale = draw(floats(min_value=scale_min, max_value=scale_max, width=32))
|
||||
|
||||
return scale, zero_point, quantized_type
|
||||
|
||||
"""Strategy to create different shapes.
|
||||
Args:
|
||||
min_dims / max_dims: minimum and maximum rank.
|
||||
min_side / max_side: minimum and maximum dimensions per rank.
|
||||
|
||||
Generates:
|
||||
Possible shapes for a tensor, constrained to the rank and dimensionality.
|
||||
|
||||
Example:
|
||||
# Generates 3D and 4D tensors.
|
||||
@given(Q = qtensor(shapes=array_shapes(min_dims=3, max_dims=4))
|
||||
some_test(self, Q):...
|
||||
"""
|
||||
@st.composite
|
||||
def array_shapes(draw, min_dims=1, max_dims=None, min_side=1, max_side=None, max_numel=None):
|
||||
"""Return a strategy for array shapes (tuples of int >= 1)."""
|
||||
if min_dims >= 32:
|
||||
raise AssertionError(f"Expected min_dims < 32, got {min_dims}")
|
||||
if max_dims is None:
|
||||
max_dims = min(min_dims + 2, 32)
|
||||
if max_dims >= 32:
|
||||
raise AssertionError(f"Expected max_dims < 32, got {max_dims}")
|
||||
if max_side is None:
|
||||
max_side = min_side + 5
|
||||
candidate = st.lists(st.integers(min_side, max_side), min_size=min_dims, max_size=max_dims)
|
||||
if max_numel is not None:
|
||||
candidate = candidate.filter(lambda x: reduce(int.__mul__, x, 1) <= max_numel)
|
||||
return draw(candidate.map(tuple))
|
||||
|
||||
|
||||
"""Strategy for generating test cases for tensors.
|
||||
The resulting tensor is in float32 format.
|
||||
|
||||
Args:
|
||||
shapes: Shapes under test for the tensor. Could be either a hypothesis
|
||||
strategy, or an iterable of different shapes to sample from.
|
||||
elements: Elements to generate from for the returned data type.
|
||||
If None, the strategy resolves to float within range [-1e6, 1e6].
|
||||
qparams: Instance of the qparams strategy. This is used to filter the tensor
|
||||
such that the overflow would not happen.
|
||||
|
||||
Generates:
|
||||
X: Tensor of type float32. Note that NaN and +/-inf is not included.
|
||||
qparams: (If `qparams` arg is set) Quantization parameters for X.
|
||||
The returned parameters are `(scale, zero_point, quantization_type)`.
|
||||
(If `qparams` arg is None), returns None.
|
||||
"""
|
||||
@st.composite
|
||||
def tensor(draw, shapes=None, elements=None, qparams=None, dtype=np.float32):
|
||||
if isinstance(shapes, SearchStrategy):
|
||||
_shape = draw(shapes)
|
||||
else:
|
||||
_shape = draw(st.sampled_from(shapes))
|
||||
if qparams is None:
|
||||
if elements is None:
|
||||
elements = floats(-1e6, 1e6, allow_nan=False, width=32)
|
||||
X = draw(stnp.arrays(dtype=dtype, elements=elements, shape=_shape))
|
||||
assume(not (np.isnan(X).any() or np.isinf(X).any()))
|
||||
return X, None
|
||||
qparams = draw(qparams)
|
||||
if elements is None:
|
||||
min_value, max_value = _get_valid_min_max(qparams)
|
||||
elements = floats(min_value, max_value, allow_infinity=False,
|
||||
allow_nan=False, width=32)
|
||||
X = draw(stnp.arrays(dtype=dtype, elements=elements, shape=_shape))
|
||||
# Recompute the scale and zero_points according to the X statistics.
|
||||
scale, zp = _calculate_dynamic_qparams(X, qparams[2])
|
||||
enforced_zp = _ENFORCED_ZERO_POINT.get(qparams[2], None)
|
||||
if enforced_zp is not None:
|
||||
zp = enforced_zp
|
||||
return X, (scale, zp, qparams[2])
|
||||
|
||||
@st.composite
|
||||
def per_channel_tensor(draw, shapes=None, elements=None, qparams=None):
|
||||
if isinstance(shapes, SearchStrategy):
|
||||
_shape = draw(shapes)
|
||||
else:
|
||||
_shape = draw(st.sampled_from(shapes))
|
||||
if qparams is None:
|
||||
if elements is None:
|
||||
elements = floats(-1e6, 1e6, allow_nan=False, width=32)
|
||||
X = draw(stnp.arrays(dtype=np.float32, elements=elements, shape=_shape))
|
||||
assume(not (np.isnan(X).any() or np.isinf(X).any()))
|
||||
return X, None
|
||||
qparams = draw(qparams)
|
||||
if elements is None:
|
||||
min_value, max_value = _get_valid_min_max(qparams)
|
||||
elements = floats(min_value, max_value, allow_infinity=False,
|
||||
allow_nan=False, width=32)
|
||||
X = draw(stnp.arrays(dtype=np.float32, elements=elements, shape=_shape))
|
||||
# Recompute the scale and zero_points according to the X statistics.
|
||||
scale, zp = _calculate_dynamic_per_channel_qparams(X, qparams[2])
|
||||
enforced_zp = _ENFORCED_ZERO_POINT.get(qparams[2], None)
|
||||
if enforced_zp is not None:
|
||||
zp = enforced_zp
|
||||
# Permute to model quantization along an axis
|
||||
axis = int(np.random.randint(0, X.ndim, 1))
|
||||
permute_axes = np.arange(X.ndim)
|
||||
permute_axes[0] = axis
|
||||
permute_axes[axis] = 0
|
||||
X = np.transpose(X, permute_axes)
|
||||
|
||||
return X, (scale, zp, axis, qparams[2])
|
||||
|
||||
"""Strategy for generating test cases for tensors used in Conv.
|
||||
The resulting tensors is in float32 format.
|
||||
|
||||
Args:
|
||||
spatial_dim: Spatial Dim for feature maps. If given as an iterable, randomly
|
||||
picks one from the pool to make it the spatial dimension
|
||||
batch_size_range: Range to generate `batch_size`.
|
||||
Must be tuple of `(min, max)`.
|
||||
input_channels_per_group_range:
|
||||
Range to generate `input_channels_per_group`.
|
||||
Must be tuple of `(min, max)`.
|
||||
output_channels_per_group_range:
|
||||
Range to generate `output_channels_per_group`.
|
||||
Must be tuple of `(min, max)`.
|
||||
feature_map_range: Range to generate feature map size for each spatial_dim.
|
||||
Must be tuple of `(min, max)`.
|
||||
kernel_range: Range to generate kernel size for each spatial_dim. Must be
|
||||
tuple of `(min, max)`.
|
||||
max_groups: Maximum number of groups to generate.
|
||||
elements: Elements to generate from for the returned data type.
|
||||
If None, the strategy resolves to float within range [-1e6, 1e6].
|
||||
qparams: Strategy for quantization parameters. for X, w, and b.
|
||||
Could be either a single strategy (used for all) or a list of
|
||||
three strategies for X, w, b.
|
||||
Generates:
|
||||
(X, W, b, g): Tensors of type `float32` of the following drawen shapes:
|
||||
X: (`batch_size, input_channels, H, W`)
|
||||
W: (`output_channels, input_channels_per_group) + kernel_shape
|
||||
b: `(output_channels,)`
|
||||
groups: Number of groups the input is divided into
|
||||
Note: X, W, b are tuples of (Tensor, qparams), where qparams could be either
|
||||
None or (scale, zero_point, quantized_type)
|
||||
|
||||
|
||||
Example:
|
||||
@given(tensor_conv(
|
||||
spatial_dim=2,
|
||||
batch_size_range=(1, 3),
|
||||
input_channels_per_group_range=(1, 7),
|
||||
output_channels_per_group_range=(1, 7),
|
||||
feature_map_range=(6, 12),
|
||||
kernel_range=(3, 5),
|
||||
max_groups=4,
|
||||
elements=st.floats(-1.0, 1.0),
|
||||
qparams=qparams()
|
||||
))
|
||||
"""
|
||||
@st.composite
|
||||
def tensor_conv(
|
||||
draw, spatial_dim=2, batch_size_range=(1, 4),
|
||||
input_channels_per_group_range=(3, 7),
|
||||
output_channels_per_group_range=(3, 7), feature_map_range=(6, 12),
|
||||
kernel_range=(3, 7), max_groups=1, can_be_transposed=False,
|
||||
elements=None, qparams=None
|
||||
):
|
||||
|
||||
# Resolve the minibatch, in_channels, out_channels, iH/iW, iK/iW
|
||||
batch_size = draw(st.integers(*batch_size_range))
|
||||
input_channels_per_group = draw(
|
||||
st.integers(*input_channels_per_group_range))
|
||||
output_channels_per_group = draw(
|
||||
st.integers(*output_channels_per_group_range))
|
||||
groups = draw(st.integers(1, max_groups))
|
||||
input_channels = input_channels_per_group * groups
|
||||
output_channels = output_channels_per_group * groups
|
||||
|
||||
if isinstance(spatial_dim, Iterable):
|
||||
spatial_dim = draw(st.sampled_from(spatial_dim))
|
||||
|
||||
feature_map_shape = [draw(st.integers(*feature_map_range)) for _ in range(spatial_dim)]
|
||||
|
||||
kernels = [draw(st.integers(*kernel_range)) for _ in range(spatial_dim)]
|
||||
|
||||
tr = False
|
||||
weight_shape = (output_channels, input_channels_per_group) + tuple(kernels)
|
||||
bias_shape = output_channels
|
||||
if can_be_transposed:
|
||||
tr = draw(st.booleans())
|
||||
if tr:
|
||||
weight_shape = (input_channels, output_channels_per_group) + tuple(kernels)
|
||||
bias_shape = output_channels
|
||||
|
||||
# Resolve the tensors
|
||||
if qparams is not None:
|
||||
if isinstance(qparams, (list, tuple)):
|
||||
if len(qparams) != 3:
|
||||
raise AssertionError("Need 3 qparams for X, w, b")
|
||||
else:
|
||||
qparams = [qparams] * 3
|
||||
|
||||
X = draw(tensor(shapes=(
|
||||
(batch_size, input_channels) + tuple(feature_map_shape),),
|
||||
elements=elements, qparams=qparams[0]))
|
||||
W = draw(tensor(shapes=(weight_shape,), elements=elements,
|
||||
qparams=qparams[1]))
|
||||
b = draw(tensor(shapes=(bias_shape,), elements=elements,
|
||||
qparams=qparams[2]))
|
||||
|
||||
return X, W, b, groups, tr
|
||||
|
||||
|
||||
# We set the deadline in the currently loaded profile.
|
||||
# Creating (and loading) a separate profile overrides any settings the user
|
||||
# already specified.
|
||||
hypothesis_version = tuple(map(int, version("hypothesis").split(".")[:3]))
|
||||
|
||||
if (3, 16, 0) <= hypothesis_version < (3, 27, 0):
|
||||
# Hypothesis 3.16 → 3.26: use `timeout` instead of `deadline`
|
||||
settings.register_profile("no_deadline", timeout=hypothesis.unlimited)
|
||||
else:
|
||||
# Hypothesis >=3.27: use `deadline=None`
|
||||
settings.register_profile("no_deadline", deadline=None)
|
||||
|
||||
# Activate the profile
|
||||
settings.load_profile("no_deadline")
|
||||
|
||||
|
||||
def assert_deadline_disabled():
|
||||
"""Check that deadlines are effectively disabled across Hypothesis versions."""
|
||||
if hypothesis_version < (3, 27, 0):
|
||||
import warnings
|
||||
|
||||
warning_message = (
|
||||
"Your version of hypothesis is outdated. "
|
||||
"To avoid `DeadlineExceeded` errors, please update. "
|
||||
f"Current hypothesis version: {hypothesis.__version__}"
|
||||
)
|
||||
warnings.warn(warning_message, stacklevel=2)
|
||||
else:
|
||||
if settings().deadline is not None:
|
||||
raise AssertionError("Expected settings().deadline to be None")
|
||||
@@ -0,0 +1,470 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import unittest
|
||||
from subprocess import CalledProcessError
|
||||
|
||||
import torch
|
||||
import torch._inductor.async_compile # noqa: F401 required to warm up AsyncCompile pools
|
||||
import torch._inductor.config as config
|
||||
from torch._inductor.codecache import CppCodeCache
|
||||
from torch._inductor.codegen.common import (
|
||||
get_custom_backend_config_for_device,
|
||||
get_custom_backend_pass_for_device,
|
||||
get_scheduling_for_device,
|
||||
get_wrapper_codegen_for_device,
|
||||
init_backend_registration,
|
||||
register_backend_for_device,
|
||||
)
|
||||
from torch._inductor.codegen.wrapper import PythonWrapperCodegen
|
||||
from torch._inductor.compile_fx import shape_env_from_inputs
|
||||
from torch._inductor.custom_graph_pass import CustomGraphModulePass, CustomGraphPass
|
||||
from torch._inductor.graph import GraphLowering
|
||||
from torch._inductor.utils import (
|
||||
get_gpu_shared_memory,
|
||||
get_gpu_type,
|
||||
GPU_TYPES,
|
||||
is_big_gpu,
|
||||
is_gpu,
|
||||
OrderedSet,
|
||||
)
|
||||
from torch.fx.experimental.proxy_tensor import make_fx
|
||||
from torch.utils._helion import has_helion
|
||||
from torch.utils._pallas import has_pallas_package, has_tpu_pallas
|
||||
from torch.utils._triton import has_triton
|
||||
from torch.utils._config_module import ConfigModule
|
||||
from torch.testing._internal.common_device_type import (
|
||||
get_desired_device_type_test_bases,
|
||||
)
|
||||
from torch.testing._internal.common_utils import (
|
||||
IS_CI,
|
||||
IS_WINDOWS,
|
||||
LazyVal,
|
||||
TestCase,
|
||||
)
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
log: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def test_cpu():
|
||||
try:
|
||||
CppCodeCache.load("")
|
||||
return True
|
||||
except (
|
||||
CalledProcessError,
|
||||
OSError,
|
||||
torch._inductor.exc.InvalidCxxCompiler,
|
||||
torch._inductor.exc.CppCompileError,
|
||||
):
|
||||
return False
|
||||
|
||||
|
||||
HAS_CPU = LazyVal(test_cpu)
|
||||
|
||||
HAS_TRITON = has_triton()
|
||||
|
||||
HAS_PALLAS = has_pallas_package()
|
||||
|
||||
HAS_HELION = has_helion()
|
||||
|
||||
if HAS_TRITON:
|
||||
import triton
|
||||
|
||||
TRITON_HAS_CPU = "cpu" in triton.backends.backends
|
||||
else:
|
||||
TRITON_HAS_CPU = False
|
||||
|
||||
|
||||
HAS_CUDA_AND_TRITON = torch.cuda.is_available() and HAS_TRITON
|
||||
|
||||
HAS_XPU_AND_TRITON = torch.xpu.is_available() and HAS_TRITON
|
||||
|
||||
HAS_MPS = torch.mps.is_available()
|
||||
|
||||
HAS_GPU = HAS_CUDA_AND_TRITON or HAS_XPU_AND_TRITON
|
||||
HAS_GPU_AND_TRITON = HAS_GPU
|
||||
|
||||
GPU_TYPE = get_gpu_type()
|
||||
|
||||
HAS_MULTIGPU = any(
|
||||
getattr(torch, gpu).is_available() and getattr(torch, gpu).device_count() >= 2
|
||||
for gpu in GPU_TYPES
|
||||
)
|
||||
|
||||
_desired_test_bases = get_desired_device_type_test_bases(allow_xpu=True)
|
||||
RUN_GPU = HAS_GPU and any(
|
||||
is_gpu(getattr(x, "device_type", "")) for x in _desired_test_bases
|
||||
)
|
||||
|
||||
RUN_CPU = HAS_CPU and any(
|
||||
getattr(x, "device_type", "") == "cpu" for x in _desired_test_bases
|
||||
)
|
||||
|
||||
HAS_TPU = has_tpu_pallas()
|
||||
# TPU is a privateuse1 backend that isn't in _desired_test_bases (it requires
|
||||
# runtime initialization before the test base is registered). Check the env var
|
||||
# directly, matching the same semantics as RUN_CPU/RUN_GPU: when the env var is
|
||||
# unset, run if the hardware is available; when set, only run if "tpu" is listed.
|
||||
_only_for = os.environ.get("PYTORCH_TESTING_DEVICE_ONLY_FOR", "")
|
||||
RUN_TPU = HAS_TPU and ("tpu" in _only_for.split(",") if _only_for else True)
|
||||
|
||||
|
||||
def _check_has_dynamic_shape(
|
||||
self: TestCase,
|
||||
code,
|
||||
):
|
||||
for_loop_found = False
|
||||
has_dynamic = False
|
||||
lines = code.split("\n")
|
||||
for line in lines:
|
||||
if "for(" in line:
|
||||
for_loop_found = True
|
||||
if re.search(r";.*ks.*;", line) is not None:
|
||||
has_dynamic = True
|
||||
break
|
||||
self.assertTrue(
|
||||
has_dynamic, msg=f"Failed to find dynamic for loop variable\n{code}"
|
||||
)
|
||||
self.assertTrue(for_loop_found, f"Failed to find for loop\n{code}")
|
||||
|
||||
|
||||
def skipDeviceIf(cond, msg, *, device):
|
||||
if cond:
|
||||
|
||||
def decorate_fn(fn):
|
||||
@functools.wraps(fn)
|
||||
def inner(self, *args, **kwargs):
|
||||
if not hasattr(self, "device"):
|
||||
warn_msg = (
|
||||
"Expect the test class to have attribute device but not found. "
|
||||
)
|
||||
if hasattr(self, "device_type"):
|
||||
warn_msg += "Consider using the skip device decorators in common_device_type.py"
|
||||
log.warning(warn_msg)
|
||||
if self.device == device:
|
||||
raise unittest.SkipTest(msg)
|
||||
return fn(self, *args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
else:
|
||||
|
||||
def decorate_fn(fn):
|
||||
return fn
|
||||
|
||||
return decorate_fn
|
||||
|
||||
|
||||
def skip_windows_ci(name: str, file: str) -> None:
|
||||
if IS_WINDOWS and IS_CI:
|
||||
module = os.path.basename(file).strip(".py")
|
||||
sys.stderr.write(
|
||||
f"Windows CI does not have necessary dependencies for {module} tests yet\n"
|
||||
)
|
||||
if name == "__main__":
|
||||
sys.exit(0)
|
||||
raise unittest.SkipTest("requires sympy/functorch/filelock")
|
||||
|
||||
|
||||
# TODO: Remove HAS_MPS condition when `HAS_GPU` includes HAS_MPS
|
||||
requires_gpu = functools.partial(
|
||||
unittest.skipIf, not (HAS_GPU or HAS_MPS), "requires gpu"
|
||||
)
|
||||
requires_triton = functools.partial(unittest.skipIf, not HAS_TRITON, "requires triton")
|
||||
requires_helion = functools.partial(unittest.skipIf, not HAS_HELION, "requires helion")
|
||||
|
||||
|
||||
def requires_gpu_with_enough_memory(min_mem_required):
|
||||
def inner(fn):
|
||||
total_memory = sys.maxsize
|
||||
if torch.xpu.is_available():
|
||||
total_memory = torch.xpu.get_device_properties().total_memory
|
||||
elif torch.cuda.is_available():
|
||||
total_memory = torch.cuda.get_device_properties().total_memory
|
||||
if (
|
||||
not (torch.cuda.is_available() or torch.xpu.is_available())
|
||||
or total_memory < min_mem_required
|
||||
):
|
||||
return unittest.skip(
|
||||
f"Only if the GPU device has at least {min_mem_required / 1e9:.3f}GB memory to be safe"
|
||||
)(fn)
|
||||
else:
|
||||
return fn
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
skipCUDAIf = functools.partial(skipDeviceIf, device="cuda")
|
||||
skipXPUIf = functools.partial(skipDeviceIf, device="xpu")
|
||||
skipCPUIf = functools.partial(skipDeviceIf, device="cpu")
|
||||
|
||||
IS_A100 = LazyVal(lambda: HAS_CUDA_AND_TRITON and get_gpu_shared_memory() == 166912)
|
||||
|
||||
IS_H100 = LazyVal(lambda: HAS_CUDA_AND_TRITON and get_gpu_shared_memory() == 232448)
|
||||
|
||||
IS_BIG_GPU = LazyVal(lambda: HAS_GPU_AND_TRITON and is_big_gpu())
|
||||
|
||||
|
||||
def dummy_graph() -> GraphLowering:
|
||||
"""
|
||||
Create a graph. This is useful for unit testing code which accesses
|
||||
V.graph.sizevars.
|
||||
"""
|
||||
example_inputs = [torch.randn(10) for _ in range(2)]
|
||||
gm = make_fx(torch.add, tracing_mode="fake")(*example_inputs)
|
||||
shape_env = shape_env_from_inputs(example_inputs)
|
||||
graph = GraphLowering(
|
||||
gm,
|
||||
shape_env=shape_env,
|
||||
)
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
def maybe_skip_size_asserts(op):
|
||||
"""
|
||||
For certain ops, there meta and eager implementation returns different
|
||||
strides. This cause size/strides assert fail. Skip adding those
|
||||
asserts for now.
|
||||
"""
|
||||
if (
|
||||
op.aten_name
|
||||
in (
|
||||
"fft_hfftn",
|
||||
"fft_hfft",
|
||||
"fft_hfft2",
|
||||
"fft_ihfftn",
|
||||
"fft_fft",
|
||||
"fft_fft2",
|
||||
"fft_fftn",
|
||||
"fft_ifft",
|
||||
"fft_ifft2",
|
||||
"fft_ifftn",
|
||||
"fft_irfft",
|
||||
"fft_irfft2",
|
||||
"fft_irfftn",
|
||||
"fft_ihfft",
|
||||
"fft_ihfft2",
|
||||
"fft_rfft",
|
||||
"fft_rfft2",
|
||||
"fft_rfftn",
|
||||
"linalg_eig",
|
||||
"linalg_eigvals",
|
||||
)
|
||||
and "TORCHINDUCTOR_SIZE_ASSERTS" not in os.environ
|
||||
):
|
||||
return torch._inductor.config.patch(size_asserts=False)
|
||||
else:
|
||||
return contextlib.nullcontext()
|
||||
|
||||
|
||||
def get_func_call() -> str:
|
||||
return (
|
||||
"void inductor_entry_impl("
|
||||
if torch._inductor.config.cpp_wrapper
|
||||
else "def call("
|
||||
)
|
||||
|
||||
|
||||
def get_kernel_launch() -> str:
|
||||
return "call_triton_" if torch._inductor.config.cpp_wrapper else ".run("
|
||||
|
||||
|
||||
def clone_preserve_strides_offset(x, device=None):
|
||||
if not isinstance(x, torch.Tensor):
|
||||
return x
|
||||
buffer = torch.as_strided(
|
||||
x, (x.untyped_storage().size() // x.element_size(),), (1,), 0
|
||||
)
|
||||
if not device:
|
||||
buffer = buffer.clone()
|
||||
else:
|
||||
buffer = buffer.to(device, copy=True)
|
||||
out = torch.as_strided(buffer, x.size(), x.stride(), x.storage_offset())
|
||||
return out
|
||||
|
||||
|
||||
# define the e4m3/e5m2 constants
|
||||
E4M3_MAX_POS = torch.finfo(torch.float8_e4m3fn).max
|
||||
E5M2_MAX_POS = torch.finfo(torch.float8_e5m2).max
|
||||
E4M3FNUZ_MAX_POS = torch.finfo(torch.float8_e4m3fnuz).max
|
||||
E5M2FNUZ_MAX_POS = torch.finfo(torch.float8_e5m2fnuz).max
|
||||
|
||||
FP16_MAX_POS: float = torch.finfo(torch.float16).max
|
||||
EPS: float = 1e-12
|
||||
|
||||
Tensor = torch.Tensor
|
||||
|
||||
|
||||
def _to_fp8_saturated(x: Tensor, float8_dtype: torch.dtype) -> Tensor:
|
||||
# The default behavior in PyTorch for casting to `float8_e4m3fn`
|
||||
# and `e5m2` is to not saturate. In this context, we should saturate.
|
||||
# A common case where we want to saturate is when the history of a
|
||||
# tensor has a maximum value of `amax1`, and the current amax value
|
||||
# is `amax2`, where `amax1 < amax2`. This is common when using delayed
|
||||
# scaling.
|
||||
if float8_dtype == torch.float8_e4m3fn:
|
||||
x = x.clamp(min=-1 * E4M3_MAX_POS, max=E4M3_MAX_POS)
|
||||
elif float8_dtype == torch.float8_e5m2:
|
||||
x = x.clamp(min=-1 * E5M2_MAX_POS, max=E5M2_MAX_POS)
|
||||
elif float8_dtype == torch.float8_e4m3fnuz:
|
||||
x = x.clamp(min=-1 * E4M3FNUZ_MAX_POS, max=E4M3FNUZ_MAX_POS)
|
||||
elif float8_dtype == torch.float8_e5m2fnuz:
|
||||
x = x.clamp(min=-1 * E5M2FNUZ_MAX_POS, max=E5M2FNUZ_MAX_POS)
|
||||
else:
|
||||
raise TypeError(f"Unsupported float8_dtype: {float8_dtype}")
|
||||
return x.to(float8_dtype)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _amax_to_scale(
|
||||
amax: torch.Tensor, float8_dtype: torch.dtype, orig_dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
# To make scale dtype to be fp32 for accuracy
|
||||
amax = amax.float()
|
||||
if float8_dtype == torch.float8_e4m3fn:
|
||||
res = E4M3_MAX_POS / torch.clamp(amax, min=EPS)
|
||||
else: # e5m2
|
||||
res = E5M2_MAX_POS / torch.clamp(amax, min=EPS)
|
||||
|
||||
# Ensure that the scale is representable in float16,
|
||||
# this helps when amax is small. We are assuming that we don't need
|
||||
# to care about this for float32/bfloat16.
|
||||
if orig_dtype is torch.float16:
|
||||
res = torch.clamp(res, max=FP16_MAX_POS)
|
||||
return res
|
||||
|
||||
|
||||
def _quantize_tensorwise(x: Tensor, float8_dtype: torch.dtype):
|
||||
amax = torch.max(torch.abs(x))
|
||||
scale = _amax_to_scale(amax, float8_dtype, x.dtype)
|
||||
x_fp8 = _to_fp8_saturated(x * scale, float8_dtype)
|
||||
inverse_scale = scale.reciprocal()
|
||||
return x_fp8, inverse_scale
|
||||
|
||||
|
||||
def _quantize_rowwise(x: Tensor, float8_dtype: torch.dtype):
|
||||
amax = torch.max(torch.abs(x), dim=1, keepdim=True).values
|
||||
scale = _amax_to_scale(amax, float8_dtype, x.dtype)
|
||||
x_fp8 = _to_fp8_saturated(x * scale, float8_dtype)
|
||||
inverse_scale = scale.reciprocal()
|
||||
return x_fp8, inverse_scale
|
||||
|
||||
|
||||
def _quantize_blockwise(
|
||||
x: Tensor, float8_dtype: torch.dtype, block_outer: int, block_inner: int
|
||||
):
|
||||
min_outer = min(block_outer, x.shape[0])
|
||||
min_inner = min(block_inner, x.shape[1])
|
||||
x = x.unflatten(1, (-1, min_inner)).unflatten(0, (-1, min_outer))
|
||||
amax = x.abs().amax(dim=[1, 3], keepdim=True).float()
|
||||
scale = _amax_to_scale(amax, float8_dtype, x.dtype)
|
||||
x = x.flatten(2, 3).flatten(0, 1)
|
||||
scale = scale.flatten(2, 3).flatten(0, 1)
|
||||
scale_expanded = scale.repeat_interleave(min_outer, dim=0).repeat_interleave(
|
||||
min_inner, dim=1
|
||||
)
|
||||
x_fp8 = _to_fp8_saturated(
|
||||
x / scale_expanded, # Ensures that scaling doesn't cause inf/nan values
|
||||
float8_dtype,
|
||||
)
|
||||
inverse_scale = scale.reciprocal()
|
||||
return x_fp8, inverse_scale
|
||||
|
||||
|
||||
class MockGraphHandler(GraphLowering):
|
||||
"""Minimal mock graph handler for testing virtualized context."""
|
||||
|
||||
def __init__(self, name_to_buffer=None):
|
||||
import torch._inductor.sizevars
|
||||
|
||||
self.sizevars = torch._inductor.sizevars.SizeVarAllocator()
|
||||
self.name_to_buffer = name_to_buffer or {}
|
||||
self.graph_inputs = {}
|
||||
self.mutated_buffers = OrderedSet()
|
||||
self.removed_buffers = OrderedSet()
|
||||
self.constants = {}
|
||||
self.scheduler = None
|
||||
|
||||
def get_dtype(self, buffer_name: str) -> torch.dtype: # noqa: ARG002
|
||||
"""Return default dtype for any buffer (for testing)."""
|
||||
return torch.float32
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def patch_inductor_backend(
|
||||
device: str,
|
||||
python_wrapper_codegen: PythonWrapperCodegen = None,
|
||||
custom_pass: CustomGraphModulePass = None,
|
||||
custom_backend_config: ConfigModule = None,
|
||||
):
|
||||
"""
|
||||
Patch the inductor backend for a specific device.
|
||||
"""
|
||||
# Make sure the backend is already registered
|
||||
init_backend_registration()
|
||||
|
||||
# Get the original registration parameters
|
||||
original_scheduling = get_scheduling_for_device(device)
|
||||
original_python_wrapper = get_wrapper_codegen_for_device(device, False)
|
||||
original_cpp_wrapper = get_wrapper_codegen_for_device(device, True)
|
||||
original_fx_wrapper = get_wrapper_codegen_for_device(device, fx_wrapper=True)
|
||||
original_custom_pass = get_custom_backend_pass_for_device(device)
|
||||
original_custom_backend_config = get_custom_backend_config_for_device(device)
|
||||
|
||||
try:
|
||||
# Register modified backend for the device
|
||||
register_backend_for_device(
|
||||
device,
|
||||
original_scheduling,
|
||||
(
|
||||
python_wrapper_codegen
|
||||
if python_wrapper_codegen is not None
|
||||
else original_python_wrapper
|
||||
),
|
||||
original_cpp_wrapper,
|
||||
original_fx_wrapper,
|
||||
custom_pass if custom_pass is not None else original_custom_pass,
|
||||
(
|
||||
custom_backend_config
|
||||
if custom_backend_config is not None
|
||||
else original_custom_backend_config
|
||||
),
|
||||
)
|
||||
yield
|
||||
finally:
|
||||
# Restore the original backend
|
||||
register_backend_for_device(
|
||||
device,
|
||||
original_scheduling,
|
||||
original_python_wrapper,
|
||||
original_cpp_wrapper,
|
||||
original_fx_wrapper,
|
||||
original_custom_pass,
|
||||
original_custom_backend_config,
|
||||
)
|
||||
|
||||
def patch_custom_fallback_pass(predicate: Callable[[torch.fx.Node], bool]) -> contextlib.ContextDecorator:
|
||||
"""
|
||||
Create a custom pass which falls back based on the provided predicate. For example,
|
||||
we could provide a predicate which returns True for all aten.add.default nodes.
|
||||
Returns a context activating the pass.
|
||||
"""
|
||||
class Pass(CustomGraphPass):
|
||||
def __call__(self, graph: torch.fx.Graph):
|
||||
for node in graph.nodes:
|
||||
if predicate(node):
|
||||
node.meta["should_fallback"] = True
|
||||
|
||||
def uuid(self):
|
||||
return None
|
||||
|
||||
|
||||
return config.patch(post_grad_custom_pre_pass=Pass())
|
||||
+726
@@ -0,0 +1,726 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
# Torch
|
||||
from torch.jit.annotations import BroadcastingList2, BroadcastingList3 # noqa: F401
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
import torch.cuda
|
||||
import torch.jit
|
||||
import torch.jit._logging
|
||||
import torch.jit.frontend
|
||||
from torch.testing._internal.common_nn import module_tests, get_new_module_tests
|
||||
from torch.testing._internal.common_utils import is_iterable_of_tensors, noncontiguous_like
|
||||
|
||||
import collections
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
import math # noqa: F401
|
||||
|
||||
# Testing utils
|
||||
from torch import inf
|
||||
|
||||
if torch.get_default_dtype() != torch.float32:
|
||||
raise AssertionError(f"Expected torch.get_default_dtype() == torch.float32, got {torch.get_default_dtype()}")
|
||||
|
||||
L = 20
|
||||
M = 10
|
||||
S = 5
|
||||
|
||||
|
||||
def unpack_variables(args):
|
||||
if isinstance(args, tuple):
|
||||
return tuple(unpack_variables(elem) for elem in args)
|
||||
else:
|
||||
return args
|
||||
|
||||
class dont_convert(tuple):
|
||||
__slots__ = ()
|
||||
|
||||
non_differentiable = collections.namedtuple('non_differentiable', ['tensor'])
|
||||
|
||||
def create_input(call_args, requires_grad=True, non_contiguous=False, call_kwargs=None, dtype=torch.float, device=None):
|
||||
if not isinstance(call_args, tuple):
|
||||
call_args = (call_args,)
|
||||
|
||||
def map_arg(arg):
|
||||
def maybe_non_contig(tensor):
|
||||
if not non_contiguous or tensor.numel() < 2:
|
||||
return tensor.clone()
|
||||
|
||||
return noncontiguous_like(tensor)
|
||||
|
||||
def conjugate(tensor):
|
||||
return tensor.conj()
|
||||
|
||||
if isinstance(arg, (torch.Size, dont_convert)):
|
||||
return arg
|
||||
elif isinstance(arg, tuple) and len(arg) == 0:
|
||||
var = conjugate(torch.randn((), dtype=dtype, device=device))
|
||||
var.requires_grad = requires_grad
|
||||
return var
|
||||
elif isinstance(arg, tuple) and not isinstance(arg[0], torch.Tensor):
|
||||
return conjugate(maybe_non_contig(torch.randn(*arg, dtype=dtype, device=device))).requires_grad_(requires_grad)
|
||||
# double check casting
|
||||
elif isinstance(arg, non_differentiable):
|
||||
if isinstance(arg.tensor, torch.Tensor):
|
||||
return conjugate(maybe_non_contig(arg.tensor.to(device=device)))
|
||||
return conjugate(maybe_non_contig(arg.tensor.to(device=device)))
|
||||
elif isinstance(arg, torch.Tensor):
|
||||
if arg.is_complex() != dtype.is_complex:
|
||||
raise RuntimeError("User provided tensor is real for a test that runs with complex dtype, ",
|
||||
"which is not supported for now")
|
||||
# NOTE: We do clone() after detach() here because we need to be able to change size/storage of v afterwards
|
||||
v = conjugate(maybe_non_contig(arg)).detach().to(device=device).clone()
|
||||
v.requires_grad = requires_grad and (v.is_floating_point() or v.is_complex())
|
||||
return v
|
||||
elif callable(arg):
|
||||
return map_arg(arg(dtype=dtype, device=device))
|
||||
else:
|
||||
return arg
|
||||
args_out = tuple(map_arg(arg) for arg in call_args)
|
||||
kwargs_out = {k: map_arg(v) for k, v in call_kwargs.items()} if call_kwargs else {}
|
||||
return args_out, kwargs_out
|
||||
|
||||
# NB: JIT script tests for all nn functional interfaces, script mode does
|
||||
# not support in_place operations yet, so no inplace operation tests added.
|
||||
# removed all the deprecated functions
|
||||
#
|
||||
# (
|
||||
# method name,
|
||||
# input size/constructing fn,
|
||||
# args (tuple represents shape of a tensor arg),
|
||||
# test variant name(will be used at test name suffix,
|
||||
# 'inplace' skips grad tests), // optional
|
||||
# (True, nonfusible_nodes, fusible_nodes) for autodiff // optional
|
||||
# fn to determine if test should be skipped, // optional
|
||||
# fn mapping output to part that should be gradcheck'ed, // optional
|
||||
# kwargs for function, // optional
|
||||
# )
|
||||
def get_nn_functional_tests():
|
||||
nn_functional_tests = [
|
||||
('conv1d', (S, S, S), ((S, S, S),)),
|
||||
('conv2d', (S, S, S, S), ((S, S, S, S),)),
|
||||
('conv3d', (S, S, S, S, S), ((S, S, S, S, S),)),
|
||||
('conv_transpose1d', (S, S, S), ((S, S, S),)),
|
||||
('conv_transpose2d', (S, S, S, S), ((S, S, S, S),)),
|
||||
('conv_transpose3d', (S, S, S, S, S), ((S, S, S, S, S),)),
|
||||
('conv_tbc', (S, S, S), ((S, S, S), (S,), 2)),
|
||||
('avg_pool1d', (S, S, S), (3,)),
|
||||
('avg_pool2d', (S, S, S, S), (3,), '', (True,)),
|
||||
('avg_pool3d', (S, S, S, S, S), (3,)),
|
||||
('fractional_max_pool2d', (S, S, S, S), (3, [2, 3],)),
|
||||
('max_pool1d', (S, S, S), (2, 1)),
|
||||
('max_pool1d', (S, S, S), (2, 1, 1, 1, False, True), 'with_indices'),
|
||||
('max_pool2d', (S, S, S, S), (2, 1), '', (True, 'aten::max_pool2d_with_indices')),
|
||||
('max_pool2d', (S, S, S, S), (2, 1, 1, 1, False, True), 'with_indices', (True, 'aten::max_pool2d_with_indices')),
|
||||
('max_pool3d', (S, S, S, S, S), (2, 1)),
|
||||
('max_unpool1d', torch.tensor([[[2., 4]]]), (torch.tensor([[[1, 3]]]), 2, 2, 0)),
|
||||
('max_unpool2d', torch.tensor([[[[2., 4]]]]), (torch.tensor([[[[1, 3]]]]), 2, 2, 0)),
|
||||
('max_unpool3d', torch.tensor([[[[[2., 4]]]]]), (torch.tensor([[[[[1, 3]]]]]), 2, 2, 0)),
|
||||
('lp_pool1d', (S, S, S), (2., 3, 2,)),
|
||||
('lp_pool2d', (S, S, S, S), (2., 3, 2,)),
|
||||
('lp_pool3d', (S, S, S, S, S), (2., 3, 2,)),
|
||||
('adaptive_max_pool1d', (S, S, S), (5,)),
|
||||
('adaptive_max_pool2d', (S, S, S, S), ([5, 7],)),
|
||||
('adaptive_max_pool3d', (S, S, S, S, S), ([3, 2, 2],)),
|
||||
('adaptive_avg_pool1d', (S, S, S), (5,), '', (True,)),
|
||||
('adaptive_avg_pool2d', (S, S, S, S), ([5, 7],), '', (True,)),
|
||||
('adaptive_avg_pool3d', (S, S, S, S, S), ([3, 2, 2],), '', (True,)),
|
||||
('dropout', (S, S, S), (0.5,), '', (True, 'aten::native_dropout')),
|
||||
('alpha_dropout', (S, S, S), (0.5,)),
|
||||
('dropout2d', (S, S, S), (0.5,)),
|
||||
('dropout2d', (S, S, S, S), (0.5,), 'batched'),
|
||||
('dropout3d', (S, S, S, S), (0.5,)),
|
||||
('dropout3d', (S, S, S, S, S), (0.5,), 'batched'),
|
||||
('feature_alpha_dropout', (S, S, S), (0.5,)),
|
||||
('threshold', (S, S, S), (0.1, 2.), '', (True,)),
|
||||
('threshold', (S, S, S), (0.1, 2., True), 'inplace'),
|
||||
('relu', (S, S, S), (), '', (True,)),
|
||||
('relu', (S, S, S), (), 'inplace'),
|
||||
('glu', (S - 1, S - 1, S - 1), (),),
|
||||
('hardtanh', (S, S, S), (-0.5, 0.5), '', (True,)),
|
||||
('hardtanh', (S, S, S), (-0.5, 0.5, True), 'inplace'),
|
||||
('relu6', (S, S, S), (), '', (True,)),
|
||||
('relu6', (S, S, S), (True), 'inplace'),
|
||||
('elu', (S, S, S), (0.9,),),
|
||||
('elu', (S, S, S), (0.9, True), 'inplace'),
|
||||
('selu', (S, S, S), (),),
|
||||
('selu', (S, S, S), (True), 'inplace'),
|
||||
('celu', (S, S, S), (0.9,),),
|
||||
('celu', (S, S, S), (0.9, True), 'inplace'),
|
||||
('leaky_relu', (S, S, S), (0.02,), '', (True,)),
|
||||
('leaky_relu', (S, S, S), (0.02,), 'inplace'),
|
||||
('rrelu', (S, S), (0.1, 0.3, False),),
|
||||
('rrelu', (S, S), (0.1, 0.3, False, True), 'inplace'),
|
||||
('hardshrink', (S, S, S), (0.4,), '', (True,)),
|
||||
('tanhshrink', (S, S, S), (),),
|
||||
('softsign', (S, S, S), (),),
|
||||
('softplus', (S, S, S), (), '', (True,)),
|
||||
('softmin', (S, S, S), (0,),),
|
||||
('softmax', (S, S, S), (0,), '', (True,)),
|
||||
('softmax', (S, S, S), (0, 3, torch.double), 'with_all_args', (True,)),
|
||||
('tanh', (S, S, S), (), '', (True,)),
|
||||
('sigmoid', (S, S, S), (), '', (True,)),
|
||||
('silu', (S, S, S), (), '', (True,)),
|
||||
('log_softmax', (S, S, S), (0,), '', (True,)),
|
||||
('linear', (S, S), ((M, S),), '', (True, ['aten::linear'])),
|
||||
('linear', (S, S), ((M, S), (M,)), 'addmm', (True, ['aten::linear'])),
|
||||
('bilinear', (S, S, S), ((S, S, M), torch.zeros(M, S, M),),),
|
||||
('embedding', torch.tensor([[1, 2, 4, 5], [4, 3, 2, 5]]), (torch.rand(6, 3), ), '', (True,)),
|
||||
('embedding_bag', torch.tensor([1, 2, 4, 2]), (torch.rand(5, 3), torch.tensor([0, 4]),),),
|
||||
('batch_norm', (S, S),
|
||||
(non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), None, None, True, ),
|
||||
'training', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (0, S, S, S),
|
||||
(non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), True, ),
|
||||
'size_zero', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (0, S, S, S),
|
||||
(non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), True, ),
|
||||
'size_zero_inference', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (S, S),
|
||||
(non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), True, ),
|
||||
'with_weight_and_bias_training', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
None, non_differentiable(torch.ones(S)), True, ),
|
||||
'with_only_bias_training', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
non_differentiable(torch.randn(S)), None, True, ),
|
||||
'with_only_weight_training', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
None, None, False, ),
|
||||
'inference', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), False, ),
|
||||
'with_weight_and_bias_inference', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
None, non_differentiable(torch.ones(S)), False, ),
|
||||
'with_only_bias_inference', (True, 'aten::_batch_norm_impl_index')),
|
||||
('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)),
|
||||
non_differentiable(torch.randn(S)), None, False, ),
|
||||
'with_only_weight_inference', (True, 'aten::_batch_norm_impl_index')),
|
||||
('instance_norm', (S, S, S), (non_differentiable(torch.zeros(S)), non_differentiable(torch.ones(S))),),
|
||||
('layer_norm', (S, S, S, S), ([5],), '',
|
||||
(False, ['aten::contiguous', 'aten::_batch_norm_impl_index'])),
|
||||
('layer_norm', (S, S, S, S), ([5], non_differentiable(torch.rand(S)),), 'with_only_weight',
|
||||
(False, ['aten::contiguous', 'aten::_batch_norm_impl_index'])),
|
||||
('layer_norm', (S, S, S, S), ([5], None, non_differentiable(torch.rand(S)),), 'with_only_bias',
|
||||
(False, ['aten::contiguous', 'aten::_batch_norm_impl_index'])),
|
||||
('layer_norm', (S, S, S, S), ([5], non_differentiable(torch.rand(S)),
|
||||
non_differentiable(torch.rand(S))), 'with_weight_and_bias',
|
||||
(False, ['aten::contiguous', 'aten::_batch_norm_impl_index', 'aten::addcmul'])),
|
||||
('group_norm', (S, S, S), (1, torch.rand(5),),),
|
||||
('local_response_norm', (S, S, S), (2, ),),
|
||||
('nll_loss', F.log_softmax(torch.randn(3, 5), dim=0), (torch.tensor([1, 0, 4]),), '',),
|
||||
('poisson_nll_loss', torch.rand(S, 2), (torch.rand(S, 2),),),
|
||||
('poisson_nll_loss', torch.rand(S, 2), (torch.rand(S, 2), True, True), 'full'),
|
||||
('kl_div', F.log_softmax(torch.randn(S, 10), 1), (F.softmax(torch.randn(S, 10), 1),),),
|
||||
('cross_entropy', (3, S), (torch.randint(S, (3,), dtype=torch.int64),),),
|
||||
('binary_cross_entropy_with_logits', (3,), (torch.empty(3).random_(2), ),),
|
||||
('smooth_l1_loss', (3, S), (non_differentiable(torch.rand(3, S)),),),
|
||||
('huber_loss', (3, S), (non_differentiable(torch.rand(3, S)),),),
|
||||
('l1_loss', (3, S), (non_differentiable(torch.rand(3, S)),),),
|
||||
('mse_loss', (3, S), (non_differentiable(torch.rand(3, S)),),),
|
||||
('smooth_l1_loss', (3, S), ((torch.rand(3, S)),), 'with_grad'),
|
||||
('huber_loss', (3, S), ((torch.rand(3, S)),), 'with_grad'),
|
||||
('l1_loss', (3, S), ((torch.rand(3, S)),), 'with_grad'),
|
||||
('mse_loss', (3, S), ((torch.rand(3, S)),), 'with_grad'),
|
||||
('margin_ranking_loss', (S,), ((S,), (S,)),),
|
||||
('hinge_embedding_loss', (3, S), (non_differentiable(torch.rand(3, S)),),),
|
||||
('soft_margin_loss', (3, S), (non_differentiable(torch.rand(3, S)),),),
|
||||
('multilabel_soft_margin_loss', (3, S), (non_differentiable(torch.rand(3, S)),),),
|
||||
('cosine_embedding_loss', (S, S), ((S, S), non_differentiable(torch.rand(S,))),),
|
||||
('pixel_shuffle', (1, 9, 4, 4), (3,),),
|
||||
('pixel_unshuffle', (1, 1, 12, 12), (3,),),
|
||||
('affine_grid', (S, 2, 3), (torch.Size([S, 1, 7, 7]),),),
|
||||
('pad', (3, 3, 4, 2), ([1, 1],),),
|
||||
('pairwise_distance', (S, S), ((S, S),),),
|
||||
('pdist', (S, S), (),),
|
||||
('cosine_similarity', (S, S), ((S, S),),),
|
||||
('triplet_margin_loss', (S, S), ((S, S), (S, S)),),
|
||||
('normalize', (S, S, S), (),),
|
||||
('unfold', (S, S, S, S), ([2, 3]),),
|
||||
('fold', (1, 3 * 2 * 2, 12), ([4, 5], [2, 2]),),
|
||||
('grid_sample', (S, S, S, S), (non_differentiable(torch.rand(S, S, S, 2)),),),
|
||||
('gumbel_softmax', (S, S), (2.,), '', (True, ['aten::softmax', 'aten::add', 'aten::div'], ['aten::neg'])),
|
||||
('gumbel_softmax', (S, S), (2., True,), 'hard', (True, ['aten::softmax', 'aten::add', 'aten::div'], ['aten::neg'])),
|
||||
('multilabel_margin_loss', torch.tensor([[0.2, -0.2, 0.07]]), (torch.tensor([[0, 0, 1]]),),),
|
||||
('multi_margin_loss', (S, S), (non_differentiable(torch.randint(S, (S, ), dtype=torch.int64)),
|
||||
1, 1., non_differentiable(torch.randn(S))),),
|
||||
('binary_cross_entropy', torch.randn(3, 2).sigmoid(), (non_differentiable(torch.rand(3, 2)),
|
||||
non_differentiable(torch.randn(3, 2))),),
|
||||
('binary_cross_entropy', torch.randn(3, 2).sigmoid(),
|
||||
(non_differentiable(torch.rand(3, 2)),
|
||||
non_differentiable(torch.randn(3, 2)), None, None, 'mean'), 'size_average'),
|
||||
('ctc_loss', torch.rand(S, S, S).log_softmax(2).detach().requires_grad_(),
|
||||
(torch.randint(1, S, (S, S), dtype=torch.long), torch.full((S,), S, dtype=torch.long),
|
||||
torch.randint(1, S, (S,), dtype=torch.long))),
|
||||
('upsample', torch.randn(S, S, M, M), (None, 2.), 'with_scale'),
|
||||
('upsample', torch.randn(S, S, M, M), (4,), 'with_size'),
|
||||
('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2,), 'nearest_4d'),
|
||||
('interpolate', torch.randn(S, S, M, M), (None, 2.), 'nearest_4d_with_scale'),
|
||||
('interpolate', torch.randn(S, S, M, M), (4,), 'nearest_4d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2,), 'area_4d'),
|
||||
('interpolate', torch.randn(S, S, M, M), (None, 2.), 'area_4d_with_scale'),
|
||||
('interpolate', torch.randn(S, S, M, M), (4,), 'area_4d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2,), 'bilinear_4d'),
|
||||
('interpolate', torch.randn(S, S, M, M), (None, 2.), 'bilinear_4d_with_scale'),
|
||||
('interpolate', torch.randn(S, S, M, M), (4,), 'bilinear_4d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2,), 'bicubic_4d'),
|
||||
('interpolate', torch.randn(S, S, M, M), (None, 2.), 'bicubic_4d_with_scale'),
|
||||
('interpolate', torch.randn(S, S, M, M), (4,), 'bicubic_4d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3).view(1, 3, 3), (2,), 'nearest_3d'),
|
||||
('interpolate', torch.randn(S, M, M), (None, 2.), 'nearest_3d_with_scale'),
|
||||
('interpolate', torch.randn(S, M, M), (4,), 'nearest_3d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3).view(1, 3, 3), (2,), 'area_3d'),
|
||||
('interpolate', torch.randn(S, M, M), (None, 2.), 'area_3d_with_scale'),
|
||||
('interpolate', torch.randn(S, M, M), (4,), 'area_3d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3).view(1, 3, 3), (2,), 'linear_3d'),
|
||||
('interpolate', torch.randn(S, M, M), (None, 2.), 'linear_3d_with_scale'),
|
||||
('interpolate', torch.randn(S, M, M), (4,), 'linear_3d_with_size'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (None, 2.), 'nearest_5d_with_scale'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (4,), 'nearest_5d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3, 3).view(1, 1, 3, 3, 3), (2,), 'area_5d'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (None, 2.), 'area_5d_with_scale'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (4,), 'area_5d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3, 3).view(1, 1, 3, 3, 3), (2,), 'trilinear_5d'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (None, 2.), 'trilinear_5d_with_scale'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (4,), 'trilinear_5d_with_size'),
|
||||
('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2, None, 'nearest', None, False),
|
||||
'nearest_4d_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, S, M, M), (4, None, 'nearest', None, False),
|
||||
'nearest_4d_with_size_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, S, M, M), (None, 2., 'bilinear', None, False),
|
||||
'bilinear_4d_with_scale_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, S, M, M), (4, None, 'bilinear', None, False),
|
||||
'bilinear_4d_with_size_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, S, M, M), (None, 2., 'bicubic', None, False),
|
||||
'bicubic_4d_with_scale_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, S, M, M), (4, None, 'bicubic', None, False),
|
||||
'bicubic_4d_with_size_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, M, M), (None, 2., 'nearest', None, False),
|
||||
'nearest_3d_with_scale_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, M, M), (4, None, 'nearest', None, False),
|
||||
'nearest_3d_with_size_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, M, M), (None, 2., 'linear', None, False),
|
||||
'linear_3d_with_scale_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, M, M), (4, None, 'linear', None, False),
|
||||
'linear_3d_with_size_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (None, 2., 'nearest', None, False),
|
||||
'nearest_5d_with_scale_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (4, None, 'nearest', None, False),
|
||||
'nearest_5d_with_size_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (None, 2., 'trilinear', None, False),
|
||||
'trilinear_5d_with_scale_not_recompute_scale_factor'),
|
||||
('interpolate', torch.randn(S, M, M, M, M), (4, None, 'trilinear', None, False),
|
||||
'trilinear_5d_with_size_not_recompute_scale_factor'),
|
||||
]
|
||||
return nn_functional_tests
|
||||
|
||||
script_template = '''
|
||||
def the_method({}):
|
||||
return {}
|
||||
'''
|
||||
|
||||
def value_to_literal(value):
|
||||
if isinstance(value, str):
|
||||
# Quotes string and escapes special characters
|
||||
return ascii(value)
|
||||
if isinstance(value, torch.Tensor):
|
||||
return 'torch.' + str(value)
|
||||
else:
|
||||
return str(value)
|
||||
|
||||
def get_call(method_name, func_type, args, kwargs):
|
||||
kwargs_str = ', '.join([k + '=' + value_to_literal(v) for k, v in kwargs.items()])
|
||||
self_arg = args[0]
|
||||
if func_type == 'method':
|
||||
args = args[1:]
|
||||
|
||||
argument_str = ', '.join(args)
|
||||
argument_str += ', ' if len(args) and len(kwargs) else ''
|
||||
argument_str += kwargs_str
|
||||
|
||||
if func_type == 'functional' or func_type == 'function':
|
||||
call = f'torch.{method_name}({argument_str})'
|
||||
elif func_type == 'method':
|
||||
call = f'{self_arg}.{method_name}({argument_str})'
|
||||
elif func_type == 'nn_functional':
|
||||
call = f'torch.nn.functional.{method_name}({argument_str})'
|
||||
else:
|
||||
raise TypeError('Unsupported function type')
|
||||
|
||||
return call
|
||||
|
||||
def get_constant(x):
|
||||
if x == inf:
|
||||
return 'math.inf'
|
||||
if x == -inf:
|
||||
return '-math.inf'
|
||||
return x
|
||||
|
||||
def get_script_args(args):
|
||||
formals: list[str] = []
|
||||
tensors: list[torch.Tensor | list[torch.Tensor]] = []
|
||||
actuals: list[str] = []
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor):
|
||||
name = f'i{len(formals)}'
|
||||
formals.append(name)
|
||||
actuals.append(name)
|
||||
tensors.append(arg)
|
||||
elif is_iterable_of_tensors(arg):
|
||||
name = f'i{len(formals)}'
|
||||
formals.append(name + ': List[torch.Tensor]')
|
||||
actuals.append(name)
|
||||
tensors.append(list(arg))
|
||||
elif isinstance(arg, str):
|
||||
actuals.append(f"'{arg}'")
|
||||
else:
|
||||
actuals.append(str(get_constant(arg)))
|
||||
return (formals, tensors, actuals)
|
||||
|
||||
# create a script function from (name, func_type, output_process_fn),
|
||||
# and returns the compiled function and example inputs
|
||||
def gen_script_fn_and_args(method_name, func_type, *args, **kwargs):
|
||||
formals, tensors, actuals = get_script_args(args)
|
||||
call = get_call(method_name, func_type, actuals, kwargs)
|
||||
script = script_template.format(', '.join(formals), call)
|
||||
CU = torch.jit.CompilationUnit(script)
|
||||
return CU.the_method, tensors
|
||||
|
||||
# create a script function from (name, func_type),
|
||||
# returns a function takes in (args, kwargs) and runs the compiled function
|
||||
def create_script_fn(self, method_name, func_type):
|
||||
# function returns tuple containing original output and
|
||||
# filtered output to be used in checking gradients
|
||||
def script_fn(*args, **kwargs):
|
||||
fn, tensors = gen_script_fn_and_args(method_name, func_type, *args, **kwargs)
|
||||
self.assertExportImport(fn.graph, tensors)
|
||||
output = fn(*tensors)
|
||||
# skip type annotate function attributes for now, see: https://github.com/python/mypy/issues/2087
|
||||
script_fn.last_graph = fn.graph_for(*tensors) # type: ignore[attr-defined]
|
||||
return output
|
||||
return script_fn
|
||||
|
||||
class SplitInputs:
|
||||
all_tensors: list[Any]
|
||||
tensor_args: list[Any]
|
||||
nontensor_args: list[Any]
|
||||
arg_types: list[str]
|
||||
tensor_kwargs: dict[str, Any]
|
||||
kwarg_order: list[str]
|
||||
nontensor_kwargs: dict[str, Any]
|
||||
kwarg_types: dict[str, Any]
|
||||
|
||||
@staticmethod
|
||||
def _is_tensor_input(arg):
|
||||
return isinstance(arg, torch.Tensor) or is_iterable_of_tensors(arg)
|
||||
|
||||
def __init__(self, args, kwargs):
|
||||
self.arg_types = ['t' if self._is_tensor_input(arg) else 's' for arg in args]
|
||||
self.kwarg_types = {k: 't' if self._is_tensor_input(v) else 's' for k, v in kwargs.items()}
|
||||
self.tensor_args = [arg for arg in args if self._is_tensor_input(arg)]
|
||||
self.nontensor_args = [arg for arg in args if not self._is_tensor_input(arg)]
|
||||
self.tensor_kwargs = {k: v for k, v in kwargs.items() if self._is_tensor_input(v)}
|
||||
self.nontensor_kwargs = {k: v for k, v in kwargs.items() if not self._is_tensor_input(v)}
|
||||
self.all_tensors = [*self.tensor_args, *[v for k, v in self.tensor_kwargs.items()]]
|
||||
self.kwarg_order = [k for k, v in kwargs.items()]
|
||||
|
||||
def nontensors_match(self, other: 'SplitInputs'):
|
||||
if self.arg_types != other.arg_types:
|
||||
return False
|
||||
if self.kwarg_types != other.kwarg_types:
|
||||
return False
|
||||
if self.kwarg_order != other.kwarg_order:
|
||||
return False
|
||||
if self.nontensor_args != other.nontensor_args:
|
||||
return False
|
||||
if self.nontensor_kwargs != other.nontensor_kwargs:
|
||||
return False
|
||||
return True
|
||||
|
||||
# make a new function where all non-tensor arguments in 'args' have been partially
|
||||
# applied, and all tensor arguments remain.
|
||||
# used to trace functions when some arguments are not tensors
|
||||
def partial_apply_nontensors(fn, args, kwargs):
|
||||
inputs = SplitInputs(args, kwargs)
|
||||
|
||||
def new_fn(*tensors_):
|
||||
tensors = iter(tensors_)
|
||||
full_args = [args[i] if s == 's' else next(tensors) for i, s in enumerate(inputs.arg_types)]
|
||||
full_kwargs = {k: kwargs[k] if s == 's' else next(tensors) for k, s in inputs.kwarg_types.items()}
|
||||
return fn(*full_args, **full_kwargs)
|
||||
|
||||
return new_fn, inputs
|
||||
|
||||
# create a trace function from input fn
|
||||
def create_traced_fn(self, fn, cache_traced_fn=False):
|
||||
def traced_fn(*inputs, **kwargs):
|
||||
# `check_trace` is set to False because check_trace is run with @no_grad
|
||||
# Also, `check_against_reference` already does all the checks
|
||||
# against python function
|
||||
fn_tensors, split_inputs = partial_apply_nontensors(fn, inputs, kwargs)
|
||||
if not cache_traced_fn or not hasattr(traced_fn, 'traced'):
|
||||
traced = torch.jit.trace(fn_tensors, split_inputs.all_tensors, check_trace=False)
|
||||
self.assertExportImport(traced.graph, split_inputs.all_tensors)
|
||||
output = traced(*split_inputs.all_tensors)
|
||||
if cache_traced_fn:
|
||||
traced_fn.traced = traced
|
||||
traced_fn.split_inputs = split_inputs
|
||||
else:
|
||||
# Guard to check that nontensor inputs are the same as during tracing
|
||||
self.assertTrue(traced_fn.split_inputs.nontensors_match(split_inputs))
|
||||
output = traced_fn.traced(*split_inputs.all_tensors)
|
||||
traced = traced_fn.traced
|
||||
# skip type annotate function attributes for now, see: https://github.com/python/mypy/issues/2087
|
||||
traced_fn.last_graph = traced.graph_for(*split_inputs.all_tensors) # type: ignore[attr-defined]
|
||||
traced_fn.graph = traced.graph # type: ignore[attr-defined]
|
||||
return output
|
||||
return traced_fn
|
||||
|
||||
# known to be failing in script
|
||||
EXCLUDE_SCRIPT = {
|
||||
'test_norm_fro_default',
|
||||
'test_norm_fro_cpu',
|
||||
'test_norm_nuc',
|
||||
'test_norm_fro',
|
||||
'test_norm_nuc_batched',
|
||||
|
||||
# aten op has additional cudnn argument
|
||||
'test_nn_unfold',
|
||||
|
||||
# flaky test - TODO fix
|
||||
'test_nn_ctc_loss',
|
||||
|
||||
# unknown builtin op
|
||||
'test_nn_fold',
|
||||
|
||||
# jit doesn't support sparse tensors.
|
||||
'test_to_sparse',
|
||||
'test_to_sparse_dim',
|
||||
}
|
||||
|
||||
# generates a script function and set of example inputs
|
||||
# from a specified test in the format of nn_functional_tests
|
||||
def get_nn_functional_compiled_fn_and_inputs(name, self_size, args, variant_name='', *extra_args):
|
||||
test_name = 'test_nn_' + name
|
||||
|
||||
if variant_name != '':
|
||||
test_name = test_name + '_' + variant_name
|
||||
|
||||
self_variable = create_input((self_size,))[0][0]
|
||||
|
||||
# need to record this because methods can change the size (e.g. unsqueeze)
|
||||
args_variable, _kwargs_variable = create_input(args)
|
||||
|
||||
self_tensor = deepcopy(self_variable.data)
|
||||
args_tensor = deepcopy(unpack_variables(args_variable))
|
||||
|
||||
f_args_variable = (self_variable,) + args_variable
|
||||
f_args_tensor = (self_tensor,) + args_tensor # noqa: F841
|
||||
with torch._jit_internal._disable_emit_hooks():
|
||||
script_fn, inputs = gen_script_fn_and_args(name, "nn_functional", *f_args_variable)
|
||||
return script_fn, inputs
|
||||
|
||||
|
||||
|
||||
EXCLUDE_SCRIPT_MODULES = {
|
||||
'test_nn_AdaptiveAvgPool2d_tuple_none',
|
||||
'test_nn_AdaptiveAvgPool3d_tuple_none',
|
||||
'test_nn_AdaptiveMaxPool2d_tuple_none',
|
||||
'test_nn_AdaptiveMaxPool3d_tuple_none',
|
||||
|
||||
# Doesn't use future division, so this is not supported
|
||||
'test_nn_CrossMapLRN2d',
|
||||
# Derivative for aten::_scaled_dot_product_flash_attention_backward is not implemented
|
||||
'test_nn_TransformerDecoderLayer_gelu_activation',
|
||||
'test_nn_TransformerDecoderLayer_relu_activation',
|
||||
'test_nn_TransformerEncoderLayer_gelu_activation',
|
||||
'test_nn_TransformerEncoderLayer_relu_activation',
|
||||
'test_nn_Transformer_multilayer_coder',
|
||||
}
|
||||
|
||||
script_method_template = '''
|
||||
def forward({}):
|
||||
return {}
|
||||
'''
|
||||
|
||||
def create_script_module(self, nn_module, constructor_args, *args, **kwargs):
|
||||
def script_module(*args, **kwargs):
|
||||
_formals, tensors, actuals = get_script_args(args)
|
||||
|
||||
method_args = ', '.join(['self'] + actuals)
|
||||
call_args_str = ', '.join(actuals)
|
||||
call = f"self.submodule({call_args_str})"
|
||||
script = script_method_template.format(method_args, call)
|
||||
|
||||
submodule_constants = []
|
||||
if kwargs.get('is_constant'):
|
||||
submodule_constants = ['submodule']
|
||||
|
||||
# Create module to use the script method
|
||||
class TheModule(torch.jit.ScriptModule):
|
||||
__constants__ = submodule_constants
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.submodule = nn_module(*constructor_args)
|
||||
|
||||
def make_module(script):
|
||||
module = TheModule()
|
||||
# check __repr__
|
||||
str(module)
|
||||
module.define(script)
|
||||
return module
|
||||
|
||||
module = make_module(script)
|
||||
if self:
|
||||
self.assertExportImportModule(module, tensors)
|
||||
module(*args)
|
||||
# skip type annotate function attributes for now, see: https://github.com/python/mypy/issues/2087
|
||||
create_script_module.last_graph = module.graph # type: ignore[attr-defined]
|
||||
return module
|
||||
return script_module
|
||||
|
||||
def check_alias_annotation(method_name, args, kwargs, *, aten_name, func_type='method'):
|
||||
formals, tensors, actuals = get_script_args(args)
|
||||
call = get_call(method_name, func_type, actuals, kwargs)
|
||||
script = script_template.format(', '.join(formals), call)
|
||||
CU = torch.jit.CompilationUnit(script)
|
||||
# to clean up IR
|
||||
torch._C._jit_pass_inline(CU.the_method.graph)
|
||||
torch._C._jit_pass_constant_propagation(CU.the_method.graph)
|
||||
torch._C._jit_check_alias_annotation(CU.the_method.graph, tuple(tensors), aten_name)
|
||||
|
||||
def get_nn_module_name_from_kwargs(**kwargs):
|
||||
if 'module_name' in kwargs:
|
||||
return kwargs['module_name']
|
||||
elif 'fullname' in kwargs:
|
||||
return kwargs['fullname']
|
||||
elif 'constructor' in kwargs:
|
||||
return kwargs['constructor'].__name__
|
||||
|
||||
def get_nn_mod_test_name(**kwargs):
|
||||
if 'fullname' in kwargs:
|
||||
test_name = kwargs['fullname']
|
||||
else:
|
||||
test_name = get_nn_module_name_from_kwargs(**kwargs)
|
||||
if 'desc' in kwargs:
|
||||
test_name = f"{test_name}_{kwargs['desc']}"
|
||||
return f'test_nn_{test_name}'
|
||||
|
||||
def get_nn_module_class_from_kwargs(**kwargs):
|
||||
name = get_nn_module_name_from_kwargs(**kwargs)
|
||||
index = name.find("_")
|
||||
if index == -1:
|
||||
return name
|
||||
else:
|
||||
return name[0:name.find("_")]
|
||||
|
||||
def try_get_nn_module_compiled_mod_and_inputs(*args, **kwargs):
|
||||
name = get_nn_module_name_from_kwargs(**kwargs)
|
||||
|
||||
if 'desc' in kwargs and 'eval' in kwargs['desc']:
|
||||
# eval() is not supported, so skip these tests
|
||||
return
|
||||
|
||||
test_name = name
|
||||
if 'desc' in kwargs:
|
||||
test_name = f"{test_name}_{kwargs['desc']}"
|
||||
test_name = get_nn_mod_test_name(**kwargs)
|
||||
|
||||
if test_name in EXCLUDE_SCRIPT_MODULES:
|
||||
return
|
||||
if 'constructor' in kwargs:
|
||||
nn_module = kwargs['constructor']
|
||||
else:
|
||||
nn_module = getattr(torch.nn, name)
|
||||
|
||||
if "FunctionalModule" in str(nn_module):
|
||||
return
|
||||
|
||||
if 'constructor_args_fn' in kwargs:
|
||||
constructor_args = kwargs['constructor_args_fn']()
|
||||
else:
|
||||
constructor_args = kwargs.get('constructor_args', ())
|
||||
|
||||
# Set up inputs from tuple of sizes or constructor fn
|
||||
input_dtype = torch.double
|
||||
if 'input_fn' in kwargs:
|
||||
input = kwargs['input_fn']()
|
||||
if isinstance(input, torch.Tensor):
|
||||
input = (input,)
|
||||
|
||||
if all(tensor.is_complex() for tensor in input):
|
||||
input_dtype = torch.cdouble
|
||||
else:
|
||||
input = (kwargs['input_size'],)
|
||||
|
||||
# Extra parameters to forward()
|
||||
if 'extra_args' in kwargs:
|
||||
input = input + kwargs['extra_args']
|
||||
|
||||
if 'target_size' in kwargs:
|
||||
input = input + (kwargs['target_size'],)
|
||||
elif 'target_fn' in kwargs:
|
||||
if torch.is_tensor(input):
|
||||
input = (input,)
|
||||
input = input + (kwargs['target_fn'](),)
|
||||
|
||||
args_variable, _kwargs_variable = create_input(input, dtype=input_dtype)
|
||||
f_args_variable = deepcopy(unpack_variables(args_variable))
|
||||
out_var = deepcopy(f_args_variable)
|
||||
|
||||
|
||||
_args, mod = f_args_variable, create_script_module(
|
||||
None, nn_module, constructor_args, *f_args_variable
|
||||
)(*f_args_variable)
|
||||
|
||||
return mod, out_var
|
||||
|
||||
|
||||
def get_all_nn_module_tests():
|
||||
# additional modules test
|
||||
# TODO: delete this list once we make all nn_tests work
|
||||
additional_module_tests = [
|
||||
{
|
||||
'module_name': 'Bilinear',
|
||||
'constructor_args': (S, S, M),
|
||||
'input_size': (S, S),
|
||||
'extra_args': ((S, S),)
|
||||
},
|
||||
{
|
||||
'module_name': 'RNNCell',
|
||||
'constructor_args': (S, S),
|
||||
'input_size': (S, S),
|
||||
},
|
||||
{
|
||||
'module_name': 'LSTMCell',
|
||||
'constructor_args': (S, S),
|
||||
'input_size': (S, S),
|
||||
},
|
||||
{
|
||||
'module_name': 'GRUCell',
|
||||
'constructor_args': (S, S),
|
||||
'input_size': (S, S),
|
||||
},
|
||||
{
|
||||
'module_name': 'MultiheadAttention',
|
||||
'constructor_args': (128, 8),
|
||||
'input_size': (10, 8, 128),
|
||||
'extra_args': (torch.randn(10, 8, 128), torch.randn(10, 8, 128)),
|
||||
'slowTest': True
|
||||
},
|
||||
{
|
||||
'module_name': 'Transformer',
|
||||
'constructor_args': (1, 1, 1, 1, 2),
|
||||
'input_size': (3, 1, 1),
|
||||
'extra_args': (torch.randn(1, 1, 1),),
|
||||
'slowTest': True
|
||||
}
|
||||
]
|
||||
|
||||
return module_tests + get_new_module_tests() + additional_module_tests
|
||||
@@ -0,0 +1,897 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
# Torch
|
||||
from torch.autograd import Variable
|
||||
from torch.autograd.function import _nested_map
|
||||
from torch.jit.annotations import BroadcastingList2, BroadcastingList3 # noqa: F401
|
||||
|
||||
from torch.onnx import OperatorExportTypes
|
||||
import torch
|
||||
import torch.cuda
|
||||
import torch.jit
|
||||
import torch.jit._logging
|
||||
import torch.jit.frontend
|
||||
import torch.jit.quantized
|
||||
import zipfile
|
||||
import functools
|
||||
|
||||
# Testing utils
|
||||
from torch.testing import FileCheck
|
||||
from torch.testing._internal.common_utils import IS_WINDOWS, \
|
||||
freeze_rng_state, enable_profiling_mode_for_profiling_tests, ProfilingMode, TEST_BAILOUTS, \
|
||||
is_iterable_of_tensors
|
||||
from torch.testing._internal.common_jit import JitCommonTestCase
|
||||
from torch.testing._internal.common_utils import enable_profiling_mode # noqa: F401
|
||||
|
||||
# Standard library
|
||||
from contextlib import contextmanager
|
||||
from functools import reduce
|
||||
from io import StringIO
|
||||
from collections import defaultdict
|
||||
|
||||
import importlib.util
|
||||
import inspect
|
||||
import io
|
||||
import math
|
||||
import os
|
||||
import pickle
|
||||
import sys
|
||||
import tempfile
|
||||
import textwrap
|
||||
from importlib.abc import Loader
|
||||
from typing import Any
|
||||
|
||||
RUN_CUDA = torch.cuda.is_available()
|
||||
RUN_CUDA_MULTI_GPU = RUN_CUDA and torch.cuda.device_count() > 1
|
||||
RUN_CUDA_HALF = RUN_CUDA
|
||||
# HIP supports half, no version check necessary
|
||||
if torch.cuda.is_available() and not torch.version.hip:
|
||||
CUDA_VERSION = torch._C._cuda_getCompiledVersion()
|
||||
for d in range(torch.cuda.device_count()):
|
||||
major = torch.cuda.get_device_capability(d)[0]
|
||||
if (major < 6):
|
||||
RUN_CUDA_HALF = False
|
||||
|
||||
def execWrapper(code, glob, loc):
|
||||
exec(code, glob, loc)
|
||||
|
||||
def do_input_map(fn, input):
|
||||
return _nested_map(lambda t: isinstance(t, torch.Tensor), fn)(input)
|
||||
|
||||
def clear_class_registry():
|
||||
torch._C._jit_clear_class_registry()
|
||||
torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()
|
||||
torch.jit._state._clear_class_state()
|
||||
|
||||
def get_execution_plan(graph_executor_state):
|
||||
execution_plans = list(graph_executor_state.execution_plans.values())
|
||||
num_plans = len(execution_plans)
|
||||
if num_plans != 1:
|
||||
raise RuntimeError('This test assumes this GraphExecutor should '
|
||||
f'only have one execution plan, got: {num_plans}')
|
||||
return execution_plans[0]
|
||||
|
||||
class _AssertRaisesRegexWithHighlightContext:
|
||||
"""
|
||||
A context manager that is useful for checking that error messages highlight
|
||||
the correct part of the source code.
|
||||
"""
|
||||
|
||||
def __init__(self, test_case, exception, regex, highlight):
|
||||
self.test_case = test_case
|
||||
self.exception_type = exception
|
||||
self.regex = regex
|
||||
self.highlight = highlight
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, type, value, traceback):
|
||||
with self.test_case.assertRaisesRegex(self.exception_type, self.regex):
|
||||
if type:
|
||||
raise value
|
||||
|
||||
if self.highlight:
|
||||
FileCheck().check_source_highlighted(self.highlight).run(str(value))
|
||||
|
||||
return True
|
||||
|
||||
FUSION_GROUP = "prim::TensorExprGroup"
|
||||
|
||||
class JitTestCase(JitCommonTestCase):
|
||||
_do_cuda_memory_leak_check = True
|
||||
_restored_warnings = False
|
||||
|
||||
class capture_stdout(list):
|
||||
"""
|
||||
Replace sys.stdout with a temporary StringIO
|
||||
"""
|
||||
def __enter__(self):
|
||||
self.sys_stdout = sys.stdout
|
||||
self.stringio = StringIO()
|
||||
sys.stdout = self.stringio
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
self.append(str(self.stringio.getvalue()))
|
||||
del self.stringio
|
||||
sys.stdout = self.sys_stdout
|
||||
|
||||
class capture_stderr(list):
|
||||
"""
|
||||
Replace sys.stderr with a temporary StringIO
|
||||
"""
|
||||
def __enter__(self):
|
||||
self.sys_stderr = sys.stderr
|
||||
self.stringio = StringIO()
|
||||
sys.stderr = self.stringio
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
self.append(str(self.stringio.getvalue()))
|
||||
del self.stringio
|
||||
sys.stderr = self.sys_stderr
|
||||
|
||||
def setHooks(self):
|
||||
torch._C._jit_set_emit_hooks(self.emitModuleHook, self.emitFunctionHook)
|
||||
|
||||
def clearHooks(self):
|
||||
torch._C._jit_set_emit_hooks(None, None)
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
# unittest overrides all warning filters and forces all of them to show up
|
||||
# after we install our own to silence those coming from inside PyTorch.
|
||||
# This will ensure that our filter still takes precedence.
|
||||
if not JitTestCase._restored_warnings:
|
||||
torch.jit.TracerWarning.ignore_lib_warnings()
|
||||
JitTestCase._restored_warnings = True
|
||||
self.setHooks()
|
||||
|
||||
def tearDown(self):
|
||||
super().tearDown()
|
||||
# needs to be cleared because python might be unloaded before
|
||||
# the callback gets destructed
|
||||
self.clearHooks()
|
||||
clear_class_registry()
|
||||
|
||||
def assertAllFused(self, graph, except_for=()):
|
||||
|
||||
# note this helper collects nodes on 'fast path' only
|
||||
# i.e. the true blocks of specialized checks
|
||||
def get_nodes_and_parents_recursively(block, kind, acc):
|
||||
for node in block.nodes():
|
||||
if node.kind() == kind:
|
||||
acc[block].append(node)
|
||||
elif node.kind() == 'prim::DifferentiableGraph':
|
||||
get_nodes_and_parents_recursively(node.g('Subgraph'), kind, acc)
|
||||
elif node.kind() == 'prim::If' and (node.inputs().__next__().node().kind() == 'aten::all' or
|
||||
node.inputs().__next__().node().kind() == 'prim::TypeCheck' or
|
||||
node.inputs().__next__().node().kind() == 'prim::RequiresGradCheck'):
|
||||
get_nodes_and_parents_recursively(node.blocks().__next__(), kind, acc)
|
||||
else:
|
||||
for inner_block in node.blocks():
|
||||
get_nodes_and_parents_recursively(inner_block, kind, acc)
|
||||
|
||||
allowed_nodes = {'prim::Constant', FUSION_GROUP, 'prim::BailoutTemplate',
|
||||
'prim::TupleConstruct', 'prim::If', 'prim::TypeCheck', 'prim::RequiresGradCheck'} | set(except_for)
|
||||
|
||||
fusion_groups : dict[torch._C.Block, list[torch._C.Node]] = defaultdict(list)
|
||||
get_nodes_and_parents_recursively(graph, FUSION_GROUP, fusion_groups)
|
||||
self.assertTrue(len(fusion_groups) == 1, f'got {graph}')
|
||||
(graph, fusion_nodes) = next(iter(fusion_groups.items()))
|
||||
# the block contains one FUSION_GROUP and the rest of nodes are `allowed_nodes`
|
||||
self.assertTrue(len(fusion_nodes) == 1, f'got {graph}')
|
||||
self.assertTrue(all(node.kind() in allowed_nodes for node in graph.nodes()),
|
||||
f'got {graph}')
|
||||
|
||||
def _isHookExceptionOk(self, e):
|
||||
se = str(e)
|
||||
allowed = ("Could not export Python function",
|
||||
"closures are not exportable")
|
||||
for a in allowed:
|
||||
if a in se:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _compared_saved_loaded(self, m):
|
||||
def extract_files(buffer):
|
||||
# crack open the zip format to get at the main module code
|
||||
with zipfile.ZipFile(buffer) as archive:
|
||||
# check that we have no duplicate names
|
||||
self.assertEqual(len(set(archive.namelist())), len(archive.namelist()))
|
||||
files = list(filter(lambda x: x.startswith('archive/code/'), archive.namelist()))
|
||||
# unwrap all the code files into strings
|
||||
code_files_str = filter(lambda x: x.endswith('.py'), files)
|
||||
code_files = []
|
||||
for f in code_files_str:
|
||||
with archive.open(f) as stream:
|
||||
code_files.append("".join([line.decode() for line in stream]))
|
||||
|
||||
# unpickled all the debug files
|
||||
debug_files_str = filter(lambda f: f.endswith('.debug_pkl'), files)
|
||||
debug_files = []
|
||||
for f in debug_files_str:
|
||||
with archive.open(f) as stream:
|
||||
debug_files.append(pickle.load(stream))
|
||||
return code_files, debug_files
|
||||
|
||||
# disable the hook while we parse code, otherwise we will re-enter the hook
|
||||
with torch._jit_internal._disable_emit_hooks():
|
||||
try:
|
||||
# short-circuit if this is an empty function or module
|
||||
if len(m.code) == 0:
|
||||
return
|
||||
if isinstance(m, torch._C.ScriptModule):
|
||||
if len(m._method_names()) == 0:
|
||||
return
|
||||
|
||||
# save the module to a buffer
|
||||
buffer = io.BytesIO()
|
||||
torch.jit.save(m, buffer)
|
||||
# copy the data in the buffer so we can restore it later. This
|
||||
# is because py2 and py3 have different semantics with zipfile
|
||||
# and it's easier to just work with a fresh copy each time.
|
||||
buffer_copy = buffer.getvalue()
|
||||
|
||||
code_files, _debug_files = extract_files(buffer)
|
||||
|
||||
except RuntimeError as e:
|
||||
if not self._isHookExceptionOk(e):
|
||||
raise
|
||||
else:
|
||||
return
|
||||
|
||||
# import the model again (from a the copy we made of the original)
|
||||
buffer2 = io.BytesIO(buffer_copy)
|
||||
imported = torch.jit.load(buffer2)
|
||||
|
||||
# save it again
|
||||
saved_module_buffer_2 = io.BytesIO()
|
||||
torch.jit.save(imported, saved_module_buffer_2)
|
||||
|
||||
saved_module_buffer_2.seek(0)
|
||||
code_files_2, _debug_files_2 = extract_files(saved_module_buffer_2)
|
||||
|
||||
for a, b in zip(code_files, code_files_2, strict=True):
|
||||
self.assertMultiLineEqual(a, b)
|
||||
|
||||
if isinstance(m, torch._C.ScriptModule):
|
||||
self.assertTrue(torch._C._ivalue_tags_match(m, imported._c))
|
||||
|
||||
|
||||
def emitFunctionHook(self, func):
|
||||
# func has invalid names for export, skip the jitter check
|
||||
if func.name == "<lambda>" or "aten::" in func.name:
|
||||
return
|
||||
self._compared_saved_loaded(func)
|
||||
|
||||
def emitModuleHook(self, module):
|
||||
self._compared_saved_loaded(module)
|
||||
|
||||
|
||||
def getExportImportCopyWithPacking(self, m, also_test_file=True, map_location=None):
|
||||
buffer = io.BytesIO()
|
||||
m.apply(lambda s: s._pack() if s._c._has_method('_pack') else None)
|
||||
torch.jit.save(m, buffer)
|
||||
m.apply(lambda s: s._unpack() if s._c._has_method('_unpack') else None)
|
||||
buffer.seek(0)
|
||||
imported = torch.jit.load(buffer, map_location=map_location)
|
||||
imported.apply(lambda s: s._unpack() if s._c._has_method('_unpack') else None)
|
||||
|
||||
if not also_test_file:
|
||||
return imported
|
||||
|
||||
# Ideally we would like to not have to manually delete the file, but NamedTemporaryFile
|
||||
# opens the file, and it cannot be opened multiple times in Windows. To support Windows,
|
||||
# close the file after creation and try to remove it manually
|
||||
with tempfile.NamedTemporaryFile(delete=False) as f:
|
||||
try:
|
||||
f.close()
|
||||
imported.save(f.name)
|
||||
result = torch.jit.load(f.name, map_location=map_location)
|
||||
finally:
|
||||
os.unlink(f.name)
|
||||
|
||||
result.apply(lambda s: s._unpack() if s._c._has_method('_unpack') else None)
|
||||
return result
|
||||
|
||||
def assertGraphContains(self, graph, kind, consider_subgraphs=False):
|
||||
|
||||
if consider_subgraphs:
|
||||
strgraph = str(graph)
|
||||
count = strgraph.count(kind) - strgraph.count(f'with {kind}')
|
||||
self.assertTrue(count > 0)
|
||||
return
|
||||
|
||||
def nodes(block):
|
||||
out = []
|
||||
for node in block.nodes():
|
||||
if node.kind() == kind:
|
||||
out.append(node)
|
||||
for block in node.blocks():
|
||||
out += nodes(block)
|
||||
return out
|
||||
|
||||
out_nodes = nodes(graph)
|
||||
self.assertTrue(len(out_nodes) > 0)
|
||||
|
||||
def assertGraphContainsExactly(self, graph, kind, num_kind_nodes, consider_subgraphs=False):
|
||||
def perform_assert(graph, kind, actual, expected, consider_subgraphs):
|
||||
if actual == expected:
|
||||
return
|
||||
subgraph = 'including' if consider_subgraphs else 'excluding'
|
||||
raise AssertionError(
|
||||
f'{graph}\nError: graph contains {actual} {kind} nodes ({subgraph} subgraphs) but expected {expected}')
|
||||
|
||||
if consider_subgraphs:
|
||||
strgraph = str(graph)
|
||||
count = strgraph.count(kind) - strgraph.count(f'with {kind}')
|
||||
perform_assert(graph, kind, count, num_kind_nodes,
|
||||
consider_subgraphs)
|
||||
return
|
||||
|
||||
def nodes(block):
|
||||
out = []
|
||||
for node in block.nodes():
|
||||
if node.kind() == kind:
|
||||
out.append(node)
|
||||
for block in node.blocks():
|
||||
out += nodes(block)
|
||||
return out
|
||||
|
||||
out_nodes = nodes(graph)
|
||||
perform_assert(graph, kind, len(out_nodes), num_kind_nodes,
|
||||
consider_subgraphs)
|
||||
|
||||
def assertExpectedONNXGraph(self, g, *args, **kwargs):
|
||||
g = torch.onnx._optimize_trace(g, operator_export_type=OperatorExportTypes.ONNX)
|
||||
self.assertExpectedGraph(g, *args, **kwargs)
|
||||
|
||||
def assertExpectedGraph(self, trace, *args, **kwargs):
|
||||
if isinstance(trace, torch._C.Graph):
|
||||
graph = trace
|
||||
else:
|
||||
graph = trace.graph()
|
||||
|
||||
torch._C._jit_pass_lint(graph)
|
||||
torch._C._jit_pass_dce(graph)
|
||||
torch._C._jit_pass_lint(graph)
|
||||
graph = torch._C._jit_pass_canonicalize(graph)
|
||||
torch._C._jit_pass_lint(graph)
|
||||
self.assertExpected(str(graph), *args, **kwargs)
|
||||
|
||||
def run_pass(self, name, trace):
|
||||
if isinstance(trace, torch._C.Graph):
|
||||
graph = trace
|
||||
set_graph = False
|
||||
else:
|
||||
set_graph = True
|
||||
graph = trace.graph()
|
||||
|
||||
torch._C._jit_pass_lint(graph)
|
||||
result = getattr(torch._C, '_jit_pass_' + name)(graph)
|
||||
if result is not None and not isinstance(result, bool):
|
||||
graph = result
|
||||
torch._C._jit_pass_lint(graph)
|
||||
|
||||
if set_graph:
|
||||
trace.set_graph(graph)
|
||||
return graph
|
||||
|
||||
def get_frame_vars(self, frames_up):
|
||||
frame = inspect.currentframe()
|
||||
if not frame:
|
||||
raise RuntimeError("failed to inspect frame")
|
||||
i = 0
|
||||
while i < frames_up + 1:
|
||||
frame = frame.f_back
|
||||
if not frame:
|
||||
raise RuntimeError("failed to get frame")
|
||||
i += 1
|
||||
defined_vars: dict[str, Any] = {}
|
||||
defined_vars.update(frame.f_locals)
|
||||
defined_vars.update(frame.f_globals)
|
||||
return defined_vars
|
||||
|
||||
def assertRaisesRegexWithHighlight(self, exception, regex, highlight):
|
||||
return _AssertRaisesRegexWithHighlightContext(self, exception, regex, highlight)
|
||||
|
||||
def checkScriptRaisesRegex(self, script, inputs, exception, regex,
|
||||
name=None, outputs=None, capture_output=False,
|
||||
frames_up=1, profiling=ProfilingMode.PROFILING):
|
||||
"""
|
||||
Checks that a given function will throw the correct exception,
|
||||
when executed with normal python, the string frontend, and the
|
||||
AST frontend. Logic taken from `checkScript` (see comments there
|
||||
for details)
|
||||
"""
|
||||
with enable_profiling_mode_for_profiling_tests():
|
||||
# Normal Python
|
||||
with self.assertRaisesRegex(exception, regex):
|
||||
if isinstance(script, str):
|
||||
frame = self.get_frame_vars(frames_up)
|
||||
the_locals: dict[str, Any] = {}
|
||||
execWrapper(script, glob=frame, loc=the_locals)
|
||||
frame.update(the_locals)
|
||||
|
||||
python_fn = frame[name]
|
||||
else:
|
||||
python_fn = script
|
||||
|
||||
python_fn(*inputs)
|
||||
|
||||
# String frontend
|
||||
with self.assertRaisesRegex(exception, regex):
|
||||
if isinstance(script, str):
|
||||
cu = torch.jit.CompilationUnit(script, _frames_up=frames_up)
|
||||
string_frontend = getattr(cu, name)
|
||||
else:
|
||||
source = textwrap.dedent(inspect.getsource(script))
|
||||
cu = torch.jit.CompilationUnit(source, _frames_up=frames_up)
|
||||
string_frontend = getattr(cu, script.__name__)
|
||||
|
||||
string_frontend(*inputs)
|
||||
|
||||
# Python AST frontend
|
||||
if not isinstance(script, str):
|
||||
with self.assertRaisesRegex(exception, regex):
|
||||
ge = torch.jit.script(python_fn)
|
||||
ge(*inputs)
|
||||
|
||||
def checkBailouts(self, model, inputs, expected):
|
||||
state = model.get_debug_state()
|
||||
plan = get_execution_plan(state)
|
||||
num_bailouts = plan.code.num_bailouts()
|
||||
for i in range(num_bailouts):
|
||||
plan.code.request_bailout(i)
|
||||
bailout_outputs = model(*inputs)
|
||||
self.assertEqual(bailout_outputs, expected)
|
||||
|
||||
def checkScript(self,
|
||||
script,
|
||||
inputs,
|
||||
name='func',
|
||||
optimize=True,
|
||||
inputs_requires_grad=False,
|
||||
capture_output=False,
|
||||
frames_up=1,
|
||||
profiling=ProfilingMode.PROFILING,
|
||||
atol=None,
|
||||
rtol=None):
|
||||
"""
|
||||
Checks that a given script generates the same output as the Python
|
||||
version using the given inputs.
|
||||
"""
|
||||
with torch.jit.optimized_execution(optimize), enable_profiling_mode_for_profiling_tests():
|
||||
extra_profile_runs = any(isinstance(x, torch.Tensor) and x.requires_grad for x in inputs)
|
||||
if isinstance(script, str):
|
||||
# Compile the string to a Script function
|
||||
# with enable_profiling_mode():
|
||||
cu = torch.jit.CompilationUnit(script, _frames_up=frames_up)
|
||||
|
||||
# Execute the Python function so we can run it later and get its
|
||||
# outputs
|
||||
|
||||
frame = self.get_frame_vars(frames_up)
|
||||
the_locals: dict[str, Any] = {}
|
||||
execWrapper(script, glob=frame, loc=the_locals)
|
||||
frame.update(the_locals)
|
||||
|
||||
python_fn = frame[name]
|
||||
scripted_fn = getattr(cu, name)
|
||||
else:
|
||||
|
||||
# Check the string frontend first
|
||||
source = textwrap.dedent(inspect.getsource(script))
|
||||
self.checkScript(
|
||||
source,
|
||||
inputs,
|
||||
script.__name__,
|
||||
optimize=optimize,
|
||||
inputs_requires_grad=inputs_requires_grad,
|
||||
capture_output=capture_output,
|
||||
profiling=profiling,
|
||||
frames_up=2)
|
||||
|
||||
# Continue checking the Python frontend
|
||||
scripted_fn = torch.jit.script(script, _frames_up=1)
|
||||
python_fn = script
|
||||
|
||||
if inputs_requires_grad:
|
||||
recording_inputs = do_input_map(lambda t: t.detach().requires_grad_(), inputs)
|
||||
else:
|
||||
recording_inputs = inputs
|
||||
|
||||
if capture_output:
|
||||
with self.capture_stdout() as script_stdout:
|
||||
script_outputs = scripted_fn(*recording_inputs)
|
||||
with self.capture_stdout():
|
||||
opt_script_outputs = scripted_fn(*recording_inputs)
|
||||
with self.capture_stdout():
|
||||
python_outputs = python_fn(*inputs)
|
||||
if not IS_WINDOWS:
|
||||
self.assertExpected(script_stdout[0], subname='stdout')
|
||||
self.assertEqual(python_outputs, opt_script_outputs, atol=atol, rtol=rtol)
|
||||
else:
|
||||
# profiling run
|
||||
script_outputs = scripted_fn(*recording_inputs)
|
||||
if inputs_requires_grad or extra_profile_runs:
|
||||
opt_script_outputs = scripted_fn(*recording_inputs)
|
||||
# optimized run
|
||||
opt_script_outputs = scripted_fn(*recording_inputs)
|
||||
if TEST_BAILOUTS:
|
||||
self.checkBailouts(scripted_fn, inputs, opt_script_outputs)
|
||||
python_outputs = python_fn(*inputs)
|
||||
self.assertEqual(python_outputs, script_outputs, atol=atol, rtol=rtol)
|
||||
self.assertEqual(script_outputs, opt_script_outputs, atol=atol, rtol=rtol)
|
||||
return scripted_fn
|
||||
|
||||
def checkTrace(self, func, reference_tensors, input_tensors=None,
|
||||
drop=None, allow_unused=False, verbose=False,
|
||||
inputs_require_grads=True, check_tolerance=1e-5, export_import=True,
|
||||
_force_outplace=False, grad_atol=None, grad_rtol=None):
|
||||
|
||||
# TODO: check gradients for parameters, not just inputs
|
||||
def allSum(vs):
|
||||
# drop allows us to remove some values from ever being used
|
||||
# to test unused outputs
|
||||
if drop is not None:
|
||||
vs = vs[:-drop]
|
||||
# we don't want all the grad for all the outputs to be the same
|
||||
# so we multiply each by a constant
|
||||
return sum(math.log(i + 2) * v.sum() for i, v in enumerate(vs) if v is not None)
|
||||
if input_tensors is None:
|
||||
input_tensors = reference_tensors
|
||||
|
||||
def flatten_inputs(inputs):
|
||||
def input_reduce(input, fn, acc):
|
||||
if isinstance(input, torch.Tensor):
|
||||
fn(input, acc)
|
||||
elif isinstance(input, dict):
|
||||
reduce(lambda acc, key: input_reduce(input[key], fn, acc), input, acc)
|
||||
else:
|
||||
reduce(lambda acc, val: input_reduce(val, fn, acc), input, acc)
|
||||
return acc
|
||||
return tuple(input_reduce(recording_inputs, lambda t, acc: acc.append(t), []))
|
||||
|
||||
nograd_inputs = reference_tensors
|
||||
if inputs_require_grads:
|
||||
recording_inputs = do_input_map(lambda t: t.clone().requires_grad_(), reference_tensors)
|
||||
flattened_recording_inputs = flatten_inputs(recording_inputs)
|
||||
else:
|
||||
recording_inputs = reference_tensors
|
||||
|
||||
# `check_trace` is set to False because check_trace is run with @no_grad
|
||||
# Also, `checkTrace` already does all the checks
|
||||
# against python function
|
||||
ge = torch.jit.trace(func, input_tensors, check_tolerance=check_tolerance,
|
||||
_force_outplace=_force_outplace, check_trace=False)
|
||||
|
||||
if export_import:
|
||||
ge = self.getExportImportCopy(ge)
|
||||
|
||||
if verbose:
|
||||
print(ge.graph)
|
||||
|
||||
# test no gradients case
|
||||
outputs = func(*nograd_inputs)
|
||||
outputs_ge = ge(*nograd_inputs)
|
||||
self.assertEqual(outputs, outputs_ge)
|
||||
|
||||
# test gradients case
|
||||
outputs = func(*recording_inputs)
|
||||
if inputs_require_grads:
|
||||
grads = torch.autograd.grad(allSum(outputs), flattened_recording_inputs,
|
||||
allow_unused=allow_unused)
|
||||
|
||||
outputs_ge = ge(*recording_inputs)
|
||||
if inputs_require_grads:
|
||||
grads_ge = torch.autograd.grad(allSum(outputs_ge), flattened_recording_inputs,
|
||||
allow_unused=allow_unused)
|
||||
self.assertEqual(outputs, outputs_ge)
|
||||
if inputs_require_grads:
|
||||
self.assertEqual(grads, grads_ge, atol=grad_atol, rtol=grad_rtol)
|
||||
|
||||
# test the grad grad case
|
||||
outputs = func(*recording_inputs)
|
||||
l1 = allSum(outputs)
|
||||
if inputs_require_grads:
|
||||
grads = torch.autograd.grad(l1, flattened_recording_inputs, create_graph=True,
|
||||
allow_unused=allow_unused)
|
||||
if inputs_require_grads:
|
||||
l2 = (allSum(grads) * l1)
|
||||
grads2 = torch.autograd.grad(l2, flattened_recording_inputs, allow_unused=allow_unused)
|
||||
|
||||
if inputs_require_grads:
|
||||
recording_inputs = do_input_map(lambda t: Variable(t, requires_grad=True), reference_tensors)
|
||||
flattened_recording_inputs = flatten_inputs(recording_inputs)
|
||||
|
||||
outputs_ge = ge(*recording_inputs)
|
||||
l1_ge = allSum(outputs_ge)
|
||||
if inputs_require_grads:
|
||||
grads_ge = torch.autograd.grad(
|
||||
l1_ge, flattened_recording_inputs, create_graph=True, allow_unused=allow_unused)
|
||||
|
||||
if inputs_require_grads:
|
||||
l2_ge = (allSum(grads_ge) * l1_ge)
|
||||
grads2_ge = torch.autograd.grad(l2_ge, flattened_recording_inputs, allow_unused=allow_unused)
|
||||
|
||||
self.assertEqual(outputs, outputs_ge)
|
||||
if inputs_require_grads:
|
||||
self.assertEqual(grads, grads_ge, atol=grad_atol, rtol=grad_rtol)
|
||||
for g2, g2_ge in zip(grads2, grads2_ge, strict=True):
|
||||
if g2 is None and g2_ge is None:
|
||||
continue
|
||||
self.assertEqual(g2, g2_ge, atol=8e-4, rtol=8e-4)
|
||||
|
||||
return ge
|
||||
|
||||
def checkModule(self, nn_module, args):
|
||||
"""
|
||||
Check that a nn.Module's results in Script mode match eager and that it
|
||||
can be exported
|
||||
"""
|
||||
sm = torch.jit.script(nn_module)
|
||||
|
||||
with freeze_rng_state():
|
||||
eager_out = nn_module(*args)
|
||||
|
||||
with freeze_rng_state():
|
||||
script_out = sm(*args)
|
||||
|
||||
self.assertEqual(eager_out, script_out)
|
||||
self.assertExportImportModule(sm, args)
|
||||
|
||||
return sm
|
||||
|
||||
class NoTracerWarnContextManager:
|
||||
def __enter__(self):
|
||||
self.prev = torch._C._jit_get_tracer_state_warn()
|
||||
torch._C._jit_set_tracer_state_warn(False)
|
||||
|
||||
def __exit__(self, *args):
|
||||
torch._C._jit_set_tracer_state_warn(self.prev)
|
||||
|
||||
@contextmanager
|
||||
def inline_everything_mode(should_inline):
|
||||
old = torch._C._jit_get_inline_everything_mode()
|
||||
torch._C._jit_set_inline_everything_mode(should_inline)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch._C._jit_set_inline_everything_mode(old)
|
||||
|
||||
@contextmanager
|
||||
def set_fusion_group_inlining(inlining):
|
||||
old = torch._C._debug_get_fusion_group_inlining()
|
||||
torch._C._debug_set_fusion_group_inlining(inlining)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch._C._debug_set_fusion_group_inlining(old)
|
||||
|
||||
# note: not re-entrant, use unnested only
|
||||
@contextmanager
|
||||
def disable_autodiff_subgraph_inlining(enabled=True):
|
||||
torch._C._debug_set_autodiff_subgraph_inlining(not enabled)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch._C._debug_set_autodiff_subgraph_inlining(True)
|
||||
|
||||
def _inline_everything(fn):
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
with inline_everything_mode(True):
|
||||
fn(*args, **kwargs)
|
||||
return wrapper
|
||||
|
||||
# this exists for forward compatibility reasons temporarily.
|
||||
# TODO(suo) remove
|
||||
def _tmp_donotuse_dont_inline_everything(fn):
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
with inline_everything_mode(False):
|
||||
fn(*args, **kwargs)
|
||||
return wrapper
|
||||
|
||||
# make it easy to quickly define/trace a function for these tests
|
||||
def _trace(*args, **kwargs):
|
||||
def wrapper(func):
|
||||
return torch.jit.trace(func, args, **kwargs)
|
||||
return wrapper
|
||||
|
||||
|
||||
def enable_cpu_fuser(fn):
|
||||
def wrapper(*args, **kwargs):
|
||||
torch._C._jit_override_can_fuse_on_cpu_legacy(True)
|
||||
torch._C._jit_override_can_fuse_on_cpu(True)
|
||||
torch._C._jit_set_te_must_use_llvm_cpu(False)
|
||||
try:
|
||||
fn(*args, **kwargs)
|
||||
finally:
|
||||
torch._C._jit_override_can_fuse_on_cpu_legacy(False)
|
||||
torch._C._jit_override_can_fuse_on_cpu(False)
|
||||
torch._C._jit_set_te_must_use_llvm_cpu(True)
|
||||
return wrapper
|
||||
|
||||
|
||||
def enable_cpu_fuser_if(cond):
|
||||
if cond:
|
||||
return enable_cpu_fuser
|
||||
else:
|
||||
def noop_fuser(fn):
|
||||
def wrapper(*args, **kwargs):
|
||||
return fn(*args, **kwargs)
|
||||
return wrapper
|
||||
return noop_fuser
|
||||
|
||||
def get_forward(c):
|
||||
return c._get_method('forward')
|
||||
|
||||
def get_forward_graph(c):
|
||||
return c._get_method('forward').graph
|
||||
|
||||
def get_module_method(m, module, method):
|
||||
return m._c.getattr(module)._get_method(method)
|
||||
|
||||
def attrs_with_prefix(module, prefix):
|
||||
return [x for x, _ in module._modules._c.items()
|
||||
if x.startswith(prefix)]
|
||||
|
||||
def warmup_backward(f, *args):
|
||||
profiling_count = 3
|
||||
results = []
|
||||
for _ in range(profiling_count):
|
||||
if len(args) > 0:
|
||||
r = torch.autograd.grad(f, *args)
|
||||
results.append(r)
|
||||
else:
|
||||
f.backward(retain_graph=True)
|
||||
|
||||
return results
|
||||
|
||||
# TODO: Remove me once https://bugs.python.org/issue42666 is resolved
|
||||
def make_global(*args):
|
||||
for arg in args:
|
||||
setattr(sys.modules[arg.__module__], arg.__name__, arg)
|
||||
|
||||
# Helper function to eval Python3 code without causing a syntax error for
|
||||
# this file under py2
|
||||
def _get_py3_code(code, fn_name):
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
script_path = os.path.join(tmp_dir, 'script.py')
|
||||
with open(script_path, 'w') as f:
|
||||
f.write(code)
|
||||
spec = importlib.util.spec_from_file_location(fn_name, script_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
loader = spec.loader
|
||||
if not isinstance(loader, Loader): # Assert type to meet MyPy requirement
|
||||
raise AssertionError(f"Expected loader to be Loader, got {type(loader)}")
|
||||
loader.exec_module(module)
|
||||
fn = getattr(module, fn_name)
|
||||
return fn
|
||||
|
||||
class TensorExprTestOptions:
|
||||
def __init__(self) -> None:
|
||||
self.old_profiling_executor = torch._C._jit_set_profiling_executor(True)
|
||||
self.old_profiling_mode = torch._C._get_graph_executor_optimize(True)
|
||||
|
||||
self.old_cpu_fuser_state = torch._C._jit_can_fuse_on_cpu()
|
||||
self.old_gpu_fuser_state = torch._C._jit_can_fuse_on_gpu()
|
||||
torch._C._jit_override_can_fuse_on_cpu(True)
|
||||
torch._C._jit_override_can_fuse_on_gpu(True)
|
||||
self.texpr_fuser_state = torch._C._jit_texpr_fuser_enabled()
|
||||
torch._C._jit_set_texpr_fuser_enabled(True)
|
||||
self.old_fusion_inlining = torch._C._debug_get_fusion_group_inlining()
|
||||
torch._C._debug_set_fusion_group_inlining(False)
|
||||
self.old_te_must_use_llvm_cpu = torch._C._jit_get_te_must_use_llvm_cpu()
|
||||
torch._C._jit_set_te_must_use_llvm_cpu(False)
|
||||
|
||||
def restore(self):
|
||||
torch._C._jit_set_profiling_executor(self.old_profiling_executor)
|
||||
torch._C._get_graph_executor_optimize(self.old_profiling_mode)
|
||||
|
||||
torch._C._jit_set_texpr_fuser_enabled(self.texpr_fuser_state)
|
||||
torch._C._jit_override_can_fuse_on_gpu(self.old_gpu_fuser_state)
|
||||
torch._C._jit_override_can_fuse_on_cpu(self.old_cpu_fuser_state)
|
||||
torch._C._debug_set_fusion_group_inlining(self.old_fusion_inlining)
|
||||
torch._C._jit_set_te_must_use_llvm_cpu(self.old_te_must_use_llvm_cpu)
|
||||
|
||||
def clone_inputs(args):
|
||||
inputs: list[torch.Tensor | list[torch.Tensor]] = []
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor):
|
||||
inputs.append(arg.detach().clone())
|
||||
elif is_iterable_of_tensors(arg):
|
||||
inputs.append([t.detach().clone() for t in arg])
|
||||
else:
|
||||
inputs.append(arg)
|
||||
|
||||
return inputs
|
||||
|
||||
def get_traced_sample_variant_pairs(device, dtype, op):
|
||||
# tuples of (variant, sample)
|
||||
outputs: list[tuple[Any, Any]] = []
|
||||
|
||||
samples = op.sample_inputs(device, dtype)
|
||||
|
||||
# Acquires variants to test
|
||||
func = op.get_op()
|
||||
method = op.get_method()
|
||||
variants = {
|
||||
# TODO: inplace tests currently fail, fix and add inplace variant
|
||||
'function': func, 'method': method,
|
||||
}
|
||||
|
||||
# TODO: find better way to standardize on op registration itself..
|
||||
has_fake_function = op.name in ["resize_", 'resize_as_']
|
||||
|
||||
if has_fake_function:
|
||||
variants = {'method': getattr(torch.Tensor, op.name)}
|
||||
|
||||
# In eager mode, these ops can take (Tensor, bool) args; but in
|
||||
# JIT they can only take (Tensor, Scalar), and bool is not a
|
||||
# scalar in the JIT type system. So to test these in JIT, the bool
|
||||
# is converted to an int for the test.
|
||||
ops_with_unsupported_bool_args = [
|
||||
{
|
||||
"name": "div_floor_rounding",
|
||||
"arg_idx": [0],
|
||||
},
|
||||
{
|
||||
"name": "div_no_rounding_mode",
|
||||
"arg_idx": [0],
|
||||
},
|
||||
{
|
||||
"name": "div_trunc_rounding",
|
||||
"arg_idx": [0],
|
||||
},
|
||||
{
|
||||
"name": "index_fill",
|
||||
"arg_idx": [2],
|
||||
},
|
||||
{
|
||||
"name": "full_like",
|
||||
"arg_idx": [0],
|
||||
},
|
||||
{
|
||||
"name": "mul",
|
||||
"arg_idx": [0],
|
||||
},
|
||||
{
|
||||
"name": "new_full",
|
||||
"arg_idx": [1],
|
||||
},
|
||||
]
|
||||
|
||||
# doesn't support tracing
|
||||
if has_fake_function:
|
||||
return outputs
|
||||
|
||||
for sample in samples:
|
||||
for variant in variants.values():
|
||||
if variant is None:
|
||||
continue
|
||||
|
||||
if is_lambda(variant):
|
||||
continue
|
||||
|
||||
matching_ops = filter(lambda x: op.formatted_name == x["name"], ops_with_unsupported_bool_args)
|
||||
for op_data in matching_ops:
|
||||
for idx in op_data["arg_idx"]:
|
||||
args = list(sample.args)
|
||||
if len(sample.args) > idx and isinstance(sample.args[idx], bool):
|
||||
args[idx] = int(args[idx])
|
||||
sample.args = tuple(args)
|
||||
|
||||
outputs.append((variant, sample))
|
||||
|
||||
return outputs
|
||||
|
||||
# types.LambdaType gave false positives
|
||||
def is_lambda(lamb):
|
||||
LAMBDA = lambda: 0 # noqa: E731
|
||||
return isinstance(lamb, type(LAMBDA)) and lamb.__name__ == LAMBDA.__name__
|
||||
@@ -0,0 +1,167 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch
|
||||
from torch.utils._pytree import tree_map
|
||||
from collections.abc import Iterator
|
||||
import logging
|
||||
import contextlib
|
||||
import itertools
|
||||
from torch.utils._dtype_abbrs import dtype_abbrs as _dtype_abbrs
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
from torch.utils.weak import WeakTensorKeyDictionary
|
||||
import functools
|
||||
from torch._C._profiler import gather_traceback, symbolize_tracebacks
|
||||
|
||||
logger = logging.getLogger("LoggingTensor")
|
||||
|
||||
# How the chain of calls works for LoggingTensor:
|
||||
# 1. Call torch.sin
|
||||
# 2. Attempt __torch_function__. In LoggingTensor torch function is disabled so we bypass it entirely
|
||||
# 3. Enter dispatcher, wind your way through Autograd
|
||||
# 4. Hit Python dispatch key, call __torch_dispatch__
|
||||
|
||||
# This Tensor can work with autograd in two ways:
|
||||
# - The wrapped Tensor does not require gradients. In that case, the LoggingTensor
|
||||
# can require gradients if the user asks for it as a constructor kwarg.
|
||||
# - The wrapped Tensor can require gradients. In that case autograd will be tracked
|
||||
# for the wrapped Tensor and the LoggingTensor itself cannot require gradients.
|
||||
# WARNING: We allow these two possibilities for testing purposes. You should NEVER use both in a single
|
||||
# test or you might get surprising behavior.
|
||||
|
||||
# TODO: TensorBase should work
|
||||
class LoggingTensor(torch.Tensor):
|
||||
elem: torch.Tensor
|
||||
|
||||
__slots__ = ['elem']
|
||||
|
||||
context = contextlib.nullcontext
|
||||
|
||||
@staticmethod
|
||||
def __new__(cls, elem, *args, **kwargs):
|
||||
# The wrapping tensor (LoggingTensor) shouldn't hold any
|
||||
# memory for the class in question, but it should still
|
||||
# advertise the same device as before
|
||||
r = torch.Tensor._make_wrapper_subclass(
|
||||
cls, elem.size(),
|
||||
strides=elem.stride(), storage_offset=elem.storage_offset(),
|
||||
# TODO: clone storage aliasing
|
||||
dtype=elem.dtype, layout=elem.layout,
|
||||
device=elem.device, requires_grad=kwargs.get("requires_grad", False)
|
||||
)
|
||||
# ...the real tensor is held as an element on the tensor.
|
||||
r.elem = elem.detach() if r.requires_grad else elem
|
||||
return r
|
||||
|
||||
def __repr__(self):
|
||||
return super().__repr__(tensor_contents=f"{self.elem}")
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
|
||||
def unwrap(e):
|
||||
return e.elem if isinstance(e, cls) else e
|
||||
|
||||
def wrap(e):
|
||||
return cls(e) if isinstance(e, torch.Tensor) else e
|
||||
|
||||
with cls.context():
|
||||
rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs)))
|
||||
logging.getLogger("LoggingTensor").info(f"{func.__module__}.{func.__name__}", args, kwargs, rs) # noqa: G004
|
||||
return rs
|
||||
|
||||
class LoggingTensorMode(TorchDispatchMode):
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
rs = func(*args, **kwargs)
|
||||
logging.getLogger("LoggingTensor").info(f"{func.__module__}.{func.__name__}", args, kwargs, rs) # noqa: G004
|
||||
return rs
|
||||
|
||||
class LoggingTensorReentrant(LoggingTensor):
|
||||
context = torch.overrides.enable_reentrant_dispatch
|
||||
|
||||
# https://stackoverflow.com/questions/36408496/python-logging-handler-to-append-to-list
|
||||
class LoggingTensorHandler(logging.Handler):
|
||||
def __init__(
|
||||
self, log_list: list[str], use_shortid_for_all_tensors: bool,
|
||||
with_type: bool, tracebacks_list: list | None) -> None:
|
||||
logging.Handler.__init__(self)
|
||||
self.log_list = log_list
|
||||
self.use_shortid_for_all_tensors = use_shortid_for_all_tensors
|
||||
self.tracebacks_list = tracebacks_list
|
||||
self.memo = WeakTensorKeyDictionary()
|
||||
self.next_id = 0
|
||||
self.with_type = with_type
|
||||
|
||||
def _shortid(self, t: torch.Tensor) -> int:
|
||||
if t not in self.memo:
|
||||
self.memo[t] = self.next_id
|
||||
self.next_id += 1
|
||||
return self.memo[t]
|
||||
|
||||
def _fmt(self, a: object, with_type: bool = False) -> str:
|
||||
cond_cls = torch.Tensor if self.use_shortid_for_all_tensors else LoggingTensor
|
||||
if isinstance(a, cond_cls):
|
||||
maybe_type = ""
|
||||
if with_type and self.with_type:
|
||||
maybe_type = f": {_dtype_abbrs[a.dtype]}[{', '.join(map(str, a.shape))}]"
|
||||
x = f"${self._shortid(a)}{maybe_type}"
|
||||
return x
|
||||
else:
|
||||
return repr(a)
|
||||
|
||||
def emit(self, record):
|
||||
fmt_args = ", ".join(
|
||||
itertools.chain(
|
||||
(str(tree_map(self._fmt, a)) for a in record.args[0]),
|
||||
(f"{k}={str(tree_map(self._fmt, v))}" for k, v in record.args[1].items()),
|
||||
)
|
||||
)
|
||||
fmt_rets = tree_map(functools.partial(self._fmt, with_type=True), record.args[2])
|
||||
self.log_list.append(f'{fmt_rets} = {record.msg}({fmt_args})')
|
||||
if self.tracebacks_list is not None:
|
||||
self.tracebacks_list.append(record.traceback)
|
||||
|
||||
def log_input(name: str, var: object) -> None:
|
||||
logger.info("input", (name,), {}, var) # noqa: PLE1205
|
||||
|
||||
class GatherTraceback(logging.Filter):
|
||||
def __init__(self, python=True, script=True, cpp=False):
|
||||
self.python = python
|
||||
self.script = script
|
||||
self.cpp = cpp
|
||||
|
||||
def filter(self, record):
|
||||
record.traceback = gather_traceback(python=self.python, script=self.script, cpp=self.cpp)
|
||||
return True
|
||||
|
||||
@contextlib.contextmanager
|
||||
def capture_logs(is_mode=False, python_tb=False, script_tb=False, cpp_tb=False) -> Iterator[list[str]]:
|
||||
collect_traceback = python_tb or script_tb or cpp_tb
|
||||
log_list: list[str] = []
|
||||
tracebacks_list: list[str] = []
|
||||
handler = LoggingTensorHandler(
|
||||
log_list,
|
||||
with_type=True,
|
||||
use_shortid_for_all_tensors=is_mode,
|
||||
tracebacks_list=tracebacks_list if collect_traceback else None
|
||||
)
|
||||
logger.addHandler(handler)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.propagate = False
|
||||
if collect_traceback:
|
||||
logger.addFilter(GatherTraceback(python=python_tb, script=script_tb, cpp=cpp_tb))
|
||||
try:
|
||||
if collect_traceback:
|
||||
yield log_list, tracebacks_list
|
||||
else:
|
||||
yield log_list
|
||||
finally:
|
||||
symbolized_tracebacks = symbolize_tracebacks(tracebacks_list)
|
||||
tracebacks_list.clear()
|
||||
tracebacks_list.extend(symbolized_tracebacks)
|
||||
logger.removeHandler(handler)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def capture_logs_with_logging_tensor_mode(python_tb=False, script_tb=False, cpp_tb=False):
|
||||
with LoggingTensorMode(), capture_logs(True, python_tb, script_tb, cpp_tb) as logs:
|
||||
yield logs
|
||||
@@ -0,0 +1,243 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch._dynamo.test_case
|
||||
import unittest.mock
|
||||
import os
|
||||
import contextlib
|
||||
import torch._logging
|
||||
import torch._logging._internal
|
||||
from contextlib import AbstractContextManager
|
||||
from collections.abc import Callable
|
||||
from torch._dynamo.utils import LazyString
|
||||
from torch._inductor import config as inductor_config
|
||||
import logging
|
||||
import io
|
||||
|
||||
@contextlib.contextmanager
|
||||
def preserve_log_state():
|
||||
prev_state = torch._logging._internal._get_log_state()
|
||||
torch._logging._internal._set_log_state(torch._logging._internal.LogState())
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch._logging._internal._set_log_state(prev_state)
|
||||
torch._logging._internal._init_logs()
|
||||
|
||||
def log_settings(settings):
|
||||
exit_stack = contextlib.ExitStack()
|
||||
settings_patch = unittest.mock.patch.dict(os.environ, {"TORCH_LOGS": settings})
|
||||
exit_stack.enter_context(preserve_log_state())
|
||||
exit_stack.enter_context(settings_patch)
|
||||
torch._logging._internal._init_logs()
|
||||
return exit_stack
|
||||
|
||||
def log_api(**kwargs):
|
||||
exit_stack = contextlib.ExitStack()
|
||||
exit_stack.enter_context(preserve_log_state())
|
||||
torch._logging.set_logs(**kwargs)
|
||||
return exit_stack
|
||||
|
||||
|
||||
def kwargs_to_settings(**kwargs):
|
||||
INT_TO_VERBOSITY = {10: "+", 20: "", 40: "-"}
|
||||
|
||||
settings = []
|
||||
|
||||
def append_setting(name, level):
|
||||
if isinstance(name, str) and isinstance(level, int) and level in INT_TO_VERBOSITY:
|
||||
settings.append(INT_TO_VERBOSITY[level] + name)
|
||||
return
|
||||
else:
|
||||
raise ValueError("Invalid value for setting")
|
||||
|
||||
for name, val in kwargs.items():
|
||||
if isinstance(val, bool):
|
||||
settings.append(name)
|
||||
elif isinstance(val, int):
|
||||
append_setting(name, val)
|
||||
elif isinstance(val, dict) and name == "modules":
|
||||
for module_qname, level in val.items():
|
||||
append_setting(module_qname, level)
|
||||
else:
|
||||
raise ValueError("Invalid value for setting")
|
||||
|
||||
return ",".join(settings)
|
||||
|
||||
|
||||
# Note on testing strategy:
|
||||
# This class does two things:
|
||||
# 1. Runs two versions of a test:
|
||||
# 1a. patches the env var log settings to some specific value
|
||||
# 1b. calls torch._logging.set_logs(..)
|
||||
# 2. patches the emit method of each setup handler to gather records
|
||||
# that are emitted to each console stream
|
||||
# 3. passes a ref to the gathered records to each test case for checking
|
||||
#
|
||||
# The goal of this testing in general is to ensure that given some settings env var
|
||||
# that the logs are setup correctly and capturing the correct records.
|
||||
def make_logging_test(**kwargs):
|
||||
def wrapper(fn):
|
||||
@inductor_config.patch({"fx_graph_cache": False})
|
||||
def test_fn(self):
|
||||
|
||||
torch._dynamo.reset()
|
||||
records = []
|
||||
# run with env var
|
||||
if len(kwargs) == 0:
|
||||
with self._handler_watcher(records):
|
||||
fn(self, records)
|
||||
else:
|
||||
with log_settings(kwargs_to_settings(**kwargs)), self._handler_watcher(records):
|
||||
fn(self, records)
|
||||
|
||||
# run with API
|
||||
torch._dynamo.reset()
|
||||
records.clear()
|
||||
with log_api(**kwargs), self._handler_watcher(records):
|
||||
fn(self, records)
|
||||
|
||||
|
||||
return test_fn
|
||||
|
||||
return wrapper
|
||||
|
||||
def make_settings_test(settings):
|
||||
def wrapper(fn):
|
||||
def test_fn(self):
|
||||
torch._dynamo.reset()
|
||||
records = []
|
||||
# run with env var
|
||||
with log_settings(settings), self._handler_watcher(records):
|
||||
fn(self, records)
|
||||
|
||||
return test_fn
|
||||
|
||||
return wrapper
|
||||
|
||||
class LoggingTestCase(torch._dynamo.test_case.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls._exit_stack.enter_context(
|
||||
unittest.mock.patch.dict(os.environ, {"___LOG_TESTING": ""})
|
||||
)
|
||||
cls._exit_stack.enter_context(
|
||||
unittest.mock.patch("torch._dynamo.config.suppress_errors", True)
|
||||
)
|
||||
cls._exit_stack.enter_context(
|
||||
unittest.mock.patch("torch._dynamo.config.verbose", False)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls._exit_stack.close()
|
||||
torch._logging._internal.log_state.clear()
|
||||
torch._logging._init_logs()
|
||||
|
||||
def hasRecord(self, records, m):
|
||||
return any(m in r.getMessage() for r in records)
|
||||
|
||||
def getRecord(self, records, m):
|
||||
record = None
|
||||
for r in records:
|
||||
# NB: not r.msg because it looks like 3.11 changed how they
|
||||
# structure log records
|
||||
if m in r.getMessage():
|
||||
self.assertIsNone(
|
||||
record,
|
||||
msg=LazyString(
|
||||
lambda: f"multiple matching records: {record} and {r} among {records}"
|
||||
),
|
||||
)
|
||||
record = r
|
||||
if record is None:
|
||||
self.fail(f"did not find record with {m} among {records}")
|
||||
return record
|
||||
|
||||
# This patches the emit method of each handler to gather records
|
||||
# as they are emitted
|
||||
def _handler_watcher(self, record_list):
|
||||
exit_stack = contextlib.ExitStack()
|
||||
|
||||
def emit_post_hook(record):
|
||||
nonlocal record_list
|
||||
record_list.append(record)
|
||||
|
||||
# registered logs are the only ones with handlers, so patch those
|
||||
for log_qname in torch._logging._internal.log_registry.get_log_qnames():
|
||||
logger = logging.getLogger(log_qname)
|
||||
num_handlers = len(logger.handlers)
|
||||
self.assertLessEqual(
|
||||
num_handlers,
|
||||
2,
|
||||
"All pt2 loggers should only have at most two handlers (debug artifacts and messages above debug level).",
|
||||
)
|
||||
|
||||
self.assertGreater(num_handlers, 0, "All pt2 loggers should have more than zero handlers")
|
||||
|
||||
for handler in logger.handlers:
|
||||
old_emit = handler.emit
|
||||
|
||||
def new_emit(record):
|
||||
old_emit(record)
|
||||
emit_post_hook(record)
|
||||
|
||||
exit_stack.enter_context(
|
||||
unittest.mock.patch.object(handler, "emit", new_emit)
|
||||
)
|
||||
|
||||
return exit_stack
|
||||
|
||||
|
||||
def logs_to_string(module, log_option):
|
||||
"""Example:
|
||||
logs_to_string("torch._inductor.compile_fx", "post_grad_graphs")
|
||||
returns the output of TORCH_LOGS="post_grad_graphs" from the
|
||||
torch._inductor.compile_fx module.
|
||||
"""
|
||||
log_stream = io.StringIO()
|
||||
handler = logging.StreamHandler(stream=log_stream)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def tmp_redirect_logs():
|
||||
try:
|
||||
logger = torch._logging.getArtifactLogger(module, log_option)
|
||||
logger.addHandler(handler)
|
||||
yield
|
||||
finally:
|
||||
logger.removeHandler(handler)
|
||||
|
||||
def ctx_manager():
|
||||
exit_stack = log_settings(log_option)
|
||||
exit_stack.enter_context(tmp_redirect_logs())
|
||||
return exit_stack
|
||||
|
||||
return log_stream, ctx_manager
|
||||
|
||||
|
||||
def multiple_logs_to_string(module: str, *log_options: str) -> tuple[list[io.StringIO], Callable[[], AbstractContextManager[None]]]:
|
||||
"""Example:
|
||||
multiple_logs_to_string("torch._inductor.compile_fx", "pre_grad_graphs", "post_grad_graphs")
|
||||
returns the output of TORCH_LOGS="pre_graph_graphs, post_grad_graphs" from the
|
||||
torch._inductor.compile_fx module.
|
||||
"""
|
||||
log_streams = [io.StringIO() for _ in range(len(log_options))]
|
||||
handlers = [logging.StreamHandler(stream=log_stream) for log_stream in log_streams]
|
||||
|
||||
@contextlib.contextmanager
|
||||
def tmp_redirect_logs():
|
||||
loggers = [torch._logging.getArtifactLogger(module, option) for option in log_options]
|
||||
try:
|
||||
for logger, handler in zip(loggers, handlers, strict=True):
|
||||
logger.addHandler(handler)
|
||||
yield
|
||||
finally:
|
||||
for logger, handler in zip(loggers, handlers, strict=True):
|
||||
logger.removeHandler(handler)
|
||||
|
||||
def ctx_manager() -> AbstractContextManager[None]:
|
||||
exit_stack = log_settings(", ".join(log_options))
|
||||
exit_stack.enter_context(tmp_redirect_logs())
|
||||
return exit_stack # type: ignore[return-value]
|
||||
|
||||
return log_streams, ctx_manager
|
||||
@@ -0,0 +1,4 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch.testing._internal.opinfo.core
|
||||
import torch.testing._internal.opinfo.definitions
|
||||
File diff suppressed because it is too large
Load Diff
+26
@@ -0,0 +1,26 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from torch.testing._internal.opinfo.core import OpInfo
|
||||
from torch.testing._internal.opinfo.definitions import (
|
||||
_masked,
|
||||
fft,
|
||||
linalg,
|
||||
signal,
|
||||
special,
|
||||
)
|
||||
|
||||
|
||||
# Operator database
|
||||
op_db: list[OpInfo] = [
|
||||
*fft.op_db,
|
||||
*linalg.op_db,
|
||||
*signal.op_db,
|
||||
*special.op_db,
|
||||
*_masked.op_db,
|
||||
]
|
||||
|
||||
python_ref_db: list[OpInfo] = [
|
||||
*fft.python_ref_db,
|
||||
*linalg.python_ref_db,
|
||||
*special.python_ref_db,
|
||||
]
|
||||
+1386
File diff suppressed because it is too large
Load Diff
+939
@@ -0,0 +1,939 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import unittest
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch.testing import make_tensor
|
||||
from torch.testing._internal.common_device_type import precisionOverride
|
||||
from torch.testing._internal.common_dtype import (
|
||||
all_types_and,
|
||||
all_types_and_complex_and,
|
||||
)
|
||||
from torch.testing._internal.common_utils import TEST_SCIPY, TEST_WITH_ROCM
|
||||
from torch.testing._internal.opinfo.core import (
|
||||
DecorateInfo,
|
||||
ErrorInput,
|
||||
OpInfo,
|
||||
sample_inputs_spectral_ops,
|
||||
SampleInput,
|
||||
SpectralFuncInfo,
|
||||
SpectralFuncType,
|
||||
)
|
||||
from torch.testing._internal.opinfo.refs import (
|
||||
_find_referenced_opinfo,
|
||||
_inherit_constructor_args,
|
||||
PythonRefInfo,
|
||||
)
|
||||
|
||||
|
||||
has_scipy_fft = False
|
||||
if TEST_SCIPY:
|
||||
try:
|
||||
import scipy.fft
|
||||
|
||||
has_scipy_fft = True
|
||||
except ModuleNotFoundError:
|
||||
pass
|
||||
|
||||
|
||||
class SpectralFuncPythonRefInfo(SpectralFuncInfo):
|
||||
"""
|
||||
An OpInfo for a Python reference of an elementwise unary operation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name, # the stringname of the callable Python reference
|
||||
*,
|
||||
op=None, # the function variant of the operation, populated as torch.<name> if None
|
||||
torch_opinfo_name, # the string name of the corresponding torch opinfo
|
||||
torch_opinfo_variant="",
|
||||
**kwargs,
|
||||
): # additional kwargs override kwargs inherited from the torch opinfo
|
||||
self.torch_opinfo_name = torch_opinfo_name
|
||||
self.torch_opinfo = _find_referenced_opinfo(
|
||||
torch_opinfo_name, torch_opinfo_variant, op_db=op_db
|
||||
)
|
||||
if not isinstance(self.torch_opinfo, SpectralFuncInfo):
|
||||
raise AssertionError(
|
||||
f"Expected torch_opinfo to be SpectralFuncInfo, got {type(self.torch_opinfo)}"
|
||||
)
|
||||
|
||||
inherited = self.torch_opinfo._original_spectral_func_args
|
||||
ukwargs = _inherit_constructor_args(name, op, inherited, kwargs)
|
||||
|
||||
super().__init__(**ukwargs)
|
||||
|
||||
|
||||
def error_inputs_fft(op_info, device, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=torch.float32)
|
||||
# Zero-dimensional tensor has no dimension to take FFT of
|
||||
yield ErrorInput(
|
||||
SampleInput(make_arg()),
|
||||
error_type=IndexError,
|
||||
error_regex="Dimension specified as -1 but tensor has no dimensions",
|
||||
)
|
||||
|
||||
|
||||
def error_inputs_fftn(op_info, device, **kwargs):
|
||||
make_arg = partial(make_tensor, device=device, dtype=torch.float32)
|
||||
# Specifying a dimension on a zero-dimensional tensor
|
||||
yield ErrorInput(
|
||||
SampleInput(make_arg(), dim=(0,)),
|
||||
error_type=IndexError,
|
||||
error_regex="Dimension specified as 0 but tensor has no dimensions",
|
||||
)
|
||||
|
||||
|
||||
def sample_inputs_fft_with_min(
|
||||
op_info, device, dtype, requires_grad=False, *, min_size, **kwargs
|
||||
):
|
||||
yield from sample_inputs_spectral_ops(
|
||||
op_info, device, dtype, requires_grad, **kwargs
|
||||
)
|
||||
if TEST_WITH_ROCM:
|
||||
# FIXME: Causes floating point exception on ROCm
|
||||
return
|
||||
|
||||
# Check the "Invalid number of data points" error isn't too strict
|
||||
# https://github.com/pytorch/pytorch/pull/109083
|
||||
a = make_tensor(min_size, dtype=dtype, device=device, requires_grad=requires_grad)
|
||||
yield SampleInput(a)
|
||||
|
||||
|
||||
def sample_inputs_fftshift(op_info, device, dtype, requires_grad, **kwargs):
|
||||
def mt(shape, **kwargs):
|
||||
return make_tensor(
|
||||
shape, device=device, dtype=dtype, requires_grad=requires_grad, **kwargs
|
||||
)
|
||||
|
||||
yield SampleInput(mt((9, 10)))
|
||||
yield SampleInput(mt((50,)), kwargs=dict(dim=0))
|
||||
yield SampleInput(mt((5, 11)), kwargs=dict(dim=(1,)))
|
||||
yield SampleInput(mt((5, 6)), kwargs=dict(dim=(0, 1)))
|
||||
yield SampleInput(mt((5, 6, 2)), kwargs=dict(dim=(0, 2)))
|
||||
|
||||
|
||||
# Operator database
|
||||
op_db: list[OpInfo] = [
|
||||
SpectralFuncInfo(
|
||||
"fft.fft",
|
||||
aten_name="fft_fft",
|
||||
decomp_aten_name="_fft_c2c",
|
||||
ref=np.fft.fft,
|
||||
ndimensional=SpectralFuncType.OneD,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=1),
|
||||
error_inputs_func=error_inputs_fft,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.fft2",
|
||||
aten_name="fft_fft2",
|
||||
ref=np.fft.fft2,
|
||||
decomp_aten_name="_fft_c2c",
|
||||
ndimensional=SpectralFuncType.TwoD,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
decorators=[precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4})],
|
||||
skips=(
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestCommon",
|
||||
"test_complex_half_reference_testing",
|
||||
device_type="cuda",
|
||||
dtypes=[torch.complex32],
|
||||
active_if=TEST_WITH_ROCM,
|
||||
),
|
||||
# RuntimeError: [srcBuf length] > 0 INTERNAL ASSERT FAILED
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out",
|
||||
device_type="mps",
|
||||
),
|
||||
# AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 3, 10]) != torch.Size([5, 3, 11]).
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out_warning",
|
||||
device_type="mps",
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.fftn",
|
||||
aten_name="fft_fftn",
|
||||
decomp_aten_name="_fft_c2c",
|
||||
ref=np.fft.fftn,
|
||||
ndimensional=SpectralFuncType.ND,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(
|
||||
torch.bool,
|
||||
torch.half,
|
||||
torch.complex32,
|
||||
),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
decorators=[precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4})],
|
||||
skips=(
|
||||
# RuntimeError: [srcBuf length] > 0 INTERNAL ASSERT FAILED
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out",
|
||||
device_type="mps",
|
||||
),
|
||||
# AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 3, 10]) != torch.Size([5, 3, 11]).
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out_warning",
|
||||
device_type="mps",
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.hfft",
|
||||
aten_name="fft_hfft",
|
||||
decomp_aten_name="_fft_c2r",
|
||||
ref=np.fft.hfft,
|
||||
ndimensional=SpectralFuncType.OneD,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(
|
||||
torch.bool,
|
||||
torch.half,
|
||||
torch.complex32,
|
||||
),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=2),
|
||||
error_inputs_func=error_inputs_fft,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
skips=(
|
||||
# Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestSchemaCheckModeOpInfo",
|
||||
"test_schema_correctness",
|
||||
dtypes=(torch.complex64, torch.complex128),
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.hfft2",
|
||||
aten_name="fft_hfft2",
|
||||
decomp_aten_name="_fft_c2r",
|
||||
ref=scipy.fft.hfft2 if has_scipy_fft else None,
|
||||
ndimensional=SpectralFuncType.TwoD,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(
|
||||
torch.bool,
|
||||
torch.half,
|
||||
torch.complex32,
|
||||
),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(2, 2)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
check_batched_gradgrad=False,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 2e-4, torch.cfloat: 2e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
),
|
||||
],
|
||||
skips=(
|
||||
# Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestSchemaCheckModeOpInfo",
|
||||
"test_schema_correctness",
|
||||
),
|
||||
# FIXME: errors are too large; needs investigation
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestCommon",
|
||||
"test_complex_half_reference_testing",
|
||||
device_type="cuda",
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.hfftn",
|
||||
aten_name="fft_hfftn",
|
||||
decomp_aten_name="_fft_c2r",
|
||||
ref=scipy.fft.hfftn if has_scipy_fft else None,
|
||||
ndimensional=SpectralFuncType.ND,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(2, 2)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
check_batched_gradgrad=False,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 2e-4, torch.cfloat: 2e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
),
|
||||
],
|
||||
skips=(
|
||||
# Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestSchemaCheckModeOpInfo",
|
||||
"test_schema_correctness",
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.rfft",
|
||||
aten_name="fft_rfft",
|
||||
decomp_aten_name="_fft_r2c",
|
||||
ref=np.fft.rfft,
|
||||
ndimensional=SpectralFuncType.OneD,
|
||||
dtypes=all_types_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and(torch.bool, torch.half),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=1),
|
||||
error_inputs_func=error_inputs_fft,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.rfft2",
|
||||
aten_name="fft_rfft2",
|
||||
decomp_aten_name="_fft_r2c",
|
||||
ref=np.fft.rfft2,
|
||||
ndimensional=SpectralFuncType.TwoD,
|
||||
dtypes=all_types_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and(torch.bool, torch.half),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
decorators=[
|
||||
precisionOverride({torch.float: 1e-4}),
|
||||
],
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.rfftn",
|
||||
aten_name="fft_rfftn",
|
||||
decomp_aten_name="_fft_r2c",
|
||||
ref=np.fft.rfftn,
|
||||
ndimensional=SpectralFuncType.ND,
|
||||
dtypes=all_types_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and(torch.bool, torch.half),
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
decorators=[
|
||||
precisionOverride({torch.float: 1e-4}),
|
||||
],
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.ifft",
|
||||
aten_name="fft_ifft",
|
||||
decomp_aten_name="_fft_c2c",
|
||||
ref=np.fft.ifft,
|
||||
ndimensional=SpectralFuncType.OneD,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=1),
|
||||
error_inputs_func=error_inputs_fft,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32),
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out",
|
||||
device_type="mps",
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.ifft2",
|
||||
aten_name="fft_ifft2",
|
||||
decomp_aten_name="_fft_c2c",
|
||||
ref=np.fft.ifft2,
|
||||
ndimensional=SpectralFuncType.TwoD,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32),
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
skips=(
|
||||
# RuntimeError: [srcBuf length] > 0 INTERNAL ASSERT FAILED
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out",
|
||||
device_type="mps",
|
||||
),
|
||||
# AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 3, 10]) != torch.Size([5, 3, 11]).
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out_warning",
|
||||
device_type="mps",
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.ifftn",
|
||||
aten_name="fft_ifftn",
|
||||
decomp_aten_name="_fft_c2c",
|
||||
ref=np.fft.ifftn,
|
||||
ndimensional=SpectralFuncType.ND,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(
|
||||
torch.bool,
|
||||
torch.half,
|
||||
torch.complex32,
|
||||
),
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
skips=(
|
||||
# RuntimeError: [srcBuf length] > 0 INTERNAL ASSERT FAILED
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
# AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 3, 10]) != torch.Size([5, 3, 11]).
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out_warning",
|
||||
device_type="mps",
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.ihfft",
|
||||
aten_name="fft_ihfft",
|
||||
decomp_aten_name="_fft_r2c",
|
||||
ref=np.fft.ihfft,
|
||||
ndimensional=SpectralFuncType.OneD,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fft,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and(torch.bool, torch.half),
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
check_batched_grad=False,
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.ihfft2",
|
||||
aten_name="fft_ihfft2",
|
||||
decomp_aten_name="_fft_r2c",
|
||||
ref=scipy.fft.ihfftn if has_scipy_fft else None,
|
||||
ndimensional=SpectralFuncType.TwoD,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and(torch.bool, torch.half),
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
decorators=(
|
||||
# The values for attribute 'shape' do not match: torch.Size([5, 6, 5]) != torch.Size([5, 6, 6]).
|
||||
DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out_warning"),
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 2e-4}), "TestFFT", "test_reference_nd"
|
||||
),
|
||||
# Mismatched elements!
|
||||
DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out"),
|
||||
DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out_warnings"),
|
||||
),
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.ihfftn",
|
||||
aten_name="fft_ihfftn",
|
||||
decomp_aten_name="_fft_r2c",
|
||||
ref=scipy.fft.ihfftn if has_scipy_fft else None,
|
||||
ndimensional=SpectralFuncType.ND,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archss
|
||||
dtypesIfCUDA=all_types_and(torch.bool, torch.half),
|
||||
check_batched_grad=False,
|
||||
check_batched_gradgrad=False,
|
||||
decorators=[
|
||||
# The values for attribute 'shape' do not match: torch.Size([5, 6, 5]) != torch.Size([5, 6, 6]).
|
||||
DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out_warning"),
|
||||
# Mismatched elements!
|
||||
DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out"),
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 2e-4}), "TestFFT", "test_reference_nd"
|
||||
),
|
||||
],
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.irfft",
|
||||
aten_name="fft_irfft",
|
||||
decomp_aten_name="_fft_c2r",
|
||||
ref=np.fft.irfft,
|
||||
ndimensional=SpectralFuncType.OneD,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 2)),
|
||||
error_inputs_func=error_inputs_fft,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32),
|
||||
check_batched_gradgrad=False,
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.irfft2",
|
||||
aten_name="fft_irfft2",
|
||||
decomp_aten_name="_fft_c2r",
|
||||
ref=np.fft.irfft2,
|
||||
ndimensional=SpectralFuncType.TwoD,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 2)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32),
|
||||
check_batched_gradgrad=False,
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
),
|
||||
SpectralFuncInfo(
|
||||
"fft.irfftn",
|
||||
aten_name="fft_irfftn",
|
||||
decomp_aten_name="_fft_c2r",
|
||||
ref=np.fft.irfftn,
|
||||
ndimensional=SpectralFuncType.ND,
|
||||
sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 2)),
|
||||
error_inputs_func=error_inputs_fftn,
|
||||
# https://github.com/pytorch/pytorch/issues/80411
|
||||
gradcheck_fast_mode=True,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
# See https://github.com/pytorch/pytorch/pull/78358
|
||||
check_batched_forward_grad=False,
|
||||
dtypes=all_types_and_complex_and(torch.bool),
|
||||
# CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs
|
||||
dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32),
|
||||
check_batched_gradgrad=False,
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
),
|
||||
OpInfo(
|
||||
"fft.fftshift",
|
||||
dtypes=all_types_and_complex_and(
|
||||
torch.bool, torch.bfloat16, torch.half, torch.chalf
|
||||
),
|
||||
sample_inputs_func=sample_inputs_fftshift,
|
||||
supports_out=False,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
),
|
||||
OpInfo(
|
||||
"fft.ifftshift",
|
||||
dtypes=all_types_and_complex_and(
|
||||
torch.bool, torch.bfloat16, torch.half, torch.chalf
|
||||
),
|
||||
sample_inputs_func=sample_inputs_fftshift,
|
||||
supports_out=False,
|
||||
supports_forward_ad=True,
|
||||
supports_fwgrad_bwgrad=True,
|
||||
),
|
||||
]
|
||||
|
||||
python_ref_db: list[OpInfo] = [
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.fft",
|
||||
torch_opinfo_name="fft.fft",
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.ifft",
|
||||
torch_opinfo_name="fft.ifft",
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.rfft",
|
||||
torch_opinfo_name="fft.rfft",
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.irfft",
|
||||
torch_opinfo_name="fft.irfft",
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.hfft",
|
||||
torch_opinfo_name="fft.hfft",
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.ihfft",
|
||||
torch_opinfo_name="fft.ihfft",
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.fftn",
|
||||
torch_opinfo_name="fft.fftn",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.ifftn",
|
||||
torch_opinfo_name="fft.ifftn",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.rfftn",
|
||||
torch_opinfo_name="fft.rfftn",
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestCommon",
|
||||
"test_out",
|
||||
device_type="mps",
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.irfftn",
|
||||
torch_opinfo_name="fft.irfftn",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.hfftn",
|
||||
torch_opinfo_name="fft.hfftn",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 2e-4, torch.cfloat: 2e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.ihfftn",
|
||||
torch_opinfo_name="fft.ihfftn",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 2e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
),
|
||||
# AssertionError: Reference result was farther (0.09746177145360499) from the precise
|
||||
# computation than the torch result was (0.09111555632069855)
|
||||
# See https://github.com/pytorch/pytorch/pull/170856 for more details.
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestCommon",
|
||||
"test_python_ref_torch_fallback",
|
||||
dtypes=(torch.float16,),
|
||||
device_type="cuda",
|
||||
),
|
||||
# AssertionError: Reference result was farther (0.0953431016138116) from the precise
|
||||
# computation than the torch result was (0.09305490684430734)
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestCommon",
|
||||
"test_python_ref_executor",
|
||||
dtypes=(torch.float16,),
|
||||
device_type="cuda",
|
||||
),
|
||||
],
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.fft2",
|
||||
torch_opinfo_name="fft.fft2",
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.ifft2",
|
||||
torch_opinfo_name="fft.ifft2",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.rfft2",
|
||||
torch_opinfo_name="fft.rfft2",
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.irfft2",
|
||||
torch_opinfo_name="fft.irfft2",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.hfft2",
|
||||
torch_opinfo_name="fft.hfft2",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 2e-4, torch.cfloat: 2e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
)
|
||||
],
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
SpectralFuncPythonRefInfo(
|
||||
"_refs.fft.ihfft2",
|
||||
torch_opinfo_name="fft.ihfft2",
|
||||
decorators=[
|
||||
DecorateInfo(
|
||||
precisionOverride({torch.float: 2e-4}),
|
||||
"TestFFT",
|
||||
"test_reference_nd",
|
||||
),
|
||||
# FIXME:
|
||||
# Reference result was farther (0.0953431016138116) from the precise computation
|
||||
# than the torch result was (0.09305490684430734)!
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestCommon",
|
||||
"test_python_ref_executor",
|
||||
device_type="cuda",
|
||||
),
|
||||
],
|
||||
skips=(
|
||||
# AssertionError: Resizing an out= argument with no elements threw a resize warning!
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestCommon", "test_out", device_type="mps"
|
||||
),
|
||||
),
|
||||
),
|
||||
PythonRefInfo(
|
||||
"_refs.fft.fftshift",
|
||||
op_db=op_db,
|
||||
torch_opinfo_name="fft.fftshift",
|
||||
),
|
||||
PythonRefInfo(
|
||||
"_refs.fft.ifftshift",
|
||||
op_db=op_db,
|
||||
torch_opinfo_name="fft.ifftshift",
|
||||
),
|
||||
]
|
||||
+2853
File diff suppressed because it is too large
Load Diff
+1613
File diff suppressed because it is too large
Load Diff
+459
@@ -0,0 +1,459 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import unittest
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
from itertools import product
|
||||
|
||||
import numpy
|
||||
|
||||
import torch
|
||||
from torch.testing._internal.common_dtype import floating_types
|
||||
from torch.testing._internal.common_utils import TEST_SCIPY
|
||||
from torch.testing._internal.opinfo.core import (
|
||||
DecorateInfo,
|
||||
ErrorInput,
|
||||
OpInfo,
|
||||
SampleInput,
|
||||
)
|
||||
|
||||
|
||||
if TEST_SCIPY:
|
||||
import scipy.signal
|
||||
|
||||
|
||||
def sample_inputs_window(op_info, device, dtype, requires_grad, *args, **kwargs):
|
||||
r"""Base function used to create sample inputs for windows.
|
||||
|
||||
For additional required args you should use *args, as well as **kwargs for
|
||||
additional keyword arguments.
|
||||
"""
|
||||
|
||||
# Remove include_conjugated_inputs from kwargs
|
||||
kwargs.pop("include_conjugated_inputs", None)
|
||||
# Tests window sizes up to 5 samples.
|
||||
for size, sym in product(range(6), (True, False)):
|
||||
yield SampleInput(
|
||||
size,
|
||||
*args,
|
||||
sym=sym,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
requires_grad=requires_grad,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def reference_inputs_window(op_info, device, dtype, requires_grad, *args, **kwargs):
|
||||
r"""Reference inputs function to use for windows which have a common signature, i.e.,
|
||||
window size and sym only.
|
||||
|
||||
Implement other special functions for windows that have a specific signature.
|
||||
See exponential and gaussian windows for instance.
|
||||
"""
|
||||
yield from sample_inputs_window(
|
||||
op_info, device, dtype, requires_grad, *args, **kwargs
|
||||
)
|
||||
|
||||
cases = (8, 16, 32, 64, 128, 256)
|
||||
|
||||
for size in cases:
|
||||
yield SampleInput(size, sym=False)
|
||||
yield SampleInput(size, sym=True)
|
||||
|
||||
|
||||
def reference_inputs_exponential_window(
|
||||
op_info, device, dtype, requires_grad, **kwargs
|
||||
):
|
||||
yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs)
|
||||
|
||||
cases = (
|
||||
(8, {"center": 4, "tau": 0.5}),
|
||||
(16, {"center": 8, "tau": 2.5}),
|
||||
(32, {"center": 16, "tau": 43.5}),
|
||||
(64, {"center": 20, "tau": 3.7}),
|
||||
(128, {"center": 62, "tau": 99}),
|
||||
(256, {"tau": 10}),
|
||||
)
|
||||
|
||||
for size, kw in cases:
|
||||
yield SampleInput(size, sym=False, **kw)
|
||||
kw["center"] = None
|
||||
yield SampleInput(size, sym=True, **kw)
|
||||
|
||||
|
||||
def reference_inputs_gaussian_window(op_info, device, dtype, requires_grad, **kwargs):
|
||||
yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs)
|
||||
|
||||
cases = (
|
||||
(8, {"std": 0.1}),
|
||||
(16, {"std": 1.2}),
|
||||
(32, {"std": 2.1}),
|
||||
(64, {"std": 3.9}),
|
||||
(128, {"std": 4.5}),
|
||||
(256, {"std": 10}),
|
||||
)
|
||||
|
||||
for size, kw in cases:
|
||||
yield SampleInput(size, sym=False, **kw)
|
||||
yield SampleInput(size, sym=True, **kw)
|
||||
|
||||
|
||||
def reference_inputs_kaiser_window(op_info, device, dtype, requires_grad, **kwargs):
|
||||
yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs)
|
||||
|
||||
cases = (
|
||||
(8, {"beta": 2}),
|
||||
(16, {"beta": 12}),
|
||||
(32, {"beta": 30}),
|
||||
(64, {"beta": 35}),
|
||||
(128, {"beta": 41.2}),
|
||||
(256, {"beta": 100}),
|
||||
)
|
||||
|
||||
for size, kw in cases:
|
||||
yield SampleInput(size, sym=False, **kw)
|
||||
yield SampleInput(size, sym=True, **kw)
|
||||
|
||||
|
||||
def reference_inputs_general_cosine_window(
|
||||
op_info, device, dtype, requires_grad, **kwargs
|
||||
):
|
||||
yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs)
|
||||
|
||||
cases = (
|
||||
(8, {"a": [0.5, 0.5]}),
|
||||
(16, {"a": [0.46, 0.54]}),
|
||||
(32, {"a": [0.46, 0.23, 0.31]}),
|
||||
(64, {"a": [0.5]}),
|
||||
(128, {"a": [0.1, 0.8, 0.05, 0.05]}),
|
||||
(256, {"a": [0.2, 0.2, 0.2, 0.2, 0.2]}),
|
||||
)
|
||||
|
||||
for size, kw in cases:
|
||||
yield SampleInput(size, sym=False, **kw)
|
||||
yield SampleInput(size, sym=True, **kw)
|
||||
|
||||
|
||||
def reference_inputs_general_hamming_window(
|
||||
op_info, device, dtype, requires_grad, **kwargs
|
||||
):
|
||||
yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs)
|
||||
|
||||
cases = (
|
||||
(8, {"alpha": 0.54}),
|
||||
(16, {"alpha": 0.5}),
|
||||
(32, {"alpha": 0.23}),
|
||||
(64, {"alpha": 0.8}),
|
||||
(128, {"alpha": 0.9}),
|
||||
(256, {"alpha": 0.05}),
|
||||
)
|
||||
|
||||
for size, kw in cases:
|
||||
yield SampleInput(size, sym=False, **kw)
|
||||
yield SampleInput(size, sym=True, **kw)
|
||||
|
||||
|
||||
def error_inputs_window(op_info, device, *args, **kwargs):
|
||||
# Tests for windows that have a negative size
|
||||
yield ErrorInput(
|
||||
SampleInput(-1, *args, dtype=torch.float32, device=device, **kwargs),
|
||||
error_type=ValueError,
|
||||
error_regex="requires non-negative window length, got M=-1",
|
||||
)
|
||||
|
||||
# Tests for window tensors that are not torch.strided, for instance, torch.sparse_coo.
|
||||
yield ErrorInput(
|
||||
SampleInput(
|
||||
3,
|
||||
*args,
|
||||
layout=torch.sparse_coo,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
**kwargs,
|
||||
),
|
||||
error_type=ValueError,
|
||||
error_regex="is implemented for strided tensors only, got: torch.sparse_coo",
|
||||
)
|
||||
|
||||
# Tests for window tensors that are not floating point dtypes, for instance, torch.long.
|
||||
yield ErrorInput(
|
||||
SampleInput(3, *args, dtype=torch.long, device=device, **kwargs),
|
||||
error_type=ValueError,
|
||||
error_regex="expects float32 or float64 dtypes, got: torch.int64",
|
||||
)
|
||||
|
||||
# Tests for window tensors that are bfloat16
|
||||
yield ErrorInput(
|
||||
SampleInput(3, *args, dtype=torch.bfloat16, device=device, **kwargs),
|
||||
error_type=ValueError,
|
||||
error_regex="expects float32 or float64 dtypes, got: torch.bfloat16",
|
||||
)
|
||||
|
||||
# Tests for window tensors that are float16
|
||||
yield ErrorInput(
|
||||
SampleInput(3, *args, dtype=torch.float16, device=device, **kwargs),
|
||||
error_type=ValueError,
|
||||
error_regex="expects float32 or float64 dtypes, got: torch.float16",
|
||||
)
|
||||
|
||||
|
||||
def error_inputs_exponential_window(op_info, device, **kwargs):
|
||||
# Yield common error inputs
|
||||
yield from error_inputs_window(op_info, device, **kwargs)
|
||||
|
||||
# Tests for negative decay values.
|
||||
yield ErrorInput(
|
||||
SampleInput(3, tau=-1, dtype=torch.float32, device=device, **kwargs),
|
||||
error_type=ValueError,
|
||||
error_regex="Tau must be positive, got: -1 instead.",
|
||||
)
|
||||
|
||||
# Tests for symmetric windows and a given center value.
|
||||
yield ErrorInput(
|
||||
SampleInput(3, center=1, sym=True, dtype=torch.float32, device=device),
|
||||
error_type=ValueError,
|
||||
error_regex="Center must be None for symmetric windows",
|
||||
)
|
||||
|
||||
|
||||
def error_inputs_gaussian_window(op_info, device, **kwargs):
|
||||
# Yield common error inputs
|
||||
yield from error_inputs_window(op_info, device, std=0.5, **kwargs)
|
||||
|
||||
# Tests for negative standard deviations
|
||||
yield ErrorInput(
|
||||
SampleInput(3, std=-1, dtype=torch.float32, device=device, **kwargs),
|
||||
error_type=ValueError,
|
||||
error_regex="Standard deviation must be positive, got: -1 instead.",
|
||||
)
|
||||
|
||||
|
||||
def error_inputs_kaiser_window(op_info, device, **kwargs):
|
||||
# Yield common error inputs
|
||||
yield from error_inputs_window(op_info, device, beta=12, **kwargs)
|
||||
|
||||
# Tests for negative beta
|
||||
yield ErrorInput(
|
||||
SampleInput(3, beta=-1, dtype=torch.float32, device=device, **kwargs),
|
||||
error_type=ValueError,
|
||||
error_regex="beta must be non-negative, got: -1 instead.",
|
||||
)
|
||||
|
||||
|
||||
def error_inputs_general_cosine_window(op_info, device, **kwargs):
|
||||
# Yield common error inputs
|
||||
yield from error_inputs_window(op_info, device, a=[0.54, 0.46], **kwargs)
|
||||
|
||||
# Tests for negative beta
|
||||
yield ErrorInput(
|
||||
SampleInput(3, a=None, dtype=torch.float32, device=device, **kwargs),
|
||||
error_type=TypeError,
|
||||
error_regex="Coefficients must be a list/tuple",
|
||||
)
|
||||
|
||||
yield ErrorInput(
|
||||
SampleInput(3, a=[], dtype=torch.float32, device=device, **kwargs),
|
||||
error_type=ValueError,
|
||||
error_regex="Coefficients cannot be empty",
|
||||
)
|
||||
|
||||
|
||||
def reference_signal_window(fn: Callable):
|
||||
r"""Wrapper for scipy signal window references.
|
||||
|
||||
Discards keyword arguments for window reference functions that don't have a matching signature with
|
||||
torch, e.g., gaussian window.
|
||||
"""
|
||||
|
||||
def _fn(
|
||||
*args,
|
||||
dtype=numpy.float64,
|
||||
device=None,
|
||||
layout=torch.strided,
|
||||
requires_grad=False,
|
||||
**kwargs,
|
||||
):
|
||||
r"""The unused arguments are defined to disregard those values"""
|
||||
return fn(*args, **kwargs).astype(dtype)
|
||||
|
||||
return _fn
|
||||
|
||||
|
||||
def make_signal_windows_opinfo(
|
||||
name: str,
|
||||
ref: Callable,
|
||||
sample_inputs_func: Callable,
|
||||
reference_inputs_func: Callable,
|
||||
error_inputs_func: Callable,
|
||||
*,
|
||||
skips: tuple[DecorateInfo, ...] = (),
|
||||
):
|
||||
r"""Helper function to create OpInfo objects related to different windows."""
|
||||
return OpInfo(
|
||||
name=name,
|
||||
ref=ref if TEST_SCIPY else None,
|
||||
dtypes=floating_types(),
|
||||
sample_inputs_func=sample_inputs_func,
|
||||
reference_inputs_func=reference_inputs_func,
|
||||
error_inputs_func=error_inputs_func,
|
||||
supports_out=False,
|
||||
supports_autograd=False,
|
||||
skips=(
|
||||
# TODO: same as this?
|
||||
# https://github.com/pytorch/pytorch/issues/81774
|
||||
# also see: arange, new_full
|
||||
# fails to match any schemas despite working in the interpreter
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure,
|
||||
"TestOperatorSignatures",
|
||||
"test_get_torch_func_signature_exhaustive",
|
||||
),
|
||||
# fails to match any schemas despite working in the interpreter
|
||||
DecorateInfo(
|
||||
unittest.expectedFailure, "TestJit", "test_variant_consistency_jit"
|
||||
),
|
||||
# skip these tests since we have non tensor input
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"), "TestCommon", "test_noncontiguous_samples"
|
||||
),
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestCommon",
|
||||
"test_variant_consistency_eager",
|
||||
),
|
||||
DecorateInfo(unittest.skip("Skipped!"), "TestMathBits", "test_conj_view"),
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"), "TestMathBits", "test_neg_conj_view"
|
||||
),
|
||||
DecorateInfo(unittest.skip("Skipped!"), "TestMathBits", "test_neg_view"),
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestVmapOperatorsOpInfo",
|
||||
"test_vmap_exhaustive",
|
||||
),
|
||||
DecorateInfo(
|
||||
unittest.skip("Skipped!"),
|
||||
"TestVmapOperatorsOpInfo",
|
||||
"test_op_has_batch_rule",
|
||||
),
|
||||
DecorateInfo(
|
||||
unittest.skip("Buggy on MPS for now (mistakenly promotes to float64)"),
|
||||
"TestCommon",
|
||||
"test_numpy_ref_mps",
|
||||
),
|
||||
*skips,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
op_db: list[OpInfo] = [
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.hamming",
|
||||
ref=reference_signal_window(scipy.signal.windows.hamming)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=sample_inputs_window,
|
||||
reference_inputs_func=reference_inputs_window,
|
||||
error_inputs_func=error_inputs_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.hann",
|
||||
ref=reference_signal_window(scipy.signal.windows.hann) if TEST_SCIPY else None,
|
||||
sample_inputs_func=sample_inputs_window,
|
||||
reference_inputs_func=reference_inputs_window,
|
||||
error_inputs_func=error_inputs_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.bartlett",
|
||||
ref=reference_signal_window(scipy.signal.windows.bartlett)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=sample_inputs_window,
|
||||
reference_inputs_func=reference_inputs_window,
|
||||
error_inputs_func=error_inputs_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.blackman",
|
||||
ref=reference_signal_window(scipy.signal.windows.blackman)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=sample_inputs_window,
|
||||
reference_inputs_func=reference_inputs_window,
|
||||
error_inputs_func=error_inputs_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.cosine",
|
||||
ref=reference_signal_window(scipy.signal.windows.cosine)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=sample_inputs_window,
|
||||
reference_inputs_func=reference_inputs_window,
|
||||
error_inputs_func=error_inputs_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.exponential",
|
||||
ref=reference_signal_window(scipy.signal.windows.exponential)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=partial(sample_inputs_window, tau=2.78),
|
||||
reference_inputs_func=partial(reference_inputs_exponential_window, tau=2.78),
|
||||
error_inputs_func=error_inputs_exponential_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.gaussian",
|
||||
ref=reference_signal_window(scipy.signal.windows.gaussian)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=partial(sample_inputs_window, std=1.92),
|
||||
reference_inputs_func=partial(reference_inputs_gaussian_window, std=1.92),
|
||||
error_inputs_func=error_inputs_gaussian_window,
|
||||
skips=(
|
||||
DecorateInfo(
|
||||
unittest.skip("Buggy on MPS for now (mistakenly promotes to float64)"),
|
||||
"TestCommon",
|
||||
"test_numpy_ref_mps",
|
||||
),
|
||||
),
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.kaiser",
|
||||
ref=reference_signal_window(scipy.signal.windows.kaiser)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=partial(sample_inputs_window, beta=12.0),
|
||||
reference_inputs_func=partial(reference_inputs_kaiser_window, beta=12.0),
|
||||
error_inputs_func=error_inputs_kaiser_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.general_cosine",
|
||||
ref=reference_signal_window(scipy.signal.windows.general_cosine)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=partial(sample_inputs_window, a=[0.54, 0.46]),
|
||||
reference_inputs_func=partial(
|
||||
reference_inputs_general_cosine_window, a=[0.54, 0.46]
|
||||
),
|
||||
error_inputs_func=error_inputs_general_cosine_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.general_hamming",
|
||||
ref=reference_signal_window(scipy.signal.windows.general_hamming)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=partial(sample_inputs_window, alpha=0.54),
|
||||
reference_inputs_func=partial(
|
||||
reference_inputs_general_hamming_window, alpha=0.54
|
||||
),
|
||||
error_inputs_func=error_inputs_window,
|
||||
),
|
||||
make_signal_windows_opinfo(
|
||||
name="signal.windows.nuttall",
|
||||
ref=reference_signal_window(scipy.signal.windows.nuttall)
|
||||
if TEST_SCIPY
|
||||
else None,
|
||||
sample_inputs_func=sample_inputs_window,
|
||||
reference_inputs_func=reference_inputs_window,
|
||||
error_inputs_func=error_inputs_window,
|
||||
),
|
||||
]
|
||||
+931
@@ -0,0 +1,931 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torch.testing import make_tensor # noqa: F401
|
||||
from torch.testing._internal.common_dtype import highest_precision_float
|
||||
from torch.testing._internal.opinfo.core import ( # noqa: F401
|
||||
BinaryUfuncInfo,
|
||||
ErrorInput,
|
||||
generate_elementwise_binary_tensors,
|
||||
ReductionOpInfo,
|
||||
sample_inputs_reduction,
|
||||
SampleInput,
|
||||
)
|
||||
|
||||
|
||||
def _check_validate(op_info, sample):
|
||||
def _check_fail(sample):
|
||||
try:
|
||||
op_info(
|
||||
sample.sample_input.input,
|
||||
*sample.sample_input.args,
|
||||
**sample.sample_input.kwargs,
|
||||
)
|
||||
except sample.error_type:
|
||||
pass
|
||||
except Exception as msg:
|
||||
raise AssertionError( # noqa: B904
|
||||
f"{op_info.name} on {sample.sample_input=} expected exception "
|
||||
f"{sample.error_type}: {sample.error_regex}, got {type(msg).__name__}: {msg}"
|
||||
)
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"{op_info.name} on {sample.sample_input=} expected exception "
|
||||
f"{sample.error_type}: {sample.error_regex}, got none."
|
||||
)
|
||||
|
||||
def _check_success(sample):
|
||||
try:
|
||||
op_info(sample.input, *sample.args, **sample.kwargs)
|
||||
except Exception as msg:
|
||||
raise AssertionError( # noqa: B904
|
||||
f"{op_info.name} on {sample=} expected to succeed "
|
||||
f", got {type(msg).__name__}: {msg}"
|
||||
)
|
||||
|
||||
if isinstance(sample, ErrorInput):
|
||||
_check_fail(sample)
|
||||
else:
|
||||
_check_success(sample)
|
||||
|
||||
|
||||
def _sample_inputs_sparse(
|
||||
sample_inputs,
|
||||
maybe_failing_sample_inputs,
|
||||
validate_sample_input,
|
||||
op_info,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
check_validate = (
|
||||
os.environ.get("PYTORCH_TEST_CHECK_VALIDATE_SPARSE_SAMPLES", "0") == "1"
|
||||
)
|
||||
for sample in sample_inputs(op_info, *args, **kwargs):
|
||||
sample = validate_sample_input(op_info, sample, check_validate=check_validate)
|
||||
if isinstance(sample, SampleInput):
|
||||
yield sample
|
||||
# Error inputs are handled in error_inputs_sparse
|
||||
|
||||
for sample in maybe_failing_sample_inputs(op_info, *args, **kwargs):
|
||||
sample = validate_sample_input(op_info, sample, check_validate=check_validate)
|
||||
if isinstance(sample, SampleInput):
|
||||
yield sample
|
||||
|
||||
|
||||
def _error_inputs_sparse(
|
||||
maybe_failing_sample_inputs, validate_sample_input, op_info, *args, **kwargs
|
||||
):
|
||||
check_validate = (
|
||||
os.environ.get("PYTORCH_TEST_CHECK_VALIDATE_SPARSE_SAMPLES", "0") == "1"
|
||||
)
|
||||
for sample in maybe_failing_sample_inputs(op_info, *args, **kwargs):
|
||||
sample = validate_sample_input(op_info, sample, check_validate=check_validate)
|
||||
if isinstance(sample, ErrorInput):
|
||||
yield sample
|
||||
# Sample inputs are handled in sample_inputs_sparse
|
||||
|
||||
|
||||
def _apply_requires_grad_to_samples(sample_inputs):
|
||||
"""Decorator to _maybe_failing_sample_inputs_... generator functions
|
||||
that clones and sets requires_grad argument to tensors in sample
|
||||
input arguments. This is needed when the generated samples share
|
||||
tensor instances.
|
||||
"""
|
||||
|
||||
def wrapper(op_info, device, dtype, requires_grad, layout, **kwargs):
|
||||
def apply_requires_grad(x):
|
||||
if (
|
||||
not isinstance(x, torch.Tensor)
|
||||
or x.requires_grad
|
||||
or not requires_grad
|
||||
or not (x.is_floating_point() or x.is_complex())
|
||||
):
|
||||
return x
|
||||
return x.detach().clone().requires_grad_(requires_grad)
|
||||
|
||||
if requires_grad:
|
||||
for sample_input in sample_inputs(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
):
|
||||
yield sample_input.transform(apply_requires_grad)
|
||||
else:
|
||||
yield from sample_inputs(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def sample_inputs_sparse_reduction(
|
||||
op_info, device, dtype, requires_grad, layout, blocksize=None, **kwargs
|
||||
):
|
||||
"""Sample inputs for reduction operations on sparse tensors."""
|
||||
layout_name = str(layout).split(".", 1)[-1].rsplit("_coo", 1)[0]
|
||||
op_supports_layout = getattr(op_info, "supports_" + layout_name)
|
||||
if not op_supports_layout:
|
||||
return
|
||||
|
||||
for sample_input in sample_inputs_reduction(
|
||||
op_info, device, dtype, requires_grad, **kwargs
|
||||
):
|
||||
if sample_input.input.ndim == 0:
|
||||
# scalar sparse tensors are not supported
|
||||
continue
|
||||
|
||||
if layout in {
|
||||
torch.sparse_csr,
|
||||
torch.sparse_csc,
|
||||
torch.sparse_bsr,
|
||||
torch.sparse_bsc,
|
||||
}:
|
||||
if sample_input.input.ndim < 2:
|
||||
# conversion to sparse compressed tensors requires at
|
||||
# least 2 dimensional tensors
|
||||
continue
|
||||
if sample_input.input.ndim > 2 and (sample_input.input == 0).any():
|
||||
# Skip batched sparse compressed samples that contain
|
||||
# explicit zeros because to_sparse(layout=..) will
|
||||
# fail, see gh-98495.
|
||||
# TODO: remove this if-block after gh-98495 is fixed.
|
||||
continue
|
||||
|
||||
if layout in {torch.sparse_bsr, torch.sparse_bsc} and blocksize is None:
|
||||
blocksize = (1, 1)
|
||||
|
||||
yield SampleInput(
|
||||
sample_input.input.detach()
|
||||
.to_sparse(layout=layout, blocksize=blocksize)
|
||||
.requires_grad_(requires_grad),
|
||||
args=sample_input.args,
|
||||
kwargs=sample_input.kwargs,
|
||||
)
|
||||
|
||||
if layout is torch.sparse_coo and (dtype.is_floating_point or dtype.is_complex):
|
||||
# uncoalesced samples
|
||||
inp = sample_input.input.detach().to_sparse(layout=layout)
|
||||
inp = torch.sparse_coo_tensor(
|
||||
inp.indices().repeat(1, 2),
|
||||
inp.values().repeat(2),
|
||||
inp.shape,
|
||||
dtype=inp.dtype,
|
||||
device=inp.device,
|
||||
)
|
||||
if inp.is_coalesced():
|
||||
raise AssertionError("Expected inp to not be coalesced")
|
||||
yield SampleInput(
|
||||
inp.requires_grad_(requires_grad),
|
||||
args=sample_input.args,
|
||||
kwargs=sample_input.kwargs,
|
||||
)
|
||||
|
||||
if sample_input.input.ndim > 2:
|
||||
# hybrid samples
|
||||
yield SampleInput(
|
||||
sample_input.input.detach()
|
||||
.to_sparse(
|
||||
layout=layout,
|
||||
blocksize=blocksize,
|
||||
dense_dim=sample_input.input.ndim - 2,
|
||||
)
|
||||
.requires_grad_(requires_grad),
|
||||
args=sample_input.args,
|
||||
kwargs=sample_input.kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _validate_sample_input_sparse_reduction(op_info, sample, check_validate=False):
|
||||
"""Return the specified sample when it is valid and supported by the
|
||||
operation. Otherwise, return the sample as ErrorInput instance.
|
||||
|
||||
When check_validate is True, the result is validated against
|
||||
calling the op on the sample.
|
||||
"""
|
||||
UNSPECIFIED = object()
|
||||
if op_info.name == "sum":
|
||||
sample = _validate_sample_input_sparse_reduction_sum(sample)
|
||||
|
||||
if op_info.name == "masked.sum":
|
||||
mask = sample.kwargs.get("mask", UNSPECIFIED)
|
||||
if (
|
||||
mask not in {None, UNSPECIFIED}
|
||||
and mask.ndim > 2
|
||||
and mask.layout is torch.strided
|
||||
and (mask == 0).any()
|
||||
):
|
||||
# TODO: remove this if-block after gh-98495 is fixed.
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_regex="Expect the same number of specified elements per batch.",
|
||||
)
|
||||
elif not sample.kwargs.get("keepdim"):
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_type=(AssertionError, RuntimeError),
|
||||
error_regex="reduction operations on (CSR|CSC) tensors with keepdim=False is unsupported",
|
||||
)
|
||||
elif mask is UNSPECIFIED:
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_type=ValueError,
|
||||
error_regex="masked (.*) expects explicit mask for sparse_csr tensor input",
|
||||
)
|
||||
elif sample.input.ndim > 2:
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_regex="crow_indices is supposed to be a vector, but got 3 dimensional tensor.",
|
||||
)
|
||||
|
||||
if op_info.name in {"masked.amax", "masked.amin", "masked.mean", "masked.prod"}:
|
||||
t_inp = sample.input
|
||||
mask = sample.kwargs.get("mask")
|
||||
if (
|
||||
mask is not None
|
||||
and mask.ndim > 2
|
||||
and mask.layout is torch.strided
|
||||
and (mask == 0).any()
|
||||
):
|
||||
# TODO: remove this if-block after gh-98495 is fixed.
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_regex="Expect the same number of specified elements per batch.",
|
||||
)
|
||||
elif mask is None:
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_type=ValueError,
|
||||
error_regex="masked (.*) expects explicit mask for sparse_csr tensor input",
|
||||
)
|
||||
elif (
|
||||
mask.layout is sample.input.layout
|
||||
and mask.ndim > 2
|
||||
and op_info.name == "masked.mean"
|
||||
):
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_type=TypeError,
|
||||
error_regex=(
|
||||
"where[(][)] received an invalid combination of arguments"
|
||||
" - got [(]Tensor, Tensor, NoneType[)]"
|
||||
),
|
||||
)
|
||||
elif not sample.kwargs.get("keepdim"):
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_type=(AssertionError, RuntimeError),
|
||||
error_regex="reduction operations on (CSR|CSC) tensors with keepdim=False is unsupported",
|
||||
)
|
||||
elif (
|
||||
sample.input.ndim > 2
|
||||
and (sample.kwargs.get("dim") not in {0, 1})
|
||||
and mask.ndim > 2
|
||||
and mask.layout is not torch.strided
|
||||
):
|
||||
if sample.kwargs.get("dim") == (0, -1):
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_regex="tensor dimensionality must be sum of batch, base, and dense dimensionalities",
|
||||
)
|
||||
elif op_info.name == "masked.prod":
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_regex="input_dim == 2 INTERNAL ASSERT FAILED at",
|
||||
)
|
||||
else:
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_type=AssertionError,
|
||||
error_regex="Sparse CSR tensors are 2D and only support reduction along dim 0 or 1.",
|
||||
)
|
||||
elif sample.input.ndim > 2:
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_regex="crow_indices is supposed to be a vector, but got 3 dimensional tensor.",
|
||||
)
|
||||
elif (
|
||||
mask.layout is t_inp.layout
|
||||
and mask._nnz() != t_inp._nnz()
|
||||
and t_inp.dense_dim() > 0
|
||||
):
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_regex="Index tensor must have the same number of dimensions as src tensor",
|
||||
)
|
||||
|
||||
if check_validate:
|
||||
_check_validate(op_info, sample)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
def _validate_sample_input_sparse_reduction_sum(sample, check_validate=False):
|
||||
# NOTE: When fixing a failing sample case, remove the
|
||||
# corresponding if-block
|
||||
t_inp, t_kwargs = sample.input, sample.kwargs
|
||||
dim = t_kwargs.get("dim")
|
||||
keepdim = t_kwargs.get("keepdim")
|
||||
layout = t_inp.layout
|
||||
if isinstance(dim, (int, list, tuple)):
|
||||
if layout in {
|
||||
torch.sparse_csr,
|
||||
torch.sparse_csc,
|
||||
torch.sparse_bsr,
|
||||
torch.sparse_bsc,
|
||||
}:
|
||||
if layout in {torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc}:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=(
|
||||
"Currently the only compressed sparse format supported for sum.dim_IntList is CSR, but got layout"
|
||||
),
|
||||
)
|
||||
if layout in {torch.sparse_csr, torch.sparse_csc} and not keepdim:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=(
|
||||
"reduction operations on CSR tensors with keepdim=False is unsupported"
|
||||
),
|
||||
)
|
||||
if t_inp.dim() != 2:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=("input_dim == 2 INTERNAL ASSERT"),
|
||||
)
|
||||
if layout == torch.sparse_csr:
|
||||
if t_inp.dtype == torch.bool:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=("_sparse_csr_sum_cpu not implemented for 'Bool'"),
|
||||
)
|
||||
if t_inp.dtype == torch.complex32:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=(
|
||||
"_sparse_csr_sum_cuda not implemented for 'ComplexHalf'"
|
||||
),
|
||||
)
|
||||
return sample
|
||||
|
||||
|
||||
def _maybe_failing_sample_inputs_sparse_reduction_sum(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
):
|
||||
"""Generator of samples that are known to fail or that were failing in past."""
|
||||
# NOTE: When fixing a failing case, remove the Exception comment
|
||||
# but keep the `yield sample` statement.
|
||||
if layout in [
|
||||
torch.sparse_csr,
|
||||
torch.sparse_csc,
|
||||
]:
|
||||
# NotImplementedError: Could not run 'aten::sum.IntList_out' with arguments from the 'SparseCsrCPU' backend.
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype)
|
||||
.to_sparse(layout=layout)
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=0, keepdim=True),
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor([[[0, 1]], [[2, 3]]], dtype=dtype)
|
||||
.to_sparse(layout=layout, dense_dim=1)
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=0),
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype)
|
||||
.to_sparse(layout=layout)
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=(0,)),
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype)
|
||||
.to_sparse(layout=layout)
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=(0,), keepdim=True),
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor([[[0, 1]], [[2, 3]]], dtype=dtype)
|
||||
.to_sparse(layout=layout, dense_dim=1)
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=(0,)),
|
||||
)
|
||||
|
||||
# RuntimeError: torch.empty: Only batched sparse compressed (non-block) tensors are supported, but got size [2]
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype)
|
||||
.to_sparse(layout=layout)
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=0),
|
||||
)
|
||||
|
||||
if layout in [
|
||||
torch.sparse_bsr,
|
||||
torch.sparse_bsc,
|
||||
]:
|
||||
# RuntimeError: empty_sparse_compressed expected sparse compressed (non-block) tensor layout but got SparseBsr
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype)
|
||||
.to_sparse(layout=layout, blocksize=(2, 2))
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=0, keepdim=True),
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor([[[0, 1]], [[2, 3]]], dtype=dtype)
|
||||
.to_sparse(layout=layout, dense_dim=1, blocksize=(1, 1))
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=0),
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype)
|
||||
.to_sparse(layout=layout, blocksize=(1, 1))
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=(0,)),
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype)
|
||||
.to_sparse(layout=layout, blocksize=(1, 1))
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=(0,), keepdim=True),
|
||||
)
|
||||
yield SampleInput(
|
||||
torch.tensor([[[0, 1]], [[2, 3]]], dtype=dtype)
|
||||
.to_sparse(layout=layout, blocksize=(1, 1), dense_dim=1)
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=(0,)),
|
||||
)
|
||||
|
||||
# RuntimeError: torch.empty: Only batched sparse compressed (non-block) tensors are supported, but got size [2]
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype)
|
||||
.to_sparse(layout=layout, blocksize=(1, 1))
|
||||
.requires_grad_(requires_grad),
|
||||
kwargs=dict(dim=0),
|
||||
)
|
||||
|
||||
|
||||
def sample_inputs_sparse_reduction_sum(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
):
|
||||
"""Sample inputs for sum on sparse tensors."""
|
||||
yield from _sample_inputs_sparse(
|
||||
sample_inputs_sparse_reduction,
|
||||
_maybe_failing_sample_inputs_sparse_reduction_sum,
|
||||
_validate_sample_input_sparse_reduction,
|
||||
op_info,
|
||||
device,
|
||||
dtype,
|
||||
requires_grad,
|
||||
layout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def error_inputs_sparse_reduction_sum(op_info, device, layout, **kwargs):
|
||||
"""Error inputs for sum on sparse tensors."""
|
||||
dtype = torch.float64
|
||||
requires_grad = False
|
||||
yield from _error_inputs_sparse(
|
||||
_maybe_failing_sample_inputs_sparse_reduction_sum,
|
||||
_validate_sample_input_sparse_reduction,
|
||||
op_info,
|
||||
device,
|
||||
dtype,
|
||||
requires_grad,
|
||||
layout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def sample_inputs_sparse_elementwise_binary_operation(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
):
|
||||
"""Sample inputs for elementwise binary operations on sparse tensors.
|
||||
|
||||
The samples include regular, zero-sized, batched, and hybrid
|
||||
sparse tensors as well as rhs scalars. All tensors are full tensors.
|
||||
"""
|
||||
|
||||
def _to_sparse(tensor, **kwargs):
|
||||
return tensor.detach().to_sparse(**kwargs).requires_grad_(requires_grad)
|
||||
|
||||
for sample_input in generate_elementwise_binary_tensors(
|
||||
op_info,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
requires_grad=requires_grad,
|
||||
exclude_zero=True,
|
||||
**kwargs,
|
||||
):
|
||||
lhs, rhs = sample_input.input, sample_input.args[0]
|
||||
min_dense_dim = 0
|
||||
max_dense_dim = lhs.ndim - 1
|
||||
if layout in {
|
||||
torch.sparse_csr,
|
||||
torch.sparse_csc,
|
||||
torch.sparse_bsr,
|
||||
torch.sparse_bsc,
|
||||
}:
|
||||
if lhs.ndim < 2:
|
||||
# sparse compressed tensors sparse_dim must be 2
|
||||
continue
|
||||
max_dense_dim = lhs.ndim - 2
|
||||
|
||||
for dense_dim in range(min_dense_dim, max_dense_dim + 1):
|
||||
if layout in {torch.sparse_bsr, torch.sparse_bsc}:
|
||||
blocksizes = [(1, 1)]
|
||||
if lhs.numel() > 0:
|
||||
blocksizes.append(
|
||||
(
|
||||
lhs.shape[lhs.ndim - 2 - dense_dim],
|
||||
lhs.shape[lhs.ndim - 1 - dense_dim],
|
||||
)
|
||||
)
|
||||
else:
|
||||
blocksizes = [None]
|
||||
for blocksize in blocksizes:
|
||||
to_sparse_kwargs = dict(
|
||||
layout=layout, dense_dim=dense_dim, blocksize=blocksize
|
||||
)
|
||||
lhs_sparse = _to_sparse(lhs, **to_sparse_kwargs)
|
||||
rhs_sparse = _to_sparse(rhs, **to_sparse_kwargs)
|
||||
# op(sparse, sparse)
|
||||
yield SampleInput(
|
||||
lhs_sparse,
|
||||
args=(rhs_sparse, *sample_input.args[1:]),
|
||||
kwargs=sample_input.kwargs,
|
||||
)
|
||||
# op(sparse, scalar)
|
||||
yield SampleInput(
|
||||
lhs_sparse,
|
||||
args=(
|
||||
make_tensor(
|
||||
(), dtype=dtype, device=device, requires_grad=requires_grad
|
||||
),
|
||||
*sample_input.args[1:],
|
||||
),
|
||||
kwargs=sample_input.kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _validate_sample_input_elementwise_binary_sparse_mul(sample):
|
||||
# NOTE: When fixing a failing sample case, remove the
|
||||
# corresponding if-block
|
||||
t_inp, t_args = sample.input, sample.args
|
||||
batch_dim = t_inp.dim() - t_inp.dense_dim() - t_inp.sparse_dim()
|
||||
layout = t_inp.layout
|
||||
dtype = t_inp.dtype
|
||||
if layout is torch.sparse_csr and batch_dim > 0 and t_args[0].ndim > 0:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=(
|
||||
"coo_to_sparse_csr: conversion from Sparse to SparseCsr for input"
|
||||
" tensors with sparse_dim[(][)]!=2 is not supported"
|
||||
),
|
||||
)
|
||||
elif layout is torch.sparse_csc and t_args[0].ndim > 0:
|
||||
return ErrorInput(
|
||||
sample, error_regex="Expected result Tensor to be of format CSR"
|
||||
)
|
||||
elif layout is torch.sparse_bsr and t_args[0].ndim > 0:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex="empty_sparse_compressed expected sparse compressed [(]non-block[)] tensor layout but got SparseBsr",
|
||||
)
|
||||
elif layout is torch.sparse_bsc and t_args[0].ndim > 0:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex="empty_sparse_compressed expected sparse compressed [(]non-block[)] tensor layout but got SparseBsc",
|
||||
)
|
||||
elif (
|
||||
layout is torch.sparse_coo
|
||||
and dtype is torch.bool
|
||||
and t_args[0].ndim > 0
|
||||
and t_inp.is_cpu
|
||||
and t_inp.numel() > 0
|
||||
and t_inp.dense_dim() > 0
|
||||
):
|
||||
return ErrorInput(
|
||||
sample, error_regex="\"addcmul_cpu_out\" not implemented for 'Bool'"
|
||||
)
|
||||
elif (
|
||||
layout in {torch.sparse_coo, torch.sparse_csr}
|
||||
and dtype is torch.bool
|
||||
and t_inp._nnz() > 0
|
||||
and t_args[0].ndim > 0
|
||||
and t_inp.is_cpu
|
||||
and t_inp.numel() > 0
|
||||
):
|
||||
return ErrorInput(
|
||||
sample, error_regex="\"mul_out_sparse\" not implemented for 'Bool'"
|
||||
)
|
||||
elif (
|
||||
layout is torch.sparse_csr
|
||||
and t_args[0].layout is torch.strided
|
||||
and 0 < t_args[0].ndim
|
||||
and t_args[0].ndim < t_inp.ndim
|
||||
):
|
||||
return ErrorInput(
|
||||
sample, error_regex="sparse_mask_sparse_csr expects self to be 2D"
|
||||
)
|
||||
elif layout is torch.sparse_csr and (
|
||||
(t_args[0].layout is torch.strided and 0 < t_args[0].ndim)
|
||||
or (t_args[0].layout is layout and t_inp.shape != t_args[0].shape)
|
||||
):
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=(
|
||||
"expects sparse inputs with equal dimensionality, number of sparse dimensions,"
|
||||
" and shape of sparse dimensions"
|
||||
),
|
||||
)
|
||||
elif (
|
||||
layout is torch.sparse_csr
|
||||
and t_inp.dense_dim() > 0
|
||||
and t_inp._nnz() > 0
|
||||
and t_inp.is_cpu
|
||||
and dtype is torch.float16
|
||||
and t_args[0].ndim > 0
|
||||
):
|
||||
return ErrorInput(
|
||||
sample, error_regex="\"addcmul_cpu_out\" not implemented for 'Half'"
|
||||
)
|
||||
return sample
|
||||
|
||||
|
||||
@_apply_requires_grad_to_samples
|
||||
def _maybe_failing_sample_inputs_sparse_elementwise_binary_mul(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
):
|
||||
"""Generator of samples that are known to fail or that were failing in past."""
|
||||
# NOTE: When fixing a failing case, remove the Exception comment
|
||||
# but keep the `yield sample` statement.
|
||||
|
||||
blocksize = (1, 1) if layout in {torch.sparse_bsr, torch.sparse_bsc} else None
|
||||
regular = torch.tensor([[1, 2], [3, 4]], device=device, dtype=dtype).to_sparse(
|
||||
layout=layout, dense_dim=0, blocksize=blocksize
|
||||
)
|
||||
batch = torch.tensor(
|
||||
[[[1, 2], [3, 4]], [[4, 5], [6, 7]]], device=device, dtype=dtype
|
||||
).to_sparse(layout=layout, dense_dim=0, blocksize=blocksize)
|
||||
hybrid = torch.tensor(
|
||||
[[[1], [2]], [[3], [4]]], device=device, dtype=dtype
|
||||
).to_sparse(layout=layout, dense_dim=1, blocksize=blocksize)
|
||||
|
||||
if layout is torch.sparse_csr:
|
||||
# RuntimeError: crow_indices is supposed to be a vector, but got 2 dimensional tensor
|
||||
yield SampleInput(batch, args=(batch,))
|
||||
# RuntimeError: Only tensors with two sparse dimensions can be
|
||||
# converted to the SparseCsr layout, got self with 3 sparse
|
||||
# dimensions.
|
||||
yield SampleInput(
|
||||
torch.zeros_like(hybrid).requires_grad_(requires_grad),
|
||||
args=(torch.zeros_like(hybrid).requires_grad_(requires_grad),),
|
||||
)
|
||||
if dtype is torch.complex32:
|
||||
# RuntimeError: "mul_out_sparse" not implemented for 'ComplexHalf'
|
||||
yield SampleInput(regular, args=(regular,))
|
||||
if dtype is torch.bool and regular.is_cpu:
|
||||
# RuntimeError: "mul_out_sparse" not implemented for 'Bool'
|
||||
yield SampleInput(regular, args=(regular,))
|
||||
if layout is torch.sparse_csc:
|
||||
# RuntimeError: Expected result Tensor to be of format CSR
|
||||
yield SampleInput(regular, args=(regular,))
|
||||
if layout is torch.sparse_bsr:
|
||||
# RuntimeError: empty_sparse_compressed expected sparse compressed (non-block) tensor layout but got SparseBsr
|
||||
yield SampleInput(regular, args=(regular,))
|
||||
if layout is torch.sparse_bsc:
|
||||
# RuntimeError: empty_sparse_compressed expected sparse compressed (non-block) tensor layout but got SparseBsc
|
||||
yield SampleInput(regular, args=(regular,))
|
||||
if layout is torch.sparse_coo:
|
||||
if dtype is torch.complex32:
|
||||
# RuntimeError: "mul_out_sparse" not implemented for 'ComplexHalf'
|
||||
yield SampleInput(regular, args=(regular,))
|
||||
if dtype is torch.bool and regular.is_cpu:
|
||||
# RuntimeError: "mul_out_sparse" not implemented for 'Bool'
|
||||
yield SampleInput(regular, args=(regular,))
|
||||
if dtype in {torch.bool, torch.float16} and regular.is_cpu:
|
||||
# RuntimeError: "addcmul_cpu_out" not implemented for '(Bool|Half)'
|
||||
yield SampleInput(hybrid, args=(hybrid,))
|
||||
|
||||
|
||||
def _validate_sample_input_sparse_elementwise_binary_operation(
|
||||
op_info, sample, check_validate=False
|
||||
):
|
||||
if op_info.name == "mul":
|
||||
sample = _validate_sample_input_elementwise_binary_sparse_mul(sample)
|
||||
|
||||
if check_validate:
|
||||
_check_validate(op_info, sample)
|
||||
return sample
|
||||
|
||||
|
||||
def sample_inputs_sparse_mul(op_info, device, dtype, requires_grad, layout, **kwargs):
|
||||
"""Sample inputs for mul operation on sparse tensors."""
|
||||
yield from _sample_inputs_sparse(
|
||||
sample_inputs_sparse_elementwise_binary_operation,
|
||||
_maybe_failing_sample_inputs_sparse_elementwise_binary_mul,
|
||||
_validate_sample_input_sparse_elementwise_binary_operation,
|
||||
op_info,
|
||||
device,
|
||||
dtype,
|
||||
requires_grad,
|
||||
layout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def error_inputs_sparse_mul(op_info, device, layout, **kwargs):
|
||||
"""Error inputs for mul operation on sparse tensors."""
|
||||
dtype = torch.float64
|
||||
requires_grad = False
|
||||
yield from _error_inputs_sparse(
|
||||
_maybe_failing_sample_inputs_sparse_elementwise_binary_mul,
|
||||
_validate_sample_input_sparse_elementwise_binary_operation,
|
||||
op_info,
|
||||
device,
|
||||
dtype,
|
||||
requires_grad,
|
||||
layout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _sample_inputs_sparse_like_fns(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
):
|
||||
from torch.testing._internal.common_utils import TestCase
|
||||
|
||||
for tensor in TestCase().generate_simple_inputs(
|
||||
layout,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
enable_batch=True,
|
||||
enable_hybrid=True,
|
||||
enable_zero_sized=True,
|
||||
enable_non_contiguous_indices=False,
|
||||
enable_non_contiguous_values=False,
|
||||
):
|
||||
yield SampleInput(tensor, args=(), kwargs={})
|
||||
yield SampleInput(
|
||||
tensor, args=(), kwargs=dict(device=device, dtype=dtype, layout=layout)
|
||||
)
|
||||
|
||||
hpf = highest_precision_float(device)
|
||||
if dtype is not hpf:
|
||||
yield SampleInput(tensor, args=(), kwargs=dict(dtype=hpf))
|
||||
|
||||
if torch.cuda.is_available():
|
||||
other_device = "cuda" if tensor.device.type == "cpu" else "cpu"
|
||||
yield SampleInput(tensor, args=(), kwargs=dict(device=other_device))
|
||||
|
||||
if layout is torch.sparse_csr:
|
||||
other_layout = torch.sparse_csc
|
||||
elif layout is torch.sparse_csc:
|
||||
other_layout = torch.sparse_csr
|
||||
elif layout is torch.sparse_bsr:
|
||||
other_layout = torch.sparse_bsc
|
||||
elif layout is torch.sparse_bsc:
|
||||
other_layout = torch.sparse_bsr
|
||||
else:
|
||||
other_layout = torch.strided
|
||||
yield SampleInput(tensor, args=(), kwargs=dict(layout=other_layout))
|
||||
|
||||
if layout is not torch.sparse_coo:
|
||||
yield SampleInput(tensor, args=(), kwargs=dict(layout=torch.sparse_coo))
|
||||
|
||||
|
||||
def _validate_sample_input_sparse_like_fns(op_info, sample, check_validate=False):
|
||||
if (
|
||||
sample.input.layout
|
||||
in {
|
||||
torch.sparse_csr,
|
||||
torch.sparse_csc,
|
||||
torch.sparse_bsr,
|
||||
torch.sparse_bsc,
|
||||
}
|
||||
and op_info.name != "zeros_like"
|
||||
):
|
||||
if sample.kwargs.get("layout", sample.input.layout) != sample.input.layout:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=(
|
||||
"empty_like with different sparse layout is not supported"
|
||||
" \\(self is Sparse(Csc|Csr|Bsc|Bsr) but you requested Sparse(Csr|Csc|Bsr|Bsc)\\)"
|
||||
),
|
||||
)
|
||||
if sample.input.layout is torch.sparse_coo:
|
||||
return ErrorInput(
|
||||
sample,
|
||||
error_regex=(
|
||||
"Could not run 'aten::normal_' with arguments from the 'Sparse(CPU|CUDA)' backend."
|
||||
),
|
||||
)
|
||||
if check_validate:
|
||||
_check_validate(op_info, sample)
|
||||
return sample
|
||||
|
||||
|
||||
def _maybe_failing_sample_inputs_sparse_like_fns(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
):
|
||||
if torch.cuda.is_available() and layout is not torch.sparse_coo:
|
||||
other_device = "cuda" if torch.device(device).type == "cpu" else "cpu"
|
||||
if layout is torch.sparse_csr:
|
||||
other_layout = torch.sparse_csc
|
||||
elif layout is torch.sparse_csc:
|
||||
other_layout = torch.sparse_csr
|
||||
elif layout is torch.sparse_bsr:
|
||||
other_layout = torch.sparse_bsc
|
||||
elif layout is torch.sparse_bsc:
|
||||
other_layout = torch.sparse_bsr
|
||||
else:
|
||||
other_layout = torch.strided
|
||||
|
||||
blocksize = (1, 1) if layout in {torch.sparse_bsr, torch.sparse_bsc} else None
|
||||
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype, device=device).to_sparse(
|
||||
layout=layout, blocksize=blocksize
|
||||
),
|
||||
kwargs=dict(device=other_device),
|
||||
)
|
||||
|
||||
yield SampleInput(
|
||||
torch.tensor([[0, 1], [2, 3]], dtype=dtype, device=device).to_sparse(
|
||||
layout=layout, blocksize=blocksize
|
||||
),
|
||||
kwargs=dict(layout=other_layout),
|
||||
)
|
||||
|
||||
|
||||
def sample_inputs_sparse_like_fns(
|
||||
op_info, device, dtype, requires_grad, layout, **kwargs
|
||||
):
|
||||
"""Sample inputs for like-functions on sparse tensors."""
|
||||
yield from _sample_inputs_sparse(
|
||||
_sample_inputs_sparse_like_fns,
|
||||
_maybe_failing_sample_inputs_sparse_like_fns,
|
||||
_validate_sample_input_sparse_like_fns,
|
||||
op_info,
|
||||
device,
|
||||
dtype,
|
||||
requires_grad,
|
||||
layout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def error_inputs_sparse_like_fns(op_info, device, layout, **kwargs):
|
||||
"""Error inputs for like-functions on sparse tensors."""
|
||||
dtype = torch.float64
|
||||
requires_grad = False
|
||||
yield from _error_inputs_sparse(
|
||||
_maybe_failing_sample_inputs_sparse_like_fns,
|
||||
_validate_sample_input_sparse_like_fns,
|
||||
op_info,
|
||||
device,
|
||||
dtype,
|
||||
requires_grad,
|
||||
layout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _validate_sample_input_sparse_default(op_info, sample, check_validate=False):
|
||||
if op_info.name == "to_sparse":
|
||||
if (
|
||||
sample.input.layout
|
||||
in {torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc}
|
||||
and len(sample.args) == 1
|
||||
and isinstance(sample.args[0], int)
|
||||
and sample.args[0] != 2
|
||||
):
|
||||
sample = ErrorInput(
|
||||
sample,
|
||||
error_regex="sparse dim argument must be 2 for sparse_compressed_to_sparse",
|
||||
)
|
||||
|
||||
if check_validate:
|
||||
_check_validate(op_info, sample)
|
||||
return sample
|
||||
|
||||
|
||||
def validate_sample_input_sparse(op_info, sample, check_validate=False):
|
||||
"""Return the specified sample when it is valid and supported by the
|
||||
operation. Otherwise, return the sample as ErrorInput instance.
|
||||
|
||||
When check_validate is True, the result is validated against
|
||||
calling the op on the sample.
|
||||
"""
|
||||
if isinstance(op_info, ReductionOpInfo):
|
||||
return _validate_sample_input_sparse_reduction(
|
||||
op_info, sample, check_validate=check_validate
|
||||
)
|
||||
elif isinstance(op_info, BinaryUfuncInfo):
|
||||
return _validate_sample_input_sparse_elementwise_binary_operation(
|
||||
op_info, sample, check_validate=check_validate
|
||||
)
|
||||
else:
|
||||
return _validate_sample_input_sparse_default(
|
||||
op_info, sample, check_validate=check_validate
|
||||
)
|
||||
+1057
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,219 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from torch.testing._internal.opinfo.core import (
|
||||
BinaryUfuncInfo,
|
||||
OpInfo,
|
||||
ReductionOpInfo,
|
||||
UnaryUfuncInfo,
|
||||
)
|
||||
|
||||
|
||||
# NOTE [Python References]
|
||||
# Python References emulate existing PyTorch operations, but can ultimately
|
||||
# be expressed in terms of "primitive" operations from torch._prims.
|
||||
#
|
||||
# These references are experimental.
|
||||
# See https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-0/577
|
||||
# for additional context.
|
||||
#
|
||||
# Python Reference OpInfos should be added to the python_ref_db list below.
|
||||
# Tests can opt-into running on these references by including
|
||||
# that list in the Sequence they pass to the @ops decorator.
|
||||
#
|
||||
# When a Python Reference OpInfo is constructed a pointer to an
|
||||
# existing OpInfo must be provided using the torch_opinfo_name kwarg.
|
||||
# The existing OpInfo with that name and no variant will be found
|
||||
# to inherit from.
|
||||
#
|
||||
# Instead of just inheriting the existing OpInfo's metadata, the
|
||||
# Python Reference OpInfos inherit the existing OpInfo's
|
||||
# construction arguments. These arguments can be overridden
|
||||
# by adding kwargs to the constructor.
|
||||
|
||||
|
||||
def _find_referenced_opinfo(referenced_name, variant_name, *, op_db=None):
|
||||
"""
|
||||
Finds the OpInfo with the given name that has no variant name.
|
||||
"""
|
||||
# NOTE: searching the global op_db doesn't work when OpInfos are split into
|
||||
# different modules, as otherwise the op_db will not be fully constructed
|
||||
# yet. So, instead the local op_db must be passed in explicitly.
|
||||
if op_db is None:
|
||||
from torch.testing._internal.common_methods_invocations import op_db
|
||||
|
||||
for opinfo in op_db:
|
||||
if opinfo.name == referenced_name and opinfo.variant_test_name == variant_name:
|
||||
return opinfo
|
||||
|
||||
|
||||
def _inherit_constructor_args(name, op, inherited, overrides):
|
||||
# inherits metadata
|
||||
common_kwargs = {
|
||||
"name": name,
|
||||
"op": op,
|
||||
"aliases": None, # TODO add a check for alias coverage
|
||||
"method_variant": None,
|
||||
"inplace_variant": None, # TODO: add a check for inplace coverage
|
||||
"supports_scripting": False,
|
||||
}
|
||||
|
||||
# Acquires inherited kwargs
|
||||
kwargs = inherited.copy()
|
||||
|
||||
# Fixes metadata
|
||||
if "kwargs" in kwargs:
|
||||
kwargs.update(kwargs["kwargs"])
|
||||
del kwargs["kwargs"]
|
||||
if "self" in kwargs:
|
||||
del kwargs["self"]
|
||||
if "__class__" in kwargs:
|
||||
del kwargs["__class__"]
|
||||
if "skips" in kwargs:
|
||||
del kwargs["skips"]
|
||||
if "decorators" in kwargs:
|
||||
del kwargs["decorators"]
|
||||
|
||||
# Overrides metadata
|
||||
kwargs.update(common_kwargs)
|
||||
kwargs.update(overrides)
|
||||
|
||||
# At the moment no prims support autograd, so we must not run autograd
|
||||
# tests e.g. when testing dtype support. Once we start writing autograd
|
||||
# formulas for prims this can be removed.
|
||||
kwargs["supports_autograd"] = False
|
||||
kwargs["supports_gradgrad"] = False
|
||||
kwargs["supports_fwgrad_bwgrad"] = False
|
||||
kwargs["supports_inplace_autograd"] = False
|
||||
kwargs["supports_forward_ad"] = False
|
||||
|
||||
return kwargs
|
||||
|
||||
|
||||
class PythonRefInfo(OpInfo):
|
||||
"""
|
||||
An OpInfo for a Python reference of an OpInfo base class operation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name, # the stringname of the callable Python reference
|
||||
*,
|
||||
op=None, # the function variant of the operation, populated as torch.<name> if None
|
||||
op_db=None, # The database of opinfos to search for the parent opinfo
|
||||
torch_opinfo_name, # the string name of the corresponding torch opinfo
|
||||
torch_opinfo_variant_name="", # the variant name for corresponding torch opinfo
|
||||
validate_view_consistency=True,
|
||||
**kwargs,
|
||||
): # additional kwargs override kwargs inherited from the torch opinfo
|
||||
self.torch_opinfo_name = torch_opinfo_name
|
||||
self.torch_opinfo_variant_name = torch_opinfo_variant_name
|
||||
self.torch_opinfo = _find_referenced_opinfo(
|
||||
torch_opinfo_name, torch_opinfo_variant_name, op_db=op_db
|
||||
)
|
||||
self.validate_view_consistency = validate_view_consistency
|
||||
if not isinstance(self.torch_opinfo, OpInfo):
|
||||
raise AssertionError(
|
||||
f"Expected torch_opinfo to be OpInfo, got {type(self.torch_opinfo)}"
|
||||
)
|
||||
|
||||
inherited = self.torch_opinfo._original_opinfo_args
|
||||
ukwargs = _inherit_constructor_args(name, op, inherited, kwargs)
|
||||
super().__init__(**ukwargs)
|
||||
|
||||
|
||||
class ReductionPythonRefInfo(ReductionOpInfo):
|
||||
"""
|
||||
An OpInfo for a Python reference of an elementwise unary operation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name, # the stringname of the callable Python reference
|
||||
*,
|
||||
op=None, # the function variant of the operation, populated as torch.<name> if None
|
||||
op_db=None, # The database of opinfos to search for the parent opinfo
|
||||
torch_opinfo_name, # the string name of the corresponding torch opinfo
|
||||
torch_opinfo_variant_name="", # the variant name for corresponding torch opinfo
|
||||
**kwargs,
|
||||
): # additional kwargs override kwargs inherited from the torch opinfo
|
||||
self.torch_opinfo_name = torch_opinfo_name
|
||||
self.torch_opinfo_variant_name = torch_opinfo_variant_name
|
||||
self.torch_opinfo = _find_referenced_opinfo(
|
||||
torch_opinfo_name, torch_opinfo_variant_name, op_db=op_db
|
||||
)
|
||||
if not isinstance(self.torch_opinfo, ReductionOpInfo):
|
||||
raise AssertionError(
|
||||
f"Expected torch_opinfo to be ReductionOpInfo, got {type(self.torch_opinfo)}"
|
||||
)
|
||||
|
||||
inherited = self.torch_opinfo._original_reduction_args
|
||||
ukwargs = _inherit_constructor_args(name, op, inherited, kwargs)
|
||||
|
||||
# See https://github.com/pytorch/pytorch/issues/77216
|
||||
self.validate_view_consistency = False
|
||||
|
||||
super().__init__(**ukwargs)
|
||||
|
||||
|
||||
class ElementwiseUnaryPythonRefInfo(UnaryUfuncInfo):
|
||||
"""
|
||||
An OpInfo for a Python reference of an elementwise unary operation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name, # the stringname of the callable Python reference
|
||||
*,
|
||||
op=None, # the function variant of the operation, populated as torch.<name> if None
|
||||
op_db=None, # The database of opinfos to search for the parent opinfo
|
||||
torch_opinfo_name, # the string name of the corresponding torch opinfo
|
||||
torch_opinfo_variant_name="", # the variant name for corresponding torch opinfo
|
||||
validate_view_consistency=True,
|
||||
**kwargs,
|
||||
): # additional kwargs override kwargs inherited from the torch opinfo
|
||||
self.torch_opinfo_name = torch_opinfo_name
|
||||
self.torch_opinfo_variant_name = torch_opinfo_variant_name
|
||||
self.torch_opinfo = _find_referenced_opinfo(
|
||||
torch_opinfo_name, torch_opinfo_variant_name, op_db=op_db
|
||||
)
|
||||
self.validate_view_consistency = validate_view_consistency
|
||||
if not isinstance(self.torch_opinfo, UnaryUfuncInfo):
|
||||
raise AssertionError(
|
||||
f"Expected torch_opinfo to be UnaryUfuncInfo, got {type(self.torch_opinfo)}"
|
||||
)
|
||||
|
||||
inherited = self.torch_opinfo._original_unary_ufunc_args
|
||||
ukwargs = _inherit_constructor_args(name, op, inherited, kwargs)
|
||||
|
||||
super().__init__(**ukwargs)
|
||||
|
||||
|
||||
class ElementwiseBinaryPythonRefInfo(BinaryUfuncInfo):
|
||||
"""
|
||||
An OpInfo for a Python reference of an elementwise binary operation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name, # the stringname of the callable Python reference
|
||||
*,
|
||||
op=None, # the function variant of the operation, populated as torch.<name> if None
|
||||
op_db=None, # The database of opinfos to search for the parent opinfo
|
||||
torch_opinfo_name, # the string name of the corresponding torch opinfo
|
||||
torch_opinfo_variant_name="", # the variant name for corresponding torch opinfo
|
||||
**kwargs,
|
||||
): # additional kwargs override kwargs inherited from the torch opinfo
|
||||
self.torch_opinfo_name = torch_opinfo_name
|
||||
self.torch_opinfo_variant_name = torch_opinfo_variant_name
|
||||
self.torch_opinfo = _find_referenced_opinfo(
|
||||
torch_opinfo_name, torch_opinfo_variant_name, op_db=op_db
|
||||
)
|
||||
if not isinstance(self.torch_opinfo, BinaryUfuncInfo):
|
||||
raise AssertionError(
|
||||
f"Expected torch_opinfo to be BinaryUfuncInfo, got {type(self.torch_opinfo)}"
|
||||
)
|
||||
|
||||
inherited = self.torch_opinfo._original_binary_ufunc_args
|
||||
ukwargs = _inherit_constructor_args(name, op, inherited, kwargs)
|
||||
|
||||
super().__init__(**ukwargs)
|
||||
@@ -0,0 +1,282 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import collections
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from functools import partial, wraps
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
import torch
|
||||
from torch.testing._internal.common_cuda import TEST_CUDA
|
||||
from torch.testing._internal.common_dtype import (
|
||||
_dispatch_dtypes,
|
||||
all_types,
|
||||
all_types_and,
|
||||
all_types_and_complex,
|
||||
all_types_and_complex_and,
|
||||
all_types_and_half,
|
||||
complex_types,
|
||||
floating_and_complex_types,
|
||||
floating_and_complex_types_and,
|
||||
floating_types,
|
||||
floating_types_and,
|
||||
floating_types_and_half,
|
||||
integral_types,
|
||||
integral_types_and,
|
||||
)
|
||||
from torch.testing._internal.common_utils import torch_to_numpy_dtype_dict
|
||||
|
||||
|
||||
COMPLETE_DTYPES_DISPATCH = (
|
||||
all_types,
|
||||
all_types_and_complex,
|
||||
all_types_and_half,
|
||||
floating_types,
|
||||
floating_and_complex_types,
|
||||
floating_types_and_half,
|
||||
integral_types,
|
||||
complex_types,
|
||||
)
|
||||
|
||||
EXTENSIBLE_DTYPE_DISPATCH = (
|
||||
all_types_and_complex_and,
|
||||
floating_types_and,
|
||||
floating_and_complex_types_and,
|
||||
integral_types_and,
|
||||
all_types_and,
|
||||
)
|
||||
|
||||
# Better way to acquire devices?
|
||||
DEVICES = ["cpu"] + (["cuda"] if TEST_CUDA else [])
|
||||
|
||||
|
||||
class _dynamic_dispatch_dtypes(_dispatch_dtypes):
|
||||
# Class to tag the dynamically generated types.
|
||||
pass
|
||||
|
||||
|
||||
def get_supported_dtypes(op, sample_inputs_fn, device_type):
|
||||
# Returns the supported dtypes for the given operator and device_type pair.
|
||||
if device_type not in ["cpu", "cuda"]:
|
||||
raise AssertionError(
|
||||
f"Expected device_type in ['cpu', 'cuda'], got {device_type!r}"
|
||||
)
|
||||
if not TEST_CUDA and device_type == "cuda":
|
||||
warnings.warn(
|
||||
"WARNING: CUDA is not available, empty_dtypes dispatch will be returned!",
|
||||
stacklevel=2,
|
||||
)
|
||||
return _dynamic_dispatch_dtypes(())
|
||||
|
||||
supported_dtypes = set()
|
||||
for dtype in all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half):
|
||||
try:
|
||||
samples = sample_inputs_fn(op, device_type, dtype, False)
|
||||
except RuntimeError:
|
||||
# If `sample_inputs_fn` doesn't support sampling for a given
|
||||
# `dtype`, we assume that the `dtype` is not supported.
|
||||
# We raise a warning, so that user knows that this was the case
|
||||
# and can investigate if there was an issue with the `sample_inputs_fn`.
|
||||
warnings.warn(
|
||||
f"WARNING: Unable to generate sample for device:{device_type} and dtype:{dtype}",
|
||||
stacklevel=2,
|
||||
)
|
||||
continue
|
||||
|
||||
# We assume the dtype is supported
|
||||
# only if all samples pass for the given dtype.
|
||||
supported = True
|
||||
for sample in samples:
|
||||
try:
|
||||
op(sample.input, *sample.args, **sample.kwargs)
|
||||
except RuntimeError:
|
||||
# dtype is not supported
|
||||
supported = False
|
||||
break
|
||||
|
||||
if supported:
|
||||
supported_dtypes.add(dtype)
|
||||
|
||||
return _dynamic_dispatch_dtypes(supported_dtypes)
|
||||
|
||||
|
||||
def dtypes_dispatch_hint(dtypes):
|
||||
# Function returns the appropriate dispatch function (from COMPLETE_DTYPES_DISPATCH and EXTENSIBLE_DTYPE_DISPATCH)
|
||||
# and its string representation for the passed `dtypes`.
|
||||
return_type = collections.namedtuple("return_type", "dispatch_fn dispatch_fn_str")
|
||||
|
||||
# CUDA is not available, dtypes will be empty.
|
||||
if len(dtypes) == 0:
|
||||
return return_type((), "()")
|
||||
|
||||
set_dtypes = set(dtypes)
|
||||
for dispatch in COMPLETE_DTYPES_DISPATCH:
|
||||
# Short circuit if we get an exact match.
|
||||
if set(dispatch()) == set_dtypes:
|
||||
return return_type(dispatch, dispatch.__name__ + "()")
|
||||
|
||||
chosen_dispatch = None
|
||||
chosen_dispatch_score = 0.0
|
||||
for dispatch in EXTENSIBLE_DTYPE_DISPATCH:
|
||||
dispatch_dtypes = set(dispatch())
|
||||
if not dispatch_dtypes.issubset(set_dtypes):
|
||||
continue
|
||||
|
||||
score = len(dispatch_dtypes)
|
||||
if score > chosen_dispatch_score:
|
||||
chosen_dispatch_score = score
|
||||
chosen_dispatch = dispatch
|
||||
|
||||
# If user passed dtypes which are lower than the lowest
|
||||
# dispatch type available (not likely but possible in code path).
|
||||
if chosen_dispatch is None:
|
||||
return return_type((), str(dtypes))
|
||||
|
||||
return return_type(
|
||||
partial(dispatch, *tuple(set(dtypes) - set(dispatch()))),
|
||||
dispatch.__name__ + str(tuple(set(dtypes) - set(dispatch()))),
|
||||
)
|
||||
|
||||
|
||||
def is_dynamic_dtype_set(op):
|
||||
# Detect if the OpInfo entry acquired dtypes dynamically
|
||||
# using `get_supported_dtypes`.
|
||||
return op.dynamic_dtypes
|
||||
|
||||
|
||||
def str_format_dynamic_dtype(op):
|
||||
fmt_str = f"""
|
||||
OpInfo({op.name},
|
||||
dtypes={dtypes_dispatch_hint(op.dtypes).dispatch_fn_str},
|
||||
dtypesIfCUDA={dtypes_dispatch_hint(op.dtypesIfCUDA).dispatch_fn_str},
|
||||
)
|
||||
"""
|
||||
|
||||
return fmt_str
|
||||
|
||||
|
||||
def np_unary_ufunc_integer_promotion_wrapper(fn):
|
||||
# Wrapper that passes PyTorch's default scalar
|
||||
# type as an argument to the wrapped NumPy
|
||||
# unary ufunc when given an integer input.
|
||||
# This mimics PyTorch's integer->floating point
|
||||
# type promotion.
|
||||
#
|
||||
# This is necessary when NumPy promotes
|
||||
# integer types to double, since PyTorch promotes
|
||||
# integer types to the default scalar type.
|
||||
|
||||
# Helper to determine if promotion is needed
|
||||
def is_integral(dtype):
|
||||
return dtype in [
|
||||
np.bool_,
|
||||
bool,
|
||||
np.uint8,
|
||||
np.int8,
|
||||
np.int16,
|
||||
np.int32,
|
||||
np.int64,
|
||||
]
|
||||
|
||||
@wraps(fn)
|
||||
def wrapped_fn(x):
|
||||
# As the default dtype can change, acquire it when function is called.
|
||||
# NOTE: Promotion in PyTorch is from integer types to the default dtype
|
||||
np_dtype = torch_to_numpy_dtype_dict[torch.get_default_dtype()]
|
||||
|
||||
if is_integral(x.dtype):
|
||||
return fn(x.astype(np_dtype))
|
||||
return fn(x)
|
||||
|
||||
return wrapped_fn
|
||||
|
||||
|
||||
def reference_reduction_numpy(f, supports_keepdims=True):
|
||||
"""Wraps a NumPy reduction operator.
|
||||
|
||||
The wrapper function will forward dim, keepdim, mask, and identity
|
||||
kwargs to the wrapped function as the NumPy equivalent axis,
|
||||
keepdims, where, and initiak kwargs, respectively.
|
||||
|
||||
Args:
|
||||
f: NumPy reduction operator to wrap
|
||||
supports_keepdims (bool, optional): Whether the NumPy operator accepts
|
||||
keepdims parameter. If it does not, the wrapper will manually unsqueeze
|
||||
the reduced dimensions if it was called with keepdim=True. Defaults to True.
|
||||
|
||||
Returns:
|
||||
Wrapped function
|
||||
|
||||
"""
|
||||
|
||||
@wraps(f)
|
||||
def wrapper(x: npt.NDArray, *args, **kwargs):
|
||||
# Copy keys into a set
|
||||
keys = set(kwargs.keys())
|
||||
|
||||
dim = kwargs.pop("dim", None)
|
||||
keepdim = kwargs.pop("keepdim", False)
|
||||
|
||||
if "dim" in keys:
|
||||
dim = tuple(dim) if isinstance(dim, Sequence) else dim
|
||||
|
||||
# NumPy reductions don't accept dim=0 for scalar inputs
|
||||
# so we convert it to None if and only if dim is equivalent
|
||||
if x.ndim == 0 and dim in {0, -1, (0,), (-1,)}:
|
||||
kwargs["axis"] = None
|
||||
else:
|
||||
kwargs["axis"] = dim
|
||||
|
||||
if "keepdim" in keys and supports_keepdims:
|
||||
kwargs["keepdims"] = keepdim
|
||||
|
||||
if "mask" in keys:
|
||||
mask = kwargs.pop("mask")
|
||||
if mask is not None:
|
||||
if mask.layout != torch.strided:
|
||||
raise AssertionError(
|
||||
f"Expected mask.layout == torch.strided, got {mask.layout}"
|
||||
)
|
||||
kwargs["where"] = mask.cpu().numpy()
|
||||
|
||||
if "identity" in keys:
|
||||
identity = kwargs.pop("identity")
|
||||
if identity is not None:
|
||||
if identity.dtype is torch.bfloat16:
|
||||
identity = identity.cpu().to(torch.float32)
|
||||
else:
|
||||
identity = identity.cpu()
|
||||
kwargs["initial"] = identity.numpy()
|
||||
|
||||
result = f(x, *args, **kwargs)
|
||||
|
||||
# Unsqueeze reduced dimensions if NumPy does not support keepdims
|
||||
if keepdim and not supports_keepdims and x.ndim > 0:
|
||||
dim = list(range(x.ndim)) if dim is None else dim
|
||||
result = np.expand_dims(result, dim)
|
||||
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def prod_numpy(a, *args, **kwargs):
|
||||
"""
|
||||
The function will call np.prod with type as np.int64 if the input type
|
||||
is int or uint64 if is uint. This is necessary because windows np.prod uses by default
|
||||
int32 while on linux it uses int64.
|
||||
This is for fixing integer overflow https://github.com/pytorch/pytorch/issues/77320
|
||||
|
||||
Returns:
|
||||
np.prod of input
|
||||
"""
|
||||
if "dtype" not in kwargs:
|
||||
if np.issubdtype(a.dtype, np.signedinteger):
|
||||
a = a.astype(np.int64)
|
||||
elif np.issubdtype(a.dtype, np.unsignedinteger):
|
||||
a = a.astype(np.uint64)
|
||||
|
||||
fn = reference_reduction_numpy(np.prod)
|
||||
return fn(a, *args, **kwargs)
|
||||
@@ -0,0 +1,7 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
from .make_fx import make_fx_check
|
||||
from .aot_autograd import aot_autograd_check, _test_aot_autograd_forwards_backwards_helper
|
||||
from .fake_tensor import fake_check
|
||||
from .autograd_registration import autograd_registration_check
|
||||
from .generate_tests import generate_opcheck_tests, opcheck, OpCheckError, dontGenerateOpCheckTests, is_inside_opcheck_mode
|
||||
+176
@@ -0,0 +1,176 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch.testing._utils import wrapper_set_seed
|
||||
from functorch.compile import compiled_function, min_cut_rematerialization_partition, default_partition, nop
|
||||
from .make_fx import randomize
|
||||
import re
|
||||
|
||||
|
||||
class assert_raises_regex:
|
||||
def __init__(self, exception_cls, regex):
|
||||
self.exception_cls = exception_cls
|
||||
self.regex = regex
|
||||
|
||||
def __enter__(self):
|
||||
pass
|
||||
|
||||
def __exit__(self, exc_type, exc_val, traceback):
|
||||
if exc_type == self.exception_cls:
|
||||
msg = str(exc_val)
|
||||
if not re.search(self.regex, msg):
|
||||
raise AssertionError(
|
||||
f"Expected exception to match regex. regex: {self.regex}, exception: {msg}")
|
||||
return True # Squashes the exception
|
||||
if exc_type is not None:
|
||||
raise AssertionError(
|
||||
f"Expected {self.exception_cls} to be raised, instead got exception {exc_type}")
|
||||
raise AssertionError("Expected exception to be raised but none was")
|
||||
|
||||
|
||||
def aot_autograd_check(
|
||||
func,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic,
|
||||
assert_raises_regex_fn=assert_raises_regex,
|
||||
assert_equals_fn=torch.testing.assert_close,
|
||||
check_gradients=True,
|
||||
try_check_data_specialization=False,
|
||||
skip_correctness_check=False,
|
||||
disable_functionalization=False):
|
||||
"""Compares func(*args, **kwargs) in eager-mode to under AOTAutograd.
|
||||
|
||||
Compares outputs and (if check_gradients=True) gradients produced by
|
||||
AOTAutograd against eager-mode PyTorch.
|
||||
|
||||
We assume that func(*args, **kwargs) succeeds in eager-mode PyTorch.
|
||||
|
||||
"""
|
||||
flat_args, args_spec = pytree.tree_flatten((args, kwargs))
|
||||
args = [arg for arg in flat_args if isinstance(arg, torch.Tensor)]
|
||||
|
||||
# We construct a new function that only accepts Tensors as inputs
|
||||
def func_no_tensors(args):
|
||||
reconstructed_flat_args = []
|
||||
args = iter(args)
|
||||
for v in flat_args:
|
||||
if isinstance(v, torch.Tensor):
|
||||
reconstructed_flat_args.append(next(args))
|
||||
else:
|
||||
reconstructed_flat_args.append(v)
|
||||
|
||||
c_args, c_kwargs = pytree.tree_unflatten(reconstructed_flat_args, args_spec)
|
||||
return func(*c_args, **c_kwargs)
|
||||
|
||||
# cannot use the min cut partitioner without functionalization
|
||||
if disable_functionalization:
|
||||
compiled_f = compiled_function(
|
||||
func_no_tensors,
|
||||
nop,
|
||||
nop,
|
||||
dynamic=dynamic,
|
||||
partition_fn=default_partition,
|
||||
keep_inference_input_mutations=True,
|
||||
disable_functionalization=True
|
||||
)
|
||||
else:
|
||||
compiled_f = compiled_function(
|
||||
func_no_tensors,
|
||||
nop,
|
||||
nop,
|
||||
dynamic=dynamic,
|
||||
partition_fn=min_cut_rematerialization_partition,
|
||||
keep_inference_input_mutations=True,
|
||||
disable_functionalization=False
|
||||
)
|
||||
|
||||
out = wrapper_set_seed(func_no_tensors, args)
|
||||
if check_gradients == "auto":
|
||||
any_tensor_requires_grad = pytree.tree_any_only(torch.Tensor, lambda x: x.requires_grad, args)
|
||||
any_output_requires_grad = pytree.tree_any_only(torch.Tensor, lambda x: x.requires_grad, out)
|
||||
check_gradients = any_tensor_requires_grad and any_output_requires_grad
|
||||
if not check_gradients:
|
||||
compiled_out = wrapper_set_seed(compiled_f, args)
|
||||
if not skip_correctness_check:
|
||||
assert_equals_fn(compiled_out, out, msg=outputs_msg)
|
||||
return
|
||||
_test_aot_autograd_forwards_backwards_helper(
|
||||
func_no_tensors, compiled_f, args, assert_raises_regex_fn, assert_equals_fn,
|
||||
try_check_data_specialization, skip_correctness_check)
|
||||
|
||||
outputs_msg = (
|
||||
"Outputs of the operator are different in eager-mode PyTorch vs "
|
||||
"AOTDispatcher tracing. This means the operator will have incorrect output "
|
||||
"underneath torch.compile. This could be because the operator's "
|
||||
"implementation not traceable."
|
||||
)
|
||||
|
||||
|
||||
def _test_aot_autograd_forwards_backwards_helper(
|
||||
f, compiled_f, args, assert_raises_regex_fn, assert_equals_fn,
|
||||
try_check_data_specialization, skip_correctness_check=False):
|
||||
# Verify grads are equal between compiled and non-compiled versions of f.
|
||||
|
||||
def call_forwards_backwards(f, args):
|
||||
flat_args = pytree.arg_tree_leaves(*args)
|
||||
diff_args = [arg for arg in flat_args if isinstance(arg, torch.Tensor) and
|
||||
arg.requires_grad]
|
||||
out = wrapper_set_seed(f, args)
|
||||
flat_out = pytree.tree_leaves(out)
|
||||
|
||||
sm = 0
|
||||
for i in flat_out:
|
||||
if isinstance(i, torch.Tensor):
|
||||
# We need to call .abs() because it is possible that the output of the
|
||||
# operator is a complex Tensor and autograd will yell at autograd.grad
|
||||
# on a complex Tensor unless we manually provide the grad_output flag.
|
||||
sm += i.sum().abs()
|
||||
if not isinstance(sm, torch.Tensor):
|
||||
raise AssertionError(f"Expected sm to be a Tensor, got {type(sm)}")
|
||||
return out, torch.autograd.grad(sm, diff_args, allow_unused=True)
|
||||
|
||||
def check(args, ignore_failure=False):
|
||||
try:
|
||||
orig_out, orig_grad = call_forwards_backwards(f, args)
|
||||
except Exception:
|
||||
if ignore_failure:
|
||||
return
|
||||
raise
|
||||
|
||||
# See https://github.com/pytorch/pytorch/pull/98960#issuecomment-1505962215
|
||||
tensor_args = [x for x in pytree.tree_flatten(args)[0] if isinstance(x, torch.Tensor)]
|
||||
any_non_leaves = any(x.grad_fn is not None for x in tensor_args)
|
||||
if all(x is None for x in orig_grad) and any_non_leaves:
|
||||
with assert_raises_regex_fn(RuntimeError, 'does not require grad and does not have a grad_fn'):
|
||||
call_forwards_backwards(compiled_f, args)
|
||||
return
|
||||
|
||||
msg = (
|
||||
"Gradients of the operator are different in eager-mode PyTorch vs "
|
||||
"AOTDispatcher. This means the operator will have incorrect gradients "
|
||||
"underneath torch.compile. This could be because the operator's "
|
||||
"backward is incorrectly registered or not traceable."
|
||||
)
|
||||
|
||||
compiled_out, compiled_grad = call_forwards_backwards(compiled_f, args)
|
||||
if not skip_correctness_check:
|
||||
try:
|
||||
assert_equals_fn(compiled_out, orig_out)
|
||||
except Exception as e:
|
||||
raise type(e)(outputs_msg) from e
|
||||
try:
|
||||
assert_equals_fn(compiled_grad, orig_grad)
|
||||
except Exception as e:
|
||||
raise type(e)(msg) from e
|
||||
|
||||
check(args, ignore_failure=False)
|
||||
|
||||
# Randomize the data and run the traced graph with it, to catch bugs
|
||||
# where we may have baked in Tensor data into the trace.
|
||||
# This is not guaranteed to succeed, because `f` might have preconditions
|
||||
# on the values of the inputs, so we just ignore if this test fails.
|
||||
if try_check_data_specialization:
|
||||
args = randomize(args)
|
||||
check(args, ignore_failure=True)
|
||||
+135
@@ -0,0 +1,135 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import contextlib
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def set_autograd_fallback_mode(mode):
|
||||
prev = torch._C._get_autograd_fallback_mode()
|
||||
try:
|
||||
torch._C._set_autograd_fallback_mode(mode)
|
||||
yield
|
||||
finally:
|
||||
torch._C._set_autograd_fallback_mode(prev)
|
||||
|
||||
|
||||
def autograd_registration_check(op, args, kwargs):
|
||||
"""Check if autograd was registered correctly (for the operator).
|
||||
|
||||
Operators should have "autograd support" registered directly to an
|
||||
autograd dispatch key.
|
||||
An incorrect registration may lead to unexpected silent incorrectness.
|
||||
Note that this check won't catch all problems but will catch
|
||||
the most common ones.
|
||||
|
||||
Example usage:
|
||||
>>> x = torch.randn(3, requires_grad=True)
|
||||
>>> autograd_registration_check(torch.ops.aten.sin.default, (x,), {})
|
||||
|
||||
Here are some best practices if you do find your autograd is
|
||||
registered incorrectly:
|
||||
- If the operator is composite (i.e. consists of other PyTorch ops)
|
||||
and you wish the operator to decompose and get autograd support
|
||||
that way, then please register the implementation to
|
||||
DispatchKey::CompositeImplicitAutograd
|
||||
- If you're adding an autograd formula for the operator, the correct
|
||||
thing to do is to register an autograd.Function to
|
||||
DispatchKey::Autograd (preferred) or one of the
|
||||
DispatchKey::Autograd<BACKEND> keys. It is NOT OK to register
|
||||
an autograd.Function to a backend (e.g. CPU/CUDA) key.
|
||||
- If your operator is non-differentiable, then you should register
|
||||
an implementation to the Autograd key that uses
|
||||
AutoDispatchBelowAutograd and re-invokes the operator.
|
||||
|
||||
"""
|
||||
if not isinstance(op, torch._ops.OpOverload):
|
||||
raise AssertionError(f"Expected op to be OpOverload, got {type(op)}")
|
||||
# Implementation details
|
||||
# -----------------------------------------------
|
||||
# If an operator doesn't have an autograd kernel at an autograd key,
|
||||
# and the operator does not return inputs as-is, then all of
|
||||
# the outputs should have requires_grad=False before we apply
|
||||
# special behaviors of our default autograd fallback.
|
||||
# (The default autograd fallback may set requires_grad=True on output
|
||||
# tensors in certain modes so that when they are backpropped through,
|
||||
# they raise an error).
|
||||
#
|
||||
# Our strategy for detecting if an operator doesn't have an autograd
|
||||
# kernel at the autograd key is:
|
||||
# - set the autograd fallback mode to "nothing" (so it does not change
|
||||
# the required-gradness of outputs)
|
||||
# - run the operator
|
||||
# - Check if any outputs of the operator (that are not inputs) require
|
||||
# grad. This would only happen if the user calls regular PyTorch
|
||||
# operations in their backend key (this op should instead be
|
||||
# CompositeImplicitAutograd or not an op) or if the user invokes
|
||||
# an autograd.Function in the backend key.
|
||||
#
|
||||
# Note that it's already likely a bug if the operator directly returns
|
||||
# an input as output (because custom ops don't have a good way of
|
||||
# constructing true in-place or out variants), but we defer that
|
||||
# responsibility to a different test (schema_check).
|
||||
|
||||
flat_args = pytree.arg_tree_leaves(*args, **kwargs)
|
||||
all_tensors = [arg for arg in flat_args if isinstance(arg, torch.Tensor)]
|
||||
if not any(t.requires_grad for t in all_tensors):
|
||||
raise RuntimeError(
|
||||
"autograd_registration_check: no inputs have requires_grad=True so "
|
||||
"we are unable to actually perform this test. Please pass inputs "
|
||||
"that do require grad."
|
||||
)
|
||||
|
||||
# Determine which AutogradBACKEND key to check
|
||||
all_device_types = {arg.device.type for arg in all_tensors}
|
||||
if not all_device_types.issubset(["cpu", "cuda", "xpu"]):
|
||||
# Don't want to support other keys yet
|
||||
raise NotImplementedError(
|
||||
f"autograd_registration_check: NYI devices other than CPU/CUDA/XPU, got {all_device_types}"
|
||||
)
|
||||
if "cuda" in all_device_types:
|
||||
key = "AutogradCUDA"
|
||||
elif "cpu" in all_device_types:
|
||||
key = "AutogradCPU"
|
||||
elif "xpu" in all_device_types:
|
||||
key = "AutogradXPU"
|
||||
|
||||
if torch._C._dispatch_has_kernel_for_dispatch_key(op.name(), key):
|
||||
return
|
||||
if torch._C._dispatch_has_kernel_for_dispatch_key(op.name(), "Autograd"):
|
||||
return
|
||||
if torch._C._dispatch_has_kernel_for_dispatch_key(
|
||||
op.name(), "CompositeImplicitAutograd"
|
||||
):
|
||||
return
|
||||
|
||||
# At this point, we know the operator doesn't have a kernel registered to an
|
||||
# autograd key. Let's proceed with our test.
|
||||
with set_autograd_fallback_mode("nothing"):
|
||||
all_outs = op(*args, **kwargs)
|
||||
|
||||
inp_ids = {id(arg) for arg in flat_args}
|
||||
|
||||
def not_an_input_and_requires_grad(tensor):
|
||||
if not tensor.requires_grad:
|
||||
return False
|
||||
if id(tensor) in inp_ids:
|
||||
return False
|
||||
return True
|
||||
|
||||
if not pytree.tree_any_only(torch.Tensor, not_an_input_and_requires_grad, all_outs):
|
||||
return
|
||||
|
||||
raise AssertionError(
|
||||
f"{op.name()}: at least one output of this operator has requires_grad=True "
|
||||
f"but the operator does not have an autograd kernel defined at an autograd "
|
||||
f"key (e.g. DispatchKey::Autograd). This could mean that you have "
|
||||
f"incorrectly registered an autograd kernel to a non-Autograd DispatchKey, "
|
||||
f"which may lead to silently incorrect results. If your operator consists "
|
||||
f"of regular PyTorch operations, consider not using an operator at all "
|
||||
f"or registering your operator as CompositeImplicitAutograd. If you have "
|
||||
f"an autograd.Function registered to a backend (CPU/CUDA/XPU) key, the correct "
|
||||
f"location for it is the Autograd key."
|
||||
)
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch._subclasses
|
||||
|
||||
|
||||
def is_builtin(op):
|
||||
return op.namespace in ('aten', 'prims', 'prim')
|
||||
|
||||
|
||||
def fake_check(op, args, kwargs):
|
||||
with torch._subclasses.CrossRefFakeMode(ignore_op_fn=is_builtin):
|
||||
op(*args, **kwargs)
|
||||
+866
@@ -0,0 +1,866 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import datetime
|
||||
import difflib
|
||||
import functools
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch._dynamo
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._dynamo.utils import clone_input
|
||||
from torch._library.custom_ops import CustomOpDef
|
||||
from torch._subclasses.schema_check_mode import SchemaCheckMode
|
||||
from torch._utils_internal import get_file_path_2
|
||||
from torch.overrides import TorchFunctionMode
|
||||
from torch.testing._internal.optests import (
|
||||
aot_autograd_check,
|
||||
autograd_registration_check,
|
||||
fake_check,
|
||||
)
|
||||
|
||||
|
||||
def dontGenerateOpCheckTests(reason: str):
|
||||
def inner(fun):
|
||||
fun._torch_dont_generate_opcheck_tests = True
|
||||
return fun
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
def is_abstract(tensor: torch.Tensor) -> bool:
|
||||
if tensor.is_meta:
|
||||
return True
|
||||
if torch._subclasses.fake_tensor.is_fake(tensor):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def safe_schema_check(
|
||||
op: torch._ops.OpOverload,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any],
|
||||
*,
|
||||
copy_inputs: bool = True,
|
||||
rtol: float | None = None,
|
||||
atol: float | None = None,
|
||||
) -> Any:
|
||||
if copy_inputs:
|
||||
args, kwargs = deepcopy_tensors((args, kwargs))
|
||||
if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)):
|
||||
return None
|
||||
with SchemaCheckMode():
|
||||
result = op(*args, **kwargs)
|
||||
return result
|
||||
|
||||
|
||||
def safe_autograd_registration_check(
|
||||
op: torch._ops.OpOverload,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any],
|
||||
*,
|
||||
copy_inputs: bool = True,
|
||||
rtol: float | None = None,
|
||||
atol: float | None = None,
|
||||
) -> None:
|
||||
if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)):
|
||||
return
|
||||
if copy_inputs:
|
||||
args, kwargs = deepcopy_tensors((args, kwargs))
|
||||
# Don't perform autograd_registration_check if none of the inputs require grad.
|
||||
if not pytree.tree_any_only(
|
||||
torch.Tensor, lambda x: x.requires_grad, (args, kwargs)
|
||||
):
|
||||
return
|
||||
return autograd_registration_check(op, args, kwargs)
|
||||
|
||||
|
||||
def safe_fake_check(
|
||||
op: torch._ops.OpOverload,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any],
|
||||
*,
|
||||
copy_inputs: bool = True,
|
||||
rtol: float | None = None,
|
||||
atol: float | None = None,
|
||||
) -> None:
|
||||
if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)):
|
||||
return None
|
||||
if copy_inputs:
|
||||
args, kwargs = deepcopy_tensors((args, kwargs))
|
||||
return fake_check(op, args, kwargs)
|
||||
|
||||
|
||||
def safe_aot_autograd_check(
|
||||
op: torch._ops.OpOverload,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any],
|
||||
dynamic: bool,
|
||||
*,
|
||||
copy_inputs: bool = True,
|
||||
rtol: float | None = None,
|
||||
atol: float | None = None,
|
||||
) -> Any:
|
||||
# NB: copy_inputs does nothing for aot_autograd_check: it always needs to copy
|
||||
# inputs.
|
||||
if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)):
|
||||
return None
|
||||
|
||||
def func(*args, **kwargs):
|
||||
args, kwargs = pytree.tree_map_only(torch.Tensor, torch.clone, (args, kwargs))
|
||||
return op(*args, **kwargs)
|
||||
|
||||
# aot_autograd_check runs func(*args, **kwargs) multiple times
|
||||
# and assumes `func` does not modify its inputs.
|
||||
if rtol and atol:
|
||||
assert_equals_fn = functools.partial(
|
||||
torch.testing.assert_close, rtol=rtol, atol=atol
|
||||
)
|
||||
else:
|
||||
assert_equals_fn = torch.testing.assert_close
|
||||
return aot_autograd_check(
|
||||
func,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic,
|
||||
check_gradients="auto",
|
||||
assert_equals_fn=assert_equals_fn,
|
||||
)
|
||||
|
||||
|
||||
def deepcopy_tensors(inputs: Any) -> Any:
|
||||
return pytree.tree_map_only(torch.Tensor, clone_input, inputs)
|
||||
|
||||
|
||||
# Test util requirements
|
||||
# - The test util must have signature (op: OpOverload, args, kwargs)
|
||||
# - The test util must NOT mutate args, kwargs.
|
||||
# - The test utils in this list must not be prefixes of each other. For example,
|
||||
# having both "test_schema" and "test_schema_is_functional" is NOT OK.
|
||||
# - The order of items in this dict matters (for opcheck), we'll run them
|
||||
# in order.
|
||||
ALL_TEST_UTILS = {
|
||||
"test_schema": safe_schema_check,
|
||||
"test_autograd_registration": safe_autograd_registration_check,
|
||||
"test_faketensor": safe_fake_check,
|
||||
"test_aot_dispatch_static": functools.partial(
|
||||
safe_aot_autograd_check,
|
||||
dynamic=False,
|
||||
),
|
||||
"test_aot_dispatch_dynamic": functools.partial(
|
||||
safe_aot_autograd_check,
|
||||
dynamic=True,
|
||||
),
|
||||
}
|
||||
|
||||
GDOC = "https://docs.google.com/document/d/1Pj5HRZvdOq3xpFpbEjUZp2hBovhy7Wnxw14m6lF2154/edit"
|
||||
|
||||
DEFAULT_TEST_UTILS = [
|
||||
"test_schema",
|
||||
"test_autograd_registration",
|
||||
"test_faketensor",
|
||||
"test_aot_dispatch_dynamic",
|
||||
]
|
||||
|
||||
DEPRECATED_DEFAULT_TEST_UTILS = DEFAULT_TEST_UTILS + [
|
||||
"test_aot_dispatch_static",
|
||||
]
|
||||
|
||||
|
||||
def generate_opcheck_tests(
|
||||
testcase: Any,
|
||||
namespaces: list[str],
|
||||
failures_dict_path: str | None = None,
|
||||
additional_decorators: dict[str, Callable] | None = None,
|
||||
test_utils: list[str] = DEFAULT_TEST_UTILS,
|
||||
) -> None:
|
||||
"""Given an existing TestCase, use the existing tests to generate
|
||||
additional validation tests for custom operators.
|
||||
|
||||
For {all existing tests in the TestCase} x {all test utils},
|
||||
we will generate one new test. The new test runs a TorchFunctionMode
|
||||
that intercepts ``op(*args, **kwargs)`` calls and invokes
|
||||
``test_util(op, *args, **kwargs)``, where ``op`` is an operator.
|
||||
|
||||
The test_util that we support are in ALL_TEST_UTILS. They are:
|
||||
- test_schema: This runs SchemaCheckMode.
|
||||
- test_autograd_registration: This runs autograd_registration_check.
|
||||
- test_faketensor: This runs CrossRefFakeMode.
|
||||
- test_aot_dispatch_static: This runs aot_autograd_check, which:
|
||||
checks that the outputs (and gradients, if they are computable)
|
||||
are the same under eager-mode PyTorch and using AOTAutograd.
|
||||
- test_aot_dispatch_dynamic: Same as aot_dispatch_static, but
|
||||
runs AOTAutograd using dynamic shapes instead of static shapes.
|
||||
|
||||
The generated test will have name ``{test_util}__{original_name}``.
|
||||
For example, if there is a method named ``test_cumsum``, then
|
||||
we will generate a ``test_schema__test_cumsum``,
|
||||
``test_faketensor__test_cumsum``, etc.
|
||||
|
||||
For more details, see https://docs.google.com/document/d/1Pj5HRZvdOq3xpFpbEjUZp2hBovhy7Wnxw14m6lF2154/edit
|
||||
|
||||
Args:
|
||||
testcase: The testcase we will modify and generate additional tests for.
|
||||
namespaces: We will only intercept calls to custom operators with these
|
||||
namespaces.
|
||||
failures_dict_path: See ``validate_failures_dict_structure`` for more details
|
||||
test_utils: a list of test_utils to generate. Example: ["test_schema", "test_faketensor"]
|
||||
"""
|
||||
if additional_decorators is None:
|
||||
additional_decorators = {}
|
||||
test_methods = [
|
||||
m
|
||||
for m in dir(testcase)
|
||||
if m.startswith("test_") and callable(getattr(testcase, m))
|
||||
]
|
||||
if failures_dict_path is None:
|
||||
# The default failures_dict_path is failures_dict.json in
|
||||
# the same directory as the test file.
|
||||
prev_frame = inspect.currentframe().f_back
|
||||
filename = inspect.getframeinfo(prev_frame)[0]
|
||||
failures_dict_path = get_file_path_2(
|
||||
os.path.dirname(filename), "failures_dict.json"
|
||||
)
|
||||
failures_dict = FailuresDict.load(
|
||||
failures_dict_path, create_file=should_update_failures_dict()
|
||||
)
|
||||
validate_failures_dict_structure(failures_dict, test_utils, testcase)
|
||||
validate_failures_dict_formatting(failures_dict_path)
|
||||
|
||||
def construct_method(attr, prefix, tester):
|
||||
method = getattr(testcase, attr)
|
||||
if getattr(method, "_torch_dont_generate_opcheck_tests", False):
|
||||
return
|
||||
new_method_name = prefix + "__" + attr
|
||||
|
||||
@functools.wraps(method)
|
||||
def new_method(*args, **kwargs):
|
||||
with OpCheckMode(
|
||||
namespaces,
|
||||
prefix,
|
||||
tester,
|
||||
failures_dict,
|
||||
f"{testcase.__name__}.{new_method_name}",
|
||||
failures_dict_path,
|
||||
):
|
||||
result = method(*args, **kwargs)
|
||||
return result
|
||||
|
||||
if pytestmark := new_method.__dict__.get("pytestmark"):
|
||||
import pytest
|
||||
|
||||
# check if we need to simplify the parametrize marks
|
||||
# NB: you need to add this mark to your pytest.ini
|
||||
opcheck_only_one = False
|
||||
for mark in pytestmark:
|
||||
if isinstance(mark, pytest.Mark) and mark.name == "opcheck_only_one":
|
||||
opcheck_only_one = True
|
||||
|
||||
if opcheck_only_one:
|
||||
new_pytestmark = []
|
||||
for mark in pytestmark:
|
||||
if isinstance(mark, pytest.Mark) and mark.name == "parametrize":
|
||||
argnames, argvalues = mark.args
|
||||
if mark.kwargs:
|
||||
raise AssertionError("NYI: mark.kwargs is not empty")
|
||||
# Special case for device, we want to run on all
|
||||
# devices
|
||||
if argnames != "device":
|
||||
new_pytestmark.append(
|
||||
pytest.mark.parametrize(
|
||||
argnames, (next(iter(argvalues)),)
|
||||
)
|
||||
)
|
||||
continue
|
||||
new_pytestmark.append(mark)
|
||||
new_method.__dict__["pytestmark"] = new_pytestmark
|
||||
|
||||
if new_method_name in additional_decorators:
|
||||
for dec in additional_decorators[new_method_name]:
|
||||
new_method = dec(new_method)
|
||||
|
||||
if hasattr(testcase, new_method_name):
|
||||
raise RuntimeError(
|
||||
f"Tried to autogenerate {new_method_name} but {testcase} already "
|
||||
f"has method named {new_method_name}. Please rename the original "
|
||||
f"method on the TestCase."
|
||||
)
|
||||
setattr(testcase, new_method_name, new_method)
|
||||
|
||||
test_utils = {name: ALL_TEST_UTILS[name] for name in test_utils}
|
||||
for attr in test_methods:
|
||||
for prefix, tester in test_utils.items():
|
||||
construct_method(attr, prefix, tester)
|
||||
|
||||
generate_tag_tests(testcase, failures_dict, additional_decorators)
|
||||
|
||||
|
||||
def generate_tag_tests(testcase, failures_dict, additional_decorators):
|
||||
def generate_test(qualname, definitely_not_pt2_compliant, xfailed_tests):
|
||||
def inner(self):
|
||||
try:
|
||||
op = torch._library.utils.lookup_op(qualname)
|
||||
except AttributeError as e:
|
||||
# Operator not importable in this test file
|
||||
raise unittest.SkipTest(f"Can't import operator {qualname}") from e
|
||||
op_marked_as_compliant = torch.Tag.pt2_compliant_tag in op.tags
|
||||
if not op_marked_as_compliant:
|
||||
return
|
||||
if not definitely_not_pt2_compliant:
|
||||
return
|
||||
raise AssertionError(
|
||||
f"op '{qualname}' was tagged with torch.Tag.pt2_compliant_tag "
|
||||
f"but it failed some of the generated opcheck tests "
|
||||
f"({xfailed_tests}). This may lead to silent correctness issues, "
|
||||
f"please fix this."
|
||||
)
|
||||
|
||||
return inner
|
||||
|
||||
for qualname, test_dict in failures_dict.data.items():
|
||||
xfailed_tests = [
|
||||
test
|
||||
for test, status_dict in test_dict.items()
|
||||
# We're about to delete the following test after Ed's PR
|
||||
# to specialize on C++ .size() calls
|
||||
if "test_aot_dispatch_static" not in test
|
||||
and status_dict["status"] == "xfail"
|
||||
]
|
||||
definitely_not_pt2_compliant = len(xfailed_tests) > 0
|
||||
generated = generate_test(qualname, definitely_not_pt2_compliant, xfailed_tests)
|
||||
|
||||
# Could result in collisions, but unlikely. We'll raise if we see one below.
|
||||
mangled_qualname = qualname.replace("::", "_").replace(".", "_")
|
||||
test_name = "test_pt2_compliant_tag_" + mangled_qualname
|
||||
|
||||
# You can skip this test via the additional_decorators argument
|
||||
# in generate_opcheck_tests
|
||||
if test_name in additional_decorators:
|
||||
for decorator in additional_decorators[test_name]:
|
||||
generated = decorator(generated)
|
||||
|
||||
if hasattr(testcase, test_name):
|
||||
raise RuntimeError(
|
||||
f"Tried to generate a test named {test_name}, but it exists "
|
||||
f"already. This could be because of a name collision (where "
|
||||
f"we generated two tests with the same name), or where we "
|
||||
f"generated a test with the same name as an existing test."
|
||||
)
|
||||
setattr(testcase, test_name, generated)
|
||||
|
||||
|
||||
TEST_OPTIONS = ("xfail", "skip", "xsuccess")
|
||||
|
||||
|
||||
def validate_failures_dict_formatting(failures_dict_path: str) -> None:
|
||||
with open(failures_dict_path) as fp:
|
||||
actual = fp.read()
|
||||
failures_dict = FailuresDict.load(failures_dict_path)
|
||||
expected = failures_dict._save(to_str=True)
|
||||
if actual == expected:
|
||||
return
|
||||
if should_update_failures_dict():
|
||||
failures_dict = FailuresDict.load(failures_dict_path)
|
||||
failures_dict.save()
|
||||
return
|
||||
expected = expected.splitlines(1)
|
||||
actual = actual.splitlines(1)
|
||||
diff = difflib.unified_diff(actual, expected)
|
||||
diff = "".join(diff)
|
||||
raise RuntimeError(
|
||||
f"\n{diff}\n\nExpected the failures dict to be formatted "
|
||||
f"a certain way. Please see the above diff; you can correct "
|
||||
f"this either manually or by re-running the test with "
|
||||
f"PYTORCH_OPCHECK_ACCEPT=1"
|
||||
)
|
||||
|
||||
|
||||
def validate_failures_dict_structure(
|
||||
failure_dict: "FailuresDict", test_utils: list[str], testcase: Any
|
||||
) -> None:
|
||||
"""Validates the failures dict.
|
||||
|
||||
The failure dict looks something like the following.
|
||||
It maps operator name (qualname) to a list of autogenerated tests.
|
||||
Each autogenerated test may have a check for the operator (if the operator is
|
||||
called by the test); the dictionary specifies if we should skip the check,
|
||||
or if we expect some check to fail.
|
||||
|
||||
{
|
||||
"fbgemm::split_lengths": {
|
||||
"test_schema__test_split_lengths": {
|
||||
"comment": "you can put whatever you want into the comment section",
|
||||
"status": "xfail",
|
||||
}
|
||||
"test_schema__test_split_lengths_empty": {
|
||||
"comment": "",
|
||||
"status": "skip",
|
||||
},
|
||||
},
|
||||
"fbgemm::gather_lengths": {
|
||||
"test_schema__test_gather_lengths": {
|
||||
"comment": "",
|
||||
"status": "skip",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
"""
|
||||
failure_dict = failure_dict.data
|
||||
for test_to_option in failure_dict.values():
|
||||
for test_name, test_dict in test_to_option.items():
|
||||
if set(test_dict.keys()) != set({"comment", "status"}):
|
||||
raise RuntimeError(
|
||||
"in failures_dict, expected sub-dict to have keys 'comment' and 'status'"
|
||||
)
|
||||
test_option = test_dict["status"]
|
||||
if test_option not in TEST_OPTIONS:
|
||||
raise RuntimeError(
|
||||
f"In failures_dict, got status={test_option} but it needs to be in {TEST_OPTIONS}"
|
||||
)
|
||||
test_class, actual_test_name = test_name.split(".")
|
||||
if not any(actual_test_name.startswith(test) for test in test_utils):
|
||||
raise RuntimeError(
|
||||
f"In failures_dict, test name '{test_name}' should begin with one of {test_utils}"
|
||||
)
|
||||
for test in test_utils:
|
||||
if not actual_test_name.startswith(test):
|
||||
continue
|
||||
base_test_name = actual_test_name[len(test) + 2 :]
|
||||
# remove potential pytest parametrization suffix
|
||||
base_test_name = re.sub(r"\[.*\]", "", base_test_name)
|
||||
if testcase.__name__ != test_class:
|
||||
continue
|
||||
if hasattr(testcase, base_test_name):
|
||||
continue
|
||||
raise RuntimeError(
|
||||
f"In failures dict, got test name '{test_name}'. We parsed this as "
|
||||
f"running test '{test}' on '{base_test_name}', but "
|
||||
f"{base_test_name} does not exist on the TestCase '{testcase.__name__}]. "
|
||||
f"Maybe you need to change the test name?"
|
||||
)
|
||||
|
||||
|
||||
def should_update_failures_dict() -> bool:
|
||||
key = "PYTORCH_OPCHECK_ACCEPT"
|
||||
return key in os.environ and os.environ[key] == "1"
|
||||
|
||||
|
||||
_is_inside_opcheck_mode = threading.local()
|
||||
_is_inside_opcheck_mode.value = False
|
||||
|
||||
|
||||
def is_inside_opcheck_mode():
|
||||
return _is_inside_opcheck_mode.value
|
||||
|
||||
|
||||
class OpCheckMode(TorchFunctionMode):
|
||||
"""
|
||||
For a given test, OpCheckMode intercepts calls to operators and runs
|
||||
test_util(op, args, kwargs) for each intercepted (op, args, kwargs).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
namespaces: list[str],
|
||||
test_util_name: str,
|
||||
test_util: Callable,
|
||||
failures_dict: "FailuresDict",
|
||||
test_name: str,
|
||||
failures_dict_path: str,
|
||||
):
|
||||
# We will intercept calls to ops with these namespaces
|
||||
self.namespaces = namespaces
|
||||
# The test utility function. Its signature should be (op, args, kwargs) -> None.
|
||||
# Examples of test utilities are: schema_check, make_fx_check
|
||||
self.test_util = test_util
|
||||
self.test_util_name = test_util_name
|
||||
# The name of the test that is running this OpCheckMode.
|
||||
self.test_name = test_name
|
||||
# Maps qualname -> test_name -> skip/xfail
|
||||
# Tells us if we should skip a test or assert that there is a failure.
|
||||
self.failures_dict = failures_dict
|
||||
# Location of the failures dict. Makes it so that the error message is better.
|
||||
self.failures_dict_path = failures_dict_path
|
||||
|
||||
# OpCheckMode suppresses errors, collects them here, and then raises them on exit.
|
||||
# Maps qualname -> List[(Exception, func, maybe args, maybe kwargs)]
|
||||
self.seen_ops_to_errors = {}
|
||||
|
||||
def maybe_raise_errors_on_exit(self) -> None:
|
||||
# Check expected failures first
|
||||
for qualname in self.seen_ops_to_errors:
|
||||
option = self.failures_dict.get_status(qualname, self.test_name)
|
||||
if len(self.seen_ops_to_errors[qualname]) == 0:
|
||||
if should_update_failures_dict():
|
||||
self.failures_dict.set_status(
|
||||
qualname, self.test_name, "xsuccess", comment=""
|
||||
)
|
||||
else:
|
||||
if option == "xfail":
|
||||
raise OpCheckError(
|
||||
f"generate_opcheck_tests: Unexpected success for operator "
|
||||
f"{qualname} on test {self.test_name}. This may mean that "
|
||||
f"you have fixed this test failure. Please rerun the test with "
|
||||
f"PYTORCH_OPCHECK_ACCEPT=1 to automatically update the test runner "
|
||||
f"or manually remove the "
|
||||
f"expected failure in the failure dict at "
|
||||
f"{self.failures_dict_path}"
|
||||
f"For more details, see "
|
||||
f"{GDOC}"
|
||||
)
|
||||
continue
|
||||
failed_ops = []
|
||||
for qualname in self.seen_ops_to_errors:
|
||||
option = self.failures_dict.get_status(qualname, self.test_name)
|
||||
if option != "xsuccess":
|
||||
continue
|
||||
if len(self.seen_ops_to_errors[qualname]) == 0:
|
||||
continue
|
||||
failed_ops.append(qualname)
|
||||
if not failed_ops:
|
||||
return
|
||||
|
||||
if should_update_failures_dict():
|
||||
for op in failed_ops:
|
||||
self.failures_dict.set_status(op, self.test_name, "xfail")
|
||||
return
|
||||
|
||||
# Raise from the first error but also report about all of them to make
|
||||
# recording xfails easier.
|
||||
ex, op, args, kwargs = self.seen_ops_to_errors[failed_ops[0]][0]
|
||||
repro_command = generate_repro(
|
||||
self.test_util_name, op, args, kwargs, save_data=should_print_better_repro()
|
||||
)
|
||||
raise OpCheckError(
|
||||
f"Test generated by `generate_opcheck_tests`, {self.test_name}, "
|
||||
f"failed on operators {failed_ops}. This usually means that the "
|
||||
f"operators are not implemented correctly and may lead to silently "
|
||||
f"incorrect behavior. Set PYTORCH_OPCHECK_PRINT_BETTER_REPRO=1 for a standalone repro, "
|
||||
f"or please see "
|
||||
f"{GDOC} "
|
||||
f"for more recommendations. "
|
||||
f"To reproduce this problem locally, try to run the following:\n{repro_command}"
|
||||
) from ex
|
||||
|
||||
def __enter__(self, *args, **kwargs):
|
||||
self.prev_is_opcheck_mode = _is_inside_opcheck_mode.value
|
||||
self.prev_dynamo_disable = os.environ.get("TORCHDYNAMO_DISABLE", "")
|
||||
_is_inside_opcheck_mode.value = True
|
||||
os.environ["TORCHDYNAMO_DISABLE"] = "1"
|
||||
# When running this test mode, we want to disable
|
||||
# default torch.compile custom op checker
|
||||
self.prev_functorch_config_for_checking_custom_op = (
|
||||
torch._functorch.config.check_custom_op_aliasing
|
||||
)
|
||||
torch._functorch.config.check_custom_op_aliasing = False
|
||||
return super().__enter__(*args, **kwargs)
|
||||
|
||||
def __exit__(self, *args, **kwargs):
|
||||
_is_inside_opcheck_mode.value = self.prev_is_opcheck_mode
|
||||
os.environ["TORCHDYNAMO_DISABLE"] = self.prev_dynamo_disable
|
||||
torch._functorch.config.check_custom_op_aliasing = (
|
||||
self.prev_functorch_config_for_checking_custom_op
|
||||
)
|
||||
try:
|
||||
self.maybe_raise_errors_on_exit()
|
||||
if should_update_failures_dict():
|
||||
self.failures_dict.save()
|
||||
finally:
|
||||
result = super().__exit__(*args, **kwargs)
|
||||
return result
|
||||
|
||||
def run_test_util(self, op, args, kwargs):
|
||||
try:
|
||||
self.test_util(op, args, kwargs, copy_inputs=False)
|
||||
except torch._subclasses.fake_tensor.UnsupportedFakeTensorException:
|
||||
# We might get here if the input is already a FakeTensor
|
||||
# or if we're in a torch.compile block. Just ignore these
|
||||
# since we can't handle them and reporting them as failures
|
||||
# is too noisy.
|
||||
pass
|
||||
|
||||
def __torch_function__(self, func, types, args=(), kwargs=None):
|
||||
kwargs = kwargs if kwargs else {}
|
||||
|
||||
# Only intercept calls to operators
|
||||
if not isinstance(func, (torch._ops.OpOverloadPacket, torch._ops.OpOverload)):
|
||||
return func(*args, **kwargs)
|
||||
if (
|
||||
torch.jit.is_tracing()
|
||||
or torch.jit.is_scripting()
|
||||
or torch._dynamo.is_compiling()
|
||||
):
|
||||
return func(*args, **kwargs)
|
||||
# Pre-existing code may not use the .default overload. If we see an
|
||||
# OpOverloadPacket and we cannot resolve the overload, then we just throw
|
||||
# and ask the user to clarify. Otherwise, we attempt to resolve the overload.
|
||||
if isinstance(func, torch._ops.OpOverloadPacket):
|
||||
func = resolve_unique_overload_or_throw(func)
|
||||
qualname = func.name()
|
||||
ns = qualname.split("::")[0]
|
||||
if ns not in self.namespaces:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
args_c, kwargs_c = deepcopy_tensors((args, kwargs))
|
||||
result = func(*args, **kwargs)
|
||||
|
||||
option = self.failures_dict.get_status(qualname, self.test_name)
|
||||
if option == "xsuccess" or option == "xfail":
|
||||
# Suppress all errors during execution. Raise them during __exit__.
|
||||
try:
|
||||
if qualname not in self.seen_ops_to_errors:
|
||||
self.seen_ops_to_errors[qualname] = []
|
||||
self.run_test_util(func, args_c, kwargs_c)
|
||||
except Exception as ex:
|
||||
if should_print_better_repro():
|
||||
self.seen_ops_to_errors[qualname].append((ex, func, args, kwargs))
|
||||
else:
|
||||
self.seen_ops_to_errors[qualname].append((ex, func, None, None))
|
||||
elif option == "skip":
|
||||
pass
|
||||
return result
|
||||
|
||||
|
||||
def should_print_better_repro() -> None:
|
||||
"""If set, the tests generated by `generate_opcheck_tests` will print a
|
||||
repro command on failure.
|
||||
|
||||
In order to print the repro command, we need to save some tensors to disk.
|
||||
These will be saved under the following directory:
|
||||
{tempfile.gettempdir()}/pytorch_opcheck_safe_to_delete/.
|
||||
|
||||
Although this is a temp folder, it will usually not automatically get cleaned
|
||||
up, so you'll need to manually delete it.
|
||||
"""
|
||||
key = "PYTORCH_OPCHECK_PRINT_BETTER_REPRO"
|
||||
if key not in os.environ:
|
||||
return False
|
||||
value = os.environ[key]
|
||||
return value == "1" or value == 1
|
||||
|
||||
|
||||
def opcheck(
|
||||
op: torch._ops.OpOverload | torch._ops.OpOverloadPacket | CustomOpDef,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
*,
|
||||
test_utils: str | Sequence[str] = DEFAULT_TEST_UTILS,
|
||||
raise_exception: bool = True,
|
||||
rtol: float | None = None,
|
||||
atol: float | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""See torch.library.opcheck for docstring"""
|
||||
|
||||
if (rtol is None) ^ (atol is None):
|
||||
raise ValueError(
|
||||
"opcheck(op, ...): if you specify one of rtol/atol, you must specify both"
|
||||
)
|
||||
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
if isinstance(op, CustomOpDef):
|
||||
op = op._opoverload
|
||||
if isinstance(op, torch._ops.OpOverloadPacket):
|
||||
op = resolve_unique_overload_or_throw(op)
|
||||
if not isinstance(op, torch._ops.OpOverload):
|
||||
raise ValueError(
|
||||
f"opcheck(op, ...): op must be instance of torch._ops.OpOverload, "
|
||||
f"e.g. torch.ops.aten.sin.default, got {type(op)}"
|
||||
)
|
||||
if test_utils == "ALL":
|
||||
test_utils = tuple(ALL_TEST_UTILS.keys())
|
||||
if isinstance(test_utils, str):
|
||||
test_utils = (test_utils,)
|
||||
if not isinstance(test_utils, (tuple, list)) or not set(test_utils).issubset(
|
||||
ALL_TEST_UTILS.keys()
|
||||
):
|
||||
raise ValueError(
|
||||
f"opcheck(op, ..., test_utils={test_utils}), expected test_utils "
|
||||
f"to be subset of {tuple(ALL_TEST_UTILS.keys())} but it was not"
|
||||
)
|
||||
|
||||
results_dict = {}
|
||||
for test_util in test_utils:
|
||||
tester = ALL_TEST_UTILS[test_util]
|
||||
try:
|
||||
tester(op, args, kwargs, rtol=rtol, atol=atol)
|
||||
results_dict[test_util] = "SUCCESS"
|
||||
except Exception as ex:
|
||||
if raise_exception:
|
||||
raise OpCheckError(
|
||||
f"opcheck(op, ...): {test_util} failed with {ex} "
|
||||
f"(scroll up for stack trace)"
|
||||
) from ex
|
||||
results_dict[test_util] = ex
|
||||
return results_dict
|
||||
|
||||
|
||||
class OpCheckError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def generate_repro(
|
||||
test: str,
|
||||
op: torch._ops.OpOverload,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any],
|
||||
*,
|
||||
save_data: bool,
|
||||
dry_run: bool = False,
|
||||
) -> str:
|
||||
if save_data:
|
||||
now = datetime.datetime.now()
|
||||
path = os.path.join(tempfile.gettempdir(), "pytorch_opcheck_safe_to_delete")
|
||||
unix_timestamp = datetime.datetime.timestamp(now) * 100000
|
||||
filepath = os.path.join(path, f"repro_{unix_timestamp}.pt")
|
||||
if not dry_run:
|
||||
os.makedirs(path, exist_ok=True)
|
||||
torch.save((args, kwargs), filepath)
|
||||
args_kwargs = f'args, kwargs = torch.load("{filepath}")'
|
||||
else:
|
||||
args_kwargs = (
|
||||
"# If you rerun your test with PYTORCH_OPCHECK_PRINT_BETTER_REPRO=1\n"
|
||||
"# we will fill them in same (args, kwargs) as in your test\n"
|
||||
"args = () # args to the operator\n"
|
||||
"kwargs = {} # kwargs to the operator"
|
||||
)
|
||||
|
||||
ns, name = op._schema.name.split("::")
|
||||
overload = op._overloadname
|
||||
|
||||
repro_command = (
|
||||
f"# =========================================================\n"
|
||||
f"# BEGIN REPRO SCRIPT\n"
|
||||
f"# =========================================================\n"
|
||||
f"import torch\n"
|
||||
f"from torch.testing._internal.optests import opcheck\n"
|
||||
f"\n"
|
||||
f"# Make sure you have loaded the library that contains the op\n"
|
||||
f"# via an import or torch.ops.load_library(...)\n"
|
||||
f"op = torch.ops.{ns}.{name}.{overload}\n"
|
||||
f"\n"
|
||||
f"{args_kwargs}\n"
|
||||
f'opcheck(op, args, kwargs, test_utils="{test}")\n'
|
||||
f"# =========================================================\n"
|
||||
f"# END REPRO SCRIPT\n"
|
||||
f"# =========================================================\n"
|
||||
)
|
||||
return repro_command
|
||||
|
||||
|
||||
def resolve_unique_overload_or_throw(
|
||||
op: torch._ops.OpOverloadPacket,
|
||||
) -> torch._ops.OpOverload:
|
||||
all_schemas = torch._C._jit_get_schemas_for_operator(op._qualified_op_name)
|
||||
if len(all_schemas) != 1:
|
||||
raise RuntimeError(
|
||||
f"opcheck can only test operators without overloads. "
|
||||
f"Got the following overloads for {op._qualified_op_name}: "
|
||||
f"{[schema.overload_name for schema in all_schemas]}"
|
||||
)
|
||||
|
||||
overload_name = all_schemas[0].overload_name
|
||||
if overload_name == "":
|
||||
return op.default
|
||||
return getattr(op, overload_name)
|
||||
|
||||
|
||||
DUMP_OPTIONS = {"indent": 2, "sort_keys": True}
|
||||
|
||||
|
||||
FailuresDictData = dict[str, dict[str, dict[str, str]]]
|
||||
|
||||
|
||||
VERSION = 1
|
||||
DESCRIPTION = (
|
||||
f"This is a dict containing failures for tests autogenerated by "
|
||||
f"generate_opcheck_tests. "
|
||||
f"For more details, please see {GDOC}"
|
||||
)
|
||||
|
||||
|
||||
class FailuresDict:
|
||||
def __init__(self, path: str, data: FailuresDictData):
|
||||
self.path = path
|
||||
self.data = data
|
||||
|
||||
@staticmethod
|
||||
def load(path, *, create_file=False) -> "FailuresDict":
|
||||
if create_file and not os.path.exists(path):
|
||||
result = FailuresDict(path, {})
|
||||
FailuresDict.save()
|
||||
return result
|
||||
with open(path) as fp:
|
||||
contents = fp.read()
|
||||
if contents.strip() == "":
|
||||
dct = {
|
||||
"_description": DESCRIPTION,
|
||||
"data": {},
|
||||
"_version": VERSION,
|
||||
}
|
||||
else:
|
||||
dct = json.loads(contents)
|
||||
if "data" not in dct:
|
||||
raise AssertionError("Expected 'data' in dct")
|
||||
if "_version" not in dct or dct["_version"] != VERSION:
|
||||
raise AssertionError(
|
||||
f"Expected '_version' in dct with value {VERSION}"
|
||||
)
|
||||
return FailuresDict(path, dct["data"])
|
||||
|
||||
def _save(self, to_str=False) -> str | None:
|
||||
to_dump = {
|
||||
"_description": DESCRIPTION,
|
||||
"data": self.data,
|
||||
"_version": VERSION,
|
||||
}
|
||||
# json.dumps doesn't end with a newline. Let's add one because files
|
||||
# should end in newlines.
|
||||
serialized = json.dumps(to_dump, **DUMP_OPTIONS) + "\n"
|
||||
if to_str:
|
||||
return serialized
|
||||
with open(self.path, "w") as fp:
|
||||
fp.write(serialized)
|
||||
return None
|
||||
|
||||
def save(self) -> None:
|
||||
return self._save()
|
||||
|
||||
def get_status(self, qualname: str, test_name: str) -> str:
|
||||
if qualname not in self.data:
|
||||
return "xsuccess"
|
||||
dct = self.data[qualname]
|
||||
if test_name not in dct:
|
||||
return "xsuccess"
|
||||
return dct[test_name]["status"]
|
||||
|
||||
def set_status(
|
||||
self,
|
||||
qualname: str,
|
||||
test_name: str,
|
||||
status: str,
|
||||
*,
|
||||
comment: str | None = None,
|
||||
):
|
||||
if qualname not in self.data:
|
||||
self.data[qualname] = {}
|
||||
dct = self.data[qualname]
|
||||
if test_name not in dct:
|
||||
dct[test_name] = {"status": None, "comment": ""}
|
||||
|
||||
if status == "xsuccess":
|
||||
# The default status is "xsuccess".
|
||||
del dct[test_name]
|
||||
else:
|
||||
dct[test_name]["status"] = status
|
||||
if comment is not None:
|
||||
dct[test_name]["comment"] = comment
|
||||
@@ -0,0 +1,89 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import torch
|
||||
from torch.fx.experimental.proxy_tensor import make_fx
|
||||
from torch.testing._utils import wrapper_set_seed
|
||||
import torch.utils._pytree as pytree
|
||||
|
||||
|
||||
def make_fx_check(
|
||||
func,
|
||||
args,
|
||||
kwargs,
|
||||
tracing_mode,
|
||||
assert_close=torch.testing.assert_close,
|
||||
randomize_data=False,
|
||||
):
|
||||
f, *new_args = handle_sizes_for_dynamic_shapes(func, args, kwargs)
|
||||
|
||||
def run(f, *args, **kwargs):
|
||||
return wrapper_set_seed(f, *args, **kwargs)
|
||||
|
||||
traced_f = make_fx(f, tracing_mode=tracing_mode)(*new_args)
|
||||
|
||||
msg = (
|
||||
"op(*args, **kwargs) and make_fx(op)(*args, **kwargs) produced different "
|
||||
"values. This could mean that your abstract impls (meta/FakeTensor impls) "
|
||||
"are incorrect, that your operator is not completely traceable (e.g., "
|
||||
"it relies on some global state), or that there is a bug in make_fx. "
|
||||
"Note that if you passed a python function (and not an operator) to "
|
||||
"make_fx_check, it is still possible that the python function will still "
|
||||
"work with torch.compile because it handles capturing pieces of "
|
||||
"your python code to compile."
|
||||
)
|
||||
|
||||
# Randomize the data and run the traced graph with it, to catch bugs
|
||||
# where we may have baked in Tensor data into the trace.
|
||||
# This is not guaranteed to succeed, because `f` might have preconditions
|
||||
# on the values of the inputs, so we just ignore if we used
|
||||
# random data and it fails.
|
||||
if randomize_data:
|
||||
new_args = randomize(new_args)
|
||||
try:
|
||||
expected = run(f, *new_args)
|
||||
except Exception:
|
||||
if randomize_data:
|
||||
return
|
||||
raise
|
||||
result = run(traced_f, *new_args)
|
||||
assert_close(result, expected, msg=msg)
|
||||
|
||||
|
||||
# Arguably we should make make_fx promote torch.Size() objects to symbolic shapes.
|
||||
# Absent that, here is our strategy:
|
||||
#
|
||||
# If any argument is a torch.Size(), maybe get dynamic shapes for it by:
|
||||
# - Create a temporary Tensor whose size is the torch.Size() we want. Note that
|
||||
# we use an expanded Tensor as we cannot pass "meta" Tensors to make_fx.
|
||||
# - Pass it to make_fx such that it is converted to a proxy Tensor
|
||||
# - Unpack the size in the wrapper to get a torch.Size with dynamic shapes (in
|
||||
# symbolic mode, a no-op otherwise)
|
||||
def handle_sizes_for_dynamic_shapes(func, args, kwargs):
|
||||
def f(args, kwargs, extra_args, extra_kwargs):
|
||||
if extra_args:
|
||||
for i, t in extra_args:
|
||||
args[i] = t.size()
|
||||
if extra_kwargs:
|
||||
for k, t in extra_kwargs.items():
|
||||
kwargs[k] = t.size()
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
extra_args = []
|
||||
extra_kwargs = {}
|
||||
for i, arg in enumerate(args):
|
||||
if isinstance(arg, torch.Size):
|
||||
extra_args.append((i, torch.empty(arg, device="cpu")))
|
||||
for key, value in kwargs.items():
|
||||
if isinstance(value, torch.Size):
|
||||
extra_kwargs[key] = torch.empty(value, device="cpu")
|
||||
|
||||
return f, args, kwargs, extra_args, extra_kwargs
|
||||
|
||||
|
||||
def randomize(args):
|
||||
def transform(x):
|
||||
if not x.dtype.is_floating_point:
|
||||
return x
|
||||
return x.detach().clone().uniform_(0, 1).requires_grad_(x.requires_grad)
|
||||
return pytree.tree_map_only(torch.Tensor, transform, args)
|
||||
@@ -0,0 +1,29 @@
|
||||
import torch
|
||||
|
||||
|
||||
# This is a very odd way to define a test case that uses CALL_INTRINSIC_1 4 and
|
||||
# CALL_INTRINSIC_2 7. The problem is those intrinsics only exists on Python
|
||||
# 3.12+ and the Python code to produce them is not backwards compatible with
|
||||
# previous versions.
|
||||
class Foo:
|
||||
@classmethod
|
||||
def _default_update(cls):
|
||||
def f[T](a: "This is a new annotation"): # noqa: F722
|
||||
"""This is a test"""
|
||||
|
||||
f.attr = "This is also a test"
|
||||
f.__wrapped__ = "This is a bald faced lie"
|
||||
|
||||
def wrapper(b: "This is the prior annotation"): # noqa: F722
|
||||
pass
|
||||
|
||||
return wrapper, f
|
||||
|
||||
def test_default_update(self):
|
||||
@torch.compile(backend="eager", fullgraph=True)
|
||||
def fn(x):
|
||||
wrapper, f = Foo._default_update()
|
||||
return x + 1
|
||||
|
||||
x = torch.randn(2)
|
||||
fn(x)
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class LinearReluFunctionalChild(nn.Module):
|
||||
def __init__(self, N):
|
||||
super().__init__()
|
||||
self.w1 = nn.Parameter(torch.empty(N, N))
|
||||
self.b1 = nn.Parameter(torch.zeros(N))
|
||||
torch.nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5))
|
||||
|
||||
def forward(self, x):
|
||||
x = torch.nn.functional.linear(x, self.w1, self.b1)
|
||||
x = torch.nn.functional.relu(x)
|
||||
return x
|
||||
|
||||
class LinearReluFunctional(nn.Module):
|
||||
def __init__(self, N):
|
||||
super().__init__()
|
||||
self.child = LinearReluFunctionalChild(N)
|
||||
self.w1 = nn.Parameter(torch.empty(N, N))
|
||||
self.b1 = nn.Parameter(torch.zeros(N))
|
||||
torch.nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5))
|
||||
|
||||
def forward(self, x):
|
||||
x = self.child(x)
|
||||
x = torch.nn.functional.linear(x, self.w1, self.b1)
|
||||
x = torch.nn.functional.relu(x)
|
||||
return x
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user