[Paper Review] Dynamic Automatic Differentiation of GPU Broadcast Kernels

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

Parallel Computing and Optimization Techniques22 references4 citations

TL;DR

This paper presents a mixed-mode automatic differentiation technique that combines forward-mode and reverse-mode AD to efficiently differentiate GPU-optimized broadcast operations. By exploiting the sparsity of the broadcast Jacobian and avoiding reverse-mode backpropagation over broadcasted subgraphs, the method enables efficient differentiation through data-dependent control flow and outperforms pure reverse-mode approaches on GPU workloads, particularly in HM-LSTM cell computations.

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.

Motivation & Objective

• To address the limitations of pure reverse-mode automatic differentiation in handling broadcast operations with data-dependent control flow on GPUs.
• To develop a method that leverages forward-mode AD within reverse-mode computations to reduce memory and computational overhead.
• To enable broadcast fusion in primal code despite the presence of data-dependent control flow, which is typically restricted in pure reverse-mode frameworks.
• To demonstrate that mixed-mode AD can achieve superior performance on GPU workloads compared to pure reverse-mode implementations.
• To motivate the design of next-generation differentiating compilers that dynamically select optimal AD modes per subprogram based on local structure.

Proposed method

• The method employs a reverse-mode-interleavable forward-mode AD approach that computes Jacobians for broadcast operations by exploiting their inherent sparsity structure.
• It uses multidimensional dual numbers in forward mode to compute directional derivatives, with perturbation vectors scaled to the input arity of the broadcasted function.
• The approach avoids storing intermediate values for reverse-mode backpropagation by directly computing forward-mode derivatives, thus eliminating the need for reverse-mode-specific restrictions.
• The implementation dynamically injects code transformation passes into Julia’s JIT compilation pipeline, enabling context-aware AD behavior without modifying the core compiler.
• The technique supports fusion of broadcast operations even when control flow depends on input data, which is infeasible in pure reverse-mode due to the need for a full backward pass.
• The method is integrated into Julia’s ecosystem via packages like ReverseDiff, Flux, and Zygote, and has been extended to support TPUs.

Experimental results

Research questions

• RQ1 Can forward-mode AD be effectively interleaved with reverse-mode AD to differentiate GPU-optimized broadcast operations more efficiently than pure reverse-mode?
• RQ2 Does exploiting the sparsity of the broadcast Jacobian enable performance improvements while avoiding reverse-mode backpropagation overhead?
• RQ3 Can this mixed-mode approach support broadcast fusion in the presence of data-dependent control flow, which is typically prohibited in pure reverse-mode AD?
• RQ4 How does the performance of this method scale with increasing input arity in forward-mode differentiation on GPU architectures?
• RQ5 To what extent does this approach outperform pure reverse-mode implementations in real-world deep learning workloads like HM-LSTM cell updates?

Key findings

• The mixed-mode AD approach outperformed pure reverse-mode TensorFlow and Julia implementations in differentiating broadcast operations within an HM-LSTM cell update calculation on GPU.
• The method enables broadcast fusion in primal code even when control flow depends on input data, which is not supported by pure reverse-mode AD due to backpropagation constraints.
• Forward-mode differentiation of high-arity broadcast operations on GPU suffers from reduced occupancy and hardware utilization when input arity exceeds 5, due to increased register pressure.
• At 18 input arguments, both compute and bandwidth utilization dropped below 60%, indicating the kernel became latency-bound due to low occupancy.
• The technique is already deployed in production Julia-based ML frameworks like Flux and Zygote, demonstrating practical viability and performance gains.
• The method has been extended to support TPUs, indicating broader applicability beyond GPU architectures.