后向重计算在OneFlow中的实现：以时间换空间，大幅降低显存占用

撰文 | 赵露阳

2016年，陈天奇团队提出了亚线性内存优化相关的“gradient/activation checkpointing（后向重计算）”等技术[1]，旨在降低深度学习训练过程中的中间激活(activation)带来的显存占用。Checkpointing技术属于亚线性内存优化的一种，除此之外还有CPU offload等技术（CPU offload在微软Deepspeed框架中被广泛使用）。 

CPU offload将暂时用不到的GPU内存换入到CPU内存中存储，待需要时再取出，主要开销是来自于CPU和GPU之间的拷贝，会占用传输带宽（PCIE带宽），属于以传输换空间。而Checkpointing核心在于以时间换空间：通过计算图分析技术来实施Inplace操作以及内存共享优化（Memory sharing），在每个mini-batch的前向过程中删除一些暂时用不到的中间激活特征以降低内存占用，并在后向过程中需要时借助额外的前向计算恢复它们。

在OneFlow中，Checkpointing的实现主要是通过静态内存复用的方式，前向Tensor的生命周期结束后，其余Tensor可以复用这块内存，从而起到内存复用、节省内存的效果。

OneFlow目前支持了“gradient/activation checkpointing”（后向重计算）以实现亚线性内存优化，且对算法开发者非常友好，使用方式很简单：针对需要优化的网络部分，用一行代码将其包裹在“Checkpointing”的scope范围内即可，系统内部会针对此scope区域内的网络做分析并在训练过程中自动进行Checkpointing内存优化。

本文主要内容为以下3点：

• 1.亚线性内存优化的用法
• 2.亚线性内存优化的设计
• 3.代码解读

其中：1.将介绍如何在OneFlow中开启试用亚线性内存优化；2.将介绍OneFlow中亚线性内存优化是如何设计的及其原理；3.将从代码入手，剖析具体实现过程。

1

亚线性内存优化的用法

OneFlow中开启亚线性内存优化的方式如下：

# 用法：
with flow.experimental.scope.config(checkpointing=True):
# your net work, such as :
# input layernorm
norm1 = layernorm("layernorm_1", h)
# attention
h = h + self.attn(norm1)
# output layernorm
norm2 = layernorm("layernorm_2", h)
# mlp
h = h + self.mlp(norm2)

用上述代码包裹后，此scope区域内的网络，在整个前向过程中只会保存一份input tensor的内存，从input到最后输出h，这之间所有中间特征tensor的内存都不会被保存，后向过程需要时从input开始进行（前向的）重计算。

我们在多个网络上进行了开启/关闭checkpointing的显存占用测试，以GPT-2为例，具体是在每个Transformer Layer内都使用checkpointing = True scope标记重计算的部分。

可以看见，开启checkpointing后会大幅降低GPT-2训练时的显存占用，在batch size = 4 时，内存节省超过50+%。

2

亚线性内存优化的设计

在系列文章《深度解析：让你掌握OneFlow框架的系统设计（上篇、中篇、下篇）》中，我们介绍了OneFlow中的OpNode/OpGragh抽象以及建立在这之上的Actor、SBP抽象等系统设计，正是这些良好的系统设计和抽象使得OneFlow在多种任务下都有着优秀的表现。

OneFlow的Job任务在逻辑图编译期会基于由OpNode构成的Job逻辑图(OpGragh），进行一系列pass的系统优化过程，每个pass对逻辑图进行了一次图修改/重写（对逻辑图中的节点和连边进行了增删操作)，这些优化对性能的提升至关重要。

Activation Checkpointing 在OneFlow中的实现，也是通过一个Checkpointing的pass对Job逻辑图实现修改/重写来实现的（见https://github.com/Oneflow-Inc/oneflow/pull/3976）。

主要原理

如图所示： 

1.上半部分为正常情况下的逻辑子图。T1、T2为Transformer Layer的前向计算部分、子图中每个op计算完成后得到的中间激活特征将持续占用内存，当计算进行到反向时（T1_grad、T2_grad），再利用这些中间激活进行反向的计算； 

2.下半部分为开启Activation Checkpointing后的逻辑子图。可以看到，中间部分增加了虚线框住，用于重计算的fake子图，由于fake子图的存在，正常forward子图在进行前向时，就无须保存中间激活了，当backward计算需要用到时，再临时根据fake子图进行前向的重计算。

在OneFlow中，Activation Checkpointing的细节流程如下：

‍‍1.收集checkpointing作用域包裹下的所有前向pass下的ops

