Skip to content

Latest commit

 

History

History
187 lines (139 loc) · 7.85 KB

README.ko.md

File metadata and controls

187 lines (139 loc) · 7.85 KB
Raw

torchgpipe

PyPI Build Status Coverage Status Documentation Status English README

PyTorch 용 GPipe 구현입니다. TPU 대신 CUDA를 활용합니다.

from torchgpipe import GPipe
model = nn.Sequential(a, b, c, d)
model = GPipe(model, balance=[1, 1, 1, 1], chunks=8)
output = model(input)

GPipe란?

GPipe는 Google Brain에서 발표한 학습 기법으로, 메모리를 많이 차지하는 큰 모델을 효율적으로 학습시키는 데 유용합니다. Google이 공개한 논문의 벤치마크에 따르면 기준보다 8배 많은 장치(TPU)로 25배 큰 모델을 학습시킬 수 있고, 기준보다 4배 많은 장치에서 3.5배 빨리 학습시킬 수 있다고 합니다.

GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

Google은 GPipe를 이용해 5.6억개의 패러미터를 가지는 AmoebaNet-B 모델을 학습시켰습니다. 이 모델은 ImageNet에서 top-1 정확도 84.3%, top-5 정확도 97.0%로 SOTA를 기록하고 있습니다. (2019년 5월 기준)

GPipe는 Pipeline Parallelism과 Checkpointing, 두 가지 방법으로 가능한 큰 모델을 학습시킵니다.

Pipeline Parallelism
우선 GPipe는 모델을 여러 파티션으로 나눠 각각 서로 다른 장치에 배치해 더 많은 메모리를 사용할 수 있게 한다. 그리고 여러 파티션이 최대한 병렬적으로 작동할 수 있도록, 모델에 입력되는 미니배치를 여러 마이크로배치로 나눠서 모델에 흘려보낸다.
Checkpointing
각 파티션엔 체크포인트를 만들어 메모리 가용량을 극대화한다. 순전파(forward propagation) 때 파티션 경계의 입출력만 기억하고 내부의 히든레이어는 휘발시킨다. 휘발된 히든레이어는 역전파(backpropagation) 때 다시 계산된다.

사용법

현재 torchgpipe는 다음 환경을 지원합니다:

  • Python 3.6 이상
  • PyTorch 1.1 이상

우선 torchgpipe를 PyPI에서 설치합니다:

$ pip install torchgpipe

임의의 nn.Sequential 모듈을 torchgpipe.GPipe로 감싸면 GPipe가 적용됩니다. balance 인자는 각 파티션의 레이어 개수를 정합니다. chunks 인자는 마이크로배치 개수를 설정합니다. 모듈의 입출력과 각 파티션 경계의 입출력은 모두 Tensor 혹은 Tuple[Tensor, ...] 형식이어야 합니다.

다음 예제코드는 총 4층으로 이뤄진 모듈을 각각 1층씩 지니는 4개의 파티션으로 나누는 방법을 보여줍니다. 마이크로배치 개수는 8개로 설정했습니다:

from torchgpipe import GPipe

model = nn.Sequential(a, b, c, d)
model = GPipe(model, balance=[1, 1, 1, 1], chunks=8)

for input in data_loader:
    output = model(input)

문서화

API 문서를 비롯한 자세한 문서는 torchgpipe.readthedocs.io에서 확인할 수 있습니다.

벤치마크

각 벤치마크의 자세한 내용과 추가적인 벤치마크는 torchgpipe.readthedocs.io에서 확인할 수 있습니다.

ResNet-101 정확도 벤치마크

배치크기 torchgpipe nn.DataParallel Goyal et al.
256 21.99±0.13 22.02±0.11 22.08±0.06
1K 22.24±0.19 22.04±0.24 N/A
4K 22.13±0.09 N/A N/A

GPipe를 사용해 학습할 때 추가적인 하이퍼파라미터 조정이 없길 바랍니다. 이를 검증하기 위해 Accurate, Large Minibatch SGD 논문의 표2(c)에 보고된 ResNet-101 정확도(오답률) 벤치마크를 재현했습니다.

U-Net (B, C) 메모리 벤치마크

실험 U-Net (B, C) 파라미터 메모리 사용량
baseline (6, 72) 362.2M 20.3 GiB
pipeline-1 (11, 128) 2.21B 20.5 GiB
pipeline-2 (24, 128) 4.99B 43.4 GiB
pipeline-4 (24, 160) 7.80B 79.1 GiB
pipeline-8 (48, 160) 15.82B 154.1 GiB

GPipe로 얼마나 큰 U-Net 모델을 학습시킬 수 있는지 측정했습니다. baseline은 GPipe를 적용하지 않은 경우를 나타내고, pipeline-1, -2, -4, -8은 GPipe를 적용했을 때 GPU 수에 따른 경우를 나타냅니다.

이 벤치마크엔 간략화한 U-Net 구조를 사용했습니다. 모델의 크기는 하이퍼파라미터 B와 C로 결정합니다. 각각 레이어 수와 필터 수에 비례합니다.

U-Net (5, 64) 속도 벤치마크

실험 처리량 속도향상
baseline 28.500/s 1×
pipeline-1 24.456/s 0.858×
pipeline-2 35.502/s 1.246×
pipeline-4 67.042/s 2.352×
pipeline-8 88.497/s 3.105×

U-Net 구조는 여러 롱스킵커넥션을 포함합니다. 스킵커넥션에서의 효율성을 검증하기 위해 U-Net에서 GPU 수에 따른 처리량을 측정했습니다.

AmoebaNet-D (18, 256) Speed Benchmark

실험 처리량 torchgpipe Huang et al.
n=2, m=1 26.733/s 1× 1×
n=2, m=4 41.133/s 1.546× 1.07×
n=2, m=32 47.386/s 1.780× 1.21×
n=4, m=1 26.827/s 1.006× 1.13×
n=4, m=4 44.543/s 1.680× 1.26×
n=4, m=32 72.412/s 2.711× 1.84×
n=8, m=1 24.918/s 0.932× 1.38×
n=8, m=4 70.065/s 2.625× 1.72×
n=8, m=32 132.413/s 4.966× 3.48×

(n: 파티션 수, m: 마이크로배치 수)

GPipe 논문의 표2에 보고된 AmoebaNet-D (18, 256) 학습 속도 벤치마크를 재현했습니다. 논문의 K를 n으로 바꿔 표기했습니다.

참고사항

이 프로젝트는 개발진이 의도한대로 동작하나, 아직 인터페이스가 확정되지 않았습니다. v0.1.0 전까지는 공개된 API가 경고 없이 바뀔 수 있습니다.

개발진 및 사용권

torchgpipe 프로젝트는 카카오브레인의 이흥섭, 정명룡, 김치헌이 개발하고 임성빈, 김일두, 백운혁, 윤부근의 도움을 받았습니다. BSD-3-Clause 사용권으로 배포됩니다.

인용

해당 라이브러리를 연구용으로 사용할 경우, 아래 BibTeX 링크를 인용해야 합니다.

@article{kim2020torchgpipe,
    title={torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models},
    author={Chiheon Kim and Heungsub Lee and Myungryong Jeong and Woonhyuk Baek and Boogeon Yoon and Ildoo Kim and Sungbin Lim and Sungwoong Kim},
    year={2020},
    eprint={2004.09910},
    archivePrefix={arXiv}
}