k8s+docker+deepSpeed多机多卡微调实验记录

在本地三台服务器上进行了多机多卡微调的基础上，准备使用docker将其容器化，然后使用k8s的StatefulSet进行容器的创建和训练。

1、为什么选择StatefuleSet

容器化的时候，由于docker镜像每次run一个容器IP地址都会变化，所以需要使用StatefulSet提供稳定的网络标识、存储和节点顺序。如果使用Deployment或者Job，Pod可能会因重启而重新分配新的名称（例如bert-train-abc123)，而StatefulSet会确保：bert-train-0、bert-train-1、bert-train-2，Pod名称不会变，保证了RANK解析正确，训练节点编号稳定。

如果想更多的了解StatefulSet，请参考k8s官方网站。

2、下一步开始Dockerfile的编写，如果还没有安装Docker，请参考mac安装docker和配置教程，linux安装docker和配置教程。Dockerfile如下，如果对于容器构建流程不了解，请参考文章：docker容器构建教程 ，Dockerfile如下：

FROM deepspeed/gh-builder:py3.11

# 设置工作目录
WORKDIR /workspace
# 设置环境变量以防止交互式安装
ENV DEBIAN_FRONTEND=noninteractive

# 更新 apt-get 并安装 Python、pip 以及一些基础工具
RUN apt-get update --allow-unauthenticated && apt-get install -y \
    python3 \
    python3-pip \
    python3-dev \
    libopenmpi-dev \
    git \
    tzdata \
    && rm -rf /var/lib/apt/lists/*

# 强制设置时区，避免交互式选择
RUN ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime && \
    dpkg-reconfigure --frontend noninteractive tzdata

# 安装 accelerate 库
RUN pip install --no-cache-dir 'accelerate>=0.26.0'
  
# 设置 NCCL 变量，优化通信
ENV NCCL_DEBUG=INFO

# 安装 Python 依赖
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

# 复制训练代码到容器中
COPY . .

# 使用 Kubernetes Headless Service 解决主机解析
#RUN echo "bert-training-0.bert-service.default.svc.cluster.local slots=1" > /workspace/hostfile && \
    #echo "bert-training-1.bert-service.default.svc.cluster.local slots=1" >> /workspace/hostfile && \
    #echo "bert-training-2.bert-service.default.svc.cluster.local slots=1" >> /workspace/hostfile

# 入口命令，默认使用 DeepSpeed 启动
# 设置 `hostfile`
RUN chmod +x entrypoint.sh

# 使用 bash 来解析环境变量，并启动 DeepSpeed
CMD ["bash", "-c", "deepspeed --hostfile=/workspace/hostfile --no_ssh --node_rank=$RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT train.py --deepspeed --deepspeed_config ds_config.json"]

3、现在回过头来解决如何设置hostfile正确解析节点，在这里我们使用一个init.py脚本来实现，这个脚本主要实现三个功能: (1)获取参数配置，比如RANK。 (2)书写hostfile，让deepSpeed能够正确解析到容器节点。 (3)deepSpeed启动训练脚本。 init.py脚本如下：

import os
import re 
import subprocess

hostname = os.environ.get("HOSTNAME","bert-train-0")
WORK_SIZE = int(os.environ.get("WORK_SIZE",3))
RANK=0
NAME=""
def get_rank():
    global RANK
    global NAME
    pattern = "(.+)-([0-9]+)$"
    groups = re.match(pattern,hostname)
    NAME=groups[1]
    RANK=int(groups[2])

def write_hostfile():
    hostfile = ""
    for i in range(WORK_SIZE):
        hostfile += "{NAME}-{i}.{SVC} slots=1 \n".format(NAME=NAME,i=i,SVC=NAME)
    with open("/workspace/hostfile","w") as f:
        f.write(hostfile)
    print("Hostfile:")
    print(hostfile)

Command = "cd /workspace && deepspeed --hostfile=/workspace/hostfile --no_ssh --node_rank={RANK} --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT train.py --deepspeed --deepspeed_config ds_config.json"

def startup():
    cmd = Command.format(RANK=RANK)
    print("RUN COMMAND:\n{cmd}\n".format(cmd=cmd))
    proc = subprocess.Popen(cmd,shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    print("Start")
    while True:
        line = proc.stdout.readline()
        print(line)
        if not line:
            break
    retval = proc.wait()
    print("Command End:{}".format(retval))
    exit(retval)
    
get_rank()
write_hostfile()
startup()

4、接下来是编写k8s的yaml文件：

apiVersion: v1
kind: Service
metadata:
  name: bert-training
  labels:
    app: bert-training
spec:
  clusterIP: None  # Headless Service
  selector:
    app: bert-training
  ports:
    - name: nccl
      protocol: TCP
      port: 29500
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: bert-training
spec:
  serviceName: bert-training
  replicas: 3  # 3 个节点
  selector:
    matchLabels:
      app: bert-training
  template:
    metadata:
      labels:
        app: bert-training
    spec:
      # restartPolicy: Always
      containers:
        - name: bert-container
          image: 192.168.2.186:5000/deepspeed:v9-test2  # 你的 DeepSpeed 镜像
          imagePullPolicy: Always
          command: ["/bin/bash"]
          args: ["-c","python3 /init.py && tail -f /dev/null"]
          ports:
          - containerPort: 29500
            name:  nccl
            protocol: TCP
          env:
            - name: POD_NAME
              valueFrom:
                fieldRef:
                  fieldPath: metadata.name  # 获取 pod 的名称
            - name: WORLD_SIZE
              value: "3"
            - name: MASTER_ADDR
              value: "bert-training-0.bert-training"
            - name: MASTER_PORT
              value: "29500"
          # volumeMounts:
          #   - name: model-data
          #     mountPath: /workspace
      volumes:
        - name: model-data
          persistentVolumeClaim:
            claimName: bert-pvc
  volumeClaimTemplates:
    - metadata:
        name: bert-pvc
      spec:
        accessModes: [ "ReadWriteOnce" ]
        resources:
          requests:
            storage: 10Gi

当一切就绪就可以部署statefulset.yaml文件：

kubectl apply -f statefulset.yaml

备注：此次容器化的多机多卡训练涉及很多技术，比如K8s安装、集群搭建，可以参考文章: k3s集群搭建教学上面涉及的文件，需要的朋友可以点个收藏加关注，私信给我，我百度网盘发你。