2.收集ops下所有的子图subgraphs

3.遍历子图subgraphs，并对所有需要做后向的subgraph做如下操作:

• 生成fake子图，并将其作为后向消费者的输入（而不是真实子图）

• 在fake子图中增加由end op连向所有源节点source nodes的控制边

• 将fake子图添加至job builder（被其管理）

4.在job builder中更新所有后向消费者ops

代码实现：

https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/job/checkpointing_config_def.cpp
https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/job/job_build_and_infer_ctx.cpp#L989
https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/job_rewriter/checkpointing_pass.cpp

3

代码解读

3.1 收集所有前向pass下的ops

由于activation checkpointing的设计主要是节省前向计算过程中的内存，即在checkpointing作用范围内，将前向计算过程中op nodes产生的activation显存释放掉，后向backward时，重新进行此部分的前向计算，得到所需的activation。

由于我们所有的操作是在逻辑图层面，操作的对象为每个op node节点，所以首先需要标记、筛选出checkpointing作用范围内所有前向的op nodes。此部分主要通过CollectAllCheckpointingOpsInForwardPass()方法实现：

https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/job_rewriter/checkpointing_pass.cpp#L65

void CollectAllCheckpointingOpsInForwardPass(
// 收集所有属于前向pass下，且符合条件的op nodes，存放至HashMap中
const OpGraph& op_graph, HashMap<std::string, const OpNode*>* checkpointing_op_name2op_node) {
// NOTE(chengcheng):
// ignore batch_norm ops because of recompute bn will repeat the calculation of 'm' and 'v'.
// in the future, we need to support the recomputation version of batch_norm which do NOT
// update forward variables.
HashSet<std::string> ignore_op_type_names = {"normalization", "normalization_add_relu",
"cudnn_fused_normalization_add_relu"};
op_graph.ForEachNode([&](const OpNode* op_node) {
const OperatorConf& op_conf = op_node->op().op_conf();
// 跳过不包含user_conf以及ignore_op_type_names指定的op_node
if (!op_conf.has_user_conf()) { return; }
if (ignore_op_type_names.find(op_conf.user_conf().op_type_name())
!= ignore_op_type_names.end()) {
return;
}
// 对scope范围内开启了checkpointing且标记为ForwardPass的op_node，则为目标node,将其插入HashMap中
if (IsForwardPass7CheckpointingScope(Scope4OpNode(op_node))) {
CHECK(checkpointing_op_name2op_node->emplace(op_conf.name(), op_node).second);
}
});
}


其中，主要通过IsForwardPass7CheckpointingScope()方法，来对符合条件的op node进行筛选：

bool IsForwardPassScope(const Scope& scope) {
// scope中，calculation_pass_name属性为kForwardPass的node，则为参与前向计算的目标node
return scope.scope_proto().calculation_pass_name() == kForwardPass;
}

bool IsForwardPass7CheckpointingScope(const Scope& scope) {
// True if 属性为kForwardPass的node且scope开启了checkpointing
return IsForwardPassScope(scope) && scope.Bool("checkpointing");
}


IsForwardPass7CheckpointingScope()方法通过node的scope来判断该op node是否属于直接参与前向计算的node（scope中包含kForwardPass），且是否开启了“checkpointing”，同时满足则为目标node，将其插入hashmap(checkpointing_op_name2op_node)中。

3.2 收集ops下所有的subgraphs

筛选出checkpointing作用区域内所有的op nodes后，需要根据这些nodes生成所有子图subgraghs，这些子图有些是和后向重计算无关、有些则是后向重计算所需的目标子图，它们的输出作为后向op node的输入被消费，这些子图是实现activation checkpointing设计中前向重计算的最小单位。

生成子图的代码如下：

// 根据ops生成所有subgraphs子图，并将其存放在vector中
// step 2. get all connected subgraphs in checkpointing ops.
std::vector<HashSet<const OpNode*>> checkpointing_subgraphs;
GenConnectedCheckpointingSubgraphs(checkpointing_op_name2op_node, &checkpointing_subgraphs);


其中，主要通过GenConnectedCheckpointingSubgraphs()方法生成subgraphs：

