pytorch分布式训练之training-operator和pytorch-distributed里关于RANK变量不统一的问题解决

语言: CN / TW / HK

趁热记录下,给未来的自己

我们在使用 training-operator 框架来实现 pytorch 分布式任务时,发现一个变量不统一的问题:在使用 pytorch 的分布式 launch 时,需要指定一个变量是 node_rank 。同时,在 OpenMMLab 框架的 dist_train.sh 里,读取的系统环境变量是 NODE_RANK(如果系统里 NODE_RANK 没有被指定,则用默认值0)。

dist_train.sh

```

!/usr/bin/env bash

CONFIG=$1 GPUS=$2 NNODES=${NNODES:-1} NODE_RANK=${NODE_RANK:-0} # 如果NODE_RANK没有被设置为系统变量,则使用默认值0 PORT=${PORT:-29500} MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}

PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ python -m torch.distributed.launch \ --nnodes=$NNODES \ --node_rank=$NODE_RANK \ # 作为torch.distributed.launch参数的一部分 --master_addr=$MASTER_ADDR \ --nproc_per_node=$GPUS \ --master_port=$PORT \ $(dirname "$0")/train.py \ $CONFIG \ --seed 0 \ --launcher pytorch ${@:3} ```

而在 training-operator 里,NODE_RANK 这个环境变量是以 RANK 的形式出现的。

这就会导致:通过 training-operator 启动的训练 pod 里只有 RANK 变量,没有 NODE_RANK 变量,那么, dist_train.sh 里的 $NODE_RANK 变量是一个默认值 0,每一个被启动的训练 pod 里的 NODE_RANK 也是 0。这会让每个pod都认为自己是第 0 个,每个 pod 都无法感知到别的 pod 的存在,那就会各自为政,在自己的 NODE 节点上重复性的做单机多卡的分布式训练

那么,为了实现多机多卡的训练,就势必需要解决 training-operator 提供的环境变量 RANK 与 torch.distributed.launch 需要的环境变量 NODE_RANK 的不统一的问题。

解决的思路有两个方向:

  1. 保持 training-operator 的 RANK 变量不变,在训练的 pod 容器里,将 RANK 变量赋值给 NODE_RANK
  2. 修改 training-operator,添加 NODE_RANK 变量,并将 NODE_RANK 变量的值设为 RANK 的值

这里选第二个,因为第一个方案没走通。。。

  • 然后,添加一个 name=NODE_RANKvalue= strconv.Itoa(rank) 的环境变量

``` func setPodEnv(obj interface{}, podTemplateSpec corev1.PodTemplateSpec, rtype, index string) error { pytorchjob, ok := obj.(kubeflowv1.PyTorchJob) if !ok { return fmt.Errorf("%+v is not a type of PyTorchJob", obj) }

for i := range podTemplateSpec.Spec.Containers { // Initialize the environment variables. if len(podTemplateSpec.Spec.Containers[i].Env) == 0 { podTemplateSpec.Spec.Containers[i].Env = make([]corev1.EnvVar, 0) } // Set PYTHONUNBUFFERED to true, to disable output buffering. // Ref http://stackoverflow.com/questions/59812009/what-is-the-use-of-pythonunbuffered-in-docker-file. podTemplateSpec.Spec.Containers[i].Env = append( podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{ Name: "PYTHONUNBUFFERED", Value: "0", })

  // If the master is not null, then we need to set the MASTER_ADDR and RANK.
  if pytorchjob.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeMaster] != nil {
     envVars, err := GetMasterEnvVarGenerator().Generate(pytorchjob)
     if err != nil {
        return err
     }

     // Set master related environment variables.
     podTemplateSpec.Spec.Containers[i].Env = append(
        podTemplateSpec.Spec.Containers[i].Env, envVars...)

     // Set world size and rank.
     rank, err := strconv.Atoi(index)
     if err != nil {
        return err
     }

     if rtype == strings.ToLower(string(kubeflowv1.PyTorchJobReplicaTypeWorker)) {
        rank = rank + 1
     }

     totalReplicas := getTotalReplicas(pytorchjob)

     podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
        Name:  "WORLD_SIZE",
        Value: strconv.Itoa(int(totalReplicas)),
     })

     podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
        Name:  "RANK",
        Value: strconv.Itoa(rank),
     })

     // 新增一个名为NODE_RANK的环境变量
     podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
        Name:  "NODE_RANK",
        Value: strconv.Itoa(rank),
     })
  }

  // Set the elastic environment variables if the elasticPolicy is not null.
  if pytorchjob.Spec.ElasticPolicy != nil {
     envVars, err := GetElasticEnvVarGenerator().Generate(pytorchjob)
     if err != nil {
        return err
     }
     // Set elastic related environment variables.
     podTemplateSpec.Spec.Containers[i].Env = append(
        podTemplateSpec.Spec.Containers[i].Env, envVars...)
  }

}

return nil } ```

  • 重新编译:go build & docker build

# Build manager binary. go build -o bin/manager cmd/training-operator.v1/main.go # Build docker image with the manager. docker build -t ${IMG} -f build/images/training-operator/Dockerfile . # Push docker image with the manager. docker push ${IMG}

  • 替换掉默认的镜像,在./manifests/base/deployment.yaml里修改镜像地址为上一步骤docker push的地址

image.png

  • 重新部署, 在./manifests/overlays/standalone目录下

kubectl apply -k .

获得 NODE_RANK变量,如下: image.png

以上