TL;DR

Megatron-LM은 transformer 아키텍처의 고유한 구조를 활용하여 수십억 파라미터 transformer 모델을 학습하기 위한 간단하고 효율적인 레이어 내 모델 병렬화 접근법을 제시하며, 네이티브 PyTorch에서 몇 개의 all-reduce 연산만을 추가하여 구현합니다. 핵심 혁신은 MLP와 self-attention 블록을 모두 분할하여 GPU 간 동기화를 최소화하고, GEMM 쌍을 융합하여 중간 통신 지점을 제거하는 것입니다. 이 접근법은 512개의 GPU에서 83억 파라미터 GPT-2 모델을 학습할 때 76%의 스케일링 효율을 달성하며, WikiText103(10.8 perplexity), LAMBADA(66.5% 정확도), RACE(90.9% 정확도)에서 당시 최고 성능을 기록했습니다. 흔히 간과되지만 중요한 발견은 BERT 계열 모델에서 layer normalization 배치를 재배치하는 것이 대규모 학습의 안정성에 필수적이라는 점이며, 이를 적용하지 않으면 모델 크기가 증가함에 따라 성능이 저하됩니다.

• Paper Link: https://arxiv.org/pdf/1909.08053

Related Papers

모델 병렬화 기법:

• GPipe - 파이프라인 병렬화를 통한 대규모 모델 훈련
• Reducing Activation Recomputation in Large Transformer Models - 메모리 효율적인 병렬 훈련 기법
• Ring Self-Attention - 시퀀스 차원 병렬화

분산 훈련 시스템:

• DeepSpeed Ulysses - 시퀀스 병렬화와 텐서 병렬화 결합
• USP - 통합 시퀀스 병렬화 프레임워크

Takeaways

1. Contribution

Megatron-LM이 해결하는 핵심 과제는 현대 하드웨어에서 매우 큰 transformer 언어 모델을 학습하는 것을 막는 근본적인 메모리 제약입니다. 2019년까지 NLP 커뮤니티는 transformer 모델의 규모를 키우면 질의응답에서 자연어 추론까지 다양한 작업에서 상당한 성능 향상이 나타난다는 강력한 경험적 증거를 축적했습니다. 그러나 학습 가능한 최대 모델은 개별 GPU의 메모리 용량(당시 일반적으로 16 또는 32GB)에 의해 제한되었습니다. 표준 데이터 병렬화는 계산을 분산시키는 데 효과적이지만, 각 워커가 모델의 완전한 복사본을 보유해야 하므로 이 제약을 완화하지 못합니다. Activation checkpointing은 메모리 사용량을 줄일 수 있지만, 모델 파라미터와 optimizer 상태가 단일 디바이스에 상주해야 한다는 근본적인 제약을 극복하지 못합니다.

기존 모델 병렬화 접근법은 존재했지만 심각한 실용적 단점을 겪었습니다. GPipe는 신중한 마이크로 배치 스케줄링이 필요한 파이프라인 스타일 실행을 요구했고, 효율을 감소시키는 파이프라인 버블로 인해 어려움을 겪었으며, optimizer 또는 학습 절차의 수정이 필요했습니다. Mesh-TensorFlow는 분산 텐서 계산을 위한 일반 프레임워크를 제공했지만 커스텀 컴파일러, 병렬화를 명시하기 위한 새로운 프로그래밍 언어가 필요했으며, 기존 코드베이스에 쉽게 통합할 수 없었습니다. 두 접근법 모두 상당한 엔지니어링 노력을 요구하여 모델을 확장하고자 하는 연구자들에게 높은 진입 장벽을 만들었습니다. 커스텀 컴파일러와 프레임워크가 필요하다는 것은 이러한 도구들이 기반이 되는 딥러닝 라이브러리의 빠른 개발에 뒤처진다는 것을 의미했습니다.

