正文
我们在使用 training-operator 框架来实现 pytorch 分布式任务时,发现一个变量不统一的问题:在使用 pytorch 的分布式 launch 时,需要指定一个变量是 node_rank 。同时,在 OpenMMLab 框架的 dist_train.sh 里,读取的系统环境变量是 NODE_RANK(如果系统里 NODE_RANK 没有被指定,则用默认值0)。
dist_train.sh
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
|
#!/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)
的环境变量
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
|
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
1
2
3
4
5
6
|
# 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目录下
1
|
kubectl apply -k . |
获得 NODE_RANK变量
如下:
以上就是分布式训练training-operator和pytorch-distributed RANK变量不统一解决的详细内容,更多关于pytorch RANK变量不统一的资料请关注服务器之家其它相关文章!
原文链接:https://juejin.cn/post/7127079132822945800