mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
Fix Fun-ASR-Nano NPU autocast device selection (#3113)
Preserve supported accelerator device types, including npu, when choosing Fun-ASR-Nano autocast device type. This prevents npu devices from falling back to CPU autocast before real torch_npu execution can be validated.
This commit is contained in:
@@ -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],
|
||||
):
|
||||
|
||||
@@ -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"
|
||||
@@ -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],
|
||||
):
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user