理解 PyTorch 分布式 Autograd 设计

Autograd 是一个反向自动微分系统(或梯度计算引擎),基于记录所有的操作来构建一个有向无环图——Autograd 计算图,其中叶子节点是输入 Tensor,根节点 root 是输出 Tensor,通过跟踪图中从根节点 root 到叶子节点的路径上的操作,能够自动地计算出梯度。 在 PyTorch 中,模型训练的每一轮迭代,都会创建对应的 Autograd 计算图:在前向传播阶段动态地创建 Autograd 计算图,在反向传播阶段根据 Autograd 计算图来进行梯度的计算。 构建分布式 Autograd 计算图 对于分布式模型训练环境下,需要在各个节点(主机)之间进行大量的 RPC 调用,统一协调各个过程来完成模型的训练。PyTorch 实现的分布式 Autograd,在前向传播过程中构建 Autograd 计算图,并且基于 Autograd 计算图在反向传播过程中计算梯度。在前向传播过程中,PyTorch 持续跟踪各个 RPC 调用的情况,必须确保在反向传播过程中计算是正确的,所以 PyTorch 在实现过程中使用了 send、recv 这一对函数来进行跟踪,当执行 RPC 调用时将 send 和 recv 绑定到 Autograd 计算图上。 send 函数被绑定到 RPC

理解 PyTorch DDP 分布式数据并行模式

PyTorch 使用 DDP(Distributed Data Parallel) 实现了真正的分布式数据并行,在下面的两个场景下都可以使用 DDP 实现模型的分布式训练: 单机、多 GPU(单进程多线程的伪分布式) 多机、多 GPU(多机多进程的真正分布式) 上面第一种方式,就是类似使用简单的 DP 数据并行模式,但是 DP 使用的单进程、多线程的范式来实现的;而 DDP 完全使用了多进程的方式,包括单机多进程、多机多进程,如果是多机的情形则对应着物理上的分布式多进程模式。为了获得更好的性能,最好是使用 DDP 模式来训练模型,即使是在单机、多 GPU 的情况下,也建议使用 DDP 模式来实现基于数据并行的模型训练,使用单机 DDP 模式训练模型的性能要比 DP 模式好很多。 DDP 基于集合通信(Collective Communications)来实现分布式训练过程中的梯度同步。在反向传播过程中,DDP 使用 AllReduce 来实现分布式梯度计算和同步。 下面,我们从与 DDP 相关的几个方面来理解 DDP 的设计与实现,包括: 集合通信(Collective Communication) 通信后端(Communication Backend) DDP 内部实现概览

PyTorch 使用 DP 模式实现数据并行

PyTorch 使用 torch.nn.DataParallel 来实现基于数据并行的单机、多 GPU 并行训练。使用 DP 方式在多 GPU 上训练模型,需要保证模型能够在每个 GPU 上放得下,训练过程中会把整个模型都复制到每个 GPU 内,通过数据并行的方式提高训练模型的速度。 虽然 DataParallel 使用起来非常容易,但是通常不能够获得最佳的性能。DataParallel 在每一轮的前向传播过程中,会复制一遍模型,同时这种基于单进程、多线程的并行方式也会存在 GIL(Global Interpreter Lock) 竞争问题。 DP 数据并行训练流程 下面我们分析一下 DP 数据并行模式在多 GPU 的情况下训练模型的基本流程,如下图所示: 基于 DP 模式,模型训练的基本过程分为三个阶段,描述如下: 前向传播计算过程 1.Scatter mini-batch inputs to GPUs 我们通过指定 batch_size 大小,对输入的训练数据集进行了分割,在训练神经网络模型的前向传播过程中,会将每一个小批数据(Mini-Batch)分发到每一个 GPU 上。具体过程是:将 4 个小批数据 i1 ~ i4 复制到 GPU-1 上,再将 4 个小批数据分别发送到 GPU-1 ~ GPU-4

PyTorch 流水线并行模式设计分析

