From d41d1c9c7cddef2cf5af31c00cc92af390ed9758 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Tue, 15 Sep 2026 19:09:15 +0000 Subject: [PATCH 1/2] fun_asr_nano: opt-in scoped SDPA backends and compiled decoder for fine-tuning --- .../fun_asr_nano/docs/finetune.md | 7 ++ .../fun_asr_nano/finetune.sh | 2 + .../fun_asr_nano/model.py | 6 ++ .../models/fun_asr_nano/llm_forward_opts.py | 66 +++++++++++++++++++ funasr/models/fun_asr_nano/model.py | 6 ++ 5 files changed, 87 insertions(+) create mode 100644 funasr/models/fun_asr_nano/llm_forward_opts.py diff --git a/examples/industrial_data_pretraining/fun_asr_nano/docs/finetune.md b/examples/industrial_data_pretraining/fun_asr_nano/docs/finetune.md index 62a1419817..af244343ae 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/docs/finetune.md +++ b/examples/industrial_data_pretraining/fun_asr_nano/docs/finetune.md @@ -87,6 +87,13 @@ For more detailed parameters, refer to: [SenseVoice Model Training and Testing]( bash finetune.sh ``` +### Training-speed options + +Two `llm_conf` keys, both off by default; `finetune.sh` turns them on. + +- `++llm_conf.torch_compile=true` runs the LLM decoder stack through `torch.compile(dynamic=True)` for inputs on a CUDA GPU (a CPU decoder and single-sequence batches stay eager); the first compiled step pays a one-time compile. +- `++llm_conf.sdpa_backends=[flash,efficient,math]` runs the LLM forward under `torch.nn.attention.sdpa_kernel` with these attention backends. It matters on sm_90 / sm_100 GPUs, where torch prefers cuDNN and cuDNN builds an execution plan for every new batch shape. The flags it sets are process-wide while the forward runs (restored on return), so a thread running attention concurrently in the same process sees the same selection; leave it unset when several models share one process. + ### Recommended Configuration - For training data less than 1000 hours, it is recommended to fine-tune the audio_adaptor. diff --git a/examples/industrial_data_pretraining/fun_asr_nano/finetune.sh b/examples/industrial_data_pretraining/fun_asr_nano/finetune.sh index 0e14453ad0..60053ec195 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/finetune.sh +++ b/examples/industrial_data_pretraining/fun_asr_nano/finetune.sh @@ -62,4 +62,6 @@ ${train_tool} \ ++audio_encoder_conf.freeze=true \ ++audio_adaptor_conf.freeze=true \ ++llm_conf.freeze=false \ +++llm_conf.torch_compile=true \ +++llm_conf.sdpa_backends=[flash,efficient,math] \ ++output_dir="${output_dir}" &> ${log_file} diff --git a/examples/industrial_data_pretraining/fun_asr_nano/model.py b/examples/industrial_data_pretraining/fun_asr_nano/model.py index af3513ce47..85ca785375 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/model.py +++ b/examples/industrial_data_pretraining/fun_asr_nano/model.py @@ -20,6 +20,7 @@ normalize_checkpoint_state, ) from funasr.models.fun_asr_nano.device_utils import resolve_autocast_device_type +from funasr.models.fun_asr_nano.llm_forward_opts import configure_llm_forward from ctc import CTC @@ -93,6 +94,11 @@ def __init__( self.llm_dtype = llm_conf.get("llm_dtype", "fp32") self.llm = model.to(dtype_map[self.llm_dtype]) + # Opt-in training-speed switches, both off by default (finetune.sh turns them on): + # llm_conf.sdpa_backends sets the process-wide SDPA backend flags while this forward + # runs and restores them on return; llm_conf.torch_compile compiles the decoder stack + # for inputs on a CUDA device (a CPU decoder stays eager). See llm_forward_opts.py. + configure_llm_forward(self.llm, llm_conf) llm_dim = model.get_input_embeddings().weight.shape[-1] # adaptor diff --git a/funasr/models/fun_asr_nano/llm_forward_opts.py b/funasr/models/fun_asr_nano/llm_forward_opts.py new file mode 100644 index 0000000000..52812e155e --- /dev/null +++ b/funasr/models/fun_asr_nano/llm_forward_opts.py @@ -0,0 +1,66 @@ +"""Opt-in switches for the Fun-ASR-Nano LLM forward, read from ``llm_conf``; both default to off. + +``llm_conf.sdpa_backends``: a list of ``flash`` / ``efficient`` / ``math``. The decoder forward runs +under ``torch.nn.attention.sdpa_kernel`` with these backends: the process-wide SDPA flags are set on +entry and restored on return, so another thread calling scaled_dot_product_attention meanwhile sees +the same selection (leave it unset when several models share one process concurrently). +``llm_conf.torch_compile``: run the decoder stack through ``torch.compile(dynamic=True)`` for +inputs on a CUDA device, decided per call; a CPU decoder and a 1-sequence batch stay eager. +""" + +import torch + +SDPA_BACKENDS = {"flash": "FLASH_ATTENTION", "efficient": "EFFICIENT_ATTENTION", "math": "MATH"} + + +def resolve_sdpa_backends(names): + """``llm_conf.sdpa_backends`` (a list of names, or ``None``) -> ``SDPBackend`` members.""" + if names is None: + return None + from torch.nn.attention import SDPBackend # a pybind enum: getattr, not subscript + + backends = [] + for name in names: + member = SDPA_BACKENDS.get(name.lower()) + if member is None: + raise ValueError( + f"llm_conf.sdpa_backends: unknown backend {name!r}, " + f"choose from {sorted(SDPA_BACKENDS)}" + ) + backends.append(getattr(SDPBackend, member)) + return backends + + +def configure_llm_forward(llm, llm_conf): + """Install the switches ``llm_conf`` asks for on ``llm.model.forward`` (nothing by default).""" + backends = resolve_sdpa_backends(llm_conf.get("sdpa_backends", None)) + compile_requested = bool(llm_conf.get("torch_compile", False)) + if backends is None and not compile_requested: + return + + decoder = llm.model + forward = eager = decoder.forward + + if compile_requested: + compiled = torch.compile(eager, dynamic=True) + + def forward(*args, **kw): + # Per call, not at construction: the trainer builds the model on the CPU and moves it + # to the GPU afterwards. A 1-sequence batch would get its own specialised graph. + x = kw.get("inputs_embeds", kw.get("input_ids")) + if x is None and args: + x = args[0] + if x is None or x.device.type != "cuda" or x.shape[0] == 1: + return eager(*args, **kw) + return compiled(*args, **kw) + + if backends is not None: + from torch.nn.attention import sdpa_kernel + + inner = forward + + def forward(*args, **kw): + with sdpa_kernel(backends): # wraps the compiled call, so Dynamo traces under it + return inner(*args, **kw) + + decoder.forward = forward diff --git a/funasr/models/fun_asr_nano/model.py b/funasr/models/fun_asr_nano/model.py index 609fe1cb96..01a1f88d84 100644 --- a/funasr/models/fun_asr_nano/model.py +++ b/funasr/models/fun_asr_nano/model.py @@ -24,6 +24,7 @@ from .ctc import CTC from .checkpoint_utils import disable_incomplete_ctc, normalize_checkpoint_state from .device_utils import resolve_autocast_device_type +from .llm_forward_opts import configure_llm_forward from .tools.utils import forced_align dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} @@ -127,6 +128,11 @@ def __init__( self.llm_dtype = llm_conf.get("llm_dtype", "fp32") self.llm = model.to(dtype_map[self.llm_dtype]) + # Opt-in training-speed switches, both off by default (finetune.sh turns them on): + # llm_conf.sdpa_backends sets the process-wide SDPA backend flags while this forward + # runs and restores them on return; llm_conf.torch_compile compiles the decoder stack + # for inputs on a CUDA device (a CPU decoder stays eager). See llm_forward_opts.py. + configure_llm_forward(self.llm, llm_conf) llm_dim = model.get_input_embeddings().weight.shape[-1] # lora: inject LoRA adapters into the LLM target Linear layers From 2c1a0c82a3be01f4bba115e7d74d95702ce6a22d Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Tue, 15 Sep 2026 19:09:16 +0000 Subject: [PATCH 2/2] tests: FunASRNano opt-ins leave defaults untouched --- tests/test_fun_asr_nano_train_opts.py | 136 ++++++++++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 tests/test_fun_asr_nano_train_opts.py diff --git a/tests/test_fun_asr_nano_train_opts.py b/tests/test_fun_asr_nano_train_opts.py new file mode 100644 index 0000000000..895bdd9ded --- /dev/null +++ b/tests/test_fun_asr_nano_train_opts.py @@ -0,0 +1,136 @@ +"""llm_conf.sdpa_backends / llm_conf.torch_compile on both copies of FunASRNano (the package +class and the recipe's model.py, imported by path). CPU-only, no weights: tiny stand-ins.""" + +import importlib.util +import os +import sys +import types + +import pytest +import torch +import torch.nn as nn + +from funasr.models.fun_asr_nano import model as funasr_model +from funasr.register import tables + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +RECIPE_DIR = os.path.join(REPO, "examples", "industrial_data_pretraining", "fun_asr_nano") + + +class _TinyDecoder(nn.Module): + """Stands in for the HF decoder stack (``llm.model``); records every eager call.""" + + def __init__(self, dim): + super().__init__() + self.proj = nn.Linear(dim, dim) + self.eager_calls = [] + + def forward(self, input_ids=None, inputs_embeds=None, **kwargs): + if not torch.compiler.is_compiling(): + cudnn = torch.backends.cuda.cudnn_sdp_enabled() + self.eager_calls.append({"batch": int(inputs_embeds.shape[0]), "cudnn": cudnn}) + return self.proj(inputs_embeds) + + +class _TinyLLM(nn.Module): + def __init__(self, dim=8, vocab=16): + super().__init__() + self.model, self.embed = _TinyDecoder(dim), nn.Embedding(vocab, dim) + + def get_input_embeddings(self): + return self.embed + + +class _TinyEncoder(nn.Module): + def __init__(self, input_size=80, **kwargs): + super().__init__() + self.lin = nn.Linear(input_size, 4) + + def output_size(self): + return 4 + + +class _TinyAdaptor(nn.Module): + def __init__(self, encoder_dim=4, llm_dim=8, **kwargs): + super().__init__() + self.lin = nn.Linear(encoder_dim, llm_dim) + + +def _load_recipe_model(): + previous = tables.model_classes.get("FunASRNano") + if RECIPE_DIR not in sys.path: + sys.path.insert(0, RECIPE_DIR) # the recipe imports ``ctc`` and ``tools.utils`` bare + path = os.path.join(RECIPE_DIR, "model.py") + spec = importlib.util.spec_from_file_location("fun_asr_nano_recipe_model", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module # tables.register() calls inspect.getfile on the class + try: + spec.loader.exec_module(module) + except ImportError as exc: + sys.modules.pop(spec.name, None) + pytest.skip(f"recipe model.py not importable here: {exc}") + finally: + tables.model_classes["FunASRNano"] = previous # the package class, registered at import + return module.FunASRNano + + +@pytest.fixture(scope="module", params=["funasr", "recipe"]) +def model_class(request): + return funasr_model.FunASRNano if request.param == "funasr" else _load_recipe_model() + + +@pytest.fixture +def build(monkeypatch, model_class): + """``build(llm_conf) -> FunASRNano`` on the CPU, with the transformers loaders stubbed.""" + fake = types.ModuleType("transformers") + fake.AutoConfig = types.SimpleNamespace(from_pretrained=lambda path, **kw: {}) + fake.AutoModelForCausalLM = types.SimpleNamespace(from_config=lambda config, **kw: _TinyLLM()) + monkeypatch.setattr(funasr_model, "AutoConfig", fake.AutoConfig) # bound at import time + monkeypatch.setattr(funasr_model, "AutoModelForCausalLM", fake.AutoModelForCausalLM) + monkeypatch.setitem(sys.modules, "transformers", fake) # the recipe imports in __init__ + monkeypatch.setitem(tables.encoder_classes, "TinyEnc", _TinyEncoder) + monkeypatch.setitem(tables.adaptor_classes, "TinyAdp", _TinyAdaptor) + kw = dict(audio_encoder="TinyEnc", audio_adaptor="TinyAdp", llm="tiny") + return lambda c: model_class(audio_encoder_conf={}, audio_adaptor_conf={}, llm_conf=c, **kw) + + +def _run(model, batch): + device = next(model.llm.parameters()).device + return model.llm.model.forward(inputs_embeds=torch.zeros(batch, 3, 8, device=device)) + + +def test_default_leaves_sdpa_flag_and_forward_alone(build): + flag = torch.backends.cuda.cudnn_sdp_enabled() + model = build({}) + assert torch.backends.cuda.cudnn_sdp_enabled() == flag + assert "forward" not in model.llm.model.__dict__ # still the class method + _run(model, 2) + assert model.llm.model.eager_calls[-1]["cudnn"] == flag + + +def test_sdpa_backends_disable_cudnn_only_inside_forward(build): + flag = torch.backends.cuda.cudnn_sdp_enabled() + model = build({"sdpa_backends": ["flash", "efficient", "math"]}) + assert torch.backends.cuda.cudnn_sdp_enabled() == flag # construction changed nothing + _run(model, 2) + assert model.llm.model.eager_calls[-1]["cudnn"] is False + assert torch.backends.cuda.cudnn_sdp_enabled() == flag # restored on return + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_torch_compile_wraps_cuda_llm_and_keeps_batch_one_eager(build): + model = build({"torch_compile": True}).cuda() # built on the CPU and moved, as the trainer does + decoder = model.llm.model + assert "forward" in decoder.__dict__ + _run(model, 1) + assert [c["batch"] for c in decoder.eager_calls] == [1] + out = _run(model, 2) # compiled: the stand-in does not record under Dynamo + assert [c["batch"] for c in decoder.eager_calls] == [1] + torch.testing.assert_close(out, decoder.proj(torch.zeros(2, 3, 8, device=out.device))) + + +def test_torch_compile_keeps_cpu_decoder_eager(build): + model = build({"torch_compile": True}) + _run(model, 1) + _run(model, 4) + assert [c["batch"] for c in model.llm.model.eager_calls] == [1, 4]