使用Kubernetes跑Stable-diffusion分布式训练作业，流程示意：

一、准备环境

1. GPU集群

创建带GPU的K8s集群，安装GPU插件（GPU卡调度），RoCE网卡插件（RoCE网卡调度），Volcano插件（容器组调度）。

    确认集群可用：启动一个Pod，使用 kubectl exec 进去容器，确保GPU卡都正常（nvidia-smi命令OK）。

2. 共享存储

因为是K8s集群，所以需要创建对应的PV+PVC。并启动pod验证容器可以正常挂载该共享存储。

二、准备模型

因为是增量训练，所以需要下载基础模型。这里选择：stable-diffusion-v1-4 为例

https://huggingface.co/CompVis/stable-diffusion-v1-4

把这个目录下载下来（文件有点大）。

然后放在比如 /home/stable-diffusion-v1-4 这个目录。

三、准备代码

这里以text2image为例（旁边菜单就是LoRA，这个貌似更火），参考

https://huggingface.co/docs/diffusers/v0.20.0/en/training/text2image

代码：

https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py

如下：

export MODEL_NAME="/home/stable-diffusion-v1-4"
export dataset_name="/home/pokemon_dataset"

accelerate launch --config_file vjacc.yaml train_text_to_image.py \
 --pretrained_model_name_or_path=$MODEL_NAME \
 --dataset_name=$dataset_name \
 --use_ema \
 --resolution=1024 --center_crop --random_flip \
 --train_batch_size=3 \
 --dataloader_num_workers=8 \
 --mixed_precision="fp16" \
 --max_train_steps=15000 \
 --learning_rate=1e-05 \
 --max_grad_norm=1 \
 --lr_scheduler="constant" --lr_warmup_steps=0 \
 --output_dir="sd-pokemon-model"

把这段脚本写入 run.sh 文件（记得加 chmod +x 可执行权限）

放在如 /home/run.sh 这个地方。

四、准备分布式配置文件

上一节写到，训练启动是利用 accelerate 工具（一个huggingface搞的东西）

https://github.com/huggingface/accelerate

它默认的单机训练的，要配置多机训练，需要给 accelerate 一个配置文件，告诉它，多机多卡的信息（通过 --config_file 指定）。

参见：https://huggingface.co/docs/accelerate/basic_tutorials/launch#custom-configurations

如我们这里：

accelerate launch --config_file vjacc.yaml train_text_to_image.py

所以这里，需要设置 vjacc.yaml 配置文件的内容：

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: MULTI_GPU
downcast_bf16: 'no'
gpu_ids: all
machine_rank: 0
main_process_ip: vj-sd-worker-0.vj-sd
main_process_port: 10086
main_training_function: main
mixed_precision: 'no'
num_machines: 2
num_processes: 16
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

 注意其中的几个参数：

machine_rank：分布式训练兄弟中，自己的“序号”。

main_process_ip：老大的IP地址。这里填“域名”就行。（注，域名是volcano帮忙自动生成的）

num_machines：总兄弟数量。

num_processes：总兄弟数量 x 8卡

这几个参数，可以提前写好。

也可以根据env环境变量，来获得（Volcano会为每个容器自动注入env，来告诉你，在分布式训练中，你的序号是几，总共兄弟有多少个）。

五、准备训练样本

这里用“宝可梦”为例，

https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions

把这些文件下载下来。

Ps，找训练同事帮忙，说为了避免训练中去连huggingface外网，推荐将parquet文件格式，转成csv格式。目录如下：

/home/pokemon_dataset/train

其中 metadata.csv 内容：

没有空格，逗号分隔。

六、准备新镜像

因为pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel镜像跑 train_text_to_image.py 训练，需要很多依赖包

pip install diffusers==0.20.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install transformers -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install datasets -i https://pypi.tuna.tsinghua.edu.cn/simple

每次调试都需要提前安装，很麻烦。

所以这里可以打个新的Docker镜像。

