已合并
【fix】batcnnorm_check_fix #36321
cuiduo创建于 5月21日
【fix】batcnnorm_check_fix #36321
已合并
共 5 个文件变更+11-6
| @@ -8,7 +8,6 @@ from testutils import TestUtils | |||
| 8 | import torch_npu | 8 | import torch_npu |
| 9 | import torch_npu._inductor | 9 | import torch_npu._inductor |
| 10 | 10 | ||
| 11 | - | ||
| 12 | class TestRunWithRngState(TestUtils): | 11 | class TestRunWithRngState(TestUtils): |
| 13 | def op_calc(self, current_state, like, device, dtype): | 12 | def op_calc(self, current_state, like, device, dtype): |
| 14 | res1 = torch._prims.rng_prims.run_with_rng_state( | 13 | res1 = torch._prims.rng_prims.run_with_rng_state( |
| @@ -1,4 +1,4 @@ | |||
| 1 | -import os | 1 | +import os |
| 2 | import copy | 2 | import copy |
| 3 | import dataclasses | 3 | import dataclasses |
| 4 | from typing import Optional, Union | 4 | from typing import Optional, Union |
| @@ -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,9 @@ 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)) | ||
| 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 |
| @@ -175,7 +175,7 @@ def _do_wrapper_libraries_func(json_dict): | |||
| 175 | 175 | ||
| 176 | def _wrapper_cuda(fn): | 176 | def _wrapper_cuda(fn): |
| 177 | 177 | ||
| 178 | - def decorated(*args, **kwargs): | 178 | + def decorated(*args, **kwargs): |
| 179 | replace_int = fn.__name__ in ['to', 'to_empty'] | 179 | replace_int = fn.__name__ in ['to', 'to_empty'] |
| 180 | if args: | 180 | if args: |
| 181 | args_new = list(args) | 181 | args_new = list(args) |
| @@ -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,10 @@ 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 | + if tensor.is_nested and not isinstance(tensor, NestedTensor): | ||
| 101 | + return reduce_nested_tensor(tensor) | ||
| 102 | + | ||
| 98 | storage = tensor._typed_storage() | 103 | storage = tensor._typed_storage() |
| 99 | 104 | ||
| 100 | if tensor.requires_grad and not tensor.is_leaf: | 105 | if tensor.requires_grad and not tensor.is_leaf: |