基于MindSpore的Whisper模型推理优化实践

5月 8, 2025 Large Language Model

在使用 MindSpore NLP 的 Whisper 模型进行推理时，模力方舟团队发现推理性能存在瓶颈。例如，对于一段 91 秒的音频，模型的识别时间达到 95 秒，推理开销较大。当前环境为 mindspore==2.5.0 与 mindnlp==0.4.0，且模型仅支持 Eager 模式，无法充分发挥硬件性能。

为解决此问题，我们系统分析了注意力机制的执行方式，并尝试引入高性能的 FlashAttention2 实现。同时，我们也借助 MindSpore Profiler 定位性能瓶颈，最终通过引入原生的 Conv1D 算子，进一步提升整体推理性能。

一、三种注意力机制对比

模式	特点	性能表现
Eager	默认执行方式，算子按顺序执行，调试友好	无融合优化，推理慢
SDPA（Scaled Dot-Product Attention）	经典注意力实现	适用于普通推理需求
FlashAttention2	高性能实现，显著减少显存和加速计算	长序列推理性能大幅提升

二、接入 FlashAttention2 模式

为提升推理速度，我们在原始 eager 模式基础上引入了 FlashAttention2 。改动包括：

1. 适配 `flash-attn` 库

适配 flash-attn 库中 bert-padding.py 中的关键方法，包括：

index_put_first_axis
index_first_axis
unpad_input
pad_input

2. 新增支撑模块 `modeling_flash_attention_utils.py`

新增辅助函数实现：

_get_unpad_data
_unpad_input
_prepare_fa2_from_position_ids
_fa_peft_integration_check
_flash_attention_forward

3. 修改 Whisper 模型支持 FlashAttention2

在 modeling_whisper.py 中加入 WhisperFlashAttention2 逻辑，允许初始化模型时通过：

1attn_implementation="flash_attention_2"

来启用新模式。

4. 性能初步评估

推理耗时从 95s 降至约 85s，性能提升约 10% 性能。但通过 Profiler 分析，发现瓶颈仍存在。

三、进一步优化：替换 Conv1D 实现

1. Profiler 分析结果

我们使用如下脚本采集性能数据：

 1import mindspore
 2from mindnlp.transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 3from mindspore import Profiler
 4from mindspore.profiler import ProfilerLevel, schedule, tensorboard_trace_handler
 5
 6model_id = ""openai/whisper-large-v3""
 7model = AutoModelForSpeechSeq2Seq.from_pretrained(
 8    model_id, 
 9    ms_dtype=mindspore.float16, 
10    low_cpu_mem_usage=True,
11    use_safetensors=True,
12    attn_implementation="flash_attention_2",
13)
14processor = AutoProcessor.from_pretrained(model_id)
15pipe = pipeline(
16    "automatic-speech-recognition",
17    model=model,
18    tokenizer=processor.tokenizer,
19    feature_extractor=processor.feature_extractor,
20    ms_dtype=mindspore.float16,
21    return_timestamps=True,
22)
23# Profiler 数据默认存储在路径：
24# ./data/modelfoundry-prod-node-xxx/ASCEND_PROFILER_OUTPUT
25with Profiler(
26        profiler_level=ProfilerLevel.Level0,
27        schedule=schedule(wait=0, warmup=0, active=1, repeat=1, skip_first=0),
28        on_trace_ready=tensorboard_trace_handler
29     ) as prof:
30	pipe("/path/to/yourself.mp3")
31	prof.step()

通过使用 MindStudio Insight 对性能数据进行分析后发现，模型中 Conv1D 实现依赖 Conv2D 间接方式，且运行在 CPU 侧，成为主要的性能瓶颈。

2. 原因定位

旧版本的 Conv1D 是通过 Conv2D 间接构造实现，存在：

多余维度转换
CPU 执行
内存频繁拷贝

3. 解决方案

自 MindSpore 2.6.0rc1 起，框架已提供高效的原生 Conv1D 实现，支持图模式和硬件加速。在此基础上，我们重构模型中的 Conv1D 调用，显著缩短推理时间。

4. 最终效果

结合 FlashAttention2 和 Conv1D 及其它算子优化后，最终在 MindSpore 2.6.0rc1 版本下，91 秒音频推理时间缩短至约 57 秒，整体性能提升 40%，CPU 占用率显著下降。

四、手把手推理教程

1. 下载镜像

执行以下Shell命令，拉取 MindSpore 容器镜像：