创建Dockerfile文件，

vi Dockerfile

内容如下

FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
USER root
RUN pip install diffusers==0.20.0 -i https://pypi.tuna.tsinghua.edu.cn/simple; \
pip install accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple; \
pip install transformers -i https://pypi.tuna.tsinghua.edu.cn/simple; \
pip install datasets -i https://pypi.tuna.tsinghua.edu.cn/simple

然后执行：

docker build -t pytorch/pytorch:tsj .

可以生成一个包含依赖库的新的镜像了。

六、启动分布式训练容器组

我们使用volcano作业来启动分布式训练，所以需要一个 volcano的yaml，如下：

apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata:
  name: vj-sd
spec:
  minAvailable: 2
  schedulerName: volcano
  plugins:
    env: []
    svc: []
  tasks:
    - replicas: 2
      name: worker
      policies:
        - action: TerminateJob
          event: PodFailed
        - action: CompleteJob
          event: TaskCompleted
        - action: TerminateJob
          event: PodEvicted
      template:
        spec:
          imagePullSecrets:
            - name: default-secret
          containers:
            - command: ['bash', '-c', 'sleep 3600'] # 填入启动训练的命令
              image: pytorch/pytorch:tsj #镜像地址
              name: pytorch
              resources:
                requests:
                  cpu: '80'
                  mellanox.com/InfiniBand: '1'
                  memory: 800G
                  nvidia.com/gpu: '8'
                limits:
                  cpu: '80'
                  mellanox.com/InfiniBand: '1'
                  memory: 800G
                  nvidia.com/gpu: '8'
              volumeMounts:
                - mountPath: /dev/shm
                  name: shm
                - mountPath: /home
                  name: home
              env:
                - name: NCCL_IB_TC
                  value: "128"
                - name: NCCL_IB_HCA
                  value: "mlx5"
                - name: NCCL_IB_GID_INDEX
                  value: "3"
          restartPolicy: Never
          volumes:
            - name: shm
              emptyDir:
                medium: Memory
                sizeLimit: 8G
            - name: home
              hostPath:
                path: /home
                type: Directory
          hostNetwork: true #由于Ant1用的RoCE网卡，需要将容器设置为主机网络才能进行RDMA通信
          dnsPolicy: ClusterFirstWithHostNet #打开主机网络后需要增加该配置才能解析k8s svc

都是K8s标准字段，一般都能理解。特别注意PVC挂载（上述为了方便，直接挂载Hostpath。 但是节点多了以后，每台节点都得复制，肯定“共享存储”挂载方式更简单。）

假设上面这个volcano作业文件名叫 sd.yaml

启动分布式作业：

kubectl apply -f sd.yaml

这样就可以启动一组容器。 可以使用

kubectl get vj

查看，也可以使用

kubectl get pod

查看容器是否都启动正常。（不正常，解决之，属于k8s基本功能）

七、调试

因为我们上面，容器的启动命令是“sleep 600”，所以容器启动后什么也没干。需要我们进入容器进行触发执行训练作业。

kubectl exec -it 容器名 -- /bin/bash
#然后
cd /home
./run.sh

因为是分布式作业，兄弟没启动的话，这边会等待，只有大家都启动，才会开始训练。

所以分别进入各个容器，执行如上命令。

“老大”会保存新的模型：

日志可以在容器里面看，如：

./run.sh 2>&1 | tee 1.log

这样

或者

kubectl logs 容器名

在容器外面看。

八、自动化

每个容器手动启动进程，比较繁琐。后续自动化思路：

1. Accelerate多机配置文件，自动更新
2. Volcano作业启动后，自动拉起训练进程

Volcano会给每个容器都注入env，告诉你在分布式训练兄弟中的“序号”

所以，可以利用这个信息，在 run.sh 前面写的脚本，去更新 accelerate 配置文件。

然后，修改volcano-job.yaml 的 cmd启动命令，更新为直接执行run.sh。