[논문 리뷰] Memory-Efficient Pipeline-Parallel DNN Training

Deepak Narayanan, Amar Phanishayee외 3명|arXiv (Cornell University)|2020. 06. 16.

Topic Modeling참고 문헌 22인용 수 62

한 줄 요약

PipeDream-2BW는 이중 버퍼 가중치 업데이트와 자동 파티셔닝으로 메모리 효율적인 파이프라인 병렬화를 도입하여 기존 정확도와 비슷한 수준으로 거대 트랜스포머 모델의 훈련 속도를 높입니다.

ABSTRACT

Many state-of-the-art ML results have been obtained by scaling up the number of parameters in existing models. However, parameters and activations for such large models often do not fit in the memory of a single accelerator device; this means that it is necessary to distribute training of large models over multiple accelerators. In this work, we propose PipeDream-2BW, a system that supports memory-efficient pipeline parallelism. PipeDream-2BW uses a novel pipelining and weight gradient coalescing strategy, combined with the double buffering of weights, to ensure high throughput, low memory footprint, and weight update semantics similar to data parallelism. In addition, PipeDream-2BW automatically partitions the model over the available hardware resources, while respecting hardware constraints such as memory capacities of accelerators and interconnect topologies. PipeDream-2BW can accelerate the training of large GPT and BERT language models by up to 20$ imes$ with similar final model accuracy.

연구 동기 및 목표

• 가속기 메모리 한도 내에서 수십억 매개변수 DNN의 훈련이 가지는 도전 과제를 제시합니다.
• 플러시 유발 없는 처리량 향상을 위한 메모리 효율적 파이프라인 병렬 시스템을 제안합니다.
• 이질적 인터커넥트에 걸쳐 모델을 자동으로 파티셔닝하면서 디바이스 메모리 제약을 준수합니다.
• 대 baselines와 비교하여 큰 트랜스포머 모델(GPT/BERT)의 수렴 및 처리량 이점을 입증합니다.

제안 방법

• 가중치 버전 두 개를 유지하고 마이크로배치에서 그래디언트를 축적하여 메모리 사용량을 줄이는 2BW(이중 버퍼 가중치 업데이트) 도입.
• 계획 알고리즘을 사용해 DNN 연산자를 등분으로 복제된 스테이지(equi-replicated)로 분할하고, 마이크로배치와 그래디언트 축적을 할당해 처리량과 메모리 사이의 균형을 맞춥니다.
• 메모리 발자국을 더 낮추기 위해 단일 가중치 버전과 파이프라인 플러시를 사용하는 PipeDream-Flush 변형 제공.
• 1단계 지연으로 가중치 업데이트 의미를 형식화: W(t+1)=W(t)−η·∇f(W(t−1)) 및 일반적인 업데이트와 수렴 유사성 논의.
• activation 재계산을 도입해 활성화 저장을 줄이고 GPU당 실행 가능한 마이크로배치 크기를 늘립니다.
• AWS에서 NVLink를 사용한 PyTorch로 대형 GPT/BERT 모델을 평가하고 Megatron 스타일 텐서 모델 병렬화 및 GPipe와 비교합니다.

실험 결과

연구 질문

• RQ1 파이프라인 병렬 훈련에서 메모리와 처리량을 수십억 매개변수 모델에 대해 어떻게 균형 잡을 수 있을까요?
• RQ2 이중 버퍼 가중치 업데이트 체계가 파이프라인 플러시 없이 높은 처리량과 관리 가능한 메모리 발자국을 달성할 수 있을까요?
• RQ3 자동 분할 전략 중 어떤 것이 이질적 인터커넥트에 걸친 트랜스포머 유사 아키텍처를 수렴을 유지하면서 최적으로 매핑할 수 있을까요?
• RQ4 2BW 및 PipeDream-Flush 의미가 대형 NLP 모델에서 vanilla 옵티마이저와 비교 가능한 수렴/정확도를 유지합니까?

주요 결과

• 다양한 GPT 모델에 대해 최적화된 모델 병렬 baselines 대비 엔드-투-엔드 속도 향상 1.3x에서 20x.
• 파이프라인 플러시가 없고 메모리 풋프린트가 더 낮아 GPipe보다 최대 3.2배 빠름.
• 64개 8개 또는 16GPU 구성으로 거의 300억 매개변수까지 대형 트랜스포머 모델 학습 가능, 메가트론식 구성이 메모리 제한 상황에서 더 나은 성능.
• Adam 옵티마이저와 2BW는 BERT와 GPT 프리트레이닝 작업에서 vanilla Adam과 유사한 학습 손실 궤적 및 다운스트림 미세조정 정확도를 보임.
• PipeDream-Flush는 처리량 손실의 대가로 메모리 발자국이 더 낮고, 2BW는 두 개의 가중치 버전과 활성화 재계산으로 더 높은 처리량 달성.
• 메모리 발자국 비교에서 GPipe의 메모리 문제, PipeDream-Flush의 중간 발자국, PipeDream-2BW의 중간 이상 발자국.