Tensor Parallel

TL;DR

메가트론-LM은 여러 GPU에 걸쳐 개별 모델 레이어를 분할하여 거대한 트랜스포머 언어 모델을 훈련시키는 우아하고 효율적인 접근 방식을 소개합니다. 이 “텐서 모델 병렬화” 기술을 사용하여 연구자들은 83억 개의 파라미터를 가진 모델(GPT-2보다 5.6배 큰)을 높은 계산 효율성(512 GPU에서 76% 스케일링 효율성)으로 성공적으로 훈련시켰습니다. 이 접근 방식은 기존 PyTorch 코드에 최소한의 변경만 필요하며, 다른 병렬화 기술과 결합될 수 있고, 언어 모델링 성능을 향상시킵니다. 연구자들은 또한 더 큰 모델로 확장할 때 레이어 정규화 배치가 중요해진다는 것을 발견했습니다.

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

모델 병렬화 기법:

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

분산 훈련 시스템:

DeepSpeed Ulysses - 시퀀스 병렬화와 텐서 병렬화 결합
USP - 통합 시퀀스 병렬화 프레임워크
LoongTrain - 하이브리드 병렬화 접근법

Takeaways

1. 동기 및 배경

문제: 대형 모델에 대한 메모리 제한

이 연구의 주요 동기는 간단했습니다: 연구자들은 더 큰 트랜스포머 모델이 자연어 작업에서 성능을 꾸준히 향상시키지만, GPU 메모리 제한에 부딪히고 있다는 것을 관찰했습니다.

메가트론-LM 이전에, BERT(3억 4천만 파라미터)와 GPT-2(15억 파라미터)와 같은 언어 모델은 모델 크기를 확장하면 더 나은 결과를 얻을 수 있음을 보여주었습니다. 그러나 추가 확장은 개별 GPU의 메모리 용량에 의해 제한되었습니다. 32GB 메모리를 가진 고성능 NVIDIA V100을 사용하더라도, 연구자들은 훈련 중에 약 12억 파라미터의 모델만 맞출 수 있었습니다(메모리 최적화 기술을 사용해도).

근본적인 병목 현상은 분명했습니다: 더 큰 모델을 훈련시키기 위해서는 데이터뿐만 아니라 모델 자체를 여러 GPU에 분할하는 방법이 필요했습니다.

기존 접근 방식과 그 한계

이 연구 이전에는 분산 훈련에 두 가지 주요 접근 방식이 있었습니다:

데이터 병렬화: 각 GPU가 모델의 완전한 복사본을 가지고, 배치를 GPU에 걸쳐 분할합니다. 이것은 대형 모델의 메모리 제한에 도움이 되지 않습니다.
파이프라인 병렬화(예: GPipe): 서로 다른 레이어를 서로 다른 GPU에 할당하여 모델을 분할합니다. 효과적이지만, 이 접근 방식은 효율성을 떨어뜨리는 파이프라인 버블(GPU 유휴 시간)을 도입합니다.

부족했던 것은 레이어 간이 아닌 각 레이어 내에서 세밀한 병렬화를 가능하게 하는, 트랜스포머 모델의 독특한 계산 구조를 특별히 다루는 효율적인 접근 방식이었습니다.

2. 핵심 가정 및 성공 조건

메가트론-LM의 접근 방식이 효과적으로 작동하기 위해서는 몇 가지 핵심 가정과 조건이 필요했습니다:

