Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 18 additions & 10 deletions funasr/models/fun_asr_nano/inference_vllm.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,19 @@
dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}


def _resolve_vllm_dtype(dtype: str) -> str:
"""Use a numerically stable dtype for the Qwen3 language model."""
if dtype == "fp16":
logger.warning(
"Fun-ASR-Nano's Qwen3 language model is numerically unstable in "
"float16; vLLM will use bfloat16 while audio components remain in "
"float16. This keeps the same memory footprint and avoids degraded "
"or repetitive transcription."
)
return "bfloat16"
return {"bf16": "bfloat16", "fp32": "auto"}.get(dtype, dtype)


def prepare_vllm_model_dir(model_dir: str, output_dir: str = None) -> str:
"""Extract LLM weights from Fun-ASR-Nano model.pt and save in HuggingFace format.

Expand Down Expand Up @@ -141,7 +154,8 @@ class FunASRNanoVLLM:
Args:
model_dir: Path to the Fun-ASR-Nano model directory.
device: Device for audio encoder/adaptor (e.g. "cuda:0").
dtype: Dtype for audio processing ("bf16", "fp16", "fp32").
dtype: Dtype for audio processing ("bf16", "fp16", "fp32"). The Qwen3
language model uses bf16 when audio processing uses fp16.
tensor_parallel_size: Number of GPUs for vLLM tensor parallelism.
gpu_memory_utilization: Fraction of GPU memory for vLLM KV cache.
max_model_len: Maximum sequence length for vLLM.
Expand Down Expand Up @@ -177,13 +191,6 @@ def __init__(
self.device = device
self.dtype = dtype
self.torch_dtype = dtype_map.get(dtype, torch.bfloat16)
if self.torch_dtype == torch.float16:
logger.warning(
"dtype='fp16' can produce degraded or garbage transcription for "
"Fun-ASR-Nano (numerical overflow in the audio embedding path). "
"Use dtype='bf16' (recommended) or dtype='fp32'. On GPUs without "
"bfloat16 support (e.g. NVIDIA V100), use 'fp32'."
)
self.model_dir = model_dir

# Step 1: Prepare LLM weights for vLLM (extract from model.pt if needed)
Expand All @@ -205,7 +212,7 @@ def __init__(
gpu_memory_utilization=gpu_memory_utilization,
max_model_len=max_model_len,
enforce_eager=enforce_eager,
dtype={"bf16": "bfloat16", "fp16": "float16", "fp32": "auto"}.get(dtype, dtype),
dtype=_resolve_vllm_dtype(dtype),
trust_remote_code=True,
**vllm_kwargs,
)
Expand Down Expand Up @@ -712,7 +719,8 @@ def from_pretrained(
model: Model name or local directory path.
hub: "ms" (ModelScope) or "hf" (HuggingFace).
device: Device for audio encoder/adaptor.
dtype: Compute dtype ("bf16", "fp16", "fp32").
dtype: Audio compute dtype ("bf16", "fp16", "fp32"). The Qwen3
language model uses bf16 when audio compute uses fp16.
tensor_parallel_size: GPUs for vLLM tensor parallel.
gpu_memory_utilization: GPU memory fraction for vLLM.
max_model_len: Maximum sequence length.
Expand Down
10 changes: 7 additions & 3 deletions funasr/models/fun_asr_nano/inference_vllm_streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,8 @@ class FunASRNanoStreamingVLLM:
Args:
model_dir: Path to Fun-ASR-Nano model directory.
device: Device for audio encoder/adaptor.
dtype: Compute dtype ("bf16", "fp16", "fp32").
dtype: Audio compute dtype ("bf16", "fp16", "fp32"). The Qwen3
language model uses bf16 when audio compute uses fp16.
tensor_parallel_size: GPUs for vLLM tensor parallelism.
gpu_memory_utilization: GPU memory fraction for KV cache.
max_model_len: Maximum sequence length.
Expand All @@ -69,7 +70,10 @@ def __init__(self, model_dir, device="cuda:0", dtype="bf16",
max_model_len=2048, enforce_eager=False,
chunk_ms=720, rollback_chars=8, **kwargs):
from vllm import LLM
from funasr.models.fun_asr_nano.inference_vllm import prepare_vllm_model_dir
from funasr.models.fun_asr_nano.inference_vllm import (
_resolve_vllm_dtype,
prepare_vllm_model_dir,
)

self.device = device
self.dtype = dtype
Expand All @@ -87,7 +91,7 @@ def __init__(self, model_dir, device="cuda:0", dtype="bf16",
tensor_parallel_size=tensor_parallel_size,
gpu_memory_utilization=gpu_memory_utilization,
max_model_len=max_model_len, enforce_eager=enforce_eager,
dtype={"bf16": "bfloat16", "fp16": "float16", "fp32": "auto"}.get(dtype, dtype),
dtype=_resolve_vllm_dtype(dtype),
trust_remote_code=True, **vllm_kwargs,
)
self.tokenizer = self.vllm_engine.get_tokenizer()
Expand Down
78 changes: 78 additions & 0 deletions tests/test_fun_asr_nano_vllm_dtype.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
import logging
import sys
import types

import torch

from funasr.models.fun_asr_nano import inference_vllm
from funasr.models.fun_asr_nano import inference_vllm_streaming


class _FakeLLM:
calls = []

def __init__(self, **kwargs):
self.calls.append(kwargs)

def get_tokenizer(self):
return object()


def _install_fake_vllm(monkeypatch):
vllm = types.ModuleType("vllm")
vllm.LLM = _FakeLLM
vllm.SamplingParams = object
inputs = types.ModuleType("vllm.inputs")
inputs.EmbedsPrompt = object
monkeypatch.setitem(sys.modules, "vllm", vllm)
monkeypatch.setitem(sys.modules, "vllm.inputs", inputs)


def test_fp16_keeps_audio_compute_but_promotes_vllm_to_bf16(monkeypatch, caplog):
_FakeLLM.calls.clear()
_install_fake_vllm(monkeypatch)
monkeypatch.setattr(inference_vllm, "prepare_vllm_model_dir", lambda path: path)
monkeypatch.setattr(
inference_vllm.FunASRNanoVLLM,
"_load_audio_components",
lambda self, model_dir, **kwargs: None,
)
monkeypatch.setattr(
inference_vllm.FunASRNanoVLLM,
"_load_embedding_layer",
lambda self, model_dir: None,
)

def load_streaming_audio(self, model_dir):
self.frontend = types.SimpleNamespace(fs=16000)

monkeypatch.setattr(
inference_vllm_streaming.FunASRNanoStreamingVLLM,
"_load_audio_components",
load_streaming_audio,
)
monkeypatch.setattr(
inference_vllm_streaming.FunASRNanoStreamingVLLM,
"_load_embedding_layer",
lambda self, model_dir: None,
)

with caplog.at_level(logging.WARNING):
offline = inference_vllm.FunASRNanoVLLM("/tmp/model", dtype="fp16")
streaming = inference_vllm_streaming.FunASRNanoStreamingVLLM(
"/tmp/model", dtype="fp16"
)

assert offline.torch_dtype is torch.float16
assert streaming.torch_dtype is torch.float16
assert [call["dtype"] for call in _FakeLLM.calls] == ["bfloat16", "bfloat16"]
assert "audio components remain in float16" in caplog.text


def test_vllm_dtype_mapping_preserves_supported_values(caplog):
with caplog.at_level(logging.WARNING):
assert inference_vllm._resolve_vllm_dtype("bf16") == "bfloat16"
assert inference_vllm._resolve_vllm_dtype("fp32") == "auto"
assert inference_vllm._resolve_vllm_dtype("custom") == "custom"

assert not caplog.text