关于Pytorch 分布式训练local_rank的坑

先是说明一下这个问题：

我们在看很多大佬写的代码时候，特别是涉及到分布式训练的时候会发现他们在 argsparse 中添加了这样一个参数“--loacl_rank”，比如下面是Swin-Transformer官方训练源码中cmd参数声明的一个部分，可以看到第70行添加了一个"“--local_rank”。

可是我们知道local_rank一般是需要动态获取的（也就是获取当前进程在本机上的rank），而不是由用户设置为固定值，所以它这个参数声明是什么意思呢。

问题解决：

其实这个问题在官方的说明文档上已经给出了答案：

大概内容就是，这个命令行参数“--loacl_rank”是必须声明的，但它不是由用户填写的，而是由pytorch为用户填写，也就是说这个值是会被自动赋值为当前进程在本机上的rank。但是有的新手（比如我），就会把这个参数理解为是需要用户声明，而用户声明的值会覆盖pytorch为用户生成的值，因此就会产生莫名其妙的错误。

问题扩展

现在命令行参数“--loacl_rank”的问题解决了，还以一个问题，就是还有很多大佬的代码在分布式训练中并没有声明命令行参数“--loacl_rank”，但程序同样可以运行，这是为什么呢？

回答这个问题首先需要解释一下pytorch分布式训练的启动方式（当然这种方式官方已经建议废弃，但在很多SOTA论文的代码中都使用这种方式，所以有必要了解）：

python -m torch.distributed.launch  --nproc_per_node 3  --use_env main.py

其中“--nproc_per_node”是每个节点的进程数量，“main.py”是程序的入口脚本。但“--use_env”这个参数，很多人拿来就用了，并没有注意它是干什么的，官方是这样解释的：

大概意思就是说，声明“--use_env”后，pytorch会将当前进程在本机上的rank添加到环境变量“LOCAL_RANK”中，而不再添加到args.local_rank。大家可以看一下下面的代码就理解了：

# d.py
import os
import argparse
def main(args):
    local_rank = args.local_rank
    print(local_rank, os.environ['LOCAL_RANK'])


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int)
    args = parser.parse_args()
    main(args)

输出：

也就是说如果声明“--use_env”那么pytorch就会把当前进程的在本机上的rank放到环境变量中，而不会放在args.local_rank中。

同时上面的输出大家可能也也注意到了，官方现在已经建议废弃使用torch.distributed.launch，转而使用torchrun，而这个torchrun已经把“--use_env”这个参数废弃了，转而强制要求用户从环境变量LOACL_RANK里获取当前进程在本机上的rank（一般就是本机上的gpu编号）。关于这样更改后的新写法大家可以参考：
官方文档，以及github上facebookai实现的detr