Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts

1조 규모 AI 모델의 숨은 병목: 디바이스 간 통신이 실행 시간의 47%를 차지

아주 큰 규모의 인공지능(AI) 모델을 운용할 때, 모델 내부에서 발생하는 ‘통신’이 생각보다 많은 시간을 차지해 전체 속도를 저하시킨다는 문제가 있어 왔다. 바이트댄스(ByteDance) 씨드(Seed) 연구팀은 이러한 문제를 해결하기 위해 코멧(Comet)이라는 새로운 기술을 개발했으며, 이를 통해 평균 최대 1.96배 더 빠른 실행이 가능하다고 밝혔다.

인공지능 모델의 규모가 매우 커지면 여러 대의 GPU(그래픽 처리 장치)를 사용해야 하고, 이 때 서로 다른 GPU 사이에 데이터를 주고받는 통신 과정이 필수적으로 발생한다. 특히 혼합 전문가 모델(Mixture-of-Experts, MoE)은 ‘전문가’라는 모듈을 여러 개 두어, 각 입력 데이터에 맞춰 필요한 것만 활성화한다. 이렇게 하면 모델 전체 파라미터(학습에 필요한 숫자)는 매우 많더라도 실제로는 일부만 계산하므로 효율적이다.

그러나 전문가 모듈이 여러 GPU에 분산되어 있으면, 이들 GPU 간 통신이 빈번히 일어난다. 실제로 널리 쓰이는 MoE 모델에서는 통신이 전체 실행 시간의 최대 47%를 차지할 정도로 크게 나타난다고 한다.

청크 단위 처리의 함정: GPU 계산 효율성 저하를 극복한 코멧의 미세 단위 중첩 기술

그동안은 통신과 계산을 중첩하기 위해 입력 데이터를 청크(chunk)로 나누어, 한쪽에서는 계산을, 다른 한쪽에서는 통신을 병렬로 수행하는 파이프라인 방식을 시도해 왔다. 그러나 청크를 잘게 쪼개면 전문가들이 나눠진 데이터를 처리하게 되면서 GPU가 낼 수 있는 최대 계산 성능을 충분히 발휘하기 어려워진다. 또한 MoE는 입력이 유동적으로 변하기 때문에 통신량과 계산량이 항상 동일한 비율로 나뉘지는 않으며, 그 결과 균형 잡힌 중첩이 쉽지 않다.

코멧은 이 문제를 극복하기 위해 훨씬 미세한 단위에서 통신과 계산을 겹치도록 새로 설계한 시스템이다. 통신과 계산이 서로 얽혀 있는 복잡한 의존성을 해결하고, GPU 리소스를 유연하게 배분해 미세 단위 중첩을 실현한다.

텐서 분해와 데이터 의존성 해결

코멧은 텐서(모델이 주고받는 다차원 배열)를 특정 축(차원)에 따라 잘게 나누고, 이 데이터를 처리하는 순서를 재구성함으로써 통신과 계산이 동시에 일어나는 데 필요한 구조를 만든다. 이러한 방법으로, 기존에는 “계산을 시작하기 전 통신이 끝나야 한다”와 같은 복잡한 제약이 사라지고, 통신과 계산 사이의 미세 단위 중첩이 가능해진다.

이를 위해 공유되는 텐서에 대한 의존성 관계를 분석해, 특정 차원을 따라 텐서를 분해하고 연산 순서를 재배치하는 두 가지 과정을 거친다. 이를 통해 세분화된 데이터 의존성을 깨뜨리고, 효율적으로 통신과 계산을 재스케줄링해 더욱 자유로운 실행 순서를 확보한다.

GPU 스레드 블록 배분의 최적화: 통신과 계산 워크로드를 실시간으로 조율하는 적응형 기술

코멧은 이렇게 단위가 잘게 쪼개진 통신과 계산 작업에 대해 GPU의 스레드 블록을 동적으로 배분한다. 예컨대 통신이 집중적으로 필요한 순간에는 통신에 더 많은 스레드 블록을 배정하고, 계산이 주가 될 때는 계산 쪽에 배정해 통신이 계산을 가로막지 않도록 한다.