Megatron-LM의 주요 기여는 transformer 아키텍처의 특정 구조를 활용하여 최소한의 코드 변경으로 효율적인 병렬화를 달성하는 모델 병렬화 접근법입니다. 범용 병렬화 프레임워크를 구축하는 대신, 저자들은 transformer 모델이 병렬화에 적합한 자연스러운 분해를 가지고 있음을 관찰합니다. Multi-head attention 메커니즘은 이미 계산을 독립적인 head들로 나누고, 2개 레이어의 MLP 블록은 첫 번째 GEMM에 대해 열 방향 분할을, 두 번째에 대해 행 방향 분할을 선택하여 중간 동기화 지점을 제거함으로써 분할될 수 있습니다. 이러한 구조적 통찰을 통해 전체 접근법은 기존 PyTorch 코드에 몇 개의 all-reduce 연산만 삽입하여 구현될 수 있으며, 커스텀 C++ 커널, 컴파일러, 또는 프레임워크 수정이 필요하지 않습니다.

이 연구의 실질적 중요성은 여러 측면에서 상당합니다. 첫째, 접근법의 단순성은 더 넓은 연구 커뮤니티가 쉽게 채택할 수 있음을 의미합니다. 통신을 처리하는 $f$와 $g$ 켤레 연산자는 각각 10줄 미만의 PyTorch 코드로 구현 가능합니다. 둘째, 강력한 단일 GPU 기준선(이 자체로 피크 FLOPS의 30%를 달성)에 대해 512개 GPU에서 76%의 스케일링 효율은 접근법이 최소한의 오버헤드를 도입함을 보여줍니다. 셋째, 논문은 단순히 시스템 기여를 보여주는 것이 아니라 당시 가장 큰 언어 모델을 학습하고 여러 벤치마크에서 최고 성능을 달성함으로써 end-to-end로 검증합니다.

부차적이지만 매우 중요한 기여는 BERT 계열 모델에서 layer normalization 배치가 대규모 학습 가능성에 깊은 영향을 미친다는 발견입니다. Lan et al. (2019)의 이전 연구는 BERT를 Large 구성(336M 파라미터) 이상으로 확장하면 모델 성능이 저하됨을 관찰하고 파라미터 공유를 해결책으로 제안했습니다. Megatron-LM은 근본 원인이 residual connection에 대한 layer normalization 배치임을 식별합니다. Layer normalization을 self-attention과 MLP 하위 레이어 이전에 배치함으로써(원래의 post-norm이 아닌 pre-norm 아키텍처), 저자들은 최대 39억 파라미터의 BERT 모델에 대해 안정적이고 단조롭게 개선되는 학습을 가능하게 합니다. 이 발견은 지속적인 영향을 미쳤으며, pre-norm 아키텍처는 이후 거의 모든 대규모 transformer 학습의 표준이 되었습니다.

논문은 또한 Wikipedia, CC-Stories, RealNews, OpenWebText를 결합하고 locality-sensitive hashing으로 중복 제거한 174GB의 신중하게 큐레이션된 학습 데이터셋을 기여합니다. Residual layer 이전에 $\frac{1}{\sqrt{2N}}$로 초기화를 스케일링하고, 혼합 정밀도 학습, 그래디언트 클리핑을 포함한 학습 레시피는 후속 대규모 학습 작업에 영향을 미친 포괄적인 모범 사례 세트를 나타냅니다. Megatron-LM 코드베이스의 오픈소스 공개는 Turing-NLG 등을 포함한 많은 후속 대규모 모델 학습 프로젝트의 기초 빌딩 블록이 되면서 논문의 영향을 더욱 증폭시켰습니다.

2. Methodology

2.1 핵심 직관

Megatron-LM의 모델 병렬화 접근법의 근본적인 통찰은 transformer 아키텍처가 GPU 경계와 일치하는 자연스러운 분할 지점을 포함하여, 최소한의 통신으로 효율적인 계산 분산을 가능하게 한다는 것입니다. 임의의 계산 그래프에 대해 추론해야 하는 일반적인 텐서 분해 접근법과 달리, 저자들은 transformer의 두 가지 특정 구조적 속성을 활용합니다.