고대역폭 GPU 인터커넥트: 이 접근 방식은 GPU 간의 빠른 통신을 가정합니다. 연구자들은 NVLink(GPU 간 300+ GB/s 대역폭)를 갖춘 NVIDIA의 DGX SuperPOD를 사용했는데, 이는 표준 PCIe 연결(~32 GB/s)보다 훨씬 빠릅니다.
은닉 차원의 균등 분할 가능성: 모델의 은닉 차원은 텐서 병렬화에 사용되는 GPU 수로 균등하게 나누어져야 합니다. 이는 가중치 행렬이 나머지 없이 분할될 수 있도록 합니다.
병렬화의 이점을 얻기에 충분히 큰 행렬: 이 접근 방식은 행렬 연산을 분할하는 계산상의 이점이 통신 오버헤드보다 크다고 가정합니다. 이는 충분히 큰 모델에서만 유효합니다.
NCCL 통신 라이브러리: 구현은 GPU 간의 효율적인 집합 통신 작업을 위해 NVIDIA의 NCCL 라이브러리에 의존합니다.
혼합 정밀도 훈련: 메모리 효율성과 계산 처리량을 최대화하기 위해, 혼합 정밀도 훈련(FP16/FP32)이 실질적으로 필요합니다.
대형 모델의 안정적인 훈련: 모델이 커질수록 수치적 안정성이 더 어려워집니다. 연구자들은 표준 트랜스포머 아키텍처에 대한 수정, 특히 레이어 정규화 배치가 필요하다는 것을 발견했습니다.

3. 메가트론-LM 방법 설명

메가트론-LM의 핵심 통찰은 트랜스포머 모델 중 최소한의 통신 오버헤드로 GPU에 걸쳐 효율적으로 병렬화할 수 있는 부분을 식별하는 것입니다.

3.1 핵심 개념: 행과 열 분할을 통한 텐서 병렬화

주요 혁신은 트랜스포머 레이어 내의 큰 가중치 행렬을 여러 GPU에 걸쳐 분할하는 것입니다. 두 가지 주요 유형의 행렬 분할이 있습니다:

열-병렬 선형 레이어:

각 GPU는 동일한 입력을 받지만 출력 뉴런의 일부만 담당합니다. 이는 가중치 행렬을 열 방향으로 분할하는 것과 같습니다.

class ColumnParallelLinear(torch.nn.Module):
    def __init__(self, input_size, output_size, bias=True, gather_output=True):
        super(ColumnParallelLinear, self).__init__()
        # 모델 병렬 그룹에서 월드 사이즈와 랭크 가져오기
        world_size = get_model_parallel_world_size()
        rank = get_model_parallel_rank()

        # GPU에 걸쳐 출력 크기 분할
        self.output_size_per_partition = output_size // world_size

        # 로컬 선형 레이어 생성 (출력 뉴런의 일부만)
        self.local_linear = torch.nn.Linear(
            input_size, self.output_size_per_partition, bias=bias)

        # GPU에 걸쳐 출력을 수집할지 여부를 나타내는 플래그
        self.gather_output = gather_output

    def forward(self, input_):
        # 로컬 순방향 계산
        local_output = self.local_linear(input_)

        # 필요한 경우, 모든 GPU에서 출력 수집
        if self.gather_output:
            # 모든 출력을 수집하는 all-gather 연산
            output = all_gather_from_model_parallel(local_output)
            return output
        else:
            return local_output

행-병렬 선형 레이어:

각 GPU는 입력의 일부만 받아 전체 출력의 부분적인 결과를 계산합니다. 이는 가중치 행렬을 행 방향으로 분할하는 것과 같습니다.

class RowParallelLinear(torch.nn.Module):
    def __init__(self, input_size, output_size, bias=True, input_is_parallel=False):
        super(RowParallelLinear, self).__init__()
        # 모델 병렬 그룹에서 월드 사이즈와 랭크 가져오기
        world_size = get_model_parallel_world_size()
        rank = get_model_parallel_rank()

        # GPU에 걸쳐 입력 크기 분할
        self.input_size_per_partition = input_size // world_size

        # 로컬 선형 레이어 생성 (입력 뉴런의 일부만)
        self.local_linear = torch.nn.Linear(
            self.input_size_per_partition, output_size, bias=False)

        # 입력이 이미 분할되어 있는지 나타내는 플래그
        self.input_is_parallel = input_is_parallel

        # 바이어스는 all-reduce 후에 적용되므로 한 개만 필요
        self.bias = torch.nn.Parameter(torch.empty(output_size)) if bias else None

    def forward(self, input_):
        # 입력이 아직 분할되지 않았다면, 분산시킴
        if not self.input_is_parallel:
            input_parallel = scatter_to_model_parallel(input_)
        else:
            input_parallel = input_

        # 로컬 순방향 계산
        local_output = self.local_linear(input_parallel)

        # 완전한 출력을 얻기 위한 all-reduce
        output = all_reduce_from_model_parallel(local_output)

        # all-reduce 후 바이어스 적용
        if self.bias is not None:
            output = output + self.bias

        return output

