Stable diffusion inference 多卡并行

stable diffusion 推理过程多卡并行注意事项

以SDXL为例，指定GPU，添加device_map参数信息

python 复制代码

device_map = {
    'add_embedding': 1,
    'decoder': 1,
    'encoder': 1,
    'conv_in': 1,
    'conv_out': 1,
    'post_quant_conv': 1,
    'text_model': 6,
    'conv_norm_out': 1,
    'quant_conv': 1,
    'time_embedding': 1,
    'text_projection': 6,
    'up_blocks': 1,
    'mid_block': 6,
    'down_blocks': 6,
}
from diffusers import StableDiffusionXLPipeline
model = StableDiffusionXLPipeline.from_pretrained(model_path, torch_dtype=torch.float16, device_map=device_map)

需要注意的是，diffusers的版本太高时，device_map的参数类型有要求，0.28.0时候，只能输入str，而这个str有几个指定的选项，然而，我在尝试使用"auto"时，提示只能输入"balanced"

"auto " 和 "balanced" 将会在所有的GPU上平衡切分模型。主要是有可能发现更高效的分配策略。"balanced" 参数的功能则保持稳定。
"balanced_low_0" 会在除了第一个GPU上的其它GPU上平衡划分模型，并且在第一个 GPU 上占据较少资源。这个选项符合需要在第一个 GPU 上进行额外操作的需求，例如需要在第一个 GPU 执行 generate 函数（迭代过程）。
"sequential" 按照GPU的顺序分配模型分片，从 GPU 0 开始，直到最后的 GPU（那么最后的 GPU 往往不会被占满，和 "balanced_low_0" 的区别就是第一个还是最后一个，以及非均衡填充），但是我在实际使用当中GPU 0 会直接爆显存了

diffusers版本低于0.28.0时候才支持其他选项，同时支持自定义device_map，如0.27.2是可以的

同时，有另一个非常值得注意的事情，上述模型分片要注意不能随意分片，至少同一个模型需要在一个GPU上，否则会报如下错误

powershell 复制代码

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:6 and cuda:1! (when checking argument for argument mat2 in method wrapper_mm)

然而不同模块占用显存是不一样的，因此可以多试试多个方案

同时，持续更新问题

问题1：怎么根据no_split_module_classes创建device_map呢

answer: 不用写device_map自动拆分（本人未验证过，参考下面第一个链接中的提问，仅作记录）

python 复制代码

from accelerate import load_checkpoint_and_dispatch
    model = load_checkpoint_and_dispatch(model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block'])

本文参考过如下链接，然后血泪教训就是，已有参考信息可以先尝试别人的方案，不要一开始就按照自己的理解瞎冲，否则原理理解不对会浪费很多时间：