Large-Scale Parallelism Coordinate System

DP/FSDP까지 이해했다면 다음 질문은 단순히 “GPU를 더 붙이면 빨라지는가?”가 아니다.

모델 state는 어디에 있는가?
layer 내부 행렬곱은 어디서 나뉘는가?
layer depth는 어디서 나뉘는가?
sequence length는 어디서 나뉘는가?
MoE expert는 어디서 나뉘는가?

이 path에서는 이 질문을 하나의 좌표계로 본다.

world_size = DP x TP x PP x CP x EP

DP PP CP TP EP

world size 8 2 x 2 x 1 x 2 x 1

왜 좌표계가 먼저인가

대규모 학습 코드는 병렬화 기법 하나만 따로 쓰지 않는다. Megatron 문서도 DP, TP, PP, CP, EP, Megatron-FSDP를 조합 가능한 전략으로 설명한다.

그래서 코드를 읽을 때 먼저 확인해야 하는 것은 “이 함수가 무슨 최적화를 하는가?”가 아니라 “이 rank가 어떤 좌표에 있는가?”이다.

rank 7
  dp rank = ?
  tp rank = ?
  pp rank = ?
  cp rank = ?
  ep rank = ?

이 좌표가 정해져야 내가 속한 process group도 정해진다.

세 프로젝트의 역할

Picotron은 이 좌표계를 가장 노골적으로 보여준다. ProcessGroupManager는 world size가 TP x CP x PP x DP와 맞는지 확인하고, rank들을 4D grid로 만든 뒤 TP/CP/PP/DP group을 만든다. 그래서 이 path의 출발점은 Picotron이다.

Nanotron은 같은 아이디어를 실제 training framework 설정으로 올린다. ParallelismArgs에는 dp, pp, tp, expert_parallel_size, context_parallel_size가 들어간다. 여기서부터는 알고리즘만 보는 것이 아니라 config, trainer, checkpoint, pipeline engine까지 같이 봐야 한다.

Megatron은 production reference다. TP/PP/DP/CP/EP를 모두 지원하고, 여기에 FSDP, distributed optimizer, communication overlap, virtual pipeline, MoE token dispatch 같은 실제 운영 복잡도가 붙는다. 처음부터 Megatron으로 배우면 좌표계보다 예외 처리와 최적화가 먼저 보인다.

따라서 읽는 순서는 다음이 좋다.

Picotron  -> 좌표계와 최소 구현
Nanotron  -> framework 구조와 설정
Megatron  -> production 조합과 최적화

이 카드에서 기억할 것

TP/PP/CP/EP는 “DDP보다 어려운 기능 목록”이 아니다. 서로 다른 병목을 서로 다른 축으로 나누는 방법이다.

DP: batch를 나눈다
TP: layer 내부 tensor 계산을 나눈다
PP: layer depth를 나눈다
CP: sequence length를 나눈다
EP: MoE expert와 token routing을 나눈다

첫 번째 실습은 실제 GPU 통신 전에 rank grid를 손으로 만든다. 이 좌표계가 없으면 Picotron의 간단한 코드도, Megatron의 parallel_state.py도 읽기 어렵다.

확인

• world_size = DP x TP x PP x CP x EP에서 각 축은 무엇을 나누는가?
• Picotron을 먼저 읽는 이유는 무엇인가?
• Megatron을 첫 교재가 아니라 마지막 reference로 두는 이유는 무엇인가?