How Deep Learning Computes

Picotron to Megatron Reading Map

마지막 수정:

trainingdistributedpicotronnanotronmegatronstrategy

이 path의 목적은 Megatron을 처음부터 끝까지 읽는 것이 아니다. 목표는 대규모 학습 병렬화를 읽는 좌표계를 만드는 것이다.

Parallelism

여러 GPU가 일을 나눠 수행한다.

GPU 0
data A GPU 1
data B GPU 2
data C
예: DP는 같은 모델 복사본들이 서로 다른 데이터를 처리한다.

Sharding

큰 tensor나 state를 조각내 저장한다.

GPU 0
state 0 GPU 1
state 1 GPU 2
state 2
예: ZeRO는 optimizer state, gradient, parameter를 DP rank에 나눠 저장한다.

세 코드베이스의 역할은 다르다.

Picotron:
원리를 숨기지 않는 작은 구현

Nanotron:
실제 trainer/config/checkpoint 구조가 있는 중간 규모 framework

Megatron:
production-scale reference와 최적화의 압축본

Picotron은 첫 번째 답을 준다

Picotron은 “이 병렬화가 본질적으로 무엇을 하는가?”에 답하기 좋다.

TP: linear weight를 column/row로 나눈다
PP: layer ownership과 microbatch schedule을 만든다
CP: K/V chunk를 ring으로 돌리고 online softmax로 합친다
EP: 구현이 거의 비어 있으므로 빈칸으로 인식한다

Picotron에서 얻어야 하는 것은 API 사용법이 아니다.

process group slicing
local tensor shape
communication primitive
forward/backward dataflow

이 네 가지를 말할 수 있으면 다음 코드베이스로 넘어갈 준비가 된 것이다.

Nanotron은 framework 경계를 보여준다

Nanotron은 Picotron보다 복잡하지만, Megatron보다 읽기 쉽다.

DistributedTrainer는 시작부터 병렬 컨텍스트를 만든다.

self.parallel_context = ParallelContext(
    tensor_parallel_size=self.config.parallelism.tp,
    pipeline_parallel_size=self.config.parallelism.pp,
    data_parallel_size=self.config.parallelism.dp,
    expert_parallel_size=self.config.parallelism.expert_parallel_size,
    context_parallel_size=self.config.parallelism.context_parallel_size,
)

Config 출력도 같은 축을 드러낸다.

tp | pp | dp | cp | ep

Nanotron에서 봐야 할 것은 “구현 원리”보다 “시스템 경계”다.

config -> ParallelContext -> model build -> optimizer/checkpoint -> trainer loop

예를 들어 PP는 PipelineEngine과 PipelineTrainBatchState로 분리되고, TP는 linear layer의 pg와 mode로 제어되며, EP는 MoE layer와 checkpoint naming까지 영향을 준다.

Megatron은 왜 production 코드가 복잡해지는지 보여준다

Megatron은 처음부터 읽으면 압도된다. 대신 질문을 들고 들어가야 한다.

TP 질문:
ColumnParallelLinear/RowParallelLinear에 production option이 왜 이렇게 많은가?

PP 질문:
1F1B, interleaving, virtual pipeline, activation deallocation은 어떤 병목을 줄이나?

CP 질문:
cp_comm_type, hierarchical CP, Transformer Engine integration은 topology와 어떤 관련이 있나?

EP 질문:
token dispatcher가 왜 all-gather/all-to-all/flex/DeepEP로 나뉘나?

Megatron 문서의 전략도 보수적이다.

1. DP로 시작한다
2. 모델이 안 맞으면 TP를 추가한다
3. 더 큰 모델에는 PP를 추가한다
4. 긴 sequence에는 CP를 추가한다

여기에 MoE 모델이면 EP가 들어온다.

dense model:
DP/FSDP + TP + PP + CP

MoE model:
DP/FSDP + TP + PP + CP + EP

어떤 순서로 학습해야 하나

이 path의 카드는 다음 순서로 다시 복습하면 좋다.

  1. Large-Scale Parallelism Coordinate System: rank 좌표계를 잡는다.
  2. Picotron ProcessGroupManager: group slicing을 손으로 이해한다.
  3. Tensor Parallel Linear from Picotron: TP의 local tensor shape을 잡는다.
  4. Sequence Parallelism as TP Layout: SP를 독립 axis가 아니라 TP layout mode로 본다.
  5. Pipeline Parallelism from Picotron: layer ownership과 AFAB/1F1B schedule을 읽는다.
  6. Context Parallel Ring Attention: K/V ring과 online softmax를 이해한다.
  7. Expert Parallel MoE Dispatch: router 결과가 all-to-all plan이 되는 과정을 본다.
  8. 이 카드에서 세 코드베이스의 역할을 정리한다.

실습 코드는 무엇을 증명하나

labs/large-scale-training-parallelism의 실습은 작지만 각 카드의 핵심 불변식을 확인한다.

rank_grid.py:
world_size = DP x PP x CP x TP x EP

tp_linear_sim.py:
column/row sharding이 full linear와 같은 결과를 만든다

sp_layout_sim.py:
SP shard -> TP matmul full sequence -> SP shard 복귀

pipeline_schedule_sim.py:
AFAB와 1F1B-style schedule의 bubble/utilization 차이

cp_ring_attention_sim.py:
chunked online attention이 full causal attention과 같은 결과

moe_dispatch_sim.py:
router output이 all-to-all send buckets와 tokens_per_expert로 바뀜

이 실습들은 multi-node training을 대신하지 않는다. 대신 production code를 읽기 전에 필요한 mental model을 검증한다.

다음 확장 기준

이제 다음 단계는 “라이브러리 사용법”보다 “작은 구현을 framework로 키우는 법”이다.

1. TP linear sim -> 작은 torch.distributed TP layer
2. PP schedule sim -> 2-process pipeline toy model
3. CP attention sim -> torch.distributed ring K/V exchange
4. MoE dispatch sim -> all_to_all_single 기반 token dispatch
5. Nanotron config로 같은 병렬 축을 실제 trainer에서 실행
6. Megatron에서 같은 개념의 production 구현 위치 확인

이렇게 가면 Pico에서 Mega로 점프하지 않는다. 작은 invariant를 먼저 확인하고, framework 경계를 붙인 뒤, production 최적화를 읽는다.

최종 판단 기준

이 path를 완료했다는 기준은 용어 암기가 아니다. 다음 질문에 답할 수 있어야 한다.

TP:
내 rank의 weight shard shape은 무엇이고, forward/backward에서 어떤 collective가 필요한가?

SP:
sequence shard layout이 언제 all-gather/reduce-scatter로 바뀌는가?

PP:
내 stage가 어떤 layer를 소유하고, microbatch activation은 언제 저장/해제되는가?

CP:
local Q가 어떤 K/V chunk를 어떤 순서로 보고, dK/dV는 어디로 돌아가는가?

EP:
router의 top-k 결과가 어떻게 all-to-all communication plan으로 바뀌는가?

이 다섯 질문을 코드 위치와 함께 답할 수 있으면, Megatron 같은 production training stack을 읽을 준비가 된 것이다.

먼저 보면 좋은 카드Large-Scale Parallelism Coordinate SystemPicotron, Nanotron, Megatron을 같은 DP/TP/PP/CP/EP 좌표계로 읽기 위한 기준을 세운다.먼저 보면 좋은 카드Expert Parallel MoE DispatchPicotron에는 비어 있는 EP 축을 Nanotron의 작은 MoE 구현과 Megatron의 token dispatcher 구조로 채우며 router, permutation, all-to-all, expert compute의 흐름을 이해한다.