void GenConnectedCheckpointingSubgraphs(
// 生成Subgraphs子图
const HashMap<std::string, const OpNode*>& checkpointing_op_name2op_node,
std::vector<HashSet<const OpNode*>>* checkpointing_subgraphs) {
HashSet<const OpNode*> visited_nodes;
for (const auto& pair : checkpointing_op_name2op_node) {
const OpNode* node = pair.second;
if (visited_nodes.find(node) != visited_nodes.end()) { continue; }

// new subgraph
checkpointing_subgraphs->push_back(HashSet<const OpNode*>());
CHECK(!checkpointing_subgraphs->empty());
auto& subgraph = checkpointing_subgraphs->back();
CHECK(subgraph.empty());

// bfs search all node in checkpointing ops
CHECK(visited_nodes.insert(node).second);
std::queue<const OpNode*> queued_nodes;
queued_nodes.push(node);
while (!queued_nodes.empty()) {
const OpNode* cur_node = queued_nodes.front();
queued_nodes.pop();

CHECK(subgraph.insert(cur_node).second);
cur_node->ForEachNodeOnInOutEdge([&](const OpNode* next_node) {
const std::string& next_op_name = next_node->op().op_name();
if (checkpointing_op_name2op_node.find(next_op_name) != checkpointing_op_name2op_node.end()
&& cur_node->parallel_desc() == next_node->parallel_desc()
&& visited_nodes.find(next_node) == visited_nodes.end()) {
queued_nodes.push(next_node);
CHECK(visited_nodes.insert(next_node).second);
}
});
}
}
}


根据当前节点（cur_node），找到下一个子图节点(next_node)，采用的是BFS搜索，搜索逻辑为：以cur_node为起点，遍历其输入/输入边上有消费关系的节点next_node；对于不属于checkpointing op && 没有被当作子图node访问过 && 并行方式和cur_node一致的node；作为subgraph中的目标node（next_node），插入subgraph队列中，并将该node标记为已访问，放置到visited_nodes Set中。

3.3 遍历子图subgraphs

经历上述过程后生成了子图vector()，我们需要对其进行遍历，筛选出和activation checkpointing相关的子图subgraghs，并做如下几件事：

• 生成fake子图，并将其作为后向消费者的输入（而不是真实子图）

• 在fake子图中增加由end op连向所有源节点source nodes的控制边

• 将fake子图添加至job builder（被其管理）

对子图subgraghs的遍历主要在：

https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/job_rewriter/checkpointing_pass.cpp#L148-L290

过滤和activation checkpointing无关的subgragh子图

在[for (auto& subgraph : checkpointing_subgraphs) {}]()遍历循环的一开始，就会跳过不符合activation checkpointing条件的subgragh

