Distributed GEMM

현대 LLM 의 한 layer 는 이미 한 GPU 안에 다 안 들어간다. 70B 모델의 한 attention block 의 weight 만 해도 수십 GB. 그래서 layer 안의 GEMM 자체를 여러 GPU 로 split 한다 — 이게 tensor parallelism 의 코어. 그리고 그 split 은 communication 과 compute 의 overlap 이라는 새 문제를 만든다.

강의 transcript 가 누락되었지만 슬라이드 링크가 살아 있어 핵심 토픽이 분명하다 — distributed GEMM 의 표준 패턴, 통신 비용 분석, NCCL 매핑, production trade-off. 본 노트는 슬라이드 + Megatron-LM / NCCL 공식 자료 + 도메인 지식 위에서 정리한다.

이 강의가 가리키는 자리 — single-GPU GEMM 강의 (L045) 의 위에 한 단계 — multi-GPU GEMM. L048 의 ultra scale playbook 이 다양한 parallelism 의 결합을 다룬다면, 본 강의는 그 중 한 축 (TP) 의 구체적 GEMM 구현 메커니즘.

transformer 의 한 layer 가 안고 있는 GEMM 은 큰 것 4개. attention 의 QKV projection, attention output projection, MLP 의 두 linear (up_proj + down_proj). 이 GEMM 들이 각자 다른 split 전략을 만난다.

핵심 패턴 — Megatron-LM 의 표준 구조. QKV(column) → O(row), up_proj(column) → down_proj(row). 두 linear 를 묶으면 column-parallel 의 결과를 그대로 row-parallel 의 입력으로 흘려보낼 수 있어 — 중간에 통신 없이 두 GEMM 을 묶을 수 있다. 통신은 attention output 끝과 MLP 끝에 한 번씩.

한 GEMM Y = X @ W 를 N GPU 사이로 자르는 방법은 두 가지. output dim 으로 (column-parallel), input dim 으로 (row-parallel). 둘은 입출력의 shape 이 다르고 통신이 다르다.

두 방식의 비용 구조 차이 — 통신이 입력 쪽에서 발생하는가, 출력 쪽에서 발생하는가. 그래서 column-parallel 의 output → row-parallel 의 input으로 묶으면 중간 통신 없이 흐른다. Megatron-LM 의 핵심 설계.

distributed GEMM 의 통신은 거의 모두 NCCL 의 표준 collective 로 표현된다. 세 종류가 한 transformer block 안에서 모두 등장.

tensor parallel 의 모든 최적화가 가리키는 한 자리 — compute 와 comm 을 같은 시간에 굴린다. naive 구현은 compute 끝나고 comm 시작 = serial. 좋은 구현은 그 둘이 시간 위에서 겹친다.

all-reduce 같은 추상 연산이 실제 hardware 의 link 위에서 어떻게 실행되는지. NCCL 은 데이터 크기와 GPU 토폴로지에 따라 두 알고리즘 사이에 자동 분기.

distributed GEMM 의 비용 분석은 hardware topology 에 종속이다. 한 노드 (8 GPU) 안의 NVLink 와 노드 사이의 InfiniBand 의 대역폭 차이가 결정 요소.

이 hierarchy 가 결정을 만든다 — TP 는 보통 한 노드 안 (8 GPU) 까지. 노드를 넘어서 TP 를 키우면 IB 위 통신 = 10배 느림. 노드 사이는 PP (pipeline parallel) 또는 DP (data parallel) 로 — 이쪽은 통신 빈도가 낮다.

최근 (2024–2025) production system 에 들어간 두 추가 트릭이 강의가 다뤘을 자리. 둘 다 통신을 더 잘게 쪼개서 overlap 을 가능하게 하거나 메모리를 절약.

distributed GEMM 의 reference 구현 — NVIDIA Megatron-LM. 한 transformer block 의 forward 가 어떤 NCCL 호출 시퀀스를 거치는지 확인할 수 있는 가장 깨끗한 자료.

# Megatron-LM 의 ColumnParallelLinear 흐름 (간략)
class ColumnParallelLinear:
    def forward(self, x):
        # 1. 모든 TP rank 가 같은 x 가 필요
        x = copy_to_tensor_model_parallel_region(x)   # identity in fwd, AR in bwd

        # 2. 각 rank 의 partial weight 와 GEMM
        out = F.linear(x, self.weight_local, self.bias_local)

        # 3. output 은 column-split — 그대로 반환
        return out

class RowParallelLinear:
    def forward(self, x):
        # 1. 입력은 column-split (이전 column-par 의 output)
        out_partial = F.linear(x, self.weight_local, bias=None)

        # 2. partial sum 합치기 → all-reduce (또는 RS for sequence parallel)
        out = reduce_from_tensor_model_parallel_region(out_partial)

        # 3. bias 는 reduce 후 한 번만 적용
        return out + self.bias if self.bias is not None else out

이 두 클래스를 묶으면 — QKV(Column) → attention → O(Row + AR) → up(Column) → activation → down(Row + AR). 한 block 당 통신 2번. backward 에서 추가 2번 (column-par 의 bwd 가 AR).

손에 새기기 — 실습 시퀀스

1. 2-GPU all-reduce 측정 — torch.distributed 로 작은 tensor all-reduce. NVLink 위 latency 와 bandwidth 가 NCCL 공식 수치와 일치하는지.
2. row-parallel linear 직접 구현 — Megatron-LM 의 RowParallelLinear 를 자기 코드로 재현. forward 에서 all-reduce, backward 에서 identity.
3. 한 transformer block 의 forward — QKV(column) → attention → O(row) → up(column) → down(row) 를 8 GPU 위에서. 통신이 정확히 2회인지 확인.
4. compute-comm overlap 측정 — Nsys 로 trace 떠서 GEMM 과 NCCL 이 같은 시간에 도는지 확인.
5. SP 추가 — all-reduce 를 RS + AG 로 분해. layernorm 을 그 사이에. 메모리 사용량 확인.
6. 크기 별 strong scaling 측정 — TP=1, 2, 4, 8 에서 latency 측정. ideal scaling 과의 격차 = 통신 비용.

본 노트는 슬라이드 링크 + Megatron-LM 코드 + NCCL 문서를 베이스로 토픽을 추정. 영상 직접 확인이 필요한 사항들.

• 강의가 직접 다룬 구체적 패턴 — async TP 가 본 강의의 메인 토픽인지, fused TP 인지, 또는 일반 TP 의 introduction 인지 — 슬라이드 직접 확인 필요.
• 구체적 benchmark 수치 — H100 / A100 / 노드 수 별 throughput, scaling efficiency 의 실측. 슬라이드의 그래프와 영상의 발화 직접 확인.
• CUTLASS 3.x distributed primitive — 강의 시점(2025-03) 에 CUTLASS 가 부분적으로 distributed primitive 노출. 강의가 이 자리를 다뤘는지.
• TorchTitan / DTensor 와의 관계 — PyTorch 의 새 distributed API. 본 강의의 패턴이 이 추상 위에서 어떻게 표현되는지.
• MNNVL (multi-node NVLink) 의 영향 — GB200 NVL72 같은 시스템에서 64+ GPU 의 single NVLink domain. 기존 노드 단위 가정이 어떻게 변하는지.
• Ali Hassani 의 NATTEN 과의 관계 — Hassani 의 다른 작업이 dilated/local attention. distributed GEMM 강의가 그 작업과 어떻게 연결되는지 — 영상 직접 확인.