[论文解读] Dynamic Automatic Differentiation of GPU Broadcast Kernels

Jarrett Revels, Tim Besard+3 more|arXiv (Cornell University)|Oct 18, 2018

Parallel Computing and Optimization Techniques参考文献 22被引用 4

一句话总结

本文提出了一种混合模式自动微分技术，结合前向模式与反向模式AD，以高效地对GPU优化的广播操作进行微分。通过利用广播雅可比矩阵的稀疏性，并避免在广播子图上进行反向模式反向传播，该方法实现了对数据依赖控制流的高效微分，并在GPU工作负载上优于纯反向模式方法，尤其在HM-LSTM单元计算中表现更优。

ABSTRACT

We show how forward-mode automatic differentiation (AD) can be employed within larger reverse-mode computations to dynamically differentiate broadcast operations in a GPU-friendly manner. Our technique fully exploits the broadcast Jacobian's inherent sparsity structure, and unlike a pure reverse-mode approach, this "mixed-mode" approach does not require a backwards pass over the broadcasted operation's subgraph, obviating the need for several reverse-mode-specific programmability restrictions on user-authored broadcast operations. Most notably, this approach allows broadcast fusion in primal code despite the presence of data-dependent control flow. We discuss an experiment in which a Julia implementation of our technique outperformed pure reverse-mode TensorFlow and Julia implementations for differentiating through broadcast operations within an HM-LSTM cell update calculation.

研究动机与目标

• 解决纯反向模式自动微分在处理GPU上具有数据依赖控制流的广播操作时的局限性。
• 开发一种在反向模式计算中引入前向模式AD的方法，以减少内存和计算开销。
• 即使存在数据依赖控制流，也能在原始代码中实现广播融合，而这类限制在纯反向模式框架中通常无法实现。
• 证明混合模式AD在GPU工作负载上相比纯反向模式实现具有更优性能。
• 推动下一代微分编译器的设计，使其能够根据局部结构动态选择每个子程序的最佳AD模式。

提出的方法

• 该方法采用可与反向模式交错的前向模式AD，通过利用广播操作固有的稀疏结构来计算雅可比矩阵。
• 在前向模式中使用多维双数来计算方向导数，扰动向量按广播函数的输入arity进行缩放。
• 通过直接计算前向模式导数，避免为反向模式反向传播存储中间值，从而消除了对反向模式特有约束的需求。
• 该实现动态地将代码转换阶段注入Julia的JIT编译管道，实现在不修改核心编译器的前提下实现上下文感知的AD行为。
• 该技术即使在控制流依赖于输入数据时，也能支持广播操作的融合，这在纯反向模式中由于需要完整反向传递而不可行。
• 该方法已通过ReverseDiff、Flux和Zygote等包集成到Julia生态系统中，并已扩展以支持TPU。

实验结果

研究问题

• RQ1 前向模式AD能否有效与反向模式AD交错，以比纯反向模式更高效地微分GPU优化的广播操作？
• RQ2 利用广播雅可比矩阵的稀疏性是否能实现性能提升，同时避免反向模式反向传播的开销？
• RQ3 该混合模式方法能否在存在数据依赖控制流的情况下支持广播融合，而这类情况在纯反向模式AD中通常被禁止？
• RQ4 在GPU架构上，前向模式微分中输入arity增加时，该方法的性能如何扩展？
• RQ5 在真实世界深度学习工作负载（如HM-LSTM单元更新）中，该方法相比纯反向模式实现的性能优势有多大？

主要发现

• 在GPU上对HM-LSTM单元更新计算中的广播操作进行微分时，混合模式AD方法在TensorFlow和Julia的纯反向模式实现中表现更优。
• 即使控制流依赖于输入数据，该方法也能在原始代码中实现广播融合，而这是纯反向模式AD由于反向传播约束而无法支持的。
• 在GPU上对高arity广播操作进行前向模式微分时，当输入arity超过5时，由于寄存器压力增加，导致占用率和硬件利用率下降。
• 在18个输入参数时，计算和带宽利用率均降至60%以下，表明由于占用率低，内核已变为延迟受限。
• 该技术已在Flux和Zygote等基于Julia的生产级机器学习框架中实际部署，证明了其实际可行性和性能增益。
• 该方法已扩展以支持TPU，表明其在GPU架构之外也具有更广泛的应用潜力。