Accelerate 0.24.0文档三：超大模型推理（内存估算、Sharded checkpoints、bitsandbytes量化、分布式推理）

文章目录

- 一、内存估算
- - 1.1 Gradio Demos
  - 1.2 The Command
- 二、使用Accelerate加载超大模型
- - 2.1 模型加载的常规流程
  - 2.2 加载空模型
  - 2.3 分片检查点（Sharded checkpoints）
  - 2.4 示例：使用Accelerate推理GPT2-1.5B
  - 2.5 device_map
- 三、bitsandbytes量化
- - 3.1 环境依赖
  - 3.2 量化示例
  - 3.3 保存和加载 8 位模型
  - 3.4 微调量化模型
- 四、分布式推理
- - 4.1 使用torch.distributed进行分布式推理
  - 4.2 使用Accelerate进行分布式推理

《探索多种方案下 LLM 的预训练性能》

一、内存估算

参考《Understanding how big of a model can fit on your machine》、《Transformer 估算》

目前使用大模型一个非常困难的方面是了解当前显卡的内存可以容纳多大的模型（例如将模型加载到 CUDA 上）。为了帮助解决这个问题，Accelerate提供了accelerate estimate-memory 命令行界面，来进行估算。

1.1 Gradio Demos

打开Gradio Demos页面。输入Model Name or URL就可以进行估算。

在这里插入图片描述

该计算器将告诉您纯粹加载模型而不是执行推理需要多少内存
模型所需的最小内存建议表示为“最大层”的大小，模型的训练大约是其大小的 4 倍（对于 Adam）
在执行推理时，预计会额外增加 20%。未来将进行更多测试，以获得每个模型更准确的基准。
目前此工具支持使用 transformers 和 timm 托管的所有模型
此计算的精确度在实际值的百分之几内，例如，当以全精度加载到 CUDA 上时，加载 bert-base-cased 实际上需要 413.68 MB ，计算器估计是 413.18 MB 。

1.2 The Command

你也可以使用命令行界面来获取结果，例如计算bert-base-cased 的内存占用量：

accelerate estimate-memory bert-base-cased

这将下载 bert-based-cased 的 config.json文件，在 meta 设备上加载模型，并报告它将使用多少空间：

dtype	Largest Layer	Total Size	Training using Adam
float32	84.95 MB	418.18 MB	1.61 GB
float16	42.47 MB	206.59 MB	826.36 MB
int8	21.24 MB	103.29 MB	413.18 MB
int4	10.62 MB	51.65 MB	206.59 MB

默认情况下，它将返回所有支持的数据类型（ int4 到 float32 ），你也可以进行过滤：

accelerate estimate-memory bert-base-cased --dtypes float32 float16

dtype	Largest Layer	Total Size	Training using Adam
float32	84.95 MB	413.18 MB	1.61 GB
float16	42.47 MB	206.59 MB	826.36 MB

如果无法确定来自哪个库，可以传入库的名称：

HuggingFaceM4/idefics-80b-instruct：

accelerate estimate-memory HuggingFaceM4/idefics-80b-instruct --library_name transformers

dtype	Largest Layer	Total Size	Training using Adam
float32	3.02 GB	297.12 GB	1.16 TB
float16	1.51 GB	148.56 GB	594.24 GB
int8	772.52 MB	74.28 GB	297.12 GB
int4	386.26 MB	37.14 GB	148.56 GB

timm/resnet50.a1_in1k：

accelerate estimate-memory timm/resnet50.a1_in1k --library_name timm

dtype	Largest Layer	Total Size	Training using Adam
float32	9.0 MB	97.7 MB	390.78 MB
float16	4.5 MB	48.85 MB	195.39 MB
int8	2.25 MB	24.42 MB	97.7 MB
int4	1.12 MB	12.21 MB	48.85 MB

二、使用Accelerate加载超大模型

本章参考：

《Handling big models for inference》、《Distributed Inference with 🤗 Accelerate》
《使用HuggingFace的Accelerate库加载和运行超大模型》
有关本章用到的各种API和更多超大模型加载内容，可参考《Working with large models》

在这里插入图片描述

2.1 模型加载的常规流程

在 PyTorch 中加载预训练模型时，通常的工作流程如下所示：

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

简而言之，这些步骤是：

使用随机初始化的权重创建模型
从磁盘加载模型权重（通常称为state_dict）
将这些权重加载到模型中
将模型加载到相应设备上（比如GPU）进行模型推理

如果模型不是很大的话，使用上面的代码加载权重是可以的。但当我们处理大型模型时，此工作流程有一些明显的局限性。

