diff --git a/examples/industrial_data_pretraining/fun_asr_nano/README.md b/examples/industrial_data_pretraining/fun_asr_nano/README.md
index 2d86c78e5c..bffd425202 100644
--- a/examples/industrial_data_pretraining/fun_asr_nano/README.md
+++ b/examples/industrial_data_pretraining/fun_asr_nano/README.md
@@ -58,6 +58,8 @@ cd Fun-ASR
pip install -r requirements.txt
```
+`finetune.sh` sets `llm_conf.sdpa_backends`, which needs `torch.nn.attention`, so fine-tuning with it requires PyTorch 2.3 or newer.
+
# TODO
diff --git a/examples/industrial_data_pretraining/fun_asr_nano/README_zh.md b/examples/industrial_data_pretraining/fun_asr_nano/README_zh.md
index 6fb09df018..90616cdd81 100644
--- a/examples/industrial_data_pretraining/fun_asr_nano/README_zh.md
+++ b/examples/industrial_data_pretraining/fun_asr_nano/README_zh.md
@@ -58,6 +58,8 @@ cd Fun-ASR
pip install -r requirements.txt
```
+`finetune.sh` 会设置 `llm_conf.sdpa_backends`，该选项依赖 `torch.nn.attention`，因此使用它进行微调需要 PyTorch 2.3 或更高版本。
+
# TODO
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/examples/industrial_data_pretraining/fun_asr_nano/requirements.txt b/examples/industrial_data_pretraining/fun_asr_nano/requirements.txt
index 8bb51b590b..4934e4f146 100644
--- a/examples/industrial_data_pretraining/fun_asr_nano/requirements.txt
+++ b/examples/industrial_data_pretraining/fun_asr_nano/requirements.txt
@@ -1,4 +1,4 @@
-torch>=2.0
+torch>=2.3
funasr
websockets>=12.0
regex
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
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..df4d293889
--- /dev/null
+++ b/tests/test_fun_asr_nano_train_opts.py
@@ -0,0 +1,154 @@
+"""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 re
+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]
+
+
+def test_recipe_declares_the_torch_floor_its_options_need():
+ """sdpa_backends goes through torch.nn.attention, absent before torch 2.3; if a recipe
+ turns it on, the example's requirements.txt has to say so."""
+ for recipe in ("finetune.sh", "lora_finetune.sh"):
+ path = os.path.join(RECIPE_DIR, recipe)
+ if not os.path.exists(path):
+ continue
+ with open(path) as f:
+ if "llm_conf.sdpa_backends" not in f.read():
+ continue
+ with open(os.path.join(RECIPE_DIR, "requirements.txt")) as f:
+ declared = re.search(r"^torch>=(\d+)\.(\d+)", f.read(), re.M)
+ assert declared, f"{recipe} sets sdpa_backends but requirements.txt pins no torch floor"
+ floor = (int(declared.group(1)), int(declared.group(2)))
+ assert floor >= (2, 3), f"{recipe} sets sdpa_backends; torch floor is {floor}, needs >= 2.3"