1docker pull quay.io/ascend/mindspore:openeuler-python3.10-cann8.1.rc1-mindspore2.6.0rc1
2# 推荐国内源加速
3# docker pull quay.xzt.me/ascend/mindspore:openeuler-python3.10-cann8.1.rc1-mindspore2.6.0rc1

2. 创建并进入容器

执行以下命令创建容器，name 设置为 whisper：

 1docker run -itd --privileged  --name=whisper --net=host \
 2   --shm-size 500g \
 3   --device=/dev/davinci0 \
 4   --device=/dev/davinci1 \
 5   --device=/dev/davinci2 \
 6   --device=/dev/davinci3 \
 7   --device=/dev/davinci4 \
 8   --device=/dev/davinci5 \
 9   --device=/dev/davinci6 \
10   --device=/dev/davinci7 \
11   --device=/dev/davinci_manager \
12   --device=/dev/hisi_hdc \
13   --device /dev/devmm_svm \
14   -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
15   -v /usr/local/Ascend/firmware:/usr/local/Ascend/firmware \
16   -v /usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi \
17   -v /usr/local/sbin:/usr/local/sbin \
18   -v /etc/hccn.conf:/etc/hccn.conf \
19   quay.io/ascend/mindspore:openeuler-python3.10-cann8.1.rc1-mindspore2.6.0rc1 \
20   bash

进入容器，后续所有操作均在容器内操作

1docker exec -it whisper bash

3. 安装 MindSpore NLP 与依赖

执行以下脚本，安装 MindSpore NLP 及相关依赖包：

 1# 安装相关依赖
 2yum install ffmpeg git
 3# 配置国内源
 4pip config set global.index-url https://repo.huaweicloud.com/repository/pypi/simple/ 
 5# 升级 pip
 6pip install --upgrade pip
 7# 安装 mindnlp
 8git clone -b 0.4 https://github.com/mindspore-lab/mindnlp.git
 9cd mindnlp
10bash scripts/build_and_reinstall.sh

4. 推理代码示例

国内可以配置 hf 镜像源拉取模型：

1export HF_ENDPOINT=https://hf-mirror.com

执行以下代码进行推理：

 1import mindspore
 2from mindnlp.transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 3
 4model_id = "openai/whisper-large-v3"
 5model = AutoModelForSpeechSeq2Seq.from_pretrained(
 6    model_id, 
 7    ms_dtype=mindspore.float16, 
 8    low_cpu_mem_usage=True,
 9    use_safetensors=True,
10    attn_implementation="flash_attention_2",
11)
12processor = AutoProcessor.from_pretrained(model_id)
13pipe = pipeline(
14    "automatic-speech-recognition",
15    model=model,
16    tokenizer=processor.tokenizer,
17    feature_extractor=processor.feature_extractor,
18    ms_dtype=mindspore.float16,
19    return_timestamps=True,
20)
21
22result = pipe("/path/to/yourself.mp3")

五、相关材料

MindSpore

基于MindSpore的Whisper模型推理优化实践

一、三种注意力机制对比

二、接入 FlashAttention2 模式

1. 适配 `flash-attn` 库

2. 新增支撑模块 `modeling_flash_attention_utils.py`

3. 修改 Whisper 模型支持 FlashAttention2

4. 性能初步评估

三、进一步优化：替换 Conv1D 实现

1. Profiler 分析结果

2. 原因定位

3. 解决方案

4. 最终效果

四、手把手推理教程

1. 下载镜像

2. 创建并进入容器

3. 安装 MindSpore NLP 与依赖

4. 推理代码示例

五、相关材料

是小柒鸭

佛系·猫奴·程序媛

相关

基于MindSpore的Whisper模型推理优化实践

一、三种注意力机制对比

二、接入 FlashAttention2 模式

1. 适配 flash-attn 库

2. 新增支撑模块 modeling_flash_attention_utils.py

3. 修改 Whisper 模型支持 FlashAttention2

4. 性能初步评估

三、进一步优化：替换 Conv1D 实现

1. Profiler 分析结果

2. 原因定位

3. 解决方案

4. 最终效果

四、手把手推理教程

1. 下载镜像

2. 创建并进入容器

3. 安装 MindSpore NLP 与依赖

4. 推理代码示例

五、相关材料

是小柒鸭

佛系·猫奴·程序媛

相关

1. 适配 `flash-attn` 库

2. 新增支撑模块 `modeling_flash_attention_utils.py`