3.2 실제 예시

1024 은닉 차원과 2-way 텐서 병렬화를 사용한 트랜스포머 레이어의 구체적인 예를 들어보겠습니다:

셀프-어텐션 블록 예시:

병렬화 없이, 셀프-어텐션 계산은 다음과 같을 것입니다:

QKV 투영: 입력 [batch_size, seq_len, 1024] → 크기 [1024, 1024]의 3개 행렬 → 출력 [batch_size, seq_len, 3072]
어텐션 계산 및 출력 투영: [batch_size, seq_len, 3072] → [batch_size, seq_len, 1024]

2-way 텐서 병렬화를 사용하면:

GPU 0:

QKV 투영: 입력 [batch_size, seq_len, 1024] → 크기 [1024, 512]의 3개 행렬 → 출력 [batch_size, seq_len, 1536]
헤드의 절반에 대한 어텐션 계산
출력 투영: 입력 [batch_size, seq_len, 512] → 행렬 [512, 1024] → 부분 출력 [batch_size, seq_len, 1024]

GPU 1:

QKV 투영: 입력 [batch_size, seq_len, 1024] → 크기 [1024, 512]의 3개 행렬 → 출력 [batch_size, seq_len, 1536]
나머지 헤드의 어텐션 계산
출력 투영: 입력 [batch_size, seq_len, 512] → 행렬 [512, 1024] → 부분 출력 [batch_size, seq_len, 1024]

통신:

출력 투영 결과에 대한 all-reduce로 완전한 결과 획득

설계상, 각 GPU는 가중치에 대해 절반의 메모리를 사용하고 절반의 연산을 계산하지만, 셀프-어텐션 블록 끝에서 단 한 번의 통신만 필요합니다.

4. 실험 결과 및 분석

4.1 주요 성능 결과

모델 크기	GPU	텐서 병렬화	배치 크기	처리량	스케일링 효율성	지속 성능
1.2B	1	1	8	39 TFLOPS	100% (기준선)	피크의 30%
2.5B	2	2	8	73 TFLOPS	94%	-
4.2B	4	4	8	138 TFLOPS	88%	-
8.3B	8	8	8	254 TFLOPS	81%	-
8.3B	512	8	512	15.1 PFLOPS	76%	-

이 결과는 여러 가지 이유로 인상적입니다. 첫째, 단일 GPU에서 39 TFLOPS의 기준 성능은 NVIDIA V100 GPU의 이론적 피크의 30%로, 복잡한 딥러닝 애플리케이션에서는 꽤 좋은 수치입니다. 둘째, 텐서 병렬화가 증가해도 스케일링 효율성이 높게 유지되어 8-way 병렬화에서 81% 효율성을 보입니다.

가장 인상적인 결과는 512 GPU를 사용하여 달성한 15.1 PFLOPS로, 8-way 텐서 병렬화와 64-way 데이터 병렬화를 결합했습니다. 이는 단일 GPU 기준선 대비 76%의 스케일링 효율성을 나타내는데, 이 규모의 분산 시스템에서는 놀라운 수치입니다.

4.2 언어 모델링 성능

