자동 텐서 병렬화

자동 텐서 병렬화 (Automatic Tensor Parallelism)

과거에는 DeepSpeed 추론에서 텐서 병렬화를 쓰려면 직접 injection policy를 지정해야 했어요. 이제는 커널 인젝션을 쓰지 않고 injection policy도 제공하지 않으면, Hugging Face 모델에 대해 자동 텐서 병렬화가 기본으로 적용돼요. 이 페이지는 그런 자동 텐서 병렬화 기능을 다뤄요.

출처: DeepSpeed Automatic Tensor Parallelism (공식)

코드가 얼마나 단순해지나요

자동 모드에서는 init_inference에 mp_size만 넘겨주면 돼요. 텐서 병렬화 degree를 world size로 잡으면, 모델이 여러 GPU에 자동으로 분산돼요.

import os
import torch
import transformers
import deepspeed

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

# 모델 파이프라인 생성
pipe = transformers.pipeline(
    task="text2text-generation",
    model="google/t5-v1_1-small",
    device=local_rank
)

# DeepSpeed-Inference 엔진 초기화
pipe.model = deepspeed.init_inference(
    pipe.model,
    mp_size=world_size,
    dtype=torch.float
)
output = pipe('Input String')

예전 방식과 비교

자동 텐서 병렬화가 없던 시절엔, 커널 인젝션을 지원하지 않는 모델에 텐서 병렬화만 적용하려면 injection policy를 수동으로 넘겨줘야 했어요. Transformer의 인코더/디코더 레이어에 있는 두 개의 linear 레이어, 즉 어텐션 출력 GeMM과 레이어 출력 GeMM을 지정해 주는 방식이었죠.

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

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

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

pipe.model = deepspeed.init_inference(
    pipe.model,
    mp_size=world_size,
    dtype=torch.float,
    injection_policy={
        T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')
    }
)
output = pipe('Input String')

자동 텐서 병렬화가 도입되면서 이런 수동 설정이 불필요해졌어요. 호환이 되는 모델이라면 그냥 mp_size만 넘기면 되는 거죠.

성능 확인하기

공식 inference 테스트 스위트로 지연 시간, 대역폭, 처리량, 메모리까지 확인해 볼 수 있어요.

deepspeed --num_gpus <num_gpus> \
  DeepSpeedExamples/inference/huggingface/text-generation/inference-test.py \
  --name <model> --batch_size <batch_size> --test_performance

텐서 병렬화를 켜려면 호환 모델에 ds_inference 플래그를 추가하면 돼요.

deepspeed --num_gpus <num_gpus> \
  DeepSpeedExamples/inference/huggingface/text-generation/inference-test.py \
  --name <model> --batch_size <batch_size> --test_performance --ds_inference

공식 벤치마크(T5 11B, OPT 13B)에서 자동 텐서 병렬화를 켰을 때 지연 시간과 처리량이 개선되는 것을 확인할 수 있어요.

자동 텐서 병렬화