技术

《大数据经典论文解读》 三驾马车学习 Spark 内存管理及调优 Yarn学习 从Spark部署模式开始讲源码分析 容器狂占内存资源怎么办? 多角度理解一致性 golang io使用及优化模式 Flink学习 c++学习 学习ebpf go设计哲学 ceph学习 学习mesh kvm虚拟化 学习MQ go编译器 学习go 为什么要有堆栈 汇编语言 计算机组成原理 运行时和库 Prometheus client mysql 事务 mysql 事务的隔离级别 mysql 索引 坏味道 学习分布式 学习网络 学习Linux go 内存管理 golang 系统调用与阻塞处理 Goroutine 调度过程 重新认识cpu mosn有的没的 负载均衡泛谈 单元测试的新解读 《Redis核心技术与实现》笔记 《Prometheus监控实战》笔记 Prometheus 告警学习 calico源码分析 对容器云平台的理解 Prometheus 源码分析 并发的成本 基础设施优化 hashicorp raft源码学习 docker 架构 mosn细节 与微服务框架整合 Java动态代理 编程范式 并发通信模型 《网络是怎样连接的》笔记 go channel codereview gc分析 jvm 线程实现 go打包机制 go interface及反射 如何学习Kubernetes 《编译原理之美》笔记——后端部分 《编译原理之美》笔记——前端部分 Pilot MCP协议分析 go gc 内存管理玩法汇总 软件机制 istio流量管理 Pilot源码分析 golang io 学习Spring mosn源码浅析 MOSN简介 《datacenter as a computer》笔记 学习JVM Tomcat源码分析 Linux可观测性 学习存储 学计算 Gotty源码分析 kubernetes operator kaggle泰坦尼克问题实践 kubernetes扩缩容 神经网络模型优化 直觉上理解深度学习 如何学习机器学习 TIDB源码分析 什么是云原生 Alibaba Java诊断工具Arthas TIDB存储——TIKV 《Apache Kafka源码分析》——简介 netty中的线程池 guava cache 源码分析 Springboot 启动过程分析 Spring 创建Bean的年代变迁 Linux内存管理 自定义CNI IPAM 共识算法 spring redis 源码分析 kafka实践 spring kafka 源码分析 Linux进程调度 让kafka支持优先级队列 Codis源码分析 Redis源码分析 C语言学习 《趣谈Linux操作系统》笔记 docker和k8s安全访问机制 jvm crash分析 Prometheus 学习 Kubernetes监控 容器日志采集 Kubernetes 控制器模型 容器狂占资源怎么办? Kubernetes资源调度——scheduler 时序性数据库介绍及对比 influxdb入门 maven的基本概念 《Apache Kafka源码分析》——server Kubernetes类型系统 源码分析体会 《数据结构与算法之美》——算法新解 Kubernetes源码分析——controller mananger Kubernetes源码分析——apiserver Kubernetes源码分析——kubelet Kubernetes介绍 ansible学习 Kubernetes源码分析——从kubectl开始 jib源码分析之Step实现 jib源码分析之细节 线程排队 跨主机容器通信 jib源码分析及应用 为容器选择一个合适的entrypoint kubernetes yaml配置 《持续交付36讲》笔记 mybatis学习 程序猿应该知道的 无锁数据结构和算法 CNI——容器网络是如何打通的 为什么很多业务程序猿觉得数据结构和算法没用? 串一串一致性协议 当我在说PaaS时,我在说什么 《数据结构与算法之美》——数据结构笔记 PouchContainer技术分享体会 harbor学习 用groovy 来动态化你的代码 精简代码的利器——lombok 学习 《深入剖析kubernetes》笔记 编程语言那些事儿 rxjava3——背压 rxjava2——线程切换 spring cloud 初识 《深入拆解java 虚拟机》笔记 《how tomcat works》笔记 hystrix 学习 rxjava1——概念 Redis 学习 TIDB 学习 如何分发计算 Storm 学习 AQS1——论文学习 Unsafe Spark Stream 学习 linux vfs轮廓 《自己动手写docker》笔记 java8 实践 中本聪比特币白皮书 细读 区块链泛谈 比特币 大杂烩 总纲——如何学习分布式系统 hbase 泛谈 forkjoin 泛谈 看不见摸不着的cdn是啥 《jdk8 in action》笔记 程序猿视角看网络 bgp初识 calico学习 AQS——粗略的代码分析 我们能用反射做什么 web 跨域问题 《clean code》笔记 《Elasticsearch权威指南》笔记 mockito简介及源码分析 2017软件开发小结—— 从做功能到做系统 《Apache Kafka源码分析》——clients dns隐藏的一个坑 《mysql技术内幕》笔记 log4j学习 为什么netty比较难懂? 回溯法 apollo client源码分析及看待面向对象设计 学习并发 docker运行java项目的常见问题 OpenTSDB 入门 spring事务小结 分布式事务 javascript应用在哪里 《netty in action》读书笔记 netty对http2协议的解析 ssl证书是什么东西 http那些事 苹果APNs推送框架pushy apple 推送那些事儿 编写java框架的几大利器 java内存模型 java exception Linux IO学习 netty内存管理 测试环境docker化实践 netty在框架中的使用套路 Nginx简单使用 《Linux内核设计的艺术》小结 Go并发机制及语言层工具 Linux网络源代码学习——数据包的发送与接收 《docker源码分析》小结 docker namespace和cgroup Linux网络源代码学习——整体介绍 zookeeper三重奏 数据库的一些知识 Spark 泛谈 链式处理的那些套路 netty回顾 Thrift基本原理与实践(二) Thrift基本原理与实践(一) 回调 异步执行抽象——Executor与Future Docker0.1.0源码分析 java gc Jedis源码分析 深度学习泛谈 Linux网络命令操作 JTA与TCC 换个角度看待设计模式 Scala初识 向Hadoop学习NIO的使用 以新的角度看数据结构 并发控制相关的硬件与内核支持 systemd 简介 quartz 源码分析 基于docker搭建测试环境(二) spring aop 实现原理简述 自己动手写spring(八) 支持AOP 自己动手写spring(七) 类结构设计调整 分析log日志 自己动手写spring(六) 支持FactoryBean 自己动手写spring(九) 总结 自己动手写spring(五) bean的生命周期管理 自己动手写spring(四) 整合xml与注解方式 自己动手写spring(三) 支持注解方式 自己动手写spring(二) 创建一个bean工厂 自己动手写spring(一) 使用digester varnish 简单使用 关于docker image的那点事儿 基于docker搭建测试环境 分布式配置系统 JVM执行 git maven/ant/gradle/make使用 再看tcp kv系统 java nio的多线程扩展 《Concurrency Models》笔记 回头看Spring IOC IntelliJ IDEA使用 Java泛型 vagrant 使用 Go常用的一些库 Python初学 Goroutine 调度模型 虚拟网络 《程序员的自我修养》小结 Kubernetes存储 访问Kubernetes上的Service Kubernetes副本管理 Kubernetes pod 组件 Go基础 JVM类加载 硬币和扑克牌问题 LRU实现 virtualbox 使用 ThreadLocal小结 docker快速入门

