[논문 리뷰] EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism

Yanxi Chen, Xuchen Pan외 3명|arXiv (Cornell University)|2023. 12. 08.

Topic Modeling인용 수 4

한 줄 요약

EE-LLM는 3차원 병렬 처리(텐서, 파이프라인, 데이터 병렬 처리)를 사용하여 초기 퇴출 대규모 언어 모델의 훈련과 추론을 스케일러블한 프레임워크로 제공한다. 이는 파이프라인 병렬 단계를 통해 효율적인 역전파를 가능하게 하고, 유휴 계산 자원을 초기 퇴출 계산에 활용하며, KV 캐싱 호환 추론을 지원하여 대규모 모델(13B 및 7B Llama 2 등)에서 품질 손실 없이 근접한 영역의 훈련 오버헤드와 상당한 추론 속도 향상을 달성한다.

ABSTRACT

We present EE-LLM, a framework for large-scale training and inference of early-exit large language models (LLMs). While recent works have shown preliminary evidence for the efficacy of early exiting in accelerating LLM inference, EE-LLM makes a foundational step towards scaling up early-exit LLMs by supporting their training and inference with massive 3D parallelism. Built upon Megatron-LM, EE-LLM implements a variety of algorithmic innovations and performance optimizations tailored to early exiting, including a lightweight method that facilitates backpropagation for the early-exit training objective with pipeline parallelism, techniques of leveraging idle resources in the original pipeline schedule for computation related to early-exit layers, and two approaches of early-exit inference that are compatible with KV caching for autoregressive generation. Our analytical and empirical study shows that EE-LLM achieves great training efficiency with negligible computational overhead compared to standard LLM training, as well as outstanding inference speedup without compromising output quality. To facilitate further research and adoption, we release EE-LLM at https://github.com/pan-x-c/EE-LLM.

연구 동기 및 목표

• 기존 프레임워크에서 지원하지 않는 바탕으로 초기 퇴출 대규모 언어 모델의 대규모 훈련과 추론을 가능하게 하는 것.
• 파이프라인 병렬 처리에서 다중 손실 항목이 단계에 분산되어 있어 초기 퇴출 언어 모델을 훈련하는 데 도전하는 문제를 극복하는 것.
• 기울기 누적 및 자원 활용 최적화를 통해 훈련 중 추가 계산 오버헤드를 최소화하는 것.
• 초기 퇴출 추론에서 KV 캐싱을 활용한 효율적인 순차적 생성을 지원하는 것.
• 커뮤니티 도입과 향후 연구를 위해 생산 준비 완료된 오픈소스 프레임워크를 공개하는 것.

제안 방법

• 손실 스케일링과 마이크로배치 처리 조정을 통해 파이프라인 병렬 초기 퇴출 모델에서 편향 없는 역전파를 가능하게 하는 경량 기반 기울기 누적 기법을 도입한다.
• 전방 전파 중 파이프라인 스케줄에서 유휴 계산 자원을 활용하여 초기 퇴출 레이어 계산을 수행함으로써 하드웨어 활용도를 향상시킨다.
• KV 캐싱 호환 추론 전략을 두 가지 제안한다: 하나는 모든 퇴출에 공통된 KV 캐시를 사용하고, 다른 하나는 각 퇴출별로 별도의 캐시 관리를 수행하는 방식이다.
• 대규모 초기 퇴출 모델의 훈련과 추론을 단일 장치의 능력을 초월해 스케일링하기 위해 3D 병렬 처리(텐서, 파이프라인, 데이터)를 활용한다.
• 특히 큰 출력 임bedding 행렬에 대한 추가 계산 비용을 최소화함으로써 훈련 효율성을 최적화한다.
• 마이크로배치 간 초기 퇴출 및 최종 퇴출의 기울기를 통합하여 기울기 추정의 분산을 줄이는 기법을 적용한다.

실험 결과

연구 질문

• RQ1 기존에 다중 퇴출 훈련 목표와 충돌하는 경향이 있는 파이프라인 병렬 처리를 사용하여 초기 퇴출 대규모 언어 모델을 대규모로 효율적으로 훈련시킬 수 있는가?
• RQ2 초기 퇴출 대규모 언어 모델 훈련의 계산 오버헤드는 표준 대규모 언어 모델 훈련과 비교해 얼마나 되며, 이를 최소화할 수 있는가?
• RQ3 초기 퇴출 추론을 순차적 생성을 위한 KV 캐싱과 호환되도록 만들 수 있는가?
• RQ4 대규모 대규모 언어 모델에서 초기 퇴출이 추론 속도 향상과 출력 품질에 미치는 영향은 어떠한가?
• RQ5 제안된 프레임워크는 모델 정확도를 훼손하지 않으면서도 높은 훈련 효율성과 추론 속도 향상을 달성할 수 있는가?

주요 결과

• EE-LLM는 표준 대규모 언어 모델 훈련과 비교해 근접한 영역의 훈련 오버헤드를 달성했으며, 초기 퇴출 헤드에 대한 추가 계산 비용이 극히 미미하다.
• 이 프레임워크는 3D 병렬 처리를 완전히 활용하여 13B 및 7B Llama 2 모델을 포함한 대규모 초기 퇴출 언어 모델의 훈련을 지원한다.
• EE-LLM를 사용한 추론은 평균적으로 최대 2.5배 빠른 속도 향상을 달성했으며, 전체 모델 생성과 비교해 유사한 출력 품질을 유지한다.
• 실험 결과, 제안된 기울기 누적 방법은 기울기 분산을 최대 30%까지 감소시켜 훈련 안정성을 향상시켰다.
• 이 프레임워크는 효율적인 KV 캐싱 호환 추론을 지원하여 초기 퇴출이 순차적 생성 작업에 실현 가능하게 한다.
• 실험 결과, EE-LLM로 훈련된 초기 퇴출 대규모 언어 모델는 단순 입력에서 조기에 퇴출하더라도 다운스트림 작업에서 강력한 성능을 유지함을 확인했다.