已合并
【fix】batcnnorm_check_fix #36322
cuiduo创建于 5月21日
【fix】batcnnorm_check_fix #36322
已合并
共 4 个文件变更+13-4
| @@ -9,6 +9,7 @@ import torch_npu | |||
| 9 | import torch_npu._inductor | 9 | import torch_npu._inductor |
| 10 | 10 | ||
| 11 | 11 | ||
| 12 | + | ||
| 12 | class TestRunWithRngState(TestUtils): | 13 | class TestRunWithRngState(TestUtils): |
| 13 | def op_calc(self, current_state, like, device, dtype): | 14 | def op_calc(self, current_state, like, device, dtype): |
| 14 | res1 = torch._prims.rng_prims.run_with_rng_state( | 15 | res1 = torch._prims.rng_prims.run_with_rng_state( |
| @@ -2,7 +2,6 @@ import os | |||
| 2 | import copy | 2 | import copy |
| 3 | from typing import Any, Callable, Optional, TYPE_CHECKING, Union | 3 | from typing import Any, Callable, Optional, TYPE_CHECKING, Union |
| 4 | import hashlib | 4 | import hashlib |
| 5 | -import sympy | ||
| 6 | 5 | ||
| 7 | import torch | 6 | import torch |
| 8 | from torch._inductor import config | 7 | from torch._inductor import config |
| @@ -36,6 +35,7 @@ class _NPUKernelCodegenMixin: | |||
| 36 | from leaking into subgraphs. | 35 | from leaking into subgraphs. |
| 37 | """ | 36 | """ |
| 38 | 37 | ||
| 38 | + | ||
| 39 | # generate numel expr for range_tree_node | 39 | # generate numel expr for range_tree_node |
| 40 | def generate_node_numel_expr(self, kernel_name: str, node, numel_expr): | 40 | def generate_node_numel_expr(self, kernel_name: str, node, numel_expr): |
| 41 | expr = f"{kernel_name}_{node.name}_numel" | 41 | expr = f"{kernel_name}_{node.name}_numel" |
| @@ -270,9 +270,7 @@ class FastSyncBatchNorm(_BatchNorm): | |||
| 270 | ) | 270 | ) |
| 271 | 271 | ||
| 272 | def forward(self, input1: Tensor) -> Tensor: | 272 | def forward(self, input1: Tensor) -> Tensor: |
| 273 | - # currently NPU or GPU input1 is supported | 273 | + |
| 274 | - if not input1.is_cuda and not input1.is_npu: | ||
| 275 | - raise ValueError("SyncBatchNorm expected input1 tensor to be on NPU or GPU" + ops_error(ErrCode.VALUE)) | ||
| 276 | 274 | ||
| 277 | self._check_input_dim(input1) | 275 | self._check_input_dim(input1) |
| 278 | self._check_non_zero_input_channels(input1) | 276 | self._check_non_zero_input_channels(input1) |
| @@ -324,6 +322,10 @@ class FastSyncBatchNorm(_BatchNorm): | |||
| 324 | and torch.distributed.is_initialized() | 322 | and torch.distributed.is_initialized() |
| 325 | ) | 323 | ) |
| 326 | if need_sync: | 324 | if need_sync: |
| 325 | + # currently NPU or GPU input1 is supported | ||
| 326 | + if not input1.is_cuda and not input1.is_npu: | ||
| 327 | + raise ValueError("SyncBatchNorm expected input1 tensor to be on NPU or GPU" + ops_error(ErrCode.VALUE)) | ||
| 328 | + | ||
| 327 | process_group = torch.distributed.group.WORLD | 329 | process_group = torch.distributed.group.WORLD |
| 328 | if self.process_group: | 330 | if self.process_group: |
| 329 | process_group = self.process_group | 331 | process_group = self.process_group |
| @@ -12,6 +12,7 @@ from torch.multiprocessing.reductions import ( | |||
| 12 | rebuild_tensor, | 12 | rebuild_tensor, |
| 13 | storage_from_cache, | 13 | storage_from_cache, |
| 14 | rebuild_meta_tensor, | 14 | rebuild_meta_tensor, |
| 15 | + reduce_nested_tensor, | ||
| 15 | ) | 16 | ) |
| 16 | 17 | ||
| 17 | import torch_npu | 18 | import torch_npu |
| @@ -95,6 +96,11 @@ def rebuild_npu_tensor( | |||
| 95 | 96 | ||
| 96 | 97 | ||
| 97 | def _npu_reduce_tensor(tensor): | 98 | def _npu_reduce_tensor(tensor): |
| 99 | + from torch.nested._internal.nested_tensor import NestedTensor | ||
| 100 | + | ||
| 101 | + # 特殊处理：嵌套张量（Nested Tensor）走专用归约 | ||
| 102 | + if tensor.is_nested and not isinstance(tensor, NestedTensor): | ||
| 103 | + return reduce_nested_tensor(tensor) | ||
| 98 | storage = tensor._typed_storage() | 104 | storage = tensor._typed_storage() |
| 99 | 105 | ||
| 100 | if tensor.requires_grad and not tensor.is_leaf: | 106 | if tensor.requires_grad and not tensor.is_leaf: |