AI推理计算框架中的内存优化

背景

内存管理是AI计算中非常重要的一部分。我们希望模型计算时占用内存尽可能小，这样我们训练或推理时就可以用更大的batch size使其尽快收敛，或者提高吞吐率。又或者让我们可以使用参数更多、或更复杂的模型从而达到更好的准确率。由于现代深度学习模型大多在GPU上运行，而GPU的显存相比CPU小得多，因此我们这里主要关注的是GPU memory。

首先看下我们需要重点关注哪些GPU memory。对于模型在计算中需要用到的GPU memory，论文《Estimating GPU Memory Consumption of Deep Learning Models》做了比较具体的总结。对推理计算而言，主要有这么几类：

• 权重参数（Weight Parameter）：这个不用多说，模型中的参数。
• 中间张量（Intermediate Tensor）：网络中每层的输出与输入张量。
• 其它：计算中需要的临时内存（如kernel中使用的一些memory，cuDNN的workspace），还有一些常驻的内存（如CUDA context）。

其中第三类本身占的空间不大，也比较难优化。第一类减少权重内存占用的话可以通过一些模型压缩方法，如量化，剪枝。之前写过一些相关文章如《闲话模型压缩之量化（Quantization）篇》和《闲话模型压缩之网络剪枝（Network Pruning）篇》，有兴趣可以参考。相比之下，第二类，即中间张量的优化空间更大，因此很多业界的工作也是针对它来优化。优化的思路有很多，比如：

• 内存重用（Memory reuse）：由于有些中间张量的生命周期间互不重叠，因此可以reuse。MegEngine, IREE, TensorRT等框架都做了memory usage相关的优化。
• 重计算（Recomputation）：该技术主要用于训练中。它将模型中的一些节点作为checkpoint，其它节点的输出可丢弃，当在计算梯度时需要时通过最近的checkpoint重新计算生成。因此被称为checkpointing技术。该问题也被称为tensor rematerialization优化问题。相关论文如适用于静态网络的offline方法《Training Deep Nets with Sublinear Memory Cost》，适用于动态网络的online方法《Dynamic Tensor Rematerialization》，建模为MILP进行求解的《Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization》等。
• 交换（Swap）：基本思想是将显存中的数据交换到CPU上，相当于把GPU memory当成CPU memory的cache。一些塞不进GPU显存的层，如DLRM模型的embedding层可能会用到这种技术。相关的论文如《Supporting Massive DLRM Inference Through Software Defined Memory》，《vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design》等。
• 压缩（Compression）：将数据进行压缩，如《Gist: Efficient Data Encoding for Deep Neural Network Training》，根据特定层输出特点对层的输出，即feature map数据进行编码压缩。
• …

后面几种都会一定程度上牺牲性能，这里我们主要关注第一种。它在对延迟关注的推理场景用得尤其多一些。比如TensorFlow Lite利用该技术可以显著减少内存占用（详见https://blog.tensorflow.org/2020/10/optimizing-tensorflow-lite-runtime.html）。

对于网络前面层的输出，到计算后面的层时可能已经不再使用了。换句话说，对于中间层的输出张量，它们的生命周期可能是没有重叠的，对于它们便可以进行重用。下面是最简单情况下的示意图：

其中Task 1写数据到Tensor 1交给Task 2，Task 2处理后将结果写入Tensor 2，交给Task 3。因为Tensor 1与Tensor 2生命周期并不重叠，所以它们可以重用同一个Tensor。典型的如神经网络中的一些element wise操作。

但实际中的情况远没有这么简单。网络中不总是线性结构，另外张量的大小可能各不相同，给重用带来困难。要得到其最优的分配策略，是一个NP-complete问题。因此实际当中，我们可以倾向于一些heuristic方法，这样可以在合理时间内得到一个近似最优解。

那如何优化呢？在不少地方，如TensorRT会提到基于register allocation的思想。任何一本编译器的教材中都会介绍register allocation，在此不展开。大体会先通过liveness analysis得到变量的live range，然后根据它生成inference graph，转为着色问题来解。业界也有采用这种方法的做法，如《Memory Allocation for Neural Networks using Graph Coloring》。但是，memory planning与register allocation所面临的问题还是有所区别的，比如：