流水线并行(Pipeline Parallelism)最早是 Google 在 Gpipe 论文中提出的,这种并行训练模式能够充分利用多 GPU 的资源高效地训练评估大模型。目前 PyTorch 最新版本是 2.2,流水线并行的功能是基于 torchgpipe 论文中的设计来实现的,该功能当前还处于试验阶段。 问题背景 大模型无法直接放到单个 GPU 中进行训练,通过模型并行(Model Parallelism)的方法可以把模型进行分片,每一个分片放置到一个 GPU 上,这样能够很好实现模型并行且利用多 GPU 的资源。虽然使用这种较为初级的方式能够实现大模型的训练,但在训练的过程中并不能充分利用 GPU 资源,因为对顺序(Sequential)模型来说它每次只能激活一个 GPU 来进行训练,其它的 GPU 此时是闲置的,所以在底层设备上其实仍然是顺序执行。 例如,对一个有 4 层的顺序(Sequential)神经网络模型,经过模型分片后,训练过程中每一层(或 Subnetwork)放在一个 GPU 上,先进行前向传播计算得到 Loss,然后反向传播计算梯度,如下图所示: 使用这种方式利用 GPU 训练,我们可以看到在训练过程中 GPU 完全没有被充

基于 PyTorch 编程使用预训练模型

使用预训练模型有两种方式:一种是直接使用得到的预训练模型进行推理,并满足应用的需要,使用起来非常简单;另一种是在预训练模型的基础上,进行微调,使得到的新模型能够更好地满足我们解决问题的需要,这种方式需要能够对模型进行调优有一定门槛。这里,我们尝试第一种方式直接使用预训练模型,着重关注使用预训练模型处理图片分类的过程,从而熟悉在实际应用中都需要做哪些处理工作。 预训练模型 预训练模型(Pre-trained Models,PTMs)是一种深度学习架构,它在大规模数据集上进行训练,以获取丰富的特征表示。训练得到的模型可以进行复用,不仅能够适用于最初要解决的问题,还可以迁移到其他类似的应用场景中,从而提高在这些新领域的应用的性能。 预训练模型通常具有较大的参数规模,需要使用海量的数据和高昂的计算资源代价,才能完成模型训练并最终得到模型参数,这对于一些不具备基于超大规模数据训练能力的使用者来说,就无法发挥模型的作用,而且也不能很方便地在特定应用领域内探索并验证一些应用的想法。 例如,在 NLP 领域,预训练模型应用的特别广

使用 PyTorch SWA 优化模型训练入门实践

PyTorch 实现了 SWA(Stochastic Weight Averaging,随机加权平均),相比于传统的 SGD,使用 SWA 能够明显改善一些深度神经网络模型的测试精度(Test Accuracy)。而且,SWA 使用起来非常简单,能够加速模型训练,并提高模型的泛化能力。 SWA 基本原理 SWA 依赖两个重要的因素: 第一个是,SWA 使用一个不断修改的 LR 调节器(Learning Rate Schedule),使得 SGD 能够在最优值附近进行调整,并评估最优解附近的值对应的模型的精度,而不是只选取最优解对应的模型。因为,最优解对应的模型不一定是最优的,而且泛化能力可能也不一定最好。比如,在 75% 的训练时间里,可以使用一个标准的衰减学习率(Decaying Learning Rate)策略,然后在剩余 25% 的训练时间里将学习率设置为一个比较高的固定值。如下图所示: 第二个是,SWA 计算的是 SGD 遍历过的神经网络权重的平均值。例如,上面提到模型训练的后 25% 时间,我们可以在这 25% 时间里的每一轮训练(every epoch)后,计算一个权重的 running 平均值,在训练结束后再设置网络模型的权重为 SWA 权重平均值。 SWA

PyTorch 使用 TensorBoard 实现可视化

在 PyTorch 中使用 TensorBoard,可以实现样本数据、模型、训练过程的可视化,能够非常直观地查看在整个训练过程中产生的效果,方便分析和解决一些问题。本文完全根据 PyTorch 官方的 TensorBoard Tutorial 文档进行实践操作,体验 TensorBoard 的基本可视化功能。 首先,需要下载 tensorboard 模块: pip3 install tensorboard 这样就可以使用 TensorBoard 了,需要通过如下命令,启动 TensorBoard 服务: tensorboard --logdir=runs 上面 runs 目录是我们写数据的目录,可以根据自己的需要设置目录名称。在使用 PyTorch 过程中需要可视化的话,就把对应的数据写入到这个目录下面的指定文件中,TensorBoard 会直接从该目录下读取并进行可视化。 启动成功后,可以打开浏览器链接 http://localhost:6006/,查看 TensorBoard UI 界面。 我们向 runs 目录中写入数据,直接使用 torch.utils.tensorboard.SummaryWriter 就可以实现,所以先要创建可用的 SummaryWriter 对象,代码如下: from torch.utils.tensorboard import SummaryWriter writer = Summar

