Transformers.js项目中的WebGPU内存泄漏问题分析与解决方案

问题背景

在Transformers.js项目中使用Whisper模型进行语音转文字任务时，开发人员发现了一个严重的内存泄漏问题。当处理较长的音频文件（如10分钟以上的录音）时，GPU内存使用量会持续增长，最终导致设备内存耗尽或GPU设备丢失错误。这个问题在WebGPU后端尤为明显，特别是在NVIDIA GTX 1080（8GB）等显存较小的GPU上表现更为突出。

问题现象

多位开发者在不同环境下重现了这一问题：

1. 使用Whisper模型进行长时间音频转录时，GPU内存持续增加
2. 即使转录任务完成，内存也不会被释放
3. 最终可能导致两种结果：
   • 显存较小的GPU直接耗尽内存
   • 计算过程中GPU设备丢失，产生"device is lost"错误

技术分析

经过深入调查，发现问题根源在于Key-Value（KV）缓存的内存管理机制上。Transformers.js为了提升性能，将KV缓存保留在GPU缓冲区中，这使得ONNX Runtime不再拥有相关张量的所有权，而需要开发者手动释放这些缓冲区。

在Whisper模型的编码器-解码器架构中：

1. 编码器PKV（Past Key Values）在生成过程中被重复使用
2. 解码器PKV在每次生成步骤后都会被释放
3. 但编码器PKV在最后一个token生成后没有被正确释放

解决方案

项目维护者提出了以下修复方案：

1. 在生成过程结束后释放编码器PKV
2. 禁用编码器输出到GPU的选项（因为这是内存泄漏的主要来源）
3. 遵循ONNX Runtime WebGPU的最佳实践，明确管理GPU缓冲区生命周期

正确的KV缓存管理流程应该是：

1. 首次运行：输入位于CPU，提供初始数据
2. 第二次运行：将输入（past_key_values）赋值为第一次运行的输出（此时输入位于GPU）
3. 后续运行：释放前一次运行的输入，并将当前运行输入赋值为前一次运行的输出
4. 最终运行：释放前一次输入和当前运行的输出

开发者建议

对于遇到类似问题的开发者，可以采取以下临时解决方案：

1. 对于浏览器扩展等场景，可以在后台worker中运行模型，定期重启worker
2. 合理分割长音频，使用语音活动检测(VAD)进行批处理
3. 监控GPU内存使用情况，在达到阈值时主动释放资源

未来优化方向

ONNX Runtime WebGPU实现还有进一步优化的空间：

1. 改进GPU缓冲区重用机制
2. 提供更明确的缓冲区重用标识机制
3. 优化编码器-解码器模型的静态和动态图支持

结论

内存管理是WebGPU应用开发中的关键挑战，特别是在处理大型模型如Whisper时。通过正确理解和管理KV缓存的生命周期，开发者可以有效避免内存泄漏问题。Transformers.js团队正在积极解决这一问题，未来版本将提供更稳定和高效的内存管理机制。