架构

实时训练 分布式链路追踪 helm tensorflow原理——python层分析 如何学习tensorflow 数据并行——allreduce 数据并行——ps 机器学习中的python调用c 机器学习训练框架概述 embedding的原理及实践 tensornet源码分析 大模型训练 X的生成——特征工程 tvm tensorflow原理——core层分析 模型演变 《深度学习推荐系统实战》笔记 keras 和 Estimator tensorflow分布式训练 分布式训练的一些问题 基于Volcano的弹性训练 图神经网络 pytorch弹性分布式训练 在离线业务混部 RNN pytorch分布式训练 CNN 《动手学深度学习》笔记 pytorch与线性回归 多活 volcano特性源码分析 推理服务 kubebuilder 学习 mpi 学习pytorch client-go学习 tensorflow学习 提高gpu 利用率 GPU与容器的结合 GPU入门 AI云平台 tf-operator源码分析 k8s批处理调度 喜马拉雅容器化实践 Kubernetes 实践 学习rpc BFF 生命周期管理 openkruise学习 可观察性和监控系统 基于Kubernetes选主及应用 《许式伟的架构课》笔记 Kubernetes webhook 发布平台系统设计 k8s水平扩缩容 Scheduler如何给Node打分 Scheduler扩展 controller 组件介绍 openkruise cloneset学习 controller-runtime源码分析 pv与pvc实现 csi学习 client-go源码分析 kubelet 组件分析 调度实践 Pod是如何被创建出来的? 《软件设计之美》笔记 mecha 架构学习 Kubernetes events学习及应用 CRI 资源调度泛谈 业务系统设计原则 grpc学习 元编程 以应用为中心 istio学习 下一代微服务Service Mesh 《实现领域驱动设计》笔记 serverless 泛谈 概率论 《架构整洁之道》笔记 处理复杂性 那些年追过的并发 服务器端编程 网络通信协议 架构大杂烩 如何学习架构 《反应式设计模式》笔记 项目的演化特点 反应式架构摸索 函数式编程的设计模式 服务化 ddd反模式——CRUD的败笔 研发效能平台 重新看面向对象设计 业务系统设计的一些体会 函数式编程 《左耳听风》笔记 业务程序猿眼中的微服务管理 DDD实践——CQRS 项目隔离——案例研究 《编程的本质》笔记 系统故障排查汇总及教训 平台支持类系统的几个点 代码腾挪的艺术 abtest 系统设计汇总 《从0开始学架构》笔记 初级权限系统设计 领域驱动理念入门 现有上传协议分析 移动网络下的文件上传要注意的几个问题 推送系统的几个基本问题 用户登陆 做配置中心要想好的几个基本问题 不同层面的异步 分层那些事儿 性能问题分析 当我在说模板引擎的时候,我在说什么 用户认证问题 资源的分配与回收——池 消息/任务队列

