From 7376eddb110b4554594dc1e4290451d260af3f9d Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Sat, 12 Sep 2026 08:59:12 +0000 Subject: [PATCH 1/4] fun_asr_nano: disable the cuDNN SDPA backend for the LLM (per-shape plan builds cost ~1 s on 1 step in 6) Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BfjQZ3d4EfZYoEKtDbB9ms --- examples/industrial_data_pretraining/fun_asr_nano/model.py | 7 +++++++ funasr/models/fun_asr_nano/model.py | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/examples/industrial_data_pretraining/fun_asr_nano/model.py b/examples/industrial_data_pretraining/fun_asr_nano/model.py index af3513ce47..5675d43bbb 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/model.py +++ b/examples/industrial_data_pretraining/fun_asr_nano/model.py @@ -93,6 +93,13 @@ def __init__( self.llm_dtype = llm_conf.get("llm_dtype", "fp32") self.llm = model.to(dtype_map[self.llm_dtype]) + # torch >= 2.5 on sm_90+/sm_100 puts cuDNN first in the SDPA backend order. With the + # explicit padding mask the LLM passes, cuDNN builds a new execution plan (fwd + bwd, + # ~1 s on B200 / torch 2.11) for every new (batch, seq_len) pair, and token-bucketed + # ASR batches produce new pairs on ~1 step in 6. Flash / mem-efficient handle the + # masked case without per-shape plans, so take cuDNN out of the order. + if torch.cuda.is_available() and hasattr(torch.backends.cuda, "enable_cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(False) llm_dim = model.get_input_embeddings().weight.shape[-1] # adaptor diff --git a/funasr/models/fun_asr_nano/model.py b/funasr/models/fun_asr_nano/model.py index 609fe1cb96..85ac46db1d 100644 --- a/funasr/models/fun_asr_nano/model.py +++ b/funasr/models/fun_asr_nano/model.py @@ -127,6 +127,13 @@ def __init__( self.llm_dtype = llm_conf.get("llm_dtype", "fp32") self.llm = model.to(dtype_map[self.llm_dtype]) + # torch >= 2.5 on sm_90+/sm_100 puts cuDNN first in the SDPA backend order. With the + # explicit padding mask the LLM passes, cuDNN builds a new execution plan (fwd + bwd, + # ~1 s on B200 / torch 2.11) for every new (batch, seq_len) pair, and token-bucketed + # ASR batches produce new pairs on ~1 step in 6. Flash / mem-efficient handle the + # masked case without per-shape plans, so take cuDNN out of the order. + if torch.cuda.is_available() and hasattr(torch.backends.cuda, "enable_cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(False) llm_dim = model.get_input_embeddings().weight.shape[-1] # lora: inject LoRA adapters into the LLM target Linear layers From 6f1654626b2dae392066376e1238e4fa0f39679e Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Sat, 12 Sep 2026 09:32:34 +0000 Subject: [PATCH 2/4] fun_asr_nano: run the Qwen3 decoder stack through torch.compile (dynamic shapes; 1-sequence batches eager) Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BfjQZ3d4EfZYoEKtDbB9ms --- .../fun_asr_nano/model.py | 18 ++++++++++++++++++ funasr/models/fun_asr_nano/model.py | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/examples/industrial_data_pretraining/fun_asr_nano/model.py b/examples/industrial_data_pretraining/fun_asr_nano/model.py index 5675d43bbb..befe905666 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/model.py +++ b/examples/industrial_data_pretraining/fun_asr_nano/model.py @@ -100,6 +100,24 @@ def __init__( # masked case without per-shape plans, so take cuDNN out of the order. if torch.cuda.is_available() and hasattr(torch.backends.cuda, "enable_cudnn_sdp"): torch.backends.cuda.enable_cudnn_sdp(False) + # The 28 decoder layers are ~2,900 small kernels per step (fwd + bwd) under eager and + # the step is CPU launch-bound: compile the decoder stack (embeddings come in as + # inputs_embeds; lm_head and the loss stay eager). Bound-method assignment keeps the + # module tree and the state_dict keys as they are. dynamic=True: a new (B, L) every + # batch; a 1-sequence batch would be specialised into its own graph, so it runs eagerly. + if llm_conf.get("torch_compile", True) and torch.cuda.is_available(): + _eager_decoder = self.llm.model.forward + _compiled_decoder = torch.compile(_eager_decoder, dynamic=True) + + def _decoder_forward(*args, _eager=_eager_decoder, _compiled=_compiled_decoder, **kw): + x = kw.get("inputs_embeds", kw.get("input_ids")) + if x is None and args: + x = args[0] + if x is not None and x.shape[0] == 1: + return _eager(*args, **kw) + return _compiled(*args, **kw) + + self.llm.model.forward = _decoder_forward llm_dim = model.get_input_embeddings().weight.shape[-1] # adaptor diff --git a/funasr/models/fun_asr_nano/model.py b/funasr/models/fun_asr_nano/model.py index 85ac46db1d..9c7f21ab8c 100644 --- a/funasr/models/fun_asr_nano/model.py +++ b/funasr/models/fun_asr_nano/model.py @@ -134,6 +134,24 @@ def __init__( # masked case without per-shape plans, so take cuDNN out of the order. if torch.cuda.is_available() and hasattr(torch.backends.cuda, "enable_cudnn_sdp"): torch.backends.cuda.enable_cudnn_sdp(False) + # The 28 decoder layers are ~2,900 small kernels per step (fwd + bwd) under eager and + # the step is CPU launch-bound: compile the decoder stack (embeddings come in as + # inputs_embeds; lm_head and the loss stay eager). Bound-method assignment keeps the + # module tree and the state_dict keys as they are. dynamic=True: a new (B, L) every + # batch; a 1-sequence batch would be specialised into its own graph, so it runs eagerly. + if llm_conf.get("torch_compile", True) and torch.cuda.is_available(): + _eager_decoder = self.llm.model.forward + _compiled_decoder = torch.compile(_eager_decoder, dynamic=True) + + def _decoder_forward(*args, _eager=_eager_decoder, _compiled=_compiled_decoder, **kw): + x = kw.get("inputs_embeds", kw.get("input_ids")) + if x is None and args: + x = args[0] + if x is not None and x.shape[0] == 1: + return _eager(*args, **kw) + return _compiled(*args, **kw) + + self.llm.model.forward = _decoder_forward llm_dim = model.get_input_embeddings().weight.shape[-1] # lora: inject LoRA adapters into the LLM target Linear layers From 96c30ae727b5179b008354129a006fc6dd88c0df Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Tue, 15 Sep 2026 07:00:52 +0000 Subject: [PATCH 3/4] fun_asr_nano: make the SDPA backend choice and the compiled decoder explicit opt-ins Review of #3705: the constructor changed the process-wide cuDNN SDPA flag for every caller, and torch_compile defaulted to on for every model on a CUDA host, including CPU eval decoders and batched inference. - llm_conf.sdpa_backends (default None = torch's own selection): when set, e.g. [flash, efficient, math], the LLM decoder forward runs under torch.nn.attention.sdpa_kernel(...) with exactly those backends; the process flags are restored when the forward returns. torch.backends.cuda.enable_cudnn_sdp is no longer called. - llm_conf.torch_compile now defaults to False. The compiled graph is used only for inputs on a CUDA device, decided per call (funasr-train-ds builds the model on the CPU and moves it to the GPU afterwards, so a construction-time device test would never see CUDA; a decoder running on the CPU stays eager). The 1-sequence-batch eager path is unchanged. When both switches are on, the SDPA context wraps the compiled call. - Both model copies call one helper, funasr/models/fun_asr_nano/llm_forward_opts.py, so the two blocks cannot drift; the recipe already imports funasr.models.fun_asr_nano.* helpers. - finetune.sh turns both on (++llm_conf.torch_compile=true, ++llm_conf.sdpa_backends="[flash,efficient,math]"); docs/finetune.md documents the keys, the first-step compile cost and where sdpa_backends matters. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_012cdPcxsNZ5nE3qys9t14sj --- .../fun_asr_nano/docs/finetune.md | 22 +++ .../fun_asr_nano/finetune.sh | 2 + .../fun_asr_nano/model.py | 31 +--- .../models/fun_asr_nano/llm_forward_opts.py | 133 ++++++++++++++++++ funasr/models/fun_asr_nano/model.py | 31 +--- 5 files changed, 169 insertions(+), 50 deletions(-) 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..176e58939f 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,28 @@ For more detailed parameters, refer to: [SenseVoice Model Training and Testing]( bash finetune.sh ``` +### Training-speed options + +`finetune.sh` turns on two optional switches under `llm_conf`; both default to off, so +every other way of building the model (inference, `decode.py`, custom configs) runs +exactly as before. + +- `++llm_conf.torch_compile=true` — wraps the LLM decoder stack in + `torch.compile(dynamic=True)` for inputs on a CUDA GPU (a decoder running on the CPU + and batches with a single sequence still run eagerly). The first training step then + pays a one-time compile of + about 1 minute with a warm Inductor cache and about 2.5 minutes cold (measured on one + B200 with torch 2.11), repaid after a few thousand steps. +- `++llm_conf.sdpa_backends="[flash,efficient,math]"` — runs the LLM forward under + `torch.nn.attention.sdpa_kernel` with exactly these attention backends, scoped to this + model's forward (no process-wide torch flag is changed). It matters on sm_90 / sm_100 + GPUs (H100, B200), where torch puts cuDNN first in the backend order and cuDNN builds a + ~1 s execution plan for every new `(batch, padded length)` pair of token-bucketed + batches; elsewhere it changes nothing. Accepted names: `flash`, `efficient`, `math`, + `cudnn`. + +Remove either line from `finetune.sh` to go back to the default behaviour. + ### 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..27c5097f50 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 befe905666..e152381aab 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,31 +94,11 @@ def __init__( self.llm_dtype = llm_conf.get("llm_dtype", "fp32") self.llm = model.to(dtype_map[self.llm_dtype]) - # torch >= 2.5 on sm_90+/sm_100 puts cuDNN first in the SDPA backend order. With the - # explicit padding mask the LLM passes, cuDNN builds a new execution plan (fwd + bwd, - # ~1 s on B200 / torch 2.11) for every new (batch, seq_len) pair, and token-bucketed - # ASR batches produce new pairs on ~1 step in 6. Flash / mem-efficient handle the - # masked case without per-shape plans, so take cuDNN out of the order. - if torch.cuda.is_available() and hasattr(torch.backends.cuda, "enable_cudnn_sdp"): - torch.backends.cuda.enable_cudnn_sdp(False) - # The 28 decoder layers are ~2,900 small kernels per step (fwd + bwd) under eager and - # the step is CPU launch-bound: compile the decoder stack (embeddings come in as - # inputs_embeds; lm_head and the loss stay eager). Bound-method assignment keeps the - # module tree and the state_dict keys as they are. dynamic=True: a new (B, L) every - # batch; a 1-sequence batch would be specialised into its own graph, so it runs eagerly. - if llm_conf.get("torch_compile", True) and torch.cuda.is_available(): - _eager_decoder = self.llm.model.forward - _compiled_decoder = torch.compile(_eager_decoder, dynamic=True) - - def _decoder_forward(*args, _eager=_eager_decoder, _compiled=_compiled_decoder, **kw): - x = kw.get("inputs_embeds", kw.get("input_ids")) - if x is None and args: - x = args[0] - if x is not None and x.shape[0] == 1: - return _eager(*args, **kw) - return _compiled(*args, **kw) - - self.llm.model.forward = _decoder_forward + # Opt-in training-speed switches, both off by default (finetune.sh turns them on): + # llm_conf.sdpa_backends scopes the SDPA backend choice to this forward (no + # process-wide flag is changed); 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..16a11f0c5a --- /dev/null +++ b/funasr/models/fun_asr_nano/llm_forward_opts.py @@ -0,0 +1,133 @@ +"""Opt-in switches for the Fun-ASR-Nano LLM forward, read from ``llm_conf``. + +Both are off unless the config sets them (``finetune.sh`` turns them on for the +fine-tuning recipe); with the defaults nothing here touches the model or any +process-wide torch setting. + +``llm_conf.sdpa_backends`` (default ``None``): a list of scaled-dot-product-attention +backend names, e.g. ``[flash, efficient, math]``. When set, the LLM decoder forward runs +under ``torch.nn.attention.sdpa_kernel(...)`` with exactly these backends enabled, and the +process-wide flags (``torch.backends.cuda.enable_cudnn_sdp`` and friends) are restored as +soon as the forward returns; other models in the same process are not affected. The +reason to set it: torch >= 2.5 on sm_90 / sm_100 puts cuDNN first in the backend order, +and with the explicit padding mask the LLM passes, cuDNN builds a new execution plan +(fwd + bwd, about 1 s on B200 / torch 2.11) for every new (batch, seq_len) pair; +token-bucketed ASR batches produce a new pair on about 1 step in 6. Flash and +mem-efficient attention handle the masked case without per-shape plans. + +``llm_conf.torch_compile`` (default ``False``): wrap the decoder stack (``llm.model``, +the layers between the input embeddings and ``lm_head``) in +``torch.compile(dynamic=True)``. The compiled graph is used only for inputs on a CUDA +device (decided per call, because the trainer builds the model on the CPU and moves it +to the GPU afterwards; a decoder running on the CPU stays eager), and a batch with a +single sequence runs the eager forward because Dynamo would specialise a size-1 +dimension into its own graph. The 28 decoder layers are about +2,900 small kernels per step (fwd + bwd) under eager and the fine-tuning step is +CPU-launch-bound, so compiling them roughly halves the launch count. The first +compiled step costs about 1 minute with a warm Inductor cache and 2.5 minutes cold +(B200, torch 2.11). + +The switches are installed by bound-method assignment on ``llm.model.forward``, so the +module tree and the state_dict keys are unchanged. When both are on, the SDPA context +wraps the compiled call, so the backend choice is in effect while Dynamo traces the +decoder (the compiled graph then keeps that backend for forward and backward). +""" + +import torch + +# Short names accepted in ``llm_conf.sdpa_backends`` -> ``torch.nn.attention.SDPBackend`` +# member. The full member names (``FLASH_ATTENTION`` ...) are accepted as well. +SDPA_BACKEND_ALIASES = { + "flash": "FLASH_ATTENTION", + "efficient": "EFFICIENT_ATTENTION", + "math": "MATH", + "cudnn": "CUDNN_ATTENTION", +} + + +def resolve_sdpa_backends(names): + """Map ``llm_conf.sdpa_backends`` to a list of ``SDPBackend`` members. + + ``names`` may be ``None`` (returns ``None``: leave torch's backend selection alone), + a list / tuple / ``ListConfig`` of names, or one comma-separated string such as + ``"flash,efficient,math"`` (surrounding brackets are ignored, so the string form of + a Hydra list override works too). Names are case-insensitive. Raises ``ValueError`` + for an empty list or an unknown name. + """ + if names is None: + return None + if isinstance(names, str): + names = names.strip().strip("[]()").split(",") + names = [str(name).strip() for name in names] + names = [name for name in names if name] + if not names: + raise ValueError( + "llm_conf.sdpa_backends must name at least one backend " + f"(choose from {sorted(SDPA_BACKEND_ALIASES)}) or be left unset" + ) + try: + from torch.nn.attention import SDPBackend + except ImportError as exc: # torch < 2.3 + raise ValueError( + "llm_conf.sdpa_backends needs torch.nn.attention.sdpa_kernel (torch >= 2.3)" + ) from exc + backends = [] + for name in names: + member = SDPA_BACKEND_ALIASES.get(name.lower(), name.upper()) + if member == "ERROR" or not hasattr(SDPBackend, member): + raise ValueError( + f"llm_conf.sdpa_backends: unknown backend {name!r}; " + f"choose from {sorted(SDPA_BACKEND_ALIASES)}" + ) + 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``. + + Returns the installed forward, or ``None`` when nothing was installed (the default). + The installed forward carries a ``funasr_nano_opts`` dict (``torch_compile``: bool, + ``sdpa_backends``: list or None) describing what is in effect. + """ + 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 None + + decoder = llm.model + eager = decoder.forward + forward = eager + compiled = False + + if compile_requested: + # torch.compile is lazy (nothing is traced until the first call), so the device + # test is made per call on the decoder's input: funasr-train-ds builds the model on + # the CPU and moves it to the GPU afterwards, and a decoder that runs on the CPU + # must stay eager. + compiled_decoder = torch.compile(eager, dynamic=True) + + def forward(*args, _eager=eager, _compiled=compiled_decoder, **kw): + 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) + + compiled = True + + if backends is not None: + from torch.nn.attention import sdpa_kernel + + inner = forward + + def forward(*args, _inner=inner, _backends=backends, **kw): + with sdpa_kernel(_backends): + return _inner(*args, **kw) + + if forward is eager: + return None + forward.funasr_nano_opts = {"torch_compile": compiled, "sdpa_backends": backends} + decoder.forward = forward + return forward diff --git a/funasr/models/fun_asr_nano/model.py b/funasr/models/fun_asr_nano/model.py index 9c7f21ab8c..db3f5ee2a3 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,31 +128,11 @@ def __init__( self.llm_dtype = llm_conf.get("llm_dtype", "fp32") self.llm = model.to(dtype_map[self.llm_dtype]) - # torch >= 2.5 on sm_90+/sm_100 puts cuDNN first in the SDPA backend order. With the - # explicit padding mask the LLM passes, cuDNN builds a new execution plan (fwd + bwd, - # ~1 s on B200 / torch 2.11) for every new (batch, seq_len) pair, and token-bucketed - # ASR batches produce new pairs on ~1 step in 6. Flash / mem-efficient handle the - # masked case without per-shape plans, so take cuDNN out of the order. - if torch.cuda.is_available() and hasattr(torch.backends.cuda, "enable_cudnn_sdp"): - torch.backends.cuda.enable_cudnn_sdp(False) - # The 28 decoder layers are ~2,900 small kernels per step (fwd + bwd) under eager and - # the step is CPU launch-bound: compile the decoder stack (embeddings come in as - # inputs_embeds; lm_head and the loss stay eager). Bound-method assignment keeps the - # module tree and the state_dict keys as they are. dynamic=True: a new (B, L) every - # batch; a 1-sequence batch would be specialised into its own graph, so it runs eagerly. - if llm_conf.get("torch_compile", True) and torch.cuda.is_available(): - _eager_decoder = self.llm.model.forward - _compiled_decoder = torch.compile(_eager_decoder, dynamic=True) - - def _decoder_forward(*args, _eager=_eager_decoder, _compiled=_compiled_decoder, **kw): - x = kw.get("inputs_embeds", kw.get("input_ids")) - if x is None and args: - x = args[0] - if x is not None and x.shape[0] == 1: - return _eager(*args, **kw) - return _compiled(*args, **kw) - - self.llm.model.forward = _decoder_forward + # Opt-in training-speed switches, both off by default (finetune.sh turns them on): + # llm_conf.sdpa_backends scopes the SDPA backend choice to this forward (no + # process-wide flag is changed); 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 8482dd2a92b77cc2a6c82d01fad54b75ccf1756a Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Tue, 15 Sep 2026 07:00:52 +0000 Subject: [PATCH 4/4] tests: FunASRNano default leaves the SDPA flag and llm.model.forward alone; opt-ins install scoped wrappers CPU-only, no weights: tiny stand-ins for the LLM (a .model.forward and get_input_embeddings), encoder and adaptor, so FunASRNano.__init__ runs end to end for both the built-in class and the recipe's model.py (loaded by path). - default / explicitly disabled llm_conf: torch.backends.cuda.cudnn_sdp_enabled() is unchanged after construction and llm.model.forward is the class method - torch_compile=true with the decoder running on the CPU (an eval decoder on a GPU host) takes the eager path on every call; built on the CPU and moved to a CUDA device (the trainer's order) it uses the compiled graph, keeps 1-sequence batches eager and matches the eager result - sdpa_backends=[flash, efficient, math]: cuDNN is off and flash on inside the forward only; the flags are restored on return and unchanged by construction - both switches compose on CUDA; unknown or empty backend names raise ValueError - the two ++llm_conf.* lines of finetune.sh parse through Hydra's override parser into the values the constructor expects Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_012cdPcxsNZ5nE3qys9t14sj --- tests/test_fun_asr_nano_train_opts.py | 295 ++++++++++++++++++++++++++ 1 file changed, 295 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..3733bcf1f2 --- /dev/null +++ b/tests/test_fun_asr_nano_train_opts.py @@ -0,0 +1,295 @@ +"""The llm_conf.sdpa_backends / llm_conf.torch_compile switches of FunASRNano. + +CPU-only, no weights: the LLM loader is replaced by a tiny stand-in and the encoder / +adaptor registries get tiny modules, so FunASRNano.__init__ runs end to end. Both copies +of the class are exercised: funasr.models.fun_asr_nano.model and the recipe's +examples/industrial_data_pretraining/fun_asr_nano/model.py (loaded by path). +""" + +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 llm_forward_opts +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``).""" + + 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(): + self.eager_calls.append( + { + "batch": int(inputs_embeds.shape[0]), + "cudnn": torch.backends.cuda.cudnn_sdp_enabled(), + "flash": torch.backends.cuda.flash_sdp_enabled(), + } + ) + return self.proj(inputs_embeds) + + +class _TinyLLM(nn.Module): + """Stands in for AutoModelForCausalLM.from_config(...).""" + + def __init__(self, dim=8, vocab=16): + super().__init__() + self.model = _TinyDecoder(dim) + self.embed = nn.Embedding(vocab, dim) + + def get_input_embeddings(self): + return self.embed + + def gradient_checkpointing_enable(self): + pass + + +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, **kwargs): + super().__init__() + self.lin = nn.Linear(kwargs.get("encoder_dim", 4), kwargs.get("llm_dim", 8)) + + +def _load_recipe_model_module(): + """Import examples/.../fun_asr_nano/model.py by path, keeping the registry entry.""" + 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 + spec = importlib.util.spec_from_file_location( + "fun_asr_nano_recipe_model_for_test", os.path.join(RECIPE_DIR, "model.py") + ) + 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: # e.g. torchaudio / soundfile missing in a minimal CI env + sys.modules.pop(spec.name, None) + pytest.skip(f"recipe model.py not importable here: {exc}") + finally: + if previous is not None: + tables.model_classes["FunASRNano"] = previous + return module + + +@pytest.fixture(scope="module", params=["funasr", "recipe"]) +def model_class(request): + if request.param == "funasr": + return funasr_model.FunASRNano + return _load_recipe_model_module().FunASRNano + + +@pytest.fixture +def build(monkeypatch, model_class): + """Return ``build(llm_conf, device="cpu") -> FunASRNano`` with the loaders stubbed.""" + device = {"value": "cpu"} + + class _FakeAutoConfig: + @staticmethod + def from_pretrained(path, **kwargs): + return {"path": path} + + class _FakeAutoModelForCausalLM: + @staticmethod + def from_config(config, **kwargs): + return _TinyLLM().to(device["value"]) + + fake_transformers = types.ModuleType("transformers") + fake_transformers.AutoConfig = _FakeAutoConfig + fake_transformers.AutoModelForCausalLM = _FakeAutoModelForCausalLM + # the built-in class binds the names at import time, the recipe imports inside __init__ + monkeypatch.setattr(funasr_model, "AutoConfig", _FakeAutoConfig) + monkeypatch.setattr(funasr_model, "AutoModelForCausalLM", _FakeAutoModelForCausalLM) + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + for table in ("encoder_classes", "adaptor_classes"): + if not hasattr(tables, table): + monkeypatch.setattr(tables, table, {}, raising=False) + monkeypatch.setitem(tables.encoder_classes, "TinyEncoderForTest", _TinyEncoder) + monkeypatch.setitem(tables.adaptor_classes, "TinyAdaptorForTest", _TinyAdaptor) + + def _build(llm_conf, device_type="cpu"): + device["value"] = device_type + return model_class( + audio_encoder="TinyEncoderForTest", + audio_encoder_conf={}, + audio_adaptor="TinyAdaptorForTest", + audio_adaptor_conf={}, + llm="tiny", + llm_conf=dict(llm_conf), + ) + + return _build + + +def _forward_is_plain(model): + """True when llm.model.forward is the class method, i.e. nothing was installed.""" + return "forward" not in model.llm.model.__dict__ + + +def _run(model, batch): + x = torch.zeros(batch, 3, 8, device=next(model.llm.parameters()).device) + return model.llm.model.forward(inputs_embeds=x) + + +def test_default_llm_conf_leaves_sdpa_flag_and_forward_alone(build): + flag_before = torch.backends.cuda.cudnn_sdp_enabled() + flash_before = torch.backends.cuda.flash_sdp_enabled() + model = build({}) + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before + assert torch.backends.cuda.flash_sdp_enabled() == flash_before + assert _forward_is_plain(model) + _run(model, 2) + assert model.llm.model.eager_calls[-1]["cudnn"] == flag_before + + +def test_explicitly_disabled_matches_default(build): + flag_before = torch.backends.cuda.cudnn_sdp_enabled() + model = build({"torch_compile": False, "sdpa_backends": None}) + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before + assert _forward_is_plain(model) + + +def test_torch_compile_keeps_cpu_decoder_eager(build): + # A decoder running on the CPU (e.g. an eval decoder on a GPU host) stays eager even + # when asked: the wrapper is installed, but every CPU call takes the eager path. + flag_before = torch.backends.cuda.cudnn_sdp_enabled() + model = build({"torch_compile": True}) + assert model.llm.model.forward.funasr_nano_opts == { + "torch_compile": True, + "sdpa_backends": None, + } + decoder = model.llm.model + _run(model, 1) + _run(model, 4) + assert [c["batch"] for c in decoder.eager_calls] == [1, 4] + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before + + +@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): + # Built on the CPU and moved afterwards, the way funasr-train-ds does it. + flag_before = torch.backends.cuda.cudnn_sdp_enabled() + model = build({"torch_compile": True}).cuda() + assert not _forward_is_plain(model) + opts = model.llm.model.forward.funasr_nano_opts + assert opts == {"torch_compile": True, "sdpa_backends": None} + decoder = model.llm.model + _run(model, 1) # 1-sequence batch: the eager path + assert [c["batch"] for c in decoder.eager_calls] == [1] + out = _run(model, 2) # compiled path: 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))) + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before + + +def test_sdpa_backends_scoped_to_forward(build): + from torch.nn.attention import SDPBackend + + flag_before = torch.backends.cuda.cudnn_sdp_enabled() + flash_before = torch.backends.cuda.flash_sdp_enabled() + model = build({"sdpa_backends": ["flash", "efficient", "math"]}) + assert not _forward_is_plain(model) + opts = model.llm.model.forward.funasr_nano_opts + assert opts["torch_compile"] is False + assert opts["sdpa_backends"] == [ + SDPBackend.FLASH_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.MATH, + ] + # construction changed nothing process-wide + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before + assert torch.backends.cuda.flash_sdp_enabled() == flash_before + _run(model, 2) + seen = model.llm.model.eager_calls[-1] + # inside the forward: cuDNN off, flash on ... + assert seen["cudnn"] is False + assert seen["flash"] is True + # ... and restored on return + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before + assert torch.backends.cuda.flash_sdp_enabled() == flash_before + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_both_switches_compose_on_cuda(build): + flag_before = torch.backends.cuda.cudnn_sdp_enabled() + model = build( + {"torch_compile": True, "sdpa_backends": ["flash", "efficient", "math"]}, + device_type="cuda", + ) + opts = model.llm.model.forward.funasr_nano_opts + assert opts["torch_compile"] is True and len(opts["sdpa_backends"]) == 3 + decoder = model.llm.model + _run(model, 1) # eager path, still under the context + assert decoder.eager_calls[-1] == {"batch": 1, "cudnn": False, "flash": True} + _run(model, 2) # compiled path + assert len(decoder.eager_calls) == 1 + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before + + +def test_invalid_backend_name_raises(build): + flag_before = torch.backends.cuda.cudnn_sdp_enabled() + with pytest.raises(ValueError, match="unknown backend 'cudnnn'"): + build({"sdpa_backends": ["flash", "cudnnn"]}) + with pytest.raises(ValueError, match="at least one backend"): + build({"sdpa_backends": []}) + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before + + +def test_resolve_sdpa_backends_accepts_list_string_and_full_names(): + from torch.nn.attention import SDPBackend + + expect = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH] + assert llm_forward_opts.resolve_sdpa_backends(None) is None + assert llm_forward_opts.resolve_sdpa_backends(["flash", "efficient", "math"]) == expect + assert llm_forward_opts.resolve_sdpa_backends("[flash, efficient, math]") == expect + assert llm_forward_opts.resolve_sdpa_backends(("FLASH_ATTENTION", "Math")) == [ + SDPBackend.FLASH_ATTENTION, + SDPBackend.MATH, + ] + assert llm_forward_opts.resolve_sdpa_backends(["cudnn"]) == [SDPBackend.CUDNN_ATTENTION] + + +def test_finetune_sh_overrides_parse_and_enable_both_switches(build): + """The two ``++llm_conf.*`` lines of finetune.sh, parsed the way funasr-train-ds does.""" + hydra_parser = pytest.importorskip("hydra.core.override_parser.overrides_parser") + with open(os.path.join(RECIPE_DIR, "finetune.sh")) as f: + script = f.read() + lines = re.findall(r'(\+\+llm_conf\.(?:torch_compile|sdpa_backends)=\S+)', script) + assert len(lines) == 2, lines + # the shell strips the double quotes before hydra sees the argument + overrides = [line.replace('"', "") for line in lines] + parsed = hydra_parser.OverridesParser.create().parse_overrides(overrides) + llm_conf = {o.key_or_group.split(".", 1)[1]: o.value() for o in parsed} + assert llm_conf["torch_compile"] is True + assert llm_conf["sdpa_backends"] == ["flash", "efficient", "math"] + + flag_before = torch.backends.cuda.cudnn_sdp_enabled() + model = build(llm_conf) + assert model.llm.model.forward.funasr_nano_opts["torch_compile"] is True + assert len(model.llm.model.forward.funasr_nano_opts["sdpa_backends"]) == 3 + _run(model, 2) # on the CPU: eager path, inside the SDPA scope + assert model.llm.model.eager_calls[-1] == {"batch": 2, "cudnn": False, "flash": True} + assert torch.backends.cuda.cudnn_sdp_enabled() == flag_before