모델	파라미터	언어 모델링 퍼플렉시티	LAMBADA 정확도
GPT-2	1.5B	~35.7	45.9%
메가트론-LM	1.5B	~35.0	50.6%
메가트론-LM	8.3B	~29.8	63.2%

이 결과는 더 큰 모델이 더 나은 언어 모델링 능력을 제공한다는 가설을 검증합니다. 83억 파라미터 모델은 15억 모델보다 퍼플렉시티(낮을수록 좋음)와 LAMBADA 정확도(높을수록 좋음) 모두에서 크게 성능이 향상되었습니다.

특히 흥미로운 점은 크기가 동일한 15억 메가트론-LM 모델도 GPT-2보다 약간 더 나은 성능을 보였다는 것인데, 이는 그들의 훈련 접근 방식이나 아키텍처 수정이 규모와 상관없이 유익했을 수 있음을 시사합니다.

4.3 애블레이션 연구: 레이어 정규화 배치

레이어 정규화 위치	1B 모델 퍼플렉시티	4B 모델 퍼플렉시티
Pre-LN (원래)	3.77	발산
Post-LN (수정됨)	3.69	3.31

이 애블레이션 연구는 중요한 점을 밝혀냈습니다: 트랜스포머 모델이 커질수록, 작은 모델에서 잘 작동했던 아키텍처 선택이 더 이상 실행 가능하지 않을 수 있습니다. 구체적으로, 표준 레이어 정규화 배치(Pre-LN, 어텐션/FFN 블록 이전에 레이어 정규화 적용)는 더 큰 모델에서 훈련 발산을 일으켰습니다.

레이어 정규화를 어텐션/FFN 블록 이후로 이동시킴으로써(Post-LN), 그들은 안정적인 훈련과 더 나은 성능을 달성했습니다. 이 발견은 꽤 중요했으며 이후 대형 모델 설계에 영향을 미쳤습니다.

5. 영향 및 유산

메가트론-LM의 영향은 상당했습니다:

더 큰 모델 가능: 이 접근 방식은 이전에 가능했던 것보다 훨씬 더 큰 모델을 훈련시킬 수 있게 하여, 수천억 개의 파라미터를 가진 GPT-3, Gopher, PaLM 등의 후속 모델을 위한 기반을 마련했습니다.
아키텍처 혁신: 레이어 정규화 배치 및 기타 훈련 안정화 기술에 대한 그들의 발견은 후속 모델 설계에 영향을 미쳤습니다.
대규모 훈련의 민주화: PyTorch 코드에 최소한의 변경으로 접근 방식을 구현함으로써, 특수한 컴파일러나 하드웨어 없이도 연구자들이 대규모 훈련에 접근할 수 있게 했습니다.
다른 접근 방식과의 보완성: 아마도 가장 중요한 것은, 그들의 텐서 병렬화 접근 방식이 파이프라인 병렬화 및 데이터 병렬화와 보완적임이 증명되어, 가장 큰 모델 훈련에서 표준이 된 3D 병렬화를 가능하게 했다는 점입니다.

메가트론-LM 연구는 언어 모델 확장의 중요한 순간을 나타내며, 적절한 병렬화 전략을 통해 수십억 개의 파라미터를 가진 모델을 훈련시키는 것이 가능할 뿐만 아니라 높은 계산 효율성으로 수행될 수 있음을 보여주었습니다. 이는 최근 언어 모델의 발전을 정의한 모델 규모의 폭발적인 증가를 위한 기반을 직접적으로 마련했습니다.

결론적으로, 메가트론-LM의 모델 병렬화 접근 방식은 오늘날의 대형 언어 모델을 향한 중요한 디딤돌이었습니다. 회고적으로 보면 비교적 간단한 기술처럼 보일 수 있지만, 기존 코드에 최소한의 변경만으로 중요한 스케일링 병목 현상을 우아하게 해결했습니다. 이 효과성과 접근성의 조합은 대규모 언어 모델링 분야의 발전을 가속화하는 데 특히 영향력이 컸습니다.