Stable diffusion inference 多卡并行
stable diffusion 推理过程 多卡并行 注意事项
- 以SDXL为例，指定GPU，添加device_map参数信息
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上，否则会报如下错误
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自动拆分 （本人未验证过，参考下面第一个链接中的提问，仅作记录）
from accelerate import load_checkpoint_and_dispatchmodel = load_checkpoint_and_dispatch(model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block'])
本文参考过如下链接，然后血泪教训就是，已有参考信息可以先尝试别人的方案，不要一开始就按照自己的理解瞎冲，否则原理理解不对会浪费很多时间：
- https://blog.csdn.net/weixin_43457608/article/details/136071789
- https://blog.csdn.net/CSTGYinZong/article/details/137506083
- https://huggingface.co/docs/diffusers/training/distributed_inference