利用 Flash Attention 进行打包以提升训练效率
机器学习
2024-09-04 v6 人工智能
摘要
在调优 LLM 模型时，常会通过添加特殊标记到较短的训练示例中，以匹配每个 batch 中最长序列的长度。虽然这确保了 batch 处理的一致性，但会因包含无关的填充标记而引入效率问题，浪费 GPU 资源。Hugging Face SFT trainer 一直提供使用打包将多个训练示例组合在一起的选项，以实现最大化的 GPU 资源利用率。然而，迄今为止，该选项没有为每个打包的训练示例提供proper的掩码处理。Hugging Face Transformers 4.44 版本已添加了此功能。我们分析了这一新特性，展示了不同打包变体的优势。
关键词
引用
@article{arxiv.2407.09105,
title = {Enhancing Training Efficiency Using Packing with Flash Attention},
author = {Achintya Kundu and Rhui Dih Lee and Laura Wynter and Raghu Kiran Ganti and Mayank Mishra},
journal= {arXiv preprint arXiv:2407.09105},
year = {2024}
}