[Github]
[arXiv](Current version v5)
Abstract
신경망 계층 여러 가속기에 분산하여 다양한 작업에 대한 효율적인 규모 확장을 가능케 하는 라이브러리인 GPipe 제안
The GPipe Library
Interface
네트워크는 L 레이어를 가지고 있고, 각 레이어는 순전파 함수 f, 피라미터 w로 이루어져 있다.
GPipe는 L개의 레이어를 K개의 파티션으로 분할하고, i ~ j 레이어를 나타내는 파티션 p는 피라미터 wi, wi+1,..., wj, 순전파 함수 F, 역전파 함수 B, 비용 추정기 C를 가진다.
Algorithm
k번째 셀을 k번째 가속기에 배치.
미니 배치를 나눈 각 마이크로 배치에 대해 gradient가 누적되고 미니 배치가 완전히 끝났을 때 피라미터 업데이트.
배치 정규화가 필요할 경우 마이크로 배치에 대해 계산.
Performance Optimization
활성화 메모리를 줄이기 위해 순전파 중에는 파티션의 경계에서만 활성화를 저장하고 역전파에서 F를 다시 계산한다.
그림 (c)를 보면 각 가속기의 유휴 시간이 있는데, 이를 bubble overhead라고 한다. 마이크로 배치 M이 클수록 bubble overhead가 줄어든다고 한다.
