torchgpipe

토치 파이프

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에서 게시 한 확장 가능한 파이프 라인 병렬 처리 라이브러리로, 메모리 소비 모델을 효율적으로 교육 할 수 있습니다. 이 논문에 따르면, GPIPE는 8 배의 장치 (TPU)를 사용하여 25 배 더 큰 모델을 훈련시키고 4 배의 장치를 사용하여 3.5 배 더 빨리 모델을 훈련시킬 수 있습니다.

GPIPE : 파이프 라인 병렬 처리를 사용한 거대한 신경망의 효율적인 교육

Google은 GPIPE를 통해 557m 매개 변수로 Amoebanet-B를 훈련 시켰습니다. 이 모델은 Imagenet Classification 벤치 마크 (2019 년 5 월 현재 최신 성과)에서 84.3% Top-1 및 97.0% Top-5 정확도를 달성했습니다.

GPIPE는 (a) 파이프 라인 병렬 처리 및 (b) 역전 전파 동안 전방 전파의 자동 재 계산을 사용하므로 대형 모델 교육을 활용합니다. 우리는 (b)를 Pytorch 커뮤니티의 잘 알려진 용어에 따라 검사 점으로 언급합니다.

파이프 라인 병렬 처리

GPIPE는 모델을 여러 파티션으로 분할하고 다른 장치에 각 파티션을 배치하여 더 많은 메모리 용량을 차지합니다. 그리고 파티션이 가능한 한 평행하게 작동하도록하기 위해 미니 배치를 여러 마이크로 배치로 나눕니다.

검사 점

모델별로 전체 메모리 소비를 최소화하기 위해 각 파티션에 검사 점이 적용됩니다. 전파 전파 동안, 파티션 사이의 경계에있는 텐서만이 기억됩니다. 다른 모든 중간 텐서는 휘발되고 필요할 때 역설 중에 재 계산됩니다.

용법

현재 Torchgpipe에는 다음과 같은 환경이 필요합니다.

• 파이썬 3.6+
• Pytorch 1.1+

Torchgpipe를 사용하려면 PYPI를 통해 설치하십시오.

$ pip install torchgpipe

GPIPE로 모듈을 훈련 시키려면 torchgpipe.GPipe 로 마무리하십시오. GPIPE가 연속 레이어가있는 파티션으로 모듈을 자동으로 분할하므로 모듈은 nn.Sequential 에 반항해야합니다. balance 인수는 각 파티션의 레이어 수를 결정합니다. chunks 인수는 마이크로 배치 수를 지정합니다. 입력, 출력 및 중간 텐서는 Tensor 또는 Tuple[Tensor, ...] .

아래 예제 코드는 4 개의 레이어가있는 모듈을 각각 단일 레이어를 갖는 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 정확도 벤치 마크

배치 크기	토치 파이프	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는 추가 하이퍼 파라미터 튜닝을 도입하지 않도록 투명해야합니다. 투명성을 확인하기 위해 Goyal et al.에 의해 정확하고 큰 미니 배트 SGD의 표 2 (c) 에보 고 된 바와 같이 Imagenet에서 RESNET-101의 상위 1 오류율을 재현했습니다.

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

실험	U-Net (B, C)	매개 변수	메모리 사용
기준선	(6, 72)	362.2m	20.3 gib
파이프 라인 -1	(11, 128)	2.21B	20.5 gib
파이프 라인 -2	(24, 128)	4.99b	43.4 gib
파이프 라인 -4	(24, 160)	7.80b	79.1 gib
파이프 라인 -8	(48, 160)	15.82b	154.1 gib

이 표는 GPIPE가 U-NET 모델 스케일링을 용이하게하는 방법을 보여줍니다. 기준선은 파이프 라인 병렬 처리 나 체크 포인팅이없는 기준선을 나타내며, 파이프 라인 -1 , -2 , -4 , -8은 모델이 해당 수의 파티션과 함께 GPIPE로 훈련된다는 것을 나타냅니다.

여기서 우리는 단순화 된 U-Net 아키텍처를 사용했습니다. 모델의 크기는 각각 층 및 필터의 수에 비례하는 하이퍼 파라미터 B 및 C에 의해 결정됩니다.

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

실험	처리량	속도를 높이십시오
기준선	28.500/s	1 ×
파이프 라인 -1	24.456/s	0.858 ×
파이프 라인 -2	35.502/s	1.246 ×
파이프 라인 -4	67.042/s	2.352 ×
파이프 라인 -8	88.497/s	3.105 ×

Skip Connections의 효율성을 확인하기 위해 다양한 수의 장치로 U-NET의 처리량을 측정했습니다. U-NET에는 여러 개의 긴 건너 뛰기 연결이 있으므로 사용하기로 결정했습니다.

Amoebanet-D (18, 256) 속도 벤치 마크

실험	처리량	토치 파이프	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 : 마이크로 배치 수)

이 표는 Huang et al. 논문의 k를 n 으로 교체했습니다.

메모

이 프로젝트는 기능적이지만 인터페이스는 아직 확인되지 않았습니다. 모든 공개 API는 V0.1.0까지 경고없이 변경 될 수 있습니다.

저자와 라이센스

Torchgpipe Project는 Kakao Brain의 Heungsub Lee, Myungryong Jeong 및 Chheon Kim이 Sungbin Lim, Ildoo Kim, Woonhyuk Baek 및 Boogeon Yoon의 도움으로 개발했습니다. 3 명 BSD 라이센스에 따라 배포됩니다.

소환

이 라이브러리를 프로젝트 및 연구에 적용하면 코드를 인용하십시오.

 @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}
}