DeepSpeed 추론

DeepSpeed 추론 (DeepSpeed-Inference)

대규모 Transformer 모델을 서빙할 때도 메모리와 지연 시간이 문제예요. DeepSpeed-Inference는 그런 모델을 효율적으로 서빙하기 위한 기능 모음이에요. 모델 병렬화(MP)를 지원해서 GPU 메모리에 안 들어가는 큰 모델도 배포할 수 있고, 작은 모델에서도 지연 시간을 줄이는 데 도움을 줘요.

출처: DeepSpeed 추론 튜토리얼 (공식)

추론 모드 초기화

init_inference API로 모델을 추론 모드로 로드해요. 여기서 MP degree를 지정할 수 있고, 체크포인트를 미리 로드하지 않았다면 체크포인트 경로나 json 파일로 설명을 넘겨줄 수 있어요.

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)

replace_with_kernel_inject=True를 주면 컨볼루션·GeMM 같은 연산을 커널 단위로 교체해 성능을 높여요. 추론은 Hugging Face의 pipeline과 함께 쓸 수 있어서, 기존 서빙 코드에 자연스럽게 녹아들어요.

커널 인젝션 정책 지정하기

커널 인젝션을 지원하지 않는 모델이라면 injection_policy로 특정 linear 레이어를 지정해 텐서 병렬화를 적용할 수 있어요. Transformer의 인코더/디코더 레이어에서 두 개의 linear 레이어(어텐션 출력 GeMM, 레이어 출력 GeMM)를 지정해 주면 돼요.

from transformers.models.t5.modeling_t5 import T5Block
import deepspeed

pipe = pipeline("text2text-generation", model="google/t5-v1_1-small", device=local_rank)

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')
    }
)

체크포인트 로드하기

Hugging Face로 학습한 모델은 위처럼 from_pretrained로 그대로 로드하면 돼요. 모델 병렬로 학습한 Megatron-LM 모델이라면, 모든 모델 병렬 체크포인트 목록을 JSON 설정으로 넘겨줘야 해요.

{
  "type": "Megatron",
  "version": 0.0,
  "checkpoints": [
    "mp_rank_00/model_optim_rng.pt",
    "mp_rank_01/model_optim_rng.pt"
  ]
}

DeepSpeed로 학습한 체크포인트라면 경로만 저장하면 돼요.

{
  "type": "ds_model",
  "version": 0.0,
  "checkpoints": "path_to_checkpoints"
}

흥미로운 점은, 추론의 MP degree가 학습 때와 달라도 된다는 거예요. 예를 들어 MP 없이 학습한 모델을 MP=2로 추론하거나, MP=4로 학습한 모델을 MP 없이 추론해도 돼요. 초기화 과정에서 DeepSpeed가 체크포인트를 자동으로 병합하거나 분할해 줘요.

멀티 GPU로 실행하기

여러 GPU에서 추론을 띄우려면 DeepSpeed 런처를 쓰면 돼요.

deepspeed --num_gpus 2 inference.py

이렇게 하면 지정한 GPU 수만큼 모델이 분산되어 서빙돼요.

DeepSpeed 추론