DeepSpeed 인퍼런스 시작하기 — 트랜스포머 모델 추론 (Inference)

DeepSpeed 인퍼런스 시작하기 — 트랜스포머 모델 추론 (Inference)

한 줄로 요약하면, DeepSpeed-Inference는 PyTorch 트랜스포머 모델을 훨씬 효율적으로 서빙하기 위해 몇 가지 기법을 함께 제공해요. 모델 병렬화(MP)로 GPU 메모리에 안 들어가는 큰 모델을 나눠 담고, 작은 모델도 지연 시간을 줄이는 데 쓰죠. 여기에 추론 전용 커널을 주입하고, 모델을 압축하는 MoQ 양자화까지 얹으면 비용까지 줄일 수 있어요. 참고로 이 튜토리얼의 공식 문서는 DeepSpeed-Inference v2를 가리키며, 최신 버전은 곧 DeepSpeed-FastGen 이야기라는 점도 짚어 둘게요.

출처: DeepSpeed 공식 문서 — Getting Started with DeepSpeed for Inferencing Transformer based Models

추론 초기화: init_inference

DeepSpeed로 추론하려면 init_inference API로 모델을 추론 모드로 불러와요. 여기서 모델 병렬(MP) 차수를 지정하고, 아직 체크포인트를 못 불러온 상태라면 JSON 파일이나 체크포인트 경로로 설명을 넘겨줄 수 있어요. 고성능 커널을 주입하려면 호환 모델에 대해 replace_with_kernel_inject를 True로 설정하면 돼요.

핵심은 모델링 쪽을 전혀 바꿀 필요가 없다는 거예요. DeepSpeed, Megatron, HuggingFace로 학습한 호환 트랜스포머 모델이면 내보내기나 별도 체크포인트 생성 없이 그대로 추론 모드로 전환돼요. 다중 GPU 추론이면 모델 병렬 차수와 체크포인트 정보만 주면 되고, 나머지(모델 분할, 커널 주입, GPU 간 통신)는 DeepSpeed가 알아서 처리해요. 커널 주입이 가능한 모델 목록은 repo 안의 replace_policy.py에서 확인할 수 있어요.

import deepspeed

# DeepSpeed-Inference 엔진 초기화
ds_engine = deepspeed.init_inference(
    model,
    tensor_parallel={"tp_size": world_size},
    dtype=torch.half,
    checkpoint=None if args.pre_load_checkpoint else args.checkpoint_json,
    replace_with_kernel_inject=True,
)
model = ds_engine.module

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
output = pipe('Input String')

커널을 지원하지 않는 모델이라 모델 병렬만 쓰고 싶다면, 트랜스포머 인코더/디코더 레이어의 두 특정 선형 레이어(1) 어텐션 출력 GeMM과 (2) 레이어 출력 GeMM을 가리키는 injection_policy를 넘겨줘요. 이 두 부분이 GPU 간 all-reduce 통신으로 부분 결과를 합쳐야 하는 지점이라서 꼭 필요해요. T5로 쓰는 예시를 보면 이렇습니다.

# create the model
import transformers
from transformers.models.t5.modeling_t5 import T5Block
import deepspeed

pipe = pipeline("text2text-generation", model="google/t5-v1_1-small", device=local_rank)
# DeepSpeed-Inference 엔진 초기화
pipe.model = deepspeed.init_inference(
    pipe.model,
    tensor_parallel={"tp_size": world_size},
    dtype=torch.float,
    injection_policy={T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')},
)
output = pipe('Input String')

체크포인트 불러오기

HuggingFace로 학습한 모델은 위처럼 from_pretrained로 체크포인트를 미리 불러올 수 있어요. Megatron-LM으로 모델 병렬로 학습한 경우엔 병렬 체크포인트 전체 목록을 JSON 설정으로 넘겨야 해요. 재미있는 점은 추론 MP 차수가 학습 때와 달라도 된다는 거예요 — MP 없이 학습한 모델을 MP=2로 돌리거나, 반대로 MP=4로 학습한 모델을 MP 없이 추론해도 DeepSpeed가 초기화 과정에서 체크포인트를 자동으로 합치고 나눠요.

실행하기

여러 GPU에서 추론을 실행할 땐 DeepSpeed 런처를 사용해요.

deepspeed --num_gpus 2 inference.py

GPT-NEO 2.7B 엔드투엔드 예제

DeepSpeed 추론은 HuggingFace pipeline과 함께 쓸 수 있어요. 아래는 GPT-NEO-2.7B로 텍스트를 생성하는 엔드투엔드 클라이언트 코드예요.

import os
import deepspeed
import torch
from transformers import pipeline

local_rank = int(os.getenv('LOCAL_RANK', '0'))
world_size = int(os.getenv('WORLD_SIZE', '1'))

generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B', device=local_rank)
generator.model = deepspeed.init_inference(
    generator.model,
    tensor_parallel={"tp_size": world_size},
    dtype=torch.float,
    replace_with_kernel_inject=True,
)
string = generator("DeepSpeed is", do_sample=True, min_length=50)

여기서 눈여겨볼 점은, 원래 모델이 모델 병렬 없이 학습됐고 체크포인트도 단일 GPU용인데도 모델 병렬 텐서 슬라이싱으로 여러 GPU에 걸쳐 추론을 돌린다는 거예요. deepspeed 런처 없이도 환경변수 LOCAL_RANK, WORLD_SIZE만 제대로 세팅되면 동작해요.

데이터 타입과 양자화 모델

DeepSpeed 추론은 fp32, fp16, int8 파라미터를 지원해요. init_inference의 dtype으로 적절한 타입을 지정하면 그 타입에 최적화된 커널을 자동으로 골라 줘요. int8 양자화 모델이라면, DeepSpeed의 MoQ 방식으로 양자화했을 경우 그 양자화 설정(그룹 수, 트랜스포머 MLP 부분의 추가 그룹화 여부)을 init_inference에 넘겨줘야 해요.

import deepspeed
model = deepspeed.init_inference(
    model,
    checkpoint='./checkpoint.json',
    dtype=torch.int8,
    quantization_setting=(quantize_groups, mlp_extra_grouping),
)

이렇게 하면 DeepSpeed 인퍼런스 튜토리얼이 끝나요. 핵심만 다시 말하면 — 모델 병렬로 큰 모델을 분할하고, 호환 모델엔 커널을 주입하며, 필요하면 MoQ로 양자화해서 비용까지 낮추는 흐름이에요.

더 알아보기 (Learn more)