Wenet项目中新IO多数据训练卡死问题分析与解决

问题背景

在语音识别领域，Wenet作为一个端到端的语音识别工具包，近期引入了新的IO数据加载机制以支持更大规模的数据集训练。然而，在实际使用过程中，部分开发者在多GPU环境下使用新IO加载多个数据集（如AISHELL2、Librispeech和CSTAL组合）时遇到了训练卡死的问题。

问题现象

具体表现为：

1. 在8卡GPU环境中训练时，其中一张GPU的利用率降为0
2. 其余GPU保持100%利用率
3. 训练进程看似正常但实际上已经停滞
4. 仅使用单个数据集（如AISHELL）时不会出现此问题

根本原因分析

经过深入排查，发现问题主要由以下几个因素导致：

1. 数据压缩格式选择不当：开发者将shard格式数据打包为bz2格式，这种压缩格式虽然节省存储空间，但解压速度较慢，容易成为数据加载的瓶颈。

2. 音频解码方式不兼容：在自定义修改的代码中，音频解码部分没有完全适配新的IO机制。原始代码使用直接加载方式：

   with stream.extractfile(tar_info) as file_obj:
       waveform, sample_rate = torchaudio.load(file_obj)
   

   而新IO机制需要改为：

   with stream.extractfile(tar_info) as file_obj:
       waveform, sample_rate = torchaudio.load(io.BytesIO(file_obj.read()))
   

3. 数据加载不均衡：由于上述两个原因，导致不同GPU上的数据加载速度不一致，进而引发梯度同步问题，最终造成训练卡死。

解决方案

针对上述问题，建议采取以下措施：

1. 使用推荐的压缩格式：避免使用bz2这种高压缩比但解压慢的格式，可以考虑使用更高效的压缩方式或直接使用未压缩数据。

2. 严格遵循新IO接口规范：确保音频解码部分完全按照新IO机制的要求实现，正确处理字节流数据。

3. 保持代码同步：建议基于最新的master分支进行开发，避免因版本差异导致的不兼容问题。

4. 完善错误处理：在数据加载部分增加适当的异常处理和日志记录，便于快速定位问题。

技术细节补充

在Wenet的新IO机制中，decode_wav函数设计为同时支持两种输入类型：

• 字节流数据（从tar文件中提取）
• 文件路径

正确的实现应该将波形数据和采样率的赋值放在条件判断之外，确保无论哪种输入路径都能正确设置这些值。当前master分支中存在一个小bug需要修复，即将：

sample['wav'] = waveform
sample['sample_rate'] = sample_rate

移到if-else条件判断的外部。

最佳实践建议

1. 在多数据集训练时，建议先进行小规模测试，确保数据加载正常后再进行完整训练。

2. 监控各GPU的利用率，发现不均衡情况及时检查数据加载部分。

3. 对于自定义修改，建议与官方代码保持同步更新，定期合并最新改动。

4. 关注Wenet官方文档更新，特别是关于新IO机制的详细说明。

通过以上分析和解决方案，开发者可以避免在多数据集训练时出现卡死问题，充分利用新IO机制带来的性能优势。