已合并
add autocast dtype get #20627
AtomGit-Bot创建于 2025年4月23日
add autocast dtype get #20627
已合并
从refs/pull/20627/head合入到master
共 2 个文件变更+6-1
| @@ -455,6 +455,7 @@ class TestAmp(TestCase): | |||
| 455 | 455 | ||
| 456 | 456 | ||
| 457 | def forward(ctx, a, b): | 457 | def forward(ctx, a, b): |
| 458 | + self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu")) | ||
| 458 | self.assertTrue(a.dtype is torch.float32) | 459 | self.assertTrue(a.dtype is torch.float32) |
| 459 | self.assertTrue(b.dtype is torch.float32) | 460 | self.assertTrue(b.dtype is torch.float32) |
| 460 | self.assertTrue(torch.npu.is_autocast_enabled()) | 461 | self.assertTrue(torch.npu.is_autocast_enabled()) |
| @@ -464,6 +465,7 @@ class TestAmp(TestCase): | |||
| 464 | 465 | ||
| 465 | 466 | ||
| 466 | def backward(ctx, grad): | 467 | def backward(ctx, grad): |
| 468 | + self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu")) | ||
| 467 | self.assertTrue(torch.npu.is_autocast_enabled()) | 469 | self.assertTrue(torch.npu.is_autocast_enabled()) |
| 468 | a, b = ctx.saved_tensors | 470 | a, b = ctx.saved_tensors |
| 469 | return grad.mm(b.t()), a.t().mm(grad) | 471 | return grad.mm(b.t()), a.t().mm(grad) |
| @@ -484,6 +486,7 @@ class TestAmp(TestCase): | |||
| 484 | 486 | ||
| 485 | 487 | ||
| 486 | def forward(ctx, a, container, expect_type): | 488 | def forward(ctx, a, container, expect_type): |
| 489 | + self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu")) | ||
| 487 | b = container[1][0] | 490 | b = container[1][0] |
| 488 | self.assertTrue(a.dtype is expect_type) | 491 | self.assertTrue(a.dtype is expect_type) |
| 489 | self.assertTrue(b.dtype is expect_type) | 492 | self.assertTrue(b.dtype is expect_type) |
| @@ -494,6 +497,7 @@ class TestAmp(TestCase): | |||
| 494 | 497 | ||
| 495 | 498 | ||
| 496 | def backward(ctx, grad): | 499 | def backward(ctx, grad): |
| 500 | + self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu")) | ||
| 497 | a, b = ctx.saved_tensors | 501 | a, b = ctx.saved_tensors |
| 498 | return grad.mm(b.t()), None, None | 502 | return grad.mm(b.t()), None, None |
| 499 | 503 | ||
| @@ -108,6 +108,7 @@ def custom_fwd(fwd=None, **kwargs): | |||
| 108 | 108 | ||
| 109 | 109 | ||
| 110 | def decorate_fwd(*args, **kwargs): | 110 | def decorate_fwd(*args, **kwargs): |
| 111 | + args[0]._dtype = torch.get_autocast_dtype("npu") | ||
| 111 | if cast_inputs is None: | 112 | if cast_inputs is None: |
| 112 | args[0]._fwd_used_autocast = torch_npu._C.is_autocast_enabled() | 113 | args[0]._fwd_used_autocast = torch_npu._C.is_autocast_enabled() |
| 113 | return fwd(*args, **kwargs) | 114 | return fwd(*args, **kwargs) |
| @@ -136,7 +137,7 @@ def custom_bwd(bwd): | |||
| 136 | 137 | ||
| 137 | 138 | ||
| 138 | def decorate_bwd(*args, **kwargs): | 139 | def decorate_bwd(*args, **kwargs): |
| 139 | - with autocast(args[0]._fwd_used_autocast): | 140 | + with autocast(args[0]._fwd_used_autocast, dtype=args[0]._dtype): |
| 140 | return bwd(*args, **kwargs) | 141 | return bwd(*args, **kwargs) |
| 141 | 142 | ||
| 142 | return decorate_bwd | 143 | return decorate_bwd |