10. LayerGroup
10.1. 基本概念
TPU芯片分为片外内存(或称Global Memory, 简称GMEM)和片内内存(或称Local Memory, 简称LMEM)。
通常片外内存非常大(比如4GB), 片内内存非常小(比如16MB)。神经网络模型的数据量和计算量
都非常大, 通常每层的OP都需要切分后放到Local Memory进行运算, 结果再保存到Global Memory。
LayerGroup就是让尽可能多的OP经过切分后能够在Local Memroy执行, 而避免过多的Local和Grobal Memory的拷贝。
- 要解决的问题:
如何使Layer数据保持在有限的Local Memory进行运算, 而不是反复进行Local与Global Memory之间的拷贝
- 基本思路:
通过切Activation的N和H, 使每层Layer的运算始终在Local Memory中, 如图(网络切分举例)
图 10.1 网络切分举例
10.2. BackwardH
对网络进行H切分的时候, 大多数Layer输入和输出的H是一致的。但是对于Conv、Pool等等需要特别计算。
以Conv举例, 如图(卷积BackwardH举例)
图 10.2 卷积BackwardH举例
10.3. 划分Mem周期
如何划分group? 首先把每一层Layer需要的lmem罗列出来, 大体可以归为三类:
Activation Tensor, 用于保存输入输出结果, 没有使用者后直接释放
Weight, 用于保存权重, 不切的情况下用完就释放; 否则一直驻留在lmem
Buffer, 用于Layer运算保存中间结果, 用完就释放
然后依次广度优先的方式配置id, 举例如图(LMEM的ID分配)
图 10.3 LMEM的ID分配
然后再配置周期, 配置方法如图(TimeStep分配)
图 10.4 TimeStep分配
关于配置周期的细节如下:
[T2,T7], 表示在T2开始的时候就要申请lmem, 在T7结束的时候释放lmem
w4的原始周期应该是[T5,T5], 但是被修正成[T2,T5], 因为在T2做卷积运算时w4可以被同时加载
当N或者H被切分时, Weight不需要重新被加载, 它的结束点会被修正为正无穷
10.4. LMEM分配
当n或h存在切分的情况下, weight常驻LMEM, 每一个切分都可以继续使用weight。
这时候会先分配weight, 如图所示(有切分情况的分配)
图 10.5 有切分情况的分配
当n和h都没有切分的情况下, weight和activation处理过程一样, 不使用时就释放。
这时候的分配过程, 如图所示(无切分情况的分配)
图 10.6 无切分情况的分配
那么Lmem分配问题就可以转换成这些方块如何摆放问题(注意方块只能左右移动, 不能上下移动)。
另外lmem分配时优先不要跨bank。
目前策略是按照op顺序依次分配, 优先分配timestep长的, 次分配lmem大的。
10.5. 划分最优Group
图 10.7 Group流程
目前从尾部开始向头部方向划分group, 优先切N, 当N切到最小单位时还不能满足要求, 则切h。
当网络很深的时候, 因为Conv、Pool等等算子会有重复计算部分, h切的过多导致重复部分过多;
为了避免过多重复, 当backward后的layer的输入, 如果h_slice重复的部分>h/2, 则认为失败。
举例: 比如input的h = 100, 经过切分后变成2个input, h[0, 80)和h[20, 100), 则重复部分为60, 则认为失败; 2个input对应h[0, 60)和h[20, 100), 重复部分为40, 认为成功。