첫 번째 속성은 multi-head attention이 자연스럽게 head들로 분해된다는 것입니다. 각 attention head는 독립적인 key, query, value 프로젝션에서 작동하고, 독립적인 scaled dot-product attention을 수행하며, 독립적인 출력을 생성합니다. 이는 $K$, $Q$, $V$ 프로젝션을 위한 GEMM 연산을 GPU 간에 열 방향으로 분할할 수 있으며, 각 GPU는 head의 하위 집합에 대해 전체 attention을 계산함을 의미합니다. 모든 head의 출력이 최종 출력 프로젝션에서 결합되어야 할 때까지 통신이 필요하지 않습니다.

두 번째 속성은 더 미묘하며 MLP 블록과 관련됩니다. 표준 transformer MLP는 두 개의 선형 레이어와 그 사이의 비선형 활성화(GeLU)로 구성됩니다: 먼저 hidden dimension $H$에서 $4H$로의 확장, 그 다음 $4H$에서 다시 $H$로의 축소입니다. 첫 번째 가중치 행렬이 행을 따라 분할되고 입력이 열을 따라 분할되면, $\text{GeLU}(X_1 A_1 + X_2 A_2) \neq \text{GeLU}(X_1 A_1) + \text{GeLU}(X_1 A_1)$이므로 GeLU 비선형성 이전에 동기화(all-reduce)가 필요합니다. 그러나 첫 번째 가중치 행렬이 대신 열을 따라 분할되면, 각 GPU는 독립적으로 해당 파티션에 GeLU를 적용할 수 있습니다: $[Y_1, Y_2] = [\text{GeLU}(X A_1), \text{GeLU}(X A_2)]$. 이는 동기화 지점을 제거합니다. 그런 다음 두 번째 가중치 행렬은 행을 따라 분할되어 분할된 출력을 직접 받아들이며, 부분 결과를 합산하기 위해 단일 all-reduce만 필요합니다.

전반적인 원칙은 GEMM 쌍을 융합하고 중간 동기화를 피하는 분할 방식을 선택하여 통신을 최소화하는 것입니다. 이 접근법은 GPU를 통신 제한이 아닌 계산 제한 상태로 유지하며, layer normalization과 dropout 같은 저비용 연산을 한 GPU에서 계산하고 브로드캐스트하는 대신 모든 GPU에 걸쳐 중복 계산합니다. 적은 양의 중복 계산을 감소된 통신과 교환하는 이 설계 철학은 높은 스케일링 효율을 달성하는 데 핵심입니다.

2.2 모델 아키텍처

Megatron-LM에서 사용되는 transformer 아키텍처는 두 가지 변형을 갖는 표준 공식을 따릅니다: autoregressive 언어 모델링을 위한 decoder-only GPT-2 아키텍처와 masked 언어 모델링을 위한 encoder-only BERT 아키텍처입니다. 둘 다 GeLU 활성화를 사용하고 원래 Vaswani et al. (2017) transformer처럼 출력이 아닌 각 하위 레이어의 입력에 layer normalization을 적용합니다(pre-norm).

단일 transformer 레이어의 모델 병렬 분해는 다음과 같이 설명될 수 있습니다:

입력 X (모든 GPU에 복제됨)
  |
  v
[f 연산자] -- forward는 identity, backward는 all-reduce
  |
  +---> Self-Attention 블록 (모델 병렬)
  |      K,Q,V 프로젝션을 GPU 간 열 방향으로 분할
  |      각 GPU는 자신의 head에 대해 attention 계산
  |      출력 프로젝션을 행 방향으로 분할
  |
  +---> [g 연산자] -- forward는 all-reduce, backward는 identity
  |
  v
Dropout + Residual (복제됨)
  |
  v
[f 연산자]
  |
  +---> MLP 블록 (모델 병렬)
  |      첫 번째 GEMM (H -> 4H)을 열 방향으로 분할
  |      각 파티션마다 독립적으로 GeLU 적용
  |      두 번째 GEMM (4H -> H)을 행 방향으로 분할
  |
  +---> [g 연산자]
  |
  v
Dropout + Residual (복제됨)
  |
  v
출력 Y (모든 GPU에 복제됨)

