How Parallelism Works — 모델 하나를 GPU 만 대에 나눠 싣는 법

1. 왜 잘라야 하나: 405B 모델의 학습 메모리

학습할 때 GPU 메모리에는 가중치만 올라가지 않습니다. 요즘 표준인 혼합 정밀도 + Adam 학습에서는 파라미터 하나에 다음이 붙습니다.

405B × 16바이트 ≈ 6.5TB입니다. H100 한 대의 메모리는 80GB이니, 활성(순전파 중간 결과)을 빼고 상태만 담아도 GPU가 최소 81대 필요합니다. 이 16바이트 분해는 ZeRO 논문이 쓴 것과 같고(옵티마이저 몫을 K = 12로 둡니다), 아래 데모 1에서 그대로 씁니다. 이렇게 "한 대에 안 들어간다"는 사실이 모든 병렬화의 출발점입니다. 다만 어디를 자르느냐에 따라 대가가 달라집니다.

2. 네 논문, 네 가지 자르기

초거대 모델 학습의 계보를 정리할 때 함께 나오는 네 논문은 "모델 병렬화"로 한데 묶이곤 하지만, 자르는 대상이 모두 다릅니다. "순수 파이프라인 병렬화는 마지못한 최후 수단"이라는 말에서 파이프라인은 이 중 GPipe 계열 하나만 가리킵니다. 나머지 셋은 파이프라인이 아니라, 오히려 파이프라인 대신 무거운 일을 맡게 된 쪽입니다.

Megatron-LM은 2021년 후속 논문(Narayanan et al., SC21)에서 파이프라인(interleaved 1F1B)을 더해 TP × PP × DP의 3D 병렬화로 확장했습니다. GPipe 자료에 나오는 "Megatron 스케줄"이 이것입니다. 그래도 Megatron-LM의 대표 기여는 텐서 병렬로 봅니다.

3. 데이터 병렬과 ZeRO / FSDP DP

가장 단순한 방법은 모델은 통째로 복제하고 배치만 나누는 것입니다. GPU N대가 각자 다른 데이터로 기울기를 구한 뒤 all-reduce로 평균을 내고, 모두 같은 갱신을 합니다. 통신은 스텝마다 한 번이고, 다음 층의 역전파와 겹쳐서 숨기기도 쉽습니다. 문제는 메모리입니다. N대 모두가 16바이트 × 파라미터 수를 똑같이 들고 있으니, 모델이 한 대에 안 들어가면 데이터 병렬만으로는 답이 없습니다.

ZeRO는 이 중복을 겨냥합니다. 어차피 N대가 같은 값을 들고 있으니, 1/N씩 나눠 맡고 필요할 때만 모으자는 겁니다. 나누는 대상을 단계적으로 넓혀 세 단계가 됩니다.

Ψ = 405B, N = 64로 놓고 ZeRO-3를 골라 보세요. 장치당 약 101GB로, 활성을 넣기도 전에 80GB를 넘습니다. N을 512까지 늘리면 상태는 13GB 정도로 줄지만, 이번에는 매 층의 순전파·역전파마다 파라미터를 512대에서 all-gather해야 하고, 층 하나의 활성 자체는 전혀 줄지 않습니다. 데이터 병렬은 배치를 자를 뿐 층을 자르지 않기 때문입니다. 여기서 텐서 병렬이 필요해집니다.

통신량: 보통 데이터 병렬의 all-reduce는 reduce-scatter + all-gather 두 단계라 GPU당 약 2Ψ를 주고받습니다. ZeRO-1·2는 같은 2Ψ로 끝나고, ZeRO-3는 파라미터를 순전파·역전파에서 한 번씩 더 모아 약 1.5배(3Ψ)가 됩니다. PyTorch의 FSDP(Zhao et al., 2023)가 ZeRO-3를 구현한 것이고, 옵션으로 ZeRO-2처럼 동작하게 할 수도 있습니다.

4. 텐서 병렬 (Megatron-LM) TP

