怎样实现SD多卡并行推理

实现Stable Diffusion（SD）多卡并行推理，核心思路是将不同的推理任务（如不同的Prompt、不同的图片）分配到多个GPU上同时执行，或者将单个模型拆分到不同卡上（张量并行）。对于大多数用户和场景，数据并行（多任务/多请求并行） 是最简单、最实用且效果最好的方式。

以下是几种主流的实现方案，从易到难：

方法一：最简单实用 —— 多进程/多实例（数据并行）

这是最推荐的方式。利用 Python 的 multiprocessing 或者启动多个独立的推理脚本，每个进程绑定一张卡。
原理： 每个 GPU 加载一份完整的 SD 模型，各自处理不同的输入。
代码示例（Python Multiprocessing）：

import os
import torch
from diffusers import StableDiffusionPipeline
from multiprocessing import Process
def run_inference_on_gpu(gpu_id, prompt_list):
# 指定当前进程使用的GPU
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
# 加载模型（每个进程会独立加载一份到对应的GPU）
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")
print(f"GPU {gpu_id} 开始推理...")
for prompt in prompt_list:
image = pipe(prompt).images[0]
image.save(f"output_gpu{gpu_id}_{prompt[:10]}.png")
print(f"GPU {gpu_id} 完成.")
if __name__ == "__main__":
prompts = ["a cat", "a dog", "a car", "a tree", "a house", "a bird"]
num_gpus = torch.cuda.device_count()
processes = []
# 简单分配任务，假设有2张卡
chunk_size = len(prompts) // num_gpus
for i in range(num_gpus):
start = i * chunk_size
end = start + chunk_size if i != num_gpus - 1 else len(prompts)
p = Process(target=run_inference_on_gpu, args=(i, prompts[start:end]))
p.start()
processes.append(p)
for p in processes:
p.join()

注意：这种方式显存占用是 显存占用 * 显卡数量，因为每张卡都要存一份模型。

方法二：使用推理服务框架（生产环境推荐）

如果你需要搭建一个稳定的 API 服务，或者处理高并发请求，建议使用成熟的推理框架。这些框架内部已经封装好了多卡调度逻辑。

1. vLLM (主要用于LLM，但也支持部分扩散模型)

虽然 vLLM 以 LLM 闻名，但其最新版本开始支持扩散模型，利用 PagedAttention 等优化技术，吞吐量极高。

2. TensorRT / TensorRT-LLM

NVIDIA 的高性能推理方案。

• 步骤： 将 SD 模型转换为 TensorRT 引擎（Engine），然后使用 TensorRT 的 Python/C++ API 进行多卡推理。
• 优点： 速度极快。
• 缺点： 转换过程复杂，调试困难。

3. Accelerate (Hugging Face)

Hugging Face 的 accelerate 库提供了简单的 API 来实现模型并行（将模型的不同层放在不同卡上）。
代码示例：

from accelerate import infer_auto_device_map, init_empty_weights
from diffusers import StableDiffusionPipeline
import torch
# 1. 初始化空权重（不占显存）
with init_empty_weights():
model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
# 2. 自动分配设备（例如将 UNet 分到 GPU 0，Text Encoder 分到 GPU 1）
device_map = infer_auto_device_map(
model,
max_memory={0: "8GiB", 1: "8GiB"}, # 限制每张卡的最大显存
no_split_module_classes=["UNet2DConditionModel"] # 某些层不能切分
)
# 3. 加载模型
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
device_map=device_map,
torch_dtype=torch.float16
)
# 推理时，数据会在卡间自动传输
image = pipe("a photo of an astronaut").images[0]

注意：这种方式通常用于单张图的生成，利用多卡显存来容纳超大模型，但推理速度不一定变快（因为数据需要在卡间传输）。

方法三：流水线并行 (Pipeline Parallelism)

将 SD 的推理步骤（如 CLIP 编码 -> UNet 去噪 -> VAE 解码）分配到不同的 GPU 上。

• GPU 0: 负责 Text Encoder (CLIP)。
• GPU 1: 负责 UNet (最耗时的部分)。
• GPU 2: 负责 VAE Decoder。

适用场景： 当你需要生成连续不断的视频帧或者超高分辨率图片，且单卡显存不足时。这通常需要手动编写调度代码，或者使用像 DeepSpeed 或 Megatron-LM 这样的分布式训练/推理库，配置较为复杂。

方法四：使用 WebUI (Automatic1111 / ComfyUI)

如果你是在本地使用图形界面：

1. ComfyUI (推荐)：

• ComfyUI 原生支持多卡。你可以在节点中指定 device 或者利用它的 Queue 系统配合脚本实现多卡。
• 更常见的是：启动两个不同端口的 ComfyUI 实例，分别绑定不同的 GPU (--cuda-device 0 和 --cuda-device 1)，然后用负载均衡器（如 Nginx）分发请求。

1. Automatic1111:

• 原生不支持多卡并行推理单张图。
• 变通方案：在启动参数中指定 CUDA_VISIBLE_DEVICES。
• 例如：终端1运行 CUDA_VISIBLE_DEVICES=0 python launch.py，终端2运行 CUDA_VISIBLE_DEVICES=1 python launch.py。然后手动在两个网页间切换提交任务。

总结与建议

最通用的代码逻辑（伪代码）：

# 假设有2张卡
prompts = load_prompts()
split_prompts = split_into_chunks(prompts, 2)
run_process_1(gpu=0, prompts=split_prompts[0])
run_process_2(gpu=1, prompts=split_prompts[1])

注意： 模型并行（切分模型）通常是为了解决显存不足的问题，而不是为了加速。对于 SD 这种模型，如果单卡能放下，多卡切分反而会因为 PCIe 数据传输变慢。因此，多任务并行（数据并行）才是加速的王道。