TorchSpec: 추측 디코딩 초안 모델을 대규모로 학습하는 파이토치 네이티브 프레임워크

9bow · 7월 8, 2026, 12:30오후

TorchSpec 소개

TorchSpec은 추측 디코딩(speculative decoding)을 위한 초안 모델(draft model)을 학습하는 파이토치 네이티브 프레임워크입니다. 추측 디코딩은 작은 초안 모델이 여러 토큰을 미리 제안하고, 큰 대상(target) 모델이 이를 한 번의 순전파로 검증하는 방식으로 대규모 언어 모델의 생성 속도를 끌어올리는 기법입니다. 제안한 토큰이 받아들여지면 한 스텝에 여러 토큰을 확정할 수 있어 처리량과 지연 시간이 함께 개선됩니다. 잘 학습된 초안 모델일수록 이 효과가 안정적으로 나오기 때문에, 초안 모델을 어떻게 학습하느냐가 추측 디코딩의 실효성을 좌우합니다.

초안 모델 학습의 핵심은 대상 모델의 중간 은닉 상태(hidden states)를 초안 모델로 전달하는 데 있습니다. 그런데 프론티어 모델의 규모가 커지면서 이 은닉 상태를 옮기는 일 자체가 새로운 병목이 됐습니다. 예를 들어 EAGLE-3 방식은 대상 모델의 은닉 상태 3개 층을 사용하는데, Kimi K2.5 대상으로 초안 모델을 학습할 때 128K 토큰짜리 학습 샘플 하나가 약 7GB의 은닉 상태를 요구합니다(PyTorch 블로그). 데이터셋 규모로 늘어나면 이 양은 감당하기 어려운 수준이 됩니다.

TorchSpec은 은닉 상태를 만드는 추론 시스템과 그것을 소비하는 학습 시스템을 완전히 분리(disaggregation)하는 방식으로 이 문제에 접근합니다. 은닉 상태를 디스크에 쓰는 대신, 추론 엔진 그룹에서 학습 워커 그룹으로 Mooncake 스토어를 거쳐 RDMA(Remote Direct Memory Access) 또는 TCP로 곧장 스트리밍합니다. 덕분에 디스크 저장 부담을 없애면서 추론 자원과 학습 자원을 서로 독립적으로 확장할 수 있습니다. 본 게시물에서는 TorchSpec이 풀려는 병목, 분리형 아키텍처의 설계, 실제 학습 결과와 사용법을 정리합니다.

TorchSpec이 해결하는 은닉 상태 전송 병목

기존 초안 모델 학습 파이프라인은 대체로 두 가지 방식 중 하나를 따랐고, 둘 다 모델과 컨텍스트 길이가 커질수록 한계를 드러냅니다.

첫 번째는 추론과 학습을 같은 GPU에 올리는 코로케이션(co-located) 방식 입니다. 대상 모델이 순전파로 은닉 상태를 만들면 초안 모델이 곧바로 이를 소비합니다. 이 방식은 초안 모델의 병렬화 전략이 대상 모델에 묶여 유연성이 떨어지고, 추론과 학습이 같은 자원을 공유해 독립적으로 확장하기 어렵습니다. 무엇보다 대상 모델이 GPU 메모리를 크게 점유합니다. PyTorch 블로그의 분석에 따르면 1조 파라미터 MoE 모델인 Kimi K2.5(약 575GB 가중치)를 8×H100에 올리면 GPU당 남는 메모리가 약 8GB에 불과해, 컨텍스트 길이 4096으로만 학습할 수 있습니다.

두 번째는 은닉 상태를 미리 계산해 디스크에 저장해 두는 오프라인(offline) 방식 입니다. 추론과 학습을 분리할 수 있다는 장점이 있지만, 긴 컨텍스트를 다루는 대형 모델에서는 저장 용량이 폭발합니다. Kimi K2.5 기준으로 샘플 하나(컨텍스트 131,072 토큰)가 약 7GB이고, 이를 데이터셋 규모로 곱하면 10만 샘플에 700TB가 필요합니다(PyTorch 블로그). 이 규모에서는 분산 파일 시스템에 심한 I/O 부하가 걸리고, 직렬화·역직렬화 비용도 학습을 눈에 띄게 느리게 만듭니다.

