대규모 MoE 모델 추론 — DeepSpeed-MoE Inference

MoE(Mixture-of-Experts) 모델은 거대한 파라미터를 갖지만 모든 파라미터를 늘 쓰진 않아서, 추론에서도 특별한 전략이 필요해요. DeepSpeed-MoE Inference는 dense 모델의 추론 최적화(DeepSpeed-Inference) 위에 여러 병렬화를 얹어서, 낮은 지연 시간과 뛰어난 처리량 확장을 동시에 노려요. 전문가(expert)가 아닌 파라미터에는 데이터 병렬과 텐서 슬라이싱을, 전문가 파라미터에는 전문가 병렬(expert-parallelism)과 전문가 슬라이싱(expert-slicing)을 적용하죠. DeepSpeed-MoE 라이브러리로 학습한 MoE 모델이면 DeepSpeed-Inference 엔진으로 eval 모드에 올리기만 하면 바로 추론할 수 있어요.

출처: DeepSpeed 공식 문서 — Getting Started with DeepSpeed-MoE for Inferencing Large-Scale MoE Models

추론 초기화

init_inference API로 MoE 모델을 추론용으로 불러와요. 여기서 텐서 슬라이싱 차수(mp_size), 전문가 병렬 차수(ep_size), 전문가 수(moe_experts)를 지정할 수 있어요. DeepSpeed는 world_size(전체 GPU 수)와 전문가 병렬 크기 중 작은 값 기준으로 프로세스 그룹을 만들고, 이 그룹을 통해 전문가를 전문가-병렬 GPU에 나눠 담아요. 전문가 수가 GPU 수보다 적으면 전문가 슬라이싱으로 전문가 파라미터를 GPU 사이에 분할하죠. 아직 체크포인트를 못 불러온 경우엔 JSON 파일이나 체크포인트 경로로 설명을 넘겨주고, 고성능 커널을 쓰려면 replace_with_kernel_inject를 True로 설정해요.

import deepspeed
import torch.distributed as dist

# 전문가 병렬 크기 설정
world_size = dist.get_world_size()
expert_parallel_size = min(world_size, args.num_experts)

# MoE 모델 생성
moe_model = get_model(model, ep_size=expert_parallel_size)

# DeepSpeed-Inference 엔진 초기화
ds_engine = deepspeed.init_inference(
    moe_model,
    mp_size=tensor_slicing_size,
    dtype=torch.half,
    moe_experts=args.num_experts,
    checkpoint=args.checkpoint_path,
    replace_with_kernel_inject=True,
)
model = ds_engine.module
output = model('Input String')

다양한 설정 옵션

DeepSpeed 추론 엔진은 텐서 슬라이싱 차수, 전문가 수, 전체 GPU 수를 조합해 다양한 병렬화 그룹을 자동으로 만들어요. 전문가 파라미터는 먼저 전문가-병렬로 각 전문가 그룹을 GPU 하나에 배정하고, GPU 수가 전문가 수보다 많으면 전문가 슬라이싱으로 각 전문가를 GPU에 수직/수평 분할하죠. 실제 예시에 쓰이는 인자는 이렇게 생겼어요.

generate_samples_gpt.py \
  --tensor-model-parallel-size 1 \
  --num-experts ${experts} \
  --num-layers 24 \
  --hidden-size 2048 \
  --num-attention-heads 32 \
  --max-position-embeddings 1024 \
  --tokenizer-type GPT2BPETokenizer \
  --load $checkpoint_path \
  --fp16 \
  --ds-inference

여기서 마지막 ds-inference 플래그가 DeepSpeed-MoE와 PyTorch 구현을 오갈 수 있게 해주는 핵심이에요.

성능: 딥 MoE에서의 확장

52B 아키텍처(전문가 128개, dense 1.3B)를 GPU를 늘려가며 돌려보면, DeepSpeed-MoE(Generic)가 PyTorch 대비 24%~~60% 성능 향상을 보여요. 여기에 통신 최적화와 MoE 전용 커널 같은 전체 기능을 켜면 속도 향상이 2x~~3.2x까지 올라가요.

규모를 더 키운 설정은 이렇게 비교돼요 — 2.4B+MoE-128은 A100 8개로, 24B+MoE-128은 A100 64개로 추론했어요.

| 모델 | 크기(billions) | #Layers | Hidden size | MP degree | EP degree | | 2.4B+MoE-128 | 107.7 | 16 | 3584 | 1 | 64 - 128 | | 24B+MoE-128 | 1046.9 | 30 | 8192 | 8 | 64 - 128 |

표준 MoE 기준으로 DeepSpeed는 두 모델에서 PyTorch보다 각각 1.4x, 1.65x 성능을 냈고, PR-MoE 최적화를 쓰면 모델 품질을 유지하면서 1.81x, 1.87x로 끌어올릴 수 있어요.

PR-MoE 최적화로 더 낮은 비용

서로 다른 MoE 구조를 고르기 위해 추론 예시에 mlp-type 파라미터가 추가됐어요. 'standard'와 'residual' 중 선택해서 PR-MoE가 제공하는 모델링 최적화를 켤 수 있어요.

generate_samples_gpt.py \
  --tensor-model-parallel-size 1 \
  --num-experts ${experts} \
  --mlp_type 'residual' \
  --num-layers 24 \
  --hidden-size 2048 \
  --num-attention-heads 16 \
  --max-position-embeddings 1024 \
  --tokenizer-type GPT2BPETokenizer \
  ...

24B+MoE-128처럼 non-expert 부분이 커서 단일 GPU에 안 들어가는 모델은 모델 병렬 크기를 1보다 크게 잡고, 최고 성능을 내도록 텐서 슬라이싱 차수를 조정해요.

대규모 MoE 모델 추론 — DeepSpeed-MoE Inference