vTrain: A Simulation Framework
for Evaluating Cost-effective and Compute-optimal Large Language Model Training

GPU 활용률 10% 저하로 훈련 비용 수백만 달러 증가, vTrain으로 해결책 제시

대규모 언어 모델(LLM)이 다양한 응용 분야에서 널리 보급됨에 따라 인공지능 커뮤니티가 직면한 중요한 과제는 이러한 대규모 AI 모델을 비용 효율적인 방식으로 훈련하는 방법이다. 기존의 LLM 훈련 계획은 일반적으로 LLM 병렬화 공간에 대한 철저한 검토보다는 경험적 관찰에 기반한 휴리스틱 기반 병렬 훈련 전략을 채택한다. 이러한 한계로 인해 기존 시스템은 상당한 성능 향상의 여지를 남겨두게 되며, 수백만 달러에 달하는 훈련 비용이 낭비된다. 예를 들어, 1,024대의 NVIDIA A100 GPU를 사용하여 GPT-3(175B 매개변수)를 훈련할 때, GPU 컴퓨팅 활용률이 단지 10%(50%에서 40%로) 감소하더라도 훈련 시간이 8일 증가하여 수백만 달러의 추가 비용이 발생한다. 그러나 이러한 대규모 LLM을 훈련해야 하는 규모 때문에 훈련 시스템 구성의 설계 공간을 철저히 탐색하여 가장 최적의 비용 효율적인 하이퍼파라미터를 찾는 것은 극히 어려운 일이다.

수십 분 내 최적 훈련 전략 도출하는 vTrain의 혁신적 시뮬레이션 기술

카이스트가 발표한 논문에 따르면, 비용 효율적이고 컴퓨팅 최적의 LLM 훈련 시스템 평가를 안내하는 프로파일링 기반 시뮬레이터인 vTrain은 해당 문제를 해결하는 데 도움이 된다. vTrain은 AI 실무자들에게 효율적이고 비용 효율적인 LLM 훈련 시스템 구성을 결정하기 위한 빠르고 정확한 소프트웨어 프레임워크를 제공한다. vTrain의 핵심 특징은 프로파일링 기반 방법론을 사용하여 각 설계 지점의 LLM 훈련 시간을 정확하게 추정하는 것이다. 이는 고성능 멀티코어 CPU 서버에서 몇 십 분 내에 최적의 LLM 훈련 시스템 구성을 결정할 수 있게 한다.

vTrain의 설계는 LLM 훈련 시간을 정확하게 추정할 수 있게 하는 다음과 같은 핵심 관찰에 기반한다. 첫째, 최신 AI 알고리즘은 각 그래프 노드가 신경망 레이어를 나타내는 비순환 그래프로 표현된다. 둘째, LLM 추론과 달리, 훈련을 위한 LLM 그래프 노드의 실행 순서는 컴파일 시간에 정확하게 정의되므로 vTrain은 얼마나 많은 LLM 그래프 노드를 실행해야 하는지와 그 실행 순서를 정적으로 결정할 수 있다. 셋째, 대상 GPU 아키텍처에서 각 개별 LLM 그래프 노드(각 레이어)의 실행 시간은 매우 결정적이며 서로 다른 실행 간에 거의 변동이 없다.

경험적 방식 대비 10% 적은 GPU로 5% 비용 절감, vTrain의 사례 연구

vTrain의 실용성을 입증하기 위해 여러 사례 연구를 실시했다.

첫 번째 사례는 비용 효율적인 LLM 훈련 계획이다. 주어진 LLM, 훈련 토큰 크기 및 컴퓨팅 예산(즉, 총 GPU 수)이 주어졌을 때, 벽시계 훈련 시간과 그에 관련된 훈련 비용을 최소화하는 가장 최적의 훈련 병렬화 전략을 결정할 수 있다.

두 번째는 비용 효율적인 멀티테넌트 LLM 스케줄링으로, 여러 LLM 훈련 작업이 GPU 클러스터를 공유할 때, GPU 활용률을 최대화하면서 작업 완료 시간을 최소화하는 효율적인 스케줄링 알고리즘을 식별할 수 있다.

세 번째는 컴퓨팅 최적의 LLM 모델 설계로, 고정된 컴퓨팅 및 훈련 시간 예산이 주어졌을 때, Chinchilla 스케일링 법칙을 만족하는 가장 큰 LLM을 결정할 수 있다. 예를 들어, MT-NLG(530B) 모델 훈련에서 vTrain은 기존 방식보다 10% 적은 GPU를 사용하면서 4.5% 높은 GPU 활용률을 달성하고, 훈련 비용을 5% 절감하는 훈련 계획을 도출했다.

