[논문 리뷰] AutoFreeze: Automatically Freezing Model Blocks to Accelerate Fine-tuning

Yuhan Liu, Saurabh Agarwal외 1명|arXiv (Cornell University)|2021. 02. 02.

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는 캐싱 기능을 활용해 표준 미세조정 대비 10개의 NLP 작업에서 최대 2.55배의 속도 향상을 달성합니다.
• 64-GPU 클러스터에서 AutoFreeze는 종단 간 훈련 시간을 최대 4.38배 단축시키며 총 비용을 최대 5.03배 절감합니다. 이는 정확도 저하 없이 이루어졌습니다.
• 기울기 노름 테스트는 최소한의 오버헤드를 유발합니다—구간 시간의 1% 미만이며, 에포크당 약 453MB 메모리 소비에 그치므로 실용적인 실세계 적용이 가능합니다.
• 수렴이 가장 느린 레이어만 동결하는 방식은 정확도를 유지하지만, 정적 방법은 초기 레이어를 동결함으로써 MRPC 데이터셋에서 정확도를 87%에서 76.5%로 떨어뜨립니다.
• 이 방법은 훈련 시간을 크게 줄이면서도 최신 기술 수준의 정확도를 유지하며, 정확도 저하 문제가 발생하는 이전의 방법들(예: EarlyBERT)을 뛰어넘습니다.
• 최적의 동결 주기는 에포크당 2~5회의 평가 간격 사이에 위치하며, 이는 속도 향상과 모델 안정성 간의 균형을 잘 맞춥니다.