标签


pytorch分布式训练

2021年10月30日

简介

整体思路

DDP系列第一篇:入门教程多机多卡

  1. 分布式训练的几个选择:
    1. 模型并行 vs 数据并行
    2. PS 模型 vs Ring-Allreduce 。 Allreduce operation can not start until all processes join, it is considered to be a synchronized communication, as opposed to the P2P communication used in parameter servers.
  2. 假设 有两台机子,每台8张显卡,那就是2x8=16个进程,并行数是16。DDP 推荐每个进程一张卡,在16张显卡,16的并行数下,DDP会同时启动16个进程。rank表示当前进程的序号(0~16),用于进程间通讯。local_rank 表示每台机子上的进程的序号(0~7),也用来指定机器上的gpu 序号。用起来 就是一行代码,model = DDP(model, device_ids=[local_rank], output_device=local_rank)后续的模型关于前向传播、后向传播的用法,和单机单卡完全一致DDP系列第二篇:实现原理与源代码解析
  3. 每个进程跑的是同一份代码,进程从外界(比如环境变量)获取 rank/master_addr/master_port 等参数,rank = 0 的进程为 master 进程

[源码解析] PyTorch 分布式(8) ——– DistributedDataParallel 之 论文篇应用程序开发通常从本地模型开始,然后在必要时扩展。所以需要有一个从本地模型开始,修改代码以适应分布式的过程。为了避免这个从本地模型到分布式模型的过渡期间太过麻烦,API在应用程序代码中是非侵入性的。

单机单卡

Pytorch分布式训练

单机单卡训练步骤:定义网络,定义dataloader,定义loss和optimizer,开训

BATCH_SIZE = 256
EPOCHS = 5
if __name__ == "__main__":
    # 1. define network
    device = "cuda"
    net = torchvision.models.resnet18(num_classes=10)
    net = net.to(device=device)  
    # 2. define dataloader
    trainset = torchvision.datasets.CIFAR10(...)
    train_loader = torch.utils.data.DataLoader(
        trainset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
    )
    # 3. define loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        net.parameters(),
        lr=0.01,
        momentum=0.9,
        weight_decay=0.0001,
        nesterov=True,
    )
    print("            =======  Training  ======= \n")
    # 4. start to train
    net.train()
    for ep in range(1, EPOCHS + 1):                 # 控制在全部数据上训练的次数
        for idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss = criterion(outputs, targets)      # 获得输出结果和样本值的损失函数
            optimizer.zero_grad()
            loss.backward()                         # 根据loss 进行反向梯度传播
            optimizer.step()                        # 用计算的梯度去做优化

