diff --git a/examples/industrial_data_pretraining/fun_asr_nano/model.py b/examples/industrial_data_pretraining/fun_asr_nano/model.py index 491fd0ecb..88c52e9ce 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/model.py +++ b/examples/industrial_data_pretraining/fun_asr_nano/model.py @@ -15,6 +15,7 @@ from funasr.register import tables from funasr.train_utils.device_funcs import force_gatherable, to_device from funasr.utils.datadir_writer import DatadirWriter from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video +from funasr.models.fun_asr_nano.device_utils import resolve_autocast_device_type from ctc import CTC @@ -231,9 +232,9 @@ class FunASRNano(nn.Module): stats["batch_size_real_frames"] = speech_lengths.sum().item() stats["padding_frames"] = stats["batch_size_x_frames"] - stats["batch_size_real_frames"] - device_type = next(self.parameters()).device.type + autocast_device_type = resolve_autocast_device_type(next(self.parameters()).device) with torch.autocast( - device_type=device_type if device_type in ["cuda", "xpu", "mps"] else "cpu", + device_type=autocast_device_type, enabled=True if self.llm_dtype != "fp32" else False, dtype=dtype_map[self.llm_dtype], ): @@ -650,9 +651,9 @@ class FunASRNano(nn.Module): llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype - device_type = torch.device(kwargs.get("device", "cuda")).type + autocast_device_type = resolve_autocast_device_type(kwargs.get("device", "cuda")) with torch.autocast( - device_type=device_type if device_type in ["cuda", "xpu", "mps"] else "cpu", + device_type=autocast_device_type, enabled=True if llm_dtype != "fp32" else False, dtype=dtype_map[llm_dtype], ): diff --git a/funasr/models/fun_asr_nano/device_utils.py b/funasr/models/fun_asr_nano/device_utils.py new file mode 100644 index 000000000..2e91f5cbe --- /dev/null +++ b/funasr/models/fun_asr_nano/device_utils.py @@ -0,0 +1,34 @@ +# Copyright FunASR (https://github.com/modelscope/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +"""Device helpers for Fun-ASR-Nano runtime paths.""" + +_SUPPORTED_AUTOCAST_DEVICE_TYPES = {"cuda", "xpu", "mps", "npu"} + + +def _device_type_from_value(device): + """Resolve a device type without requiring optional backend registration.""" + if device is None: + return "cpu" + + device_type = getattr(device, "type", None) + if device_type: + return str(device_type).lower() + + if isinstance(device, str): + return device.split(":", 1)[0].lower() + + return str(device).split(":", 1)[0].lower() + + +def resolve_autocast_device_type(device): + """Return the torch.autocast device_type for a Fun-ASR-Nano device. + + PyTorch builds without torch_npu may reject ``torch.device("npu:0")`` before + torch_npu registers the backend. Parse strings directly so NPU requests do + not fall back to CPU autocast, which only supports bf16 and caused #3034. + """ + device_type = _device_type_from_value(device) + if device_type in _SUPPORTED_AUTOCAST_DEVICE_TYPES: + return device_type + return "cpu" diff --git a/funasr/models/fun_asr_nano/model.py b/funasr/models/fun_asr_nano/model.py index a8f95f783..d9dd2a7cb 100644 --- a/funasr/models/fun_asr_nano/model.py +++ b/funasr/models/fun_asr_nano/model.py @@ -22,6 +22,7 @@ except ImportError: AutoModelForCausalLM = None from .ctc import CTC +from .device_utils import resolve_autocast_device_type from .tools.utils import forced_align dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} @@ -279,9 +280,9 @@ class FunASRNano(nn.Module): stats["batch_size_real_frames"] = speech_lengths.sum().item() stats["padding_frames"] = stats["batch_size_x_frames"] - stats["batch_size_real_frames"] - device_type = next(self.parameters()).device.type + autocast_device_type = resolve_autocast_device_type(next(self.parameters()).device) with torch.autocast( - device_type=device_type if device_type in ["cuda", "xpu", "mps"] else "cpu", + device_type=autocast_device_type, enabled=True if self.llm_dtype != "fp32" else False, dtype=dtype_map[self.llm_dtype], ): @@ -758,9 +759,9 @@ class FunASRNano(nn.Module): padded[i, Tmax - Ti :, :] = e[0].to(dt) # left padding attn[i, Tmax - Ti :] = 1 - device_type = torch.device(kwargs.get("device", "cuda")).type + autocast_device_type = resolve_autocast_device_type(kwargs.get("device", "cuda")) with torch.autocast( - device_type=device_type if device_type in ["cuda", "xpu", "mps"] else "cpu", + device_type=autocast_device_type, enabled=True if llm_dtype != "fp32" else False, dtype=dt, ): @@ -852,9 +853,9 @@ class FunASRNano(nn.Module): llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype - device_type = torch.device(kwargs.get("device", "cuda")).type + autocast_device_type = resolve_autocast_device_type(kwargs.get("device", "cuda")) with torch.autocast( - device_type=device_type if device_type in ["cuda", "xpu", "mps"] else "cpu", + device_type=autocast_device_type, enabled=True if llm_dtype != "fp32" else False, dtype=dtype_map[llm_dtype], ): diff --git a/tests/test_fun_asr_nano_autocast_device.py b/tests/test_fun_asr_nano_autocast_device.py new file mode 100644 index 000000000..cf68f97e0 --- /dev/null +++ b/tests/test_fun_asr_nano_autocast_device.py @@ -0,0 +1,54 @@ +import importlib.util +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +DEVICE_UTILS_PATH = REPO_ROOT / "funasr" / "models" / "fun_asr_nano" / "device_utils.py" +PACKAGE_MODEL_PATH = REPO_ROOT / "funasr" / "models" / "fun_asr_nano" / "model.py" +EXAMPLE_MODEL_PATH = ( + REPO_ROOT / "examples" / "industrial_data_pretraining" / "fun_asr_nano" / "model.py" +) + + +def _load_device_utils(): + spec = importlib.util.spec_from_file_location("fun_asr_nano_device_utils", DEVICE_UTILS_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_resolve_autocast_device_type_keeps_npu_backend(): + device_utils = _load_device_utils() + + assert device_utils.resolve_autocast_device_type("npu:1") == "npu" + assert device_utils.resolve_autocast_device_type("npu") == "npu" + assert device_utils.resolve_autocast_device_type("cuda:0") == "cuda" + assert device_utils.resolve_autocast_device_type("xpu:0") == "xpu" + assert device_utils.resolve_autocast_device_type("mps") == "mps" + assert device_utils.resolve_autocast_device_type("cpu") == "cpu" + assert device_utils.resolve_autocast_device_type("unknown:0") == "cpu" + + +def test_device_type_from_value_treats_none_as_cpu(): + device_utils = _load_device_utils() + + assert device_utils._device_type_from_value(None) == "cpu" + + +def test_resolve_autocast_device_type_accepts_device_like_objects(): + device_utils = _load_device_utils() + + class DeviceLike: + type = "npu" + + assert device_utils.resolve_autocast_device_type(DeviceLike()) == "npu" + + +def test_fun_asr_nano_autocast_calls_use_shared_resolver(): + old_inline_fallback = 'device_type if device_type in ["cuda", "xpu", "mps"] else "cpu"' + + for path in (PACKAGE_MODEL_PATH, EXAMPLE_MODEL_PATH): + source = path.read_text(encoding="utf-8") + assert old_inline_fallback not in source + assert source.count("resolve_autocast_device_type(") >= 2