[2022-03-13 업데이트: 전문가 선택 라우팅 추가]

[2022-06-10 업데이트]: Greg 와 저는 이 게시글의 간결하고 업그레이드된 버전을 작성하여 OpenAI Blog 에서 "Large Neural Networks 훈련을 위한 기술"이라는 제목으로 발표했습니다.

최근 몇 년 동안 우리는 더 큰 사전 학습 언어 모델 (pre-trained language models) 을 사용하여 많은 NLP 벤치마크 작업에서 더 나은 결과를 보고 있습니다. 초대규모 신경망 훈련은 GPU 메모리 요구량이 개별 GPU 머신에 호스팅될 수 있는 것보다 훨씬 많고, 긴 훈련 시간 지평을 필요로 하기 때문에 도전적입니다.

그러나 개별 GPU 워커는 제한된 메모리를 가지고 있으며, 많은 초대규모 모델의 크기는 단일 GPU 를 넘어섰습니다. 여러 GPU 를 통해 모델 훈련을 가능하게 하는 몇 가지 병렬성 패러다임이 있고, 매우 큰 신경망을 훈련할 수 있도록 도와주는 다양한 모델 아키텍처 및 메모리 절약 설계가 있습니다.

훈련 병렬성

초대규모 신경망 모델 훈련의 주요 병목 현상은 개별 GPU 머신에 호스팅될 수 있는 것보다 훨씬 많은 GPU 메모리에 대한 강렬한 요구입니다. 모델 가중치 (예: 수십 억 개의 부동소수점 숫자) 외에도 중간 계산 출력인 경계 (gradients) 와 옵티마이저 상태 (예: Adam 의 모멘텀 및 변이) 를 저장하는 것은 더욱 비쌉니다. 또한 큰 모델을 훈련하는 것은 종종 큰 훈련 코퍼스 (training corpus) 와 짝을 이루며, 따라서 단일 프로세스는 영원히 걸릴 수 있습니다.

결과적으로 병렬성은 필수적입니다. 병렬성은 데이터, 모델 아키텍처, 텐서 연산 등 다른 차원에서 발생할 수 있습니다.

데이터 병렬성

데이터 병렬성 (DP) 의 가장 단순한 방법은 동일한 모델 가중치를 여러 워커에 복사하고 각 워커에 데이터를 할당하여 동시에 처리하는 것입니다.

모델 크기가 단일 GPU 노드의 메모리보다 큰 경우, Naive DP 는 잘 작동하지 않습니다. GeePS (Cui et al. 2016) 와 같은 방법은 모델이 한 대의 머신에 들어맞지 않을 때 제한된 GPU 메모리를 사용하도록 사용할 수 없는 파라미터를 일시적으로 CPU 로 오프로드합니다. 데이터 스왑 전송은 백엔드에서 발생해야 하며 훈련 계산과 간섭하지 않아야 합니다.

각 미니 배치 (minibatch) 의 끝에는 경계 또는 가중치를 동기화하여 신선도를 방지해야 합니다. 두 가지 주요 동기화 접근 방식이 있고 각각 명확한 장단점이 있습니다.

Bulk synchronous parallels (BSP): 워커는 각 미니 배치의 끝에서 데이터를 동기화합니다. 이는 모델 가중치 신선도를 방지하고 좋은 학습 효율성을 보장하지만, 각 머신은 다른 머신이 경계를 보내기 위해 멈추고 기다려야 합니다.Asynchronous parallel (ASP): 모든 GPU 워커는 비동기적으로 데이터를 처리하며, 대기나 정지 없이 진행됩니다. 그러나 이는 신선한 가중치를 사용하도록 쉽게 유도할 수 있으며, 따라서 통계적 학습 효율성을 낮출 수 있습니다. 계산 시간을 증가시키더라도 수렴까지 훈련 시간을 가속화할 수는 없습니다.

그 어느 것도 중간에 경계를 전역적으로 $x$ 번의 반복 ($x > 1$) 에 한 번 동기화하는 것입니다. 이 기능은 Pytorch v1.5 (Li et al. 2021) 에서 Distribution Data Parallel (DDP) 에서 "gradient accumulation"으로 불립니다. Bucketing 경계는 즉시 AllReduce 연산을 수행하지 않고 여러 경계를 하나의 AllReduce 로 버킷화하여 성능을 향상시킵니다. 계산 그래프에 기반한 계산 및 통신 스케줄링 최적화를 할 수 있습니다.

모델 병렬성

Model Parallelism (MP) 는 모델 가중치가 단일 노드에 들어가지 않을 경우를 해결하기 위해 설계되었습니다. 계산과 모델 파라미터는 여러 머신에 걸쳐 분할됩니다. 데이터 병렬성 (Data Parallelism) 과 달리, 각 작업자는 전체 모델의 완전한 사본을 호스팅하지 않고, 대신 하나의 작업자마다 모델 파라미터의 일부만 할당되므로 메모리 사용량과 계산량이 모두 감소합니다.

