[논문 리뷰] Reducing Activation Recomputation in Large Transformer Models

Vijay Anand Korthikanti, Jared Casper외 5명|arXiv (Cornell University)|2022. 05. 10.

Parallel Computing and Optimization Techniques인용 수 52

한 줄 요약

본 논문은 시퀀스 병렬성(sequence parallelism)과 선택적 활성 재계산(selective activation recomputation)을 도입해 활성화 메모리와 재계산 오버헤드를 대폭 줄이고, 더 적은 메모리로 매우 큰 transformers를 훈련시키며 대략 30%의 처리량 증가를 가능하게 한다.

ABSTRACT

Training large transformer models is one of the most important computational challenges of modern AI. In this paper, we show how to significantly accelerate training of large transformer models by reducing activation recomputation. Activation recomputation is commonly used to work around memory capacity constraints. Rather than storing activations for backpropagation, they are traditionally recomputed, which saves memory but adds redundant compute. In this work, we show most of this redundant compute is unnecessary because we can reduce memory consumption sufficiently without it. We present two novel yet very simple techniques: sequence parallelism and selective activation recomputation. In conjunction with tensor parallelism, these techniques almost eliminate the need to recompute activations. We evaluate our approach on language models up to one trillion parameters in scale and show that our method reduces activation memory by 5x, while reducing execution time overhead from activation recomputation by over 90%. For example, when training a 530B parameter GPT-3 style model on 2240 NVIDIA A100 GPUs, we achieve a Model Flops Utilization of 54.2%, which is 29% faster than the 42.1% we achieve using recomputation. Our implementation will be available in both Megatron-LM and NeMo-Megatron.

연구 동기 및 목표

• 대형 transformer 훈련에서 활성화 저장소의 메모리 병목을 동기 부여한다.
• 높은 계산 오버헤드 없이 활성화 메모리를 줄이는 방법을 개발한다.
• 텐서 병렬성과 시퀀스 병렬성을 결합해 재계산을 최소화한다.
• 메모리와 오버헤드를 추가로 줄이기 위해 선택적 활성 재계산을 도입한다.
• 1조 파라미터에 근접한 모델들까지의 경험적 평가를 제공한다.

제안 방법

• 트랜스포머 레이어의 활성화 메모리 추정치를 도출해 메모리 필요량을 정량화한다(식 1).
• 어텐션 및 MLP 블록 내부의 활성화를 분할하기 위해 텐서 병렬성을 적용한다(식 2).
• 시퀀스 차원에 따라 비 텐서 병렬 영역을 분할하는 시퀀스 병렬성을 도입한다(식 4로 이어짐).
• 텐서 및 시퀀스 병렬성과 파이프라인 병렬성을 결합하고 변환기 g와 g를 도출해 추가 통신을 최소화한다(도표 5 및 6).
• 높은 활성화 영역(예: 어텐션 관련 연산)에 목표를 두고 재계산을 선택적으로 적용해 FLOP 오버헤드는 작게 유지하면서 메모리를 줄인다(식 6).
• Megatron-LM 및 NeMo-Megatron 구현을 사용해 1T 파라미터에 이르는 모델에서 메모리 및 실행 시간 영향력을 평가한다.

실험 결과

연구 질문

• RQ1 표준 메모리 레이아웃으로 대형 transformer 모델을 훈련할 때 필요한 활성화 메모리는 얼마나 되는가?
• RQ2 하이브리드 텐서/시퀀스/파이프라인 병렬성이 큰 계산/통신 페널티를 크게 유발하지 않으면서 활성화 메모리를 줄일 수 있는가?
• RQ3 선택적 활성 재계산이 FLOP 오버헤드를 낮게 유지하면서 메모리를 실질적으로 줄일 수 있는가?
• RQ4 이 기술들을 적용할 때 트릴리언-파라미터 transformer의 엔드 투 엔드 처리량은 어떤 수준으로 향상되는가?

주요 결과

• 시퀀스 병렬성과 선택적 활성 재계산을 함께 적용하면 활성화 메모리를 약 5배 줄일 수 있다.
• 선택적 활성 재계산은 대형 모델에서 재계산 FLOP 오버헤드를 3% 이하로 낮추고, 전체 재계산의 30-40%에 비해 낮은 수준으로 유지한다.
• 전체 재계산 대비 테스트 구성에서 엔드-투-엔드 반복 처리량이 약 29-32% 향상된다.
• GPT-3 530B 및 MT-NLG에서 이 접근법은 트랜스포머 레이어별 메모리를 크게 줄이고 heavy recomputation 없이 학습을 가능하게 한다.
• 530B GPT-3 스타일 모델에 대한 2240-A100 구성에서 Model FLOPs Utilization은 54.2%로, 기준치 42.1%보다 29% 빠르다.
• 모델 규모가 커질수록 메모리 절감 효과가 확장되며, 전체적으로 약 5x 감소를 달성해 기본 활성화가 디바이스 메모리를 초과하는 경우를 가능하게 한다.