单机多卡

如果是单机多卡,定义网络时加入 net = nn.DataParallel(net)

PyTorch 分布式(2) —– DataParallel(上) 缺点很多。

分布式训练

pytorch ddp论文:During distributed training, each pro- cess has its own local model replica and local optimizer. In terms of correctness, distributed data parallel training and local training must be mathematically equivalent. DDP guarantees the correctness by making sure that all model replicas start from the exact same model state, and see the same parameter gradients after every backward pass. Therefore, even though optimizers from different processes are all independent, they should be able to bring their local model replicas to the same state at the end of every iteration.

PyTorch 分布式(5) —— DistributedDataParallel 总述&如何使用

  1. 加载模型阶段。每个GPU都拥有模型的一个副本,所以不需要拷贝模型。rank为0的进程会将网络初始化参数broadcast到其它每个进程中,确保每个进程中的模型都拥有一样的初始化值。
  2. 加载数据阶段。DDP 不需要广播数据,而是使用多进程并行加载数据。在 host 之上,每个worker进程都会把自己负责的数据从硬盘加载到 page-locked memory。DistributedSampler 保证每个进程加载到的数据是彼此不重叠的。
  3. 前向传播阶段。在每个GPU之上运行前向传播,计算输出。每个GPU都执行同样的训练,所以不需要有主 GPU。
  4. 计算损失。在每个GPU之上计算损失。
  5. 反向传播阶段。运行后向传播来计算梯度,在计算梯度同时也对梯度执行all-reduce操作。每一层的梯度不依赖于前一层,所以梯度的All-Reduce和后向过程同时计算,以进一步缓解网络瓶颈。在后向传播的最后,每个节点都得到了平均梯度,这样模型参数保持同步。关于AllReduce
  6. 更新模型参数阶段。因为每个GPU都从完全相同的模型开始训练,并且梯度被all-reduced,因此每个GPU在反向传播结束时最终得到平均梯度的相同副本,所有GPU上的权重更新都相同,也就不需要模型同步了

模板代码

import torch.distributed as dist
import torch.utils.data.distributed

parser = argparse.ArgumentParser(description='PyTorch distributed training on cifar-10')
parser.add_argument('--rank', default=0,help='rank of current process')
parser.add_argument('--word_size', default=2,help="word size")
parser.add_argument('--init_method', default='tcp://127.0.0.1:23456',help="init-method")
args = parser.parse_args()
...
dist.init_process_group(backend='nccl', init_method=args.init_method, rank=args.rank, world_size=args.word_size)
...
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
## 数据并行需要进行数据切片
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset,num_replicas=args.world_size,rank=rank)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)
...
net = Net() 
net = Net().to(device) # device 代表到某个 gpu 或cpu 上
## 使用DistributedDataParallel 修饰模型
net = torch.nn.parallel.DistributedDataParallel(net)

PyTorch分布式训练简明教程

  1. 用dist.init_process_group初始化分布式环境
    1. 一般来说,nccl 用于 GPU 分布式训练,gloo 用于 CPU 进行分布式训练。
    2. 这个调用是阻塞的,必须等待所有进程来同步,如果任何一个进程出错,就会失败。
  2. 数据侧,我们nn.utils.data.DistributedSampler来给各个进程切分数据,只需要在dataloader中使用这个sampler就好
    1. 使用 DDP 时,不再是从主 GPU 分发数据到其他 GPU 上,而是各 GPU 从自己的硬盘上读取属于自己的那份数据。
    2. 训练循环过程的每个epoch开始时调用train_sampler.set_epoch(epoch),(主要是为了保证每个epoch的划分是不同的)其它的训练代码都保持不变。
    3. 举个例子:假设10w 数据,10个实例。单机训练任务一般采用batch 梯度下降,比如batch=100。分布式后,10个实例每个训练1w,这个时候每个实例 batch 是100 还是 10 呢?是100,batch 的大小取决于 每个实例 GPU 显存能放下多少数据。
  3. 模型侧,我们只需要用DistributedDataParallel包装一下原来的model

