핵심 요약
대형 언어 모델(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를 거치지 않는 제로 카피 전송을 구현했다.
학습 GPU의 메모리를 온전히 드래프트 모델에만 할당할 수 있어 B200 GPU 기준 최대 20만 토큰의 시퀀스 학습이 가능하며 이는 기존 공동 배치 방식으로는 불가능했던 수준이다.
vLLM 및 SGLang과 같은 상용 추론 엔진을 직접 활용하여 학습과 배포 환경 간의 토크나이저 및 커널 격차를 없앴으며 새로운 모델 아키텍처에 대한 지원도 엔진 업데이트만으로 즉시 가능하다.
1조 파라미터 규모의 Kimi K2.5 모델을 위한 EAGLE-3 드래프트 모델을 1,500 H200 GPU 시간 동안 학습한 결과 배치 사이즈 1에서 처리량이 60% 이상 향상되는 성과를 거두었다.
실무 Takeaway
- 대규모 LLM의 추측 디코딩 성능을 극대화하려면 TorchSpec과 같은 분산형 학습 프레임워크를 도입하여 메모리 제약 없이 10만 토큰 이상의 긴 컨텍스트 데이터를 학습해야 한다.
- Mooncake 엔진의 RDMA 스트리밍을 활용하면 수백 TB 규모의 중간 데이터 저장 비용을 제거하고 디스크 I/O 병목 없이 학습 속도를 높일 수 있다.
- vLLM 등 실제 서비스용 추론 엔진을 학습 파이프라인에 직접 통합함으로써 배포 환경과 동일한 데이터 분포에서 모델을 최적화하고 구현 복잡도를 낮출 수 있다.
언급된 리소스
GitHubkimi-k2.5-eagle3
AI 분석 전체 내용 보기
AI 요약 · 북마크 · 개인 피드 설정 — 무료