자동 텐서 병렬화(AutoTP)로 HuggingFace 모델 추론 가속하기

원래 DeepSpeed에서 텐서 병렬화를 쓰려면 커널 주입을 지원하지 않는 모델에 대해 injection_policy를 직접 넘겨야 했어요. 그런데 이제는 커널 주입을 끄고 injection policy를 주지 않으면, HuggingFace 모델이 기본으로 자동 텐서 병렬화(AutoTP)를 지원해요. 덕분에 커널 주입으로는 아직 지원되지 않는 모델도 injection policy를 하나하나 만들 필요 없이 성능을 끌어올릴 수 있게 됐죠. 이 문서는 추론용 AutoTP를 다루고, 학습용 텐서 병렬화는 별도 문서를 참고하면 돼요.

출처: DeepSpeed 공식 문서 — Automatic Tensor Parallelism for HuggingFace Models

AutoTP로 달라진 점

AutoTP를 쓰면 injection policy가 필요 없어요. 런타임에 자동으로 결정돼서 적용되거든요. T5로 비교하는 가장 간단한 형태를 볼게요.

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

예전 방식은 이렇게 인코더/디코더 레이어의 두 특정 레이어(어텐션 출력 GeMM, 레이어 출력 GeMM)를 가리키는 policy를 넘겨야 했어요. 이 두 부분이 GPU 간 all-reduce 통신으로 부분 결과를 합쳐야 하는 지점이라서요.

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

자동 텐서 병렬화를 쓰면 지원 모델이라면 이 policy 작성이 통째로 사라져요.

성능 비교 (T5 11B, V100 32GB)

실제 성능을 보면 텐서 병렬화가 메모리와 처리량을 동시에 개선하는 걸 확인할 수 있어요. GPU당 메모리 할당이 줄고, 최대 배치 크기와 GPU당 최대 처리량이 함께 늘어나요.

| 테스트 | GPU당 메모리 | 최대 배치 | GPU당 최대 처리량 | | No TP / 1 GPU | 21.06 GB | 64 | 9.29 TFLOPS | | 2 GPU TP | 10.56 GB | 320 | 13.04 TFLOPS | | 4 GPU TP | 5.31 GB | 768 | 14.04 TFLOPS |

OPT 13B도 마찬가지로, 4 GPU TP에서 GPU당 메모리가 23.94 GB에서 6.36 GB로 줄고 처리량도 1.65 → 4.90 TFlops로 늘어나는 걸 확인할 수 있어요.

지원/미지원 모델

AutoTP는 다음 모델 계열에서 검증됐어요(그 외 모델은 될 수도 있지만 아직 테스트되지 않음): albert, arctic, baichuan, bert, bigbird_pegasus, bloom, camembert, chatglm2, chatglm3, codegen, codellama, deberta_v2, electra, ernie, esm, falcon, glm, gpt-j, gpt-neo, gpt-neox, longt5, luke, llama, llama2, m2m_100, marian, mistral, mixtral, mpt, mvp, nezha, openai, opt, pegasus, perceiver, phi, plbart, qwen, qwen2, qwen2-moe, qwen2.5, qwen3, reformer, roberta, roformer, splinter, starcode, t5, xglm, xlm_roberta, yoso, yuan.

반면 AutoTP에 지원되지 않는 모델도 있어요 (deberta, flaubert, fsmt, gpt2, led, longformer, xlm, xlnet). 이들은 커널 주입(Bloom 등) 등 다른 DeepSpeed 기능과는 호환될 수 있으니 참고하세요.

자동 텐서 병렬화(AutoTP)로 HuggingFace 모델 추론 가속하기