[论文解读] AutoFreeze: Automatically Freezing Model Blocks to Accelerate Fine-tuning

Yuhan Liu, Saurabh Agarwal+1 more|arXiv (Cornell University)|Feb 2, 2021

Topic Modeling参考文献 72被引用 16

一句话总结

AutoFreeze 通过自适应地冻结早期收敛的层来加速 BERT 微调，采用基于梯度范数的准则识别稳定层。结合激活缓存和分布式执行，它在单张 GPU 上实现了最高 2.55 倍的加速，在 64 张 GPU 集群上实现了最高 5.03 倍的成本降低，且未造成模型准确率损失，达到最先进水平。

ABSTRACT

With the rapid adoption of machine learning (ML), a number of domains now use the approach of fine tuning models which were pre-trained on a large corpus of data. However, our experiments show that even fine-tuning on models like BERT can take many hours even when using modern accelerators like GPUs. While prior work proposes limiting the number of layers that are fine-tuned, e.g., freezing all layers but the last layer, we find that such static approaches lead to reduced accuracy. We propose, AutoFreeze, a system that uses an adaptive approach to choose which layers are trained and show how this can accelerate model fine-tuning while preserving accuracy. We also develop mechanisms to enable efficient caching of intermediate activations which can reduce the forward computation time when performing fine-tuning. We extend AutoFreeze to perform distributed fine-tuning and design two execution modes that minimize cost and running time respectively. Our evaluation on ten NLP tasks shows that AutoFreeze, with caching enabled, can improve fine-tuning on a single GPU by up to 2.55x. On a 64 GPU cluster, for fine-tuning on the AG's news dataset, AutoFreeze is able to achieve up to 4.38x speedup when optimizing for end-to-end training time and 5.03x reduction in total cost when optimizing for efficiency, without affecting model accuracy.

研究动机与目标

• 为解决微调大型预训练模型（如 BERT）带来的高计算成本问题，其在 4 张 GPU 集群上可能耗时长达 27 小时。
• 克服静态微调策略（如固定冻结层或使用彩票剪枝）的局限性，这些策略常导致模型准确率下降。
• 开发一种自适应的、在运行时确定的微调过程中冻结模型块的方法，以在保持准确率的同时加速训练。
• 集成高效的激活缓存和分布式执行模式，进一步降低多 GPU 环境下的训练时间和成本。

提出的方法

• AutoFreeze 使用基于梯度范数的测试，按训练过程中各层变化速率对层进行排序，识别出早期收敛的层并予以冻结。
• 选择梯度范数衰减最慢的层进行冻结，避免反向传播计算，从而减少计算和通信开销。
• 系统采用动态激活缓存机制，避免在前向传播中重复计算中间输出，从而减少前向计算时间。
• AutoFreeze 支持两种分布式执行模式：一种针对端到端训练时间进行优化，另一种针对最小化总成本，实现集群中的高效扩展。
• 该方法已集成至 PyTorch，支持单 GPU 和多 GPU 微调，且运行时开销极低。
• 冻结决策在训练过程中定期做出，评估间隔时间经过调优，以在加速效果与准确率之间取得平衡。

实验结果

研究问题

• RQ1 在微调过程中自适应地冻结层是否能显著减少训练时间而不降低模型准确率？
• RQ2 与静态冻结或彩票剪枝相比，基于梯度范数的准则在加速效果和准确率保持方面表现如何？
• RQ3 在多 GPU 环境下，激活缓存和分布式执行能在多大程度上进一步加速微调？
• RQ4 评估层收敛的最优频率是多少，才能在加速效果与模型性能之间实现最佳平衡？

主要发现

• 在单张 GPU 上，使用缓存的 AutoFreeze 在十个 NLP 任务上相比标准微调实现了最高 2.55 倍加速。
• 在 64 张 GPU 集群上，AutoFreeze 将端到端训练时间减少最多 4.38 倍，总成本减少最多 5.03 倍，且无准确率下降。
• 梯度范数测试引入的开销极低——每轮次不足间隔时间的 1%，每轮次内存占用不足 453MB，适用于实际部署。
• 仅冻结变化最慢的层可保持模型准确率，而静态方法冻结早期层则导致 MRPC 数据集上的准确率从 87% 降至 76.5%。
• 该方法在显著减少训练时间的同时保持了最先进水平的准确率，优于如 EarlyBERT 等先前方法，后者存在准确率下降问题。
• 最优冻结频率介于每轮次 2 到 5 个评估间隔之间，可在加速效果与模型稳定性之间实现良好平衡。