매일 수십 기가바이트의 모델 가중치를 GPU 메모리로 옮기는 작업은 대규모 언어 모델(LLM) 학습에서 병목 현상을 일으키는 주범이다. 최근 Unsloth(LLM 학습 속도를 최적화하는 라이브러리)와 NVIDIA(그래픽 처리 장치 제조사)가 협력하여 공개한 최적화 기법은, 단순히 연산 속도를 높이는 것을 넘어 데이터 이동과 연산 사이의 간극을 메우는 데 집중하고 있다.

이중 버퍼링을 통한 데이터 전송과 연산의 병렬화

기존의 학습 방식은 직렬화 패턴(Serialization Pattern)을 따른다. 하나의 버퍼를 사용하여 활성화 값(Activation)을 불러오고(Copy) 이를 역전파 연산(Backward Compute)에 사용하는 과정을 순차적으로 반복한다. 이 경우 전체 시간은 두 작업 시간의 합과 거의 동일해진다. 연구팀은 이를 해결하기 위해 이중 버퍼링(Double Buffering)을 도입했다. 버퍼 A에서 역전파 연산이 진행되는 동안, 복사 스트림은 다음 레이어의 활성화 값을 버퍼 B에 미리 로드한다. 이 방식은 수학적 연산량을 줄이는 것이 아니라, 연산이 진행되는 동안 데이터 복사 지연(Copy Latency)을 숨기는 전략이다.

대규모 모델에서의 성능 향상과 벤치마크 수치

이 최적화는 모델의 크기가 커질수록 더 큰 효과를 발휘한다. 모델이 클수록 은닉 차원(Hidden Dimension)이 넓어져 데이터 이동량이 많아지고, 레이어 수가 늘어나 연산 뒤에 복사 작업을 숨길 기회가 많아지기 때문이다. NVIDIA B200(블랙웰 아키텍처 기반의 고성능 GPU) 환경에서 벤치마크를 수행한 결과, 최종 손실값(Final Loss)에는 변화가 없으면서도 학습 단계별 속도가 유의미하게 개선되었다. 특히 핀드 메모리(Pinned Memory, CPU와 GPU 간 데이터 전송 속도를 높이기 위해 고정된 호스트 메모리) 대역폭이 약 55.7 GB/s인 시스템에서, 수십 밀리초 단위의 복사 지연을 수백 밀리초 단위의 전체 단계 시간 절감으로 치환하는 성과를 거두었다.

MoE 모델의 라우팅 최적화와 동적 인덱싱 개선

예전에는 MoE(Mixture of Experts, 모델의 일부 파라미터만 활성화하는 구조) 모델의 라우팅 과정에서 토큰을 각 전문가에게 할당할 때 데이터 의존적인 연산이 반복되었다. 기존 방식은 다음과 같다.

python

기존의 비효율적인 라우팅 방식

for expert in experts:

tokens = torch.where(routing_mask == expert)

이 과정에서 매번 동적 인덱싱 오버헤드 발생

이제는 모든 토큰을 한 번에 그룹화하여 오프셋(Offset)을 재사용하는 방식으로 전환했다. 이는 런타임에 동적 인덱싱 질문을 던지는 횟수를 획기적으로 줄여, 전문가 수(num_experts)에 비례하여 증가하던 CPU-GPU 동기화 오버헤드를 제거한다. 이러한 최적화는 Unsloth의 native_torch 백엔드를 사용하는 모든 MoE 모델에 즉시 적용 가능하다.

최적화의 핵심은 연산 커널 외부의 병목 제거

개발자가 체감하는 변화는 핵심 연산 커널(Kernel)이 아닌, 그 주변을 감싸고 있는 접착 코드(Glue Code)에서 발생한다. 메인 연산이 빨라질수록 과거에는 보이지 않던 오버헤드가 전체 학습 시간에서 차지하는 비중이 커지기 때문이다. 이번 연구는 수학적 연산을 개선하는 단계를 넘어, 데이터 이동의 효율성과 런타임 호출 횟수를 줄이는 것이 현대적인 고성능 학습 환경의 핵심임을 증명한다.