• 大小不确定：Register的大小基本是相同，或者说是差不多的，而memory的大小可能差异很大，重用一块过大或过小的memory会产生问题。
• 拷贝成本高：Register拷贝一下还好，但memory拷贝开销比较大，尤其是大块的memory。因此理想情况下我们希望不要拷贝。
• Fallback机制：Register实在分配不了会产生spill，即放到memory中。虽然memory要不不够理论上也能往更下一层存储器上搬，业界也有这方面的研究，但很多情况下因为时延等原因不会这么做。

另外，从静态/动态角度，内存的管理方式大体有动态分配与静态规划两种：

• 动态分配（Dynamic Memory Allocation）：内存分配在运行时进行。由于从系统中分配的开销较大，通常维护一个memory pool。需要时从中分配，不再需要时放回到pool。如TensorFlow中的BFC allocator。
• 静态规划（Static Memory Planning）：内存分配在运行前进行，常见于基于编译器的计算框架。通过规划进一步减少内存使用，减少OOM带来的不确定性，同时最少化运行时内存管理的开销。如MXNet与MegEngine/MegCC中的static memory planning。

光看概念有些抽象，下面就以一个实例 - TensorFlow Lite（TF Lite）中的内存优化来看看具体的实现。其实在论文《Efficient Memory Management for Deep Neural Net Inference》与《On-Device Neural Net Inference with Mobile GPUs》中对其原理已经介绍得比较清楚了。下面主要是结合代码理解下它的实现。

代码走读

基础数据结构

先来看几个关键数据结构。它们定义在types.h文件中。结构体TensorUsageRecord即论文中提到的Tensor usage record，用于记录张量的使用记录。它主要包含三个信息：tensor size, 以及第一个与最后一个使用它的task。代表这两个task的成员first_task与last_task即它的生命周期。

using UsageGraph = std::vector<std::vector<size_t>>;template <typename TensorSizeT>
struct TensorUsageRecord {TensorSizeT tensor_size;TaskId first_task;TaskId last_task;...
};

注意它是个模板类，有针对size_t，uint2，uint3与BHWC的特化（实现在memory_management.c文件）。

结构体ObjectsAssignment与OffsetsAssignment都用于存放分配结果。

// Information about assignment of tensors to shared objects
template <typename TensorSizeT>
struct ObjectsAssignment {// shared_object_ids_[i] is ID of shared object, that tensor i will be using.std::vector<size_t> object_ids;// shared_object_sizes_[i] is a size of shared object with ID equal to i.std::vector<TensorSizeT> object_sizes;
};// Information about assignment of tensors to offsets for the case, when all of
// them are going to be allocated in one continuous memory block.
struct OffsetsAssignment {std::vector<size_t> offsets;size_t total_size;
};

它们对应后面会提到的两种分配方式。前者用于shared object（指可以用于多个tensor的内存块）分配，后者用于从大块连续内存中分配子内存区域。

为了解它的使用，可以参考它的测试用例memory_management_test.cc。比较典型的有几个case：OneRecord，ChainRecords，ComplexRecords，分别对于只有一个节点，链式（即线性）计算图，和复杂计算图。

以ChainRecords这个case为例：

