已合并
【fix】batcnnorm_check_fix #36324
cuiduo创建于 5月21日
【fix】batcnnorm_check_fix #36324
已合并
共 4 个文件变更+14-6
| @@ -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 copy | |||
| 2 | import hashlib | 2 | import hashlib |
| 3 | import os | 3 | import os |
| 4 | 4 | ||
| 5 | -import sympy | ||
| 6 | 5 | ||
| 7 | import torch | 6 | import torch |
| 8 | import torch_npu.npu.aclnn | 7 | import torch_npu.npu.aclnn |
| @@ -35,7 +34,7 @@ class _NPUKernelCodegenMixin: | |||
| 35 | (AOT debug / aclnn initialization / whole-graph benchmark harness, etc.) | 34 | (AOT debug / aclnn initialization / whole-graph benchmark harness, etc.) |
| 36 | from leaking into subgraphs. | 35 | from leaking into subgraphs. |
| 37 | """ | 36 | """ |
| 38 | - | 37 | + |
| 39 | # generate numel expr for range_tree_node | 38 | # generate numel expr for range_tree_node |
| 40 | def generate_node_numel_expr(self, kernel_name: str, node, numel_expr): | 39 | def generate_node_numel_expr(self, kernel_name: str, node, numel_expr): |
| 41 | expr = f"{kernel_name}_{node.name}_numel" | 40 | expr = f"{kernel_name}_{node.name}_numel" |
| @@ -270,10 +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 | - | ||
| 277 | self._check_input_dim(input1) | 274 | self._check_input_dim(input1) |
| 278 | self._check_non_zero_input_channels(input1) | 275 | self._check_non_zero_input_channels(input1) |
| 279 | 276 | ||
| @@ -324,6 +321,10 @@ class FastSyncBatchNorm(_BatchNorm): | |||
| 324 | and torch.distributed.is_initialized() | 321 | and torch.distributed.is_initialized() |
| 325 | ) | 322 | ) |
| 326 | if need_sync: | 323 | if need_sync: |
| 324 | + # currently NPU or GPU input1 is supported | ||
| 325 | + if not input1.is_cuda and not input1.is_npu: | ||
| 326 | + raise ValueError("SyncBatchNorm expected input1 tensor to be on NPU or GPU" + ops_error(ErrCode.VALUE)) | ||
| 327 | + | ||
| 327 | process_group = torch.distributed.group.WORLD | 328 | process_group = torch.distributed.group.WORLD |
| 328 | if self.process_group: | 329 | if self.process_group: |
| 329 | process_group = self.process_group | 330 | 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,12 @@ 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 | + if tensor.is_nested and not isinstance(tensor, NestedTensor): | ||
| 102 | + return reduce_nested_tensor(tensor) | ||
| 103 | + | ||
| 104 | + storage = tensor._typed_storage() | ||
| 98 | storage = tensor._typed_storage() | 105 | storage = tensor._typed_storage() |
 | |||
| 99 | 106 | ||
| 100 | if tensor.requires_grad and not tensor.is_leaf: | 107 | if tensor.requires_grad and not tensor.is_leaf: |
🟡 Medium Priority
第 104 行是 PR 新增的行，与第 105 行原有代码完全相同。新增行的赋值结果被第 105 行立即覆盖，属于无意义的死存储（dead store），且表明该行是误加而非有意为之。虽然不影响最终结果，但会引入不必要的计算开销且降低代码可读性。