在第一步中，需要在RAM中加载一个完整版本的模型，并花费一些时间对权重进行随机初始化（在第三步将被丢弃）。
在第二步中，需要再次在RAM中加载一个完整版本的模型，其中包含预训练的权重。

对于包含60亿参数的模型，这意味每一步都需要需要24GB的RAM（float32每个占四字节），总共需要48GB（其中一半用于在FP16中加载模型）。如果加载更大的模型，比如BLOOM或OPT-176B（1760亿个参数），需要大概1.4TB的内存空间。这样内存肯定是不够的，因此我们需要做一些技巧性的改变。

2.2 加载空模型

PyTorch 1.9版本提供了一个叫做meta的新设备。使用该设备，可以创建非常大的张量，而无需考虑内存或显存的大小。同样的，也可以方便创建模型，而无需加载权重。例如，下面的代码如果在Colab上或kaggle的kernel运行，会报错OOM，因为默认使用精度为32位的情况下，创建下述张量需要40G内存。

import torch

large_tensor = torch.randn(100000, 100000)

而如果与meta设备来创建，即可只定义张量形状，而不消耗内存。

import torch

large_tensor = torch.randn(100000, 100000, device="meta")
large_tensor

# 张量只有形状，并没有数据（不占据内存或显存）。
tensor(..., device='meta', size=(100000, 100000))

Accelerate 引入的第一个帮助处理大型模型的工具是上下文管理器 init_empty_weights()，可以使用meta初始化一个空模型（只有shape，没有数据）。

from accelerate import init_empty_weights

with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

这可以初始化一个略大于100亿参数的空模型，在init_empty_weights()下的初始化过程中，每当创建一个参数，它会立即被移到meta设备上。定义好上述模型后，可以喂输入，然后得到一个meta设备的输出张量（同样，只有形状，没有数据，其实就是进行了shape计算）。

2.3 分片检查点（Sharded checkpoints）

当模型过大，无法将其整体加载到内存中时，仍可通过分片方式加载，尤其在有一个或多个GPU的情况下，因为GPU提供了更多内存。
检查点分片：将检查点分割成多个较小的文件，称为检查点分片。Accelerate库能够处理这种检查点分片，前提是遵循特定的格式：

检查点应该存储在一个文件夹中。
文件夹中包含多个具有部分状态字典的文件，每个文件对应模型的不同部分。
有一个JSON格式的索引文件，该文件包含一个字典，将参数名称映射到包含其权重的文件。

使用 save_model() 可以轻松地将模型分片：

accelerator.wait_for_everyone()  # 确保所有进程训练完成
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)

得到的结果类似于：

first_state_dict.bin
index.json
second_state_dict.bin

其中index.json是以下文件：

{
  "linear1.weight": "first_state_dict.bin",
  "linear1.bias": "first_state_dict.bin",
  "linear2.weight": "second_state_dict.bin",
  "linear2.bias": "second_state_dict.bin"
}

加载时，你可以使用 load_checkpoint_in_model() 函数将其加载在特定设备上。

load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"": device})

也可以使用load_checkpoint_and_dispatch() 函数在空模型中加载完整检查点或分片检查点，它还会自动在您可用的设备（GPU、CPU RAM）上分配这些权重。完整的模型分片推理过程见此YouTube视频。

load_checkpoint_and_dispatch函数常用参数为model、checkpoint、device_map、max_memory、no_split_module_classes，后两个参数将在后面讲到。

2.4 示例：使用Accelerate推理GPT2-1.5B

使用 minGPT库的默认配置初始化模型

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    model = GPT(model_config)

下载模型权重并加载

pip install huggingface_hub

from huggingface_hub import snapshot_download
from accelerate import load_checkpoint_and_dispatch

checkpoint = "marcsun13/gpt2-xl-linear-sharded"
weights_location = snapshot_download(repo_id=checkpoint)
model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
)

上述代码中：

使用load_checkpoint_and_dispatch()加载检查点到空模型。
设置device_map="auto"会让Accelerate自动判断把模型的每个层放在哪里:
- 优先尽量放在GPU上
- 如果GPU空间不够，将剩余层放在CPU RAM上
- 如果CPU RAM也不够，将剩余层存储在硬盘上，以memory-mapped tensors的形式。
通过no_split_module_classes参数可以指定某些层不被分割（比如包含残差连接的Block等模块）。

memory-mapped tensors简称 mmap tensors，是PyTorch提供的一种特殊的tensors,它允许将数据存储在磁盘文件中，而不占用宝贵的RAM内存，CPU可以直接对磁盘文件中的数据进行读写操作，就像操作RAM中的tensors一样。mmap tensors既可以享受性能接近RAM的缓存系统带来的读写速度，同时也不会耗尽宝贵的RAM空间，从而使Accelerate支持超大模型的训练

