[perf] Fun-ASR-Nano fine-tuning on B200: 0.244 s to 0.068 s per step at 1 GPU, mostly from cuDNN SDPA plan builds on every new batch shape

[TarzanZhao]

Summary

Two opt-in switches for the Fun-ASR-Nano fine-tuning recipe, both read from llm_conf in FunASRNano.__init__ and both off by default, so every other way of building the model (inference, decode.py, custom configs) is unchanged. finetune.sh turns both on. 4 commits: the two original ones, then the review round (96c30ae72 scoping and defaults, 8482dd2a9 tests).

1. llm_conf.sdpa_backends (default None = torch's own backend 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-wide flags are restored when the forward returns, so the choice is scoped to this model's forward and nothing else in the process is affected. Why: on sm_90 and sm_100 with torch 2.11, the Qwen3 attention with its explicit padding mask goes to cuDNN, which builds a new execution plan for every (batch, padded length) pair, about 1 s each. Token-bucketed ASR batches produce a new pair on about 1 step in 6: 59 of 382 steps, 56 s of a 96 s training phase. Flash and mem-efficient attention handle the masked case without per-shape plans.
2. llm_conf.torch_compile (default False). When true, the decoder stack (llm.model) is wrapped in torch.compile(dynamic=True); 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 decoder running on the CPU stays eager. After the plan builds are gone, a step is 8,080 kernel launches for 37.5 ms of GPU work: the host is the bottleneck, and the 28 decoder layers are about 2,900 of those launches; the compiled decoder halves the launch count. Embeddings still come in as inputs_embeds; lm_head and the loss stay eager. A batch with one sequence runs the eager forward, because Dynamo would specialise a dimension of size 1 into its own graph.

Both model copies (funasr/models/fun_asr_nano/model.py and the recipe's examples/industrial_data_pretraining/fun_asr_nano/model.py, loaded by trust_remote_code=true) call one helper, funasr/models/fun_asr_nano/llm_forward_opts.py, installed by bound-method assignment on llm.model.forward (module tree and state_dict keys unchanged). When both switches are on, the SDPA context wraps the compiled call. docs/finetune.md documents both keys.

Result

Measured on 1 × B200; not measured elsewhere. finetune.sh as shipped (LLM fine-tuning, encoder and adaptor frozen; on this branch that includes its two new ++llm_conf.* lines), 3600 AISHELL-1 utterances = 382 steps, seed 1234, two interleaved runs per arm against the unmodified 486b4b7ce on the same setup, nothing set in the environment:

	baseline (main)	this PR, finetune.sh overrides on	change
step time, mean of steps 2 to 382	0.2429 s (0.2404, 0.2453)	0.0680 s (0.0685, 0.0676)	3.57x, −72.0 %
step time, median	0.0960 s (0.095, 0.097)	0.0645 s (0.065, 0.064)	1.49x
steps over 0.5 s (cuDNN plan builds, 1.0 to 1.13 s each)	59 of 381	0
forward / backward / optimizer, median	52-53 / 29-30 / 7 ms	40 / 11 / 7 ms
kernel launches per step, steps without a plan build	8,080	4,530
peak allocated memory	7.39 GB	6.82 GB
step 1	2.2 s	64 s (warm Inductor cache), 157 s (cold)	one-time compile
training phase, 382 steps	93.8 / 95.7 s	89.7 / 90.7 s including the compile
of which: sdpa_backends alone (two runs, same day and setup)		0.0996 s mean (0.1019, 0.0974), 0.0965 s median, 0 slow steps	2.44x on the mean
of which: torch_compile on top		0.0680 s mean, 0.0645 s median	a further 1.46x

Run-to-run spread: 4.9 ms on the baseline mean, 0.9 ms on this branch; no recompile in 382 steps. With the defaults (no overrides) this branch is the baseline computation: nothing is installed on the model.

Type of change

• Bug fix
• Documentation
• Example or demo
• Runtime or deployment
• Benchmark or evaluation
• Model/training change

Validation

Measured on 1 × B200 against the unmodified 486b4b7ce, both arms with finetune.sh's own flags and nothing set in the environment; full command under Details.

• python -m pytest tests/test_fun_asr_nano_train_opts.py — 17 passed with a CUDA device, 13 passed / 4 skipped (the CUDA-only cases) with CUDA_VISIBLE_DEVICES=""; both model copies are parametrised through every test
• python -m compileall funasr examples tests
• tests/test_fun_asr_nano_lora_injection.py, test_fun_asr_nano_vllm_dtype.py, test_fun_asr_nano_autocast_device.py still pass (14)
• Docs: examples/industrial_data_pretraining/fun_asr_nano/docs/finetune.md, new "Training-speed options" section
• Runtime/deployment command tested (training-only change)

Correctness Verification

Recorded with probe on both arms under the same setup (hooks on TarzanZhao/FunASR:perf/fun-asr-nano-training-speed-verify, this branch plus one commit, not in this PR): 6,128 checkpoints per run, baseline vs this branch with the finetune.sh overrides on.

6,105 of 6,128 checkpoints pass the declared tolerances; 23 do not. The failures are 2 of 382 per-step losses, 2 of 7 gradient norms and 19 argmax-token counts; the differences are at bf16 kernel-rounding scale (the switch moves the attention from cuDNN to the flash / mem-efficient kernels, and Inductor keeps fused intermediates in fp32, so the compiled bf16 forward rounds differently from eager), but they are failures of the gate as declared, not numerical equivalence.

recorded	difference, baseline vs this branch	tolerance	gate
per-step loss, steps 1 to 382	0.0128 max, 0.0014 mean; steps 6 (0.0128) and 41 (0.0115) over	0.01 absolute	unmet on 2 of 382
validation loss and accuracy, 200 dev utterances	0.00013 absolute; accuracy identical	0.001 absolute	met
logits abs-mean and std (fp32), every step	0.61 % / 0.67 % max, 0.16 % / 0.19 % mean	1 %	met
argmax tokens equal to the label, every step	3 of 34,886 tokens net; 23 steps differ by one token, 11 of them over the per-step budget	1 % per step	unmet on 11 steps (19 checkpoints)
trainable-parameter L2 norm, steps 1, 2, 10, 50, 100, 200, 374	3.4e-9 relative	1e-4	met
gradient norm before clipping, same steps	0 to 0.8 % on five steps; 2.0 % at steps 200 (1.9766 vs 1.9375) and 374 (1.5312 vs 1.5625), bf16	1 %	unmet on 2 of 7
batch shapes and speech-frame counts, every step	identical	exact	met

The record of this head (opt-ins through sdpa_kernel and torch_compile=true) is identical in all 6,128 values to the record of the previous revision (process-wide enable_cudnn_sdp(False) plus compile on by default), so the scoping changed nothing in what is computed. If the project wants bit-identity with eager cuDNN attention, the recipe should not set the two overrides.

User impact

Nobody who does not set the two keys. With the defaults, FunASRNano.__init__ installs nothing and changes no torch setting; the regression tests hold both copies to that. Users of finetune.sh on a CUDA GPU get the switches: on sm_90 and sm_100 with a torch that puts cuDNN first in the SDPA order (2.5 and later), sdpa_backends removes a plan build of about 1 s on every new batch shape and the two together give the table above on the one B200 measured; on GPUs where torch does not pick cuDNN the SDPA scope changes nothing and the compiled decoder alone gave about 28 % in the previous round. The cost is a one-time compile on the first step (about 1 minute with a warm Inductor cache, 2.5 minutes cold on B200); removing either line from finetune.sh turns that switch off. The encoder, the adaptor, non-CUDA runs and inference paths are untouched.

Notes for reviewers

• Review round (2026-09-13 comments). Process-wide SDPA policy: enable_cudnn_sdp(False) is removed; the choice is llm_conf.sdpa_backends, scoped with sdpa_kernel to this model's forward, and test_default_llm_conf_leaves_sdpa_flag_and_forward_alone / test_sdpa_backends_scoped_to_forward assert the flag is unchanged after construction and after a forward. Eager default: torch_compile defaults to False, the compiled graph is used only for inputs on a CUDA device (per call, not host CUDA at construction), and test_torch_compile_keeps_cpu_decoder_eager covers the CPU-decoder-on-GPU-host case, test_torch_compile_wraps_cuda_llm_and_keeps_batch_one_eager the enabled path; the same tests run against the recipe copy. Timing claims are scoped to the B200 experiment; the out-of-tolerance checkpoints are reported as unmet gates above; the frozen-encoder train/eval policy is not touched.
• Risk: rounding. Neither switch is bit-identical to eager cuDNN attention (table above).
• Risk: shapes. dynamic=True handles the varying (batch, length) of token-bucketed batches without recompiles in this run (0 in 382 steps). Batches of exactly one sequence take the eager path on purpose.
• The compile pays for itself: it saves about 32 ms per step over the SDPA scope alone, so 64 s of warm compile is repaid in about 2,000 steps and 157 s cold in about 5,000, far below the recipe's 50 epochs.
• TORCH_CUDNN_SDPA_DEPRIORITIZED=1 in the environment has the same effect as sdpa_backends=[flash,efficient,math] for anyone who prefers not to change the recipe (records bit-identical, previous round). The profile behind it is in Fine-tuning Fun-ASR-Nano on a B200: one step in six takes 1 s because cuDNN SDPA builds a plan for every new batch shape #3704.
• The helper is called from the model class, so the switches apply to any trainer path that instantiates FunASRNano with the keys set; DeepSpeed (use_deepspeed=true) was not tested.
• Other things the profile showed, none in this PR: the frozen encoder runs in train mode with dropout on every step (trainer_ds.py calls model.train() each epoch); the SANM encoder is about 2,800 launches per step and would compile too if the dropout question were settled; AdamW keeps bf16 optimizer state for the LLM (fused AdamW changes rounding, so it was left out); reduce-overhead mode re-records a CUDA graph per shape and loses.

Details: hardware, model, full command, traces

Hardware. 1 × NVIDIA B200 (sm_100, 183 GB) on an 8-GPU node, driver 580.126.20, CUDA 12.8 (the torch build), cuDNN 9.19, torch 2.11.0+cu128; the process pinned with numactl --cpunodebind=0 --membind=0 to the NUMA node of GPU 0. Weights: the ModelScope snapshot of FunAudioLLM/Fun-ASR-Nano-2512 on local disk; AISHELL-1 wavs on node-local disk.

Model. SenseVoiceEncoderSmall (70 SANM layers, 221 M, frozen, fp32 with TF32 matmuls) + Transformer adaptor (12.6 M, frozen) + Qwen3-0.6B (28 layers, 596 M, bf16, trained) + a CTC decoder (39 M, trainable, not called in training); 868.86 M parameters, 635.12 M trainable. AdamW lr 2e-4, 2500 warm-up steps, grad clip 5. Each step is a token-budget batch of 6000 (speech_length + text_length) with at most 10 utterances: on average 9.6 utterances, 640 LLM tokens and 722 fbank frames, with a new (batch, padded length) pair on about 1 step in 6.

Data (once): the recipe's tools/scp2jsonl.py on an AISHELL-1 wav.scp and text pair produces train.jsonl and val.jsonl (docs/finetune.md:15-53); 3600 training utterances (the first 3600 of a seeded sample of 36,000) and 200 dev utterances.

The measured job, both arms; the checkout under test goes first on PYTHONPATH so import funasr resolves to it. The two ++llm_conf.torch_compile=true ++llm_conf.sdpa_backends=[flash,efficient,math] lines are passed on this branch only (they are the two lines finetune.sh adds; the baseline checkout does not know the keys):

export HF_HUB_OFFLINE=1 CUDA_VISIBLE_DEVICES=0
export TORCHINDUCTOR_CACHE_DIR=<persistent dir> TRITON_CACHE_DIR=<persistent dir>
cd examples/industrial_data_pretraining/fun_asr_nano
numactl --cpunodebind=0 --membind=0 \
torchrun --nnodes 1 --nproc_per_node 1 --node_rank 0 --master_addr 127.0.0.1 --master_port 26669 \
  $(which funasr-train-ds) \
  hydra.run.dir=<out>/hydra \
  ++seed=1234 ++model=<model dir> ++trust_remote_code=true \
  ++train_data_set_list=<data>/train.jsonl ++valid_data_set_list=<data>/val.jsonl \
  ++dataset_conf.data_split_num=1 ++dataset_conf.batch_sampler=BatchSampler ++dataset_conf.batch_size=6000 \
  ++dataset_conf.sort_size=1024 ++dataset_conf.batch_type=token ++dataset_conf.num_workers=4 \
  ++train_conf.max_epoch=1 ++train_conf.log_interval=1 ++train_conf.resume=true \
  ++train_conf.validate_interval=2000 ++train_conf.save_checkpoint_interval=2000 \
  ++train_conf.effective_save_name_excludes=None ++train_conf.keep_nbest_models=20 ++train_conf.avg_nbest_model=10 \
  ++train_conf.use_deepspeed=false ++train_conf.deepspeed_config=deepspeed_conf/ds_stage1.json \
  ++train_conf.find_unused_parameters=true \
  ++optim_conf.lr=0.0002 ++audio_encoder_conf.freeze=true ++audio_adaptor_conf.freeze=true ++llm_conf.freeze=false \
  ++output_dir=<out> \
  ++llm_conf.torch_compile=true ++llm_conf.sdpa_backends=[flash,efficient,math]   # this branch only

trust_remote_code=true loads the recipe's own model.py from the working directory, which is why the change is mirrored into funasr/models/fun_asr_nano/model.py. ++train_conf.find_unused_parameters=true is a no-op on one GPU; it is there because the multi-GPU form of the script needs it (the CTC decoder is trainable but unused). The trainer has no max-steps flag, so the run length is the dataset. ++seed=1234 seeds torch, numpy and random; the batch sampler shuffles with manual_seed(epoch), so the batch order is fixed.

Correctness run. On the verify branch, the same command with PROBE=1 PROBE_OUT=<dir>/code.json; for the baseline record, check out 486b4b7ce, cherry-pick the hooks commit, run into <dir>/base.json; probe compare <dir>/base.rank0.json <dir>/code.rank0.json prints every checkpoint against its tolerance. Also swept on the baseline and rejected as noise: OMP_NUM_THREADS 8 and 1, PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, ++cudnn_benchmark=true, PYTORCH_NVML_BASED_CUDA_CHECK=1; excluded because they change the data or the numerics: ++dataset_conf.num_workers=8, ++optim_conf.fused=true. Earlier revisions of this PR: compile alone measured 1.38x (0.0985 to 0.0713 s) against a baseline with TORCH_CUDNN_SDPA_DEPRIORITIZED=1 set; the process-wide switch plus compile-on-by-default measured 0.2438 to 0.0680 s (3.59x) on the same setup as this table.

Traces (torch.profiler, rank 0, steps 4 to 9 of the unmodified script and steps 9 to 14 of the optimized arm, previous revision, same computation; open in https://ui.perfetto.dev):

• baseline: https://drive.google.com/open?id=14roiLxfbM3zmsxTgmcO2QKoiZnS4B_n_
• optimized: https://drive.google.com/open?id=1Md6OxOjDtvy0rfCtyWPrkOWYT7U-STqd

[LauraGPT]

Thanks for the detailed experiment and for explicitly reporting the measurements outside the numerical tolerances. I reviewed the two-file diff at c478743. Two constructor-scope regressions need addressing before this is ready; inline comments below.

For a bounded check, I AST-extracted the exact two added constructor blocks from both files and executed them with real PyTorch 2.11.0+cu128 on an H100-visible host, using a CPU eval-mode Linear decoder stand-in. All four cases confirmed that the cuDNN SDPA process flag becomes false even with torch_compile=false; omitting the option also replaces the CPU eval decoder's forward with the lazy compile wrapper. The original process flag was restored. This is not full model initialization, compiled execution, B200 training, or independent verification of the speedup/numerics.

Please keep the B200 timing claims scoped to that experiment, and retain the reported loss/token-accuracy/gradient-norm tolerance violations as unmet gates rather than describing numerical equivalence as passing. A documented, explicit experimental training opt-in with regression tests for unchanged defaults would make this much easier to assess. Frozen-encoder train/eval policy should remain a separate recipe decision.

Review prepared with Codex assistance.

[LauraGPT]

[P2] Do not change the process-wide SDPA policy during model construction. enable_cudnn_sdp(False) persists beyond this instance and affects unrelated models in the same process; it still runs when llm_conf.torch_compile is false. The guard checks only whether any CUDA device exists, not the model's device, training mode, or the measured backend configuration. Please move this policy to an explicit training entry-point opt-in, or otherwise avoid changing the caller's global policy. Add a default/opt-out regression asserting that constructing this model leaves the pre-existing flag unchanged. The mirrored recipe block has the same issue.

[LauraGPT]

[P2] Preserve eager execution as the default for existing callers. This condition defaults torch_compile to true and only checks host CUDA availability, so even an eval-mode CPU decoder on a GPU host is wrapped; it also changes ordinary batched inference, not just the fine-tuning recipe measured here. Batch size one is the only eager bypass. The reported cold/warm first-step compile costs therefore become an implicit behavior change for existing configurations. Please make compilation explicitly opt-in for the experimental training path and test unchanged default/disabled behavior plus the intended enabled path in both model copies.

[LauraGPT]

I checked the new shared helper at8482dd2. With real torch2.11.0+cu128 and CPU decoder stand-ins, default configuration leaves both forward and SDPA flags untouched, explicit compile on CPU takes the eager path, and an explicitly selected SDPA context restores the flags after a sequential forward. The old unconditional-constructor/default-on findings therefore should not be applied unchanged to this revision. The recipe now explicitly enables the options.

One scope claim remains too strong: sdpa_kernel is not model-local/thread-isolated in this runtime. I held a configured math-only forward open in a worker using Events and observed SDPA flags from the other thread. cuDNN/flash/mem-efficient were disabled in both threads during that forward, then restored afterward. No concurrent ASR, model-weight inference, CUDA compilation, backward pass or B200 benchmark was run; this is a real flag-observation check of the exact helper, not the submitted full-constructor test suite.

Please describe this as temporary process-wide backend selection with restoration, not 'other models in the same process are not affected' or 'no process-wide torch flag is changed'. Document the concurrency limitation for explicitly enabled use. I have not reproduced the submitted training-speed/numerical results, and the new constructor/GPU tests are not independently accepted by this limited check.

[LauraGPT]

[P2] State the process-wide scope of this context. On torch2.11.0, an unrelated thread observes the selected backend flags while this forward is active; restoration on return does not make the selection model-local. Please remove the isolation guarantee and describe the concurrency limitation for callers enabling this option.