정리하면 코로케이션은 메모리와 확장성에서, 오프라인은 저장과 I/O에서 각각 막힙니다. TorchSpec은 이 둘 사이의 절충이 아니라, 은닉 상태를 디스크에 내리지 않고 네트워크로 흘려보내는 세 번째 길을 택합니다.

TorchSpec의 분리형 학습 아키텍처

TorchSpec의 구조는 위 그림처럼 세 부분으로 나뉩니다. 왼쪽의 추론 엔진 그룹(Inference Group)이 대상 모델을 돌려 은닉 상태를 생성하고, 오른쪽의 학습 그룹(Training)이 여러 랭크(rank)로 나뉘어 초안 모델을 학습합니다. 그리고 두 그룹 사이에는 Mooncake 기반의 은닉 상태 스토어(Hidden States Store)가 자리 잡아, 텐서를 디스크에 쓰지 않고 RDMA 또는 TCP로 스트리밍합니다. 추론 컨트롤러와 학습 컨트롤러는 Ray 위에서 작업(task)과 샘플, 메타데이터를 주고받으며 양쪽을 조율합니다.

이 분리형 설계는 앞의 두 병목을 정면으로 해소합니다. 추론과 학습의 GPU 수를 독립적으로 정할 수 있어, 은닉 상태 생성 처리량을 높이려면 추론 엔진을 늘리고 더 큰 배치나 FSDP 샤딩이 필요하면 학습 GPU를 늘리면 됩니다. 학습 GPU는 온전히 초안 모델에만 쓰이므로 긴 시퀀스와 큰 배치에 필요한 메모리를 확보할 수 있고, 은닉 상태가 추론에서 학습으로 곧장 흐르기 때문에 디스크 저장과 직렬화 비용이 사라집니다.

전송 계층으로 Mooncake를 고른 이유도 분명합니다. Mooncake는 Moonshot AI와 칭화대학교가 프로덕션 LLM 서빙의 KV 캐시 관리를 위해 개발한 전송 엔진으로, RDMA와 TCP를 하나의 API로 다루고 GPU Direct RDMA로 CPU를 거치지 않고 GPU 메모리에 데이터를 직접 넣습니다. 텐서를 미리 등록한 고정(pinned) 메모리 버퍼에 담아 복사 없이(zero-copy) 옮기므로, 샘플마다 수 기가바이트의 은닉 상태를 다뤄야 하는 TorchSpec의 요구에 잘 맞습니다.

또 하나의 설계 선택은 은닉 상태를 프로덕션 추론 엔진이 직접 생성하도록 한 추론 엔진 네이티브(inference-engine-native) 방식입니다. 템플릿 포매팅과 토크나이징, 커널이 실제 배포 환경과 일치하므로 학습과 배포 사이의 간극이 없고, 추론 엔진이 지원하는 모델이라면 학습 측 변경을 거의 하지 않고 초안 모델을 학습할 수 있습니다. MoE나 멀티모달 같은 새 아키텍처, FP8·INT4 같은 양자화 모델도 엔진이 지원하는 한 그대로 다룰 수 있습니다.

TorchSpec의 학습 결과와 긴 컨텍스트 지원

TorchSpec 팀은 이 프레임워크로 Kimi K2.5용 EAGLE-3 초안 모델을 학습했습니다. 1500 H200 GPU 시간을 들여 60만 개 학습 샘플, 60억 토큰 규모로 확장한 결과입니다(PyTorch 블로그). 이렇게 학습한 초안 모델은 lookahead 3 토큰 조건에서 출력 처리량이 배치 크기 1에서 60% 이상, 배치 크기 8에서 30%, 배치 크기 16에서 26% 향상됐습니다.