pytorch 中的任何 net 都 是 torch.nn.Module 的子类,DistributedDataParallel 也是 torch.nn.Module 子类,任何 torch.nn.Module 子类 都可以覆盖 __init__forward方法 ,DistributedDataParallel 可以从 net 拿到模型数据(以及 在哪个gpu 卡上运行) ,也可以 从指定或默认的 process_group 获取信息。最后在__init__ 和 forward 中加入 同步梯度的逻辑,完美的将 同步梯度的逻辑 隐藏了起来

实际例子

BATCH_SIZE = 256
EPOCHS = 5
if __name__ == "__main__":
    # 0. set up distributed device
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(rank % torch.cuda.device_count())
    dist.init_process_group(backend="nccl")
    device = torch.device("cuda", local_rank) # 把模型放置到特定的GPU上
    print(f"[init] == local rank: {local_rank}, global rank: {rank} ==")
    # 1. define network
    net = torchvision.models.resnet18(pretrained=False, num_classes=10)
    net = net.to(device)
    # DistributedDataParallel
    net = DDP(net, device_ids=[local_rank], output_device=local_rank)
    # 2. define dataloader
    trainset = torchvision.datasets.CIFAR10(...)
    # DistributedSampler
    # we test single Machine with 2 GPUs so the [batch size] for each process is 256 / 2 = 128
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        trainset,
        shuffle=True,
    )
    train_loader = torch.utils.data.DataLoader(
        trainset,
        batch_size=BATCH_SIZE,
        num_workers=4,
        pin_memory=True,
        sampler=train_sampler,
    )
    # 3. define loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        net.parameters(),
        lr=0.01 * 2,
        momentum=0.9,
        weight_decay=0.0001,
        nesterov=True,
    )
    if rank == 0:
        print("            =======  Training  ======= \n")
    # 4. start to train
    net.train()     # 标记为训练模式,推理时使用 net.eval()
    for ep in range(1, EPOCHS + 1):
        train_loss = correct = total = 0
        # set sampler
        train_loader.sampler.set_epoch(ep)
        for idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss = criterion(outputs, targets)
            optimizer.zero_grad()       # 将梯度归零
            loss.backward()             # 反向传播计算得到每个参数的梯度值
            optimizer.step()            # 通过梯度下降执行一步参数更新,optimizer只负责通过梯度下降进行优化,而不负责产生梯度,梯度是tensor.backward()方法产生的。
            train_loss += loss.item()
            total += targets.size(0)
            correct += torch.eq(outputs.argmax(dim=1), targets).sum().item()
            # 输出loss 和 正确率
    if rank == 0:
        print("\n            =======  Training Finished  ======= \n")

每条命令表示一个进程。若已开启的进程未达到 word_size 的数量,则所有进程会一直等待

对训练的影响

假设10w 数据,单机训练任务一般采用batch 梯度下降,比如batch=100。分布式后10个实例,每个实例训练1w。

  1. 如果每个实例 batch 是100 ,根据pytorch 数据并行的原理,每个实例每次batch 训练后向传播会交换一次梯度,得到的梯度是每个实例上梯度的均值,也就是说你有10个实例,这个梯度均值就是1000个样本的梯度均值。学习率一般也要相应调整。
  2. 一般 每个实例的batch = 单机batch / world_size

DDP 总体实现

DDP implementation lives in both Python and C++ files, with Python exposing the API and composing non-performance-critical components, and C++ serving the core gradient reduction algorithm. The Python API calls into C++ core through Pybind11.

本地模型 Module 定义如下,经过 DDP 的封装,便成为一个有分布式能力的Module了。

class Module:
    dump_patches: bool = False
    _version: int = 1
    training: bool      # 本网络是否正在训练
    _is_full_backward_hook: Optional[bool]

    def __init__(self):
        """
        Initializes internal Module state, shared by both nn.Module and ScriptModule.
        """
        torch._C._log_api_usage_once("python.nn_module")
        self.training = True
        self._parameters = OrderedDict()            # 在训练过程中会随着 BP 而更新的参数
        self._buffers = OrderedDict()               # 在训练过程中不会随着 BP 而更新的参数, running variance,running mean etc
        self._non_persistent_buffers_set = set()
        self._backward_hooks = OrderedDict()
        self._is_full_backward_hook = None
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._state_dict_hooks = OrderedDict()
        self._load_state_dict_pre_hooks = OrderedDict()
        self._modules = OrderedDict()               # 本网络下属的子模块,采取迭代的方式进行定义

