已合并
perf(inductor): optimize training operators #45392
luqichao创建于 10 天前
perf(inductor): optimize training operators #45392
已合并
共 2 个文件变更+114-0
| @@ -18,6 +18,114 @@ from .lowering_common import run_once | |||
| 18 | aten = torch.ops.aten | 18 | aten = torch.ops.aten |
| 19 | npu = torch.ops.npu | 19 | npu = torch.ops.npu |
| 20 | 20 | ||
| 21 | + | ||
| 22 | +def _matmul_backward_inductor(grad, self, other, mask): | ||
| 23 | + """Decompose matmul backward locally for the Triton Inductor backend.""" | ||
| 24 | + dim_self = self.dim() | ||
| 25 | + dim_other = other.dim() | ||
| 26 | + | ||
| 27 | + size_grad = grad.size() | ||
| 28 | + size_self = self.size() | ||
| 29 | + size_other = other.size() | ||
| 30 | + grad_self = None | ||
| 31 | + grad_other = None | ||
| 32 | + | ||
| 33 | + def matmul_backward_1d_1d(): | ||
| 34 | + nonlocal grad_self, grad_other | ||
| 35 | + grad_self = other.mul(grad) if mask[0] else grad_self | ||
| 36 | + grad_other = self.mul(grad) if mask[1] else grad_other | ||
| 37 | + return grad_self, grad_other | ||
| 38 | + | ||
| 39 | + def matmul_backward_2d_1d(): | ||
| 40 | + nonlocal grad_self, grad_other | ||
| 41 | + grad_self = grad.unsqueeze(1).mm(other.unsqueeze(0)) if mask[0] else grad_self | ||
| 42 | + grad_other = ( | ||
| 43 | + self.transpose(-1, -2).mm(grad.unsqueeze(1)).squeeze_(1) | ||
| 44 | + if mask[1] | ||
| 45 | + else grad_other | ||
| 46 | + ) | ||
| 47 | + return grad_self, grad_other | ||
| 48 | + | ||
| 49 | + def matmul_backward_1d_2d(): | ||
| 50 | + nonlocal grad_self, grad_other | ||
| 51 | + grad_self = ( | ||
| 52 | + grad.unsqueeze(0).mm(other.transpose(-1, -2)).squeeze_(0) | ||
| 53 | + if mask[0] | ||
| 54 | + else grad_self | ||
| 55 | + ) | ||
| 56 | + grad_other = self.unsqueeze(1).mm(grad.unsqueeze(0)) if mask[1] else grad_other | ||
| 57 | + return grad_self, grad_other | ||
| 58 | + | ||
| 59 | + def matmul_backward_nd_lt3d(): | ||
| 60 | + nonlocal grad_self, grad_other | ||
| 61 | + view_size = 1 if dim_other == 1 else size_grad[-1] | ||
| 62 | + unfolded_grad = ( | ||
| 63 | + (grad.unsqueeze(-1) if dim_other == 1 else grad) | ||
| 64 | + .contiguous() | ||
| 65 | + .view(-1, view_size) | ||
| 66 | + ) | ||
| 67 | + if mask[0]: | ||
| 68 | + unfolded_other = ( | ||
| 69 | + other.unsqueeze(0) if dim_other == 1 else other.transpose(-1, -2) | ||
| 70 | + ) | ||
| 71 | + grad_self = unfolded_grad.mm(unfolded_other).view(size_self) | ||
| 72 | + | ||
| 73 | + if mask[1]: | ||
| 74 | + unfolded_self = self.contiguous().view(-1, size_self[-1]) | ||
| 75 | + grad_other = ( | ||
| 76 | + unfolded_self.transpose(-1, -2).mm(unfolded_grad).view(size_other) | ||
| 77 | + ) | ||
| 78 | + return grad_self, grad_other | ||
| 79 | + | ||
| 80 | + def matmul_backward_lt3d_nd(): | ||
| 81 | + nonlocal grad_self, grad_other | ||
| 82 | + view_size = 1 if dim_self == 1 else size_grad[-2] | ||
| 83 | + unfolded_grad_t = ( | ||
| 84 | + grad.view(-1, view_size) | ||
| 85 | + if dim_self == 1 | ||
| 86 | + else grad.transpose(-1, -2).contiguous().view(-1, view_size) | ||
| 87 | + ) | ||
| 88 | + if mask[0]: | ||
| 89 | + unfolded_other_t = ( | ||
| 90 | + other.transpose(-1, -2) | ||
| 91 | + .contiguous() | ||
| 92 | + .view(-1, size_other[-2]) | ||
| 93 | + .transpose(-1, -2) | ||
| 94 | + ) | ||
| 95 | + grad_self = ( | ||
| 96 | + unfolded_other_t.mm(unfolded_grad_t).transpose(-1, -2).view(size_self) | ||
| 97 | + ) | ||
| 98 | + | ||
| 99 | + if mask[1]: | ||
| 100 | + size_other_t = list(size_other[:-2]) | ||
| 101 | + size_other_t.extend([size_other[dim_other - 1], size_other[dim_other - 2]]) | ||
| 102 | + unfolded_self = self.unsqueeze(0) if dim_self == 1 else self | ||
| 103 | + grad_other = ( | ||
| 104 | + unfolded_grad_t.mm(unfolded_self).view(size_other_t).transpose(-1, -2) | ||
| 105 | + ) | ||
| 106 | + return grad_self, grad_other | ||
| 107 | + | ||
| 108 | + if dim_self == 1 and dim_other == 1: | ||
| 109 | + grad_self, grad_other = matmul_backward_1d_1d() | ||
| 110 | + elif dim_self == 2 and dim_other == 1: | ||
| 111 | + grad_self, grad_other = matmul_backward_2d_1d() | ||
| 112 | + elif dim_self == 1 and dim_other == 2: | ||
| 113 | + grad_self, grad_other = matmul_backward_1d_2d() | ||
| 114 | + elif dim_self >= 3 and (dim_other == 1 or dim_other == 2): | ||
| 115 | + grad_self, grad_other = matmul_backward_nd_lt3d() | ||
| 116 | + elif (dim_self == 1 or dim_self == 2) and dim_other >= 3: | ||
| 117 | + grad_self, grad_other = matmul_backward_lt3d_nd() | ||
| 118 | + else: | ||
| 119 | + grad_self = ( | ||
| 120 | + torch.matmul(grad, other.transpose(-1, -2)) if mask[0] else grad_self | ||
| 121 | + ) | ||
| 122 | + grad_other = ( | ||
| 123 | + torch.matmul(self.transpose(-1, -2), grad) if mask[1] else grad_other | ||
| 124 | + ) | ||
| 125 | + | ||
| 126 | + return grad_self, grad_other | ||
| 127 | + | ||
| 128 | + | ||
| 21 | def _register_triton_decompositions(): | 129 | def _register_triton_decompositions(): |
| 22 | from .config import is_ascend950, enable_fast_gelu | 130 | from .config import is_ascend950, enable_fast_gelu |
| 23 | from .lowering import _add_overload # noqa: F401 | 131 | from .lowering import _add_overload # noqa: F401 |
| @@ -30,6 +138,9 @@ def _register_triton_decompositions(): | |||
| 30 | aten.expm1, | 138 | aten.expm1, |
| 31 | aten.native_layer_norm, | 139 | aten.native_layer_norm, |
| 32 | aten.repeat_interleave.Tensor, # perf issue | 140 | aten.repeat_interleave.Tensor, # perf issue |
| 141 | + aten.slice_backward, | ||
| 142 | + aten.embedding_dense_backward, | ||
| 143 | + aten.matmul_backward.default, | ||
| 33 | ] | 144 | ] |
| 34 | 145 | ||
| 35 | if is_ascend950: | 146 | if is_ascend950: |
| @@ -64,6 +175,8 @@ def _register_triton_decompositions(): | |||
| 64 | result = x * sigmoid_z | 175 | result = x * sigmoid_z |
| 65 | return result | 176 | return result |
| 66 | 177 | ||
| 178 | + register_decomposition([aten.matmul_backward.default])(_matmul_backward_inductor) | ||
| 179 | + | ||
| 67 | 180 | ||
| 68 | _register_npu_triton_decompositions() | 181 | _register_npu_triton_decompositions() |
| 69 | 182 | ||
| @@ -607,6 +607,7 @@ TORCH_NATIVE_FALLBACK_LIST = [ | |||
| 607 | aten._embedding_bag_forward_only.out, | 607 | aten._embedding_bag_forward_only.out, |
| 608 | aten._embedding_bag_per_sample_weights_backward.default, | 608 | aten._embedding_bag_per_sample_weights_backward.default, |
| 609 | aten._embedding_bag_per_sample_weights_backward.out, | 609 | aten._embedding_bag_per_sample_weights_backward.out, |
| 610 | + aten.embedding_dense_backward.default, | ||
| 610 | aten._fft_r2c.default, | 611 | aten._fft_r2c.default, |
| 611 | aten._fft_r2c.out, | 612 | aten._fft_r2c.out, |
| 612 | aten._flash_attention_backward.default, | 613 | aten._flash_attention_backward.default, |