推理
现在我们的模型位于多个设备上，也许还有在硬盘上的部分，但它仍然可以用作常规 PyTorch 模型进行推理：

from mingpt.bpe import BPETokenizer
tokenizer = BPETokenizer()
inputs = tokenizer("Hello, my name is").to(0)

outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
tokenizer.decode(outputs.cpu().squeeze())

在幕后， Accelerate 添加了钩子（hook）到模型，以便：

在每一层，输入被放在正确的设备上，因此，即使您的模型分布在多个GPU上，它也能正常工作。
对于卸载到CPU上的权重，在前向传播之前将它们放在GPU上，然后在之后清理
对于卸载到硬盘上的权重，在前向传播之前将其加载到RAM中，然后放在GPU上，之后再清理

这样，即使您的模型在GPU或CPU上装不下，也可以进行推理！同时需要注意的是，Accelerate 通过 hook 仅支持推理而不支持训练，这是因为：

内存限制：在训练过程中,我们需要在每一层保留激活值来计算梯度,这会大大增加内存需求。但在推理时我们可以在每一层后立即释放激活,所以内存需求较小。
计算复杂性：训练需要频繁的反向传播，权重需要频繁更新。反向传播涉及到复杂的计算图跟踪，这在分布式环境下变得非常困难，特别是当layer分布在不同设备时,实现反向传播会变得极具挑战性，高效的分布式权重更新也非常困难。
精度要求：训练中需要精确计算梯度的值，但在这种分布式环境下很难保证，最终会影响模型的训练质量。相比之下，在推理过程中，不需要计算梯度。

当我们谈到深度学习中的 “hook” 时，我们指的是一种机制，允许在神经网络的不同部分插入自定义代码，以便在训练或推断的过程中执行一些额外的操作。这些操作可能包括记录中间激活值、梯度、权重等，或者进行某种修改，以适应特定需求。

hook一般包括Forward Hook（前向钩子）和Backward Hook（反向钩子）。上面的讲解中，hooks 的作用是确保输入被正确放置在合适的设备上。下面是一段关于hook的伪代码：

def forward_hook(module, input, output):
    # 在前向传播中执行的自定义操作   
def backward_hook(module, grad_input, grad_output):
    # 在反向传播中执行的自定义操作

# 注册前向和反向 hook
hook_handle = model.fc.register_forward_hook(forward_hook)
hook_handle_backward = model.fc.register_backward_hook(backward_hook)

# 运行前向传播和反向传播
output = model(input_data)
output.backward(torch.randn_like(output))

# 移除 hook
hook_handle.remove()
hook_handle_backward.remove()

2.5 device_map

查看和设置device_map
您可以通过访问模型的 hf_device_map 属性来查看 Accelerate 选择的 device_map（一个字典，包含模型模块、权重以及对应的设备），比如对于上述GPT2-1.5B模型：

model.hf_device_map

{'transformer.wte': 0,
 'transformer.wpe': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 ...
 'transformer.h.21': 0, 
 'transformer.h.22': 1, 
 'transformer.h.23': 1, 
 'transformer.h.24': 1,
 ...
 'transformer.h.47': 1, 
 'transformer.ln_f': 1, 
 'lm_head': 1}

你也可以自定义设备映射，指定要使用的 GPU 设备（数字）、 “cpu” 或 “disk” 并将其传入：

device_map = {
    "transformer.wte": "cpu",
    "transformer.wpe": 0,
    "transformer.drop": "cpu",
    "transformer.h.0": "disk"
}

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map=device_map)