PyTorch 分布式训练模式 FSDP 设计分析

在 AI 大模型训练场景中,数据是海量的,模型也是超大的,对于训练大模型会带来很大挑战,比如对算力的需求,对处理大模型的工程复杂度,等等。PyTorch 给出了一种实现方式——FSDP(Fully Sharded Data Parallel),它提供了易用的 API,可以非常方便地解决大模型分布式训练的难题。 FSDP 是在 DDP(DistributedDataParallel)的基础上提出的,首先我们了解一下 PyTorch 的 DDP(DistributedDataParallel) 训练模式的一些特点: 在 DDP 中,核心的能力还是训练数据并行(Data Parallel)。以多机多卡方式为例,每个 process/worker 都会持有模型的一个副本(Replica),通过使每个 process/worker 处理一个 batch 的数据试下并行处理,最后使用 all-reduce 操作对多个不同 process/worker 计算得到的梯度进行累加求和;接着,再将优化器状态、梯度通过跨多个 process/worker 进行复制,使得每个 process/worker 上的模型参数都得到同步更新。也就是说,在 DDP 中并没有实现对模型参数的分片管理,即模型并行(Model Parallel)。 在 FSDP 中实现了模型的分片管理能

使用 PyTorch 实现并训练 LeNet-5 模型

LeNet-5 是由 Yann LeCun提出的卷积神经网络,在论文《Gradient-Based Learning Applied To Document Recognition》中可以看到 LeNet-5 网络模型的结构,如下图所示: 通过上图可以看到,从左至右网络各个层顺序连接: 输入层 :图片大小 32×32 卷积层1 :输入通道 1,输出通道 6,卷积核大小 5×5,步长 1 池化层 :输入通道 6,输出通道 6,过滤器大小 2×2,步长 2 卷积层2 :输入通道 6,输出通道 16,卷积核大小 5×5, 步长 1 池化层2 :输入通道 16,输出通道 16,过滤器大小 2×2,步长 2 全连接层1:节点数 120 全连接层2:节点数 84 全连接层3:节点数 10 我们只需要准备好数据集,并基于上图连接结构,使用 PyTorch 搭建 CNN 网络的结构并进行训练和使用。 实现 LeNet-5 模型 基本环境配置如下: Python:3.11.3 PyTorch:2.0.1(torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2) 1 准备数据集 使用经典的手写数字数据集 MNIST,可以直接通过 PyTorch 的 datasets.MNIST 下载并准备数据: import torch from to

使用 PyTorch 构建机器学习应用

通过 PyTorch 官网给出的 Quickstart,了解使用 PyTorch 完成一个模型的数据准备、模型训练和评估、模型加载并应用。在实际应用中,只需要按照这个流程来编程构建即可。 下面,我们通过分步骤来说明开发机器学习应用程序的基本流程。 我们使用 PyTorch-2.0.1,安装该版本 PyTorch 执行如下命令: pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 1 准备数据 PyTorch 使用 torch.utils.data.DataLoader 和 torch.utils.data.Dataset 来实现数据的加载,并转换成包含样本和标签的 DataSet。 我们可以在 https://pytorch.org/vision/stable/datasets.html 中找到 Torchvision 提供的大量内置 DataSet,通过这些工具类就可以方便构建并使用 DataSet。如果是使用我们自己的数据集,可以继承自 torch.utils.data.Dataset 实现我们自己的 DataSet 以及 DataLoader。 这里使用了 FashionMNIST 数据集,这个数据集是一个包含时尚衣物图像,其中包括衣物图像和它们对应的标签. 首先,下载 FashionMNIST 数据集,可以直接通过 PyTorch 的 datasets API 下