[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

When I fine-tuned Fun-ASR-Nano with finetune.sh on one B200, a training step took 0.244 s on average, and two causes accounted for most of that time. First, torch sends the Qwen3 attention to cuDNN, which builds a new execution plan whenever the batch shape changes; that happened on 59 of 382 steps and cost about 1 s each. Second, on the other steps the GPU was mostly idle, waiting for the CPU to launch 8,080 small kernels that together ran for only 37 ms. This PR adds two opt-in switches under llm_conf, both turned on in finetune.sh: sdpa_backends keeps the attention off cuDNN, and torch_compile compiles the decoder so it launches far fewer kernels. On that B200 the mean step drops to 0.068 s (3.6x); I have not measured other GPUs. The two switches:

1. llm_conf.sdpa_backends. On the B200 (sm_100) I measured, torch sends the Qwen3 attention to cuDNN because it carries a padding mask, and cuDNN rebuilds its execution plan for every new batch shape, about 1 s each, on one step in six. Setting it to [flash, efficient, math] runs the decoder forward under sdpa_kernel with those backends, which need no per-shape plans. The flags are process-wide while the forward runs and are restored on return, so leave it unset when several models share one process.
2. llm_conf.torch_compile. Once the plan builds are gone, the step is launch-bound: 8,080 kernels for 37 ms of GPU work, 2,900 of them in the decoder. Setting it wraps the decoder stack in torch.compile(dynamic=True), which halves the launch count. Inputs on the CPU and single-sequence batches take the eager path, and the first step pays a one-time compile.

Speedup Result

Measured on one B200.

	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
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

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: finetune.sh as shipped (on this branch, with its two new ++llm_conf.* lines), 3600 AISHELL-1 utterances = 382 steps, seed 1234, two interleaved runs per arm, nothing set in the environment; full command under Details.

• python -m pytest tests/test_fun_asr_nano_train_opts.py: 4 cases x 2 model copies: 8 passed with a CUDA device, 6 passed / 2 skipped (the CUDA-only case) with CUDA_VISIBLE_DEVICES=""
• 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

I trained the baseline and this branch on the same 382 batches and compared 6,128 values from each run: per-step loss, logit statistics and correct-token counts, parameter and gradient norms at seven steps, the final validation loss, and every batch shape. The tolerances in the table below were set before the runs. The recording code is on the perf/fun-asr-nano-training-speed-verify branch of my fork, not in this PR.

23 of the 6,128 checks fail the tolerances set before the runs, so numerical equivalence has not been shown. The 23 are two per-step losses at 0.0128 and 0.0115 against a 0.01 limit, two gradient norms at 2 % against 1 %, and a one-token change in the correct-token count on 11 steps. The other 6,105 are within tolerance, and the final validation loss is 0.101834 on the baseline and 0.101967 on this branch with identical accuracy. My guess is bf16 rounding, since the flash and memory-efficient kernels and the compiled decoder round differently from eager cuDNN, but I have not verified that, and it does not make the failures acceptable by itself.

recorded	baseline vs this branch	tolerance
validation loss and accuracy, 200 dev utterances	loss 0.101834 vs 0.101967; accuracy identical	0.001
trainable-parameter norm, 7 steps	3.4e-9 relative	1e-4
logits mean and std, every step	0.16 % mean, 0.67 % max	1 %
per-step loss, 382 steps	0.0014 mean, 0.0128 max; 2 steps above 0.01	0.01
correct-token count, 382 steps	3 of 34,886 tokens net; 11 steps above 1 %	1 % per step
gradient norm before clipping, 7 steps	0.8 % or less on 5 steps, 2.0 % on 2	1 %
batch shapes and frame counts	identical	exact

User impact

With the defaults the model is built exactly as before, and the new tests check that for both model copies. The two keys are read by both model constructors, so any caller that sets them gets the switches; finetune.sh is the recipe that sets them. The speedup above was measured on one B200 only. The cost is a one-time compile on the first step, about a minute with a warm Inductor cache and 2.5 minutes cold on that B200. Removing either key turns that switch off.

Notes for reviewers

• The review comments of 2026-09-13 and 2026-09-15 are addressed: both switches are opt-in and off by default, the attention setting is described as process-wide while the forward runs, and the tests cover the defaults and the enabled path for both model copies.
• Changing batch shapes caused no recompiles in 382 steps; single-sequence batches run eagerly.
• On that B200 the compile pays for itself after about 2,000 steps with a warm cache and 5,000 cold, well within the recipe's 50 epochs.
• TORCH_CUDNN_SDPA_DEPRIORITIZED=1 gives the same result as sdpa_backends without editing the recipe. The profile 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.
• Not tested: DeepSpeed (use_deepspeed=true) and more than one GPU.
• Left out of this PR: the frozen encoder runs with dropout on every step, and AdamW keeps the LLM's optimizer state in bf16.

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.

[LauraGPT]

Checked revision 2c1a0c8: the helper docstring, both model comments, training guide and PR description now accurately state that the SDPA flags are temporarily process-wide and warn against concurrent shared-process use. This addresses the wording/concurrency-boundary request in this thread.

I also reran the bounded exact-helper check with real PyTorch 2.11.0+cu128, CUDA hidden and CPU decoder stand-ins: defaults leave forward and flags untouched; explicit compile with CPU input uses eager forward; an Event-held math-only forward changes the flags observed by another thread, then restores them on return. All assertions passed. The new revision also simplifies the backend parser and reduces the submitted tests, so this is not full-constructor, combined-option, CUDA compilation/backward or B200 numerical/performance acceptance. I have not independently run the submitted constructor/GPU suite or verified the reported training tolerances. No model weights were loaded.

[LauraGPT]

The latest body edit at unchanged head 2c1a0c8 introduces two broader claims than the evidence supports:

• Please keep the 3.57x result explicitly limited to the measured single-B200 setup. The new User impact statement that training on sm_90 and sm_100 gets the speedup shown above is not supported by the one-B200 measurements; hardware/backend behavior does not establish the same speedup on H100 or other systems.
• These options are read by both model constructors, so they can affect any caller explicitly setting the keys, not only people running finetune.sh. Defaults are unchanged; finetune.sh is the recipe that enables them automatically.

The table still shows the numerical threshold exceedances, which is useful. Please also retain an explicit statement that 23/6,128 checks fail the predeclared tolerances and numerical equivalence has not passed. Attributing the discrepancies to bf16 rounding does not turn those failures into accepted correctness gates without a justified acceptance decision. My previous reply confirms only the corrected SDPA-scope wording and bounded CPU-helper behavior, not independent acceptance of the full constructor/GPU tests or training numerics/performance.

This comment reviews the body change only; the code head is unchanged, so I did not repeat the previous tests.

[TarzanZhao]

@LauraGPT Hi, thanks for the careful comments. I think I have fixed these problems and the pr description. I just make this PR ready to review.

[LauraGPT]

@TarzanZhao I checked the revised description and the transition to ready for review at head 2c1a0c82a3be01f4bba115e7d74d95702ce6a22d.

The three description points in my previous comment are now addressed: the measured speedup is limited to the single B200, the opt-in keys are correctly described as available to callers of both constructors, and the 23/6,128 failures against the predeclared tolerances are explicitly retained, with the rounding explanation marked as unverified.

The source head has not changed since the previous scoped CPU-helper check, so I have not rerun that unchanged probe. That check is not independent validation of the full constructor suite, CUDA compilation/backward, B200 performance, or numerical equivalence. In particular, the reported numerical gate remains unmet; making the PR ready for review does not itself resolve that acceptance question. This acknowledges the corrected description, not merge approval.