1 overview
AiM是一个专门用于加速GEMV算子的加速器。
从加速器通用性角度看,AiM与TPU类似,两者都是十分专用的加速器。TPU使用脉冲阵列加速计算密集性算子GEMM;AiM使用存算架构加速访存密集型算子GEMV。
然而,目前广泛使用的加速器是英伟达的通用加速器GPU。GPU是以计算为中心的架构,使用SIMT掩盖访存开销。对于GEMM算子,GPU打不过TPU;对于GEMV算子,GPU打不过AiM。但是,GPU有很强的通用性,加上完善的生态,它可以很好地完成各种算子(不论是计算密集性算子还是访存密集型算子)。
2 AiM defects
- 潜在的通信瓶颈
- AiM单片只有8GB存储,存不下整个大模型,这也就意味着需要进行多片协同处理。
- 看AiM架构,CPU和AiM使用pcle总线连接;AiM芯片之间使用类似NVLink的东西连接。
- 使用4片AiM,面向LLaMA-7B进行用例编程
- LLaMA-7B有32layer,那就直接每片AiM存储8layer模型参数
- 使用小batch流水线,4片AiM均在满负荷工作。
- hidden size = 4096,每一片AiM完成4 layer × (4096×4096×K)的计算量,向下一片传输4096的数据。K的数量级是10
- 如果没有类似NVLink的技术实现AiM间快速通信,可能会产生通信瓶颈。
- 专用的加速器
- 增大batch,计算密集型算子占比增大,性能收益下降。
- 对于transformer,假设输入token数量为N,generate token 数量为S。
- 如果N>>S,此时,计算密集型算子占比增大,性能收益下降。
- 编程
- GPU是以计算为中心的编程模型,只需要malloc内存,把数据传到显存就可以了;然而,AiM将计算和存储放在一块了。
- 如果数据划分需要用户来考虑,这跟直接写汇编差不多了。
- AiM的SDK要暴露出哪些编程接口?
- 以矩阵存储为中心的编程?
- SDK只暴露出GEMV等适合做的算子的接口?
- 在绝大多数情况下,AiM需要多卡协同工作。这对编程的挑战很大。
3 开始闲扯
- 为什么AI加速器生态牢牢被GPU占据?其实,在2014年,google就敏锐的察觉到AI加速器的前景,开始做TPU。那为什么TPU没有占据AI加速器生态?
- 为什么TPU里面的脉冲阵列没有放到GPU里面?
- 未来的神经网络模型中的访存密集型算子会更多吗?
- AI加速器到底是会更加通用还是专用?