已合并
[master]fix dropout #30017
daifu1234567创建于 1月26日
[master]fix dropout #30017
已合并
共 3 个文件变更+70-1
| @@ -116,5 +116,47 @@ class TestDTensorCustomOps(DTensorTestBase): | |||
| 116 | self.assertEqual(dist_res.to_local().shape, local_result.shape) | 116 | self.assertEqual(dist_res.to_local().shape, local_result.shape) |
| 117 | 117 | ||
| 118 | 118 | ||
| 119 | + def exec_dropout_backward(self, p, train): | ||
| 120 | + x = torch.randn(12, 5).npu() | ||
| 121 | + x.requires_grad = True | ||
| 122 | + output = torch.nn.functional.dropout(x, p=p, training=train).sum.backward() | ||
| 123 | + x_grad = x.grad | ||
| 124 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 125 | + x_dtensor = distribute_tensor(npu_input, device_mesh, [Shard(0)]) | ||
| 126 | + x_dtensor.requires_grad = True | ||
| 127 | + dist_res = torch.nn.functional.dropout(x, p=p, training=train).sum.backward() | ||
| 128 | + x_grad_dtensor = x_dtensor.grad.redistribute(device_mesh, [Replicate()]) | ||
| 129 | + self.assertEqual(x_grad_dtensor.to_local().shape, x_grad.shape) | ||
| 130 | + | ||
| 131 | + | ||
| 132 | + | ||
| 133 | + def test_dtensor_dropout_backward(self): | ||
| 134 | + p_list = [0, 0.2, 0.5, 0.7, 1] | ||
| 135 | + train_list = [True, False] | ||
| 136 | + for p in p_list: | ||
| 137 | + for train in train_list: | ||
| 138 | + self.exec_dropout_backward(p, train) | ||
| 139 | + | ||
| 140 | + | ||
| 141 | + | ||
| 142 | + def test_dtensor_npu_transpose(self): | ||
| 143 | + npu_input = torch.randn(5, 3, 6, 4).npu() | ||
| 144 | + perm = [1, 0, 2, 3] | ||
| 145 | + local_result = torch_npu.npu_transpose(npu_input, perm) | ||
| 146 | + | ||
| 147 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 148 | + shard0_spec = Shard(0) | ||
| 149 | + shard1_spec = Shard(1) | ||
| 150 | + replica_spec = Replicate() | ||
| 151 | + | ||
| 152 | + placement_specs = [shard0_spec, shard1_spec, replica_spec] | ||
| 153 | + for spec in placement_specs: | ||
| 154 | + dt_input = distribute_tensor(npu_input, device_mesh, [spec]) | ||
| 155 | + dist_res: DTensor = cast(DTensor, torch_npu.npu_transpose(dt_input, perm)).redistribute( | ||
| 156 | + device_mesh, [replica_spec] | ||
| 157 | + ) | ||
| 158 | + self.assertEqual(dist_res.to_local().shape, local_result.shape) | ||
| 159 | + | ||
| 160 | + | ||
| 119 | if __name__ == '__main__': | 161 | if __name__ == '__main__': |
| 120 | run_tests() | 162 | run_tests() |
| @@ -1,6 +1,7 @@ | |||
| 1 | import torch._ops | 1 | import torch._ops |
| 2 | from torch._inductor.decomposition import decompositions, pw_cast_for_opmath | 2 | from torch._inductor.decomposition import decompositions, pw_cast_for_opmath |
| 3 | from torch._inductor.decomposition import register_decomposition | 3 | from torch._inductor.decomposition import register_decomposition |
| 4 | +from torch._prims_common.wrappers import out_wrapper | ||
| 4 | 5 | ||
| 5 | from .lowering import _init_set | 6 | from .lowering import _init_set |
| 6 | 7 | ||
| @@ -15,7 +16,9 @@ DECOMPOSITION_OVERLOAD_OP = [ | |||
| 15 | aten._log_softmax_backward_data, | 16 | aten._log_softmax_backward_data, |
| 16 | aten.embedding_dense_backward, | 17 | aten.embedding_dense_backward, |
| 17 | aten.addmm, | 18 | aten.addmm, |
| 18 | - aten.gelu | 19 | + aten.gelu, |
| 20 | + aten.native_dropout, | ||
| 21 | + aten.native_dropout_backward | ||
| 19 | ] | 22 | ] |
| 20 | 23 | ||
| 21 | 24 | ||
| @@ -36,3 +39,25 @@ def _register_npu_inductor_decompositons(): | |||
| 36 | def erfc(x): | 39 | def erfc(x): |
| 37 | tensor = torch.ones_like(x) - torch.exp(x) | 40 | tensor = torch.ones_like(x) - torch.exp(x) |
| 38 | return tensor | 41 | return tensor |
| 42 | + | ||
| 43 | + | ||
| 44 | + | ||
| 45 | + def native_dropout(tensor_input, p, train): | ||
| 46 | + if torch._inductor.config.fallback_random: | ||
| 47 | + if train and p != 0: | ||
| 48 | + return torch.ops.npu._npu_dropout(tensor_input, p) | ||
| 49 | + return (tensor_input, torch.ones_like(tensor_input, dtype=torch.bool)) | ||
| 50 | + else: | ||
| 51 | + from torch._decomp.decompositions import native_dropout | ||
| 52 | + return native_dropout(tensor_input, p, train) | ||
| 53 | + | ||
| 54 | + | ||
| 55 | + | ||
| 56 | + def native_dropout_backward(grad_output, mask, scale): | ||
| 57 | + if torch._inductor.config.fallback_random: | ||
| 58 | + p = 1 if scale == 0 else (1 - 1 / scale) | ||
| 59 | + r = torch.ops.npu.npu_dropout_backward(grad_output, mask, p) | ||
| 60 | + return r | ||
| 61 | + else: | ||
| 62 | + from torch._decomp.decompositions import native_dropout_backward | ||
| 63 | + return native_dropout_backward(grad_output, mask, scale) | ||
代码逻辑和结构: 新增的
config.fallback_random = True配置项在代码中缺乏上下文解释和用途说明。从变量名来看，它似乎控制随机数生成的回退行为，但该配置项在后续代码中没有任何使用或引用，也没有相关的注释说明其具体作用。这可能导致开发者困惑，不清楚这个配置项的实际影响和适用场景。问题类型: 代码逻辑和结构 文件路径:
torch_npu/_inductor/config.py行号: 16 问题代码:config.fallback_random = True修改建议:
此评论由代码审查工具自动生成