텐서 병렬은 층 하나의 행렬곱 자체를 여러 GPU로 쪼갭니다. Megatron-LM의 요령은 Transformer의 MLP와 어텐션을 통신이 최소가 되게 자르는 방법을 찾은 것입니다. MLP Z = GeLU(X·A)·B에서 첫 행렬 A는 열 방향으로 자릅니다. 그러면 GPU마다 \(Y_i = \mathrm{GeLU}(X A_i)\)를 혼자 계산할 수 있습니다(GeLU는 원소별이라 열을 섞지 않습니다). 둘째 행렬 B는 행 방향으로 잘라 \(Z_i = Y_i B_i\)를 구하고, 마지막에 \(Z = \sum_i Z_i\)를 all-reduce로 한 번 더합니다.

어텐션은 헤드 단위로 나눕니다. Q·K·V 투영을 헤드별로 열 방향으로 자르면 각 GPU가 자기 헤드의 어텐션을 독립적으로 계산하고, 출력 투영을 행 방향으로 잘라 다시 all-reduce 한 번으로 합칩니다. 그래서 Transformer 층 하나에 순전파 2번, 역전파 2번의 all-reduce가 생깁니다. Megatron-LM은 이 방법으로 8.3B GPT-2 모델을 512개 GPU에서 학습했습니다.

ZeRO와 다른 점은 무엇을 주고받느냐입니다. ZeRO는 스텝마다 파라미터와 기울기를 주고받고, 계산과 겹치기 쉽습니다. 텐서 병렬은 층마다 활성을 주고받고, 그 결과가 나와야 다음 계산을 할 수 있어서 통신이 계산 경로 한가운데 놓입니다. 그래서 대역폭이 모자라면 바로 느려집니다.

NVLink에서 TP 8은 통신이 전체의 10% 남짓입니다. 같은 TP 8을 서버 사이 400Gb/s(GPU당 50GB/s) 연결로 묶으면 통신이 계산보다 길어집니다. TP 16은 서버 하나(GPU 8대)를 넘어가므로 NVLink를 고를 수 없습니다. 이것이 텐서 병렬은 보통 한 노드(NVLink로 묶인 GPU 8대) 안에서만 쓴다는 규칙의 이유입니다. 모델을 더 잘라야 하면 다른 축을 써야 합니다.

5. 레이아웃으로 일반화하기 (Mesh-TensorFlow)

Mesh-TensorFlow는 새로운 병렬화를 하나 더 만든 게 아니라, 병렬화를 기술하는 언어를 만들었습니다. 장치들을 2차원(또는 n차원) 격자, 즉 메시로 보고, 텐서의 각 차원에 이름을 붙인 뒤 "batch 차원은 메시의 행 축에, hidden 차원은 열 축에 나눈다" 같은 레이아웃만 적으면 필요한 통신은 시스템이 알아서 넣습니다. batch를 나누면 데이터 병렬, hidden을 나누면 텐서 병렬, 둘 다 나누면 둘을 섞은 것이 됩니다. 논문은 이 방법으로 최대 5B 파라미터 Transformer를 TPU 512코어에서 학습했습니다.

이 생각은 이후 Google의 GSPMD, JAX의 PartitionSpec, PyTorch의 DeviceMesh/DTensor로 이어졌습니다. 9절에서 볼 Llama 3의 4D 병렬화도 결국 [TP, CP, PP, DP] 네 축짜리 장치 메시에 무엇을 어느 축으로 나눌지 정한 레이아웃입니다.

6. 파이프라인 병렬 (GPipe 계열): 왜 "최후 수단"인가 PP

파이프라인 병렬은 모델을 층 묶음(스테이지) 단위로 잘라 GPU마다 하나씩 싣고, 미니배치를 마이크로배치로 쪼개 흘려보냅니다. 스테이지 사이에는 경계의 활성만 점대점으로 넘기면 되니 통신량이 적고, 그래서 느린 노드 사이 연결에도 둘 수 있다는 게 장점입니다. 원리와 타임라인은 How GPipe Works에서 직접 돌려 볼 수 있습니다.

그런데 파이프라인만으로 스케일하면 대가가 커집니다.

