[논문 리뷰] Dynamic Automatic Differentiation of GPU Broadcast Kernels

Jarrett Revels, Tim Besard외 3명|arXiv (Cornell University)|2018. 10. 18.

Parallel Computing and Optimization Techniques참고 문헌 22인용 수 4

한 줄 요약

이 논문은 GPU 최적화된 브로드캐스트 연산을 효율적으로 미분하기 위해 전진모드와 역행모드 AD를 조합한 혼합모드 자동미분 기법을 제안한다. 브로드캐스트 자이아드 행렬의 희박성 특성을 활용하고, 브로드캐스트된 부분 그래프에 대한 역행모드 역전파를 피하므로써, 데이터에 의존하는 제어 흐름을 거쳐도 효율적인 미분이 가능해지며, 특히 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를 통해 자이아드 행렬을 계산하는 방법을 채택한다.
• 다차원 이중수를 전진모드에서 사용하여 방향도를 계산하며, 브로드캐스트된 함수의 입력 개수에 비례하여 페르튜베이션 벡터를 스케일링한다.
• 역행모드 역전파를 위한 중간 값 저장을 피하기 위해 직접 전진모드 도함수를 계산함으로써, 역행모드 전용 제약 조건이 필요 없도록 한다.
• Julia의 JIT 컴파일 파이프라인에 동적으로 코드 변환 단계를 삽입하여, 핵심 컴파일러를 수정하지 않고도 맥락 인식 가능한 AD 동작을 가능하게 한다.
• 제어 흐름이 입력 데이터에 따라 달라지는 경우에도 브로드캐스트 연산의 융합을 지원할 수 있으며, 이는 순수 역행모드에서 역전파 전역 패스가 필요하기 때문에 비현실적이다.
• ReverseDiff, Flux, Zygote 등의 패키지를 통해 Julia 생태계에 통합되었으며, TPUs 지원도 확장되어 있다.

실험 결과

연구 질문

• RQ1 전진모드 AD를 역행모드 AD와 효과적으로 혼합하여 순수 역행모드보다 GPU 최적화된 브로드캐스트 연산을 더 효율적으로 미분할 수 있는가?
• RQ2 브로드캐스트 자이아드 행렬의 희박성을 활용하면, 역행모드 역전파 오버헤드를 피하면서도 성능 향상이 가능한가?
• RQ3 이 혼합모드 접근 방식은 일반적으로 순수 역행모드 AD에서 금지되어 있는 데이터에 의존하는 제어 흐름이 존재하는 상황에서도 브로드캐스트 융합을 지원할 수 있는가?
• RQ4 GPU 아키텍처에서 전진모드 미분의 입력 개수가 증가함에 따라 이 방법의 성능 스케일링은 어떻게 되는가?
• RQ5 실제 딥러닝 워크로드, 예를 들어 HM-LSTM 셀 업데이트에서 이 방법은 순수 역행모드 구현보다 얼마나 뛰어나게 성능을 발휘하는가?

주요 결과

• 혼합모드 AD 접근 방식은 GPU에서 HM-LSTM 셀 업데이트 계산 내의 브로드캐스트 연산을 미분할 때 순수 역행모드 TensorFlow 및 Julia 구현보다 뛰어난 성능을 보였다.
• 이 방법은 제어 흐름이 입력 데이터에 따라 달라지는 경우에도 바탕 코드에서 브로드캐스트 융합을 가능하게 하며, 이는 순수 역행모드 AD에서 역전파 제약 조건으로 인해 지원되지 않는다.
• GPU에서 고차원 입력 개수를 가진 브로드캐스트 연산의 전진모드 미분은 입력 개수가 5를 초과할 경우 레지스터 압박 증가로 인해 성능 저하를 겪는다.
• 18개의 입력 인수에서 계산 및 대역폭 활용률이 60% 이하로 떨어져, 낮은 할당률로 인해 커널이 지연 시간 기반으로 전환되었다.
• 이 기법은 이미 Flux와 Zygote와 같은 실사용 기반의 Julia 기반 머신러닝 프레임워크에 생산 환경에 도입되어 실용성과 성능 향상을 입증했다.
• TPU 지원 확장으로 인해 GPU 아키텍처를 초월한 광범위한 적용 가능성을 보였다.