torch memory planning | 实在之路

type

status

date

slug

summary

📝 Offline Momory Planning

静态神经网络编译时，我们可以对memory做一些plan，来降低网络对内存使用的大小。

如上图，横轴是时间，纵轴是地址。一个矩形代表了一个buffer，其高代表了buffer的大小。由于buffer的live range是确定的，也就是横轴坐标是确定的（不考虑swap out), 所以memory planning的过程就是确定每个矩形纵坐标的过程。

图中展示了两种不同的plan，其中右图所需要的 memory更小。

Memoy Planning带来的隐式依赖

由于memoy planning让更多的tensor在memory空间上有了更多的复用，也就导致了tensor之间有了更多的隐式依赖。

比如两个原本来不同分支的计算节点，这两个节点相互之间是没有任何依赖的，可以并行。但是可能会因为memory planning导致这两个节点的output使用同一块memory 空间，使得两者产生了隐式依赖，导致无法并行。

一种解决方法是，可以限制memory plan仅在同一分支中进行，不同分支可以并发的节点/tensor，不做memroy reuse。

Memory Planning in PyTorch

torch.compile中，会对memory对一些离线的规划。

memoy plan发生的codegen wrapper code的过程中。这里的wrapper code指的调用kernel的那些code，而非kernel code本身。memoy的申请与释放也属于wrapper code。

在为inductor ir node生成wrapper code时，不会直接生成memoy相关的code，而是将相关信息记录下来，等到后续memory plan后，再实际codegen相关代码。

torch里关于memory plan有两个分支， memory_plan和memory_plan_reuse的。

memory_plan_reuse

这个分支逻辑比较简单。

初始化 reuse_list，是一个key-value 表，其中key为: {device, dtype, size}， value为待reuse的buffer。

按顺序遍历所有的Alloc/Free 动作

Alloc buffer时

根据Alloc的信息，构造一个key，去reuse_list里查找。
有的话就复用这个buffer。将被复用的buffer从reuse_list中移除。

Free Buffer时，将buffer加入reuse_list。

memory_plan

memory_plan的能力强一些，会利用到live range，在时间和空间两个维度上对memory进行复用.

memoy_plan的核心代码都放在 _inductor/codegen/memory_planning.py下。

算法逻辑

概括一下PyTorch memoy plan的逻辑就是：

首先将所有buffer排序，也就是确定先摆放哪个矩形。

按buffer size降序
live_range跨越的range长度降序。

按序进行分配，也就是，按照顺序摆放矩形，这里可能有两种情况：

当前矩形在live_range上，和已经摆放了矩形，没有冲突，可以直接挨着横轴摆放。
如果在横轴上和已经摆放好的矩阵有overlap，则贴着已有的矩形放。（但是torch目前的实现做不到贴着放，纵方向上，矩形之间可能会有空隙。）

这里先进行排序，先分配大的tensor，应该是这样能减少后续的碎片化。

整个流程有点像俄罗斯方块，但是不能控制方块左右移动，只能控制方块的出现顺序。

AllocationTree

为了实现上面的逻辑，有一个问题需要解决：

放置矩形时，如果有overlap, 怎么计算矩形的start_address

pytorch使用了一套叫AllocationTree的数据结构来解决。

AllocationTree将所有的Allocation用树的方式组织了起来，抽象出了时间分片，和空间分片，来简化分配过程。

每个节点都是对整个memroy使用的一个抽象

每个节点都有的基本属性：

size：占用的memory size

live_ranges: 所有子节点的live_range的并集，也就是可以有多个不连续的live_range。

allocate: 成员函数，尝试在当前节点占用的资源上，进行复用分配。

AllocationTree的节点种类有：

Allocation, 代表了上图中一个最原子的矩形，也就是对应了一个tensor。一定是树的叶子节点。

Empty,占位符，代表了一块暂未分配的memory。一定是树的叶子节点

TemporalSplit，时间分片，有多个子节点

每个子节点代表了对同一块memory空间在不同时间上的使用。
每一个sub allocation的大小相同，live_range互不overlap。

SpatialSplit, 空间分片，有两个子节点， left, right

left, right将一整块大的，连续的memory划分成两块，空间上没有重叠。
left, right的类型都是TemporalSplit

AllocationTree的根节点是一个TemporalSplit，其size等于当目前为止的peak memory usage.

每次分配时：

现尝试直接在当前的root TemporalSplit节点上分配，看看能不能复用，而不扩大整体memroy size。

如果root TemporalSplit上分配失败，则将原root和当前alloc一起用一个SpatialSplit包起来，再用一个TemporalSplit将这个SpataialSplit包起来，作为新的root。本质上就是扩大了整体memory size。

SpatialSplit allocate(buffer):

left.allocate(buffer) or right.allocate(buffer)

TemporalSplit allocate:

先看和所有子节点在live_range上有没有overlap

如果没有overlap：

看buffer和这个temporal节点的size比较

如果恰好相等，直接append一个子节点
如果小于，则创建一个SpatailSplit(buffer, empty)，然后append。这里empty的size为：当前节点size - buffer size。
如果大于，则将当前temporalSplit扩容，即对每个子节点都创建一个SpatialSplit包起来。

如果有overlap，看overlap数量：

如果只有一个时间分片子节点和buffer在liverange上overlap，则尝试在用这个子节点进行allocate，因为里面可能有empty的节点。
如果有多个时间分片子节点和buffer存在overlap，则直接失败。（这里注释也写了，可以进一步做更仔细的判断，但是目前没有实现。)

📏torch memory planning