• 버블: 파이프라인을 채우고 비우는 동안 일부 GPU가 놉니다. 스테이지 p개, 마이크로배치 m개면 유휴 비율이 \((p-1)/(m+p-1)\)입니다.
• 배치 크기 제약: 버블을 줄이려면 m을 p보다 훨씬 크게 잡아야 하는데, 그러면 전역 배치가 커집니다. 배치는 학습 안정성 때문에 마음대로 늘릴 수 없습니다.
• 불균형: 가장 느린 스테이지가 전체 속도를 정합니다. 임베딩·출력층이 붙은 첫·마지막 스테이지는 무겁고, 메모리도 앞 스테이지일수록 마이크로배치 활성을 오래 들고 있습니다.
• 구현 복잡도: 스케줄(1F1B, interleaved), 스테이지 간 비동기 통신, 메모리 해제 시점을 직접 관리해야 합니다.

그래서 지금의 정리는 이렇습니다. 데이터 병렬(ZeRO/FSDP) + 텐서 병렬(Megatron식)이 주력이고, 그래도 모델이 장치에 안 들어갈 때만 파이프라인을 노드 사이에 얇게 섞습니다. "순수" 파이프라인이라는 말은 파이프라인 하나만으로 GPU 수를 늘리는 방식을 가리킵니다.

7. 컨텍스트 병렬: 긴 시퀀스 자르기 CP

문맥이 128K 토큰쯤 되면 모델은 들어가도 시퀀스 하나의 활성이 GPU 하나에 안 들어갑니다. 컨텍스트 병렬(시퀀스 병렬이라고도 합니다)은 시퀀스 길이 방향으로 토큰을 잘라 GPU마다 나눕니다. MLP나 LayerNorm은 토큰마다 독립이라 그냥 나눠도 되지만, 어텐션은 각 쿼리가 앞쪽 모든 토큰의 K·V를 봐야 해서 통신이 필요합니다. 방법은 둘입니다.

• Ring Attention (Liu et al., 2023): K·V 조각을 GPU 고리를 따라 한 칸씩 돌리며, 받는 동안 계산을 겹칩니다.
• All-gather: K·V를 한 번에 모두 모은 뒤 자기 쿼리 조각만 계산합니다. Llama 3가 고른 방식입니다. GQA 때문에 K·V가 Q보다 훨씬 작고(405B는 쿼리 헤드 128개에 K·V 헤드 8개), 문서 경계 마스크 같은 다양한 마스크를 다루기 쉽다는 이유였습니다.

또 하나의 문제는 인과 마스크입니다. 앞 토큰은 볼 것이 적고 뒤 토큰은 볼 것이 많아서, 시퀀스를 순서대로 자르면 마지막 GPU가 일을 몰아서 합니다. Llama 3는 시퀀스를 2 × CP개 조각으로 자르고 GPU i에게 i번째와 (2·CP−1−i)번째 조각을 함께 줘서 균형을 맞췄습니다.

8. 전문가 병렬: MoE의 전문가 나누기 EP

MoE(Mixture of Experts) 모델은 MLP 자리에 전문가 MLP를 여러 개 두고, 라우터가 토큰마다 몇 개만 골라 보냅니다. 전문가 병렬은 전문가들을 GPU에 나눠 배치하고, 토큰을 자기 전문가가 있는 GPU로 보냈다가(all-to-all) 결과를 다시 돌려받는(all-to-all) 방식입니다. GShard(Lepikhin et al., 2020)가 대규모로 쓴 뒤 표준이 되었습니다.

최근 대형 모델이 MoE로 많이 가면서 이 축의 비중이 커졌습니다. 예를 들어 DeepSeek-V3는 H800 2,048개에서 파이프라인 16 × 전문가 64(노드 8개에 걸침) × ZeRO-1 데이터 병렬로 학습했고 텐서 병렬은 쓰지 않았습니다. 반면 Llama 3 405B는 전문가가 없는 밀집(dense) 모델이라 전문가 병렬이 없습니다. 그래서 Llama 3는 "4D"입니다.

9. Llama 3 405B의 4D 병렬화

Meta의 「The Llama 3 Herd of Models」(2024) 3.3.2절은 405B 모델을 최대 16,384개 H100(80GB)에서 학습한 구성을 밝힙니다. 서버마다 GPU 8대가 NVLink로 묶여 있고, 서버 사이는 GPU당 400Gb/s RoCE 네트워크입니다. 여기에 네 가지 병렬화를 겹쳤습니다.

한눈에 보기 · GPU 한 대에서 16,384대까지