각 transformer 레이어는 정확히 4개의 all-reduce 연산을 필요로 합니다: forward pass에서 2개($g$ 연산자마다 1개)와 backward pass에서 2개($f$ 연산자마다 1개). 임베딩 레이어도 병렬화됩니다: 입력 임베딩 행렬 $E_{H \times v}$는 어휘 차원을 따라 열 방향으로 분할됩니다. 출력 임베딩의 경우, 전체 logit을 수집하면($b \times s \times v$ 요소를 통신) 대신, cross-entropy 손실을 병렬 GEMM과 융합하여 통신을 $b \times s$ 차원의 스칼라 손실로 줄입니다.

특히 BERT 아키텍처의 경우, 논문은 원래의 post-norm 배치가 대규모에서 학습 불안정성을 야기함을 식별합니다. 재배치된 아키텍처는 layer normalization을 self-attention과 MLP 블록 이전에 배치하며, transformer 맨 끝에 추가 layer normalization을 둡니다. 이 pre-norm 아키텍처는 336M에서 39억 파라미터까지 안정적이고 단조롭게 개선되는 학습을 가능하게 합니다.

2.3 핵심 알고리즘 및 메커니즘

Megatron-LM의 핵심 알고리즘 메커니즘은 두 개의 켤레 통신 연산자 $f$와 $g$, 그리고 이들이 분할된 가중치 행렬과 상호작용하는 방식을 중심으로 합니다. 이러한 연산자들과 그들이 분할된 가중치 행렬과 어떻게 상호작용하는지 이해하는 것은 재구현에 필수적입니다.

$f$ 연산자는 forward pass에서 identity로 작동하며, 입력을 변경하지 않고 그대로 전달합니다. Backward pass에서는 그래디언트에 대해 all-reduce 연산을 수행합니다. 이는 각 모델 병렬 영역의 시작 부분에서 사용됩니다. 각 GPU는 GEMM의 파티션을 계산하기 위해 동일한 입력 $X$가 필요하므로, forward identity pass는 이를 직접 제공합니다($X$가 이미 복제되어 있기 때문). 역전파 중에는 각 GPU의 부분 계산에서 나온 그래디언트를 합산하여 $X$에 대한 올바른 그래디언트를 생성해야 하므로, backward pass에서 all-reduce를 수행합니다.

$g$ 연산자는 $f$의 켤레입니다: forward pass에서 all-reduce를 수행하고 backward pass에서 identity로 작동합니다. 이는 각 모델 병렬 영역의 끝에서 사용됩니다. Forward pass에서 각 GPU는 부분 결과(예: MLP의 두 번째 GEMM에 대해 $Y_1 B_1$과 $Y_2 B_2$)를 계산했으며, 이를 합산하여 올바른 출력 $Z = Y_1 B_1 + Y_2 B_2$를 생성해야 합니다. Backward pass에서 all-reduce는 합산이므로 그래디언트는 각 파티션을 독립적으로 통과하며(identity), 이는 각 파티션에 대해 $\frac{\partial Z}{\partial (Y_i B_i)} = 1$이므로 수학적으로 올바릅니다.

MLP 병렬화 세부사항. 가중치 행렬 $A \in \mathbb{R}^{H \times 4H}$와 $B \in \mathbb{R}^{4H \times H}$, 그리고 입력 $X \in \mathbb{R}^{b \times s \times H}$를 갖는 MLP를 고려합니다. 여기서 $b$는 배치 크기, $s$는 시퀀스 길이, $H$는 hidden dimension입니다. $p$개의 GPU로 다음과 같이 분할합니다:

각 GPU $i$는 $Y_i = \text{GeLU}(X A_i) \in \mathbb{R}^{b \times s \times \frac{4H}{p}}$를 계산합니다. 두 번째 행렬 $B$는 행 방향으로 분할됩니다:

각 GPU $i$는 $Z_i = Y_i B_i \in \mathbb{R}^{b \times s \times H}$를 계산합니다. 최종 출력은 all-reduce를 통해 얻어집니다: $Z = \sum_{i=1}^{p} Z_i$. 이것이 $g$ 연산자의 forward all-reduce입니다. 핵심 관찰은 $Z = \sum_i Y_i B_i = \sum_i \text{GeLU}(X A_i) B_i$이며, 이는 $A$의 열 분할 다음 $B$의 행 분할이 원래 계산을 재구성하고 GeLU를 각 열 파티션에 독립적으로 적용할 수 있으므로 $\text{GeLU}(X A) B$와 수학적으로 동등하다는 것입니다.

