已合并
AddRmsNorm + quant 支持pattern算子融合 #2551
mihudan创建于 1月26日
AddRmsNorm + quant 支持pattern算子融合 #2551
已合并
共 9 个文件变更+398-14
| @@ -15,6 +15,7 @@ | |||
| 15 | | npu_add_rms_norm输出经flatten(0,1) 后作为npu_dynamic_quant（不含smooth_scales参数）输入，且npu_dynamic_quant输出的scaleOut执行view(-1,1) | npu_add_rms_norm_dynamic_quant（自动处理flatten与view操作） | | 15 | | npu_add_rms_norm输出经flatten(0,1) 后作为npu_dynamic_quant（不含smooth_scales参数）输入，且npu_dynamic_quant输出的scaleOut执行view(-1,1) | npu_add_rms_norm_dynamic_quant（自动处理flatten与view操作） | |
| 16 | | npu_add_rms_norm输出先获取最后一维尺寸h，再经view(-1, h)变形及to(torch.float32)类型转换 | npu_add_rms_norm_cast（自动处理view） | | 16 | | npu_add_rms_norm输出先获取最后一维尺寸h，再经view(-1, h)变形及to(torch.float32)类型转换 | npu_add_rms_norm_cast（自动处理view） | |
| 17 | | matmul输出作为transpose输入，transpose参数仅支持(0,1)或者(1,0) | npu_transpose_batchmatmul | | 17 | | matmul输出作为transpose输入，transpose参数仅支持(0,1)或者(1,0) | npu_transpose_batchmatmul | |
| 18 | +| npu_add_rms_norm输出作为npu_quantize输入，npu_add_rms_norm输入尾轴需32B对齐，并满足融合算子npu_add_rms_norm_quant约束条件 | npu_add_rms_norm_quant | | ||
| 18 | 19 | ||
| 19 | 另外，用户可通过[register\_replacement](register_replacement.md)接口**实现自定义算子融合Pass注册**（参见接口调用示例），注意需自行保证融合规则的正确性。 | 20 | 另外，用户可通过[register\_replacement](register_replacement.md)接口**实现自定义算子融合Pass注册**（参见接口调用示例），注意需自行保证融合规则的正确性。 |
| 20 | 21 | ||
| @@ -22,7 +23,7 @@ | |||
| 22 | 23 | ||
| 23 | - 本功能依赖PyTorch 2.6.0或更高版本。 | 24 | - 本功能依赖PyTorch 2.6.0或更高版本。 |
| 24 | - 无论是默认支持的算子融合Pass还是自定义的算子融合Pass，均可由pattern\_fusion\_pass配置。 | 25 | - 无论是默认支持的算子融合Pass还是自定义的算子融合Pass，均可由pattern\_fusion\_pass配置。 |
| 25 | -- 表1中matmul输入必须是三维，npu_transpose_batchmatmul算子融合仅max-autotune模式不生效。 | 26 | +- 表1中matmul输入必须是三维，npu_transpose_batchmatmul及npu_add_rms_norm_quant算子融合max-autotune模式不生效。 |
| 26 | 27 | ||
| 27 | ## 使用方法 | 28 | ## 使用方法 |
| 28 | 29 | ||
| @@ -7,7 +7,7 @@ | |||
| 7 | ## 函数原型 | 7 | ## 函数原型 |
| 8 | 8 | ||
| 9 | ```python | 9 | ```python |
| 10 | -register_replacement(search_fn, replace_fn, example_inputs, trace_fn=fwd_only, extra_check=_return_true, search_fn_pattern=None) | 10 | +register_replacement(search_fn, replace_fn, example_inputs, trace_fn=fwd_only, extra_check=_return_true, search_fn_pattern=None, scalar_workaround=None, skip_duplicates=False) |
| 11 | ``` | 11 | ``` |
| 12 | 12 | ||
| 13 | ## 参数说明 | 13 | ## 参数说明 |
| @@ -15,12 +15,14 @@ register_replacement(search_fn, replace_fn, example_inputs, trace_fn=fwd_only, e | |||
| 15 | 15 | ||
| 16 | | 参数名 | 输入/输出 | 说明 | 是否必选 | | 16 | | 参数名 | 输入/输出 | 说明 | 是否必选 | |
| 17 | | --- | --- | --- | --- | | 17 | | --- | --- | --- | --- | |
| 18 | -| search_fn | 输入 | 该函数是希望在FX图中识别的算子组合或计算逻辑，如需要融合的算子组合 | 是 | | 18 | +| search_fn | 输入 | 该函数是希望在FX图中识别的算子组合或计算逻辑，如需要融合的算子组合。 | 是 | |
| 19 | -| replace_fn | 输入 | 在目标图中找到search_fn对应的组合时，会用这个函数的计算逻辑替换原有子图，实现算子融合或优化 | 是 | | 19 | +| replace_fn | 输入 | 在目标图中找到search_fn对应的组合时，会用这个函数的计算逻辑替换原有子图，实现算子融合或优化。 | 是 | |
| 20 | | example_inputs | 输入 | 用于追踪search_fn和replace_fn的示例输入张量。输入的形状和dtype需与实际场景匹配。 | 是 | | 20 | | example_inputs | 输入 | 用于追踪search_fn和replace_fn的示例输入张量。输入的形状和dtype需与实际场景匹配。 | 是 | |
| 21 | | trace_fn | 输入 | 默认仅追踪前向计算图，适用于推理阶段的优化；若需支持训练场景，可传入支持反向追踪的函数。 | 否 | | 21 | | trace_fn | 输入 | 默认仅追踪前向计算图，适用于推理阶段的优化；若需支持训练场景，可传入支持反向追踪的函数。 | 否 | |
| 22 | -| extra_check | 输入 | 找到算子组合后的额外校验函数，函数的入参必须为torch._inductor.pattern_matcher中的Match对象，用于对匹配结果进行更多自定义的判断，如判断算子组合是否在同一条流上/判断设备类型/判断入参形状等 | 否 | | 22 | +| extra_check | 输入 | 找到算子组合后的额外校验函数，函数的入参必须为torch._inductor.pattern_matcher中的Match对象，用于对匹配结果进行更多自定义的判断，如判断算子组合是否在同一条流上/判断设备类型/判断入参形状等。 | 否 | |
| 23 | | search_fn_pattern | 输入 | 自定义的pattern对象，一般无需传入。定义参考PyTorch原生MultiOutputPattern对象的定义规则。传入该参数后，将不再使用search_fn来匹配算子组合，而是直接使用该参数作为匹配规则。 | 否 | | 23 | | search_fn_pattern | 输入 | 自定义的pattern对象，一般无需传入。定义参考PyTorch原生MultiOutputPattern对象的定义规则。传入该参数后，将不再使用search_fn来匹配算子组合，而是直接使用该参数作为匹配规则。 | 否 | |
| 24 | +| scalar_workaround | 输入 | 用于显示绑定search_fn中标量参数值，用于匹配FX图追踪时固化的标量。 | 否 | | ||
| 25 | +| skip_duplicates | 输入 | 用于控制注册阶段的重复检测行为，设为True时，若检测到重复的匹配模式，会跳过该重复模式的注册；设为False时，若检测到重复模式则直接抛出错误，禁止该重复模式完成注册, 该参数仅在PyTorch版本≥2.7.0时生效。 | 否 | | ||
| 24 | 26 | ||
| 25 | ## 返回值说明 | 27 | ## 返回值说明 |
| 26 | 28 | ||
| @@ -0,0 +1,145 @@ | |||
| 1 | +__all__ = [] | ||
| 2 | + | ||
| 3 | +import functools | ||
| 4 | +import sys | ||
| 5 | +import torch | ||
| 6 | + | ||
| 7 | +from torch._inductor.pattern_matcher import Match | ||
| 8 | +from torch._subclasses.fake_tensor import FakeTensorMode | ||
| 9 | + | ||
| 10 | +from torchair.core.utils import logger | ||
| 11 | +from torchair.patterns.pattern_pass_manager import _PatternPassManager, _check_pattern_stream | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +def _pattern_extra_check(match: Match) -> bool: | ||
| 15 | + | ||
| 16 | + for node in match.nodes: | ||
 | |||
| 17 | + if node.target == torch.ops.npu.npu_quantize.default: | ||
| 18 | + dtype_x0 = node.args[0].meta['val'].dtype | ||
| 19 | + # zero_points dtype only supports INT8 UINT8 | ||
| 20 | + if node.args[2].meta['val'].dtype in [torch.uint8, torch.int8]: | ||
| 21 | + logger.debug("Parameter type mismatch in pattern match for npu_add_rms_norm_quant fusion pass.") | ||
| 22 | + return False | ||
| 23 | + # div_mode=False zero_points not supports FLOAT16 | ||
| 24 | + if(len(node.args) == 6 and not node.args[5] and not (dtype_x0 == torch.bfloat16)): | ||
| 25 | + logger.debug("Parameter type mismatch in pattern match for npu_add_rms_norm_quant fusion pass.") | ||
| 26 | + return False | ||
| 27 | + # output dtype only supports INT8, QINT8 | ||
| 28 | + if node.args[3] not in [1, 12]: # 1: INT8, 12: QINT8 | ||
| 29 | + logger.debug("Output type mismatch in pattern match for npu_add_rms_norm_quant fusion pass.") | ||
| 30 | + return False | ||
| 31 | + | ||
| 32 | + elif node.target == torch.ops.npu.npu_add_rms_norm.default: | ||
| 33 | + # div_mode=True only supports BFLOAT16 and FLOAT16 | ||
| 34 | + x0_val = node.args[0].meta['val'] | ||
| 35 | + x1_val = node.args[1].meta['val'] | ||
| 36 | + if not ((x0_val.dtype == torch.float16 and x1_val.dtype == torch.float16) or | ||
| 37 | + (x0_val.dtype == torch.bfloat16 and x1_val.dtype == torch.bfloat16)): | ||
| 38 | + logger.debug("Parameter type mismatch in pattern match for npu_add_rms_norm_quant fusion pass.") | ||
| 39 | + return False | ||
| 40 | + | ||
| 41 | + # last dim size should be multiple of 32 bytes | ||
| 42 | + if x1_val.ndim == 0: | ||
| 43 | + return False | ||
| 44 | + elem_bytes = x1_val.element_size() | ||
| 45 | + last_dim_elem = x1_val.shape[-1] | ||
| 46 | + last_dim_total_bytes = last_dim_elem * elem_bytes | ||
| 47 | + if not isinstance(last_dim_total_bytes, int) or last_dim_total_bytes % 32 != 0: | ||
| 48 | + logger.debug( | ||
| 49 | + f"Input size detected in pattern match for npu_add_rms_norm_quant fusion pass." | ||
| 50 | + f"Fusion is not supported because the input of the npu_add_rms_norm_quant must be divisible by 32 bytes." | ||
| 51 | + ) | ||
| 52 | + return False | ||
| 53 | + | ||
| 54 | + return True | ||
| 55 | + | ||
| 56 | + | ||
| 57 | +def _get_inputs(): | ||
| 58 | + """ | ||
| 59 | + generate example inputs for addrmsnormquant fusion | ||
| 60 | + """ | ||
| 61 | + N, D = 2, 16 | ||
| 62 | + xshape = [N, D] | ||
| 63 | + type_input = torch.float16 # BFLOAT16、FLOAT16 | ||
| 64 | + gammashape = [D] | ||
| 65 | + | ||
| 66 | + x1 = torch.rand(xshape, dtype=type_input) | ||
| 67 | + x2 = torch.rand(xshape, dtype=type_input) | ||
| 68 | + gamma = torch.rand(gammashape, dtype=type_input) | ||
| 69 | + scales = torch.ones(gammashape, dtype=torch.float32) # FLOAT32、BFLOAT16 | ||
| 70 | + zero_points = torch.tensor(gammashape, dtype=torch.int32) # INT32、BFLOAT16 | ||
| 71 | + return (x1, x2, gamma, scales, zero_points) | ||
| 72 | + | ||
| 73 | + | ||
| 74 | + | ||
| 75 | +def _register_addrmsnormquant_pattern(pattern_pass_manager, div_mode): | ||
| 76 | + if 'torch_npu' not in sys.modules: | ||
| 77 | + logger.info(f'The addrmsnormquant fusion will only be enabled in a torch npu env.' | ||
| 78 | + 'When there is no torch_npu in the env, skip fusion.') | ||
| 79 | + return | ||
| 80 | + | ||
| 81 | + # replace the operator combination of npu_add_rms_norm and npu_quantize | ||
| 82 | + # with the npu_add_rms_norm_quant operator. | ||
| 83 | + def search_fn(x1, x2, gamma, scales, zero_points, epsilon, dtype): | ||
| 84 | + y, _, x_out = torch.ops.npu.npu_add_rms_norm.default(x1, x2, gamma, epsilon) | ||
| 85 | + y_out = torch.ops.npu.npu_quantize.default(y, scales, zero_points=zero_points, dtype=dtype, axis=-1, div_mode=div_mode) | ||
| 86 | + return y_out, x_out | ||
| 87 | + | ||
| 88 | + def replace_fn(x1, x2, gamma, scales, zero_points, epsilon, _): | ||
| 89 | + # output dtype only supports INT8, axis currently only supports -1, passing other values will not take effect | ||
| 90 | + y1, _, x_out = torch.ops.npu.npu_add_rms_norm_quant.default(x1, x2, gamma, scales, zero_points, axis=-1, epsilon=epsilon) | ||
| 91 | + return y1, x_out | ||
| 92 | + | ||
| 93 | + fake_mode = FakeTensorMode() | ||
| 94 | + with fake_mode: | ||
| 95 | + # sizes/values don't actually matter for initial trace | ||
| 96 | + # once we get a possible match we re-trace with the actual values and verify the match still holds | ||
| 97 | + pattern_pass_manager.register_pattern( | ||
| 98 | + search_fn=search_fn, | ||
| 99 | + replace_fn=replace_fn, | ||
| 100 | + example_inputs=_get_inputs(), | ||
| 101 | + extra_check=_pattern_extra_check, | ||
| 102 | + scalar_workaround={"epsilon": 2e-6, "dtype": 1}, | ||
| 103 | + skip_duplicates=True | ||
| 104 | + ) | ||
| 105 | + | ||
| 106 | + | ||
| 107 | + | ||
| 108 | +def _register_addrmsnormquant_pattern_default(pattern_pass_manager, div_mode): | ||
| 109 | + if 'torch_npu' not in sys.modules: | ||
| 110 | + logger.info(f'The addrmsnormquant fusion will only be enabled in a torch npu env.' | ||
| 111 | + 'When there is no torch_npu in the env, skip fusion.') | ||
| 112 | + return | ||
| 113 | + | ||
| 114 | + # replace the operator combination of npu_add_rms_norm and npu_quantize | ||
| 115 | + # with the npu_add_rms_norm_quant operator. | ||
| 116 | + def search_fn(x1, x2, gamma, scales, zero_points, dtype): | ||
| 117 | + y, _, x_out = torch.ops.npu.npu_add_rms_norm.default(x1, x2, gamma) | ||
| 118 | + y_out = torch.ops.npu.npu_quantize.default(y, scales, zero_points=zero_points, dtype=dtype, axis=-1, div_mode=div_mode) | ||
| 119 | + return y_out, x_out | ||
| 120 | + | ||
| 121 | + def replace_fn(x1, x2, gamma, scales, zero_points, _): | ||
| 122 | + # output dtype only supports INT8, axis currently only supports -1, passing other values will not take effect | ||
| 123 | + y1, _, x_out = torch.ops.npu.npu_add_rms_norm_quant.default(x1, x2, gamma, scales, zero_points, axis=-1) | ||
| 124 | + return y1, x_out | ||
| 125 | + | ||
| 126 | + fake_mode = FakeTensorMode() | ||
| 127 | + with fake_mode: | ||
| 128 | + # sizes/values don't actually matter for initial trace | ||
| 129 | + # once we get a possible match we re-trace with the actual values and verify the match still holds | ||
| 130 | + pattern_pass_manager.register_pattern( | ||
| 131 | + search_fn=search_fn, | ||
| 132 | + replace_fn=replace_fn, | ||
| 133 | + example_inputs=_get_inputs(), | ||
| 134 | + extra_check=_pattern_extra_check, | ||
| 135 | + scalar_workaround={"dtype": 1}, | ||
| 136 | + skip_duplicates=True | ||
| 137 | + ) | ||
| 138 | + | ||
| 139 | + | ||
| 140 | + | ||
| 141 | +def _register_addrmsnormquant_patterns(pattern_pass_manager: _PatternPassManager): | ||
| 142 | + _register_addrmsnormquant_pattern_default(pattern_pass_manager, div_mode=True) | ||
| 143 | + _register_addrmsnormquant_pattern(pattern_pass_manager, div_mode=True) | ||
| 144 | + _register_addrmsnormquant_pattern_default(pattern_pass_manager, div_mode=False) | ||
| 145 | + _register_addrmsnormquant_pattern(pattern_pass_manager, div_mode=False) | ||
| @@ -77,7 +77,7 @@ def _register_pattern_for_bmm_transpose(pattern_pass_manager: _PatternPassManage | |||
| 77 | with fake_mode: | 77 | with fake_mode: |
| 78 | # sizes/values don't actually matter for initial trace | 78 | # sizes/values don't actually matter for initial trace |
| 79 | # once we get a possible match we re-trace with the actual values and verify the match still holds | 79 | # once we get a possible match we re-trace with the actual values and verify the match still holds |
| 80 | - input_tensor = functools.partial(torch.empty, (1, 1, 1), device="npu", dtype=torch.float16) | 80 | + input_tensor = functools.partial(torch.empty, (1, 1, 1), dtype=torch.float16) |
| 81 | pattern_pass_manager.register_pattern( | 81 | pattern_pass_manager.register_pattern( |
| 82 | search_fn=search_fn, | 82 | search_fn=search_fn, |
| 83 | replace_fn=replace_fn, | 83 | replace_fn=replace_fn, |
| @@ -1,5 +1,6 @@ | |||
| 1 | __all__ = ["register_replacement"] | 1 | __all__ = ["register_replacement"] |
| 2 | 2 | ||
| 3 | +import torch | ||
| 3 | from torch._inductor.pattern_matcher import Match, PatternMatcherPass | 4 | from torch._inductor.pattern_matcher import Match, PatternMatcherPass |
| 4 | from torch._inductor.pattern_matcher import register_replacement as register | 5 | from torch._inductor.pattern_matcher import register_replacement as register |
| 5 | try: | 6 | try: |
| @@ -67,7 +68,8 @@ class _PatternPassManager: | |||
| 67 | else: | 68 | else: |
| 68 | self.pass_dict = PatternMatcherPass() | 69 | self.pass_dict = PatternMatcherPass() |
| 69 | 70 | ||
| 70 | - def register_pattern(self, search_fn, replace_fn, example_inputs, trace_fn=fwd_only, extra_check=_return_true, search_fn_pattern=None): | 71 | + def register_pattern(self, search_fn, replace_fn, example_inputs, trace_fn=fwd_only, extra_check=_return_true, |
| 72 | + search_fn_pattern=None, scalar_workaround=None, skip_duplicates=False): | ||
| 71 | """ | 73 | """ |
| 72 | Register a new pattern for matching and replacement. | 74 | Register a new pattern for matching and replacement. |
| 73 | 75 | ||
| @@ -86,7 +88,7 @@ class _PatternPassManager: | |||
| 86 | # Then run the user-provided extra_check | 88 | # Then run the user-provided extra_check |
| 87 | return extra_check(match) | 89 | return extra_check(match) |
| 88 | 90 | ||
| 89 | - if hasattr(register, '__code__') and 'pass_dicts' in register.__code__.co_varnames: | 91 | + if torch.__version__ >= "2.7.0": |
| 90 | register( | 92 | register( |
| 91 | search_fn=search_fn, | 93 | search_fn=search_fn, |
| 92 | replace_fn=replace_fn, | 94 | replace_fn=replace_fn, |
| @@ -94,7 +96,9 @@ class _PatternPassManager: | |||
| 94 | trace_fn=trace_fn, | 96 | trace_fn=trace_fn, |
| 95 | pass_dicts=self.pass_dict, | 97 | pass_dicts=self.pass_dict, |
| 96 | extra_check=add_stream_check, | 98 | extra_check=add_stream_check, |
| 97 | - search_fn_pattern=search_fn_pattern | 99 | + search_fn_pattern=search_fn_pattern, |
| 100 | + scalar_workaround=scalar_workaround, | ||
| 101 | + skip_duplicates=skip_duplicates | ||
| 98 | ) | 102 | ) |
| 99 | else: | 103 | else: |
| 100 | register( | 104 | register( |
| @@ -102,10 +106,13 @@ class _PatternPassManager: | |||
| 102 | replace_fn=replace_fn, | 106 | replace_fn=replace_fn, |
| 103 | example_inputs=example_inputs, | 107 | example_inputs=example_inputs, |
| 104 | trace_fn=trace_fn, | 108 | trace_fn=trace_fn, |
| 105 | - pass_dict=self.pass_dict, | 109 | + pass_dicts=self.pass_dict, |
| 106 | extra_check=add_stream_check, | 110 | extra_check=add_stream_check, |
| 111 | + search_fn_pattern=search_fn_pattern, | ||
| 112 | + scalar_workaround=scalar_workaround | ||
| 107 | ) | 113 | ) |
| 108 | 114 | ||
| 115 | + | ||
| 109 | def apply_pass(self, fx_graph): | 116 | def apply_pass(self, fx_graph): |
| 110 | """ | 117 | """ |
| 111 | Apply the registered pattern pass to the given FX graph. | 118 | Apply the registered pattern pass to the given FX graph. |
| @@ -123,11 +130,14 @@ def _pattern_manager(): | |||
| 123 | return _global_pattern_pass_manager | 130 | return _global_pattern_pass_manager |
| 124 | 131 | ||
| 125 | 132 | ||
| 126 | -def register_replacement(search_fn, replace_fn, example_inputs, trace_fn=fwd_only, extra_check=_return_true, search_fn_pattern=None): | 133 | +def register_replacement(search_fn, replace_fn, example_inputs, trace_fn=fwd_only, extra_check=_return_true, |
| 134 | + search_fn_pattern=None, scalar_workaround=None, skip_duplicates=False): | ||
| 127 | _global_pattern_pass_manager = _pattern_manager() | 135 | _global_pattern_pass_manager = _pattern_manager() |
| 128 | _global_pattern_pass_manager.register_pattern(search_fn=search_fn, | 136 | _global_pattern_pass_manager.register_pattern(search_fn=search_fn, |
| 129 | replace_fn=replace_fn, | 137 | replace_fn=replace_fn, |
| 130 | example_inputs=example_inputs, | 138 | example_inputs=example_inputs, |
| 131 | trace_fn=trace_fn, | 139 | trace_fn=trace_fn, |
| 132 | extra_check=extra_check, | 140 | extra_check=extra_check, |
| 133 | - search_fn_pattern=search_fn_pattern) | 141 | + search_fn_pattern=search_fn_pattern, |
| 142 | + scalar_workaround=scalar_workaround, | ||
| 143 | + skip_duplicates=skip_duplicates) | ||
| @@ -8,6 +8,7 @@ from torchair.patterns.add_rms_norm_dynamic_quant import _register_addrmsnormdyn | |||
| 8 | _register_addrmsnormdynamicquant_pattern2 | 8 | _register_addrmsnormdynamicquant_pattern2 |
| 9 | from torchair.patterns.batch_matmul_transpose import _register_batchmatmultranspose_patterns | 9 | from torchair.patterns.batch_matmul_transpose import _register_batchmatmultranspose_patterns |
| 10 | from torchair._utils.graph_utils import debug_compare_fx_graphs | 10 | from torchair._utils.graph_utils import debug_compare_fx_graphs |
| 11 | +from torchair.patterns.add_rms_norm_quant import _register_addrmsnormquant_patterns | ||
| 11 | 12 | ||
| 12 | pattern_pass_manager = _pattern_manager() | 13 | pattern_pass_manager = _pattern_manager() |
| 13 | 14 | ||
| @@ -24,6 +25,7 @@ def _apply_pattern_passes(graph_module: torch.fx.GraphModule, example_inputs=Non | |||
| 24 | _register_addrmsnormcast_pattern(pattern_pass_manager) | 25 | _register_addrmsnormcast_pattern(pattern_pass_manager) |
| 25 | if config.mode.value != "max-autotune": | 26 | if config.mode.value != "max-autotune": |
| 26 | _register_batchmatmultranspose_patterns(pattern_pass_manager) | 27 | _register_batchmatmultranspose_patterns(pattern_pass_manager) |
| 28 | + _register_addrmsnormquant_patterns(pattern_pass_manager) | ||
| 27 | 29 | ||
| 28 | # Set stream labels for all nodes before pattern pass | 30 | # Set stream labels for all nodes before pattern pass |
| 29 | from torchair._utils.graph_utils import add_stream_label_to_node_meta | 31 | from torchair._utils.graph_utils import add_stream_label_to_node_meta |
| @@ -7,6 +7,7 @@ import shutil | |||
| 7 | import torch | 7 | import torch |
| 8 | import torch_npu | 8 | import torch_npu |
| 9 | import torchair | 9 | import torchair |
| 10 | +from torch._subclasses.fake_tensor import FakeTensorMode | ||
| 10 | from torchair.configs.compiler_config import CompilerConfig | 11 | from torchair.configs.compiler_config import CompilerConfig |
| 11 | from torchair.core.utils import logger | 12 | from torchair.core.utils import logger |
| 12 | 13 | ||
| @@ -16,8 +17,36 @@ torch.npu.manual_seed_all(7) | |||
| 16 | logger.setLevel(logging.DEBUG) | 17 | logger.setLevel(logging.DEBUG) |
| 17 | 18 | ||
| 18 | 19 | ||
| 20 | +def find_op(gm, op_default): | ||
| 21 | + for node in gm.graph.nodes: | ||
| 22 | + if node.op == "call_function" and node.target == op_default: | ||
| 23 | + return True | ||
| 24 | + | ||
| 25 | + return False | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +def create_optimize_wrapper(assert_func): | ||
| 29 | + original_func = torchair.npu_fx_compiler._optimize_fx | ||
| 30 | + | ||
| 31 | + def wrapper(gm, config, observer): | ||
| 32 | + ret = original_func(gm, config, observer) | ||
| 33 | + assert_func(gm) | ||
| 34 | + return ret | ||
| 35 | + | ||
| 36 | + return wrapper | ||
| 37 | + | ||
| 38 | + | ||
| 19 | class AclgraphTest(unittest.TestCase): | 39 | class AclgraphTest(unittest.TestCase): |
| 20 | 40 | ||
| 41 | + def setUp(self) -> None: | ||
| 42 | + self.optimize_fx_bak = torchair.npu_fx_compiler._optimize_fx | ||
| 43 | + return super().setUp() | ||
| 44 | + | ||
| 45 | + def tearDown(self) -> None: | ||
| 46 | + if self.optimize_fx_bak is not None: | ||
| 47 | + torchair.npu_fx_compiler._optimize_fx = self.optimize_fx_bak | ||
| 48 | + return super().tearDown() | ||
| 49 | + | ||
| 21 | def test_aclgraph_cache_with_static_kernel(self): | 50 | def test_aclgraph_cache_with_static_kernel(self): |
| 22 | class CachedAclGraphModel(torch.nn.Module): | 51 | class CachedAclGraphModel(torch.nn.Module): |
| 23 | def __init__(self): | 52 | def __init__(self): |
| @@ -1724,5 +1753,195 @@ class AclgraphTest(unittest.TestCase): | |||
| 1724 | 1753 | ||
| 1725 | self.assertTrue(torch.allclose(eager_output, compile_output)) | 1754 | self.assertTrue(torch.allclose(eager_output, compile_output)) |
| 1726 | 1755 | ||
| 1756 | + def assert_addrmsnorm_quant(self, after_gm, expect_fused=True): | ||
| 1757 | + """ | ||
| 1758 | + Check whether the pattern fusion of add_rms_norm + quantize is successful. | ||
| 1759 | + """ | ||
| 1760 | + check_rules = [ | ||
| 1761 | + (torch.ops.npu.npu_add_rms_norm_quant.default, expect_fused), | ||
| 1762 | + (torch.ops.npu.npu_add_rms_norm.default, not expect_fused), | ||
| 1763 | + (torch.ops.npu.npu_quantize.default, not expect_fused), | ||
| 1764 | + ] | ||
| 1765 | + | ||
| 1766 | + for torch_op, expect_exist in check_rules: | ||
| 1767 | + found = find_op(after_gm, torch_op) | ||
| 1768 | + if expect_exist: | ||
| 1769 | + self.assertTrue(found, f"Expected operator '{torch_op}' but not find") | ||
| 1770 | + else: | ||
| 1771 | + self.assertFalse(found, f"Not expected operator '{torch_op}' but find") | ||
| 1772 | + | ||
| 1773 | + def get_quant_input(self, last_axis, dtype1, dtype2, dtype3): | ||
| 1774 | + """ | ||
| 1775 | + Get the input of the add_rms_norm + quantize pattern. | ||
| 1776 | + """ | ||
| 1777 | + x1 = torch.randn(1, 2, last_axis, dtype=dtype1, device='npu') | ||
| 1778 | + x2 = torch.randn(1, 2, last_axis, dtype=dtype1, device='npu') | ||
| 1779 | + gamma = torch.ones(last_axis, dtype=dtype1, device='npu') | ||
| 1780 | + scales = torch.ones(last_axis, dtype=dtype2, device='npu') | ||
| 1781 | + zero_points = torch.zeros(last_axis, dtype=dtype3, device='npu') | ||
| 1782 | + return x1, x2, gamma, scales, zero_points | ||
| 1783 | + | ||
| 1784 | + | ||
| 1785 | + def test_pattern_pass_addrmsnorm_quant(self): | ||
| 1786 | + | ||
| 1787 | + def f(x1, x2, gamma, scales, zero_points, div_mode=True): | ||
| 1788 | + x1 = x1.reshape([1, -1, 16]) | ||
| 1789 | + x2 = x2.reshape([1, -1, 16]) | ||
Y 【一般】避免功能拓展时散弹式修改，需要改用全局变量   | |||
| 1790 | + y, _, xOut = torch_npu.npu_add_rms_norm(x1, x2, gamma, 4e-6) | ||
| 1791 | + yOut = torch_npu.npu_quantize(y, scales, zero_points, torch.qint8, axis=-1, div_mode=div_mode) | ||
| 1792 | + return yOut, xOut | ||
| 1793 | + | ||
| 1794 | + def f_static(x1, x2, gamma, scales, zero_points): | ||
| 1795 | + y, _, xOut = torch_npu.npu_add_rms_norm(x1, x2, gamma, 1e-6) | ||
| 1796 | + yOut = torch_npu.npu_quantize(y, scales, zero_points=zero_points, dtype=torch.int8, axis=-1) | ||
| 1797 | + return yOut, xOut | ||
| 1798 | + | ||
| 1799 | + torchair.npu_fx_compiler._optimize_fx = create_optimize_wrapper(lambda gm: self.assert_addrmsnorm_quant(gm, True)) | ||
| 1800 | + compile_model = torch.compile(f, backend="npugraph_ex", fullgraph=True, dynamic=True) | ||
| 1801 | + | ||
| 1802 | + # test divmode=True | ||
| 1803 | + x1, x2, gamma, scales, zero_points = self.get_quant_input(16, torch.float16, torch.float, torch.int32) | ||
| 1804 | + y1, y2 = f(x1, x2, gamma, scales, zero_points) | ||
| 1805 | + y3, y4 = compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1806 | + self.assertTrue(torch.equal(y1, y3)) | ||
| 1807 | + self.assertTrue(torch.equal(y2, y4)) | ||
| 1808 | + | ||
| 1809 | + x1, x2, gamma, scales, zero_points = self.get_quant_input(16, torch.bfloat16, torch.bfloat16, torch.bfloat16) | ||
| 1810 | + y1, y2 = f(x1, x2, gamma, scales, zero_points) | ||
| 1811 | + y3, y4 = compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1812 | + self.assertTrue(torch.equal(y1, y3)) | ||
| 1813 | + self.assertTrue(torch.equal(y2, y4)) | ||
| 1814 | + | ||
| 1815 | + # test divmode=False | ||
| 1816 | + y1, y2 = f(x1, x2, gamma, scales, zero_points, False) | ||
| 1817 | + y3, y4 = compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1818 | + self.assertTrue(torch.equal(y1, y3)) | ||
| 1819 | + self.assertTrue(torch.equal(y2, y4)) | ||
| 1820 | + | ||
| 1821 | + # test static | ||
| 1822 | + compile_model = torch.compile(f_static, backend="npugraph_ex", fullgraph=True, dynamic=False) | ||
| 1823 | + y1, y2 = f_static(x1, x2, gamma, scales, zero_points) | ||
| 1824 | + y3, y4 = compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1825 | + self.assertTrue(torch.equal(y1, y3)) | ||
| 1826 | + self.assertTrue(torch.equal(y2, y4)) | ||
| 1827 | + | ||
| 1828 | + | ||
| 1829 | + | ||
| 1830 | + def test_pattern_pass_addrmsnorm_quant_mismatched(self): | ||
| 1831 | + | ||
| 1832 | + def f(x1, x2, gamma, scales, zero_points, out_dtype=torch.qint8, div_mode=True): | ||
| 1833 | + x1 = x1.reshape([1, -1, 16]) | ||
| 1834 | + x2 = x2.reshape([1, -1, 16]) | ||
| 1835 | + y, _, xOut = torch_npu.npu_add_rms_norm(x1, x2, gamma) | ||
| 1836 | + yOut = torch_npu.npu_quantize(y, scales, zero_points, out_dtype, axis=-1, div_mode = div_mode) | ||
| 1837 | + return yOut, xOut | ||
| 1838 | + | ||
| 1839 | + def f_use(x1, x2, gamma, scales, zero_points): | ||
| 1840 | + x1 = x1.reshape([1, -1, 16]) | ||
| 1841 | + x2 = x2.reshape([1, -1, 16]) | ||
| 1842 | + y, _, xOut = torch_npu.npu_add_rms_norm(x1, x2, gamma) | ||
| 1843 | + yOut = torch_npu.npu_quantize(y, scales, zero_points=zero_points, dtype=torch.qint8, axis=-1) | ||
| 1844 | + yOut = y + yOut | ||
| 1845 | + return yOut, xOut | ||
| 1846 | + | ||
| 1847 | + def f_noreshape(x1, x2, gamma, scales, zero_points): | ||
| 1848 | + y, _, xOut = torch_npu.npu_add_rms_norm(x1, x2, gamma) | ||
| 1849 | + yOut = torch_npu.npu_quantize(y, scales, zero_points=zero_points, dtype=torch.qint8, axis=-1, div_mode=True) | ||
| 1850 | + return yOut, xOut | ||
| 1851 | + | ||
| 1852 | + torchair.npu_fx_compiler._optimize_fx = create_optimize_wrapper(lambda gm: self.assert_addrmsnorm_quant(gm, False)) | ||
| 1853 | + compile_model = torch.compile(f, backend="npugraph_ex", fullgraph=True, dynamic=True) | ||
| 1854 | + | ||
| 1855 | + # test uint8 zero_poin | ||
| 1856 | + x1, x2, gamma, scales, zero_points = self.get_quant_input(16, torch.float16, torch.float16, torch.uint8) | ||
| 1857 | + f(x1, x2, gamma, scales, zero_points) | ||
| 1858 | + compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1859 | + | ||
| 1860 | + # test int8 zero_point | ||
| 1861 | + zero_points = torch.zeros(16, dtype=torch.int8, device='npu') | ||
| 1862 | + f(x1, x2, gamma, scales, zero_points) | ||
| 1863 | + compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1864 | + | ||
| 1865 | + # test out_dtype=int32 | ||
| 1866 | + x1, x2, gamma, scales, zero_points = self.get_quant_input(16, torch.bfloat16, torch.bfloat16, torch.bfloat16) | ||
| 1867 | + f(x1, x2, gamma, scales, zero_points, torch.int32) | ||
| 1868 | + compile_model(x1, x2, gamma, scales, zero_points, torch.int32) | ||
| 1869 | + | ||
| 1870 | + # test use value npu_add_rms_norm output | ||
| 1871 | + compile_model = torch.compile(f_use, backend="npugraph_ex", fullgraph=True, dynamic=True) | ||
| 1872 | + f_use(x1, x2, gamma, scales, zero_points) | ||
| 1873 | + compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1874 | + | ||
| 1875 | + # test div_mode=False type mismatch | ||
| 1876 | + x1, x2, gamma, scales, zero_points = self.get_quant_input(16, torch.float16, torch.float16, torch.float16) | ||
| 1877 | + compile_model = torch.compile(f, backend="npugraph_ex", fullgraph=True, dynamic=True) | ||
| 1878 | + f(x1, x2, gamma, scales, zero_points, div_mode=False) | ||
| 1879 | + compile_model(x1, x2, gamma, scales, zero_points, div_mode=False) | ||
| 1880 | + | ||
| 1881 | + # # test last axis not aligned 32byte | ||
| 1882 | + compile_model = torch.compile(f_noreshape, backend="npugraph_ex", fullgraph=True, dynamic=False) | ||
| 1883 | + x1, x2, gamma, scales, zero_points = self.get_quant_input(3, torch.bfloat16, torch.bfloat16, torch.bfloat16) | ||
| 1884 | + f_noreshape(x1, x2, gamma, scales, zero_points) | ||
| 1885 | + compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1886 | + | ||
| 1887 | + # test symint | ||
| 1888 | + compile_model = torch.compile(f_noreshape, backend="npugraph_ex", fullgraph=True, dynamic=True) | ||
| 1889 | + f_noreshape(x1, x2, gamma, scales, zero_points) | ||
| 1890 | + compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1891 | + | ||
| 1892 | + | ||
| 1893 | + def test_pattern_pass_addrmsnorm_quant_with_diff_stream(self): | ||
| 1894 | + class Model(torch.nn.Module): | ||
| 1895 | + def __init__(self): | ||
| 1896 | + super().__init__() | ||
| 1897 | + self.event1 = torchair.ops.npu_create_tagged_event(tag="22") | ||
| 1898 | + self.event2 = torchair.ops.npu_create_tagged_event(tag="33") | ||
| 1899 | + | ||
| 1900 | + def forward(self, x1, x2, gamma, scales, zero_points): | ||
| 1901 | + y, _, xOut = torch_npu.npu_add_rms_norm(x1, x2, gamma) | ||
| 1902 | + torchair.ops.npu_tagged_event_record(self.event1) | ||
| 1903 | + with torchair.scope.npu_stream_switch('2', 3): | ||
| 1904 | + torchair.ops.npu_tagged_event_wait(self.event1) | ||
| 1905 | + yOut = torch_npu.npu_quantize(y, scales, zero_points=zero_points, dtype=torch.qint8, axis=-1) | ||
| 1906 | + torchair.ops.npu_tagged_event_record(self.event2) | ||
| 1907 | + torchair.ops.npu_record_tagged_stream(yOut, '2') | ||
| 1908 | + torchair.ops.npu_tagged_event_wait(self.event2) | ||
| 1909 | + return yOut, xOut | ||
| 1910 | + | ||
| 1911 | + torchair.npu_fx_compiler._optimize_fx = create_optimize_wrapper(lambda gm: self.assert_addrmsnorm_quant(gm, False)) | ||
| 1912 | + | ||
| 1913 | + model = Model() | ||
| 1914 | + compile_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=True) | ||
| 1915 | + | ||
| 1916 | + x1, x2, gamma, scales, zero_points = self.get_quant_input(16, torch.bfloat16, torch.bfloat16, torch.bfloat16) | ||
| 1917 | + compile_model(x1, x2, gamma, scales, zero_points) | ||
| 1918 | + | ||
| 1919 | + | ||
| 1920 | + def test_pattern_pass_addrmsnorm_quant_skip_duplicates(self): | ||
| 1921 | + def f(x1, x2): | ||
| 1922 | + return x1 + x2 | ||
| 1923 | + def search_fn(x1, x2, gamma, scales, zero_points, epsilon, dtype): | ||
| 1924 | + y, _, x_out = torch.ops.npu.npu_add_rms_norm.default(x1, x2, gamma, epsilon) | ||
| 1925 | + y_out = torch.ops.npu.npu_quantize.default(y, scales, zero_points=zero_points, dtype=dtype, axis=-1) | ||
| 1926 | + return y_out, x_out | ||
| 1927 | + | ||
| 1928 | + def replace_fn(x1, x2, gamma, scales, zero_points, epsilon, _): | ||
| 1929 | + y1, _, x_out = torch.ops.npu.npu_add_rms_norm_quant.default(x1, x2, gamma, scales, zero_points, axis=-1, epsilon=epsilon) | ||
| 1930 | + return y1, x_out | ||
| 1931 | + | ||
| 1932 | + fake_mode = FakeTensorMode() | ||
| 1933 | + with fake_mode: | ||
| 1934 | + # sizes/values don't actually matter for initial trace | ||
| 1935 | + # once we get a possible match we re-trace with the actual values and verify the match still holds | ||
| 1936 | + torchair.register_replacement( | ||
| 1937 | + search_fn=search_fn, | ||
| 1938 | + replace_fn=replace_fn, | ||
| 1939 | + example_inputs=self.get_quant_input(16, torch.bfloat16, torch.bfloat16, torch.bfloat16), | ||
| 1940 | + scalar_workaround={"epsilon": 2e-6, "dtype": 1}, | ||
| 1941 | + skip_duplicates=True | ||
| 1942 | + ) | ||
| 1943 | + torch.compile(f, backend="npugraph_ex", fullgraph=True, dynamic=True) | ||
| 1944 | + | ||
| 1945 | + | ||
| 1727 | if __name__ == '__main__': | 1946 | if __name__ == '__main__': |
| 1728 | unittest.main() | 1947 | unittest.main() |
| @@ -3,10 +3,10 @@ | |||
| 3 | "signature": "(x: torchair.ge.TensorBase, *, dependencies=[], node_name=None)" | 3 | "signature": "(x: torchair.ge.TensorBase, *, dependencies=[], node_name=None)" |
| 4 | }, | 4 | }, |
| 5 | "torch_npu.dynamo.torchair.register_replacement": { | 5 | "torch_npu.dynamo.torchair.register_replacement": { |
| 6 | - "signature": "(search_fn, replace_fn, example_inputs, trace_fn=<function fwd_only>, extra_check=<function _return_true>, search_fn_pattern=None)" | 6 | + "signature": "(search_fn, replace_fn, example_inputs, trace_fn=<function fwd_only>, extra_check=<function _return_true>, search_fn_pattern=None, scalar_workaround=None, skip_duplicates=False)" |
| 7 | }, | 7 | }, |
| 8 | "torch_npu.dynamo.torchair.patterns.pattern_pass_manager.register_replacement": { | 8 | "torch_npu.dynamo.torchair.patterns.pattern_pass_manager.register_replacement": { |
| 9 | - "signature": "(search_fn, replace_fn, example_inputs, trace_fn=<function fwd_only>, extra_check=<function _return_true>, search_fn_pattern=None)" | 9 | + "signature": "(search_fn, replace_fn, example_inputs, trace_fn=<function fwd_only>, extra_check=<function _return_true>, search_fn_pattern=None, scalar_workaround=None, skip_duplicates=False)" |
| 10 | }, | 10 | }, |
| 11 | "torch_npu.dynamo.torchair.CompilerConfig": { | 11 | "torch_npu.dynamo.torchair.CompilerConfig": { |
| 12 | "signature": "()" | 12 | "signature": "()" |
| @@ -325,6 +325,8 @@ def _stub_npu_add_rms_norm_default(self, *args, **kwargs): | |||
| 325 | def _stub_npu_dynamic_quant_default(self, *args, **kwargs): | 325 | def _stub_npu_dynamic_quant_default(self, *args, **kwargs): |
| 326 | return torch.randn([3, 2]), torch.randn([3, 2]) | 326 | return torch.randn([3, 2]), torch.randn([3, 2]) |
| 327 | 327 | ||
| 328 | +def _stub_npu_quantize_default(self, *args, **kwargs): | ||
| 329 | + return torch.randn([3, 2]) | ||
| 328 | 330 | ||
| 329 | def _stub_npu_dtype_cast_default(self, *args, **kwargs): | 331 | def _stub_npu_dtype_cast_default(self, *args, **kwargs): |
| 330 | return torch.randn([3, 2]), torch.randn([3, 2]) | 332 | return torch.randn([3, 2]), torch.randn([3, 2]) |
| @@ -365,6 +367,9 @@ class StubNpu: | |||
| 365 | self.npu_dynamic_quant = types.SimpleNamespace( | 367 | self.npu_dynamic_quant = types.SimpleNamespace( |
| 366 | default=_stub_npu_dynamic_quant_default | 368 | default=_stub_npu_dynamic_quant_default |
| 367 | ) | 369 | ) |
| 370 | + self.npu_quantize = types.SimpleNamespace( | ||
| 371 | + default=_stub_npu_quantize_default | ||
| 372 | + ) | ||
| 368 | self._npu_dtype_cast = types.SimpleNamespace( | 373 | self._npu_dtype_cast = types.SimpleNamespace( |
| 369 | default=_stub_npu_dtype_cast_default | 374 | default=_stub_npu_dtype_cast_default |
| 370 | ) | 375 | ) |
注意编程规范，函数的命名要和函数的行为一致，且一个函数只做一件事情。这个函数既校验了跨流，又校验了算子本身逻辑，代码糅杂，函数设计有问题