[논문 리뷰] Auto-Vectorizing TensorFlow Graphs: Jacobians, Auto-Batching And Beyond

Ashish Agarwal, Igor Ganichev|arXiv (Cornell University)|2019. 03. 08.

Scientific Computing and Data Management참고 문헌 18인용 수 6

한 줄 요약

이 논문은 텐서플로의 고수준 데이터플로우 IR에서 정적 루프 벡터화 최적화를 도입하여, 새로운 기호적 병렬 for(abstraction)를 통해 자동 배치 처리, 예측별 기울기 계산, 야코비안 계산을 효율적으로 가능하게 한다. 이 방법은 텐서플로의 기본 루프 구현 대비 GPU에서 최대 60배 빠르고, DyNet의 동적 배치 처리 대비 38배 이상 빠르며, 특히 큰 배치 크기와 고차원 출력에서 뛰어난 성능을 발휘한다.

ABSTRACT

We propose a static loop vectorization optimization on top of high level dataflow IR used by frameworks like TensorFlow. A new statically vectorized parallel-for abstraction is provided on top of TensorFlow, and used for applications ranging from auto-batching and per-example gradients, to jacobian computation, optimized map functions and input pipeline optimization. We report huge speedups compared to both loop based implementations, as well as run-time batching adopted by the DyNet framework.

연구 동기 및 목표

• 자동 배치 처리, 예측별 기울기 계산, 야코비안 계산 등의 연산에서 텐서플로의 루프 기반 구현 방식이 초래하는 성능 저하 문제를 해결하기 위해.
• DyNet과 같은 프레임워크에서 동적 배치 처리와 관련된 런타임 오버헤드를 줄이기 위해 정적 컴파일 타임 벡터화를 가능하게 하기 위해.
• 하드웨어 활용도를 향상시키고 기계학습 모델에서 중첩된 제어 흐름 및 동적 제어 흐름 패턴을 효율적으로 실행하기 위해.
• 텐서플로의 고수준 IR에서 임의의 루프를 벡터화하기 위한 일반적이고 재사용 가능한 추상화를 제공하기 위해.
• 현재 최적화가 잘 되어 있거나 지원되지 않는 야코비안과 헤시안 계산을 텐서플로가 내장 기능으로 효율적으로 처리할 수 있도록 하기 위해.

제안 방법

• 그래프 생성 시 텐서플로의 고수준 데이터플로우 IR에서 작동하는 새로운 정적 벡터화 병렬 for(pfor) 추상화를 도입한다.
• 루프 본문을 벡터화된 연산으로 융합함으로써 정적 루프 벡터화를 적용하여, 비잔티노스 병렬 연산에 대해 루프를 제거한다.
• 기호 텐서를 사용해 벡터화된 루프 반복의 출력을 표현함으로써, 중첩된 제어 흐름과 동적 형태 간의 최적화를 가능하게 한다.
• 배치 요소에 대한 순전파를 루프로 감싸고 전체 루프를 벡터화함으로써, pfor 추상화를 사용해 자동 배치 처리 모델을 구현한다.
• 예측별 기울기 계산을 위해 단일 pfor 구조 내에서 순전파와 역전파를 모두 벡터화한다.
• 각 출력 요소를 개별 스칼라 연산으로 간주함으로써, 야코비안 계산을 위해 기울기 계산을 각 출력 요소에 대해 벡터화한다.

실험 결과

연구 질문

• RQ1 텐서플로의 고수준 IR에서 정적 루프 벡터화가 자동 배치 처리 및 기울기 계산에 있어 동적 루프 기반 구현 방식보다 뚜렷이 뛰어난 성능을 발휘할 수 있는가?
• RQ2 pfor 기반의 벡터화가 비스칼라 출력 및 RNN과 같은 복잡한 모델에 대해 야코비안 계산 성능 향상에 얼마나 기여하는가?
• RQ3 다양한 배치 크기와 출력 차원에서 pfor의 성능이 DyNet의 런타임 배치 처리와 비교해 어떻게 되는가?
• RQ4 변수 길이 시퀀스 또는 중첩 루프와 같은 상황에서 pfor의 현재 구현에서 성능 저하 요인은 무엇인가?
• RQ5 벡터화된 제어 흐름을 통해 이전에는 불가능하거나 비효율적이었던 연산, 예를 들어 헤시안 계산이나 입력 파이프라인 최적화를 pfor가 가능하게 할 수 있는가?

주요 결과

• 128의 높은 출력 크기에서, pfor 기반 구현은 야코비안 계산에서 텐서플로의 기본 tf.while_loop 대비 최대 60배 빠른 성능 향상을 보였다. 특히 VGG16 및 LSTM 모델에서 유의미한 성능 향상을 보였다.
• LSTM에 대한 예측별 기울기 계산에서, 배치 크기가 256일 때 pfor는 tf.while_loop 대비 38배 빠른 성능을 기록했으며, GPU 활용도가 뛰어나다는 것을 입증했다.
• MNIST 데이터셋에서 pfor는 tf.while_loop 대비 일관된 처리량 향상을 보였으며, 더 큰 배치 크기에서 더 나은 벡터화 및 메모리 접근 패턴 덕분에 성능 향상이 두드러졌다.
• 출력 차원이 커질수록 pfor와 tf.while_loop 간의 성능 격차가 점점 커지며, 이는 벡터화 접근 방식의 우수한 확장성을 시사한다.
• LSTM 모델은 변수 길이 시퀀스와 tf.while_loop를 벡터화할 때 발생하는 잔여 오버헤드로 인해 상대적 성능 향상이 낮았다. 이는 동적 형태 처리에 대한 최적화 여전히 개선이 필요함을 시사한다.
• 이전에 지원되지 않았던 연산들, 예를 들어 동적 RNN의 야코비안과 비스칼라 출력의 헤시안 계산이 pfor를 통해 가능해졌으며, 새로운 연구 방향을 열어 놓았다.