数据搬运：将文档移入内存对神经信息检索模型效率的研究

在使用大型语言模型训练神经排序器时，实践者通常会利用多个 GPU 来加速训练时间。通过使用更多设备，PyTorch 等深度学习框架允许用户大幅增加可用显存池，从而在训练时使用更大的批次，缩短训练时间。与此同时，在运行数据密集型模型时，一个通常被忽视的关键过程是数据在磁盘、主内存和显存之间的管理方式。大多数开源研究实现忽略了这一内存层次结构，而是采用将所有文档从磁盘加载到主内存，然后让框架（如 PyTorch）处理将数据移入显存的方式。因此，随着信息检索研究专用数据集规模的不断增长，一个自然的问题随之而来：这是优化训练时间的最佳方案吗？我们在此研究了三种流行的信息检索数据集文档处理方式的行为及其在多 GPU 下的扩展性。具体而言，将文档直接加载到内存、通过查找表从文本文件直接读取文档以及使用专门处理信息检索数据集的库（ir_datasets）在性能（即每秒处理样本数）和内存占用方面存在差异。我们表明，在使用最流行的神经排序器研究库（即 PyTorch 和 Hugging Face 的 Transformers）时，将所有文档加载到主内存的做法并非总是最快的选择，且对于超过几个 GPU 的配置并不可行。同时，良好的磁盘数据流式传输实现可以更快，且具有显著更好的可扩展性。我们还展示了诸如内存钉固、多工作进程和 RAMDISK 使用等提升加载速度的流行技术如何以较小的内存开销进一步减少训练时间。