[논문 리뷰] TFLMS: Large Model Support in TensorFlow by Graph Rewriting

Tung D. Le, Haruki Imai외 2명|arXiv (Cornell University)|2018. 07. 05.

Topic Modeling참고 문헌 11인용 수 28

한 줄 요약

이 논문은 TFLMS를 제안하며, 이는 텐서플로우에서 그래프 리라이팅 기반의 접근 방식으로, 중간 활성화를 CPU 메모리로 오프로드하기 위해 스왑아웃 및 스왑인 연산을 삽입함으로써 GPU 메모리 제약을 초월해 대규모 딥 뉴럴 네트워크를 훈련시킬 수 있도록 한다. 이 방법은 프로그램 변환 규칙을 사용해 이러한 연산을 형식적으로 유도하며, ResNet-50과 3DUnet에 대해 각각 최대 4.7배와 2배의 더 큰 배치 크기를 달성한다. 이는 이전에 타일링이 필요했던 3DUnet의 전체 $192^3$ 이미지에 대한 훈련을 가능하게 한다.

ABSTRACT

While accelerators such as GPUs have limited memory, deep neural networks are becoming larger and will not fit with the memory limitation of accelerators for training. We propose an approach to tackle this problem by rewriting the computational graph of a neural network, in which swap-out and swap-in operations are inserted to temporarily store intermediate results on CPU memory. In particular, we first revise the concept of a computational graph by defining a concrete semantics for variables in a graph. We then formally show how to derive swap-out and swap-in operations from an existing graph and present rules to optimize the graph. To realize our approach, we developed a module in TensorFlow, named TFLMS. TFLMS is published as a pull request in the TensorFlow repository for contributing to the TensorFlow community. With TFLMS, we were able to train ResNet-50 and 3DUnet with 4.7x and 2x larger batch size, respectively. In particular, we were able to train 3DUNet using images of size of $192^3$ for image segmentation, which, without TFLMS, had been done only by dividing the images to smaller images, which affects the accuracy.

연구 동기 및 목표

• GPU 메모리 용량을 초월하는 매우 큰 딥 뉴럴 네트워크를 훈련시키는 데 도전하는 문제를 해결하기 위해.
• 계산 그래프에 스왑아웃 및 스왑인 연산을 삽입하기 위한 형식적이고 증명 가능한 정확한 방법을 제공하기 위해.
• 그래프 리라이팅을 통해 메모리 사용을 최적화함으로써 대규모 모델의 훈련 효율성과 확장성을 향상시키기 위해.
• 이전에 타일링이 필요했던 3DUnet을 전체 $192^3$ 이미지에서 훈련시킬 수 있도록 하기 위해.

제안 방법

• 변수를 형식적 의미 체계를 가진 일급 시민으로 간주하여 계산 그래프를 재정의함으로써 텐서플로우의 그래프 모델을 확장한다.
• 정확성과 일관성을 보장하기 위해 형식적 프로그램 변환 규칙을 사용해 스왑아웃 및 스왑인 연산을 도출한다.
• 최적의 스왑인 타이밍을 유도하기 위해 제어 의존성을 정적 탐지하는 두 가지 전략인 chain_rule과 direct_order를 도입한다.
• 과도한 오버헤드를 줄이고 성능을 향상시키기 위해 스왑인 연산을 융합하는 최적화 규칙을 적용한다.
• 커뮤니티의 도입을 가능하게 하기 위해 텐서플로우 저장소에 TFLMS라는 이름의 풀 리퀘스트로 구현한다.
• 메모리 사용과 훈련 성능 간의 트레이드오프를 제어하기 위해 구성 가능한 매개변수(n_tensors, lb, swap_branches)를 사용한다.

실험 결과

연구 질문

• RQ1 계산 그래프를 어떻게 형식적으로 재정의할 수 있을까? 이를 통해 일급 변수를 지원하고 메모리 관리를 위한 정확한 그래프 리라이팅이 가능하게 하기 위해.
• RQ2 훈련 중 모델의 정확성을 유지하기 위해 스왑아웃 및 스왑인 연산을 도출하는 데 사용할 수 있는 형식적 규칙는 무엇인가?
• RQ3 제어 의존성을 탐지하는 데 사용되는 다양한 전략이 대규모 모델 훈련 성능에 어떤 영향을 미치는가?
• RQ4 메모리 사용과 배치 크기, 훈련 스루풋을 최대화하기 위해 매개변수(n_tensors, lb, swap_branches)의 최적 설정은 무엇인가?
• RQ5 TFLMS는 3DUnet을 타일링 없이 전체 $192^3$ 이미지에서 훈련시킬 수 있으며, 이에 따른 성능은 어떠한가?

주요 결과

• TFLMS를 통해 ResNet-50의 최대 배치 크기가 832로 증가하여, TFLMS 없이 기준값 176에 비해 4.7배 증가하였다.
• 3DUnet의 경우 TFLMS로 기준값 대비 최대 배치 크기가 2배 증가하여, 타일링 없이 전체 $192^3$ 이미지에서 훈련이 가능해졌다.
• ResNet-50의 경우 lb=5와 n_tensors=100 설정에서 가장 높은 훈련 스루풋이 달성되었으며, 이는 메모리 사용과 계산-통신 오버랩 간의 균형을 맞추기 때문이다.
• 스왑인 연산을 융합함으로써 성능 향상이 이루어졌지만, 중간 텐서의 GPU 메모리 유지로 인해 최대 배치 크기가 감소하였다.
• 20의 임계값을 설정해 브랜치 스왑을 활성화함으로써 3DUnet의 최대 배치 크기가 3에서 4로 증가하여 확장성 향상 효과를 입증하였다.
• 최적의 설정을 통해 3DUnet은 $192^3$ 입력에서 초당 0.17장의 이미지 훈련 스루풋을 달성하였으며, 이는 이전에는 메모리 부족 오류로 인해 불가능했던 성능이었다.