텐서, 데이터, 파이프라인 병렬화의 최적 조합으로 LLM 훈련 효율성 극대화

현대 LLM 훈련은 최첨단 3D 병렬화 방식(데이터 병렬화, 텐서 병렬화, 파이프라인 병렬화)을 적용한다. 이는 LLM과 같은 거대한 AI 모델을 분할하여 여러 GPU에서 효율적으로 학습시키기 위한 전략이다. 텐서 병렬화는 모델 가중치를 GPU 내에서 열과 행 차원으로 나누어 같은 노드 내 GPU 간에 고대역폭 통신을 활용한다. 데이터 병렬화와 파이프라인 병렬화는 주로 노드 간 병렬화에 사용되며, 상대적으로 통신 오버헤드가 적다. vTrain은 이러한 복잡한 병렬화 전략의 성능을 정확하게 모델링하고, 최적의 구성을 찾아내어 GPU 활용률을 높이고 훈련 비용을 최소화할 수 있다.

145억에서 76억 매개변수로: vTrain으로 발견한 30일 내 훈련 가능한 현실적 모델 크기

Chinchilla 스케일링 법칙에 따르면, 주어진 컴퓨팅 예산 내에서 최적의 모델 크기와 훈련 토큰 수 사이에는 균형이 필요하다. 단순히 모델 크기만 키우는 것은 과소훈련으로 이어져 알고리즘 성능을 완전히 활용하지 못한다. GPU 효율성에 대한 현실적인 평가 없이 단순히 가용 GPU 수만으로 컴퓨팅 예산을 결정하면 오해의 소지가 있다. vTrain은 실제 GPU 활용률을 고려하여 보다 현실적인 컴퓨팅 최적 모델 크기를 도출할 수 있다. 예를 들어, 420개의 NVIDIA DGX A100 서버(3,360 A100 GPU)를 30일 동안 사용한다고 가정할 때, 단순히 100% GPU 활용률을 가정하면 1,456억 매개변수의 모델을 2,912억 토큰으로 훈련할 수 있다고 예상할 수 있다.

그러나 vTrain은 실제로는 평균 35.56%의 GPU 활용률만 달성 가능하며, 이는 원래 기대했던 30일 대신 85일의 훈련 시간이 필요함을 보여준다. vTrain을 사용하면 760억 매개변수의 모델을 1,521억 토큰으로 30일 내에 훈련할 수 있는 더 현실적인 계획을 수립할 수 있다.

FAQ

Q: 대규모 언어 모델 훈련에서 GPU 활용률이 왜 그렇게 중요한가요?
A: GPU 활용률은 훈련 시간과 비용에 직접적인 영향을 미칩니다. vTrain의 연구에 따르면 GPU 활용률이 단 10% 감소하더라도(50%에서 40%로) 훈련 시간이 8일 증가하며, 이는 수백만 달러의 추가 비용을 의미합니다. 따라서 최적의 병렬화 전략을 통한 GPU 활용률 최적화는 비용 효율적인 LLM 훈련에 필수적입니다.

Q: vTrain은 어떻게 기존 LLM 훈련 방식보다 더 효율적인 방법을 찾아낼 수 있나요?
A: vTrain은 프로파일링 기반 시뮬레이션을 통해 수천 가지의 가능한 병렬화 구성을 빠르게 평가하여 최적의 훈련 계획을 도출합니다. 기존 방식은 경험적 관찰에 기반한 휴리스틱에 의존하지만, vTrain은 전체 설계 공간을 체계적으로 탐색하여 GPU 활용률과 훈련 시간 사이의 최적 균형점을 찾아냅니다.

Q: Chinchilla 스케일링 법칙이란 무엇이며 LLM 훈련에 어떤 영향을 미치나요?
A: Chinchilla 스케일링 법칙은 주어진 컴퓨팅 예산 내에서 모델 크기와 훈련 토큰 수를 균형있게 확장해야 한다는 원칙입니다. 이 법칙에 따르면, 모델을 과소훈련하면 해당 모델의 알고리즘 잠재력을 완전히 발휘할 수 없습니다. vTrain은 실제 GPU 효율성을 고려하여 이 법칙을 적용함으로써, 주어진 시간과 자원 내에서 훈련할 수 있는 최적의 모델 크기와 토큰 수를 더 정확하게 예측할 수 있습니다.

해당 기사에서 인용한 논문 원문은 링크에서 확인할 수 있다.

이미지 출처: 카이스트

기사는 클로드와 챗GPT를 활용해 작성되었습니다.