OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity
b2eb0e8c6a
pytorch/test/test_decomp.py
drisspg 1434e0b121 Add a private _safe_softmax (#131060) 
# Summary
Changes the stance of SDPA on what to do for fully masked out rows

## Current Behavior
Several PyTorch users have expressed frustration over this issue:
- https://github.com/pytorch/pytorch/issues/41508
- https://github.com/pytorch/pytorch/issues/103749
- https://github.com/pytorch/pytorch/issues/103963

These are significant issues with extensive discussion but no satisfactory resolution. The PyTorch team's consensus, as stated here:
https://github.com/pytorch/pytorch/issues/24816#issuecomment-524415617

Can be paraphrased as follows:

When passing in fully masked out rows, attention becomes ambiguous. We have two main options:

1. Uniformly attend to all values:
   ```python
   scores[masked_out_rows] = 1 / len(row)
   out[masked_out_rows] = 1 / len(row) * value
   ```

2. Decide that attention between no queries (masked) and no keys (masked) is meaningless:
   ```python
   output[fully_masked_rows] = NaN
   ```

We went with option 2. Partially because it was easier to implement, but also people argued that users can slice the output to remove the NaNs:
``` Python
>fill_value = -float("inf")
>row0 = torch.randn(4)
>row1 = torch.tensor([(fill_value for _ in range(4)])
>matrix = torch.stack([row0, row1]).requires_grad_(True)
>out = torch.softmax(matrix, 1)
>out = out[0]
>print(out)
tensor([0.5377, 0.2729, 0.0692, 0.1201])
```
Cool, problem solved. But what happends when you call backwards..
```Python
>out.backward(torch.ones_like(out))
>print(matrix.grad)
tensor([[3.0957e-08, 1.4157e-08, 7.7802e-10, 1.3713e-08],
        [       nan,        nan,        nan,        nan]])
```
Those pesky NaNs are back!

## Why do we see NaNs today?

The core of the problem revolves around using softmax function in sdpa:

```python
> row = torch.tensor([(-float("inf")) for _ in range(4)])
> torch.softmax(row, 0)
tensor([nan, nan, nan, nan])
```

## Quick Aside: Masking in Attention

Attention itself doesn't have a concept of masking. The `sdpa` function has an argument called `attn_mask`, which would be more accurately named `attn_bias`. This is because we don't actually "mask" entries when computing attention. Instead, due to implementation details([performance](https://github.com/pytorch/pytorch/issues/25110#issuecomment-524519087)), we add a value to the masked-out query/key pairs.

We use a large negative number (typically -inf) to decrease the attention weight, as softmax assigns more weight to larger values.

## Alternative Approaches

If we use a very large negative number instead of -inf:

```python
> row = torch.tensor([(-1e6) for _ in range(4)])
> torch.softmax(row, 0)
tensor([0.2500, 0.2500, 0.2500, 0.2500])
```
However if users always remembered to "slice" out their outputs i.e.:
```Python
>fill_value = -1e6
>...
>out.backward(torch.ones_like(out))
>print(matrix.grad)
tensor([[-0.0563, -0.0564,  0.1613, -0.0486],
        [ 0.0000,  0.0000,  0.0000,  0.0000]])
```
This would bring us back into a better state.

## A Third Option

We don't necessarily need to alter the behavior of softmax for -inf or very large negative numbers. The fundamental goal is to exclude certain query/key pairs from attention, regardless of the underlying implementation.

This PR implements the new semantic for masking w/ attention in fully masked-out rows:
```python
out[masked_out_rows] = 0
```

**Important Note**: This idea isn't entirely new. The [MaskedTensor](https://pytorch.org/tutorials/prototype/maskedtensor_overview#safe-softmax) prototype, a tensor subclass, was designed to handle such cases. However, it remains a prototype feature and hasn't gained widespread adoption.

## Details
This PR stack does 3 things:
1. Adds a PRIVATE _safe_softmax op
2. Updates semantic for flash_cpu fused kernel
3. Updates semantic for efficient_cuda fused kernel

_safe_softmax is not supposed to be used generically and is only meant to be used within the context of SDPA. Due to this fact instead of decomposing softmax and checking for -inf rows we instead "cheat" and use nan_to_num.

Why I think this is okay? (please find a counter point if avail)
There are multiple ways NaNs can emerge. For the fully masked out rows case nan_to_num works. But what if there were other NaNs, wouldn't this silently remove them?

The only case that this can happen is if the input itself had a NaN or an Inf
For example:
```Python
a = torch.ones([4], requires_grad=False, dtype=torch.float16)
a[1] = torch.finfo(torch.float16).max
print(a.softmax(-1))
```
Will return
`tensor([0., 1., 0., 0.], dtype=torch.float16)`

Where
```Python
a = torch.ones([4], requires_grad=False, dtype=torch.float16)
a[1] = float("inf")
a.softmax(-1)
```
returns:
`tensor([nan, nan, nan, nan], dtype=torch.float16)`

If we dont want to even allow for the possibility of "inf" or "NaN" attention scores to be converted to 0 then we can implemented it something like this

```Python
max = torch.max(a, dim=-1, keepdim=True)
exp = torch.exp(a - max.values)
denom = torch.sum(exp, dim=-1, keepdim=True)
softmax = exp / denom
softmax = torch.where(max.values == float('-inf'), 0.0, softmax)
```
however we would be paying for this in math performance.

## Why Now
I think one point that has substantially changed where PyTorch should lie on this argument is the fact that we have fused implementations for SDPA now. And these fused implementations allow us to easily and performantly support this new semantic.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/131060
Approved by: https://github.com/jbschlosser
2024-08-08 23:09:38 +00:00

1314 lines
51 KiB
Python
Raw Blame History

# Owner(s): ["module: decompositions"]
import functools
import itertools
import re
import unittest
from collections import defaultdict
from functools import partial
import torch._inductor.decomposition
import torch.autograd
from torch import Tensor
from torch._decomp import core_aten_decompositions, decomposition_table
from torch._dispatch.python import enable_python_dispatcher
from torch._ops import DispatchKey
from torch.testing import make_tensor
from torch.testing._internal.common_cuda import tf32_off
from torch.testing._internal.common_device_type import (
instantiate_device_type_tests,
onlyCPU,
onlyCUDA,
onlyNativeDeviceTypes,
ops,
)
from torch.testing._internal.common_methods_invocations import (
op_db,
skip,
skipOps,
xfail,
)
from torch.testing._internal.common_modules import module_db, modules
from torch.testing._internal.common_utils import (
is_iterable_of_tensors,
run_tests,
skipIfCrossRef,
skipIfTorchDynamo,
suppress_warnings,
TEST_WITH_ASAN,
TEST_WITH_SLOW,
TestCase,
unMarkDynamoStrictTest,
)
from torch.utils import _pytree as pytree
from torch.utils._python_dispatch import TorchDispatchMode
from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
aten = torch.ops.aten
# TODO: this isn't going to work with non-aten namespaces
def overload_to_aten_name(op):
return op._schema.name.split("::")[1]
# All operators that can have decomp tests
decomposition_names = {
overload_to_aten_name(k)
for k in decomposition_table
if isinstance(k, torch._ops.OpOverload)
}
core_decomposition_names = {
overload_to_aten_name(k)
for k in core_aten_decompositions()
if isinstance(k, torch._ops.OpOverload)
}
_decomp_test_ops = [
op
for op in op_db
if op.aten_name in decomposition_names
or op.aten_backward_name in decomposition_names
]
_decomp_test_ops_core_autograd = [
op
for op in op_db
if op.aten_name in core_decomposition_names and op.supports_autograd
]
_sdpa_op_info = [op for op in op_db if "scaled_dot_product_attention" in op.aten_name]
def diff_arg(arg, requires_grad=True):
def is_differentiable_arg(arg):
if requires_grad:
return arg.requires_grad
else:
return arg.is_floating_point() or arg.is_complex()
if is_iterable_of_tensors(arg):
if all(is_differentiable_arg(a) for a in arg):
return True
if all(not is_differentiable_arg(a) for a in arg):
return False
raise RuntimeError("NYI: The test runner can't handle this")
return isinstance(arg, Tensor) and is_differentiable_arg(arg)
# Version of autograd.grad with some differences:
# - pytree inputs is allowed (but leaves of the pytree have to all
# be tensors)
# - if an input is not used as part of derivatives, we will return a
# zero-filled tensor for the result
def _autograd_grad(
outputs, inputs, grad_outputs=None, retain_graph=False, create_graph=True
):
inputs, inputs_spec = tree_flatten(inputs)
diff_inputs = tuple(inp for inp in inputs if inp.requires_grad)
if grad_outputs is None:
diff_outputs = tuple(out for out in outputs if out.requires_grad)
else:
diff_grad_outputs = [
(out, go) for out, go in zip(outputs, grad_outputs) if out.requires_grad
]
if len(diff_grad_outputs) == 0:
diff_outputs, grad_outputs = (), ()
else:
diff_outputs, grad_outputs = zip(*diff_grad_outputs)
grad_inputs = torch.autograd.grad(
diff_outputs,
diff_inputs,
grad_outputs,
retain_graph=retain_graph,
create_graph=create_graph,
allow_unused=True,
)
result = []
grad_inputs_iter = iter(grad_inputs)
for inp in inputs:
if inp.requires_grad:
grad_input = next(grad_inputs_iter)
if grad_input is None:
result.append(torch.zeros_like(inp))
else:
result.append(grad_input)
else:
result.append(torch.zeros_like(inp))
return tree_unflatten(result, inputs_spec)
def _as_tuple(val):
if isinstance(val, tuple):
return val
return (val,)
def ref_vjp_no_create(f, *primals):
result = f(*primals)
def wrapped(cotangents):
return _autograd_grad(
_as_tuple(result),
primals,
_as_tuple(cotangents),
create_graph=False,
retain_graph=True,
)
return result, wrapped
dtype_precisions = {
torch.float16: (0.001, 1e-5),
torch.bfloat16: (0.016, 1e-4),
torch.float32: (1.3e-6, 1e-5),
torch.float64: (1e-7, 1e-7),
torch.complex32: (0.001, 1e-5),
torch.complex64: (1.3e-6, 1e-5),
torch.complex128: (1e-7, 1e-7),
}
# Returns the "default" rtol and atol for comparing scalars or
# tensors of the given dtypes.
def _getDefaultRtolAndAtol(dtype0, dtype1):
rtol = max(
dtype_precisions.get(dtype0, (0, 0))[0], dtype_precisions.get(dtype1, (0, 0))[0]
)
atol = max(
dtype_precisions.get(dtype0, (0, 0))[1], dtype_precisions.get(dtype1, (0, 0))[1]
)
return rtol, atol
def op_assert_ref(test_case, op, test_dtype, i, orig, decomp, ref, args, kwargs):
assert orig.dtype == decomp.dtype, f"{i} Operation: {op}"
if orig.numel() == 0 or decomp.numel() == 0:
assert orig.numel() == decomp.numel()
return
assert orig.shape == decomp.shape, f"{i} Operation: {op}"
tol_table = {
(torch.bfloat16, torch.ops.aten.native_layer_norm.default): 1e-5,
(torch.float16, torch.ops.aten.native_layer_norm.default): 1e-5,
(torch.float16, torch.ops.aten.native_layer_norm_backward.default): 1e-3,
(torch.bfloat16, torch.ops.aten.native_layer_norm_backward.default): 2e-2,
(torch.bfloat16, torch.ops.aten.native_batch_norm.default): 1e-5,
(torch.float16, torch.ops.aten.native_batch_norm.default): 1e-5,
(torch.bfloat16, torch.ops.aten._native_batch_norm_legit.default): 1e-5,
(torch.bfloat16, torch.ops.aten._native_batch_norm_legit.no_stats): 1e-5,
(torch.float16, torch.ops.aten._native_batch_norm_legit.default): 1e-5,
(torch.float16, torch.ops.aten._native_batch_norm_legit.no_stats): 1e-5,
(torch.bfloat16, torch.ops.aten.linalg_vector_norm.default): 1e-4,
(torch.float16, torch.ops.aten.linalg_vector_norm.default): 1e-4,
(torch.bfloat16, torch.ops.aten.var_mean.correction): 5e-7,
(torch.float16, torch.ops.aten.var_mean.correction): 5e-7,
(torch.bfloat16, torch.ops.aten.var_mean.dim): 5e-7,
(torch.float16, torch.ops.aten.var_mean.dim): 5e-7,
(torch.float16, torch.ops.aten.nll_loss_forward.default): 1e-2,
(torch.bfloat16, torch.ops.aten.nll_loss_forward.default): 1e-1,
(torch.float16, torch.ops.aten.nll_loss2d_forward.default): 1e-2,
(torch.bfloat16, torch.ops.aten.nll_loss2d_forward.default): 2e-1,
(torch.float16, torch.ops.aten.hardswish.default): 2e-7,
(torch.bfloat16, torch.ops.aten.hardswish.default): 2e-7,
(torch.float16, torch.ops.aten.multi_margin_loss.default): 3e-2,
(torch.bfloat16, torch.ops.aten.multi_margin_loss.default): 5e-2,
(torch.float16, torch.ops.aten.multilabel_margin_loss_forward.default): 3e-2,
(torch.bfloat16, torch.ops.aten.multilabel_margin_loss_forward.default): 3e-2,
(torch.float16, torch.ops.aten.reflection_pad1d_backward.default): 5e-3,
(torch.bfloat16, torch.ops.aten.reflection_pad1d_backward.default): 5e-3,
(torch.float16, torch.ops.aten.reflection_pad2d_backward.default): 5e-3,
(torch.bfloat16, torch.ops.aten.reflection_pad2d_backward.default): 5e-3,
(torch.float16, torch.ops.aten.reflection_pad3d_backward.default): 5e-3,
(torch.bfloat16, torch.ops.aten.reflection_pad3d_backward.default): 5e-2,
# see https://github.com/pytorch/pytorch/pull/96264
(torch.float16, torch.ops.aten.mv.default): 1e-5,
(torch.bfloat16, torch.ops.aten.mv.default): 1e-5,
(torch.float16, torch.ops.aten.log_sigmoid_backward.default): 2e-5,
(torch.float16, torch.ops.aten._softmax_backward_data.default): 3e-7,
}
if ref.is_floating_point():
orig_diff = (orig - ref).abs().max()
decomp_diff = (decomp - ref).abs().max()
atol = tol_table.get((test_dtype, op), 1e-7)
if decomp_diff > orig_diff + atol:
raise RuntimeError(
f"Difference from float64 is larger with decomposition {op.__name__}"
f" than original on output {i}. Original max diff: {orig_diff}, Decomp max diff: {decomp_diff}\n"
f"atol = {atol}\n"
f"args = {args}\n"
f"kwargs = {kwargs}"
)
else:
test_case.assertEqual(
orig, decomp, msg=f"{op.__name__}\nargs = {args}\nkwargs = {kwargs}"
)
def op_assert_equal(test_case, op, test_dtype, orig, decomp, args, kwargs):
test_case.assertEqual(
orig.dtype,
decomp.dtype,
f"Operation: {op}, orig.dtype: {orig.dtype}, decomp.dtype: {decomp.dtype}, {args}, {kwargs}",
)
# Before adding an entry to this table, make sure your decomposition is right :)
tol_table = {
# Due to strange epsilon behaviors, see https://github.com/pytorch/pytorch/issues/73161
(torch.float32, torch.ops.aten.native_layer_norm.default): (1e-3, 1e-3),
(torch.float32, torch.ops.aten.native_layer_norm_backward.default): (
1e-3,
1e-3,
),
(torch.float64, torch.ops.aten.native_layer_norm.default): (1e-6, 1e-6),
# This exceeds default tolerances only on CPU, on CUDA it's fine
(torch.float32, torch.ops.aten.grid_sampler_2d.default): (7e-6, 3e-5),
# Exceeds tolerances on CUDA, likely due to fma
(torch.float32, torch.ops.aten.mv.default): (1e-5, 3e-5),
(torch.complex64, torch.ops.aten.mv.default): (5e-5, 5e-5),
(torch.float64, torch.ops.aten.upsample_bicubic2d.vec): (1e-5, 5e-4),
(torch.float64, torch.ops.aten.upsample_bicubic2d.default): (1e-5, 5e-4),
# The decomposition is TOO correct. It computes everything in int64, so sometimes
# there's an off-by-one error. See
# https://github.com/pytorch/pytorch/issues/81996
# https://github.com/pytorch/pytorch/issues/82230
(torch.int8, torch.ops.aten.linspace.default): (0, 1),
(torch.uint8, torch.ops.aten.linspace.default): (0, 1),
(torch.int16, torch.ops.aten.linspace.default): (0, 1),
(torch.int32, torch.ops.aten.linspace.default): (0, 1),
(torch.int64, torch.ops.aten.linspace.default): (0, 1),
(torch.int8, torch.ops.aten.linspace.Tensor_Tensor): (0, 1),
(torch.uint8, torch.ops.aten.linspace.Tensor_Tensor): (0, 1),
(torch.int16, torch.ops.aten.linspace.Tensor_Tensor): (0, 1),
(torch.int32, torch.ops.aten.linspace.Tensor_Tensor): (0, 1),
(torch.int64, torch.ops.aten.linspace.Tensor_Tensor): (0, 1),
(torch.int8, torch.ops.aten.linspace.Tensor_Scalar): (0, 1),
(torch.uint8, torch.ops.aten.linspace.Tensor_Scalar): (0, 1),
(torch.int16, torch.ops.aten.linspace.Tensor_Scalar): (0, 1),
(torch.int32, torch.ops.aten.linspace.Tensor_Scalar): (0, 1),
(torch.int64, torch.ops.aten.linspace.Tensor_Scalar): (0, 1),
(torch.int8, torch.ops.aten.linspace.Scalar_Tensor): (0, 1),
(torch.uint8, torch.ops.aten.linspace.Scalar_Tensor): (0, 1),
(torch.int16, torch.ops.aten.linspace.Scalar_Tensor): (0, 1),
(torch.int32, torch.ops.aten.linspace.Scalar_Tensor): (0, 1),
(torch.int64, torch.ops.aten.linspace.Scalar_Tensor): (0, 1),
}
if (decomp.dtype, op) in tol_table:
rtol, atol = tol_table[(decomp.dtype, op)]
else:
rtol, atol = _getDefaultRtolAndAtol(orig.dtype, decomp.dtype)
test_case.assertEqual(
orig,
decomp,
rtol=rtol,
atol=atol,
msg=f"{op.__name__}\nargs = {args}\nkwargs = {kwargs}",
)
# Given f, returns an f' such that:
# - f' takes only positional arguments
# - All arguments to f' are floating-point Tensors
# - All outputs of f' are floating-point Tensors
def normalize_op_input_output2(
f, args, kwargs, output_process_fn_grad=None, requires_grad=True
):
flat_args, args_spec = tree_flatten(args)
diff_argnums = tuple(
i
for i, arg in enumerate(flat_args)
if diff_arg(arg, requires_grad=requires_grad)
)
assert len(diff_argnums) > 0
primals = tuple(flat_args[i] for i in diff_argnums)
@functools.wraps(f)
def wrapped(*primals):
_args = list(flat_args)
for num, arg in zip(diff_argnums, primals):
_args[num] = arg
_args = tree_unflatten(_args, args_spec)
result = f(*_args, **kwargs)
if output_process_fn_grad is not None:
result = output_process_fn_grad(result)
if isinstance(result, tuple):
# TODO We should check that the integer outputs also agree
result = tuple(
r
for r in result
if isinstance(r, Tensor) and (r.is_floating_point() or r.is_complex())
)
assert len(result) > 0
return result
return wrapped, primals
# NB: This also upcasts dtype arguments
# TODO: handle complex correctly
def upcast_tensor(x, dtype=torch.float32):
if isinstance(x, Tensor) and x.dtype.is_floating_point:
return x.to(dtype=dtype)
elif isinstance(x, torch.dtype) and x in [
torch.float16,
torch.bfloat16,
torch.float,
]:
return dtype
else:
return x
def normalize_op_input_output(f, sample, requires_grad=True):
args = tuple([sample.input] + list(sample.args))
return normalize_op_input_output2(
f,
args,
sample.kwargs,
sample.output_process_fn_grad,
requires_grad=requires_grad,
)
CROSS_REF_EXCLUDE_SET = {
# CUBLAS_STATUS_NOT_SUPPORTED when calling
# `cublasGemmStridedBatchedExFix(handle, opa, opb, (int)m, (int)n, (int)k,
# (void*)&falpha, a, CUDA_R_16BF, (int)lda, stridea, b, CUDA_R_16BF,
# (int)ldb, strideb, (void*)&fbeta, c, CUDA_R_16BF, (int)ldc, stridec,
# (int)num_batches, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)`
("cuda", torch.bfloat16, "nn.functional.bilinear"),
# randomness
(None, None, "special.ndtr"), # aten.special_ndtr was not decomposed
(None, None, "new_empty"),
(None, None, "empty_like"),
(None, None, "empty"),
# AssertionError: False is not true : aten.item was not decomposed, saw calls for: aten._local_scalar_dense.default.
(None, None, "item"),
# It's the only in-place op without an out-of-place equivalent in the Python API
# Its OpInfo wrongly registers it as `torch.zero_(x.clone())`.
(None, None, "zero_"),
# No idea what's going on here
# In the recursive test logsumexp.default fails with args = (torch.tensor(-math.inf), [])
# in the test, but it seems to pass when tested locally and in the logsumexp test
(None, torch.float32, "masked.logsumexp"),
(None, torch.float64, "masked.logsumexp"),
# exp_vml_cpu not implemented for Half
(torch.cpu, torch.float16, "signal.windows.exponential"),
(torch.cpu, torch.float16, "signal.windows.gaussian"),
# sin_vml_cpu not implemented for Half
(torch.cpu, torch.float16, "signal.windows.cosine"),
# CompositeAutogradImplicit
# See https://github.com/pytorch/pytorch/issues/81669
(None, None, "nn.functional.relu6"),
# This decomp runs before autograd.
(None, None, "nn.functional.rrelu"),
(None, None, "meshgrid"),
# Decomposition registered as Autograd
(None, None, "nn.functional.hardshrink"),
(None, None, "nn.functional.softshrink"),
# diag was not decomposed (it just registers a decomp for diag_out, torch.diag is CompImplicit)
(None, None, "diag"),
# _softmax_backward_data's CPU kernel for bfloat16 always return the grad_input as float32
("cpu", torch.bfloat16, "_softmax_backward_data"),
(None, None, "norm"),
# native_batch_norm is only implicit when python dispatcher is on (and noncomposite otherwise)
(None, None, "native_batch_norm"),
(None, None, "_upsample_bilinear2d_aa"),
(None, None, "empty_strided"), # aten.empty_strided was not decomposed
}
CROSS_REF_BACKWARD_EXCLUDE_SET = {
# Decomposed backward formula is not as precise
("cpu", torch.bfloat16, "nn.functional.hardswish"),
("cuda", torch.float16, "nn.functional.cross_entropy"),
}
all_decomposed = set()
all_called = defaultdict(int)
# Helpful snippet for testing coverage
"""
import atexit
def check_coverage():
print("missing coverage:")
print("\n".join(map(str, decomposition_table.keys() - all_decomposed)))
atexit.register(check_coverage)
"""
# Helpful snippet for Horace to create his google sheet :)
"""
import atexit
def dump_ops():
with open('run_ops.txt', 'w') as f, open('count_ops.txt', 'w') as g:
for op, count in sorted(all_called.items(), key=lambda x: x[0].__name__):
f.write(f'{op.__name__}\n')
g.write(f'{count}\n')
with open('run_decompositions.txt', 'w') as f:
for op in sorted([i.__name__ for i in all_decomposed]):
f.write(f'{op}\n')
atexit.register(dump_ops)
"""
def any_unsupported(args, kwargs):
def test_unsupported(t):
if type(t) is torch.Tensor or type(t) is torch.nn.Parameter:
# These are all things that we haven't coded decompositions
# to handle correctly. Maybe they should.
return any(
[
t.is_sparse_csr,
t.is_sparse,
t.is_mkldnn,
t.is_quantized,
t.is_nested,
torch._is_functional_tensor(t),
]
)
elif torch.overrides.is_tensor_like(t):
# Decompositions will generally change the behavior of Tensor-like
# subclasses, so bypass tests in this case too
return True
else:
return False
flat_args = pytree.arg_tree_leaves(*args, **kwargs)
return any(test_unsupported(x) for x in flat_args)
core_backward_failures = {
skip("_softmax_backward_data"), # slow: fails with --timeout=360 secs
xfail("addcdiv"),
skip("addcmul"), # slow: fails with --timeout=360 secs
skip("deg2rad"), # slow: fails with --timeout=360 secs
skip("diag_embed"), # slow: fails with --timeout=360 secs
skip("frac"), # slow: fails with --timeout=360 secs
skip("grid_sampler_2d"), # slow: fails with --timeout=360 secs
xfail("lerp"),
skip("logaddexp"), # slow: fails with --timeout=360 secs
skip("native_dropout_backward"), # slow: fails with --timeout=360 secs
xfail("nn.functional.binary_cross_entropy_with_logits"),
skip("nn.functional.glu"), # slow: fails with --timeout=360 secs
xfail("nn.functional.hardshrink"),
xfail("nn.functional.softshrink"),
skip("nn.functional.unfold"), # slow: fails with --timeout=360 secs
xfail("norm"),
xfail("norm", "fro"),
xfail("norm", "inf"),
xfail("norm", "nuc"),
skip("rad2deg"), # slow: fails with --timeout=360 secs
skip("renorm"), # slow: fails with --timeout=360 secs
skip("rot90"), # slow: fails with --timeout=360 secs
skip("rsub"), # slow: fails with --timeout=360 secs
skip("sgn"), # slow: fails with --timeout=360 secs
skip("special.xlog1py"), # slow: fails with --timeout=360 secs
xfail("stack"),
skip("tril"), # slow: fails with --timeout=360 secs
skip("triu"), # slow: fails with --timeout=360 secs
skip("unfold_copy"), # slow: fails with --timeout=360 secs
skip("xlogy"), # slow: fails with --timeout=360 secs
xfail("zero_"),
}
if not TEST_WITH_SLOW:
core_backward_failures.update(
{
skip("addr"), # slow: takes 46 sec on A100
skip("baddbmm"), # slow: takes 800+ sec on A100
skip("clamp_min"), # slow: takes 800 sec on A100
skip("clamp_max"), # slow: takes 800 sec on A100
skip("logit"), # slow: takes 44 sec on A100
skip("nn.functional.hardswish"), # slow: takes 60 sec on A100
skip("std_mean"), # slow: takes 170 sec on A100
skip("split", variant_name="list_args"), # slow: takes 118 sec on A100
skip("transpose"), # slow: takes 50 sec on A100
skip("unbind"), # slow: takes 70 sec on A100
skip("unsafe_split"), # slow: takes 49 sec on A100
}
)
comprehensive_failures = {
xfail(
"nn.functional.interpolate", "bilinear", dtypes=(torch.uint8,)
), # off by one error
xfail(
"nn.functional.interpolate", "bicubic", dtypes=(torch.uint8,)
), # off by one error
xfail(
"nn.functional.upsample_bilinear", "", dtypes=(torch.uint8,)
), # off by one error
}
@unMarkDynamoStrictTest
class TestDecomp(TestCase):
longMessage = True
# NB: This actually overlaps with test_comprehensive, but it only
# runs on things that are definitely decomposed so it's a lot faster
# to run
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@onlyNativeDeviceTypes
@skipIfCrossRef
@suppress_warnings
@ops(_decomp_test_ops)
def test_quick(self, device, dtype, op):
self.do_cross_ref(device, dtype, op, run_all=False)
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@skipOps("TestDecomp", "test_quick_core_backward", core_backward_failures)
@onlyNativeDeviceTypes
@skipIfCrossRef
@suppress_warnings
@ops(_decomp_test_ops_core_autograd, allowed_dtypes=(torch.float64,))
def test_quick_core_backward(self, device, dtype, op):
for sample_input in op.sample_inputs(device, dtype, requires_grad=True):
aten_name = op.decomp_aten_name or op.aten_name
args = [sample_input.input] + list(sample_input.args)
kwargs = sample_input.kwargs
func = partial(op.get_op(), **kwargs)
with self.DecompCrossRefMode(
self, self.precision, self.rel_tol, dtype, run_all=False
) as mode, enable_python_dispatcher():
torch.autograd.gradcheck(func, args)
self.check_decomposed(aten_name, mode)
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@onlyNativeDeviceTypes
@skipIfCrossRef
@skipOps("TestDecomp", "test_comprehensive", comprehensive_failures)
@suppress_warnings
@ops(op_db)
def test_comprehensive(self, device, dtype, op):
self.do_cross_ref(device, dtype, op, run_all=True)
def test_uniform(self, device):
size = (2, 3, 4, 5)
dtype = torch.float32
x = make_tensor(size, dtype=dtype, device=device)
low = 0.3
high = 0.9
torch.manual_seed(123)
ref = torch.ops.aten.uniform(x, low, high)
torch.manual_seed(123)
res = torch._decomp.decompositions.uniform(x, low=low, high=high)
self.assertEqual(ref, res)
def test_broadcasting_index_copy(self, device):
x = torch.zeros([1, 10], device=device)
xs = torch.ones([2, 10], device=device)
def index_copy(xs, x):
torch._decomp.decompositions.index_copy_(
xs, 0, torch.tensor(0).to(device), x
)
index_copy(xs, x)
xs_two = torch.ones([2, 10], device=device)
xs_two[0] = x
self.assertEqual(xs, xs_two)
def test_cat_single_input(self, device):
decomp_table = torch._inductor.decomposition.select_decomp_table()
cat_inductor = decomp_table[torch.ops.aten.cat.default]
inp = torch.rand([2048, 2048], device=device)
inps = [inp for _ in range(10)]
for dim in (-1, 0, 1):
self.assertEqual(torch.cat(inps, dim), cat_inductor(inps, dim))
def test_rrelu_with_noise(self, device):
# rrelu_with_noise behavior depends on a) whether elements in the input
# are <= 0, and b) whether we're in training mode. Cover all cases:
dtype = torch.float64
x = torch.tensor([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0], dtype=dtype, device=device)
lower = 1.0
upper = 4.0
training = False
torch.manual_seed(123)
noise_ref = torch.zeros(x.shape, dtype=dtype, device=device)
ref = torch.ops.aten.rrelu_with_noise(x, noise_ref, lower, upper, training)
torch.manual_seed(123)
noise_res = torch.zeros(x.shape, dtype=dtype, device=device)
res = torch._decomp.decompositions.rrelu_with_noise(
x,
noise_res,
lower,
upper,
training,
)
self.assertEqual(ref, res)
self.assertEqual(noise_ref, noise_res)
# Now with training=True:
training = True
torch.manual_seed(123)
noise_ref = torch.zeros(x.shape, dtype=dtype, device=device)
ref = torch.ops.aten.rrelu_with_noise(x, noise_ref, lower, upper, training)
torch.manual_seed(123)
noise_res = torch.zeros(x.shape, dtype=dtype, device=device)
res = torch._decomp.decompositions.rrelu_with_noise(
x,
noise_res,
lower,
upper,
training,
)
self.assertEqual(ref, res)
self.assertEqual(noise_ref, noise_res)
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@suppress_warnings
@tf32_off()
# only tests RNNs since we have py dispsatcher decomps for them
@modules(
filter(
lambda m: m.module_cls in (torch.nn.RNN, torch.nn.LSTM, torch.nn.GRU),
module_db,
)
)
def test_rnn_decomp_module(self, device, dtype, module_info, training):
module_cls = module_info.module_cls
module_inputs = module_info.module_inputs_func(
module_info,
device=device,
dtype=dtype,
requires_grad=True,
training=training,
)
for module_input in module_inputs:
if module_input.forward_input is None:
continue
args, kwargs = (
module_input.constructor_input.args,
module_input.constructor_input.kwargs,
)
m = module_cls(*args, **kwargs)
m.to(device).to(dtype)
args, kwargs = (
module_input.forward_input.args,
module_input.forward_input.kwargs,
)
with self.DecompCrossRefMode(
self, self.precision, self.rel_tol, dtype, run_all=True
), enable_python_dispatcher():
decomp_out = m(*args, **kwargs)
non_decomp_out = m(*args, **kwargs)
# without this check, incorrect decomps at the python dispatcher level can still pass because
# they're checking aten decomps at the torch_dispatch level
self.assertEqual(decomp_out, non_decomp_out)
def test_batch_norm_unflatten_weight_bias(self, device):
# https://github.com/pytorch/pytorch/issues/100970
shape = (1, 3, 2, 2)
input = torch.randn(shape, device=device)
weight = torch.randn((3, 1, 1, 1), device=device)
bias = torch.randn(3, device=device)
mean = torch.randn(3, device=device)
var = torch.randn(3, device=device)
res = torch._decomp.decompositions.native_batch_norm(
input, weight, bias, mean, var, False, 1, 1e-05
)
self.assertEqual(shape, res[0].shape)
def test_arange_graph(self, device):
from torch.fx.experimental.proxy_tensor import make_fx
def func(x, start):
le = x.shape[-1]
if start is None:
a = torch.arange(le, dtype=torch.float32, device=x.device)
else:
a = torch.arange(start, le, dtype=torch.float32, device=x.device)
return a
pattern = r", device = device\(.+\), requires_grad = False"
cfunc = make_fx(func, decomposition_table=decomposition_table)
fx_g = cfunc(torch.rand(10, device=device), None)
fx_g_code = fx_g.code.strip()
# Remove device and requires_grad
fx_g_code = re.sub(pattern, "", fx_g_code)
self.assertExpectedInline(
fx_g_code,
"""\
def forward(self, x_1, start_1):
iota = torch.ops.prims.iota.default(10, start = 0, step = 1, dtype = torch.int64)
mul = torch.ops.prims.mul.default(iota, 1); iota = None
add = torch.ops.prims.add.default(mul, 0); mul = None
convert_element_type = torch.ops.prims.convert_element_type.default(add, torch.float32); add = None
return convert_element_type""",
)
fx_g = cfunc(torch.rand(10, device=device), 1)
fx_g_code = fx_g.code.strip()
# Remove device and requires_grad
fx_g_code = re.sub(pattern, "", fx_g_code)
self.assertExpectedInline(
fx_g_code,
"""\
def forward(self, x_1, start_1):
iota = torch.ops.prims.iota.default(9, start = 0, step = 1, dtype = torch.int64)
mul = torch.ops.prims.mul.default(iota, 1); iota = None
add = torch.ops.prims.add.default(mul, 1); mul = None
convert_element_type = torch.ops.prims.convert_element_type.default(add, torch.float32); add = None
return convert_element_type""",
)
def test_masked_fill(self, device):
from torch.fx.experimental.proxy_tensor import make_fx
if torch.device(device).type not in [
"xpu",
"cuda",
torch._C._get_privateuse1_backend_name(),
]:
self.skipTest("only runs on XPU and CUDA and PrivateUse1.")
def func(scores, mask, value):
return scores.masked_fill(mask, value)
scores_t = torch.tensor([1, 2, 3, 4], device=device)
mask_t = torch.tensor([True, True, True, True], device=device)
value_t = torch.tensor(0, dtype=scores_t.dtype)
cfunc = make_fx(func, decomposition_table=decomposition_table)
fx_g = cfunc(scores_t, mask_t, value_t)
self.assertExpectedInline(
fx_g.code.strip(),
"""\
def forward(self, scores_1, mask_1, value_1):
where = torch.ops.prims.where.default(mask_1, value_1, scores_1); mask_1 = value_1 = scores_1 = None
return where""",
)
class DecompCrossRefMode(TorchDispatchMode):
def __init__(self, test_case, saved_precision, saved_rel_tol, dtype, run_all):
self.test_case = test_case
self.saved_precision = saved_precision
self.saved_rel_tol = saved_rel_tol
self.test_dtype = dtype
self.run_all = run_all
# We check the correctness of each decomposition right after running it.
# So, when we encounter a decomposition, we run the function normally, and
# then run the decomposition, and ensure they're identical.
self.called = set()
self.decomposed = set()
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
self.test_case.precision = self.saved_precision
self.test_case.rel_tol = self.saved_rel_tol
self.called.add(func)
all_called[func] += 1
# Stuff we shouldn't bother testing
# (TODO: remove detach from the decomp table?)
# N.b. Testing in-place ops would need dedicated logic
in_place = func.name()[-1] == "_"
ignored_ops = [
torch.ops.aten.detach.default,
# non-deterministic ops
torch.ops.aten.empty.memory_format,
torch.ops.aten.empty_like.default,
torch.ops.aten.new_empty.default,
torch.ops.aten.empty_strided.default,
torch.ops.aten.new_empty_strided.default,
torch.ops.aten.randn.default,
torch.ops.aten.native_dropout.default,
]
if (
func not in decomposition_table
or func in ignored_ops
or torch.Tag.nondeterministic_seeded in func.tags
or any_unsupported(args, kwargs)
or in_place
):
return func(*args, **kwargs)
self.decomposed.add(func)
all_decomposed.add(func)
# We take 2 main strategies for verifying correctness/numerical stability of decompositions
# The first one is simply tolerance checking between decomp_out and pytorch_out
# However, for fp16/bf16 and reductions, this becomes very
# finicky, as there are not many guarantees we can make.
# So, for fp16/bf16, we instead compare the difference of
# {decomp_out, pytorch_out_64} and {pytorch_out,
# pytorch_out_64}. In other words, we compare how far the
# decomposition and pytorch are from the "ground truth" (i.e.
# fp64). If the decomposition results in more error, we error
# We also decompose the decomposition recursively for
# further coverage, as some paths not be exercised directly by
# OpInfos (sadly) but just by other ops
decomposition = decomposition_table[func]
do_relative_check = self.test_dtype in [torch.float16, torch.bfloat16]
if self.run_all:
# Execute recursively via DFS, to find the root of a possible error first
with self:
decomp_out = pytree.tree_leaves(decomposition(*args, **kwargs))
else:
decomp_out = pytree.tree_leaves(decomposition(*args, **kwargs))
# At this stage we should not be decomposing an in-place op
# We'd like to have decompositions that decompose out-of-place ops into out-of-place ops
# because decompositions are run after functionalisation and we would not like them to
# de-functionalise the graph, as that would break AoTAutograd
# We run the real function *after* the decomposition to make sure that the
# decomposition does not modify any of the inputs in-place. If it does
# real_out should be differen than decom_out so we should catch this
real_out_unflat = func(*args, **kwargs)
real_out = pytree.tree_leaves(real_out_unflat)
assert len(real_out) == len(decomp_out)
if do_relative_check:
upcast = partial(upcast_tensor, dtype=torch.float64)
real_out_double, _ = tree_flatten(
func(*tree_map(upcast, args), **tree_map(upcast, kwargs))
)
for i, (orig, decomp, ref) in enumerate(
zip(real_out, decomp_out, real_out_double)
):
if not isinstance(orig, torch.Tensor):
assert type(orig) == type(decomp)
assert orig == decomp
continue
op_assert_ref(
self.test_case,
func,
self.test_dtype,
i,
orig,
decomp,
ref,
args,
kwargs,
)
else:
for orig, decomp in zip(real_out, decomp_out):
if not isinstance(orig, torch.Tensor):
assert type(orig) == type(decomp)
assert orig == decomp
continue
op_assert_equal(
self.test_case,
func,
self.test_dtype,
orig,
decomp,
args,
kwargs,
)
return real_out_unflat
def check_decomposed(self, aten_name, mode):
self.assertTrue(
any(overload_to_aten_name(c) == aten_name for c in mode.decomposed),
msg=(
f"aten.{aten_name} was not decomposed, saw calls for: "
f"{', '.join(map(str, list(mode.called)))}. If your op is "
f"CompositeImplicitAutograd you should skip this test "
f"by updating CROSS_REF_EXCLUDE_SET."
),
)
@skipIfTorchDynamo("Test does not work with TorchDynamo")
def do_cross_ref(self, device, dtype, op, *, run_all):
test_keys = [
(torch.device(device).type, dtype, op.name),
(None, dtype, op.name),
(None, None, op.name),
]
if any(key in CROSS_REF_EXCLUDE_SET for key in test_keys):
self.skipTest(f"{op.name} in {dtype} not supported")
skip_decomp_vjp = any(
key in CROSS_REF_BACKWARD_EXCLUDE_SET for key in test_keys
)
requires_grad = (
op.supports_autograd
and dtype in op.supported_backward_dtypes(torch.device(device).type)
# TODO: OpInfo really ought to error out for this case, but it's
# not exercised in test_ops_gradients atm. The problem is not
# complex32 per-se (which is supported by data movement only ops)
# but that when we do backwards we expect other ops like add to work
and not dtype == torch.complex32
)
samples = op.sample_inputs(device, dtype, requires_grad=requires_grad)
aten_name = op.decomp_aten_name or op.aten_name
func = op.get_op()
def run_without_python_dispatcher(mode):
return any(
isinstance(op, torch._ops.OpOverload)
and op.has_kernel_for_dispatch_key(
DispatchKey.CompositeImplicitAutograd
)
for op in mode.decomposed.union([func])
)
for sample_input in samples:
if requires_grad:
fn, primals = normalize_op_input_output(func, sample_input)
primals = tree_map(
lambda x: x if isinstance(x, torch.Tensor) else x, primals
)
# Once https://github.com/pytorch/pytorch/pull/75965/ I can
# store the called list on the mode object instance and no
# explicit clearing is necessary as I will create a fresh mode
# for each region
with self.DecompCrossRefMode(
self, self.precision, self.rel_tol, dtype, run_all
) as mode, enable_python_dispatcher():
decomp_out, decomp_vjp_fn = ref_vjp_no_create(fn, *primals)
if run_without_python_dispatcher(mode):
# without this check, incorrect decomps at the python dispatcher level can still pass because
# they're checking aten decomps at the torch_dispatch level.
with self.DecompCrossRefMode(
self, self.precision, self.rel_tol, dtype, run_all
) as mode:
decomp_out, decomp_vjp_fn = ref_vjp_no_create(fn, *primals)
if aten_name in decomposition_names:
self.check_decomposed(aten_name, mode)
if not skip_decomp_vjp and (
op.aten_backward_name in decomposition_names or run_all
):
cotangents = tree_map(lambda x: torch.randn_like(x), decomp_out)
with self.DecompCrossRefMode(
self, self.precision, self.rel_tol, dtype, run_all
) as mode, enable_python_dispatcher():
decomp_vjp_fn(cotangents)
if run_without_python_dispatcher(mode):
# without this check, incorrect decomps at the python dispatcher level can still pass because
# they're checking aten decomps at the torch_dispatch level.
with self.DecompCrossRefMode(
self, self.precision, self.rel_tol, dtype, run_all
) as mode:
decomp_vjp_fn(cotangents)
if not run_all:
self.check_decomposed(op.aten_backward_name, mode)
elif aten_name in decomposition_names or run_all:
args = [sample_input.input] + list(sample_input.args)
kwargs = sample_input.kwargs
# A failure here might be because the decomposition for the op is wrong or because a
# decomposition used by the particular op is wrong.
with self.DecompCrossRefMode(
self, self.precision, self.rel_tol, dtype, run_all
) as mode, enable_python_dispatcher():
func(*args, **kwargs)
if run_without_python_dispatcher(mode):
# without this check, incorrect decomps at the python dispatcher level can still pass because
# they're checking aten decomps at the torch_dispatch level.
with self.DecompCrossRefMode(
self, self.precision, self.rel_tol, dtype, run_all
) as mode:
func(*args, **kwargs)
if not run_all:
self.check_decomposed(aten_name, mode)
else:
assert op.supports_autograd
self.skipTest(
"only backwards is decomposed, but dtype doesn't support AD"
)
instantiate_device_type_tests(TestDecomp, globals())
class DecompOneOffTests(TestCase):
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@onlyNativeDeviceTypes
@skipIfCrossRef
def test_contiguous_softmax(self, device):
size = (2, 4, 3, 3)
stride = (9, 18, 3, 1)
dtype = torch.float32
x = torch.randn(size, dtype=dtype, device=device)
x = torch.as_strided(x, size, stride)
ref = torch.ops.aten._softmax(x, -1, False)
res = torch._decomp.decompositions._softmax(x, -1, False)
self.assertEqual(ref.stride(), res.stride())
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@onlyNativeDeviceTypes
@skipIfCrossRef
def test_contiguous_log_softmax(self, device):
size = (2, 4, 3, 3)
stride = (9, 18, 3, 1)
dtype = torch.float32
x = torch.randn(size, dtype=dtype, device=device)
x = torch.as_strided(x, size, stride)
ref = torch.ops.aten._log_softmax(x, -1, False)
res = torch._decomp.decompositions._log_softmax(x, -1, False)
self.assertEqual(ref.stride(), res.stride())
@onlyCUDA
def test_exponential_non_inf(self, device):
inp = torch.empty((4, 400, 256), device=device)
with torch._dynamo.utils.preserve_rng_state():
exp_ref = inp.exponential_()
exp = torch._refs.exponential(inp)
self.assertEqual(exp, exp_ref)
self.assertFalse(exp.isinf().any())
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@skipIfCrossRef
@onlyCUDA
def test_amp_batch_norm_backward(self):
device = "cuda"
grad_out = torch.randn((1, 2, 16, 16), dtype=torch.float16, device=device)
x = torch.randn((1, 2, 16, 16), dtype=torch.float16, device=device)
weight = torch.randn((2,), dtype=torch.float32, device=device)
rmean = torch.randn((2,), dtype=torch.float32, device=device)
rvar = torch.randn((2,), dtype=torch.float32, device=device)
mean = torch.randn((0,), dtype=torch.float32, device=device)
ref = torch.ops.aten.native_batch_norm_backward(
grad_out,
x,
weight,
rmean,
rvar,
mean,
mean,
False,
1e-05,
[True, True, True],
)
res = torch._decomp.decompositions.native_batch_norm_backward(
grad_out,
x,
weight,
rmean,
rvar,
mean,
mean,
False,
1e-05,
[True, True, True],
)
for a, b in zip(ref, res):
self.assertEqual(a.stride(), b.stride())
self.assertEqual(a.dtype, b.dtype)
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@onlyNativeDeviceTypes
@skipIfCrossRef
def test_elu_backward(self, device):
size = (2, 4, 3, 3)
dtype = torch.float32
grad_out = torch.randn(size, dtype=dtype, device=device)
out = torch.randn(size, dtype=dtype, device=device)
ref = torch.ops.aten.elu_backward(grad_out, 1.0, 1, 1, True, out)
res = torch._decomp.decompositions.elu_backward(grad_out, 1.0, 1, 1, True, out)
self.assertEqual(ref, res)
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@onlyNativeDeviceTypes
@skipIfCrossRef
def test_threshold_backward_dtype(self, device):
grad = torch.randint(10, (4,), device=device)
input_tensor = torch.randint(10, (4,), device=device)
ref = torch.ops.aten.threshold_backward(grad, input_tensor, 1)
res = torch._decomp.decompositions.threshold_backward(grad, input_tensor, 1)
self.assertEqual(ref.dtype, res.dtype)
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@onlyNativeDeviceTypes
@skipIfCrossRef
def test_weight_norm_interface(self, device):
g = torch.randn((3, 10, 10), device=device)
v = torch.randn((1, 1, 10), device=device)
ref = torch.ops.aten._weight_norm_interface(g, v, 2)
res = torch._decomp.decompositions._weight_norm_interface(g, v, 2)
self.assertTrue(torch.allclose(ref[0], res[0]))
self.assertTrue(torch.allclose(ref[1], res[1]))
inp = torch.rand([30, 10], device=device)
inp2 = torch.rand([30, 1], device=device)
self.assertEqual(
torch.ops.aten._weight_norm_interface(inp, inp2),
torch._decomp.decompositions._weight_norm_interface(inp, inp2),
)
@unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN")
@onlyCPU
@skipIfCrossRef
@skipOps(
"DecompOneOffTests",
"test_sdpa",
[
xfail(
"nn.functional.scaled_dot_product_attention",
dtypes=[torch.half],
),
],
)
@ops(_sdpa_op_info)
def test_sdpa(self, device, dtype, op):
# SDPA doesn't support float16, this is aligned with aten/src/ATen/native/transformers/attention.cpp. If we
# add support for float16 over there we should update this test as well.
class ScaledDotProductAttention(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
def forward(
self, query_layer, key_layer, value_layer, mask=None, is_causal=True
):
attn_output = op(
query_layer,
key_layer,
value_layer,
attn_mask=mask,
dropout_p=0.0,
is_causal=is_causal,
)
return attn_output
query_layer = torch.randn(1, 128, 100, 64, device=device, dtype=dtype)
key_layer = torch.randn(1, 128, 100, 64, device=device, dtype=dtype)
value_layer = torch.randn(1, 128, 100, 64, device=device, dtype=dtype)
masks = [None, torch.ones((1, 1, 100, 100), device=device, dtype=torch.bool)]
atol, rtol = dtype_precisions[dtype]
for mask in masks:
is_causal = mask is None
attention = ScaledDotProductAttention()
decomposed_res = (
torch._decomp.decompositions.scaled_dot_product_flash_attention_for_cpu(
query_layer, key_layer, value_layer, 0.0, is_causal, attn_mask=mask
)
)
eager_res = op(
query_layer,
key_layer,
value_layer,
attn_mask=mask,
dropout_p=0.0,
is_causal=is_causal,
)
self.assertTrue(
torch.allclose(decomposed_res[0], eager_res, atol=atol, rtol=rtol)
)
instantiate_device_type_tests(DecompOneOffTests, globals())
class HasDecompTest(TestCase):
def setUp(self):
super().setUp()
self.maxDiff = None
@staticmethod
def _can_appear_in_trace(op: torch._ops.OpOverload) -> bool:
has_tensor_arg = any(
"Tensor" in str(a.type)
for a in itertools.chain(op._schema.arguments, op._schema.returns)
)
if not has_tensor_arg:
return False
try:
# CompositeImplicitAutograd ops are transparent to the tracer, so don't need decompositions
return not op.has_kernel_for_dispatch_key(
DispatchKey.CompositeImplicitAutograd
)
except RuntimeError as e:
# has_key fails for some jit-registered ops, which shouldn't be
# relevant here anyway
if "does not exist" in str(e):
return False
raise
def test_has_decomposition(self):
def all_aten_overloads():
for name in torch._C._dispatch_get_all_op_names():
if not name.startswith("aten::"):
continue
name = name[6:]
if "." in name:
packet_name, overload_name = name.split(".")
else:
packet_name, overload_name = name, "default"
packet = getattr(aten, packet_name)
assert isinstance(packet, torch._ops.OpOverloadPacket)
op = getattr(packet, overload_name)
yield op
# This is for operators that are only registered in some CI
# configurations, so would cause the test to fail
allow_list = {aten.get_gradients.default}
overloads_wanting_decomp = {
op for op in all_aten_overloads() if self._can_appear_in_trace(op)
}
ops_missing_decomp = overloads_wanting_decomp - decomposition_table.keys()
ops_missing_decomp -= allow_list
self.assertExpected(
"".join(sorted(op.name() + "\n" for op in ops_missing_decomp))
)
def test_aten_core_operators(self):
# If a decomposition isn't included in the core decompositions,
# then it must decompose a core ATen operator.
#
# See NOTE [Core ATen Ops]
#
# If this test fails then either:
# - Add the decomposition to torch._decomp.core_aten_decompositions,
# if decomposition should be used by inductor (not a core operator).
# - Run this test again with EXPECTTEST_ACCEPT=1 to update the list of
# core ATen operators (and inductor will not use the decomposition).
# Some decompositions are registered for CompositeImplicitAutograd
# operators, which never appear in AOTAutograd's graph so are never used.
useful_decomps = {
op
for op in decomposition_table.keys()
if isinstance(op, torch._ops.OpOverload) and self._can_appear_in_trace(op)
}
core_decomps = torch._decomp.core_aten_decompositions().keys()
core_aten_ops = useful_decomps - core_decomps
self.assertExpected("".join(sorted(op.name() + "\n" for op in core_aten_ops)))
if __name__ == "__main__":
run_tests()
Reference in New Issue View Git Blame Copy Permalink