심층 신경망은 일반적으로 수직 레이어의 스택을 포함하므로, 큰 모델을 레이어별로 분할하는 것은 직관적입니다. 작은 연속된 레이어 세트가 하나의 작업자에서 하나의 파티션으로 그룹화됩니다. 그러나 순차적 의존성을 가진 여러 작업자를 통해 각 데이터 배치 (batch) 를 통과시키는 단순한 구현은 대기 시간의 큰 버블 (bubbles) 을 발생시키고 계산 자원의 심각한 비효율성을 초래합니다.

파이프라인 병렬성 (Pipeline Parallelism)

파이프라인 병렬성 (PP) 은 모델 병렬성과 데이터 병렬성을 결합하여 비효율적인 시간 "버블"을 줄입니다. 주요 아이디어는 하나의 미니배치를 여러 마이크로배치로 분할하고, 각 단계 작업자가 동시에 하나의 마이크로배치를 처리하도록 하는 것입니다. 모든 마이크로배치는 두 번의 통과 (forward 및 backward) 를 필요로 합니다. 작업자 간 통신은 활성화 (forward) 와 기울기 (backward) 만 전달합니다. 이러한 통과가 어떻게 스케줄링되고 기울기가 어떻게 집계되는지는 다양한 접근 방식에서 다릅니다. 파티션 (작업자) 의 수를 파이프라인 깊이라고도 합니다.

GPipe (Huang et al. 2019) 는 여러 마이크로배치에서 오는 기울기를 최종적으로 동기적으로 집계하고 적용합니다. 동기적 경사 하강은 작업자의 수와 무관하게 학습의 일관성과 효율성을 보장합니다. Fig. 3 에 따르면, 버블은 여전히 존재하지만 Given $m$ 개의 균등으로 분할된 마이크로배치와 $d$ 개의 파티션을 가정할 때, 마이크로배치당 forward 와 backward 가 각각 1 단위의 시간을 소요한다고 가정하면, 버블의 비율은 다음과 같습니다:

GPipe 논문에서는 마이크로배치의 수가 파티션 수보다 4 배 이상 많을 경우 ($m > 4d$, 활성화 재계산을 적용할 때), 버블 오버헤드가 거의 무시할 수 있다고 관찰했습니다.

GPipe 는 장치의 수에 따라 거의 선형적인 속도 향상을 달성하지만, 작업자 사이에 모델 파라미터가 균등하게 분배되지 않은 경우에는 항상 보장되지 않습니다.

PipeDream (Narayanan et al. 2019) 은 각 작업자가 순차적으로 forward 와 backward 통과를 처리하도록 스케줄링합니다 (1F1B).

PipeDream 은 각 모델 파티션을 "단계 (stage)"라고 명명하며, 각 단계 작업자는 데이터 병렬성을 실행하기 위해 여러 복제본을 가질 수 있습니다. 이 과정에서 PipeDream은 동일한 미니배치의 forward 와 backward 통과가 동일한 복제본에서 발생하도록 여러 단계의 복제본 간 작업을 할당하기 위해 결정론적 라운드-robin 로드 밸런싱 전략을 사용합니다.

PipeDream 은 모든 작업자 간의 배치 끝에서의 전역 기울기 동기화를 가지지 않기 때문에, 1F1B 의 네이티브 구현은 마이크로배치의 forward 와 backward 통과가 서로 다른 버전의 모델 가중치를 사용하여 수행될 수 있어 학습 효율성을 낮출 수 있습니다. PipeDream 은 이러한 문제를 해결하기 위해 몇 가지 디자인을 제안했습니다:

Weight stashing: 각 worker 는 여러 모델 버전을 추적하며, 하나의 데이터 배치에 대해 전방 (forward) 과 후방 (backward) 패스에서 동일한 버전의 가중치 (weights) 를 사용하도록 보장합니다.
Vertical sync(선택 사항): 모델 가중치의 버전은 활성화 (activations) 와 그래디언트 (gradients) 와 함께 스테이지 worker 간에 흐릅니다. 그런 다음 계산은 이전 worker 에서 전파된 해당 저장 버전을 따릅니다. 이 과정은 worker 간 버전 일관성을 유지합니다. 이는 GPipe 와는 달리 비동기적입니다.

훈련 실행의 시작에서 PipeDream 은 먼저 모델 내 각 레이어의 계산 메모리 비용과 시간을 프로파일링한 후, 레이어를 스테이지로 분할하는 솔루션을 최적화하며, 이는 동적 프로그래밍 문제입니다.

PipeDream 의 두 가지 변형이 제안되어 스테이지된 모델 버전을 통해 메모리 발자국을 줄였습니다 (Narayanan et al. 2021).

PipeDream-flush 는 GPipe 와 마찬가지로 주기적으로 전역 동기화 파이프라인 플러시를 추가합니다. 이 방식은 throughput 를 조금 희생하더라도 메모리 발자국을 크게 줄입니다 (즉, 단일 버전의 모델 가중치만 유지).

PipeDream-2BW 는 두 가지 버전의 모델 가중치를 유지하며,

한 개의 심층 신경망 (deep neural network) 내에서 엔semble 는 여러 전문가 (experts) 를 연결하는 게이트 메커니즘 (gating mechanism) 으로 구현할 수 있습니다 (Shazeer et al., 2017). 게이트 메커니즘은 출력 생성에 어떤 네트워크의 부분집합 (예: 어떤 전문가들) 을 활성화해야 하는지를 제어합니다. 이 논문에서는 이를