Self-attention 병렬화 세부사항. 총 $h$개의 head와 $p$개의 GPU를 갖는 multi-head attention의 경우, 각 GPU는 $h/p$개의 head를 처리합니다. 프로젝션 행렬 $W^Q, W^K, W^V \in \mathbb{R}^{H \times H}$는 열 방향으로 분할되며, 각 GPU는 $W^Q_i, W^K_i, W^V_i \in \mathbb{R}^{H \times \frac{H}{p}}$를 받습니다. 각 GPU $i$는 다음을 계산합니다:

그런 다음 자신의 $h/p$개 head에 대해 표준 multi-head attention을 수행하여 출력 $O_i \in \mathbb{R}^{b \times s \times \frac{H}{p}}$를 생성합니다. 출력 프로젝션 행렬 $W^O \in \mathbb{R}^{H \times H}$는 $W^O_i \in \mathbb{R}^{\frac{H}{p} \times H}$로 행 방향으로 분할되며, 각 GPU는 $Z_i = O_i W^O_i$를 계산합니다. 그런 다음 all-reduce가 이러한 부분 결과를 합산합니다.

임베딩 병렬화. 입력 임베딩 테이블 $E \in \mathbb{R}^{H \times v}$는 어휘 차원을 따라 열 방향으로 분할됩니다. 토큰 임베딩을 조회한 후, all-reduce($g$ 연산자)가 부분 결과를 결합합니다. 출력 임베딩의 경우, 순진하게 logit을 수집하면 $b \times s \times v$ 요소를 통신해야 합니다. 대신, 저자들은 cross-entropy 손실을 병렬 logit 계산과 융합합니다. 각 GPU는 자신의 어휘 하위 집합에 대한 부분 logit을 계산하고, $b \times s$ 차원의 스칼라 손실만 통신하면 되므로 통신량이 극적으로 감소합니다.

하이브리드 모델 및 데이터 병렬화. 이 접근법은 데이터 병렬화와 결합됩니다. GPU는 모델 병렬 그룹(예: NVSwitch를 통한 단일 노드 내 8개 GPU)과 데이터 병렬 그룹(InfiniBand로 연결된 서로 다른 모델 병렬 그룹의 동일한 위치에 있는 GPU)으로 구성됩니다. 역전파 중에 그래디언트 all-reduce는 데이터 병렬 그룹 내에서 발생하고, 모델 병렬 통신은 모델 병렬 그룹 내에서 발생합니다. 총 GPU 수는 모델 병렬 차수와 데이터 병렬 차수의 곱입니다. 83억 파라미터 모델의 경우 $8 \times 64 = 512$개의 GPU입니다.

난수 생성 처리. Dropout 레이어는 모델 병렬 영역 내부와 외부 모두에 존재합니다. 모델 병렬 영역 외부의 dropout(residual connection 이전)의 경우, 모든 모델 병렬 워커는 동일한 dropout 패턴을 보장하기 위해 동일한 난수 시드를 사용합니다. 모델 병렬 영역 내부의 dropout(self-attention 내)의 경우, 각 워커는 독립적인 계산을 올바르게 정규화하는 서로 다른 난수 패턴을 보장하기 위해 모델 병렬 순위에 기반한 고유 시드를 사용합니다.

2.4 구현 세부사항

모델 구성. 논문은 355M, 25억, 83억 파라미터의 GPT-2 모델과 336M, 13억, 39억 파라미터의 BERT 모델을 탐색합니다. 스케일링 연구를 위해 attention head당 hidden size를 96으로 일정하게 유지하면서 head와 레이어를 변화시킵니다. GPT-2 학습의 경우: 시퀀스 길이 1024, 배치 크기 512, 300k iteration. BERT의 경우: 배치 크기 1024, 학습률 $1.0 \times 10^{-4}$, 2M iteration.