이러한 적응형 워크로드 할당 기법은 각 시점에서 통신과 계산이 균형을 유지하도록 하며, 이를 통해 지연 시간을 줄이고 전체 실행 성능을 높인다. 코멧은 스레드 블록을 지능적으로 재분배함으로써 통신이 계산 성능에 미치는 영향을 줄이고, 높은 GPU 활용 효율을 유지한다.

1.96배 빨라진 MoE 레이어: 수만 대 GPU 클러스터에서 검증된 코멧의 실전 성능

바이트댄스 연구팀이 NVIDIA H800과 L20 등 최신 GPU 클러스터에서 코멧을 적용해 본 결과, 일반적인 MoE 레이어에서 이전 대비 약 1.96배의 속도 향상을 달성했다고 한다. 또한 모델 전체를 동작시켰을 때, 예를 들어 Mixtral-8x7B, Qwen2-MoE, Phi3.5-MoE 같은 모델들은 평균 1.71배 더 높은 속도를 보였다고 한다.

더 나아가 이 기술은 이미 만 대 이상의 GPU를 사용하는 실제 프로덕션 환경에도 적용되어, 수백만 시간에 달하는 GPU 계산 시간을 절약했다고 한다. 이는 대규모 모델을 훈련하거나 추론하는 데 필요한 막대한 비용과 시간을 줄이는 데 매우 큰 도움이 된다.

수백만 GPU 시간 절약의 혁신: 트릴리언급 AI 모델 개발 비용을 획기적으로 낮추는 코멧 기술

코멧 기술이 공개됨에 따라, 트라이턴(Triton)이나 TVM 같은 컴파일러와 결합하거나 추가 최적화를 진행하는 등 새로운 발전 가능성도 열리고 있다. 통신-계산을 미세 단위로 중첩하는 이 기법은 앞으로 트릴리언 단위의 파라미터를 갖춘 초거대 AI 모델을 훨씬 빠르고 경제적으로 운용하게 해 줄 것으로 기대된다. 결국 코멧은 대규모 AI 모델에서 통신 오버헤드를 줄이기 위해 통신과 계산을 미세 단위로 교차 중첩시키는 혁신적인 방식이다. 이를 통해 최대 1.96배까지 모델을 가속할 수 있고, 그 결과 대규모 모델 개발 비용과 시간을 대폭 절감하는 효과가 있다.

FAQ

Q: 혼합 전문가 모델(MoE)이란 무엇이며 왜 중요한가요?

A: 혼합 전문가 모델은 모든 매개변수를 활성화하지 않고 각 입력에 대해 일부 전문가만 활성화하는 희소 구조를 도입한 모델입니다. 이 방식은 계산 비용을 크게 늘리지 않으면서도 모델 크기를 수조 개 매개변수까지 확장할 수 있게 해주어, AI 모델의 성능 향상에 중요한 역할을 합니다.

Q: 코멧(Comet)은 어떻게 기존 MoE 시스템의 성능을 향상시키나요?

A: 코멧은 통신과 계산 작업 간의 미세 단위 중첩을 가능하게 하는 두 가지 핵심 기술을 도입했습니다. 첫째, 데이터 의존성 해결 방법을 통해 복잡한 의존성을 분석하고 최적화된 파이프라인을 구축합니다. 둘째, 적응형 워크로드 할당 방법으로 GPU 자원을 통신과 계산 작업에 동적으로 배분하여 최적의 성능을 이끌어냅니다.

Q: 코멧의 연구 결과가 실제 AI 개발에 어떤 영향을 미칠까요?

A: 코멧의 기술은 이미 실제 프로덕션 환경에 적용되어 수백만 GPU 시간을 절약하고 있습니다. 이는 대규모 AI 모델 개발 비용과 시간을 크게 줄이며, 더 효율적인 AI 모델 개발을 가능하게 합니다. 또한 이 연구는 미세 단위 계산-통신 최적화 분야에 새로운 접근법을 제시하여 향후 AI 시스템 개발에 중요한 기반이 될 것입니다.

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

이미지 출처: 바이트댄스

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