8K 문맥, 16,384대 구성(표 4 둘째 행)을 안쪽부터 쌓아 올리면 이렇습니다.

128K 장문맥 단계에서는 서버와 복제본 사이에 CP 16 축이 끼어듭니다(GPU 8 × CP 16 × PP 16 = 복제본 하나가 2,048대). 그만큼 DP는 8로 줄어듭니다.

그림 1 · 16대로 줄여 본 4D 배치 (Llama 3 그림 5)

논문 그림 5는 축마다 크기 2인 GPU 16대의 예로 그룹을 나누는 순서를 보여 줍니다. GPU 번호(rank)를 [TP, CP, PP, DP] 순서로, 즉 TP가 가장 빨리 바뀌게 매깁니다: rank = tp + 2·cp + 4·pp + 8·dp. 그래서 번호가 붙어 있는 GPU끼리(같은 서버 안) TP 그룹이 되고, 가장 멀리 떨어진 GPU끼리 DP 그룹이 됩니다.

TP 그룹은 서로 붙은 두 GPU, DP 그룹은 8칸 떨어진 두 GPU입니다. 축의 순서가 곧 "얼마나 가까운 GPU끼리 묶느냐"이고, 논문은 이 순서를 네트워크 통신에 맞춰 골랐다고 적습니다. 가장 안쪽 축이 가장 넓은 대역폭과 가장 짧은 지연을 요구하니 같은 서버 안에 두는 겁니다.

그림 2 · 실제 규모: 16,384개 GPU의 배치 (Llama 3 표 4)

표를 읽는 요령은 GPU 수 = TP × CP × PP × DP와 배치당 토큰 = 시퀀스 길이 × DP당 배치 × DP 두 식입니다. 세 행 모두 배치당 약 1,600만 토큰으로 같습니다. 8,192대에서 16,384대로 늘릴 때는 TP·PP는 그대로 두고 DP만 64 → 128로 늘렸고, 대신 DP당 배치가 절반이 됩니다. 문맥을 128K로 늘릴 때는 CP 16을 새로 넣고 그만큼 DP를 8로 줄였습니다. 전역 배치를 지키면서 축 사이에서 GPU를 주고받는 겁니다. GPU를 늘릴수록 GPU당 성능(MFU)이 조금씩 떨어지는 것도 보입니다.

Llama 3가 덧붙인 세부 조정

• 파이프라인 균형: 126층을 16 스테이지로 나누면 스테이지당 8층이 안 맞습니다. 임베딩이 붙은 첫 스테이지와 출력층이 붙은 마지막 스테이지에서 Transformer 층을 하나씩 빼서(7 + 8 × 14 + 7 = 126) 무게를 맞췄습니다.
• 파이프라인 스케줄: interleaved 스케줄을 쓰고, 마이크로배치 수를 자유롭게 정할 수 있게 고쳤으며, 스테이지 사이 통신을 비동기로 바꾸고, 다 쓴 텐서를 일찍 해제해 메모리를 줄였습니다.
• FSDP: 옵티마이저 상태와 기울기는 샤딩하지만, 파라미터는 순전파 뒤 다시 흩지 않고 역전파까지 들고 있어 all-gather 한 번을 아낍니다.
• 컨텍스트 병렬: 7절의 all-gather 방식과 2·CP 조각 짝짓기를 씁니다.

이 규모의 학습은 계산만큼 고장과의 싸움이기도 합니다. 논문은 사전학습 중 54일 동안 작업 중단이 466번(예상치 못한 중단 419번) 있었고, 예상치 못한 중단의 약 78%가 GPU·HBM 같은 하드웨어 문제로 확인되었거나 의심된다고 적습니다. 그런데도 유효 학습 시간을 90% 넘게 유지했습니다.

10. 정리: 축을 고르는 순서

어느 축을 몇으로 나눌지는 모델 크기, 문맥 길이, 클러스터 네트워크 구조에 따라 그때그때 정합니다. 다만 결정 순서는 대체로 같습니다.

같은 시리즈의 How GPipe Works는 파이프라인의 버블과 재계산을, How Adam Works는 학습 메모리에서 옵티마이저 상태가 차지하는 몫을, How Scaling Laws Work는 애초에 왜 모델을 이만큼 키우는지를 다룹니다.

11. 확인 퀴즈