diffusers单机多卡推理（全网首发）

半度、

已于 2024-02-08 10:57:14 修改

阅读量3.6k

点赞数 27

文章标签： AI作画

于 2024-02-08 09:44:25 首次发布

本文链接：https://blog.csdn.net/weixin_43457608/article/details/136071789

版权

起因

博主在部署InstantID项目时，显存不够，想要将模型分散在多张卡上。
翻到这篇发现是分布式推理，博主一直以为这个可以达到我想要的效果，但是效果是多线程并行推理，并不能将一个模型切片在多个GPU上。
Distributed Inference with 🤗 Accelerate
好运的是在Accelerate说明文档中我找到了下面的文章，使用init_empty_weights在加载模型时不加载权重，load_checkpoint_and_dispatch函数分发在不同机器上。
Handling big models for inference
在这里插入图片描述

经过

找打上述方法如获至宝，然后去尝试，直接改为StableDiffusionXLPipeline模型路径，不出意外失败。原因：StableDiffusionXLPipeline是多个模型组合而成，上述方法只加载unet模型是没问题的。当我已经做好准备重写StableDiffusionXLPipeline去一个个实现，突然发现from_pretrained中存在device_map参数。幸福来的太突然了，直接官方做好集成了，device_map=‘auto’。果然不出意外还是出意外了。RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! 我的理解是同一个模型不在一张卡上，auto应该是按层数拆的，但是一个transformer它是一整个块，被拆开会有问题（我是这么理解的，有懂的大神请指出谬误）。no_split_module_classes这个参数在指定单个模型时存在，但是StableDiffusionXLPipeline没有集成。
device_map 看帮助文档可以自己指定哪一层再哪个gpu，没有指定的层会报错如下异常：
在这里插入图片描述
通过异常我们再去指定层，直到没有异常，完美！但是问题来了，层数有点多的离谱不。果断放弃，这么干我宁愿改源码。改源码过程中我又发现如下代码，只用指定前缀就可以统一分配gpu！！！
请添加图片描述

结果

from diffusers import StableDiffusionXLPipeline
from accelerate import init_empty_weights,load_checkpoint_and_dispatch,load_checkpoint_in_model,dispatch_model
import torch

model_path = "/workspace/InstantID/models/YamerMIX_v8"
# with init_empty_weights():
device_map = {
    'add_embedding': 1,
    'decoder': 1,
    'encoder': 1,
    'conv_in': 1,
    'conv_out': 1,
    'post_quant_conv': 1,
    'text_model': 1,
    'conv_norm_out': 1,
    'quant_conv': 1,
    'time_embedding': 1,
    'text_projection': 1,
    'up_blocks': 2,
    'mid_block': 2,
    'down_blocks': 2,
}
model = StableDiffusionXLPipeline.from_pretrained(model_path, torch_dtype=torch.float16, device_map=device_map)
prompt = "A cinematic shot of a baby racoon wearing an intricate italian priest robe."
image = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]