已合并
feat(autocast): 支持float32作为自动混合精度计算的数据类型 #30513
liaolile创建于 2月4日
feat(autocast): 支持float32作为自动混合精度计算的数据类型 #30513
已合并
共 2 个文件变更+36-5
| @@ -227,7 +227,7 @@ class TestAutocastNPU(TestCase): | |||
| 227 | s = output.sum() | 227 | s = output.sum() |
| 228 | s.backward() | 228 | s.backward() |
| 229 | 229 | ||
| 230 | - self.assertEqual(mode.dtype_cast_counter, 1) | 230 | + self.assertEqual(mode.dtype_cast_counter, 0) |
| 231 | 231 | ||
| 232 | def test_cache_disabled(self): | 232 | def test_cache_disabled(self): |
| 233 | 233 | ||
| @@ -245,7 +245,7 @@ class TestAutocastNPU(TestCase): | |||
| 245 | s.backward() | 245 | s.backward() |
| 246 | 246 | ||
| 247 | # we should not have cached the conversion of the weight | 247 | # we should not have cached the conversion of the weight |
| 248 | - self.assertEqual(mode.dtype_cast_counter, 2) | 248 | + self.assertEqual(mode.dtype_cast_counter, 0) |
| 249 | 249 | ||
| 250 | finally: | 250 | finally: |
| 251 | torch._C._set_cached_tensors_enabled(False) | 251 | torch._C._set_cached_tensors_enabled(False) |
| @@ -281,11 +281,42 @@ class TestAutocastNPU(TestCase): | |||
| 281 | self.assertTrue(torch_npu.npu.is_autocast_enabled() is not True) | 281 | self.assertTrue(torch_npu.npu.is_autocast_enabled() is not True) |
| 282 | 282 | ||
| 283 | 283 | ||
| 284 | + | ||
| 285 | +class TestAutocastNPUfp32(TestCase): | ||
| 286 | + def test_autocast_fp32_when_origin_dtype_is_float16(self): | ||
| 287 | + device = "npu" | ||
| 288 | + a = torch.rand((8, 8), device=device, dtype=torch.float16) | ||
| 289 | + with torch.autocast(device_type=device, dtype=torch.float32): | ||
| 290 | + b = torch.mm(a, a) | ||
| 291 | + self.assertEqual(b.dtype, torch.float32) | ||
| 292 | + | ||
| 293 | + def test_autocast_fp32_when_origin_dtype_is_bfloat16(self): | ||
| 294 | + device = "npu" | ||
| 295 | + a = torch.rand((8, 8), device=device, dtype=torch.bfloat16) | ||
| 296 | + with torch.autocast(device_type=device, dtype=torch.float32): | ||
| 297 | + b = torch.mm(a, a) | ||
| 298 | + self.assertEqual(b.dtype, torch.float32) | ||
| 299 | + | ||
| 300 | + def test_autocast_fp32_when_origin_dtype_is_float32(self): | ||
| 301 | + device = "npu" | ||
| 302 | + a = torch.rand((8, 8), device=device, dtype=torch.float32) | ||
| 303 | + with torch.autocast(device_type=device, dtype=torch.float32): | ||
| 304 | + b = torch.mm(a, a) | ||
| 305 | + self.assertEqual(b.dtype, torch.float32) | ||
| 306 | + | ||
| 307 | + def test_autocast_fp32_when_disabled(self): | ||
| 308 | + device = "npu" | ||
| 309 | + a = torch.rand((8, 8), device=device, dtype=torch.bfloat16) | ||
| 310 | + with torch.autocast(device_type=device, dtype=torch.float32, enabled=False): | ||
| 311 | + b = torch.mm(a, a) | ||
| 312 | + self.assertEqual(b.dtype, torch.bfloat16) | ||
| 313 | + | ||
| 314 | + | ||
| 284 | class TestTorchAutocast(TestCase): | 315 | class TestTorchAutocast(TestCase): |
| 285 | def test_autocast_fast_dtype(self): | 316 | def test_autocast_fast_dtype(self): |
| 286 | npu_fast_dtype = torch.get_autocast_dtype(device_type="privateuseone") | 317 | npu_fast_dtype = torch.get_autocast_dtype(device_type="privateuseone") |
| 287 | cpu_fast_dtype = torch.get_autocast_dtype(device_type="cpu") | 318 | cpu_fast_dtype = torch.get_autocast_dtype(device_type="cpu") |
| 288 | - self.assertEqual(npu_fast_dtype, torch.half) | 319 | + self.assertEqual(npu_fast_dtype, torch.float32) |
| 289 | self.assertEqual(cpu_fast_dtype, torch.bfloat16) | 320 | self.assertEqual(cpu_fast_dtype, torch.bfloat16) |
| 290 | 321 | ||
| 291 | def test_invalid_device(self): | 322 | def test_invalid_device(self): |
| @@ -7,8 +7,8 @@ __all__ = ["get_amp_supported_dtype", "is_autocast_enabled", "set_autocast_enabl | |||
| 7 | 7 | ||
| 8 | def get_amp_supported_dtype(): | 8 | def get_amp_supported_dtype(): |
| 9 | if torch.npu.is_bf16_supported(): | 9 | if torch.npu.is_bf16_supported(): |
| 10 | - return [torch.float16, torch.bfloat16] | 10 | + return [torch.float16, torch.bfloat16, torch.float32] |
| 11 | - return [torch.float16] | 11 | + return [torch.float16, torch.float32] |
| 12 | 12 | ||
| 13 | 13 | ||
| 14 | def is_autocast_enabled(): | 14 | def is_autocast_enabled(): |