ByteScale: Communication-Efficient Scaling of LLM Training with a 2048K Context Length on 16384 GPUs
https://dl.acm.org/doi/10.1145/3718958.3754352
https://conferences.sigcomm.org/sigcomm/2025/program/papers-info/ >> NetAI & Wireless
https://youtu.be/cXsMyc7ROyo?si=-EAnwxwNIZbZBd4I
summary
ByteScale은 긴 시퀀스와 짧은 시퀀스가 혼합된 큰 스케일의 학습을 위해 LLM longterm context 학습 과정에서 발생하는 두 가지 문제 (불필요한 통신, 불균형한 연산)를 해결하였다.
논문은 LLM 학습에서 사용되는 기존의 static한 DP x CP 병렬화 시스템이 데이터의 heterogeneity를 고려하지 못하는 한계 때문에 발생하는 불필요한 통신과 불균형한 연산 문제에 집중하고 있다. 불필요한 통신은 짧은 길이의 시퀀스임에도 긴 시퀀스처럼 분산 처리하여 통신이 강제로 발생되는 것, 불균형한 연산은 시퀀스 길이에 따라 달라지는 계산량 때문에 GPU의 idle 상태가 발생하는 것을 의미한다.
ByteScale은 세 가지 핵심 기법(HDP, Balance scheduler, Selective offloading)을 활용하여 이를 해결하였다. HDP, Hybrid data parallelism은 기존 CP가 짧은 시퀀스임에도 긴 시퀀스처럼 같은 그룹 크기로 분할해서 발생하는 불필요한 통신 문제를 해결하였다. 기존의 DP와 CP를 통합하여 시퀀스 길이에 따라 동적으로 GPU 그룹을 조정할 수 있도록 하였다. 이때 짧은 시퀀스의 경우 DP방식처럼 독립적으로 처리하여 불필요한 통신을 제거하고 최소 통신만 이루어지게 하였고 긴 시퀀스의 경우 필요한 랭크만 묶어서 메모리 한계 문제를 해결하였다.
balance scheduler는 시퀀스 길이 S에 따라 마이크로배치별 실행시간이 달라져 attention이 O(S^2)으로 변화하기 때문에 같은 토큰 수라도 연산량의 불균형이 발생하는 문제를 해결하였다. 같은 토큰 수가 아니라 비슷한 실행시간이 될 수 있게 매이크로배치 수를 랭크별로 유연하게 조정하고 짧은 실행시간의 랭크에 더 많은 마이크로배치가 할당되도록 해서 DP/PP Bubble을 줄였다. 그래서 장치 간 idle time을 줄일 수 있었다.
selective offloading은 activation을 선택적으로 CPU로 오프로딩하여 긴 시퀀스에서 필요한 랭크 수를 줄이고 D2H와 H2D 전송을 GPU 연산과 오버랩되도록 하여 통신 비용을 추가로 줄였다. 순차적인 전송으로 인해 발생하던 기존의 통신 병목을 완화하였다.
strengths
- 초대규모 환경에서의 효과 증명 : 256K~2048K 컨텍스트 길이, 16384개의 GPU를 사용하여 초대규모 환경에서 최신 트레이닝 시스템보다 ByteScale이 더 효과적임을 실험적으로 입증하였다.
- generality : LLaMA와 같은 단일 모델 유형인 dense LLM에서만 최적화한 것이 아니라 Mistral 같은 sparse LLM에서도 성능이 향상되었음을 입증하여 ByteScale이 모델 아키텍처에 종속되지 않으면서 성능 향상을 이루었음을 입증하였다.
- 전략별 기여도 증명 : Ablation studies에서 동적 커뮤니케이션과 선택적 오프로딩, 균형 전략, 원격 데이터로더 각각의 효율성을 실험적으로 증명하여 각각의 기법별 성능에 대한 기여도를 구체적, 실험적으로 증명하였다.
weakness
- generality 확인 부족 : 평가 대상 LLM 모델의 유형이 LLaMA와 Mistral(MoE)으로 한정되어 있어서 논문에서 증명한 효과적인 성능이 encorder, decorder, multi modal 등 다른 모델 계열에 대한 효과는 본문에서 확인할 수 없었다.
- 재현성 : 16384 GPU임이 명시되었으나 세부적인 사양을 자세하게 언급하지 않아 실험을 재현하고자 하는 경우 어려울 수 있다.
- 길이 분포 의존성으로 인한 적용 범위 제한 : 장문 비중이 높은 경우는 장문 비중이 낮은 경우에 비해 평균 처리량과 스피드업이 낮아진다고 명시적으로 논문에 언급되어있다. ByteScale의 이득이 워크로드의 길이 분포에 민감하고 배포 환경이나 데이터 비율이 달라지면 성능이 하락할 수 있어 ByteScale의 적용 범위가 제한될 것이다. 분포 변화에 적응할 수 있는 스케쥴링이나 파라미터 자동화 등의 개선이 필요할 것이다.
detailed comments
Byted가 256K 이상의 토큰이 37%를 차지하고 있어 장문 비중이 크기 때문에 평균 처리량과 스피드업이 Github보다 낮음을 직접적으로 명시하였다. 이 부분에서 데이터 분포가 성능 차이를 만든다는 것을 논문 저자가 직접 언급했음이 인상적이다. 반면 github는 256K 이상의 토큰이 9.8%이기 때문에 짧은 시퀀스가 대부분이며 이러한 경우 불필요한 통신을 강제로 수행하게 되어 성능이 악화된다. 이러한 짧은 시퀀스, 긴 시퀀스의 혼재 비율과 불필요한 통신, 연산 불균형 문제, 실제 end to end 이득에 대한 내용이 자연스럽게 연결되도록 전개한 점과 메가 스케일 레벨의 GPU 클러스터 환경으로 실험한 점이 인상적이었다.