基于Nvidia A100的Whisper多请求推理瓶颈与优化方案问询

初始场景与问题 #

• 目标：基于Whisper-Large-v3构建实时语音转文本应用，支持同时处理多个请求
• 硬件：配备80GB VRAM的Nvidia A100显卡
• 现状：
  • 单个Whisper实例因KV矩阵依赖，请求只能串行执行，无法并行处理
  • 使用Gunicorn启动多worker，加载2个独立Whisper实例时可并行处理2个请求；但加载4个实例时，仍仅能并行处理2个，剩余请求需等待队列空闲
• 诉求：明确性能瓶颈所在，提供能更高效利用显卡推理能力的优化方案（当前VRAM尚有剩余）

当前实现细节 #

• 技术栈：Transformers库 + PyTorch
• 注意力机制：Flash Attention 2或SDPA（短音频场景下效果相近）
• 计算后端：PyTorch

个人猜测 #

瓶颈可能源于PyTorch后端为输出和注意力矩阵分配连续内存块，导致无法利用空闲内存并行处理更多请求；尽管VRAM充足，但矩阵可能被加载到容量有限的SRAM中，限制了并行能力。

更新（2024.06.04） #

批量处理音频可大幅提升推理速度，但仍需适配实时应用场景，希望了解除批量处理外，是否可通过流水线方式实现显卡并行计算。

测试代码 #

1. 构建Whisper实例

import torch
from transformers import pipeline
import os
import time

pipe = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-large-v3",
    torch_dtype=torch.bfloat16,  
    device="cuda:0",  
)

2. 音频生成器与主处理逻辑

# 遍历本地音频文件夹的生成器
def data_generator():
    folder_path =  r'C:\Users\audios'
    for filename in os.listdir(folder_path):
        audio = os.path.join(folder_path, filename)
        yield audio

def main():
    for out in pipe(inputs=data_generator(), batch_size=8, return_timestamps=False, 
        generate_kwargs={"language": "en", "task": "transcribe"}):
        print(out)

后续疑问 #

当前批量处理逻辑需要凑够8个音频才启动转录，但实时场景中请求是串行到来的，不想等待凑齐批量再处理，询问是否存在比顺序调用pipe更优的实时处理方法。

内容的提问来源于stack exchange，提问作者leon