긴 컨텍스트 학습에서 분리형 설계의 이점이 특히 두드러집니다. 학습 GPU 메모리를 초안 모델에만 쓸 수 있으므로, lookahead 4 조건에서 H100 한 장이 최대 44K 토큰, B200 한 장이 최대 200K 토큰 길이의 입력 시퀀스를 학습할 수 있습니다(PyTorch 블로그). 앞서 본 코로케이션 방식이 8×H100에서 4096 토큰에 묶였던 것과 대비됩니다.

TorchSpec이 지원하는 추론 백엔드와 모델

TorchSpec은 여러 추론 엔진에서 은닉 상태를 받아올 수 있으며, 지원 등급은 다음과 같습니다.

백엔드	지원 등급	상태
vLLM	First-class	사용 가능
TokenSpeed	First-class	진행 중
TensorRT-LLM	First-class	사용 가능
SGLang	커뮤니티 지원	사용 가능
HuggingFace Transformers	커뮤니티 지원	사용 가능

은닉 상태를 뽑아내는 방식은 엔진마다 다릅니다. vLLM에서는 워커 확장(Worker Extension) 기능으로 모델 순전파에 훅을 걸어 워커 프로세스 안에서 은닉 상태를 직접 잡아내 RPC 직렬화 비용을 피하고, SGLang에서는 기존 코드베이스에 패치를 적용합니다. TensorRT-LLM에서는 엔진의 SaveHiddenStates 추측 모드를 활용해 잡아낸 은닉 상태를 디스크 대신 Mooncake로 보내도록 작은 패치를 얹습니다.

학습 예제는 Qwen3-8B, Kimi-K2.5, MiniMax-M2.5 등을 대상으로 저장소의 examples/ 디렉토리에 정리돼 있습니다. TorchSpec으로 학습한 초안 모델들은 LightSeek Foundation의 Hugging Face 조직에 kimi-k2.5-eagle3를 비롯해 여러 버전이 공개돼 있습니다.

TorchSpec 설치 및 사용법

tools/build_conda.sh 스크립트로 vLLM 또는 SGLang 백엔드와 함께 설치합니다.

# vLLM과 함께 설치
./tools/build_conda.sh 1 vllm
micromamba activate torchspec

# 또는 SGLang과 함께 설치
./tools/build_conda.sh
micromamba activate torchspec

단일 노드 4 GPU(학습 2 + 추론 2) 환경에서 Qwen3-8B용 EAGLE-3 초안 모델을 학습하는 최소 예제는 다음과 같습니다.

./examples/qwen3-8b-single-node/run.sh

설정값은 명령줄에서 바로 덮어쓸 수 있습니다.

./examples/qwen3-8b-single-node/run.sh training.learning_rate=5e-5 training.num_train_steps=500

학습이 끝난 FSDP 체크포인트는 tools/convert_to_hf.py로 Hugging Face 형식으로 변환합니다. 이때 초안 모델의 lm_head를 더 작은 토큰 집합으로 줄이는 어휘 가지치기(vocabulary pruning)를 함께 적용할 수 있습니다.

python tools/convert_to_hf.py --input-dir ./outputs/my_experiment/iter_0010000/

TorchSpec의 라이선스

TorchSpec은 MIT 라이선스로 공개되어 있어 개인 및 상업적 목적으로 자유롭게 사용할 수 있습니다.

TorchSpec 소개 블로그 (PyTorch 공식 블로그)

TorchSpec 프로젝트 GitHub 저장소

TorchSpec으로 학습한 초안 모델

더 읽어보기

이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다.

파이토치 한국 사용자 모임이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일로 보내드립니다! 텔레그램(Telegram)이나 Slack/Discord/Teams/Dooray/GoogleChat 등으로도 새 글 알림을 받으실 수 있습니다.

아래쪽에 좋아요를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~