快捷導(dǎo)航

分布式訓(xùn)練training-operator和pytorch-distributed?RANK變量不統(tǒng)一解決

更新時(shí)間：2023年04月13日 15:03:52 作者：爛筆頭

這篇文章主要介紹了分布式訓(xùn)練training-operator和pytorch-distributed?RANK變量不統(tǒng)一問(wèn)題的解決方案詳解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進(jìn)步，早日升職加薪

正文

我們?cè)谑褂?training-operator 框架來(lái)實(shí)現(xiàn) pytorch 分布式任務(wù)時(shí)，發(fā)現(xiàn)一個(gè)變量不統(tǒng)一的問(wèn)題：在使用 pytorch 的分布式 launch 時(shí)，需要指定一個(gè)變量是 node_rank 。同時(shí)，在 OpenMMLab 框架的 dist_train.sh 里，讀取的系統(tǒng)環(huán)境變量是 NODE_RANK（如果系統(tǒng)里 NODE_RANK 沒(méi)有被指定，則用默認(rèn)值0）。

dist_train.sh

#!/usr/bin/env bash 
CONFIG=$1 
GPUS=$2 
NNODES=${NNODES:-1} 
NODE_RANK=${NODE_RANK:-0} # 如果NODE_RANK沒(méi)有被設(shè)置為系統(tǒng)變量，則使用默認(rèn)值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參數(shù)的一部分 
    --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 這個(gè)環(huán)境變量是以 RANK 的形式出現(xiàn)的。

這就會(huì)導(dǎo)致：通過(guò) training-operator 啟動(dòng)的訓(xùn)練 pod 里只有 RANK 變量，沒(méi)有 NODE_RANK 變量，那么， dist_train.sh 里的 $NODE_RANK 變量是一個(gè)默認(rèn)值 0，每一個(gè)被啟動(dòng)的訓(xùn)練 pod 里的 NODE_RANK 也是 0。這會(huì)讓每個(gè)pod都認(rèn)為自己是第 0 個(gè)，每個(gè) pod 都無(wú)法感知到別的 pod 的存在，那就會(huì)各自為政，在自己的 NODE 節(jié)點(diǎn)上重復(fù)性的做單機(jī)多卡的分布式訓(xùn)練。

那么，為了實(shí)現(xiàn)多機(jī)多卡的訓(xùn)練，就勢(shì)必需要解決 training-operator 提供的環(huán)境變量 RANK 與 torch.distributed.launch 需要的環(huán)境變量 NODE_RANK 的不統(tǒng)一的問(wèn)題。

解決的思路有兩個(gè)方向

保持 training-operator 的 RANK 變量不變，在訓(xùn)練的 pod 容器里，將 RANK 變量賦值給 NODE_RANK
修改 training-operator，添加 NODE_RANK 變量，并將 NODE_RANK 變量的值設(shè)為 RANK 的值

這里選第二個(gè)，因?yàn)榈谝粋€(gè)方案沒(méi)走通。。。

首先，將 training-operator 克隆到本地：GitHub - kubeflow/training-operator: Training operators on Kubernetes.
接著，全局搜索 RANK，發(fā)現(xiàn)該變量只出現(xiàn)在 ./pkg/controller.v1/pytorch/envvar.g 里：

然后，添加一個(gè) name=NODE_RANK，value= strconv.Itoa(rank) 的環(huán)境變量

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 https://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),
         })
         // 新增一個(gè)名為NODE_RANK的環(huán)境變量
         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}

替換掉默認(rèn)的鏡像，在./manifests/base/deployment.yaml里修改鏡像地址為上一步驟docker push的地址

重新部署，在./manifests/overlays/standalone目錄下

kubectl apply -k .

獲得 NODE_RANK變量

如下：

以上就是分布式訓(xùn)練training-operator和pytorch-distributed RANK變量不統(tǒng)一解決的詳細(xì)內(nèi)容，更多關(guān)于pytorch RANK變量不統(tǒng)一的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！

您可能感興趣的文章: