已合并
[inductor]constrain the usage scenarios of welford #43722
rain-666创建于 8月4日
[inductor]constrain the usage scenarios of welford #43722
已合并
共 2 个文件变更+92-8
| @@ -49,6 +49,35 @@ class TestVarMean(TestUtils): | |||
| 49 | npu_config.enable_welford = previous | 49 | npu_config.enable_welford = previous |
| 50 | torch._dynamo.reset() | 50 | torch._dynamo.reset() |
| 51 | 51 | ||
| 52 | + def test_welford_does_not_change_standalone_var_codegen(self): | ||
| 53 | + previous = npu_config.enable_welford | ||
| 54 | + npu_config.enable_welford = True | ||
| 55 | + torch._dynamo.reset() | ||
| 56 | + try: | ||
| 57 | + input_element = self._generate_tensor((8, 16, 128), "float32") | ||
| 58 | + | ||
| 59 | + def normalize(x): | ||
| 60 | + mean = torch.mean(x, dim=-1, keepdim=True) | ||
| 61 | + var = torch.var(x, dim=-1, correction=0, keepdim=True) | ||
| 62 | + return (x - mean) / torch.sqrt(var + 1e-6) | ||
| 63 | + | ||
| 64 | + expected = normalize(input_element) | ||
| 65 | + compiled = torch.compile( | ||
| 66 | + normalize, | ||
| 67 | + backend="inductor", | ||
| 68 | + dynamic=False, | ||
| 69 | + options={"unroll_reductions_threshold": 1}, | ||
| 70 | + ) | ||
| 71 | + actual, codes = run_and_get_code(compiled, input_element) | ||
| 72 | + | ||
| 73 | + self.assertEqual(expected, actual, atol=1e-1, rtol=1e-1) | ||
| 74 | + code = "\n".join(codes) | ||
| 75 | + self.assertNotIn("triton_helpers.welford_reduce", code) | ||
| 76 | + self.assertNotIn("triton_helpers.welford(", code) | ||
| 77 | + finally: | ||
| 78 | + npu_config.enable_welford = previous | ||
| 79 | + torch._dynamo.reset() | ||
| 80 | + | ||
| 52 | def test_welford_simd_codegen(self): | 81 | def test_welford_simd_codegen(self): |
| 53 | previous = npu_config.enable_welford | 82 | previous = npu_config.enable_welford |
| 54 | npu_config.enable_welford = True | 83 | npu_config.enable_welford = True |
| @@ -102,6 +131,41 @@ class TestVarMean(TestUtils): | |||
| 102 | npu_config.enable_welford = previous | 131 | npu_config.enable_welford = previous |
| 103 | torch._dynamo.reset() | 132 | torch._dynamo.reset() |
| 104 | 133 | ||
| 134 | + @parametrize("dtype", ["float16", "bfloat16"]) | ||
| 135 | + def test_welford_simd_low_precision_codegen(self, dtype): | ||
| 136 | + if not npu_config.is_ascend950: | ||
| 137 | + self.skipTest("low-precision Welford fallback is Ascend 950-specific") | ||
| 138 | + | ||
| 139 | + previous = npu_config.enable_welford | ||
| 140 | + npu_config.enable_welford = True | ||
| 141 | + torch._dynamo.reset() | ||
| 142 | + try: | ||
| 143 | + input_element = self._generate_tensor((8, 16, 64), dtype) | ||
| 144 | + weight = self._generate_tensor((64,), dtype) | ||
| 145 | + bias = self._generate_tensor((64,), dtype) | ||
| 146 | + | ||
| 147 | + def layer_norm(x, gamma, beta): | ||
| 148 | + return F.layer_norm(x, (64,), gamma, beta, 1e-6) | ||
| 149 | + | ||
| 150 | + expected = layer_norm(input_element, weight, bias) | ||
| 151 | + compiled = torch.compile( | ||
| 152 | + layer_norm, | ||
| 153 | + backend="inductor", | ||
| 154 | + dynamic=False, | ||
| 155 | + options={"unroll_reductions_threshold": 1}, | ||
| 156 | + ) | ||
| 157 | + actual, codes = run_and_get_code( | ||
| 158 | + compiled, input_element, weight, bias | ||
| 159 | + ) | ||
| 160 | + | ||
| 161 | + self.assertEqual(expected, actual, atol=1e-1, rtol=1e-1) | ||
| 162 | + code = "\n".join(codes) | ||
| 163 | + self.assertIn("'vectorized_welford_axis':", code) | ||
| 164 | + self.assertIn("npu_kernel_type': 'simd'", code) | ||
| 165 | + finally: | ||
| 166 | + npu_config.enable_welford = previous | ||
| 167 | + torch._dynamo.reset() | ||
| 168 | + | ||
| 105 | def test_welford_simd_codegen_disabled(self): | 169 | def test_welford_simd_codegen_disabled(self): |
| 106 | previous = npu_config.enable_welford | 170 | previous = npu_config.enable_welford |
| 107 | npu_config.enable_welford = False | 171 | npu_config.enable_welford = False |
| @@ -917,7 +917,9 @@ def _register_npu_inductor_fallbacks(): | |||
| 917 | return var, mean | 917 | return var, mean |
| 918 | return (var,) | 918 | return (var,) |
| 919 | 919 | ||
| 920 | - def var_mean_helper_(x, *, axis, correction, keepdim, return_mean): | 920 | + def var_mean_helper_( |
| 921 | + x, *, axis, correction, keepdim, return_mean, allow_welford | ||
| 922 | + ): | ||
| 921 | out_dtype = x.get_dtype() | 923 | out_dtype = x.get_dtype() |
| 922 | compute_dtype = get_computation_dtype(out_dtype) | 924 | compute_dtype = get_computation_dtype(out_dtype) |
| 923 | x = to_dtype(x, compute_dtype, copy=False) | 925 | x = to_dtype(x, compute_dtype, copy=False) |
| @@ -930,7 +932,8 @@ def _register_npu_inductor_fallbacks(): | |||
| 930 | ) | 932 | ) |
| 931 | output = ( | 933 | output = ( |
| 932 | var_mean_sum_(**kwargs) | 934 | var_mean_sum_(**kwargs) |
| 933 | - if not npu_config.enable_welford | 935 | + if not allow_welford |
| 936 | + or not npu_config.enable_welford | ||
| 934 | or use_two_step_variance(x, axis=axis, keepdim=keepdim) | 937 | or use_two_step_variance(x, axis=axis, keepdim=keepdim) |
| 935 | else var_mean_welford_(**kwargs) | 938 | else var_mean_welford_(**kwargs) |
| 936 | ) | 939 | ) |
| @@ -940,13 +943,25 @@ def _register_npu_inductor_fallbacks(): | |||
| 940 | @register_lowering(aten.var_mean) | 943 | @register_lowering(aten.var_mean) |
| 941 | def var_mean(x, axis=None, *, correction=None, keepdim=False): | 944 | def var_mean(x, axis=None, *, correction=None, keepdim=False): |
| 942 | return var_mean_helper_( | 945 | return var_mean_helper_( |
| 943 | - x, axis=axis, correction=correction, keepdim=keepdim, return_mean=True | 946 | + x, |
| 947 | + axis=axis, | ||
| 948 | + correction=correction, | ||
| 949 | + keepdim=keepdim, | ||
| 950 | + return_mean=True, | ||
| 951 | + allow_welford=True, | ||
| 944 | ) | 952 | ) |
| 945 | 953 | ||
| 946 | @register_lowering([aten.var, prims.var]) | 954 | @register_lowering([aten.var, prims.var]) |
| 947 | def var_(x, axis=None, *, correction=None, keepdim=False): | 955 | def var_(x, axis=None, *, correction=None, keepdim=False): |
| 956 | + # A standalone var can fuse with a separate mean reduction. Welford then | ||
| 957 | + # creates a dual-reduction kernel that falls back from the NPU linear path. | ||
| 948 | return var_mean_helper_( | 958 | return var_mean_helper_( |
| 949 | - x, axis=axis, correction=correction, keepdim=keepdim, return_mean=False | 959 | + x, |
| 960 | + axis=axis, | ||
| 961 | + correction=correction, | ||
| 962 | + keepdim=keepdim, | ||
| 963 | + return_mean=False, | ||
| 964 | + allow_welford=False, | ||
| 950 | ) | 965 | ) |
| 951 | 966 | ||
| 952 | @register_lowering(aten.index, type_promotion_kind=None) | 967 | @register_lowering(aten.index, type_promotion_kind=None) |
| @@ -1020,9 +1035,13 @@ def _register_npu_inductor_fallbacks(): | |||
| 1020 | bias=None, | 1035 | bias=None, |
| 1021 | eps=1e-5 | 1036 | eps=1e-5 |
| 1022 | ): | 1037 | ): |
| 1023 | - # Performance consideration: fallback for bfloat16 and float16 | 1038 | + # Keep the existing low-precision fallback unless Welford is explicitly |
| 1024 | - if is_ascend950 and \ | 1039 | + # enabled for FP16/BF16 on Ascend 950. |
| 1025 | - (x.dtype == torch.bfloat16 or x.dtype == torch.float16): | 1040 | + if ( |
| 1041 | + is_ascend950 | ||
| 1042 | + and x.dtype in (torch.float16, torch.bfloat16) | ||
| 1043 | + and not npu_config.enable_welford | ||
| 1044 | + ): | ||
| 1026 | return fallback_handler(aten.native_layer_norm.default)(x, normalized_shape, weight, bias, eps) | 1045 | return fallback_handler(aten.native_layer_norm.default)(x, normalized_shape, weight, bias, eps) |
| 1027 | # Validate input | 1046 | # Validate input |
| 1028 | if not isinstance(normalized_shape, (list, tuple)): | 1047 | if not isinstance(normalized_shape, (list, tuple)): |
| @@ -1040,7 +1059,8 @@ def _register_npu_inductor_fallbacks(): | |||
| 1040 | axis=reduce_dims, | 1059 | axis=reduce_dims, |
| 1041 | correction=0, # Layer normalization uses 0 correction (population variance) | 1060 | correction=0, # Layer normalization uses 0 correction (population variance) |
| 1042 | keepdim=True, # Keep dimensions for broadcasting | 1061 | keepdim=True, # Keep dimensions for broadcasting |
| 1043 | - return_mean=True | 1062 | + return_mean=True, |
| 1063 | + allow_welford=True, | ||
| 1044 | ) | 1064 | ) |
| 1045 | 1065 | ||
| 1046 | # Calculate normalized result (x - mean) / sqrt(var + eps) | 1066 | # Calculate normalized result (x - mean) / sqrt(var + eps) |