深度学习编译优化技术的探索和实践—杨军

CodeWarrior

2019/11/22 发布于 技术 分类

文字内容
1. PAI-TAO (Tensor Accelerator and Optimizer) 深度学习编译优化技术的探索和实践 计算平台事业部 穆琢
3. Outline l背景 lMotivation l业界现状 lTAO V1 lTAO V2 3
4. 背景 几个关键指标: Clock: ? MHZ CUDA cores: ? Flops(Clock * CUDA cores * 2): ? Tflops Memory Capacity: ?GB Memory Bandwidth: ?GB Transistors: ?B
5. 背景 2008 Top500 Rank1---IBM Roadrunner(1000 Tflops) 12,960 IBM PowerXCell 8i CPUs, 6,480 AMD Opteron dual-core processors, InfiniBand 2008年北京市全年用电 68971890 MWH 2.35 MW
6. 背景 Clock(MHZ) CUDA Cores Tflops/Tops Mem Capacity (GB) Mem Bandwidth (GB/s) Transistors(Billion) K40 745~875 2880 4~5 12 288 7.1 M40 948~1114 3072 5.8~6 12 288 8 P100 1328 3584 8~9 16 732 15 V100 (FP16 TensorCore) 1245~1380 5120 15~120 16 900 21 T4 (FP16/INT8/INT4 TensorCore) 585~1590 2560 8~260 16 320 13
7. Motivation l细粒度算子开销 l新硬件打开进一步优化空间
8. Motivation lPAI上作业的开销构成 • GPU计算部分 - Kernel计算开销 - GPU Op Launch开销 • • • • CPU计算部分 分布式开销 IO及数据预处理开销 框架外开销
9. Motivation l平台化的性能优化产品 • 提升模型创新迭代速度 • 降低在线预测延迟 • 降低硬件成本,提升资源利用率 l探索编译优化之前 • 若干手工优化的业务项目 • NMT,PAI-OCR,Rokid etc. • 手工优化的主要问题 • 开发成本问题 • 通用性和可维护性问题 NMT inference 业务上实际做过 的⼀个⼿⼯优化 pattern
10. Motivation l深度学习计算图的编译器 • 编译执行方式提高用户作业性能 • 解决细颗粒度框架的灵活性与性能间的矛盾 • 对用户通用/透明
11. 业界现状 1,完整的编译器框架,原⽣支持Tensorflow 2,计算密集型节点依赖厂商Library,重点支持访存密集型节点 3,实验室性质产品,在⼤多数模型上都⽆法跑通或效果不佳 4,Google内部主体投⼊在TPU backend 1,引⼊单独的Schedule层抽象,有能⼒做计算/访存密集型pattern的CodeGen和tuning l 继承XLA的整体框架 2,要求用户对底层硬件有⼀定了解,性能依赖用户脚本,非透明,有限粒度Op fusion • 整体技术栈完整 3,执⾏引擎依赖nnvm/relay,非原⽣支持Tensorflow • 通用型/透明性需求 4,基于机器学习的schedule tuning,代价较⾼ • 节省前端与Tensorflow框架对接的开发成本 1,厂商提供的优化框架 2,本质上是基于Pattern matching+代数优化⽅法寻优⼈⼯实现的⾼效⼤颗粒度算⼦,极 l 先打GPU后端的编译优化 限性能更好,通用性上存在缺陷 3,支持Inference • 核心环节且存在较大提升空间,打造差异化竞争力 • 可直接优化PAI上近40%作业,结合其它工具优化其余作业 1,Less popular, less active 2,Codegen基于Polyhedral,更为principal,研发曲线更为陡峭 3, MLIR更值得关注,潜在的compiler of compiler 11
12. 挑战 l通用优化在性能方面的挑战 • 社区XLA性能无法满足要求 - 比较简单的CodeGen模版 –固定schedule的单一parallel loop –仅支持多个相同shape的tensor输出 - 后端制约中端,导致保守的Op Fusion策略 l一个新的底层引擎的成熟度问题
13. TAO V1 l社区XLA颗粒度不佳的具体原因
14. TAO V1 l社区XLA颗粒度不佳的具体原因
15. TAO V1 l社区XLA颗粒度不佳的具体原因 此外,在反向计算图中,还观察到大量由于每个 kernel只支持一个输出tensor导致的颗粒度问题
16. TAO V1 l 症结
17. TAO V1 前端及执⾏引擎: 根据实际业务需要, 在可用性和性能等⽅ 面的多处改进 更加激进的fusion pass,去掉原fusion pass的一 些限制,包括支持多输出等 Resource Planning 根据计算图特点划分为多个独立 schedule的子图,节点角色分类 寻找合法及最优LaunchDimension 中端及后端: 针对GPU的⼀套 CodeGen⽅案 计算shared memory用量 基于依赖关系shared memory优化 Code Generatio n Root/SubRoot/Tuple等特殊角色节 点的CodeGen,包括parallel loop, shared/global memory的读写等 部分普通节点的codegen性能改进
18. TAO V1 中端及后端,针对GPU硬件⼀套CodeGen⽅案: 一些早期尝试 TAO Compiler V1 集中分析了一批业务case,V1的细节完善 TAO Compiler V2 前端,执⾏引擎针对实际应用问题⼀些⽅案性改造: Feed-back driven clustering: 优化重新编译开销等可用性问题 Proxy output of cluster: 优化计算与通信的Overlap问题 针对阿里业务特点的更多算⼦类型支持: TopK, Gather.. etc. 约30处其它细节优化和改进: The Devil is in the Detail … ⼀些⼯具: 自动化线上作业类型分析⼯具,简化业务推⼴⼈⼒负担 自动化正确性⼆分⼯具,用于提⾼计算正确性类问题的debug效率 自动化线上作业搜集及离线编译⼯具(TAO Guardian),用于加速产品成熟化过程
19. TAO V1 l解决对Schedule有特殊要求的计算节点问题
20. TAO V1 l解决冗余计算问题 32*1024*256 thread 256倍冗余计算 32*1024 threads 每个处理1个元素,⽆冗余 … … … … 32*1024*256 thread 32*1024 threads 每个处理256个元素
21. TAO V1 l解决多输出问题 ⼀倍冗余计算 ⽆冗余计算 … … … … 256*1024 thread 256*1024 thread … … 256*1024 thread … 256*1024 thread
22. TAO V1 l中端Op Fusion Pass部分 • 目前是基于局部贪婪的rule-based算法 • producer/consumer间迭代merge • merge时预先进行一遍resource planning - 防止merge后的pattern的shared memory用量超出限制 - 防止合法Blocks数过小导致occupancy过低 22
23. TAO V1 lResource Planning • • • • 标记SubRoot节点,即通过shared memory粘接的root节点 计算合法的Launch Dimension集合 寻找最优Launch Dimension 计算各个SubRoot节点的shared memory用量 lShared Memory Optimizer • liveness analysis • 对liveness确定不重叠的SubRoot节点做grouping,分配offset等 23
24. TAO V1 中端: 基于规则的贪婪算法,预执行resource planning,判断 Merge之后是否会出现shared memory用超,或者是否会出现 blocks数太少影响并行度的情况,如果没有,则执行merge
25. TAO V1 exp节点: fuse后可能引起 冗余计算,使用 shared memory 做隔离 reduce节点: 对schedule有特 殊要求的节点, 标记为使用 shared memory Resource Planning: step 1, 标记需要使用Shared Memory的节点
26. TAO V1 exp节点: 对Launch Dimension 无要求 Resource Planning: step 1, 标记需要使用Shared Memory的节点 step 2,寻找合法Launch Dimension集合 reduce节点: Blocks需要是256 的约数 取交集: Blocks需要是256 的约数
27. TAO V1 Launch Dimension的 合法集: Blocks需要是256的约 数 预估最优的Launch Dimension: (blocks:256, threads:1024) Resource Planning: step 1, 标记需要使用Shared Memory的节点 step 2,寻找合法Launch Dimension集合 step 3,估计最优的Launch Dimension 取交集: Blocks需要是256 的约数
28. TAO V1 每个block一共需要 1024*4 =4096 bytes 每个block一共需要 1*4 = 4 bytes Resource Planning: step 1, 标记需要使用Shared Memory的节点 step 2,寻找合法Launch Dimension集合 step 3,估计最优的Launch Dimension step 4,根据Launch Dimension,计算每个节 点需要的shared memory使用量
29. TAO V1 每个block一共需要 1024*4 = 4096 bytes Shared Memory Planning: 分析各个计算节点的liveness,并进行shared memory资源优化 每个block一共需要 1*4 = 4 bytes 由于reduce计算依赖于exp计算,因此两 者的liveness一定不会重叠,故可以复用 同一块shared memory,所需要的总 shared memory大小为: max(4096,4)= 4096 bytes
30. Launch Setting Tuning Locality 提升SM访存效率 Occupancy 充分利用GPU SM资源 Block数量尽可能多 Block Size尽可能小 社区XLA的策略侧重于Occupancy 最小化Block Size (32), 从而最大化Block数量 提升Cache命中率 Block Size尽可能大
31. Launch Setting Tuning 并不是Block数量越多,Occupancy越高,因为SM数量是有上限的 以output_size=32768的Element-Wise Kernel为例: • block_size = 256就已经能够激活V100所有SM • 进一步减小block_size对Occupancy没有贡献
32. Launch Setting Tuning 我们的策略: • 保证SM激活率:block数量小于SM数量时,减小block size • 增加Locality:block数量大于SM数量2倍时,增大block size • 即通过调节block size尽可能进入下图中的目标区间
33. Feedback-driven Clustering l Dynamic shape引起的JIT编译开销问题 • 第一类:主要节点的shape都在变化 • 第二类:个别节点shape变化 l 基于反馈的处理 • 记录编译次数和引起重新编译的节点 • 分析计算图得到不必要编译的节点集 • 触发重新clustering和编译 l 反馈方案可扩展处理其它问题 • memcpy问题 • deadness analysis问题 33
34. Proxy Output l解决cluster过大时计算与通信不能overlap带来 的性能问题 不打开编译优化时 34
35. Proxy Output l解决cluster过大时计算与通信不能overlap带来 的性能问题 打开编译优化后 35
36. Proxy Output l解决cluster过大时计算与通信不能overlap带来 的性能问题 解决办法 36
37. TAO Compiler V1与社区对比 社区XLA TAO Compiler V1 每个thread处理固定个element的单⼀Parallel Loop 每个thread不同节点处理不同数量的element, 借助SharedMemory缝合的复合Parallel Loop 每个Kernel仅支持多个shape相同的Tensor输出 支持任意shape的多个Tensor输出 Launch Dimension由root节点的shape决定 Launch Dimension由Resource Planning模块根 据整个pattern推导得到 Fused kernel内的节点仅支持默认实现模版,对 硬件特性利用有限 Fused kernel内的节点可支持对schedule有特殊 要求的实现模版,例如reduction的不同实现版 本等,可利用更多硬件特性 Op Fusion的颗粒度较小 Op Fusion的颗粒度较⼤ Tensorflow节点数 社区XLA kernel数 TAO V1 kernel数 LSTM Cell 前向 18 1计算密集型 3访存密集型 1计算密集型 1访存密集型 LayerNorm 前向 42 6访存密集型 1访存密集型
38. 效果 BN Codegen VS cuDNN Lib Call 瓶颈在kernel开销 瓶颈在框架开销 只打开编译优化,Benchmark模型上, TAO V1与社区XLA最新版本的对比数 字 对于GPU kernel为瓶颈的业务,编 译优化与混合精度优化叠加使用时, 1+1>2的实际性能效果
39. 效果 l与混合精度优化的组合,1+1>2 39 BERT End2end time speedup Baseline 343ms 1.00X 混合精度 220ms 1.56X 编译优化 267ms 1.28X 混合精度 +编译优化 135ms 2.54X
40. TAO V2 lTAO V1的不足 • 性能好于社区XLA,但距离硬件性能上界仍有一定空间 • 主要性能问题来自于Rule-based策略的局限性 - 中端Op Fusion策略的问题 - 后端代码生成策略的问题 40
41. TAO V2 lV1存在的问题 • 中端Fusion策略过于贪婪 • 后端Codegen Schedule基于规则,无法进行tuning • 多个节点间的schedule独立 - 冗余的load访存 - 冗余的index和data计算 • schedule独立导致无法支持ValueCache
42. TAO V2 lTAO CodeGen V1的局限性 在不同具体size的情况下,上面三种CodeGen策略存在寻优空间,⽽ Rule-based的策略则存在局限
43. TAO V2 lTAO CodeGen V1的局限性 被shared memory缝合在 ⼀起的多个⼦kernel之 间的schedule彼此独立, 带来冗余的访存和计算。 如果⽆法被L1/L2 Cache 命中,则最终体现在端 到端耗时上。
44. TAO V2 l 可枚举的CodeGenPlan l Cost Model • CodeGen Simulator - Launch Dimension - 具体访存行为 - 计算指令数估计 • 根据参数估计开销
45. Take away l Data-driven的优化方法论 l 通用优化 v.s. 手工优化 l 非计算密集算子编译优化的挑战---计算图复杂性 l V1基于片上硬件资源特性的软硬件协同优化---分而治之 l V2进一步松驰了对片上硬件资源特性的要求 l 从rule-based(V1)到principal approach(V2)的探索过程 45
46. 46
47. 47
48. Thanks We are hiringJ Industry hire/Research Intern@杭州/北京/上海 muzhuo.yj@alibaba-inc.com 48