가중치 초기화. 가중치는 $\mathcal{N}(0, 0.02)$에서 초기화되며, residual layer 직전의 가중치에는 추가 스케일링 팩터 $\frac{1}{\sqrt{2N}}$이 적용됩니다. 여기서 $N$은 transformer 레이어 수입니다. 이 스케일링은 residual stream의 분산이 깊이에 따라 증가하는 것을 방지합니다.

최적화. Adam optimizer에 weight decay $\lambda = 0.01$, 전역 그래디언트 norm 클리핑 1.0, dropout 0.1, 그리고 V100 Tensor Core를 위한 동적 손실 스케일링을 사용한 혼합 정밀도 학습을 사용합니다. GPT-2는 학습률 $1.5 \times 10^{-4}$에 3k warmup iteration과 $1 \times 10^{-5}$로의 cosine decay를 사용합니다.

어휘 패딩. 어휘 크기는 50,257에서 51,200으로 패딩됩니다(최대 8개 GPU에 걸쳐 어휘가 분할될 때 효율적인 GEMM을 보장하기 위해 $128 \times 8 = 1024$로 나누어떨어짐).

하드웨어. 최대 32개의 DGX-2H 서버(512개의 V100 SXM3 32GB GPU)를 사용하며, NVSwitch를 통한 300GB/s 노드 내 대역폭과 서버당 8개의 InfiniBand 어댑터를 통한 100GB/s 노드 간 대역폭을 갖습니다.

메모리 관리. Activation checkpointing은 메모리 풋프린트를 줄이기 위해 모든 transformer 레이어 이후에 적용되어, 더 큰 모델이 GPU 메모리 내에 맞도록 합니다.

복잡도 분석. 각 transformer 레이어는 정확히 4개의 all-reduce 연산(forward 2개, backward 2개)을 도입합니다. 각 all-reduce는 $O(b \times s \times H)$ 데이터를 통신합니다. 계산 대 통신 비율은 GEMM 계산이 $O(b \times s \times H^2)$이고 통신이 $O(b \times s \times H)$이므로 유리하게 스케일되며, 이는 hidden dimension이 증가함에 따라 통신이 더 작은 비율이 됨을 의미합니다. 이 접근법은 39 TFLOPS 단일 GPU 기준선(이론적 피크의 30%)과 512 GPU에서 15.1 PFLOPS를 달성합니다.

3. Results

Megatron-LM의 실험적 평가는 세 가지 차원에 걸쳐 있습니다: 스케일링 효율, GPT-2 언어 모델링 품질, 그리고 BERT 다운스트림 작업 성능. 각 차원은 논문의 주장에 대한 상호 보완적인 증거를 제공합니다.

스케일링 효율. 스케일링 연구는 모델 크기와 GPU 수가 함께 증가하는 weak scaling을 사용합니다. 39 TFLOPS(이론적 피크의 30%)를 달성하는 12억 파라미터의 강력한 단일 GPU 기준선에서 시작하여, 논문은 모델 병렬화를 위해 2, 4, 8 GPU에서, 그리고 결합된 모델 및 데이터 병렬화를 위해 64~512 GPU에서 효율을 측정합니다.

8 GPU에서 77%의 모델 병렬 효율은 주목할 만하며, 이는 NVSwitch 통신을 사용한 레이어 내 분할의 오버헤드를 나타냅니다. 64-way 데이터 병렬화를 추가할 때 74%로의 추가 감소는 상대적으로 미미하며, 크로스 노드 그래디언트 all-reduce가 제한된 오버헤드를 도입함을 확인합니다. Strong scaling 실험(고정된 12억 모델, GPU 증가)은 2 GPU 이후 수익 감소를 보이며, 8 GPU에서 2.98배의 속도 향상만 보입니다. 이는 모델이 단일 GPU에 맞을 만큼 충분히 작고 GPU당 계산이 감소하므로 예상되는 결과입니다.