PyTorch 分布式(9) —– DistributedDataParallel 之初始化一次DistributedDataParallel迭代中的步骤如下

  1. Prerequisite:DDP 依赖 c10dProcessGroup进行通信。因此,应用程序必须ProcessGroup在构建 DDP 之前创建实例。
  2. Constuctor
    1. rank=0 进程会引用本地模块,把模型state_dict()参数广播到所有进程之中,这样可以保证所有进程使用同样初始化数值和模型副本进行训练。
    2. 每个 DDP 进程创建一个 local Reducer,为了提高通信效率,Reducer将参数梯度组织成桶,一次reduce一个桶。Reducer还在构造期间注册 autograd 钩子,每个参数一个钩子。当梯度准备好时,将在向后传递期间触发这些钩子。
  3. Forward Pass: [源码解析] PyTorch 分布式(12)—–DistributedDataParallel 之 前向传播
    1. 每个进程读去自己的训练数据,DistributedSampler确保每个进程读到的数据不同。
    2. DDP 获取输入并将其传递给本地模型。模型进行前向计算,结果设置为 out。模型网络输出不需要gather到 rank 0进程了,这与 DP不同。
  4. Backward Pass:
    1. backward()在 loss 上直接调用该函数 Tensor,这是 DDP 无法控制的,DDP 使用构造时注册的 autograd hooks 来触发梯度同步。当一个梯度准备好时,它在该梯度累加器上的相应 DDP 钩子将触发。
    2. 在 autograd_hook 之中进行all-reduce。假设参数index是param_index,则利用param_index获取到参数,标示为ready,如果某个桶里面梯度都ready,则该桶是ready。当一个桶中的梯度都准备好时,会 在该桶上Reducer启动异步allreduce以计算所有进程的梯度平均值。如果所有桶都ready,则等待所有 all-reduce 完成。当所有桶都准备好时,Reducer将阻塞等待所有allreduce操作完成。完成此操作后,将平均梯度写入param.grad所有参数的字段。所有进程的梯度都会reduce,更新之后,大家的模型权重都相同。所以在向后传播完成之后,跨不同DDP进程的对应的相同参数上的 grad 字段应该是相等的。
    3. 不需要像 DP 那样每次迭代之后还要广播参数。但是 Buffers 还是需要在每次迭代由 rank 0 进程广播到其他进程之上。
  5. Optimizer Step: 从优化器的角度来看,它正在优化本地模型。

单机场景下,pytorch 封装了torch.autograd包,torch.autograd is PyTorch’s automatic differentiation engine that powers neural network training.torch.autograd 封装了前向后向传播逻辑(实现自动微分),所谓分布式首先是 autograd engine的分布式PyTorch 分布式 Autograd (1) —- 设计

  1. 分布式RPC框架
  2. 前向传播期间的 Autograd 记录
  3. 分布式 Autograd 上下文
  4. 分布式反向传播
  5. 分布式优化器

PyTorch 分布式 Autograd (2) —- RPC基础 是一个系列文章,有兴趣可以细读下, 比较有价值的点就是 有很多low leve api 代码示例,可以看下较为原汁原味的 分布式前后向传播过程。比如

# 代码目的是让两个 worker 之间就通过 RPC 进行协作。
def my_add(t1, t2):
  return torch.add(t1, t2)
def worker0():
    # On worker 0:
    # Setup the autograd context. Computations that take
    # part in the distributed backward pass must be within
    # the distributed autograd context manager.
    with dist_autograd.context() as context_id:
      t1 = torch.rand((3, 3), requires_grad=True)
      t2 = torch.rand((3, 3), requires_grad=True)
      # 第一阶段:RPC操作,构建依赖基础
      # Perform some computation remotely.
      t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))
      # Perform some computation locally based on remote result.
      t4 = torch.rand((3, 3), requires_grad=True)
      t5 = torch.mul(t3, t4)
      # Compute some loss.
      loss = t5.sum()

      # 第二阶段,执行后向传播
      # Run the backward pass.
      dist_autograd.backward(context_id, [loss])
      # Retrieve the gradients from the context.
      dist_autograd.get_gradients(context_id)
      print(loss)  

