Jetfire：采用INT8数据流与逐块量化的高效且准确的Transformer预训练
机器学习
2024-07-23 v2
摘要
Transformer的预训练通常非常耗时。全量化训练（FQT）是一种有望加速预训练的方法。然而，大多数FQT方法采用量化-计算-反量化流程，由于高内存访问开销和低精度计算，在用于Transformer时往往导致次优的加速效果和显著的性能下降。在这项工作中，我们提出了Jetfire，一种专为Transformer设计的高效且准确的INT8训练方法。我们的方法具有INT8数据流以优化内存访问，并采用逐块量化方法以保持预训练Transformer的准确性。大量实验表明，我们的INT8 FQT方法达到了与FP16训练基线相当的准确度，并优于现有的Transformer INT8训练工作。此外，对于一个标准的Transformer模块，与FP16基线相比，我们的方法提供了1.42倍的端到端训练加速和1.49倍的内存减少。我们的代码已在https://github.com/thu-ml/Jetfire-INT8Training开源。
引用
@article{arxiv.2403.12422,
title = {Jetfire: Efficient and Accurate Transformer Pretraining with INT8 Data Flow and Per-Block Quantization},
author = {Haocheng Xi and Yuxiang Chen and Kang Zhao and Kai Jun Teh and Jianfei Chen and Jun Zhu},
journal= {arXiv preprint arXiv:2403.12422},
year = {2024}
}
备注
15 pages, 8 figures, 11 tables