HuggingFace模型的自动张量并行

目录

介绍

本教程展示了用于推理的新自动张量并行功能。以前，用户需要向DeepSpeed提供注入策略以启用张量并行。现在，只要未启用内核注入且未提供注入策略，DeepSpeed默认支持HuggingFace模型的自动张量并行。这使得我们的用户能够提高当前不支持通过内核注入的模型的性能，而无需提供注入策略。以下是新方法的示例：

# ---------------------------------------
# New automatic tensor parallelism method
# ---------------------------------------
import os
import torch
import transformers
import deepspeed
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
# create the model pipeline
pipe = transformers.pipeline(task="text2text-generation", model="google/t5-v1_1-small", device=local_rank)
# Initialize the DeepSpeed-Inference engine
pipe.model = deepspeed.init_inference(
    pipe.model,
    mp_size=world_size,
    dtype=torch.float
)
output = pipe('Input String')

以前，对于不支持内核注入的模型，要仅使用张量并行性运行推理，您可以传递一个注入策略，该策略显示了Transformer编码器/解码器层上的两个特定线性层：1）注意力输出GeMM和2）层输出GeMM。我们需要这些层部分来添加GPU之间所需的全归约通信，以跨模型并行等级合并部分结果。下面，我们展示了这种先前方法的示例：

# ----------------------------------
# Previous tensor parallelism method
# ----------------------------------
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"))
# create the model pipeline
pipe = transformers.pipeline(task="text2text-generation", model="google/t5-v1_1-small", device=local_rank)
# Initialize the DeepSpeed-Inference engine
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')

通过自动张量并行，我们不需要为支持的模型提供注入策略。注入策略将在运行时确定并自动应用。

示例脚本

我们可以通过使用推理测试套件观察到自动张量并行带来的性能提升。该脚本用于测试文本生成模型，并包括每个令牌的延迟、带宽、吞吐量和内存检查以进行比较。更多信息请参阅README。

启动

使用以下命令运行，无需DeepSpeed和无需张量并行。设置test_performance标志以收集性能数据：

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 推理性能比较

以下结果是使用V100 SXM2 32GB GPU收集的。

延迟

T5 Latency Graph

吞吐量

T5 Throughput Graph

内存

测试 每个GPU分配的内存 最大批量大小 每个GPU的最大吞吐量
无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 推理性能比较

以下结果是使用V100 SXM2 32GB GPU收集的。

OPT Throughput Graph

测试 每个GPU分配的内存 最大批量大小 每个GPU的最大吞吐量
无TP 23.94 GB 2 1.65 TFlops
2 GPU TP 12.23 GB 20 4.61 TFlops
4 GPU TP 6.36 GB 56 4.90 TFlops

支持的模型

以下模型系列已成功通过自动张量并行测试。其他模型可能也能工作，但尚未经过测试。

  • 阿尔伯特
  • 北极
  • 百川
  • bert
  • bigbird_pegasus
  • bloom
  • 卡门贝尔
  • chatglm2
  • chatglm3
  • codegen
  • codellama
  • deberta_v2
  • electra
  • ernie
  • esm
  • 猎鹰
  • glm
  • gpt-j
  • gpt-neo
  • gpt-neox
  • longt5
  • 卢克
  • 羊驼
  • llama2
  • m2m_100
  • 玛丽安
  • mistral
  • mixtral
  • mpt
  • mvp
  • 哪吒
  • openai
  • 选项
  • pegasus
  • 感知器
  • phi
  • plbart
  • qwen
  • qwen2
  • qwen2-moe
  • 改革者
  • roberta
  • roformer
  • splinter
  • starcode
  • t5
  • xglm
  • xlm_roberta
  • yoso
  • 元

不支持的模型

以下模型目前不支持自动张量并行。它们可能仍然兼容其他DeepSpeed功能（例如，Bloom的内核注入）：

  • deberta
  • flaubert
  • fsmt
  • gpt2
  • led
  • 长形模型
  • xlm
  • xlnet

更新: