From cd6cd8220e58d1e503f15ac758bee10bd8f66ffb Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Mon, 28 Sep 2026 22:20:15 +0800 Subject: [PATCH 1/4] fix(amp): resolve autocast/GradScaler device type for non-CUDA accelerators funasr/utils/amp.py pinned device_type="cuda" in the torch.amp compatibility shim, so on Ascend NPU (and XPU/MPS) builds torch.amp only warned "CUDA is not available ... Disabling autocast" and the context became a no-op: use_fp16/use_bf16 training silently ran in float32 and GradScaler(enabled=True) was silently disabled. Resolve the device type from the running accelerator (torch.accelerator), the same policy funasr/models/fun_asr_nano/device_utils.py already applies to the inference path, and keep "cuda" as the historical no-op default when no accelerator is available so CPU behaviour is unchanged. GradScaler prefers an accelerator-specific implementation (e.g. torch.npu.amp.GradScaler) when one exists. Signed-off-by: li-lizhe <147392333@qq.com> --- funasr/utils/amp.py | 92 +++++++++++++++++-- tests/test_amp_device_type.py | 165 ++++++++++++++++++++++++++++++++++ 2 files changed, 251 insertions(+), 6 deletions(-) create mode 100644 tests/test_amp_device_type.py diff --git a/funasr/utils/amp.py b/funasr/utils/amp.py index 2469d12cc5..b28ff03e4d 100644 --- a/funasr/utils/amp.py +++ b/funasr/utils/amp.py @@ -14,26 +14,106 @@ Note: ``torch.amp.autocast`` requires a ``device_type`` argument (e.g. ``'cuda'``) that the deprecated ``torch.cuda.amp.autocast`` did not, so -the shim supplies ``device_type="cuda"`` by default for callers that use -the old signature (``with autocast(enabled=..., dtype=...)``). +the shim resolves ``device_type`` for callers that use the old signature +(``with autocast(enabled=..., dtype=...)``) — see +``resolve_amp_device_type``. Callers that know the device can still pass +``device_type=...`` explicitly. + +Device type resolution +---------------------- + +The deprecated ``torch.cuda.amp`` API could only target CUDA, so the first +version of this shim pinned ``device_type="cuda"``. Off CUDA that is not an +error: ``torch.amp`` only warns + + UserWarning: CUDA is not available or torch_xla is imported. Disabling autocast. + +and the context becomes a no-op, so ``use_fp16``/``use_bf16`` training silently +ran in float32 on Ascend NPU / XPU / MPS builds. ``GradScaler(enabled=True)`` +was silently disabled the same way, so loss scaling was missing too. + +Both now default to the accelerator that is actually active +(``torch.accelerator``) — the same device types that +``funasr/models/fun_asr_nano/device_utils.py`` accepts for the Nano inference +path — and keep the historical ``"cuda"`` (a no-op on CPU-only builds) when no +accelerator is available, so CPU behaviour is unchanged. """ import torch torch_amp = getattr(torch, "amp", None) + +# Device types that FunASR is willing to enter an AMP context on; anything else +# keeps the historical "cuda" no-op instead of risking an unsupported context. +_ACCELERATOR_DEVICE_TYPES = ("cuda", "npu", "xpu", "mps") + + +def resolve_amp_device_type(): + """Return the ``torch.amp`` device type of the active accelerator. + + ``torch.accelerator`` reports the accelerator torch was built/registered + for (``cuda``/``npu``/``xpu``/``mps``); it is ``None`` on CPU-only builds, + where ``"cuda"`` is kept as the historical no-op default. + """ + accelerator = getattr(torch, "accelerator", None) + if accelerator is not None: + try: + current = accelerator.current_accelerator() + except (AttributeError, RuntimeError): + current = None + device_type = getattr(current, "type", None) if current is not None else None + if device_type and str(device_type).lower() in _ACCELERATOR_DEVICE_TYPES: + return str(device_type).lower() + return "cuda" + + if torch_amp is not None and hasattr(torch_amp, "autocast") and hasattr( torch_amp, "GradScaler" ): + from torch.amp import GradScaler as _amp_grad_scaler from torch.amp import autocast as _amp_autocast - from torch.amp import GradScaler - def autocast(enabled=True, dtype=None, cache_enabled=True): + def autocast(enabled=True, dtype=None, cache_enabled=True, device_type=None): """torch.amp.autocast with the legacy CUDA autocast signature.""" + if device_type is None: + device_type = resolve_amp_device_type() return _amp_autocast( - "cuda", dtype=dtype, enabled=enabled, cache_enabled=cache_enabled + device_type, dtype=dtype, enabled=enabled, cache_enabled=cache_enabled ) + def _grad_scaler_class(device_type): + """GradScaler implementation to use for ``device_type``. + + Accelerators that ship their own AMP scaler (``torch.npu.amp.GradScaler`` + from torch_npu, ``torch.xpu.amp.GradScaler``) are preferred: the generic + ``torch.amp.GradScaler("npu")`` builds its scale on the device but then + fails in ``update()`` (``aclnnAmpUpdateScale`` error 561002 on CANN 9.1.0 + with torch_npu 2.15), while ``torch.npu.amp.GradScaler`` runs a full + fp16 step. CUDA and CPU stay on ``torch.amp.GradScaler``, i.e. the + non-deprecated API this shim exists for. + """ + if device_type != "cuda": + amp_module = getattr(getattr(torch, device_type, None), "amp", None) + grad_scaler = getattr(amp_module, "GradScaler", None) + if grad_scaler is not None: + return grad_scaler + return _amp_grad_scaler + + def GradScaler(*args, **kwargs): + """Return the AMP ``GradScaler`` of the active accelerator. + + ``torch.amp.GradScaler``-compatible, with the accelerator resolved the + same way as ``autocast()``. The device is only filled in for the generic + implementation; accelerator-specific scalers default to their own + device. + """ + device_type = resolve_amp_device_type() + scaler_class = _grad_scaler_class(device_type) + if scaler_class is _amp_grad_scaler and not args and "device" not in kwargs: + kwargs["device"] = device_type + return scaler_class(*args, **kwargs) + else: from torch.cuda.amp import autocast, GradScaler -__all__ = ["autocast", "GradScaler"] +__all__ = ["autocast", "GradScaler", "resolve_amp_device_type"] diff --git a/tests/test_amp_device_type.py b/tests/test_amp_device_type.py new file mode 100644 index 0000000000..13e4252c52 --- /dev/null +++ b/tests/test_amp_device_type.py @@ -0,0 +1,165 @@ +"""Regression tests for the device type used by the AMP shim. + +The shim used to pin ``device_type="cuda"``, which ``torch.amp`` turns into a +silent no-op on any other accelerator (it only warns "CUDA is not available ... +Disabling autocast"), so ``use_fp16``/``use_bf16`` training ran in float32 on +Ascend NPU / XPU / MPS builds, and ``GradScaler(enabled=True)`` was silently +disabled as well. +""" + +import importlib.util +import types +from pathlib import Path + +import pytest +import torch + +AMP_PATH = Path(__file__).resolve().parents[1] / "funasr" / "utils" / "amp.py" + + +def _load_amp(): + spec = importlib.util.spec_from_file_location("funasr_amp_device_type", AMP_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +class _Device: + """Stand-in for a ``torch.device`` of a backend that may be unavailable.""" + + def __init__(self, device_type): + self.type = device_type + + +class _Accelerator: + def __init__(self, device_type, error=None): + self._device_type = device_type + self._error = error + + def current_accelerator(self): + if self._error is not None: + raise self._error + if self._device_type is None: + return None + return _Device(self._device_type) + + +class _Recorder: + def __init__(self, result=None): + self.calls = [] + self._result = result + + def __call__(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return self._result + + +def _patch_accelerator(monkeypatch, device_type, error=None): + monkeypatch.setattr( + torch, "accelerator", _Accelerator(device_type, error=error), raising=False + ) + + +@pytest.mark.parametrize( + "device_type,expected", + [("cuda", "cuda"), ("npu", "npu"), ("xpu", "xpu"), ("mps", "mps")], +) +def test_resolve_amp_device_type_keeps_the_active_accelerator(monkeypatch, device_type, expected): + amp = _load_amp() + _patch_accelerator(monkeypatch, device_type) + + assert amp.resolve_amp_device_type() == expected + + +@pytest.mark.parametrize("device_type", [None, "mtia", "meta"]) +def test_resolve_amp_device_type_falls_back_to_cuda(monkeypatch, device_type): + """No accelerator (CPU-only builds) keeps the historical no-op default.""" + amp = _load_amp() + _patch_accelerator(monkeypatch, device_type) + + assert amp.resolve_amp_device_type() == "cuda" + + +def test_resolve_amp_device_type_survives_a_missing_or_failing_accelerator(monkeypatch): + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu", error=RuntimeError("no accelerator compiled")) + assert amp.resolve_amp_device_type() == "cuda" + + monkeypatch.delattr(torch, "accelerator", raising=False) + assert amp.resolve_amp_device_type() == "cuda" + + +def test_autocast_resolves_the_device_type_for_the_legacy_signature(monkeypatch): + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + recorder = _Recorder(result="ctx") + monkeypatch.setattr(amp, "_amp_autocast", recorder) + + assert amp.autocast(enabled=True, dtype=torch.float16) == "ctx" + assert recorder.calls == [ + (("npu",), {"dtype": torch.float16, "enabled": True, "cache_enabled": True}) + ] + + +def test_autocast_keeps_an_explicit_device_type(monkeypatch): + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + recorder = _Recorder(result="ctx") + monkeypatch.setattr(amp, "_amp_autocast", recorder) + + amp.autocast(False) + amp.autocast(True, torch.bfloat16, False, "cuda") + amp.autocast(enabled=False, device_type="cuda") + + assert [args[0] for args, _ in recorder.calls] == ["npu", "cuda", "cuda"] + + +def test_grad_scaler_prefers_the_accelerator_amp_module(monkeypatch): + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + accelerator_scaler = _Recorder() + monkeypatch.setattr( + torch, + "npu", + types.SimpleNamespace(amp=types.SimpleNamespace(GradScaler=accelerator_scaler)), + raising=False, + ) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + amp.GradScaler(enabled=True) + + assert generic_scaler.calls == [] + assert accelerator_scaler.calls == [((), {"enabled": True})] + + +def test_grad_scaler_passes_the_resolved_device_to_the_generic_implementation(monkeypatch): + amp = _load_amp() + _patch_accelerator(monkeypatch, "mps") + # a backend without its own amp module must not shadow the generic scaler + monkeypatch.setattr(torch, "mps", types.SimpleNamespace(), raising=False) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + amp.GradScaler(enabled=True) + amp.GradScaler("cuda", enabled=False) + + assert generic_scaler.calls == [ + ((), {"enabled": True, "device": "mps"}), + (("cuda",), {"enabled": False}), + ] + + +def test_grad_scaler_keeps_cuda_and_cpu_on_the_generic_implementation(monkeypatch): + amp = _load_amp() + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + _patch_accelerator(monkeypatch, "cuda") + amp.GradScaler(enabled=True) + assert generic_scaler.calls == [((), {"enabled": True, "device": "cuda"})] + + _patch_accelerator(monkeypatch, None) + amp.GradScaler(enabled=True) + assert generic_scaler.calls[-1] == ((), {"enabled": True, "device": "cuda"}) From dd8e0ff9a4a2c45b8520ed77caed6fe123ea1dd8 Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Tue, 29 Sep 2026 06:13:33 +0800 Subject: [PATCH 2/4] fix(amp): honour an explicit device before choosing the GradScaler implementation An explicit device (positional or keyword) now selects the scaler for that device instead of the active accelerator, so GradScaler(device="cpu") no longer reaches torch.npu.amp.GradScaler (its constructor has no device parameter and raised TypeError) and GradScaler("cpu") no longer consumed "cpu" as init_scale. The device argument is dropped before dispatching to an accelerator-specific scaler and forwarded unchanged to the generic one; torch.cpu.amp is no longer mistaken for a backend-specific scaler. Explicit keyword/positional/torch.device coverage added. Signed-off-by: li-lizhe <147392333@qq.com> --- funasr/utils/amp.py | 78 ++++++++++++++++++++---- tests/test_amp_device_type.py | 108 ++++++++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 11 deletions(-) diff --git a/funasr/utils/amp.py b/funasr/utils/amp.py index b28ff03e4d..f043ec3e5e 100644 --- a/funasr/utils/amp.py +++ b/funasr/utils/amp.py @@ -37,6 +37,14 @@ ``funasr/models/fun_asr_nano/device_utils.py`` accepts for the Nano inference path — and keep the historical ``"cuda"`` (a no-op on CPU-only builds) when no accelerator is available, so CPU behaviour is unchanged. + +An explicit device always wins over the detected accelerator: +``GradScaler(device="cpu", enabled=False)`` and ``GradScaler("cpu", ...)`` +keep whichever implementation understands that device, exactly as the +``torch.amp.GradScaler`` export they replace did. On an accelerator that ships +its own scaler the ``device`` argument is consumed by this shim (that +constructor takes no ``device``) while every other argument is forwarded +unchanged. """ import torch @@ -47,6 +55,11 @@ # keeps the historical "cuda" no-op instead of risking an unsupported context. _ACCELERATOR_DEVICE_TYPES = ("cuda", "npu", "xpu", "mps") +# Backends that ship their own AMP scaler. Everything else — including ``cpu``, +# whose ``torch.cpu.amp`` is only a deprecated re-export of the generic scaler — +# goes through ``torch.amp.GradScaler``. +_NATIVE_SCALER_DEVICE_TYPES = ("npu", "xpu") + def resolve_amp_device_type(): """Return the ``torch.amp`` device type of the active accelerator. @@ -67,6 +80,21 @@ def resolve_amp_device_type(): return "cuda" +def _resolve_device_type(device): + """Return the device type of an explicitly passed device, or ``None``. + + Accepts whatever ``torch.amp.GradScaler`` accepts (``str``, ``torch.device``), + so ``"npu"``, ``"cuda:0"`` and ``torch.device("cpu")`` all resolve. + """ + if device is None: + return None + device_type = getattr(device, "type", None) + if device_type is None: + device_type = str(device).split(":")[0] + device_type = str(device_type).strip().lower() + return device_type or None + + if torch_amp is not None and hasattr(torch_amp, "autocast") and hasattr( torch_amp, "GradScaler" ): @@ -90,27 +118,55 @@ def _grad_scaler_class(device_type): fails in ``update()`` (``aclnnAmpUpdateScale`` error 561002 on CANN 9.1.0 with torch_npu 2.15), while ``torch.npu.amp.GradScaler`` runs a full fp16 step. CUDA and CPU stay on ``torch.amp.GradScaler``, i.e. the - non-deprecated API this shim exists for. + non-deprecated API this shim exists for — note that ``torch.cpu.amp`` + does exist on recent torch but is a deprecated re-export of the generic + scaler with the device pinned, so it must not be picked up here. """ - if device_type != "cuda": + if device_type in _NATIVE_SCALER_DEVICE_TYPES: amp_module = getattr(getattr(torch, device_type, None), "amp", None) grad_scaler = getattr(amp_module, "GradScaler", None) if grad_scaler is not None: return grad_scaler return _amp_grad_scaler - def GradScaler(*args, **kwargs): - """Return the AMP ``GradScaler`` of the active accelerator. + def _caller_device(args, kwargs): + """Return ``(device, where)`` for the device the caller passed, if any. - ``torch.amp.GradScaler``-compatible, with the accelerator resolved the - same way as ``autocast()``. The device is only filled in for the generic - implementation; accelerator-specific scalers default to their own - device. + ``torch.amp.GradScaler``'s first positional parameter is ``device``, so + ``GradScaler("cpu", ...)`` and ``GradScaler(device="cpu", ...)`` name a + device while anything else leaves the scaler defaults in place. """ - device_type = resolve_amp_device_type() + if "device" in kwargs: + return kwargs["device"], "keyword" + if args: + return args[0], "positional" + return None, None + + def GradScaler(*args, **kwargs): + """Return the AMP ``GradScaler`` for the requested / active device. + + ``torch.amp.GradScaler``-compatible: an explicit ``device`` (positional + or keyword) selects the implementation for that device and is forwarded + untouched to the generic scaler. Without an explicit device the active + accelerator decides, and the device is filled in for the generic + implementation only — accelerator-specific scalers default to their own + device and their constructor takes no ``device`` argument. + """ + device, where = _caller_device(args, kwargs) + device_type = ( + resolve_amp_device_type() if device is None else _resolve_device_type(device) + ) scaler_class = _grad_scaler_class(device_type) - if scaler_class is _amp_grad_scaler and not args and "device" not in kwargs: - kwargs["device"] = device_type + if scaler_class is _amp_grad_scaler: + if device is None: + kwargs["device"] = device_type + return scaler_class(*args, **kwargs) + # The accelerator implementation owns its device and has no ``device`` + # parameter, so drop ours before forwarding the remaining arguments. + if where == "keyword": + kwargs.pop("device", None) + elif where == "positional": + args = args[1:] return scaler_class(*args, **kwargs) else: diff --git a/tests/test_amp_device_type.py b/tests/test_amp_device_type.py index 13e4252c52..188ee45fd9 100644 --- a/tests/test_amp_device_type.py +++ b/tests/test_amp_device_type.py @@ -163,3 +163,111 @@ def test_grad_scaler_keeps_cuda_and_cpu_on_the_generic_implementation(monkeypatc _patch_accelerator(monkeypatch, None) amp.GradScaler(enabled=True) assert generic_scaler.calls[-1] == ((), {"enabled": True, "device": "cuda"}) + + +def _patch_npu_scaler(monkeypatch): + """Register an NPU AMP scaler that records the arguments it is called with.""" + accelerator_scaler = _Recorder() + monkeypatch.setattr( + torch, + "npu", + types.SimpleNamespace(amp=types.SimpleNamespace(GradScaler=accelerator_scaler)), + raising=False, + ) + return accelerator_scaler + + +def test_grad_scaler_honours_an_explicit_keyword_device(monkeypatch): + """``GradScaler(device="cpu", ...)`` must not reach the NPU implementation. + + The accelerator scaler has no ``device`` parameter, so forwarding it would + raise ``TypeError``; the generic implementation must be selected instead. + """ + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + accelerator_scaler = _patch_npu_scaler(monkeypatch) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + amp.GradScaler(device="cpu", enabled=False) + + assert accelerator_scaler.calls == [] + assert generic_scaler.calls == [((), {"device": "cpu", "enabled": False})] + + +def test_grad_scaler_honours_an_explicit_positional_device(monkeypatch): + """``GradScaler("cpu", ...)`` keeps ``"cpu"`` as the device, not as ``init_scale``.""" + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + accelerator_scaler = _patch_npu_scaler(monkeypatch) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + amp.GradScaler("cpu", enabled=False) + + assert accelerator_scaler.calls == [] + assert generic_scaler.calls == [(("cpu",), {"enabled": False})] + + +def test_grad_scaler_honours_an_explicit_torch_device(monkeypatch): + """A ``torch.device`` is resolved by its type, like ``torch.amp.GradScaler``.""" + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + accelerator_scaler = _patch_npu_scaler(monkeypatch) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + device = torch.device("cpu") + amp.GradScaler(device, enabled=False) + + assert accelerator_scaler.calls == [] + assert generic_scaler.calls == [((device,), {"enabled": False})] + + +def test_grad_scaler_drops_the_device_for_the_accelerator_implementation(monkeypatch): + """An explicit accelerator device selects its scaler without forwarding ``device``.""" + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + accelerator_scaler = _patch_npu_scaler(monkeypatch) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + amp.GradScaler(device="npu", enabled=True) + amp.GradScaler("npu", 1024.0) + + assert generic_scaler.calls == [] + assert accelerator_scaler.calls == [ + ((), {"enabled": True}), + ((1024.0,), {}), + ] + + +def test_grad_scaler_keeps_an_explicit_cuda_device_on_the_generic_implementation(monkeypatch): + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + accelerator_scaler = _patch_npu_scaler(monkeypatch) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + amp.GradScaler(device="cuda") + amp.GradScaler("cuda:0", enabled=False) + + assert accelerator_scaler.calls == [] + assert generic_scaler.calls == [ + ((), {"device": "cuda"}), + (("cuda:0",), {"enabled": False}), + ] + + +def test_grad_scaler_without_a_device_still_prefers_the_accelerator(monkeypatch): + """No explicit device keeps the previous behaviour: the active accelerator wins.""" + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + accelerator_scaler = _patch_npu_scaler(monkeypatch) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + amp.GradScaler(enabled=True) + + assert generic_scaler.calls == [] + assert accelerator_scaler.calls == [((), {"enabled": True})] From d6b13af4f121d48184c45e9e3aec57c88fc8a905 Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Tue, 29 Sep 2026 22:11:08 +0800 Subject: [PATCH 3/4] fix(amp): bind the generic GradScaler signature before native dispatch The accelerator scalers do not share the generic torch.amp.GradScaler positional order: torch_npu inserts `dynamic` before `enabled`, so forwarding the caller's positional tail unchanged turned GradScaler("npu", 1024.0, 2.0, 0.5, 2000, False) into an enabled static scaler while the keyword form kept working. Bind a call written against the generic signature to its parameter names first, drop the `device` the accelerator constructor does not accept, and forward the rest by name. Calls that do not match the generic signature keep the previous plain re-export behaviour. Tests: cover the positional enabled=False/True tail against a strict torch_npu-signature stand-in (a *args/**kwargs recorder cannot detect a shifted argument slot). Signed-off-by: li-lizhe <147392333@qq.com> --- funasr/utils/amp.py | 70 ++++++++++++++++------ tests/test_amp_device_type.py | 108 ++++++++++++++++++++++++++++++++-- 2 files changed, 153 insertions(+), 25 deletions(-) diff --git a/funasr/utils/amp.py b/funasr/utils/amp.py index f043ec3e5e..34123fb7f5 100644 --- a/funasr/utils/amp.py +++ b/funasr/utils/amp.py @@ -43,8 +43,8 @@ keep whichever implementation understands that device, exactly as the ``torch.amp.GradScaler`` export they replace did. On an accelerator that ships its own scaler the ``device`` argument is consumed by this shim (that -constructor takes no ``device``) while every other argument is forwarded -unchanged. +constructor takes no ``device``) while every other argument is translated by +name, because those constructors do not share the generic positional order. """ import torch @@ -60,6 +60,19 @@ # goes through ``torch.amp.GradScaler``. _NATIVE_SCALER_DEVICE_TYPES = ("npu", "xpu") +# Parameter names of ``torch.amp.GradScaler``: ``device`` first, ``enabled`` +# last. A call written against the generic signature is bound by name before it +# is dispatched to an accelerator-specific scaler, whose parameter order is not +# the same (``torch_npu`` inserts ``dynamic`` before ``enabled``). +_GENERIC_SCALER_PARAMETERS = ( + "device", + "init_scale", + "growth_factor", + "backoff_factor", + "growth_interval", + "enabled", +) + def resolve_amp_device_type(): """Return the ``torch.amp`` device type of the active accelerator. @@ -129,18 +142,24 @@ def _grad_scaler_class(device_type): return grad_scaler return _amp_grad_scaler - def _caller_device(args, kwargs): - """Return ``(device, where)`` for the device the caller passed, if any. + def _generic_arguments(args, kwargs): + """Bind a ``torch.amp.GradScaler`` call to its parameter names. - ``torch.amp.GradScaler``'s first positional parameter is ``device``, so - ``GradScaler("cpu", ...)`` and ``GradScaler(device="cpu", ...)`` name a - device while anything else leaves the scaler defaults in place. + Returns the arguments keyed by generic parameter name, or ``None`` when + the call does not match the generic signature (unknown or repeated + parameters) — such calls are forwarded unchanged so that the selected + scaler raises the error ``torch.amp`` would have raised. """ - if "device" in kwargs: - return kwargs["device"], "keyword" - if args: - return args[0], "positional" - return None, None + if len(args) > len(_GENERIC_SCALER_PARAMETERS): + return None + arguments = dict(kwargs) + for name, value in zip(_GENERIC_SCALER_PARAMETERS, args): + if name in arguments: + return None + arguments[name] = value + if any(name not in _GENERIC_SCALER_PARAMETERS for name in arguments): + return None + return arguments def GradScaler(*args, **kwargs): """Return the AMP ``GradScaler`` for the requested / active device. @@ -151,8 +170,16 @@ def GradScaler(*args, **kwargs): accelerator decides, and the device is filled in for the generic implementation only — accelerator-specific scalers default to their own device and their constructor takes no ``device`` argument. + + The accelerator implementations do not share the generic positional + order (``torch_npu`` inserts ``dynamic`` before ``enabled``), so their + arguments are translated by name: ``GradScaler("npu", 1024.0, 2.0, 0.5, + 2000, False)`` reaches ``torch.npu.amp.GradScaler`` with + ``enabled=False`` instead of filling the trailing positionals into the + native parameters one slot early and enabling a static scaler. """ - device, where = _caller_device(args, kwargs) + arguments = _generic_arguments(args, kwargs) + device = None if arguments is None else arguments.get("device") device_type = ( resolve_amp_device_type() if device is None else _resolve_device_type(device) ) @@ -161,13 +188,18 @@ def GradScaler(*args, **kwargs): if device is None: kwargs["device"] = device_type return scaler_class(*args, **kwargs) + if arguments is None: + # Not the generic signature: keep the plain export behaviour and + # only drop the device the caller spelled out, if any. + if "device" in kwargs: + kwargs.pop("device", None) + elif args: + args = args[1:] + return scaler_class(*args, **kwargs) # The accelerator implementation owns its device and has no ``device`` - # parameter, so drop ours before forwarding the remaining arguments. - if where == "keyword": - kwargs.pop("device", None) - elif where == "positional": - args = args[1:] - return scaler_class(*args, **kwargs) + # parameter, so drop ours and forward every remaining argument by name. + arguments.pop("device", None) + return scaler_class(**arguments) else: from torch.cuda.amp import autocast, GradScaler diff --git a/tests/test_amp_device_type.py b/tests/test_amp_device_type.py index 188ee45fd9..3ee6debde5 100644 --- a/tests/test_amp_device_type.py +++ b/tests/test_amp_device_type.py @@ -165,16 +165,47 @@ def test_grad_scaler_keeps_cuda_and_cpu_on_the_generic_implementation(monkeypatc assert generic_scaler.calls[-1] == ((), {"enabled": True, "device": "cuda"}) -def _patch_npu_scaler(monkeypatch): - """Register an NPU AMP scaler that records the arguments it is called with.""" - accelerator_scaler = _Recorder() +class _NpuGradScalerStub: + """``torch.npu.amp.GradScaler`` stand-in with torch_npu's real signature. + + The parameter order matters: torch_npu inserts ``dynamic`` before + ``enabled`` + (https://github.com/Ascend/pytorch/blob/6d59edda7452b26267eecbb6c328d876a6521ada/torch_npu/npu/amp/grad_scaler.py#L93-L99), + so a ``*args, **kwargs`` recorder cannot tell a correctly translated call + from one whose positional tail landed one parameter too early. + """ + + def __init__( + self, + init_scale=2.0 ** 16, + growth_factor=2.0, + backoff_factor=0.5, + growth_interval=2000, + dynamic=True, + enabled=True, + ): + self.init_scale = init_scale + self.growth_factor = growth_factor + self.backoff_factor = backoff_factor + self.growth_interval = growth_interval + self.dynamic = dynamic + self._enabled = enabled + + def is_enabled(self): + return self._enabled + + +def _patch_npu_scaler(monkeypatch, scaler=None): + """Register an NPU AMP scaler (an argument recorder by default).""" + if scaler is None: + scaler = _Recorder() monkeypatch.setattr( torch, "npu", - types.SimpleNamespace(amp=types.SimpleNamespace(GradScaler=accelerator_scaler)), + types.SimpleNamespace(amp=types.SimpleNamespace(GradScaler=scaler)), raising=False, ) - return accelerator_scaler + return scaler def test_grad_scaler_honours_an_explicit_keyword_device(monkeypatch): @@ -238,10 +269,75 @@ def test_grad_scaler_drops_the_device_for_the_accelerator_implementation(monkeyp assert generic_scaler.calls == [] assert accelerator_scaler.calls == [ ((), {"enabled": True}), - ((1024.0,), {}), + ((), {"init_scale": 1024.0}), + ] + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_grad_scaler_translates_the_generic_positional_tail_for_the_npu_scaler( + monkeypatch, enabled +): + """``GradScaler("npu", 1024.0, 2.0, 0.5, 2000, enabled)`` keeps ``enabled``. + + The last parameter of the generic signature is ``enabled``, but torch_npu + inserts ``dynamic`` before it, so forwarding the positional tail unchanged + hands ``False`` to ``dynamic`` and leaves the scaler enabled — a caller + asking to disable scaling gets an enabled static scaler instead. + """ + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + _patch_npu_scaler(monkeypatch, scaler=_NpuGradScalerStub) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + scaler = amp.GradScaler("npu", 1024.0, 2.0, 0.5, 2000, enabled) + + assert generic_scaler.calls == [] + assert isinstance(scaler, _NpuGradScalerStub) + assert scaler.is_enabled() is enabled + assert scaler.init_scale == 1024.0 + assert scaler.growth_factor == 2.0 + assert scaler.backoff_factor == 0.5 + assert scaler.growth_interval == 2000 + assert scaler.dynamic is True + + +def test_grad_scaler_translates_arguments_by_name_before_dispatch(monkeypatch): + """The accelerator scaler is called with parameter names, not a shifted tail.""" + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + accelerator_scaler = _patch_npu_scaler(monkeypatch) + generic_scaler = _Recorder() + monkeypatch.setattr(amp, "_amp_grad_scaler", generic_scaler) + + amp.GradScaler("npu", 1024.0, 2.0, 0.5, 2000, False) + + assert generic_scaler.calls == [] + assert accelerator_scaler.calls == [ + ( + (), + { + "init_scale": 1024.0, + "growth_factor": 2.0, + "backoff_factor": 0.5, + "growth_interval": 2000, + "enabled": False, + }, + ) ] +def test_grad_scaler_keeps_a_keyword_enabled_flag_for_the_npu_scaler(monkeypatch): + """The keyword form keeps meaning the same thing on every scaler.""" + amp = _load_amp() + _patch_accelerator(monkeypatch, "npu") + _patch_npu_scaler(monkeypatch, scaler=_NpuGradScalerStub) + + assert amp.GradScaler(enabled=False).is_enabled() is False + assert amp.GradScaler(device="npu", enabled=False).is_enabled() is False + assert amp.GradScaler("npu", enabled=True).is_enabled() is True + + def test_grad_scaler_keeps_an_explicit_cuda_device_on_the_generic_implementation(monkeypatch): amp = _load_amp() _patch_accelerator(monkeypatch, "npu") From d60f1658641d5219386f7ed9c19d22f60cbbde81 Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Wed, 30 Sep 2026 06:07:40 +0800 Subject: [PATCH 4/4] docs(amp): narrow the _generic_arguments docstring claim Signed-off-by: li-lizhe <147392333@qq.com> --- funasr/utils/amp.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/funasr/utils/amp.py b/funasr/utils/amp.py index 34123fb7f5..e2ffe450b3 100644 --- a/funasr/utils/amp.py +++ b/funasr/utils/amp.py @@ -147,8 +147,12 @@ def _generic_arguments(args, kwargs): Returns the arguments keyed by generic parameter name, or ``None`` when the call does not match the generic signature (unknown or repeated - parameters) — such calls are forwarded unchanged so that the selected - scaler raises the error ``torch.amp`` would have raised. + parameters). Those calls keep the plain export behaviour: an explicit + ``device`` is dropped and the rest is forwarded untouched, so the + *generic* scaler rejects them as ``torch.amp`` would, while an + accelerator-specific scaler validates the remaining arguments against + its own signature (``torch_npu`` takes one parameter more, so a call + the generic scaler would reject can still bind natively). """ if len(args) > len(_GENERIC_SCALER_PARAMETERS): return None