[논문 리뷰] PipeFisher: Efficient Training of Large Language Models Using Pipelining and Fisher Information Matrices

Kazuki Osawa, Shigang Li외 1명|arXiv (Cornell University)|2022. 11. 25.

Topic Modeling인용 수 4

한 줄 요약

PipeFisher는 파이프라인 병렬 LLM 훈련 중 발생하는 유휴 시간(버블)을 활용해 K-FAC라는 피셔 정보 행렬 기반의 이阶 최적화 방법을 실행하는 혁신적인 훈련 방식을 제안한다. K-FAC의 곡률 계산 및 행렬 역행렬 계산을 이러한 유휴 기간으로 이관함으로써, GPU 활용도를 59.8%에서 97.6%로 끌어올리고, BERT-Large의 시뮬레이션 Phase 1 훈련 시간을 75.7% 감소시켰으며, 더 나은 최적화로 인해 수렴 속도 향상도 달성한다.

ABSTRACT

Pipeline parallelism enables efficient training of Large Language Models (LLMs) on large-scale distributed accelerator clusters. Yet, pipeline bubbles during startup and tear-down reduce the utilization of accelerators. Although efficient pipeline schemes with micro-batching and bidirectional pipelines have been proposed to maximize utilization, a significant number of bubbles cannot be filled using synchronous forward and backward passes. To address this problem, we suggest that extra work be assigned to the bubbles to gain auxiliary benefits in LLM training. As an example in this direction, we propose PipeFisher, which assigns the work of K-FAC, a second-order optimization method based on the Fisher information matrix, to the bubbles to accelerate convergence. In Phase 1 pretraining of BERT-Base and -Large models, PipeFisher reduces the (simulated) training time to 50-75% compared to training with a first-order optimizer by greatly improving the accelerator utilization and benefiting from the improved convergence by K-FAC.

연구 동기 및 목표

• 대규모 언어 모델(LLM)의 파이프라인 병렬 훈련에서 발생하는 유휴 '버블'으로 인한 가속기 활용도 저하 문제를 해결하기 위해.
• 이러한 유휴 기간을 활용해 단순한 활용도 향상 외에도 유의미한 부가적 이점을 제공하는 계산 작업을 할당하기 위해.
• K-FAC를 통한 이阶 최적화가 통신 오버헤드를 증가시키지 않고도 파이프라인 스케줄링에 효율적으로 통합될 수 있음을 보여주기 위해.
• 향상된 활용도와 K-FAC를 통한 더 빠른 수렴이 결합되어 훈련 시간을 크게 단축시킬 수 있음을 보여주기 위해.

제안 방법

• PipeFisher는 GPipe나 Chimera와 같은 파이프라인 병렬 훈련 스케줄링에서 발생하는 유휴 버블에 K-FAC의 곡률 계산 및 행렬 역행렬 계산을 통합한다.
• 파이프라인 버블 기간 동안 피셔 정보 행렬(A_l, B_l)의 계산 및 그 역행렬 계산과 같은 K-FAC 작업을 스케줄링하여 전진 및 역전파를 차단하지 않도록 한다.
• 초기 훈련 단계에서 수렴 안정성을 확보하기 위해 이전 단계의 오래된 역행렬 행렬을 사용하여 첫 번째 전치 조건을 제공한다.
• 기존의 파이프라인 스케줄링 프레임워크를 활용하며, 표준 파이프라인 훈련을 초과하는 최소한의 계산 오버헤드(~단계당 6.5%)만 추가로 필요로 한다.
• 모든 파이프라인 병렬 방식과 호환되며, 트랜스포머 및 비트랜스포머 아키텍처 모두에 적용 가능하지만, 비균일한 레이어의 경우 로드 밸런싱이 더 복잡하다.
• 더 빠른 훈련 주기 덕분에 다양한 학습률 스케줄을 평가하는 데 드는 비용을 줄여 하이퍼파rameter 검색을 효율적으로 가능하게 한다.

실험 결과

연구 질문

• RQ1 파이프라인 병렬 LLM 훈련에서 발생하는 유휴 시간(버블)이 계산적으로 복잡한 최적화 방법을 실행하는 데 효과적으로 활용될 수 있는가?
• RQ2 파이프라인 버블에 K-FAC 계산을 이관함으로써 전체 훈련 효율성과 수렴 속도가 향상되는가?
• RQ3 통신 비용을 증가시키지 않고도 버블에 보조 작업을 할당함으로써 GPU 활용도를 얼마나 높일 수 있는가?
• RQ4 PipeFisher를 통한 K-FAC 통합이 BERT 훈련에서 훈련 시간과 모델 성능에 어떤 영향을 미치는가?

주요 결과

• Chimera 파이프라인 스케줄링을 사용한 BERT-Large 훈련에서 PipeFisher는 GPU 활용도를 59.8%에서 97.6%로 끌어올려 유휴 시간을 크게 줄였다.
• BERT-Large Phase 1의 시뮬레이션 훈련 시간이 275.1분에서 208.3분으로 감소하여 75.7% 향상되었으며, 이는 높은 활용도와 K-FAC로 인한 더 빠른 수렴 덕분이었다.
• BERT-Base의 경우 PipeFisher로 훈련 시간이 기준선의 50%로 줄어들어, 다양한 모델 크기에서 일관된 성능 향상을 입증했다.
• PipeFisher의 계산 오버헤드는 매우 낮으며(~단계당 6.5%), 전치 조건 계산 외에는 표준 파이프라인 훈련을 초과하는 추가 비용이 거의 없다.
• 더 빠른 훈련 주기 덕분에 다양한 학습률 스케줄을 평가하는 데 소요되는 시간을 줄여 하이퍼파rameter 검색을 빠르게 할 수 있다.
• PipeFisher는 Shampoo와 같은 다른 이阶 최적화 방법과도 일반화 가능하며, 예를 들어 날카움 인식 최소화(SAM)와 같은 다른 보조 작업으로도 확장 가능하다.