device_map选项
device_map有四个可选参数。当GPU不足以容纳整个模型时，所有选项都会产生相同的结果（即优先加载在GPU上，而后分别是 CPU和磁盘）。当 GPU显存大于模型大小时，每个选项会有如下差异：
- "auto" 或 "balanced":：Accelerate将会根据所有GPU均衡切分权重，尽量均匀的切分到各个GPU上；
- "balanced_low_0":：在第一个GPU上（序号为0）会尽量节省显存，其它个GPU均匀分割权重。这种模式可以有效节省第一个GPU的显存用于模型生成等计算操作（generate函数）；
- "sequential"：Accelerate按照GPU的顺序占用显存，因此排序靠后的GPU显存占用会少一些。
"auto" 和 "balanced" 目前会产生相同的结果，但如果我们找到更有意义的策略， "auto" 的行为将来可能会发生变化，而 "balanced" 将保持不变。
accelerate.infer_auto_device_map：此函数用于为给定的模型生成设备映射，优先使用 GPU，然后是 CPU，最后是硬盘。因为所有计算都是通过分析模型参数的大小和数据类型来完成的，所以可以是meta上的空模型。以下是具体参数：
- model（torch.nn.Module）：要分析的模型。
- max_memory（可选）：设备标识符到最大内存的字典。如果未设置，将默认为可用的最大内存。
- no_split_module_classes（可选）：不应在设备之间分割的层类名称列表（例如，任何具有残差连接的层）。
- dtype（可选）：如果提供，加载权重时将其转换为该类型。例如：
```
from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("facebook/opt-13b")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)
device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")
```
- special_dtypes（可选）：如果提供，考虑某些特定权重的特殊数据类型（将覆盖作为所有权重默认值的 dtype）。
- verbose（可选，默认为 False）：是否在函数构建设备映射时提供调试语句。
max_memory 参数
您可以使用 max_memory 参数来限制每个 GPU 和CPU上使用的内存，赋予GPU应该传递标识符（例如 0,1），内存值可以是整数（以字节为单位），也可以是表示数字及其单位的字符串，例如 “10GiB” 或 “10GB” 。
需要注意的是，当 PyTorch 中发生第一次分配时，它会加载 CUDA 内核，该内核大约需要 1-2GB 内存，具体取决于 GPU。因此，可用内存总是小于 GPU 的实际大小。要查看实际使用了多少内存，请执行 torch.ones(1).cuda() 并查看内存使用情况。
- 示例一：GPU 内存不超过10GiB，CPU内存不超过30GiB：
```
from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
```
- 对于一些生成任务的模型，如果想您有许多 GPU且想使用更大的batch size进行推理，第一个GPU应该分配较少的内存。例如，在 8x80 A100 上使用 BLOOM-176B，接近理想的分配为：
```
max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}
```

三、bitsandbytes量化

参考：《Quantization（Accelerate）》、《大规模 Transformer 模型 8 比特矩阵乘简介 - 基于 Hugging Face Transformers、Accelerate 以及 bitsandbytes》

Accelerate库集成了bitsandbytes 量化功能，几行代码就可以实现4位量化和8位量化。要了解有关 bitsandbytes 量化工作原理的更多信息，请查看8 位量化和 4 位量化的博客文章。transformers中也集成了bitsandbytes 量化功能，相关内容可查看量化文档，或者我的另一篇博客《Hugging Face高性能技术五：Transformer高效推断（bitsandbytes、FlashAttention、 BetterTransformer）》

3.1 环境依赖

使用前，安装相关依赖：

pip install bitsandbytes
pip install git+https://github.com/huggingface/accelerate.git

3.2 量化示例

安装 minGPT 和 huggingface_hub 以运行示例：

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/
pip install huggingface_hub

从 minGPT 库中获取 GPT2 模型配置，然后使用 init_empty_weights().初始化一个空模型：

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    empty_model = GPT(model_config)

需要获取模型权重的路径。该路径可以是 state_dict 文件（例如“pytorch_model.bin”）或包含分片检查点的文件夹。

from huggingface_hub import snapshot_download
weights_location = snapshot_download(repo_id="marcsun13/gpt2-xl-linear-sharded")

使用 BnbQuantizationConfig 设置量化配置

from accelerate.utils import BnbQuantizationConfig
# 8位量化
bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)
# 4位量化
bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")

使用 load_and_quantize_model()量化所选配置的空模型：

from accelerate.utils import load_and_quantize_model
quantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")

量化操作的具体实现，都集成在bitsandbytes 库的Linear8bitLt 模块中，它是torch.nn.modules 的子类。与 nn.Linear 模块略有不同，其参数属于 bnb.nn.Int8Params 类而不是 nn.Parameter 类。然后我们会调用replace_8bit_linear函数，将所有 nn.Linear 模块替换为 bitsandbytes.nn.Linear8bitLt模块。

def replace_8bit_linear(model, threshold=6.0, module_to_not_convert="lm_head"):
    for name, module in model.named_children():
        if len(list(module.children())) > 0:
            replace_8bit_linear(module, threshold, module_to_not_convert)

        if isinstance(module, nn.Linear) and name != module_to_not_convert:
            with init_empty_weights():
                model._modules[name] = bnb.nn.Linear8bitLt(
                    module.in_features,
                    module.out_features,
                    module.bias is not None,
                    has_fp16_weights=False,
                    threshold=threshold,
                )
    return model