ProcessGroup

Collective communication包含多个sender和多个receiver,一般的通信原语包括 broadcast,gather,all-gather,scatter,reduce,all-reduce,reduce-scatter,all-to-all等。再对照 ProcessGroup python 定义,ProcessGroup可以看做对 Collective Communication Library 的封装。

多个process 组成一个process_group,比如每个group 持有一个tensor,想计算tensor的和,可以借助p2p先汇总到一个 process 上求和再广播(如果tensor很大呢?),但其实可以边communication边computation,来挖掘性能和节约带宽。

PyTorch 分布式(7) —– DistributedDataParallel 之进程组

如何让 worker 在建立ProcessGroup之前发现彼此?

  1. init_method,开始的前提:如何联系到其它机器上的进程。
    1. Shared file-system initialization:init_method=’file:///mnt/nfs/sharedfile’
    2. dist.init_process_group(backend, init_method=’tcp://10.1.1.20:23456’,ank=args.rank, world_size=4)
    3. dist.init_process_group(backend, init_method=’env://’,ank=args.rank, world_size=4) 这个进程会自动读取环境变量 MASTER_ADDR/MASTER_PORT/WORLD_SIZE/RANK
  2. store ,有TcpStore 和EtcdStore 可以发挥类似 etcd 用来发挥的作用:服务发现及分布式键值存储

store 和 init_method 是互斥的。store 先于 ProcessGroup 创建和初始化。PS: 随着版本的迭代,变动较大,不纠结了

init_process_group ==> ` rendezvous_iterator = rendezvous(init_method, rank, world_size, timeout=timeout)` 使用了rendezvous 机制,与elastic agent 共用 TcpStore(默认是共用的,也可以不共用), 此处的 rendezvous_handler 是一个函数(返回store/rank/world_size),与elastic agent 使用的 RendezvousHandler 不同。

# /pytorch/torch/distributed/rendezvous.py
def rendezvous(url: str, rank: int = -1, world_size: int = -1, **kwargs) -> (store, rank, world_size):
    ...
    return _rendezvous_handlers[result.scheme](url, **kwargs)
# 如果 env TORCHELASTIC_USE_AGENT_STORE=true,则此处创建的 是 TCPStore client
def _create_c10d_store(hostname, port, rank, world_size, timeout) -> Store:
def _file_rendezvous_handler(url: str, **kwargs):
    ...
    store = FileStore(path, world_size)
    ...
def _tcp_rendezvous_handler(url: str, timeout: timedelta = default_pg_timeout, **kwargs):
    ...
    _create_c10d_store(hostname, port, rank, world_size, timeout) -> Store:
    ...
def _env_rendezvous_handler(url: str, timeout: timedelta = default_pg_timeout, **kwargs):
    ...
    _create_c10d_store(hostname, port, rank, world_size, timeout) -> Store:
    ...
def register_rendezvous_handler(scheme, handler):
register_rendezvous_handler("tcp", _tcp_rendezvous_handler)
register_rendezvous_handler("env", _env_rendezvous_handler)
register_rendezvous_handler("file", _file_rendezvous_handler)

ProcessGroup 如何被使用

Collective functionsif the system we use for distributed training has 2 nodes, each of which has 8 GPUs. On each of the 16 GPUs, there is a tensor that we would like to all-reduce. The following code can serve as a reference:

# Code running on Node 0
import torch
import torch.distributed as dist
dist.init_process_group(backend="nccl",
                        init_method="file:///distributed_test",
                        world_size=2,
                        rank=0)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
    tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))

dist.all_reduce_multigpu(tensor_list)
# Code running on Node 1
import torch
import torch.distributed as dist
dist.init_process_group(backend="nccl",
                        init_method="file:///distributed_test",
                        world_size=2,
                        rank=1)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
    tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))
dist.all_reduce_multigpu(tensor_list)

