TorchServe로 대형 모델 서빙하기
TorchServe로 대형 모델 서빙하기
큰 모델을 서빙할 때 가장 먼저 부딪히는 벽은 '모델이 한 GPU에 안 들어간다'는 문제예요. 이 문서에서 말하는 대형 모델이 바로 그런 모델이에요 — GPU 하나에 담기지 않아서 여러 GPU에 나눠 담아야(partition) 하죠. TorchServe가 이렇게 모델을 여러 GPU로 분산해 서빙하는 방식을 어떻게 지원하는지, 그리고 vLLM·PiPPy·DeepSpeed 같은 프레임워크를 어떻게 얹어 쓰는지 살펴볼게요.
동작 원리
작은 모델의 GPU 추론에서 TorchServe는 워커당 하나의 프로세스를 띄우고 GPU 하나씩 배정해요. 반면 대형 모델은 여러 GPU로 나눠야 하죠. 분할 방식은 크게 파이프라인 병렬(PP, pipeline parallel), 텐서 병렬(TP, tensor parallel), 또는 이 둘의 조합으로 나뉘는데, 어떤 방식을 쓰고 어떻게 쪼갤지는 사용하는 프레임워크 구현에 달려 있어요. TorchServe는 유연한 설정으로 여러 프레임워크의 요구를 수용해요. PiPPy·DeepSpeed처럼 GPU마다 별도 프로세스를 띄워야 하는 프레임워크가 있고, vLLM처럼 모든 GPU를 한 프로세스에 배정하는 경우도 있어요.
여러 프로세스가 필요한 경우 TorchServe는 torchrun으로 워커의 분산 환경을 구성해요. 각 워커에 배정된 GPU마다 프로세스를 하나씩 시작하죠. torchrun을 쓸지는 model-config.yaml의 parallelType 파라미터로 결정돼요.
pp— 파이프라인 병렬tp— 텐서 병렬pptp— 파이프라인 + 텐서 병렬custom— 병렬화 방식을 사용자에게 맡김
처음 세 옵션은 torchrun으로 환경을 구성하고, custom은 병렬화 방식을 사용자에게 맡기되 워커에 배정된 GPU를 단일 프로세스로 전달해요. 배정되는 GPU 수는 torchrun이 시작하는 프로세스 수(nproc-per-node) 또는 parallelLevel 파라미터로 결정돼요. 즉 parallelLevel은 nproc-per-node와 함께 설정하면 안 되고 둘 중 하나만 써야 해요.
기본적으로 TorchServe는 호스트의 GPU를 라운드로빈(round-robin) 방식으로 워커에 배정해요. 대형 모델 추론에선 model_config.yaml에 명시한 GPU 수에 따라 워커별 할당 GPU가 자동 계산되고, 그 숫자에 맞춰 CUDA_VISIBLE_DEVICES가 설정돼요.
예를 들어 노드에 GPU가 8개 있고, 워커 하나가 GPU 4개가 필요하다고 해볼게요(nproc-per-node=4 또는 parallelLevel=4). 그러면 TorchServe는 worker1에 CUDA_VISIBLE_DEVICES="0,1,2,3", worker2에 CUDA_VISIBLE_DEVICES="4,5,6,7"을 배정해요.
기본 동작 외에도 워커에 쓸 GPU를 직접 지정할 수 있어요. model config YAML에 deviceIds: [2,3,4,5]를 설정하고 nproc-per-node(또는 parallelLevel)를 2로 잡으면, worker1은 CUDA_VISIBLE_DEVICES="2,3", worker2는 CUDA_VISIBLE_DEVICES="4,5"를 받아요.
PiPPy로 파이프라인 병렬 (PyTorch 네이티브)
PiPPy는 한 GPU에 안 들어가는 대형 모델을 서빙하기 위한 파이프라인 병렬 솔루션이에요. 모델을 지정한 디바이스 수에 맞춰 크기가 같은 스테이지(stage)로 쪼갠 뒤, 마이크로배칭(microbatching)으로 배치 입력을 추론해요. 배치 크기가 1보다 클 때 특히 유리하죠.
TorchServe에서 PiPPy를 쓰려면 base_pippy_handler를 상속한 커스텀 핸들러를 만들고, 설정을 model-config.yaml에 넣어야 해요. TorchServe의 커스텀 핸들러는 모델 로딩·전처리·추론·후처리 로직을 정의하는 파이썬 스크립트예요. custom_handler.py처럼 이름을 붙여 만든 핸들러는 대략 이렇게 생겼어요.
from ts.torch_handler.distributed.base_pippy_handler import BasePippyHandler
from ts.handler_utils.distributed.pt_pippy import initialize_rpc_workers, get_pipline_driver
class ModelHandler(BasePippyHandler, ABC):
def __init__(self):
super(ModelHandler, self).__init__()
self.initialized = False
def initialize(self, ctx):
model = # load your model from model_dir
self.device = self.local_rank % torch.cuda.device_count() # move model inputs to self.device
self.model = get_pipline_driver(model, self.world_size, ctx)
model-config.yaml에는 이렇게 설정을 넣어요. 이 파일은 유연해서 frontend, backend, handler 관련 설정을 모두 담을 수 있어요.
# frontend settings
minWorkers: 1
maxWorkers: 1
maxBatchDelay: 100
responseTimeout: 120
deviceType: "gpu"
parallelType: "pp" # pp(pipeline), tp(tensor), pptp(pipeline+tensor)
# 프론트엔드가 입력을 rank0 또는 전체 rank로 라우팅하는 데 쓰임
# (예: DeepSpeed는 tp, PiPPy는 pp 지원)
torchrun:
nproc-per-node: 4 # torchrun이 시작할 프로세스 수. world_size 또는
# 쪼갤 GPU 수로 설정
# backend settings
pippy:
chunks: 1 # 마이크로배치 크기 설정. microbatch = batch size / chunks
input_names: ['input_ids'] # 모델 입력 인자 이름. FX tracing에 필요
model_type: "HF" # Huggingface 모델이면 HF, 아니면 비워두거나 다른 값
rpc_timeout: 1800
num_worker_threads: 512 # rpc worker 초기화 스레드 수
handler:
max_length: 80 # 핸들러에서 토크나이저가 처리할 최대 토큰 수
핸들러에서는 이렇게 설정 값에 접근해요.
def initialize(self, ctx):
model_type = ctx.model_yaml_config["pippy"]["model_type"]
모델 패키징과 서버 시작은 TorchServe에서 평소와 같아요. 패키징할 때 반드시 model-config.yaml을 넘겨요.
torch-model-archiver --model-name bloom --version 1.0 --handler pippy_handler.py --extra-files $MODEL_CHECKPOINTS_PATH -r requirements.txt --config-file model-config.yaml --archive-format tgz
텐서 병렬(TP)은 아직 준비 중이라 준비되는 대로 추가될 예정이에요.
DeepSpeed로 텐서 병렬
DeepSpeed-Inference는 마이크로소프트의 오픈소스 프로젝트로, 한 GPU 메모리에 안 들어가는 대형 트랜스포머 기반 PyTorch 모델을 서빙하기 위한 모델 병렬을 제공해요.
TorchServe에서 DeepSpeed를 쓰려면 base_deepspeed_handler를 상속한 커스텀 핸들러를 만들고 설정을 model-config.yaml에 넣어요. 핸들러는 이렇게 생겼어요.
from ts.handler_utils.distributed.deepspeed import get_ds_engine
from ts.torch_handler.distributed.base_deepspeed_handler import BaseDeepSpeedHandler
class ModelHandler(BaseDeepSpeedHandler, ABC):
def __init__(self):
super(ModelHandler, self).__init__()
self.initialized = False
def initialize(self, ctx):
model = # load your model from model_dir
ds_engine = get_ds_engine(self.model, ctx)
self.model = ds_engine.module
self.initialized = True
model-config.yaml은 이렇게 구성해요.
# frontend settings
minWorkers: 1
maxWorkers: 1
maxBatchDelay: 100
responseTimeout: 120
deviceType: "gpu"
parallelType: "tp" # pp/tp/pptp 중 선택
torchrun:
nproc-per-node: 4 # torchrun이 시작할 프로세스 수
# backend settings
deepspeed:
config: ds-config.json # DeepSpeed config json 파일명
# 상세: https://www.deepspeed.ai/docs/config-json/
handler:
max_length: 80 # 핸들러에서 토크나이저가 처리할 최대 토큰 수
ds-config.json 예시는 다음과 같아요.
{
"dtype": "torch.float16",
"replace_with_kernel_inject": true,
"tensor_parallel": {
"tp_size": 2
}
}
DeepSpeed 설치 — 방법 1은 requirements.txt에 넣는 거고, 방법 2는 명령으로 미리 설치하는 방식이에요(모델 로딩 속도를 위해 권장).
DS_BUILD_OPS=1 pip install deepspeed
패키징 시 model-config.yaml을 반드시 넘겨요. 모델 디렉토리를 쓰거나 HuggingFace model_name을 쓰는 두 가지 방식이 있어요.
# option 1: model_dir 사용
torch-model-archiver --model-name bloom --version 1.0 --handler deepspeed_handler.py --extra-files $MODEL_CHECKPOINTS_PATH,ds-config.json -r requirements.txt --config-file model-config.yaml --archive-format tgz
# option 2: HF model_name 사용
torch-model-archiver --model-name bloom --version 1.0 --handler deepspeed_handler.py --extra-files ds-config.json -r requirements.txt --config-file model-config.yaml --archive-format tgz
DeepSpeed MII
지원하는 모델 중 하나라면 DeepSpeed MII를 활용할 수 있어요. DeepSpeed MII는 DeepSpeed Inference에 딥러닝의 발전을 더해 지연(latency)을 최소화하고 처리량(throughput)을 극대화해요. 특정 모델 유형·모델 크기·배치 크기·가용 하드웨어 자원에서 그렇게 동작하죠. 지원 모델에서 DeepSpeed MII를 쓰는 방법은 공식 문서를, TorchServe 적용 예시는 공식 예제를 참고하면 돼요.
Accelerate로 대형 Hugging Face 모델 서빙
자원이 제한적인데 큰 Hugging Face 모델을 서빙해야 한다면 accelerate를 쓸 수 있어요. setup_config.json에서 low_cpu_mem_usage=True로 두고 device_map="auto"를 설정하면 되죠.
대형 모델 추론 팁
모델 로딩 지연 줄이기
- DeepSpeed 같은 모델 병렬 라이브러리를 컨테이너나 호스트에 미리 설치해요.
- 모델 체크포인트를 미리 다운로드해요. HuggingFace를 쓴다면
Download_model.py로 사전 훈련 모델을 미리 받을 수 있어요. HUGGINGFACE_HUB_CACHE와TRANSFORMERS_CACHE환경 변수를 설정하고,Download_model.py도구로 모델을 HuggingFace 캐시 디렉토리에 받아요.
모델 config YAML 튜닝
- 추론 지연이 커서 응답 타임아웃이 나면
responseTimeout을 올려요. - 모델 로딩 지연이 커서 시작 타임아웃이 나면
startupTimeout을 올려요. torchrun파라미터를 튜닝해요. 예를 들어 기본값 1인OMP_NUMBER_THREADS를 YAML에서 조정할 수 있어요.
# frontend settings
torchrun:
nproc-per-node: 4 # torchrun이 시작할 프로세스 수
OMP_NUMBER_THREADS: 2
지연에 민감한 애플리케이션 — Job Ticket
잦은 지연을 피해야 하는 추론에는 job ticket 기능을 권장해요. 이 기능을 켜면 TorchServe가 클라이언트 요청을 처리할 활성 워커가 있는지 확인하고, 있다면 작업 큐나 동적 배칭 대기 없이 요청을 즉시 처리해요. 반대로 없으면 503 응답을 돌려주죠. 생성 모델이나 ChatGPT 같은 자가회귀(autoregressive) 디코더 모델처럼 추론 지연이 큰 유스케이스에 특히 유용해요. 거부된 요청을 다른 서버로 라우팅하거나 모델 서버 용량을 늘리는 등 비즈니스 요구에 맞게 조치를 취할 수 있어요.
minWorkers: 2
maxWorkers: 2
jobQueueSize: 2
useJobTicket: true
이 예시에서 모델은 워커 2개, 작업 큐 크기 2를 가져요. 추론 요청은 TorchServe가 즉시 처리하거나 503으로 거부돼요.
HTTP 1.1 chunked encoding 스트리밍 응답
TorchServe의 추론 API는 HTTP 1.1 chunked encoding으로 연속된 추론 응답을 보내는 스트리밍을 지원해요. 전체 응답의 추론 지연이 크고 중간 결과를 클라이언트에 전달해야 하는 경우에만 권장돼요. LLM 생성 애플리케이션이 대표적인 경우죠 — n개 토큰을 생성하는 데 지연이 클 수 있는데, 이 기능을 쓰면 전체 응답이 끝나기 전에 생성된 토큰을 하나씩 받아볼 수 있어요. 백엔드 핸들러는 send_intermediate_predict_response를 호출해 중간 결과 하나를 프론트엔드로 보내고, 마지막 결과는 기존 방식처럼 반환해요.
from ts.handler_utils.utils import send_intermediate_predict_response
# 참고: TorchServe v1.0.0부터 아래 경로는 deprecated 예정
# "from ts.protocol.otf_message_handler import send_intermediate_predict_response" 대신
# "from ts.handler_utils.utils import send_intermediate_predict_response" 사용
def handle(data, context):
if type(data) is list:
for i in range(3):
send_intermediate_predict_response(["intermediate_response"], context.request_ids, "Intermediate Prediction success", 200, context)
return ["hello world "]
클라이언트 쪽에서는 chunked 데이터를 받아요.
import test_utils
def test_echo_stream_inference():
test_utils.start_torchserve(no_config_snapshots=True, gen_mar=False)
test_utils.register_model('echo_stream', 'https://torchserve.pytorch.org/mar_files/echo_stream.mar')
response = requests.post(TF_INFERENCE_API + '/predictions/echo_stream', data="foo", stream=True)
assert response.headers['Transfer-Encoding'] == 'chunked'
prediction = []
for chunk in (response.iter_content(chunk_size=None)):
if chunk:
prediction.append(chunk.decode("utf-8"))
assert str(" ".join(prediction)) == "hello hello hello hello world "
test_utils.unregister_model('echo_stream')
gRPC 서버 스트리밍
TorchServe gRPC API는 추론 API StreamPredictions의 서버 스트리밍도 지원해요. 같은 gRPC 스트림으로 연속된 추론 응답을 보내죠. HTTP 1.1 chunked encoding과 마찬가지로, 전체 응답 지연이 크고 중간 결과를 클라이언트로 보내는 LLM 생성 유스케이스에 권장돼요. 이 API는 자동으로 batchSize를 1로 강제해요.
service InferenceAPIsService {
// TorchServe 서버 상태 확인
rpc Ping(google.protobuf.Empty) returns (TorchServeHealthResponse) {}
// 기본 모델 버전으로 추론
rpc Predictions(PredictionsRequest) returns (PredictionResponse) {}
// 추론 요청의 스트리밍 응답
rpc StreamPredictions(PredictionsRequest) returns (stream PredictionResponse) {}
}
백엔드 핸들러는 send_intermediate_predict_response를 호출해 중간 결과를 프론트엔드로 보내고 마지막 결과를 반환해요. HTTP 예시와 같은 패턴이에요.
from ts.handler_utils.utils import send_intermediate_predict_response
def handle(data, context):
if type(data) is list:
for i in range(3):
send_intermediate_predict_response(["intermediate_response"], context.request_ids, "Intermediate Prediction success", 200, context)
return ["hello world "]