训练大型AI模型（如Transformer和语言模型）已成为AI领域不可或缺的关键环节，但也面临着高昂的计算成本、内存消耗和能源需求。例如，OpenAI的GPT-3拥有1750亿个参数，需要数周的GPU训练。这种巨大的资源需求限制了这项技术在大规模计算资源充足的组织中的应用，同时也加剧了人们对能源效率和环境影响的担忧。解决这些挑战对于确保AI发展的更广泛可及性和可持续性至关重要。
传统训练方法效率低下，亟需创新解决方案
大型模型训练效率低下的主要原因在于其对稠密矩阵的依赖，这需要大量的内存和计算能力。现代GPU对优化的低精度或低秩操作支持有限，进一步加剧了这些需求。尽管已经提出了一些方法，如矩阵分解和启发式秩降低，来缓解这些问题，但它们在实际应用中仍然受到限制。例如，GaLore能够在单批次设置下进行训练，但存在不切实际的运行时开销。同样，采用低秩适配器的LTE在大型任务上的收敛性方面也存在问题。目前缺乏一种能够同时降低内存使用、计算成本和训练时间，而不损害性能的方法，这使得创新解决方案的需求变得迫切。
CoMERA框架:通过自适应张量优化实现高效训练
来自奥尔巴尼大学（纽约州立大学）、加州大学圣巴巴拉分校、亚马逊Alexa AI和Meta的研究人员共同推出了一种名为CoMERA(Computing-and Memory-Efficient training method via Rank-Adaptive tensor optimization)的新型框架。该框架结合了内存效率和计算速度，通过自适应秩张量压缩技术来实现。与传统方法仅关注压缩不同，CoMERA采用多目标优化方法来平衡压缩比和模型精度。它利用张量化嵌入和先进的张量网络收缩来优化GPU利用率，从而减少运行时开销，同时保持强大的性能。该框架还引入了CUDA图，以最大程度地减少GPU操作期间的内核启动延迟，这是传统张量压缩方法中的一个主要瓶颈。
CoMERA的基础是自适应张量表示，它允许模型层根据资源约束动态调整其秩。通过修改张量秩，该框架可以在不损害神经网络操作完整性的情况下实现压缩。这种动态优化是通过一个两阶段的训练过程实现的:
早期阶段:专注于稳定收敛。
后期阶段:微调秩以满足特定的压缩目标。
在一个六编码器Transformer模型中，CoMERA在其早期阶段实现了高达43倍的压缩比，而在其后期优化阶段，压缩比更是高达361倍。此外，与GaLore相比，它将内存消耗降低了9倍，每轮训练速度提高了2-3倍。
多项测试结果表明CoMERA性能卓越
在应用于MNLI数据集上训练的Transformer模型时，CoMERA将模型大小从256MB缩小到低至3.2MB，同时保持了精度。在诸如DLRM的大规模推荐系统中，CoMERA将模型压缩了99倍，并使峰值内存使用量减少了7倍。该框架还在预训练CodeBERT（一个特定领域的大型语言模型）方面表现出色，获得了4.23倍的整体压缩比，并在某些训练阶段实现了2倍的加速。这些结果突显了其处理各种任务和架构的能力，扩展了其在各个领域的适用性。
CoMERA框架的关键优势总结
这项研究的主要结论如下:
CoMERA为特定层实现了高达361倍的压缩比，为整个模型实现了99倍的压缩比，大大降低了存储和内存需求。
该框架将Transformer和推荐系统的每轮训练时间缩短了2-3倍，节省了计算资源和时间。
通过使用张量化表示和CUDA图，CoMERA将峰值内存消耗减少了7倍，使得在较小的GPU上进行训练成为可能。
CoMERA的方法支持包括Transformer和大型语言模型在内的多种架构，同时保持或提高了精度。
通过降低训练所需的能源和资源，CoMERA有助于实现更可持续的AI实践，并使更广泛的受众能够使用尖端模型。
