正文
我们在使用 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 的不统一的问题。
解决的思路有两个方向
- 保持 training-operator 的 RANK 变量不变，在训练的 pod 容器里，将 RANK 变量赋值给 NODE_RANK
- 修改 training-operator，添加 NODE_RANK 变量，并将 NODE_RANK 变量的值设为 RANK 的值
这里选第二个，因为第一个方案没走通。。。
- 首先，将 training-operator 克隆到本地：GitHub - kubeflow/training-operator: Training operators on Kubernetes.
- 接着，全局搜索 RANK，发现该变量只出现在
./pkg/controller.v1/pytorch/envvar.g里：
- 然后，添加一个
name=NODE_RANK，value= 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 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),
})
// 新增一个名为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的地址
- 重新部署， 在./manifests/overlays/standalone目录下
kubectl apply -k .
获得 NODE_RANK变量
如下：
以上就是分布式训练training-operator和pytorch-distributed RANK变量不统一解决的详细内容，更多关于pytorch RANK变量不统一的资料请关注编程网其它相关文章！