TEST(Model, ChainRecords) {                                                          std::vector<TensorUsageRecord<size_t>> usage_records{                              {/*size=*/16, /*first=*/0, /*last=*/1},                                        {/*size=*/8, /*first=*/1, /*last=*/2},                                         {/*size=*/64, /*first=*/2, /*last=*/3},                                        {/*size=*/32, /*first=*/3, /*last=*/4},                                        {/*size=*/8, /*first=*/4, /*last=*/5},                                         };                                                                                 ObjectsAssignment<size_t> assignment;                                              ASSERT_TRUE(                                                                       AssignObjectsToTensors(usage_records, MemoryStrategy::NAIVE, &assignment)      .ok());                                                                    EXPECT_THAT(assignment.object_ids, ElementsAre(0, 1, 2, 3, 4));                    EXPECT_THAT(assignment.object_sizes, ElementsAre(16, 8, 64, 32, 8));    ...

可以看到，其中最核心的是AssignObjectsToTensors()函数。该函数基于由TensorUsageRecord数组表示的张量使用记录（按拓扑序排列），根据指定的分配策略（由MemoryStrategy表示），计算得到分配结果（由ObjectsAssignment对象表示）。

Object分配方式

注意AssignObjectsToTensors()是个模板函数，根据TensorUsageRecord的类型不同有多种实现。以最简单的TensorUsageRecord<size>的情况（即tensor的大小以一个size_t类型表示）为例，相关代码如下：

template <>
absl::Status AssignObjectsToTensors(const std::vector<TensorUsageRecord<size_t>>& usage_records,MemoryStrategy strategy, ObjectsAssignment<size_t>* assignment,const UsageGraph* reallocation_graph) {switch (strategy) {case MemoryStrategy::NAIVE:return NaiveAssignment(usage_records, assignment);case MemoryStrategy::EQUALITY:return EqualityAssignmentWithHash(usage_records, assignment);case MemoryStrategy::GREEDY_IN_ORDER:return GreedyInOrderAssignment(usage_records, assignment,reallocation_graph);case MemoryStrategy::GREEDY_BY_BREADTH:return GreedyByBreadthAssignment(usage_records, assignment);case MemoryStrategy::GREEDY_BY_SIZE:return GreedyBySizeDistPriorityAssignment(usage_records, assignment);case MemoryStrategy::GREEDY_BEST:return BestGreedy(usage_records, assignment);case MemoryStrategy::MINCOSTFLOW:return MinCostFlowAssignment(usage_records, assignment);default:return absl::InternalError("MemoryStrategy is not supported with current tensor size type.");}return absl::OkStatus();
}

可以看到，它的主体部分就是根据指定策略调用相应的分配算法。这几种策略分别是：

NATIVE

NaiveAssignment函数将每个张量分配独立的memory object。实现在native_assignment.h文件中。该算法为每个张量分配一个新的shared object。这是最简单，也是最浪费内存的做法。代码逻辑比较好理解，不多说了。

EQUALITY：

EqualityAssignmentWithHash()函数实现于equality_assignment.h文件中。它适用于TensorSizeT为hashable type的情况（unhashable type的情况使用EqualityAssignment()函数）。

该算法维护两个关键数据结构：一个是pool，它是一个map。其键值为size，值为目前free（即空闲）的且size与键值相同的shared objects。另一个是优先队列objects_in_use，它保存目前在使用的share objects，并按size排序。整个过程遍历所有的tensor usage record，对于每个tensor usage record，先将队列objects_in_use中相对于当前tensor已不再使用的shared object弹出，放入pool。然后在pool中找有没有与当前tensor的size匹配的shared object，如有就重用，没有的话就新创建一个shared object。过程示意图如下：

注意这里只是做memory planning，所以不用真正做内存分配。

GREEDY_IN_ORDER：

函数GreedyInOrderAssignment()实现在greedy_in_order_assignment.h文件中。它维护与前一算法中相似的两个数据结构：一个是存放free shared objects的pool，以object size排序。另一个是存放正在使用的shared object的优先队列objects_in_use，以last_task排序。该算法主体部分遍历tensor usage records。在第一步中，首先看哪些object不再使用，将它们放入pool。这一步与前面算法一样。

然后尝试将pool中的shared object分配给当前tensor。前面是需要严格匹配才能重用，这里放宽了一些，允许在size不严格一致时也能重用。这里在从pool中找可用的shared object时，不是用的find()函数，而是用的是二分查找lower_bound()，即查找不小于当前tensor的size的第一个元素。这样做保证找到的shared object（如有）能容纳当前tensor，同时又使浪费的空间尽可能小。然后尝试检查前一个元素，其shared object size与tensor size的差值是否比前面找到的元素更小。如果是就选它了。这样使得对shared object的size改变尽可能得小就能满足当前tensor的需求。比如pool中有size分别为1, 3, 6的shared object，而当前tensor为5。这种情况下就会找到size为3这个shared object，因为它与5的差值是最小的。可以看到，它每一步尽可能找与当前tensor的size尽可能接近的shared object进行重用，体现了贪心的思想。

当前面的步骤中找到合适的shared object，就把它从pool中拿走，将之分配给当前tensor，分配信息写入assignment。同时将该信息记录在objects_in_use中。如果shared object的size小于tensor，会增大shared object的size。如果没有找到合适的shared object，则创建新的shared object。

GREEDY_BY_BREADTH

函数GreedyByBreadthAssignment()实现在greedy_by_breadth_assignment.cc文件中。与前面的贪心算法类似，主要区别在于它优先考虑breadth大的task。Breadth表示该task执行时所有tensor的size之和，记录于TaskProfile对象中。TaskProfile是一个vector，元素表示task执行时还在使用的张量，以size降序排序。obj_schedules是SharedObjectSchedule的vector。SharedObjectSchedule记录了shared object对应的所有tensor usage record。

// Set of usage records for all tensors assigned to the shared object, ordered
// by first_task.
using SharedObjectSchedule = std::set<TensorUsageRecord<size_t>>;               struct TaskBreadthWithId {size_t breadth;TaskId task_id;TaskBreadthWithId(size_t breadth, size_t task_id)                             : breadth(breadth), task_id(task_id) {}                                   // Default order of TaskBreadthWithId is increasing order of their breadth.   bool operator<(const TaskBreadthWithId& other) const {return breadth < other.breadth;}         
};        

算法首先调用CalculateTaskProfiles()函数计算出所有task的TaskProfile，它包含了task的breadth信息。然后以breadth非递增顺序遍历所有task。对于每个task，遍历其所有的tensor，为它们分配shared object。分配过程中考虑obj_schedules中的shared object是否可重用。遍历其中元素，如果不合适就跳过，否则用lower_bound()函数找到其中不小于当前所需size的第一个shared object。如果找到的话检查一下，如果合法，就将该shared object作为最优候选。如果没找到就创建新的shared object。

GREEDY_BY_SIZE

函数GreedyBySizeDistPriorityAssignment()实现在greedy_by_size_assignment.cc文件中。它的流程大体与前一种相似，区别在于优先考虑哪些tensor时所用的策略。该方法使用了一种更加复杂的heuristic。核心数据结构是SizeDistPriorityInfo：

struct SizeDistPriorityInfo {// - Tensor with leftmost position in positional maximums vector has higher// priority;// - If two tensors have equal position, the one, that has usage interval with// smallest positive distance (best_dist) to some of already assigned tensors,// has higher priority;// - If two tensors have equal position and best_dist, the one with greater// tensor_size has higher priority.bool operator>(const SizeDistPriorityInfo& other) const {return position < other.position ||(position == other.position &&(best_dist < other.best_dist || (best_dist == other.best_dist &&tensor_size > other.tensor_size)));}// Recalculate best distance and best object, based on precalculated distances// in vector dist.void RecalcBestDist() {best_dist = kNotAssigned;for (size_t obj_id = 0; obj_id < dist.size(); ++obj_id) {if (dist[obj_id] < best_dist) {best_dist = dist[obj_id];best_object = obj_id;}}}size_t position;size_t tensor_size;std::vector<size_t> dist;size_t best_dist;size_t best_object;size_t tensor_usage_id;
};

根据上面的注释，在positional maximum vector中位置越靠左的tensor优先级越高。如果在该vector中位置一样，则与已分配的tensor的positive distance小的优先级更高。如果前两个标准还决不出高下，那size较大的tensor优先级更高。这个优先级考虑了tensor的size，也考虑了与时序上的局部性。

主流程中首先调用CalculatePositionalMaximums()函数计算positional maximums vector。即每个shared object size的lower bound的数组。然后根据这个信息填好SizeDistPriorityInfo数组。接下来分两层循环。它会以SizeDistPriority递增顺序遍历所有tensor。即对于每个tensor，按前面计算得到的优先级信息找到优先级最高的tensor，以及相应的object。如果找到就按之进行分配，没找到就创建新的shared object。分配完成后，SizeDistPriority信息会产生变化，因此需要进行相应的更新。

GREEDY_BEST

函数BestGreedy()实现在memory_management.cc文件中。它结合了前两种贪心算法。结合方式比较直观，即把GREEDY_BY_SIZE与GREEDY_BY_BREADTH都跑下，看哪个的总用量少（即更优）就选哪个。

  RETURN_IF_ERROR(GreedyBySizeDistPriorityAssignment(usage_records, assignment));ObjectsAssignment<size_t> assignment_by_breadth;if (GreedyByBreadthAssignment(usage_records, &assignment_by_breadth).ok() &&TotalSize(assignment_by_breadth) < TotalSize(*assignment)) {std::swap(*assignment, assignment_by_breadth);}

MINCOSTFLOW

函数MinCostFlowAssignment()实现在min_cost_flow_assignment.cc文件中。
它使用Minimum-cost flow matching algorithm。它将问题建模为一个minimum-cost flow problem（MCFP）。对于原问题，它创建一个auxiliary flow graph，对于每个intermediate tensor在图中插入两个对应节点，另外创建两个特殊节点-source与sink。然后按一定规则向图中添加有向边（具体可参见论文《On-Devie Neural Net Inference with Mobile GPUs》）。创建完后，使用Shortest Path Faster Algorithm（SPFA）求解。

TF Lite中虽然提供了多种机制，但至于哪种好不一定，所以实际使用当中可能要都试一下。

Offset分配方式

除了这种以shared object分配的模式外，TF Lite还支持一种在一整块内存块分配的模式。两种模式区别示意图：

代码中检测如果支持sub buffer方式的话，会试图以memory中的切块（以offset表示）进行分配。一般使用在先分配一大块连续memory，然后在里边“切”小块的情况。这种情况下，会调用AssignOffsetsToTensors()函数，继而调用GreedyBySizeAssignment()函数。这是论文中提到的Offet Calculation方法。它可以看作是二维的strip packing problem。在主逻辑函数GreedyBySizeAssignment()中，需要注意的是两个比较关键的数据结构：ordered_records是按size排序后的TensorUsageRecord。ordered_allocs是已经分配好memory的tensor，按offset排序。主要逻辑包含两层循环。外循环按序遍历TensorUsageRecord，然后内循环中遍历已分配的内存块。如果生命周期不重叠，则会考察它前面空出的区域是否能容纳下当前的tensor。如果能容纳且浪费的区域最小，则将当前tensor塞进去。找一圈都找不到的话就在最右边新分配一段。

框架对接

看了内存分配的实现，再看下它是如何结合进框架，用于模型的计算的。我们知道TF Lite会将计算放到相应的后端，如Hexagon，OpenGL, OpenCL, Metal等。这些后端实现放在delegates目录下。以GPU的OpenCL后端为例，相应的实现主要位于tensorflow/tensorflow/lite/delegates/gpu/cl/inference_context.cc文件中。函数InitFromGpuModel()调用AllocateMemory()函数，函数AllocateMemory继而调用AllocateBufferBasedTensors()函数或者AllocateStrongShapesTensors函数。前者用于buffer based的tensor（参见IsBufferBased()函数），比如clCreateBuffer()分配的内存。其它的则用AllocateStrongShapesTensors()函数进行处理。函数AllocateStrongShapesTensors()用于那些具有特定layout的tensor（参考delegates/gpu/common/shape.h），调用AssignObjectsToTensors()函数时传入的分配策略是MemoryStrategy::EQUALIT。

下面看下AllocateBufferBasedTensors()函数。其中核心函数GetBufferAsignment()负责buffer的分配。它先调用GetUsages()函数得到网络中张量的使用信息，然后基于它构成TensorUsageRecord数组。有了这些信息，就可以调用前面提到的AssignObjectsToTensors()函数进行内存对象的分配了。默认用的分配策略是MemoryStrategy::GREEDY_BEST。

另外如果支持sub buffer的话，还会调用AssignOffsetsToTensors()进行分配。最后判断与前面以buffer为单位（AssignObjectsToTensors()函数）的分配方式得到的结果哪种好，即使用的memory少，从而选择更优的一种。

最后总结一下整个调用流程作为收尾吧：