此函数递归地将 meta 设备上初始化的给定模型的所有 nn.Linear 层替换为 Linear8bitLt 模块。这里，必须将 has_fp16_weights 属性设置为 False，以便直接将权重加载为 Int8，并同时加载其量化统计信息。另外我们放弃了对某些模块 (这里是 lm_head) 进行替换，因为我们希望保持输出层的原始精度以获得更精确、更稳定的结果。也就是说， bitsandbytes只会量化transformer结构中除首层之外的全连接层。

3.3 保存和加载 8 位模型

使用accelerate.save_model保存8位模型，至于 4 位模型序列化，目前还不支持。

from accelerate import Accelerator
accelerate = Accelerator()
new_weights_location = "path/to/save_directory"
accelerate.save_model(quantized_model, new_weights_location)

quantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")

如果 GPU 不足以存储整个模型，您可以通过传递自定义 device_map 将某些模块卸载到 cpu/磁盘。对于 8 位量化，所选模块将转换为 8 位精度。以下是一个示例：

device_map = {
    "transformer.wte": 0,
    "transformer.wpe": 0,
    "transformer.drop": 0,
    "transformer.h": "cpu",
    "transformer.ln_f": "disk",
    "lm_head": "disk",
}

完整量化代码可查看colab notebook示例《Accelerate quantization.ipynb》，示例中将GPT2模型量化为4位模型和8位模型并进行推理。

3.4 微调量化模型

参考《Quantize 🤗 Transformers models》

8 位或 4 位量化模型无法执行全量训练。但是，您可以利用参数高效微调方法 (PEFT) 来微调这些模型，详见peft Github示例：

Colab notebook《Finetune-opt-bnb-peft.ipynb》：使用peft库的LoRa功能微调量化后的OPT-6.7b模型
《Finetuning Whisper-large-V2 on Colab using PEFT-Lora + BNB INT8 training + Streaming dataset》
《Finetuning Whisper-large-V2 on Colab using PEFT-Lora + BNB INT8 training》

请注意，device_map=auto 仅用于推理。这是因为推理过程通常不需要进行梯度计算，而且模型参数已经在训练期间被优化，因此在推理时可以更灵活地选择设备。

加载模型进行训练时，不需要显式传递device_map参数。系统会自动将模型加载到GPU上进行训练。如果需要，您可以将设备映射设置为特定设备，例如cuda:0, 0, torch.device('cuda:0')。

四、分布式推理

参考《Distributed Inference with 🤗 Accelerate》

4.1 使用torch.distributed进行分布式推理

分布式推理是一种常见的用例，尤其是自然语言处理 (NLP) 模型。用户通常希望发送多个不同的提示，每个提示发送到不同的 GPU，然后返回结果。下面是一个普通示例：

import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)

使用使用torch.distributed模块进行分布式推理：

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    pipe.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    result = pipe(prompt).images[0]
    result.save(f"result_{rank}.png")

可以看到，我们需要根据进程的rank选择不同的提示进行推理，这种方式需要手动管理每个进程的提示，显得有些繁琐。

4.2 使用Accelerate进行分布式推理

通过🤗 Accelerate，我们可以通过使用Accelerator.split_between_processes()上下文管理器（也存在于PartialState和AcceleratorState中）来简化这个过程。此函数会自动将您传递给它的任何数据（无论是提示，一组张量，先前数据的字典等）在所有进程之间进行分割（可能进行填充），以便您立即使用。让我们使用上下文管理器重写上面的示例：

from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipe(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

然后使用accelerate launch启动代码（已生成配置文件）：

accelerate launch distributed_inference.py

使用的特定配置文件启动：

accelerate launch --config_file my_config.json distributed_inference.py

不使用配置文件进行启动：（会受到一些警告，可以执行 accelerate config default来解决）

accelerate launch --num_processes 2 distributed_inference.py

现在我们不需要写分布式推理的样板代码了，Accelerate会自动分配数据到各个GPU上。如果碰到数据分割不均的情况，比如我们有 3 个提示，但只有 2 个 GPU。在上下文管理器下，第一个 GPU 将接收前两个提示，第二个 GPU 将接收第三个提示，确保所有提示都被拆分并且不需要任何开销。

如果需要对所有 GPU 的结果执行一些操作（例如聚合它们并执行一些后处理），可以在 split_between_processes 中传递 apply_padding=True，以确保提示列表被填充到相同的长度，多余的数据从最后一个样本中获取。这样，所有 GPU 将具有相同数量的提示，然后可以收集结果。

from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat", "a chicken"], apply_padding=True) as prompt:
    result = pipe(prompt).images