Attention head가 스케일링 효율에 미치는 영향은 흥미로운 ablation입니다. 83억 모델에 대한 8-way 모델 병렬화에서, attention head를 16에서 32로 증가시키면(head dimension을 192에서 96으로 조정하여 총 파라미터를 일정하게 유지) 스케일링 효율이 82%에서 77%로 감소합니다. 이는 더 많은 attention head가 head당 더 작은 GEMM을 생성하고 softmax 계산을 증가시켜, 계산이 크고 효율적인 행렬 곱셈에서 통신 대 계산 비율이 더 높은 작은 연산으로 이동하기 때문입니다.

GPT-2 언어 모델링. GPT-2 결과는 규모에 따른 명확하고 일관된 개선 패턴을 보여줍니다.

83억 모델은 WikiText103 perplexity 10.81을 달성하여 이전 최고 성능인 15.79에 비해 31.6%의 상대적 개선을 보입니다. LAMBADA에서 66.51%의 정확도는 이전 최고인 63.24%에 비해 3.27 퍼센트 포인트의 개선을 나타냅니다. 검증 perplexity 곡선은 더 큰 모델이 더 낮은 perplexity로 수렴할 뿐만 아니라 iteration 수 측면에서 더 빠르게 수렴함을 보여주며, 증가된 모델 용량이 학습 데이터의 더 효율적인 사용으로 이어짐을 시사합니다. 논문은 테스트 세트 오염이 최소한임을 신중하게 검증하며, WikiText103 테스트 세트가 학습 데이터와 최대 10.8%의 8-gram 중복을 보이며 이는 이전 연구와 일치합니다.

BERT 다운스트림 작업. BERT 실험은 이중 목적을 제공합니다: 양방향 모델에 대한 모델 병렬화 접근법을 검증하고 pre-norm 아키텍처의 중요성을 보여줍니다.

모든 다운스트림 작업은 336M에서 13억, 39억까지 모델 크기에 따라 단조롭게 개선됩니다. Megatron-39억 모델은 모든 작업에서 최고 성능을 달성하며, 특히 RACE에서 주목할 만한 이득을 보입니다(단일 모델 기준 이전 최고인 86.5%에 비해 89.5%). 5-way 앙상블은 RACE를 90.9%로 끌어올립니다. 중요한 점은 13억 구성이 BERT-xlarge와 동일하며, Lan et al. (2019)이 BERT-large보다 나쁘다고 보고한 것입니다. Pre-norm 아키텍처로 Megatron-13억은 모든 작업에서 BERT-large를 능가하여, 원래의 성능 저하가 모델 용량이 아닌 아키텍처 때문임을 확인합니다. 원래 BERT 아키텍처와 재배치된 BERT 아키텍처 간의 학습 손실 비교는 752M 원래 아키텍처가 불안정성을 겪는 반면 재배치된 버전은 더 낮은 손실로 안정적으로 학습됨을 시각적으로 보여줍니다.

이러한 결과의 주목할 만한 측면은 효율성입니다: Megatron 모델은 336M 기준선의 $1\times$ 학습 토큰만 사용하는 반면, RoBERTa는 $2\times$, ALBERT는 $3\times$를 사용합니다. 이는 더 큰 모델이 데이터 효율적이며, 동일한 양의 학습 데이터에서 더 많은 정보를 추출함을 시사합니다.

4. Critical Assessment

Strengths

1. 이 접근법은 기존 PyTorch 코드에 커스텀 autograd 함수($f$와 $g$ 연산자)를 통해 몇 개의 all-reduce 연산만 삽입하면 되므로 구현이 매우 간단하며, 컴파일러나 프레임워크 의존성이 없습니다.
2. 강력한 단일 GPU 기준선(피크 FLOPS의 30%)에 대해 512 GPU에서 76%의 스케일링 효율은 접근법이 최소한의 통신 오버헤드를 도입함을 보여주며, 실제 대규모 학습에 실용적입니다.
3. Pre-norm layer normalization이 BERT 계열 모델을 확장하는 데 중요하다는 발견은 현장에 지속적인 영향을 미쳤으며, 후속 대규모 transformer 학습의 사실상 표준이 되었습니다.
4. 논문은 단순히 FLOPS 효율을 측정하는 것이 아니라 여러 벤치마크에서 완전한 모델을 최고 성능으로 학습하여 포괄적인 end-to-end 검증을 제공하며, 시스템 기여가 실제 NLP 발전으로 이어짐을 보여줍니다.
5. Megatron-LM 코드베이스의 오픈소스 공개는 Turing-NLG(17B 파라미터)를 포함한 수많은 후속 프로젝트에서 사용되는 기초 도구가 되어 대규모 모델 학습 연구를 촉진했습니다.
6. 손실 융합을 사용한 임베딩 병렬화의 원리적 분석은 출력 통신을 $O(b \times s \times v)$에서 $O(b \times s)$로 줄이며, 이는 큰 어휘를 가진 모델에 중요한 최적화입니다.

