TorchSpec: 분산형 추측 디코딩 학습을 위한 PyTorch 네이티브 프레임워크

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

대형 언어 모델(LLM)의 추론 가속을 위한 추측 디코딩 기법이 중요해짐에 따라 고성능 드래프트 모델 학습의 필요성이 커지고 있다. 기존의 공동 배치 학습이나 오프라인 데이터 준비 방식은 GPU 메모리 부족과 막대한 스토리지 비용 문제를 야기한다. 이를 해결하기 위해 PyTorch는 추론과 학습 리소스를 분리하고 데이터를 네트워크로 직접 스트리밍하는 TorchSpec 프레임워크를 공개했다. TorchSpec은 Mooncake 엔진을 사용하여 대규모 숨겨진 상태 데이터를 디스크 저장 없이 전송함으로써 학습 효율과 확장성을 획기적으로 개선하며 Kimi K2.5 모델에서 실질적인 성능 향상을 입증했다.

배경

추측 디코딩(Speculative Decoding)의 기본 원리, PyTorch 분산 학습(FSDP)에 대한 이해, RDMA 및 고성능 네트워크 인프라 지식

대상 독자

LLM 추론 최적화 및 대규모 모델 학습을 담당하는 ML 엔지니어 및 연구원

의미 / 영향

TorchSpec은 대규모 LLM의 추론 비용을 획기적으로 낮출 수 있는 고성능 드래프트 모델 학습의 표준을 제시한다. 특히 스토리지 인프라에 대한 의존도를 낮추고 GPU 메모리 활용도를 극대화함으로써 더 긴 컨텍스트를 처리하는 차세대 AI 에이전트 서비스의 실용화를 앞당길 것으로 기대된다.

섹션별 상세

기존 추측 디코딩 학습 방식은 타겟 모델과 드래프트 모델이 GPU 메모리를 공유하여 긴 컨텍스트 학습이 불가능하거나 수백 테라바이트의 데이터를 디스크에 저장해야 하는 I/O 병목 문제를 안고 있다.

TorchSpec은 타겟 모델을 실행하는 추론 엔진 그룹과 드래프트 모델을 학습하는 워커 그룹을 완전히 분리하는 분산형 아키텍처를 채택하여 각 리소스를 독립적으로 확장할 수 있도록 설계되었다.

Moonshot AI가 개발한 Mooncake 전송 엔진을 통합하여 RDMA 또는 TCP를 통해 숨겨진 상태 데이터를 실시간으로 스트리밍하며 GPU Direct RDMA를 지원해 CPU를 거치지 않는 제로 카피 전송을 구현했다.

TorchSpec의 분산형 학습 시스템 아키텍처 다이어그램이다. — Diagram추론 그룹(Inference Group)에서 생성된 숨겨진 상태 데이터가 Mooncake 저장소를 거쳐 학습 그룹(Training)으로 RDMA/TCP를 통해 실시간 스트리밍되는 구조를 보여준다. Ray를 통한 컨트롤러 관리 체계도 확인할 수 있다.

학습 GPU의 메모리를 온전히 드래프트 모델에만 할당할 수 있어 B200 GPU 기준 최대 20만 토큰의 시퀀스 학습이 가능하며 이는 기존 공동 배치 방식으로는 불가능했던 수준이다.

GPU 유형별 최대 시퀀스 길이 비교 차트이다. — Chart공동 배치 학습과 분산형 학습의 최대 지원 토큰 수를 비교한다. 분산형 학습 도입 시 H200에서 2배(128,000 토큰), B200에서 1.7배(218,000 토큰) 더 긴 시퀀스를 학습할 수 있음을 보여준다.

vLLM 및 SGLang과 같은 상용 추론 엔진을 직접 활용하여 학습과 배포 환경 간의 토크나이저 및 커널 격차를 없앴으며 새로운 모델 아키텍처에 대한 지원도 엔진 업데이트만으로 즉시 가능하다.

1조 파라미터 규모의 Kimi K2.5 모델을 위한 EAGLE-3 드래프트 모델을 1,500 H200 GPU 시간 동안 학습한 결과 배치 사이즈 1에서 처리량이 60% 이상 향상되는 성과를 거두었다.

Kimi K2.5 드래프트 모델의 데이터셋별 벤치마크 결과 차트이다. — Chart다양한 데이터셋(Math, Code, Function Call 등)에서 300k 데이터 대비 600k 데이터를 사용했을 때 수락 길이(Accept Length)가 전반적으로 향상되었음을 보여준다. 특히 CEval 데이터셋에서 55%의 높은 향상률을 기록했다.

Eagle3 적용 시 배치 사이즈에 따른 평균 처리량 변화 그래프이다. — ChartEagle3를 적용했을 때 배치 사이즈 1에서 63%, 8에서 30%, 16에서 26%의 처리량 향상이 발생함을 수치로 증명한다. 배치 사이즈가 작을수록 추측 디코딩의 가속 효과가 더 크게 나타남을 알 수 있다.

실무 Takeaway

대규모 LLM의 추측 디코딩 성능을 극대화하려면 TorchSpec과 같은 분산형 학습 프레임워크를 도입하여 메모리 제약 없이 10만 토큰 이상의 긴 컨텍스트 데이터를 학습해야 한다.
Mooncake 엔진의 RDMA 스트리밍을 활용하면 수백 TB 규모의 중간 데이터 저장 비용을 제거하고 디스크 I/O 병목 없이 학습 속도를 높일 수 있다.
vLLM 등 실제 서비스용 추론 엔진을 학습 파이프라인에 직접 통합함으로써 배포 환경과 동일한 데이터 분포에서 모델을 최적화하고 구현 복잡도를 낮출 수 있다.

언급된 리소스

GitHubkimi-600k-training-dataset

GitHubkimi-k2.5-eagle3