From 706ec48a6c26ec459f9e44b2b7c925c28b5bc12b Mon Sep 17 00:00:00 2001 From: TomTheHand Date: Sat, 17 Oct 2026 19:25:52 -0400 Subject: [PATCH] Enable bf16 autocast on Intel XPU devices Treat XPU as a bf16-native accelerator: autocast follows the configured dtype instead of silently staying fp32, and XPU never consults the CUDA/ROCm capability queries (a colocated NVIDIA card must not decide Intel precision policy, including when MINAGI_GPU_BF16=0 disables the native-bf16 predicate). MINAGI_GPU_BF16=1/0 still overrides. Validated on Arc Pro B60 with torch 2.14.1+xpu: bf16 forward/backward and an SGD update run and stay finite. Adds CPU regression tests and an opt-in MINAGI_TEST_XPU=1 hardware smoke test. --- minagi/precision.py | 12 ++++-- tests/README.md | 28 +++++++++++++ tests/test_precision.py | 84 +++++++++++++++++++++++++++++++++++++ tests/test_precision_xpu.py | 42 +++++++++++++++++++ 4 files changed, 163 insertions(+), 3 deletions(-) create mode 100644 tests/README.md create mode 100644 tests/test_precision.py create mode 100644 tests/test_precision_xpu.py diff --git a/minagi/precision.py b/minagi/precision.py index 33853e2..b59876d 100644 --- a/minagi/precision.py +++ b/minagi/precision.py @@ -107,7 +107,9 @@ def _device_type(device): def _gpu_info(device=None): - """(properties, is it AMD) for the GPU `device` names; None without one.""" + """CUDA/ROCm (properties, is it AMD); None for XPU or without a GPU.""" + if _device_type(device) == "xpu": + return None # never query a colocated CUDA GPU for XPU try: idx = getattr(device, "index", None) if idx is None: @@ -134,11 +136,15 @@ def gpu_bf16_native(device=None): Whether this GPU multiplies bfloat16 in hardware: NVIDIA from Ampere (compute capability 8) on, AMD from RDNA3 (gfx11, gfx12) and the CDNA accelerators (gfx908, gfx90a, gfx94x, gfx95x). Not RDNA1 or 2, not Vega, - not NVIDIA before Ampere. MINAGI_GPU_BF16=1 or =0 overrides the detection. + not NVIDIA before Ampere. Intel XPU uses bf16 by default (validated on + Arc Pro B60). MINAGI_GPU_BF16=1 or =0 overrides the detection; the + CUDA/ROCm memory-based fp32 fallback does not apply to XPU. """ force = os.environ.get("MINAGI_GPU_BF16") if force is not None: return force.strip().lower() not in ("0", "", "false", "no", "off") + if _device_type(device) == "xpu": + return True # use bf16 for Intel XPU; override above info = _gpu_info(device) if info is None: return True # nothing to ask: the rule as it always was @@ -171,7 +177,7 @@ def autocast_on(device=None): dev = _device_type(device) if dt is torch.float32: return False - if dev == "cuda": + if dev in ("cuda", "xpu"): return not gpu_fp32_instead(device) return dev == "cpu" and dt is torch.bfloat16 and cpu_bf16_native() diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..4c62788 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,28 @@ +# Precision regression tests + +Run from the repository root with PyTorch installed: + +```sh +python -m unittest discover -s tests -v +``` + +The policy tests use CPU-side mocks and need no accelerator. They cover +XPU device strings/objects, bf16/fp16/fp32 selection, autocast dispatch, +GPU overrides, and unchanged CPU/CUDA/ROCm precision policy. + +Optionally test a fresh 16-to-8 linear layer on `xpu:0`: + +```sh +MINAGI_TEST_XPU=1 python -m unittest discover -s tests -v +``` + +This checks bf16 output, finite fp32 gradients/parameters, and an SGD update; +it never loads model weights or training data. Validated on Arc Pro B60 +with PyTorch 2.14.1+xpu, not a claim about all Intel GPUs or training speed. + +XPU now uses the configured low-precision autocast path. XPU bf16 policy +does not use CUDA properties (including on mixed CUDA/XPU systems). +`MINAGI_GPU_BF16` still overrides the native-bf16 policy predicate, but the +CUDA/ROCm VRAM-based fp32 fallback does not apply to XPU; set the compute +dtype to fp32 to disable XPU autocast. fp16 remains an explicit choice; +this change does not add gradient scaling or recommend fp16 training. diff --git a/tests/test_precision.py b/tests/test_precision.py new file mode 100644 index 0000000..5249c7a --- /dev/null +++ b/tests/test_precision.py @@ -0,0 +1,84 @@ +"""Precision policy tests; no accelerator or saved model is required.""" +import os +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from minagi import precision as p + + +class PrecisionTests(unittest.TestCase): + def setUp(self): + self.dtype = p.compute_dtype() + self.env = patch.dict(os.environ, {}, clear=False) + self.env.start() + os.environ.pop("MINAGI_GPU_BF16", None) + os.environ.pop("MINAGI_CPU_BF16", None) + self.addCleanup(self.env.stop) + self.addCleanup(p.set_compute_dtype, self.dtype) + + def test_xpu_bf16_policy_without_cuda_queries(self): + p.set_compute_dtype("bf16") + with patch.object(torch.cuda, "current_device", side_effect=AssertionError("CUDA query")) as current, patch.object(torch.cuda, "get_device_properties", side_effect=AssertionError("CUDA query")) as props: + for device in ("xpu", "xpu:1", torch.device("xpu:1")): + with self.subTest(device=device): + self.assertTrue(p.gpu_bf16_native(device)) + self.assertTrue(p.autocast_on(device)) + self.assertFalse(p.gpu_fp32_instead(device)) + self.assertEqual(p.describe(device), "bf16") + self.assertEqual(p.dispatch_dtype(device, torch.float32), torch.bfloat16) + current.assert_not_called() + props.assert_not_called() + + def test_xpu_overrides_do_not_inspect_cuda(self): + p.set_compute_dtype("bf16") + # A colocated CUDA card must not decide XPU precision or memory policy. + cuda_props = SimpleNamespace(major=7, total_memory=24 * 2**30) + with patch.object(torch.cuda, "current_device", return_value=0) as current, patch.object(torch.cuda, "get_device_properties", return_value=cuda_props) as props, patch.dict(p._GPU_INFO, {}, clear=True): + for force, native in (("0", False), ("false", False), ("1", True)): + with self.subTest(force=force), patch.dict(os.environ, {"MINAGI_GPU_BF16": force}): + self.assertEqual(p.gpu_bf16_native(torch.device("xpu:0")), native) + self.assertIsNone(p._gpu_info(torch.device("xpu:0"))) + self.assertTrue(p.autocast_on("xpu:0")) + self.assertFalse(p.gpu_fp32_instead("xpu:0")) + self.assertEqual(p.describe("xpu:0"), "bf16") + current.assert_not_called() + props.assert_not_called() + + def test_xpu_dtype_and_amp_dispatch(self): + for name, dtype, enabled in (("fp32", torch.float32, False), ("bf16", torch.bfloat16, True), ("fp16", torch.float16, True)): + with self.subTest(dtype=name): + p.set_compute_dtype(name) + self.assertEqual(p.autocast_on("xpu:0"), enabled) + self.assertEqual(p.describe("xpu:0"), name) + self.assertEqual(p.dispatch_dtype("xpu:0", torch.float64), dtype if enabled else torch.float64) + with patch.object(torch, "autocast") as autocast: + p.amp(torch.device("xpu:0")) + autocast.assert_called_once_with("xpu", dtype=dtype, enabled=enabled) + + def test_cpu_policy_unchanged(self): + for force in ("0", "1"): + with patch.dict(os.environ, {"MINAGI_CPU_BF16": force}): + for dtype in (torch.float32, torch.float16, torch.bfloat16): + p.set_compute_dtype(dtype) + self.assertEqual(p.autocast_on("cpu"), force == "1" and dtype is torch.bfloat16) + + def test_cuda_and_rocm_policy_unchanged(self): + p.set_compute_dtype("bf16") + for major, arch, amd, native in ((7, "", False, False), (8, "", False, True), (0, "gfx1030", True, False), (0, "gfx1100", True, True), (0, "gfx90a", True, True)): + for gib in (8, 24): + with self.subTest(major=major, arch=arch, gib=gib): + props = SimpleNamespace(major=major, minor=0, gcnArchName=arch, total_memory=gib * 2**30) + with patch.object(p, "_gpu_info", return_value=(props, amd)): + self.assertEqual(p.gpu_bf16_native("cuda"), native) + self.assertEqual(p.autocast_on("cuda"), native or gib < p.FP32_MIN_GIB) + self.assertEqual(p.autocast_on(), native or gib < p.FP32_MIN_GIB) + with patch.dict(os.environ, {"MINAGI_GPU_BF16": "1"}): + self.assertTrue(p.autocast_on("cuda")) + with patch.dict(os.environ, {"MINAGI_GPU_BF16": "0"}): + self.assertEqual(p.autocast_on("cuda"), gib < p.FP32_MIN_GIB) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_precision_xpu.py b/tests/test_precision_xpu.py new file mode 100644 index 0000000..468df5c --- /dev/null +++ b/tests/test_precision_xpu.py @@ -0,0 +1,42 @@ +"""Opt in with MINAGI_TEST_XPU=1; only fresh tiny tensors, never checkpoints.""" +import os +import unittest +from unittest.mock import patch + +import torch +from minagi import precision as p + + +@unittest.skipUnless(os.environ.get("MINAGI_TEST_XPU") == "1", "opt-in XPU smoke test") +class XpuSmokeTests(unittest.TestCase): + def test_bf16_forward_backward(self): + self.assertTrue(torch.xpu.is_available(), "XPU requested but unavailable") + original = p.compute_dtype() + self.addCleanup(p.set_compute_dtype, original) + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("MINAGI_GPU_BF16", None) + p.set_compute_dtype("bf16") + device = torch.device("xpu:0") + model = torch.nn.Linear(16, 8, device=device) + inputs = torch.ones(4, 16, device=device) + before = model.weight.detach().clone() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + with p.amp(device): + output = model(inputs) + loss = output.float().square().mean() + self.assertEqual(output.dtype, torch.bfloat16) + self.assertTrue(torch.isfinite(loss).item()) + loss.backward() + for parameter in model.parameters(): + self.assertEqual(parameter.dtype, torch.float32) + self.assertEqual(parameter.grad.dtype, torch.float32) + self.assertTrue(torch.isfinite(parameter.grad).all().item()) + self.assertGreater(model.weight.grad.abs().sum().item(), 0) + optimizer.step() + torch.xpu.synchronize(device) + self.assertFalse(torch.equal(before, model.weight)) + print(f"XPU smoke: torch={torch.__version__}; device={torch.xpu.get_device_name(device)}; output={output.dtype}; finite gradients and SGD update OK") + + +if __name__ == "__main__": + unittest.main()