Limitations

1. 이 접근법은 transformer 아키텍처에 특화되어 있으며 Mesh-TensorFlow나 FlexFlow 같은 더 일반적인 프레임워크와 달리 임의의 신경망 구조로 일반화되지 않습니다.
2. 모델 병렬화는 attention head 수에 의해 제한되며, 이는 head가 적은 아키텍처에 대한 최대 병렬화 차수를 제약하고, Table 7에서 보여지듯이 더 많은 head로 스케일링 효율이 저하됩니다.
3. 논문은 단일 노드 내(최대 8 GPU)에서만 레이어 내 병렬화를 보여주며, 노드 간 스케일링에는 데이터 병렬화에 의존합니다. 8-way 이상의 모델 병렬화가 필요한 모델은 파이프라인 병렬화가 필요하지만, 논문은 이를 인정하지만 다루지 않습니다.
4. 스케일링 분석은 weak scaling만 사용하여 모델 병렬화 효율의 효과와 더 큰 모델을 학습하는 이점을 분리하기 어렵습니다.
5. Strong scaling 결과(Table 8)는 단일 GPU에 맞는 모델의 학습을 가속화하기 위해 모델 병렬화를 사용하면 수익 감소를 보이며(8 GPU로 2.98배 속도 향상만), 진정으로 여러 GPU가 필요한 모델에 대한 유용성을 제한합니다.
6. BERT 실험은 서로 다른 모델 크기에 대해 서로 다른 학습 iteration 수를 사용하며(336M/13억의 경우 2M vs 39억의 경우 1.5M, 39억이 여전히 학습 중), 가장 큰 모델이 완전히 수렴했는지 불분명합니다.
7. 다른 모델 병렬화 프레임워크(GPipe, Mesh-TensorFlow)에 대한 스케일링 효율이나 구현 용이성 측면의 비교가 제공되지 않아, 제안된 접근법의 장점을 정량화하기 어렵습니다.

Future Directions

1. 레이어 내 모델 병렬화와 파이프라인 병렬화(레이어 간)를 결합하여 단일 DGX 노드의 메모리를 초과하는 160억 파라미터 이상의 모델 학습을 가능하게 하며, 이는 후속적으로 Megatron-LM v2에서 실현되었습니다.
2. GPT-2와 BERT 이외의 다른 모델 패밀리로 접근법을 확장하며, T5 같은 encoder-decoder 아키텍처와 XLNet 같은 autoregressive 모델을 포함합니다.
3. 모델 병렬화와 large batch 학습 간의 상호작용을 조사하며, 특히 모델 병렬화가 최적화 품질을 저하시키는 배치 크기 요구사항을 줄일 수 있는지 확인합니다.
4. 이러한 대규모 사전학습된 교사 모델에서 더 작은 학생 모델로 지식 증류를 적용하여, 비용의 일부로 성능의 많은 부분을 회복할 가능성이 있습니다.
5. 파라미터당 메모리 오버헤드를 줄이는 메모리 효율적인 optimizer를 개발하여 동일한 GPU 메모리 제약 내에서 더 큰 모델을 가능하게 합니다.
6. 모델 크기, 학습 데이터, 계산 예산 간의 관계가 다운스트림 작업 성능에 어떻게 영향을 미치는지 연구하기 위해 스케일링 분석을 확장하며, 이는 후속 스케일링 법칙 연구를 예고합니다.