for (auto& subgraph : checkpointing_subgraphs) {
// step 3.1 ignore this subgraph if there is no direct edge to backward pass op.
HashSet<const OpNode*> bw_consumers;
for (const OpNode* node : subgraph) {
node->ForEachNodeOnOutEdge([&](const OpNode* out_node) {
if (!IsForwardPassScope(Scope4OpNode(out_node))) {
bw_consumers.insert(out_node);
CHECK(subgraph.find(out_node) == subgraph.end());
}
});
}
if (bw_consumers.empty()) { continue; }


具体条件，即遍历subgraph子图中的所有node节点、判断node节点的所有出边out_edges是否有出边连到后向backward消费者op，如果subgraph中所有节点均没有连到后向backward消费者，则跳过该子图（表明该子图只有只和forward有关而和backward无关，即不是activation checkpointing优化的目标子图。

生成fake子图，并将其作为后向消费者的输入（而不是真实子图）

过滤掉无效子图后，对于和activation checkpointing直接相关的目标子图，我们需要生成fake子图，其中的每个节点由fake op构成。

fake子图，即重计算的最小单位，其作用即用于取代原有真实的子图、并在后面替换这些真实子图，用于被后向op nodes消费。通过将fake子图中fake op的scope属性从kForwardPass变为kBackwardPass，实现当计算进行到该fake op时，重新运行前向计算以产生backward所需的activation数据。

生成fake 子图的主要代码在：L168-L222 

https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/job_rewriter/checkpointing_pass.cpp#L168-L222

生成fake子图后，将其设置为backward消费者的输入，主要代码如下：

const OpNode* first_bw_consumer = nullptr;
int32_t first_bw_order = std::numeric_limits<int32_t>::max();
// 将backward消费者的input更改为fake子图op（而不是真实子图）
// step 3.3 change bw consumers input from subgraph to fake subgraph
for (const OpNode* node : bw_consumers) {
std::string bw_consumer_name = node->op().op_name();
OperatorConf bw_consumer_op_conf;
// NOTE(chengcheng):
// reuse bw conumer op conf if it has been existed in map.
if (total_bw_consumers_op_name2conf.find(bw_consumer_name)
!= total_bw_consumers_op_name2conf.end()) {
bw_consumer_op_conf = total_bw_consumers_op_name2conf.at(bw_consumer_name);
} else {
bw_consumer_op_conf = node->op().op_conf();
}
CHECK_EQ(bw_consumer_name, bw_consumer_op_conf.name());

auto* user_conf = bw_consumer_op_conf.mutable_user_conf();
// 修改和subgragh相关的backward op输入的blob name
// change input lbns if in subgraph
for (auto& pair : *(user_conf->mutable_input())) {
auto& list_s = pair.second;
for (int i = 0; i < list_s.s_size(); ++i) {
std::string old_lbn = list_s.s(i);
LogicalBlobId old_lbi = GenLogicalBlobId(old_lbn);

std::string old_input_op_name = old_lbi.op_name();
if (subgraph_op_name2op_node.find(old_input_op_name) != subgraph_op_name2op_node.end()) {
list_s.set_s(i, kCheckpointingFakeOpNamePrefix + old_lbn);
}
}
}

// NOTE(chengcheng):
// emplace maybe repeated, so do not check the return value
total_bw_consumers_op_name2conf.emplace(bw_consumer_name, bw_consumer_op_conf);

CHECK(op_node2order.find(node) != op_node2order.end());
int32_t this_order = op_node2order.at(node);
if (this_order < first_bw_order) {
first_bw_consumer = node;
first_bw_order = this_order;
}
}


在fake子图中为所有source node—end node添加控制边

这一步操作的目的主要将子图subgraph中所有和backward op相连的node(source node)，添加一条控制边。控制边的添加是人为控制node间执行的时序，控制边保证了fake子图的计算尽可能晚的发生，这样才能缩短生命周期，保证内存复用的效率。

添加控制边相关的代码在L267-L284：

https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/job_rewriter/checkpointing_pass.cpp#L267-L284

// step 3.4 add control edge from End Op to all source node in fake subgraph
CHECK(first_bw_consumer != nullptr);
std::string end_op_name = kCheckpointingBadOpName;
int32_t end_order = -1;
first_bw_consumer->ForEachNodeOnInEdge([&](const OpNode* end_node) {
CHECK(op_node2order.find(end_node) != op_node2order.end());
int32_t this_order = op_node2order.at(end_node);
if (this_order > end_order) {
end_order = this_order;
end_op_name = end_node->op().op_name();
}
});
CHECK_NE(end_order, -1);
CHECK_NE(end_op_name, kCheckpointingBadOpName);
CHECK_LT(end_order, first_bw_order);
for (const auto& source_op_name : source_node_in_fake_subgraph) {
fake_op_name2conf.at(source_op_name).add_ctrl_in_op_name(end_op_name);
}


将fake子图添加至job build（被其管理）

fake subgraphs的生成、以及控制边的添加，实际上对原有job逻辑图产生了改动。改动后，需要将fake subgraphs中这些新生成的fake op nodes添加至job builder管理，正式完成了逻辑图的图改写。

主要代码如下：

// 将fake subgraph所包含的ops加入至job_builder管理（图改写）
// step 3.5 add fake subgraph ops to job builder
std::vector<OperatorConf> fake_op_confs;
for (auto& pair : fake_op_name2conf) { fake_op_confs.push_back(pair.second); }
job_builder->AddOps(parallel_conf, fake_op_confs);


3.4 更新所有后向消费者ops

最后，由于fake op nodes更新了backward op nodes的输入输出等属性，需要将更新后的backward op nodes同步至job_builder管理：

// 在job builder中更新所有backward ops
// step 4. update bw consumers in job builder only once
std::vector<OperatorConf> total_bw_consumer_op_confs;
for (auto& pair : total_bw_consumers_op_name2conf) {
total_bw_consumer_op_confs.push_back(pair.second);
}
job_builder->MutOpsOnlyOnce(total_bw_consumer_op_confs);

return Maybe<void>::Ok();


至此，通过这些fake subgraphs的插入、输入输出连边的变动等完成了整个job逻辑图的改写，改写后的逻辑图执行时即自动支持了activation checkpointing。

OneFlow最近复现了GPT-3相关的工作，其中就使用了activation checkpointing的技术，代码在OneFlow-Benchmark已开源，欢迎在GitHub下载试用：

https://github.com/Oneflow-Inc/OneFlow-Benchmark/tree/master/LanguageModeling/GPT

注：题图源自insspirito，pixabay

‍参考文献

[1] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training Deep Nets with Sublinear Memory Cost. arXiv preprint arXiv:1604.06174, 2016.

近期文章