From aa353d6d223ec2f258a91cb623c08b1dee7745f2 Mon Sep 17 00:00:00 2001 From: damnwenxi Date: Mon, 14 Sep 2026 09:20:40 +0800 Subject: [PATCH 1/2] fix: honor submodel device placement across inference calls --- docs/python_api.md | 26 ++++- docs/python_api_zh.md | 25 +++- funasr/auto/auto_model.py | 40 +++++-- tests/test_punc_model_none.py | 1 + tests/test_python_api_docs_contract.py | 2 +- tests/test_submodel_device.py | 154 +++++++++++++++++++++++++ 6 files changed, 234 insertions(+), 14 deletions(-) create mode 100644 tests/test_submodel_device.py diff --git a/docs/python_api.md b/docs/python_api.md index 26fd06e81..c4b3c72c0 100644 --- a/docs/python_api.md +++ b/docs/python_api.md @@ -23,7 +23,7 @@ The wrapper accepts `**kwargs`, not a universally validated parameter schema. A | `ngpu` | `1` | Zero selects CPU. This is not a multi-GPU serving or sharding configuration. | | `ncpu` | `4` | Positive CPU thread count; invalid values use the fallback and values below 1 clamp to 1. Sets process-wide PyTorch threads. | | `vad_model`, `punc_model`, `spk_model` | `None` | Optional components, built once. Use `vad_model`, not the rejected typo `vda_model`. | -| `vad_kwargs`, `punc_kwargs`, `spk_kwargs` | `{}` | Component configuration dictionaries. Device is inherited; hub and CPU threads are inherited unless provided in the component dictionary. | +| `vad_kwargs`, `punc_kwargs`, `spk_kwargs` | `{}` | Component configuration dictionaries. Device, hub and CPU threads are inherited unless explicitly provided in the component dictionary. Device inherits the resolved ASR device (after fallback). Caller dictionaries are copied before defaults are applied. | | `vad_model_revision`, `punc_model_revision`, `spk_model_revision` | `"master"` each | Separate component revisions, not inherited from the ASR revision. These top-level values replace `model_revision` in the component dictionaries. | | `spk_mode` | `"punc_segment"` | Use `"punc_segment"` or `"vad_segment"`. The validation also accepts legacy `"default"`, but the sentence-building branches do not implement it; do not select it. | | `disable_update` | `False` | Disables the FunASR package version check only, not model downloading. | @@ -51,6 +51,30 @@ The wrapper accepts `**kwargs`, not a universally validated parameter schema. A `generate()` restores saved construction configuration before merging call options. Repeat request-specific options on each call; do not rely on the previous call's language, batch settings, or hotwords. This configuration reset does not establish thread safety or reset every model attribute: for example, a speaker-mode fallback mutates `self.spk_mode`. Serialize access to a shared instance unless your own concurrency tests establish otherwise. +### Place submodels on separate devices + +For example, keep FSMN VAD and punctuation on CPU while ASR uses Apple Silicon MPS: + +```python +model = AutoModel( + model="paraformer-zh", + device="mps", + vad_model="fsmn-vad", + vad_kwargs={"device": "cpu"}, + punc_model="ct-punc", + punc_kwargs={"device": "cpu"}, +) +results = model.generate(input="audio.wav") +``` + +`spk_kwargs={"device": "cpu"}` works the same way when a speaker model is configured. +Each model resolves unavailable-device fallback independently at construction. +The selected devices remain fixed across calls, including when shared runtime +options are merged. A `device` passed to `generate()` is ignored for placement; +it does not move loaded weights. Construct another `AutoModel` to change devices. +Other runtime options, such as `hotword` and `batch_size_s`, continue to be merged. +No internal model replacement or `ComputeScores` patch is needed. + ## Local Files and Batching Install the SDK using the [installation guide](installation/installation.md). The three examples below are standalone scripts with command-line arguments and require prepared local model snapshots and real audio. Their syntax and wrapper contracts are checked without downloading weights; they are not reported as full inference tests. diff --git a/docs/python_api_zh.md b/docs/python_api_zh.md index 400510d17..5128c16a4 100644 --- a/docs/python_api_zh.md +++ b/docs/python_api_zh.md @@ -23,7 +23,7 @@ | `ngpu` | `1` | 设为零选择 CPU。它不是多 GPU 服务或模型分片配置。 | | `ncpu` | `4` | 正整数 CPU 线程数；非法值使用回退值，小于 1 的值限制为 1。修改进程级 PyTorch 线程数。 | | `vad_model`, `punc_model`, `spk_model` | `None` | 可选组件，在构建时加载。正确名称是 `vad_model`，拼错为 `vda_model` 会被拒绝。 | -| `vad_kwargs`, `punc_kwargs`, `spk_kwargs` | `{}` | 子组件配置字典。设备继承主模型；平台和 CPU 线程数在子字典未指定时继承。 | +| `vad_kwargs`, `punc_kwargs`, `spk_kwargs` | `{}` | 子组件配置字典。设备、平台和 CPU 线程数在子字典未指定时继承；设备继承主模型回退后实际选定的设备。填充默认值前会复制调用方字典。 | | `vad_model_revision`, `punc_model_revision`, `spk_model_revision` | 各自为 `"master"` | 分别指定子组件版本，不继承 ASR 版本。这些顶层参数会覆盖子字典中的 `model_revision`。 | | `spk_mode` | `"punc_segment"` | 使用 `"punc_segment"` 或 `"vad_segment"`。校验还接受历史值 `"default"`，但后续句子构造分支未实现它，请勿选择。 | | `disable_update` | `False` | 仅禁用 FunASR 包版本检查，不会禁用模型下载。 | @@ -51,6 +51,29 @@ `generate()` 在合并本次选项前恢复构建时保存的配置。每次调用都应明确传入本次需要的语言、批处理、热词等选项，不要依赖上一次调用遗留的值。配置重置不代表线程安全，也不重置每个模型属性，例如说话人模式回退会修改 `self.spk_mode`。除非自行验证过并发行为，否则应串行访问共享实例。 +### 为子模型分别指定设备 + +例如，在 Apple Silicon 上让 ASR 使用 MPS，而 FSMN VAD 和标点使用 CPU： + +```python +model = AutoModel( + model="paraformer-zh", + device="mps", + vad_model="fsmn-vad", + vad_kwargs={"device": "cpu"}, + punc_model="ct-punc", + punc_kwargs={"device": "cpu"}, +) +results = model.generate(input="audio.wav") +``` + +配置说话人模型时，`spk_kwargs={"device": "cpu"}` 遵循相同规则。 +每个模型在构建时独立处理不可用设备的回退。选定的设备在多次调用和共享 +运行参数合并过程中保持不变；`generate()` 中传入的 `device` 不会改变设备配置， +也不会搬移已加载的权重。需要更换设备时，请构建另一个 `AutoModel`。 +其他运行参数（如 `hotword`、`batch_size_s`）仍会正常合并。 +无需替换内部模型或修改 `ComputeScores`。 + ## 本地文件与批处理 先按[安装指南](installation/installation_zh.md) 安装 SDK。下面三个示例都是带命令行参数的独立脚本，需要提前准备完整本地模型快照和真实音频。这里只在不下载权重的条件下检查语法与包装层契约，不宣称完成了真实模型推理测试。 diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index 16eacfbe4..61f7ce8b9 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -426,9 +426,12 @@ def __init__(self, **kwargs): Falls back to CPU if specified device is unavailable. vad_model (str, optional): VAD model for long audio segmentation. Enables processing of any-length audio. - vad_kwargs (dict, optional): VAD config, e.g. {"max_single_segment_time": 60000}. + vad_kwargs (dict, optional): VAD config, e.g. {"device": "cpu"}. + Explicit device overrides the ASR device; otherwise inherits its resolved device. punc_model (str, optional): Punctuation restoration model. Not needed for Fun-ASR-Nano/SenseVoice/Qwen3-ASR (they output punctuation natively). + punc_kwargs (dict, optional): Punctuation config; device follows the same rules as vad_kwargs. + spk_kwargs (dict, optional): Speaker config; device follows the same rules as vad_kwargs. spk_model (str, optional): Speaker model for diarization ("cam++" or full model ID). Requires vad_model. For Qwen3-ASR, also requires forced_aligner. spk_mode (str, optional): Speaker diarization mode. "punc_segment" (default) or "vad_segment". @@ -458,6 +461,10 @@ def __init__(self, **kwargs): log_level = getattr(logging, kwargs.get("log_level", "INFO").upper()) logging.basicConfig(level=log_level) + # Defaults and model construction must not mutate caller-owned submodel configs. + for name in ("vad_kwargs", "punc_kwargs", "spk_kwargs"): + kwargs[name] = copy.deepcopy(kwargs.get(name) or {}) + model, kwargs = self.build_model(**kwargs) # if vad_model is not None, build vad model else None @@ -467,7 +474,7 @@ def __init__(self, **kwargs): logging.info("Building VAD model.") vad_kwargs["model"] = vad_model vad_kwargs["model_revision"] = kwargs.get("vad_model_revision", "master") - vad_kwargs["device"] = kwargs["device"] + vad_kwargs.setdefault("device", kwargs["device"]) vad_kwargs.setdefault("ncpu", kwargs.get("ncpu", 4)) if "hub" in kwargs: vad_kwargs.setdefault("hub", kwargs["hub"]) @@ -480,7 +487,7 @@ def __init__(self, **kwargs): logging.info("Building punc model.") punc_kwargs["model"] = punc_model punc_kwargs["model_revision"] = kwargs.get("punc_model_revision", "master") - punc_kwargs["device"] = kwargs["device"] + punc_kwargs.setdefault("device", kwargs["device"]) punc_kwargs.setdefault("ncpu", kwargs.get("ncpu", 4)) if "hub" in kwargs: punc_kwargs.setdefault("hub", kwargs["hub"]) @@ -496,7 +503,7 @@ def __init__(self, **kwargs): logging.info("Building SPK model.") spk_kwargs["model"] = spk_model spk_kwargs["model_revision"] = kwargs.get("spk_model_revision", "master") - spk_kwargs["device"] = kwargs["device"] + spk_kwargs.setdefault("device", kwargs["device"]) spk_kwargs.setdefault("ncpu", kwargs.get("ncpu", 4)) if "hub" in kwargs: spk_kwargs.setdefault("hub", kwargs["hub"]) @@ -688,7 +695,7 @@ def __call__(self, *args, **cfg): **cfg: Configuration overrides. """ kwargs = self.kwargs - deep_update(kwargs, cfg) + self._merge_runtime_config(kwargs, cfg) res = self.model(*args, kwargs) return res @@ -709,6 +716,8 @@ def generate(self, input, input_len=None, progress_callback=None, **cfg): input_len (tensor, optional): Length of each input sample. progress_callback (callable, optional): fn(current, total) called during processing. **cfg: Runtime parameters: + - device: Placement is selected at construction. Runtime device overrides + are ignored; create another AutoModel to use a different device. - cache (dict): State cache for streaming mode. Pass {} for first call. - hotword (str/list): Keywords to boost recognition accuracy. - postprocess_hotwords (str/list/dict): Text-level hotword correction after @@ -737,7 +746,6 @@ def generate(self, input, input_len=None, progress_callback=None, **cfg): input, input_len=input_len, progress_callback=progress_callback, **cfg ) if self.punc_model is not None: - deep_update(self.punc_kwargs, cfg) for result in results: punc_res = self.inference( result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg @@ -785,7 +793,7 @@ def inference( kwargs = self.kwargs if kwargs is None else kwargs if "cache" in kwargs: kwargs.pop("cache") - deep_update(kwargs, cfg) + self._merge_runtime_config(kwargs, cfg) model = self.model if model is None else model batch_size = kwargs.get("batch_size", 1) @@ -879,7 +887,6 @@ def inference_with_vad(self, input, input_len=None, **cfg): cfg["return_time_stamps"] = True kwargs = self.kwargs # step.1: compute the vad model - deep_update(self.vad_kwargs, cfg) beg_vad = time.time() res = self.inference( input, input_len=input_len, model=self.vad_model, kwargs=self.vad_kwargs, **cfg @@ -895,7 +902,7 @@ def inference_with_vad(self, input, input_len=None, **cfg): # step.2 compute asr model model = self.model - deep_update(kwargs, cfg) + self._merge_runtime_config(kwargs, cfg) batch_size = max(int(kwargs.get("batch_size_s", 300)) * 1000, 1) batch_size_threshold_ms = int(kwargs.get("batch_size_threshold_s", 60)) * 1000 kwargs["batch_size"] = batch_size @@ -979,7 +986,11 @@ def inference_with_vad(self, input, input_len=None, **cfg): all_segments.extend(segments) speech_b = [i[2] for i in segments] spk_res = self.inference( - speech_b, input_len=None, model=self.spk_model, kwargs=kwargs, **cfg + speech_b, + input_len=None, + model=self.spk_model, + kwargs=self.spk_kwargs, + **cfg, ) spk_embs = torch.cat([r["spk_embedding"] for r in spk_res], dim=0) results[_b]["spk_embedding"] = spk_embs @@ -1070,7 +1081,6 @@ def inference_with_vad(self, input, input_len=None, **cfg): punc_res = None punc_array = None if self.punc_model is not None and "timestamps" not in result: - deep_update(self.punc_kwargs, cfg) raw_text = copy.copy(result["text"]) punc_input_text = _join_vad_texts( item.get("text", "") for item in restored_data @@ -1322,6 +1332,14 @@ def export(self, input=None, **cfg): return export_dir + @staticmethod + def _merge_runtime_config(kwargs, cfg): + """Merge inference options without changing an already-loaded model's placement.""" + # A runtime device string does not move weights. Keep the resolved device, + # including unavailable-device fallback, for feature creation and transfer. + runtime_cfg = {key: value for key, value in cfg.items() if key != "device"} + deep_update(kwargs, runtime_cfg) + def _store_base_configs(self): """Snapshot base kwargs for all submodules to allow reset before inference.""" baseline = {} diff --git a/tests/test_punc_model_none.py b/tests/test_punc_model_none.py index e0b66f41f..c803b822c 100644 --- a/tests/test_punc_model_none.py +++ b/tests/test_punc_model_none.py @@ -22,6 +22,7 @@ def _make_auto_model(self, punc_model=None, spk_model=None, spk_mode=None): punc_model.punc_list = ["", "_", "，", "。", "？", "、"] am.punc_kwargs = {} am.spk_model = spk_model + am.spk_kwargs = {} am.cb_model = None am.spk_mode = spk_mode am.vad_kwargs = {} diff --git a/tests/test_python_api_docs_contract.py b/tests/test_python_api_docs_contract.py index cf6703343..d4a3de5de 100644 --- a/tests/test_python_api_docs_contract.py +++ b/tests/test_python_api_docs_contract.py @@ -46,7 +46,7 @@ def test_examples_compile_and_use_source_backed_keywords(self): "funasr/models/paraformer_streaming/model.py"))) for path in DOCS: blocks = python_blocks(path) - self.assertEqual(len(blocks), 4, path.name) + self.assertEqual(len(blocks), 5, path.name) for index, block in enumerate(blocks): with self.subTest(doc=path.name, block=index): compile(block, str(path), "exec") diff --git a/tests/test_submodel_device.py b/tests/test_submodel_device.py new file mode 100644 index 000000000..02b9c4a5c --- /dev/null +++ b/tests/test_submodel_device.py @@ -0,0 +1,154 @@ +"""CPU-only regressions using real AutoModel construction and inference orchestration. + +Accelerator availability is mocked only to exercise configuration labels. Stand-in +models retain CPU weights; these tests do not claim CUDA/MPS hardware execution. +""" + +import copy +import unittest +from unittest.mock import patch + +import numpy as np +import torch + +from funasr.auto.auto_model import AutoModel +from funasr.register import tables + + +class RecordingModel(torch.nn.Module): + def __init__(self, model, **kwargs): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(1)) + self.role = model + self.calls = [] + + def to(self, device): + self.placement = str(device) + return self # Configuration probe only: weights always stay on CPU. + + def inference(self, data_in, key=None, **kwargs): + if kwargs["device"] != self.placement: + raise AssertionError(f"{self.role}: feature device differs from model placement") + if kwargs["device"] == "cpu": + features = torch.zeros(1, device=kwargs["device"]) + if features.device != self.weight.device: + raise AssertionError("CPU feature/weight mismatch") + self.calls.append(copy.deepcopy(kwargs)) + results = [] + for index, _ in enumerate(data_in): + if self.role == "device-test-vad": + result = {"key": key[index], "value": [[0, 1000]]} + elif self.role == "device-test-spk": + result = {"spk_embedding": torch.ones(1, 4)} + elif self.role == "device-test-punc": + result = {"text": "hello.", "punc_array": [2]} + else: + result = {"text": "hello", "timestamp": [[0, 1000]]} + results.append(result) + return results, {"batch_data_time": 1} + + +class TestSubmodelDevice(unittest.TestCase): + def setUp(self): + self.old_threads = torch.get_num_threads() + registry = {f"device-test-{role}": RecordingModel for role in ("asr", "vad", "punc", "spk")} + self.enterContext(patch.dict(tables.model_classes, registry)) + self.enterContext(patch("torch.cuda.is_available", return_value=True)) + self.enterContext(patch("funasr.auto.auto_model.ClusterBackend")) + # Reject accidental model-hub access. model_conf={} uses the local registry. + self.enterContext( + patch("funasr.auto.auto_model.download_model", side_effect=AssertionError("network")) + ) + self.audio = np.zeros(16000, dtype=np.float32) + + def tearDown(self): + torch.set_num_threads(self.old_threads) + + def make_model(self, with_vad=True, device="cuda:0", configs=None): + configs = ( + configs + if configs is not None + else { + role: {"model_conf": {}, "device": "cpu", f"{role}_only": "kept"} + for role in ("vad", "punc", "spk") + } + ) + kwargs = dict( + model="device-test-asr", + model_conf={}, + device=device, + ncpu=1, + disable_update=True, + disable_pbar=True, + return_spk_res=False, + ) + for role in ("vad", "punc", "spk") if with_vad else ("punc",): + kwargs[f"{role}_model"] = f"device-test-{role}" + kwargs[f"{role}_kwargs"] = configs[role] + return AutoModel(**kwargs) + + def test_explicit_devices_and_caller_configs_survive_two_calls(self): + configs = { + role: {"model_conf": {"nested": [1]}, "device": "cpu", f"{role}_only": "kept"} + for role in ("vad", "punc", "spk") + } + original = copy.deepcopy(configs) + model = self.make_model(configs=configs) + self.assertEqual(configs, original) + first = model.generate(self.audio, device="cuda:7", hotword="runtime", batch_size_s=2) + second = model.generate(self.audio) + self.assertEqual(first[0]["text"], "hello.") + self.assertEqual(second[0]["text"], "hello.") + for role in ("vad", "punc", "spk"): + child = getattr(model, f"{role}_model") + self.assertEqual([call["device"] for call in child.calls], ["cpu", "cpu"]) + self.assertEqual(child.calls[0][f"{role}_only"], "kept") + self.assertEqual(child.calls[0]["hotword"], "runtime") + self.assertEqual(child.calls[0]["batch_size_s"], 2) + self.assertNotIn("hotword", child.calls[1]) + self.assertNotIn("batch_size_s", child.calls[1]) + self.assertEqual(getattr(model, f"{role}_kwargs")["device"], "cpu") + self.assertEqual([call["device"] for call in model.model.calls], ["cuda:0", "cuda:0"]) + self.assertEqual(configs, original) + + def test_punctuation_without_vad_preserves_placement(self): + model = self.make_model(with_vad=False) + model.generate(self.audio, device="cuda:7", hotword="runtime") + model.generate(self.audio) + self.assertEqual([call["device"] for call in model.punc_model.calls], ["cpu", "cpu"]) + self.assertEqual(model.punc_model.calls[0]["hotword"], "runtime") + self.assertNotIn("hotword", model.punc_model.calls[1]) + + def test_omitted_devices_inherit_resolved_asr_device(self): + for available, expected in ((True, "cuda:0"), (False, "cpu")): + with self.subTest(available=available), patch( + "torch.cuda.is_available", return_value=available + ): + configs = {role: {"model_conf": {}} for role in ("vad", "punc", "spk")} + model = self.make_model(configs=configs) + model.generate(self.audio, device="cuda:7") + model.generate(self.audio) + for role in ("vad", "punc", "spk"): + child = getattr(model, f"{role}_model") + self.assertEqual(child.placement, expected) + self.assertEqual([call["device"] for call in child.calls], [expected, expected]) + self.assertEqual(configs, {role: {"model_conf": {}} for role in configs}) + + def test_unavailable_explicit_submodel_devices_fall_back_to_cpu(self): + with patch("torch.cuda.is_available", return_value=False): + configs = { + role: {"model_conf": {}, "device": "cuda:0"} for role in ("vad", "punc", "spk") + } + model = self.make_model(device="cpu", configs=configs) + model.generate(self.audio, device="cuda:7") + model.generate(self.audio) + for role in configs: + self.assertEqual(getattr(model, f"{role}_model").placement, "cpu") + self.assertEqual( + [call["device"] for call in getattr(model, f"{role}_model").calls], ["cpu", "cpu"] + ) + self.assertEqual(configs[role]["device"], "cuda:0") + + +if __name__ == "__main__": + unittest.main() From eee63614f0346370d3918b700102590b58154a86 Mon Sep 17 00:00:00 2001 From: damnwenxi Date: Mon, 14 Sep 2026 18:42:44 +0800 Subject: [PATCH 2/2] test: avoid Python 3.11-only unittest cleanup API --- tests/test_submodel_device.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/test_submodel_device.py b/tests/test_submodel_device.py index 02b9c4a5c..7fd3e476b 100644 --- a/tests/test_submodel_device.py +++ b/tests/test_submodel_device.py @@ -51,18 +51,21 @@ def inference(self, data_in, key=None, **kwargs): class TestSubmodelDevice(unittest.TestCase): def setUp(self): self.old_threads = torch.get_num_threads() + self.addCleanup(torch.set_num_threads, self.old_threads) registry = {f"device-test-{role}": RecordingModel for role in ("asr", "vad", "punc", "spk")} - self.enterContext(patch.dict(tables.model_classes, registry)) - self.enterContext(patch("torch.cuda.is_available", return_value=True)) - self.enterContext(patch("funasr.auto.auto_model.ClusterBackend")) + self.start_patch(patch.dict(tables.model_classes, registry)) + self.start_patch(patch("torch.cuda.is_available", return_value=True)) + self.start_patch(patch("funasr.auto.auto_model.ClusterBackend")) # Reject accidental model-hub access. model_conf={} uses the local registry. - self.enterContext( + self.start_patch( patch("funasr.auto.auto_model.download_model", side_effect=AssertionError("network")) ) self.audio = np.zeros(16000, dtype=np.float32) - def tearDown(self): - torch.set_num_threads(self.old_threads) + def start_patch(self, patcher): + value = patcher.start() + self.addCleanup(patcher.stop) + return value def make_model(self, with_vad=True, device="cuda:0", configs=None): configs = (