在当今AI领域,训练大规模模型如变压器和语言模型已成为不可或缺的一环,但这些模型动辄拥有数十亿个参数,对计算能力、内存和能源消耗提出了极高要求。例如,OpenAI的GPT-3模型拥有惊人的1750亿个参数,其训练过程需耗费数周时间,并依赖大量GPU资源。这种高昂的成本不仅限制了技术的普及,还加剧了能效和环境影响方面的担忧。因此,寻找高效且可持续的AI训练方法显得尤为重要。
传统的大规模模型训练低效主要源于对密集矩阵的依赖,这导致了显著的内存和计算需求。尽管已有矩阵分解和启发式秩缩减等方法试图缓解这一问题,但它们在现实应用中的效果并不理想。例如,GaLore虽然采用单批设置进行训练,但在运行时间上却存在不切实际的开销;而LTE则在大规模任务上的收敛上遇到了难题。显然,当前急需一种能够在不牺牲性能的前提下,同时降低内存使用、计算成本和训练时间的方法。
近日,来自纽约州立大学奥尔巴尼分校、加利福尼亚大学圣芭芭拉分校、亚马逊Alexa AI和Meta的研究人员共同提出了一种名为CoMERA(基于秩自适应张量优化的计算和内存高效的训练方法)的新框架。这一框架通过秩自适应张量压缩技术,将内存效率与计算速度完美结合,为AI训练带来了革命性的改变。
与传统方法仅专注于压缩不同,CoMERA采用多目标优化方法,平衡了压缩比和模型准确性。它利用张量化嵌入和高级张量网络收缩技术,优化了GPU的利用,显著减少了运行时间开销,同时保持了卓越的性能。此外,CoMERA还引入了CUDA Graph技术,以减少GPU操作中内核启动的延迟,从而突破了传统张量压缩方法的一个重要瓶颈。
CoMERA的核心在于自适应张量表示技术,它允许模型层根据资源约束动态调整其秩。通过调整张量秩,该框架在不损害神经网络操作完整性的前提下实现了压缩。这一动态优化过程分为两个阶段:早期阶段专注于稳定收敛,后期阶段则微调秩以满足特定的压缩目标。
实验结果显示,CoMERA在多个模型和任务上均表现出色。在一个六编码器的变压器模型中,CoMERA实现的压缩比例从早期阶段的43倍优化到后期阶段的361倍。与GaLore相比,它将内存消耗减少了9倍,每个周期的训练速度提高了2-3倍。在MNLI数据集上训练的变压器模型中,CoMERA将模型大小从256MB减少至仅3.2MB,同时保持了准确性。在DLRM等大规模推荐系统中,CoMERA也实现了高达99倍的模型压缩和7倍的峰值内存使用减少。此外,在CodeBERT等大型语言模型的预训练中,CoMERA同样表现出色,获得了4.23倍的总体压缩比,并在某些训练阶段实现了2倍的速度提升。
这项研究的主要结论包括:CoMERA在特定层实现了高达361倍的压缩比,在完整模型上达到了99倍,显著降低了存储和内存需求;该框架为变压器和推荐系统提供了每个周期2-3倍的更快训练时间,节省了计算资源和时间;使用张量化表示和CUDA Graph技术,CoMERA将峰值内存消耗减少了7倍,使得小型GPU训练成为可能;CoMERA的方法支持多种架构,包括变压器和大型语言模型,同时保持了或提高了准确性;通过降低训练的能量和资源需求,CoMERA促进了更可持续的AI实践,使得尖端模型能够被更广泛的受众所使用。
CoMERA框架通过实现更快、内存高效的训练,成功解决了AI可扩展性和可及性的重要障碍。其自适应的优化能力和与现代硬件的兼容性使其成为了一个引人注目的选择,尤其对于那些希望在不产生高昂成本的情况下训练大模型的组织而言。这项研究的结果为在分布式计算和资源受限的边缘设备等领域进一步探索基于张量的优化铺平了道路,为AI技术的未来发展注入了新的活力。