如何实现Stable Diffusion的分布式部署

Stable Diffusion分布式部署实战指南
一、方案总览与适用场景

• 单机多GPU：在一台服务器上用多张GPU并行推理，适合显存受限但希望提升吞吐的场景。常见做法是基于PyTorch分布式（DDP/FSDP）或工具链（如xformers、accelerate）实现数据并行/模型并行。优点是部署简单、通信开销低；局限是受单机GPU数量与显存上限约束。
• 多机多GPU：跨多台服务器横向扩展，适合高并发与超大分辨率/批量生成。常见做法是基于NCCL的主从（rank=0）架构，结合PyTorch Lightning/DDP或accelerate进行多进程/多节点编排。优点是线性扩展能力强；局限是对10Gbps+网络与共享存储有要求。
• 微服务化生产集群：将API网关、认证、任务调度、推理Worker、对象存储解耦，容器化后在Kubernetes上编排，配合Prometheus+Grafana监控与自动扩缩容，适合企业级高可用与弹性需求。

二、单机多GPU部署步骤（以WebUI Forge为例）

• 环境准备：建议使用Python 3.10、PyTorch（CUDA 11.8）、xformers、accelerate等；创建虚拟环境并安装依赖。
• 启动参数示例（2卡）：
• 启用多GPU与设备选择：python launch.py --multi-gpu --device-id 0,1
• 显存优化：按需添加 --cpu-offload、--cuda-malloc
• 关键要点：
• 优先开启xformers与FP16推理，显著降低显存占用并提升吞吐。
• 若出现OOM，降低分辨率/步数/批量，或启用CPU卸载与显存优化策略。
• 多卡并行以数据并行为主；如需更大分辨率可探索模型并行/流水线并行。

三、多机多GPU部署步骤（主从DDP + NCCL）

• 基础环境：
• 至少2台带NVIDIA GPU的服务器，节点间10Gbps+以太网或InfiniBand；每节点内存≥32GB、GPU显存≥24GB。
• 所有节点安装相同驱动/CUDA/Python依赖，并配置SSH免密以便启动脚本分发。
• 启动配置：
• 设定环境变量：MASTER_ADDR（主节点主机名/IP）、MASTER_PORT（如12355）、WORLD_SIZE（总GPU数）、RANK（当前节点rank）。
• 使用PyTorch DDP或accelerate/PyTorch Lightning启动多进程：
• 示例（accelerate）：accelerate launch --config_file ddp.yaml your_infer.py
• 示例（原生DDP）：torchrun --nproc_per_node=2 --nnodes=2 --node_rank=$RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT your_infer.py
• 运行与验证：
• 主节点（rank=0）负责分发与汇总，工作节点执行推理；通过日志与NCCL通信状态确认集群健康。
• 性能参考：在4×V100环境下，512×512图像生成由8.2s/张降至2.3s/张（约3.57×加速）；1024×1024由单机OOM转为5.7s/张（批量32张由268s降至78s）。

四、Kubernetes生产级部署（推理/训练皆可）

• 集群准备：
• 创建带GPU的K8s集群，安装NVIDIA GPU插件（设备调度）、RoCE网卡插件（RDMA）、Volcano（批调度）。
• 准备PV/PVC共享存储（模型与数据集），验证容器内nvidia-smi可用。
• 推理服务封装：
• 基于Docker封装Stable Diffusion推理服务（Flask/FastAPI），暴露5000端口；镜像内预置模型权重或使用PVC挂载。
• 示例Dockerfile要点：FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime；安装transformers/diffusers/flask；COPY模型权重；CMD启动服务。
• 训练作业编排：
• 使用accelerate或分布式启动脚本提交作业，示例：
• accelerate launch --config_file vjacc.yaml train_text_to_image.py --pretrained_model_name_or_path=$MODEL_NAME --dataset_name=$DATASET --mixed_precision=fp16 ...
• 通过K8s Job/Operator管理生命周期，结合Prometheus+Grafana监控GPU利用率、显存、队列与延迟。

五、共享存储与异构集群实践

• 共享存储：
• 多机部署需共享模型、配置、输出目录，推荐NFS/Ceph等网络存储；在Docker Compose中通过卷映射（如./data:/data、./output:/output）实现跨节点一致视图，避免重复下载与版本不一致。
• 异构集群（CPU/GPU混合）：
• 在同一集群内同时运行GPU服务（高性能推理）与CPU服务（预处理/后处理/兼容测试），通过Docker Swarm/Kubernetes进行智能调度与弹性伸缩，兼顾成本与性能。