After the call, all 16 tensors on the two nodes will have the all-reduced value of 16

抛开概念,从代码看其本质。processgroup 就是给每一个训练的 process 建立一个Communication thread。主线程(Computation thread)在前台进行训练,这个Communication thread 在后台做通信(比如交流梯度)。

可以看下 ProcessGroup 的接口定义(从理论上,Ring AllReduce 分为Split/ScatterReudce/AllGather 三个步骤,不同框架命名会有差异)

# torch/_C/_distributed_c10d.py
class ProcessGroup(__pybind11_builtins.pybind11_object):
    def allgather(self, *args, **kwargs): 
    def allgather_coalesced(self, output_lists, *args, **kwargs): 
    def allreduce(self, *args, **kwargs): 
    def allreduce_coalesced(self, tensors, *args, **kwargs): 
    def alltoall(self, *args, **kwargs): 
    def alltoall_base(self, *args, **kwargs): 
    def barrier(self, opts, *args, **kwargs): 
    def broadcast(self, *args, **kwargs): 
    def gather(self, *args, **kwargs): 
    def monitored_barrier(self, timeout=None, *args, **kwargs): 
    def rank(self): 
    def recv(self, arg0, *args, **kwargs): 
    def recv_anysource(self, arg0, *args, **kwargs): 
    def reduce(self, *args, **kwargs): 
    def reduce_scatter(self, *args, **kwargs): 
    def scatter(self, *args, **kwargs): 
    def send(self, arg0, *args, **kwargs): 
    def size(self): 
    def _allgather_base(self, output, *args, **kwargs):  
    def _get_sequence_number_for_group(self): 
    def _set_sequence_number_for_group(self): 

ddp 初始化时 对模型参数进行广播 会调用到 process_group->broadcast(vec),反向传播 需要对梯度做 all-reduce 时候,Reducer会调用 process_group_->allreduce(tensors) 进行处理。

backend 是一个逻辑上的概念。本质上后端是一种IPC通信机制。对于用户来说,就是采用那种方式来进行集合通信,从代码上看,就是走什么流程(一系列流程),以及后端使用 ProcessGroupMPI 还是 ProcessGroupGloo …..。各主机之间需要进行通信。因此,需要指定通信的协议架构等。torch.distributed 对其进行了封装。提供了分布式通信原语

其它

DDP 在启动时 将 rank=0 的state_dict() 广播到其他worker,以保证所有worker的模型初始状态相同。需要广播的 state_dict 是什么?pytorch 的 state_dict 是一个字典对象,其将模型的每一层与它的对应参数建立映射关系,比如 model 每一层的weights及偏置等等。只有那些参数可以训练的层(比如卷积层,线性层等)才会被保存到模型的state_dict中,池化层、BN层这些本身没有参数的层就不会保存在 state_dict 之中,比如针对下面模型。

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

state_dict 如下:

self.module.state_dict() = {OrderedDict: 4} 
 'net1.weight' = {Tensor: 10} tensor([[ 0.2687,  0.0840, -0.1032,  0.3079,  0.0385, -0.0495, -0.3068, -0.1271,\n         -0.1067, -0.1966],\n        [-0.1203,  0.1789,  0.0666,  0.1882,  0.1335,  0.1921, -0.1145, -0.1781,\n          0.0661, -0.2339],\n        [ 0.1865, -0.2076,  0.2071,  0
 'net1.bias' = {Tensor: 10} tensor([ 0.2146, -0.1599,  0.2350, -0.2843, -0.0773, -0.2151,  0.1864, -0.3068,\n        -0.2093,  0.1365])
 'net2.weight' = {Tensor: 5} tensor([[ 0.1922, -0.0148, -0.1884,  0.2124, -0.1361,  0.0172, -0.2371,  0.1946,\n          0.2047, -0.2697],\n        [-0.2690,  0.1372,  0.2269,  0.0436, -0.1353, -0.2054, -0.2418, -0.2300,\n          0.1987,  0.0007],\n        [ 0.0995, -0.2659, -0.2374, -0
 'net2.bias' = {Tensor: 5} tensor([0.1488, 0.0791, 0.1667, 0.1449, 0.0545])