[源码解析] 深度学习流水线并行之PipeDream(1)--- Profile阶段
0x00 摘要
继 GPipe 之后，我们开一个流水线并行训练新系列，介绍微软的PipeDream。本文介绍其总体思路，架构和Profile阶段。
Gpipe流水线其存在两个问题：硬件利用率低，内存占用大。于是在另一篇流水并行的论文里，微软 PipeDream 针对这些问题提出了改进方法，就是1F1B （One Forward pass followed by One Backward pass）策略。这种改进策略可以解决缓存 activation 的份数问题，使得 activation 的缓存数量只跟 stage 数相关，从而进一步节省显存，可以训练更大的模型。
PipeDream 可以分为4个阶段：profile，compute partition，convert model，runtime等四个阶段。
- Profile阶段：通过小批量数据的profile推理出DNN训练时间。
- Compute Partition阶段：依据profile结果确定所有层的运行时间，然后进行优化，优化器返回一个带注释的运算符图，每个模型层映射到一个阶段ID。
- Convert model 阶段：对运算符图执行BFS遍历，为每个阶段生成一个单独的torch.nn.Module代码。PipeDream对每个阶段中的运算符进行排序，以确保它们保持与原始PyTorch模型图的输入输出依赖关系的一致性。
- Runtime 阶段：PipeDream运行时根据其1F1B-RR调度策略将每个阶段（包括复制阶段的副本）分配给单个工作进程。
本文首先看看 PipeDream 总体思路，架构和Profile阶段。
0x01 概述
1.1 前文回顾
前文提到，目前分布式模型训练有几个必要并行技术：
- 流水线并行，尤其是如何自动设定流水；
- 梯度累加（Gradient Accumulation）；
- 后向重计算；
- 1F1B 策略；
在前面几篇文章之中，我们介绍了Gpipe 如何前三种技术。本文开始，我们通过微软分布式DNNs训练系统PipeDream来看看其如何实现流水线并行和1F1B 策略。
1.2 目前问题
DNN 训练的特点是双向训练，训练在前向和后向通道中迭代进行计算，两种传播以相反顺序穿过相同层，在每轮迭代中，训练过程循环处理输入数据的一个 minibatch，并且更新模型参数。
1.2.1 数据并行
最常见的 DNN 并行化训练方法是数据并行化，这种方法把输入数据分散到各个 workers中（每个worker拥有全部模型）运行。不幸的是，尽管在加速数据并行的性能优化方面取得了一些进展，但若要在云基础设施上训练就会产生很大的通信开销。而且，随着GPU 计算速度的飞快提升，所有模型训练的耗时瓶颈会更进一步转向通信环节。
下图为数据并行，但是其对硬件利用率依然太低。因为数据并行化中的单个 worker 在交换梯度数据时不得不进行通信等待。
1.2.2 模型并行
模型并行化是另一种并行化训练形式，这种方法是把算子分散到各个 worker 上进行计算，这通常用于训练大型 DNN 模型。本文中，模型并行指的就是把模型的不同layer放在不同的机器上，不涉及将同一个layer切分到不同机器上的场景。
下图是模型并行化，显示了一个计算时间线，该示例有四台机器和一个管道。
- 在正向阶段，每个阶段对本阶段中的层的minibatch执行正向传递，并将结果发送到下一阶段。输出级在完成前向传递后，计算minibatch的损失。
- 在后向阶段，每个阶段形成后向通道，逐一将损失传播到前一阶段。
worker 之间只能同时处理一个 minibatch，系统中只有一个minibatch是活动的，这极大限制了硬件利用率。
而且，模型并行化需要程序员决定怎样按照给定的硬件资源来分割特定的模型，这其实无形之中加重了程序员的负担。
1.2.3 Gpipe
除了通用问题之外，GPipe 流水并行策略还有一个内存问题：需要缓存多份activation。
假如一个batch被分为 n 个micro-batch，则需要缓存 n 份activation。这个 n 是梯度累加的次数，为了尽可能流水，通常这个累加次数都比较大，一般大于两倍 stage 数目。那么即使只缓存少数 Tensor，这种策略依然需要较多显存。
0x02 论文
PipeDream 就针对这些问题提出了改进方法1F1B。PipeDream是第一个以自动化和通用的方式将流水线并行，模型并行和数据并行结合起来的系统。PipeDream首先使用模型并行对DNN进行划分，并将每层的子集分配给每个worker。但是与传统的模型并行不同，PipeDream对小批量数据进行流水线处理，实现了潜在的管道并行设计。在任何时刻，不同的worker处理不同的输入，从而保证了流水线的满负荷以及并行BSP。
微软在论文 PipeDream: Fast and Efficient Pipeline Parallel DNN Training 之中对于PipeDream进行了详细阐述，所以我们就基于此论文进行分析。
2.1 方案概述
2.1.1 并行方式
PipeDream 模型的基本单位是层，PipeDream将DNN的这些层划分为多个阶段。每个阶段（stage）由模型中的一组连续层组成。
PipeDream的主要并行方式就是把模型的不同层放到不同的stage之上，不同的的stage部署在不同的机器上，顺序地进行前向和反向计算，形成了一个pipeline。
每个阶段对该阶段中的所有层执行向前和向后传递。PipeDream将包含输入层的阶段称为输入阶段，将包含输出层的阶段称为输出阶段。但是每个stage可能有不同的replication，这就是数据并行。
对于使用数据并行的stage，采用 round-robin的方式将任务分配在各个设备上，需要保证一个batch的数据在前向和后向发生在同一台机器上，
2.1.2 1F1B
由于前向计算的 activation 需要等到对应的后向计算完成后才能释放（无论有没有使用 Checkpointing 技术），因此在流水并行下，如果想尽可能节省缓存 activation 的份数，就要尽量缩短每份 activation 保存的时间，也就是让每份 activation 都尽可能早的释放，所以要让每个 micro-batch 的数据尽可能早的完成后向计算，因此需要把后向计算的优先级提高，让 micro-batch 标号小的后向比 micro-batch 标号大的前向先做。因此，如果我们让最后一个 stage 在做完一次 micro-batch 的前向后，立马就做本 micro-batch 的后向，那么我们就能让其他的 stage 尽可能早的开始后向计算，这就是 1F1B 策略。
1F1B（one-forward-one-backward）的调度模式会在每台worker机器上交替进行小批次数据的前向后向计算，同时确保这些小批量在"后向传播"时可以路由到"前向传播"的相同worker。
这种方案可以使得每个GPU上都会有一个batch的数据正在被处理，使所有worker保持忙碌，而不会出现管道暂停，整个pipeline是比较均衡的。同时能确保以固定周期执行每个stage上的参数更新，也有助于防止同时处理过多小批量并确保模型收敛。
0x03 流水线
PipeDream的pipeline parallelism（PP）是一种新的并行化策略，它将批内并行与批间并行结合起来。
3.1 流水线改进
我们首先看看流水线对于模型并行的改进。
在上节的示例中，如果只有一个活动的minibatch，系统在任何给定的时间点最多有一个GPU处于活动状态。
但是我们希望所有GPU都处于活动状态。因此PipeDream 多个小批量逐一注入到流水线中，从而通过流水线来增强模型并行训练。在完成小批量的前向传递时，每个阶段都会异步地将输出激活发送到下一阶段，同时开始处理另一个小批量。类似地，在完成一个小批量的向后传递后，每个阶段都会将梯度异步发送到前一阶段，同时开始计算另一个小批量。
与普通层间并行训练相比，流水线有两个主要优点：
- 流水线通信量较少。流水线并行比数据并行的通信量要少得多。与数据并行方法（使用集体通信或参数服务器）中的做法（把所有参数聚合梯度并且将结果发送给所有worker）不同，流水线并行中的每个worker只需要在两个 stage 边界之间将梯度和输出激活的一个子集发送给另一个worker，这可能导致某些模型的通信量大幅减少。
- 流水线重叠了计算和通信。跨阶段前向输出激活和后向梯度的异步通信可以使得这些通信与后续小批量计算在时间上重叠，因为它们在不同的输入上运行，计算和通信完全独立，没有依赖边，所以可以更容易的并行化。在稳定理想状态下，所有的 workers 时刻都在运转，不像模型并行化训练中会有停下来等待的时候。
下图是实施了 1F1B 的流水线。Machine 1先计算 蓝色 1，然后把蓝色 1 发送给 Machine 2 继续计算。Machine 1 接着计算 蓝色 2。Machine 1 和 Machine 2 之间只传送模型的一个子集。计算和通讯可以并行。
3.2 挑战
PipeDream的目标是以最小化总体训练时间的方式将流水线并行，模型并行性和数据并行性结合起来。然而，要使这种方法对大型DNN模型有效，获得流水线并行化训练的潜在收益，PipeDream 必须克服几个主要挑战：
- 如何高效划分流水线。与处理器中的流水线一样，需要将DNN高效正确地划分为若干“阶段”（层序列），每个阶段部署在不同的worker上执行。
- 模型特质和硬件拓扑会降低效率，划分应该具体取决于模型体系结构和硬件部署。不好的划分（阶段性的工作量大范围倾斜）可能会导致worker长时间闲置。所以需要依据一定原则（通信和资源利用率）来划分，比如：彼此有通信的层应该分配到相邻的处理器；如果多个层操作同一数据结构，它们应该被分配到同一个处理器上，彼此独立的层可以映射到不同处理器上。所以分配算法也必须考虑模型特质和硬件拓扑。
- 机器间的过度通信会降低硬件效率。
- 在确保训练任务向前推进的同时，如何调度计算以最大化吞吐量。
- 如何防止流水线瓶颈。
- 由木桶原理我们可以知道，在稳定状态下，一个流水线管道的吞吐量由这个流水线上最慢环节的吞吐量决定。如果各个环节的处理能力彼此差别很大，会导致管道中出现空闲时间（一半称之为bubble），这样最快环节必须停下来等待其他环境，会造成饥饿现象，从而导致资源利用不足。所以需要确保流水线中所有阶段都大致花费相同的计算时间，否则最慢的阶段将会成为整个流水线的瓶颈。
- 如何在不同的输入数据之间调度工作以均衡流水线。
- 与传统的单向流水线管道不同，DNN训练是双向的：前向传播和后向传播，两种传播以相反顺序穿过相同层。如何协调流水线工作是一个问题。
- 面对流水线带来的异步性，如何确保训练有效。
- 流水线带来的一个问题就是weight版本众多。在后向传播时候如果使用比前向传播时更高版本的weight来计算，则会造成训练模型质量降低。
- PipeDream管理后向通道里的权重版本，通过为每个小批量的weight维护版本号来解决这个问题，这样在后向通道里使用的权重版本就和前向通道里使用的相同，从而在数值上能够正确计算梯度（我们后续文章会讲解）。
3.4 流水线划分算法
PipeDream基于一个短期运行分析结果来自动划分DNN的层，依据分析结果使用算法来对不同阶段之间的计算负载进行平衡，同时最小化通信。PipeDream的自动划分算法总体目标是输出一个平衡的管道，确保每个阶段大致执行相同的总工作量。同时还必须确保各阶段之间通信的数据量尽可能小，以避免通信中断。算法如下：
- 将DNN层划分为多个阶段，以便每个阶段以大致相同的速率完成，即花费大致相同的计算时间。
- 尝试以拓扑感知的方式尽量减少worker之间的通信（例如，如果可能，向更高带宽的链路发送较大的输出）。
- 因为DNN并不总可以在可用的workers做平均分配，为了进一步改进负载平衡，PipeDream允许复制一个stage，即在这个stage上使用多个worker进行数据并行。这样多个worker可以分配到流水线同一阶段，并行处理一个batch的不同的mini-batch，提高处理效率。因为数据并行采用了RR，所以这套策略也被称为 1F1B-RR(one-forward-noe-backward-round-robin)。
这个划分问题等价于最小化流水线的最慢阶段所花费的时间，并且具有最优子问题属性：在给定worker工作量前提下，吞吐量最大化的流水线由一系列子流水线构成，其中每一个子流水线针对较小worker工作量来最大化自己的输出。因此，PipeDream使用动态规划来寻找最优解。
具体如下图：
3.5 Profile
DNN训练有一个特点：不同输入的计算时间几乎没有变化。于是 PipeDream充分利用了这一事实，给定一个具有N层和M台可用机器的DNN，PipeDream首先在一台机器上分析模型，记录向前和向后过程所花费的计算时间，层输出的大小以及每个层的相关参数的大小，最后输出为一个结果文件。
分区算法不但使用profile结果文件作为输入，而且还考虑了其他限制，如硬件拓扑和带宽、工人数量和计算设备的内存容量，最终将层分为多个阶段，同时还确定每个阶段的复制因子，以最小化模型的总训练时间。
所以总体算法大致如下：
因为PipeDream借鉴了很多GPipe的思路，所以可以看到其比Gpipe的进步之处。
比如 Gpipe是通过在代码中硬性预估ops来进行流水线负载均衡，PipeDream则是先做profile，根据实际情况再做推理。
0x04 Profile阶段
Profile是PipeDream工作的第一个阶段，是分区算法的基础。PipeDream根据profiling的结果，使用动态规划对模型进行划分，将模型划分为不同的stage，以及每个stage的replication数。
这是PipeDream针对GPipe的一个改进，两者都是对每层的运行时间进行预估，然后对模型进行划分。
- GPipe是利用经验值或者数学的方法来对运行时间进行预估。
- PipeDream根据profiling的结果对运行时间进行预估。
因为有实际数据进行支撑，所以PipeDream更加准确和先进。
4.1 思路
评测机制利用了这样一个事实：DNN训练在计算和通信时间上几乎没有变化。所以我们可以通过小批量数据的profile推理出DNN训练时间。为了确定所有层的运行时间，PipeDream在其中一台机器上使用1000个小批量对DNN模型的短期（几分钟）运行进行 profile。
4.1.1 如何计算
运行时间
对于每一层的运行时间，我们可以通过 运行时间 = 计算时间 + 通信时间 来得到。
- 计算时间就是每层layer前向和后向的计算时间，这个可以从profile得出。
- 通信时间就需要根据模型大小进行估算，PipeDream 估计通信所需的时间为"需要传输的数据量"除以"通信链路上的带宽"。
通信时间
在流水线上，大多数通信都有三个步骤：
1）在发送端机器上，从GPU传输到CPU移动数据。
2）通过网络从发送者到接收者发送数据。
3）在接收端，从CPU到GPU移动数据。
而 通过网络从发送者到接收者发送数据 是最耗时的，所以PipeDream主要考虑这个因素。如果再对这个因素细分，则有：
- 对于从层 i 到 层 i + 1 传输激活值的时间，PipeDream 基于 "激活值"来估计。
- 假如配置成了数据并行（对于 层 i 使用 m 个 worker 做数据并行）的情况，做权重同步的时间使用"权重"来估计：
- 如果使用分布式参数服务器，则权重数量被预估为 4 x ( m - 1 ) x | w i | / m。
- 如果使用 all_reduce，则每个worker给其他workers发送 ( m - 1 ) x | w i | / m 个 bytes，也接受到同样数量字节。
4.1.2 Profile内容
综上所述，PipeDream在profile之中，为每个层 i 记录三个数量：
- Ti，层 i 的在GPU上向前和向后计算时间之和，即每层layer前向和后向的计算时间；
- ai，层 i 的输出激活的大小（以及向后过程中输入梯度的大小）以字节为单位，即每层layer的输出的大小；
- wi，层 i 的权重参数的大小（以字节为单位），即每层layer参数的大小；
4.2 代码
不同模型或者说不同领域有不同的profile文件。
我们以 profiler/translation/train.py 为入口进行分析。
4.2.1 训练脚本
以下我们省略了无关代码。
4.2.1.1 训练过程
class Seq2SeqTrainer:
def feed_data(self, data_loader, training=True):
"""
Runs training or validation on batches from data_loader.
:param data_loader: data loader
:param training: if True runs training else runs validation
"""
# 白名单
module_whitelist = ["EmuBidirLSTM", "RecurrentAttention", "Classifier"]
# 样本集
for i, (src, tgt) in enumerate(data_loader):
break
(src, src_length) = src
(tgt, tgt_length) = tgt
src_length = torch.LongTensor(src_length).cuda()
src = src.cuda()
tgt = tgt.cuda()
model_input = (src, src_length, tgt[:-1])
# 使用torchsummary计算网络的计算参数等信息
summary = torchsummary.summary(model=self.model, module_whitelist=module_whitelist,
model_input=model_input, verbose=True)
for i, (src, tgt) in enumerate(data_loader):
if training and i in eval_iters:
test_bleu, _ = self.translator.run(calc_bleu=True,
epoch=self.epoch,
iteration=i)
# 训练模型
self.model.train()
self.preallocate(data_loader, training=True)
# 从模型建立图
if training:
create_graph(self.model, module_whitelist, (src, tgt), summary,
os.path.join("profiles", self.arch))
4.2.1.2 计算参数
上节在训练脚本制作的时候，torchsummary 的作用是计算网络的计算参数等信息，对于 torchsummary 我们举例如下：
import torch
import torch.nn as nn
from torchsummary import summary
class SimpleConv(nn.Module):
def __init__(self):
super(SimpleConv, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
)
def forward(self, x, y):
x1 = self.features(x)
x2 = self.features(y)
return x1, x2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleConv().to(device)
summary(model, [(1, 16, 16), (1, 28, 28)])
其打印如下：
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 1, 16, 16] 10
ReLU-2 [-1, 1, 16, 16] 0
Conv2d-3 [-1, 1, 28, 28] 10
ReLU-4 [-1, 1, 28, 28] 0
================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.77
Forward/backward pass size (MB): 0.02
Params size (MB): 0.00
Estimated Total Size (MB): 0.78
----------------------------------------------------------------
4.2.1.3 创建图
create_graph 的作用就是使用torchgraph.GraphCreator创建一个图，这个图就可以理解为模型内部的DAG图，每个节点记录如下信息。
node10 -- Dropout(p=0.2) -- forward_compute_time=0.064, backward_compute_time=0.128, activation_size=6291456.0, parameter_size=0.000
具体代码如下：
def create_graph(model, module_whitelist, model_input, summary, directory):
"""Given a model, creates and visualizes the computation DAG
of the model in the passed-in directory."""
# 创建图
graph_creator = torchgraph.GraphCreator(model, summary, module_whitelist
# 构建hook
graph_creator.hook_modules(model, root=True)
(src, tgt) = model_input
(src, src_length) = src
(tgt, tgt_length) = tgt
src_length = torch.LongTensor(src_length).cuda()
src = src.cuda()
tgt = tgt.cuda()
# 运行以得到profile
model(src, src_length, tgt[:-1])
graph_creator.unhook_modules()
# 输出profile结果
graph_creator.persist_graph(directory)
4.2.2 创建图
创建图基本是在 GraphCreator 内完成。
class GraphCreator(object):
def __init__(self, model, summary, module_whitelist):
if isinstance(model, torch.nn.Module) is False:
raise Exception("Not a valid model, please provide a 'nn.Module' instance.")
self.model = model
self.module_whitelist = module_whitelist
self.summary = copy.deepcopy(summary)
self.forward_original_methods = {}
self.graph = graph.Graph()
self.inputs = {}
4.2.2.1 设置wrapper
hook_modules 的作用是给模型的forward函数设置一个wrapper，并且遍历为子模块设置，这样在模型运行时候可以跟踪模型之间的联系。
def hook_modules(self, module, root=False):
this_creator = self
sub_modules = module.__dict__['_modules']
# Wrapper function to "forward()", keeping track of dependencies.
def forward_wrapper(self, *wrapped_inputs):
input = []
wrapped_inputs_list = list(wrapped_inputs)
for i in range(len(wrapped_inputs_list)): # 遍历输入
if isinstance(wrapped_inputs_list[i], TensorWrapper):
# 如果已经被包装，则插入input
input.append(wrapped_inputs_list[i].tensor)
else:
key = wrapped_inputs_list[i]
if key in this_creator.inputs: # 如果是原始输入，则不进行包装
wrapped_inputs_list[i] = this_creator.inputs[key]
else:
j = len(this_creator.inputs)
# 如果没有被wrap, 则构建一个TensorWrapper进行包装
wrapped_inputs_list[i] = TensorWrapper(wrapped_inputs_list[i],
"Input%d" % j, this_creator)
this_creator.inputs[key] = wrapped_inputs_list[i]
input.append(wrapped_inputs_list[i].tensor) # 则插入input
result = this_creator.forward_original_methods[self](*input)
# 对结果进行包装
wrapped_result = TensorWrapper(result, str(self), this_creator)
# 把边添加进入图
for wrapped_input in wrapped_inputs_list:
this_creator.graph.add_edge(wrapped_input.node(), wrapped_result.node())
return wrapped_result
# Wrapper function to "forward()", keeping track of dependencies.
def forward_wrapper_root(self, *wrapped_inputs):
input = []
wrapped_inputs_list = list(wrapped_inputs)
for i in range(len(wrapped_inputs_list)):
if isinstance(wrapped_inputs_list[i], TensorWrapper):
input.append(wrapped_inputs_list[i].tensor)
else:
key = wrapped_inputs_list[i]
if key in this_creator.inputs:
wrapped_inputs_list[i] = this_creator.inputs[key]
else:
j = len(this_creator.inputs)
wrapped_inputs_list[i] = TensorWrapper(wrapped_inputs_list[i],
"Input%d" % j, this_creator)
this_creator.inputs[key] = wrapped_inputs_list[i]
input.append(wrapped_inputs_list[i].tensor)
result = this_creator.forward_original_methods[self](*input)
return result
# 遍历子模块，递归设置wrapper
for name, sub_module in sub_modules.items():
# nn.Module is the only thing we care about.
if sub_module is None or isinstance(sub_module, torch.nn.Module) is False:
break
sub_module_name = sub_module.__class__.__name__
sub_sub_modules = sub_module.__dict__['_modules']
if len(sub_sub_modules) == 0 or sub_module_name in self.module_whitelist:
sub_module.reset_hooks()
#
# Hook nn.Module with no descendants.
#
# Replace "forward" with "wrapped_forward".
# 使用wrapped_forward替换forward
if sub_module not in this_creator.forward_original_methods:
this_creator.forward_original_methods.update({sub_module:
sub_module.forward})
sub_module.forward = forward_wrapper.__get__(sub_module, sub_module.__class__)
if len(sub_sub_modules) >forward_compute_time 0 and sub_module_name not in self.module_whitelist:
#
# Recursively visit this module's descendants.
# 递归设置wrapper
self.hook_modules(sub_module)
if root: # 对于root进行处理
this_creator.forward_original_methods.update({module: module.forward})
module.forward = forward_wrapper_root.__get__(module, module.__class__)
4.2.2.2 TensorWrapper
TensorWrapper 就实现了wrapper功能，graph_creator.summary 就是之前torchsummary.summary得到的网络等信息。可以看到此类会遍历 summary，计算 forward_compute_time 等信息，最终构建了一个 node。
需要注意的是：activation_sizes 是根据 output_shape 来计算的。
class TensorWrapper(object):
def __init__(self, tensor, node_desc, graph_creator, activation_size=None):
self.tensor = tensor
global object_id
self.object_id = object_id
object_id += 1
self.node_desc = node_desc
i = 0
for i in range(len(graph_creator.summary)):
if str(graph_creator.summary[i]['layer_name']) == node_desc:
break
if i < len(graph_creator.summary) and node_desc == str(graph_creator.summary[i]['layer_name']):
summary_elem = graph_creator.summary.pop(i)
forward_compute_time = summary_elem['forward_time']
backward_compute_time = summary_elem['backward_time']
if isinstance(summary_elem['output_shape'][0], list):
activation_sizes = [4.0 * functools.reduce(lambda x, y: x * y, elem)
for elem in summary_elem['output_shape']]
else:
activation_sizes = 4.0 * functools.reduce(lambda x, y: x * y, summary_elem['output_shape'])
parameter_size = 4.0 * float(summary_elem['nb_params'])
self._node = graph.Node("node%d" % object_id, node_desc=node_desc,
forward_compute_time=forward_compute_time,
backward_compute_time=backward_compute_time,
activation_size=activation_sizes,
parameter_size=parameter_size)
elif activation_size is not None:
self._node = graph.Node("node%d" % object_id, node_desc=node_desc,
activation_size=activation_size)
else:
self._node = graph.Node("node%d" % object_id, node_desc=node_desc)
self.graph_creator = graph_creator
对于某些内置方法，则也会相应处理，比如如下。
def __iadd__(self, other):
self_activation_size = self.node().activation_size
other_activation_size = other.node().activation_size
assert(self_activation_size == other_activation_size)
wrapped_result = TensorWrapper(self.tensor, "Add(inplace)", self.graph_creator,
activation_size=self_activation_size)
self.tensor += other.tensor
self.graph_creator.graph.add_edge(self._node, wrapped_result.node())
self.graph_creator.graph.add_edge(other.node(), wrapped_result.node())
return wrapped_result
最终对应：
node58 -- Add(inplace) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=102760448.000, parameter_size=0.000
4.2.3 持久化
persist_graph 就是把profile结果输出到文件。
def persist_graph(self, directory):
self.graph.to_dot(os.path.join(directory, "graph.dot"))
with open(os.path.join(directory, "graph.txt"), 'w') as f:
f.write(str(self.graph))
self.graph.render_bar_graphs_and_cdfs(directory)
具体调用了 graph.py 的函数完成，这里摘录 to_dot函数如下：
def to_dot(self, arch):
dot = graphviz.Digraph()
for node in self.nodes.values():
node_desc = "%s
[forward_compute_time=%.3f,backward_compute_time=%.3f,activation_size=%s,parameter_size=%.1f]" % (
node.node_desc, node.forward_compute_time, node.backward_compute_time,
node.activation_size, node.parameter_size)
if node.stage_id is not None:
color = self._colors[node.stage_id % len(self._colors)]
dot.node(node.node_id, node_desc,
color=color, style='filled')
else:
dot.node(node.node_id, node_desc)
for node in self.nodes.values():
if node.node_id not in self.edges:
continue
for out_node in self.edges[node.node_id]:
dot.edge(node.node_id, out_node.node_id)
dot.render(arch)
4.3 结果
我们使用源码中的结果为例 pipedream-pipedream/profiler/translation/profiles/gnmt/graph.txt，给大家展示下具体结果。
node1 -- Input0 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node4 -- Embedding(32320, 1024, padding_idx=0) -- forward_compute_time=0.073, backward_compute_time=6.949, activation_size=6291456.0, parameter_size=132382720.000
node5 -- EmuBidirLSTM( (bidir): LSTM(1024, 1024, bidirectional=True) (layer1): LSTM(1024, 1024) (layer2): LSTM(1024, 1024)) -- forward_compute_time=5.247, backward_compute_time=0.016, activation_size=12582912.0, parameter_size=67174400.000
node2 -- Input1 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node6 -- Dropout(p=0.2) -- forward_compute_time=0.077, backward_compute_time=0.196, activation_size=12582912.0, parameter_size=0.000
node7 -- LSTM(2048, 1024) -- forward_compute_time=3.190, backward_compute_time=5.348, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=50364416.000
node8 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node9 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node10 -- Dropout(p=0.2) -- forward_compute_time=0.064, backward_compute_time=0.128, activation_size=6291456.0, parameter_size=0.000
node11 -- LSTM(1024, 1024) -- forward_compute_time=2.491, backward_compute_time=4.203, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=33587200.000
node12 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node13 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node14 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node15 -- Dropout(p=0.2) -- forward_compute_time=0.059, backward_compute_time=0.121, activation_size=6291456.0, parameter_size=0.000
node16 -- LSTM(1024, 1024) -- forward_compute_time=2.492, backward_compute_time=4.201, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=33587200.000
node17 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node18 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node19 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node3 -- Input2 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node21 -- Embedding(32320, 1024, padding_idx=0) -- forward_compute_time=0.066, backward_compute_time=0.328, activation_size=6291456.0, parameter_size=132382720.000
node20 -- hidden -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node22 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node23 -- RecurrentAttention( (rnn): LSTM(1024, 1024) (attn): BahdanauAttention( (linear_q): Linear(in_features=1024, out_features=1024, bias=False) (linear_k): Linear(in_features=1024, out_features=1024, bias=False) (dropout): Dropout(p=0) ) (dropout): Dropout(p=0)) -- forward_compute_time=4.546, backward_compute_time=6.141, activation_size=[6160384.0; 131072.0; 131072.0; 6160384.0; 288768.0], parameter_size=41979904.000
node24 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node25 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node26 -- __getitem__(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node27 -- __getitem__(3) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node28 -- Dropout(p=0.2) -- forward_compute_time=0.058, backward_compute_time=0.176, activation_size=6160384.0, parameter_size=0.000
node29 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node30 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node31 -- LSTM(2048, 1024) -- forward_compute_time=3.151, backward_compute_time=5.288, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node32 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node33 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node34 -- Dropout(p=0.2) -- forward_compute_time=0.061, backward_compute_time=0.174, activation_size=6160384.0, parameter_size=0.000
node35 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node36 -- __getitem__(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node37 -- LSTM(2048, 1024) -- forward_compute_time=3.145, backward_compute_time=5.306, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node38 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node39 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node40 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node41 -- Dropout(p=0.2) -- forward_compute_time=0.055, backward_compute_time=0.198, activation_size=6160384.0, parameter_size=0.000
node42 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node43 -- __getitem__(3) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node44 -- LSTM(2048, 1024) -- forward_compute_time=3.149, backward_compute_time=15.883, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node45 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node46 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node47 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node48 -- Classifier( (classifier): Linear(in_features=1024, out_features=32320, bias=True)) -- forward_compute_time=5.609, backward_compute_time=1.227, activation_size=194437120.0, parameter_size=132512000.000
node1 -- node4
node4 -- node5
node2 -- node5
node5 -- node6
node6 -- node7
node7 -- node8
node7 -- node9
node8 -- node10
node10 -- node11
node11 -- node12
node11 -- node13
node12 -- node14
node8 -- node14
node14 -- node15
node15 -- node16
node16 -- node17
node16 -- node18
node17 -- node19
node14 -- node19
node3 -- node21
node20 -- node22
node21 -- node23
node22 -- node23
node19 -- node23
node2 -- node23
node23 -- node24
node23 -- node25
node23 -- node26
node23 -- node27
node24 -- node28
node28 -- node29
node26 -- node29
node20 -- node30
node29 -- node31
node30 -- node31
node31 -- node32
node31 -- node33
node32 -- node34
node34 -- node35
node26 -- node35
node20 -- node36
node35 -- node37
node36 -- node37
node37 -- node38
node37 -- node39
node38 -- node40
node32 -- node40
node40 -- node41
node41 -- node42
node26 -- node42
node20 -- node43
node42 -- node44
node43 -- node44
node44 -- node45
node44 -- node46
node45 -- node47
node40 -- node47
node47 -- node48
至此，我们知道了Profile阶段的内容，就是：运行训练脚本，依据运行结果来计算参数，建立一个模型内部的DAG图，然后把参数和DAG图持久化到文件之中，后续阶段会使用这个文件的内容。
下一篇我们分析如何计算自动分区。
0xFF 参考
Tensorflow实现先累加多个minibatch计算的梯度，再反向传播
PipeDream: Fast and Efficient Pipeline Parallel DNN Training
论文解读系列第五篇：微软斯坦福等PipeDream快速训练大规模神经网络
https://cs231n.github.io/neural-networks-3/#gradcheck
https://www.cnblogs.com/geekfx/p/14182048.html
训练时显存优化技术——OP合并与gradient checkpoint
Pytorch笔记04-自定义torch.autograd.Function
pytorch的自定义拓展之（三）——torch.autograd.Function的简单定义与案例
pytorch的自定义拓展之（二）——torch.autograd.Function完成自定义层