Picotron Process Group Manager
마지막 수정:
Picotron에서 가장 먼저 읽을 파일은 ProcessGroupManager다.
이 파일이 중요한 이유는 단순하다. TP, PP, CP, DP는 모두 “rank 일부를 묶은 process group”이고, Picotron은 그 group을 어떻게 자르는지 거의 그대로 보여준다.
grid = ranks.view(dp, pp, cp, tp)
[2, 3] Picotron의 핵심 아이디어
Picotron은 world size를 먼저 검증한다.
world_size == tp_size * cp_size * pp_size * dp_size
그다음 rank를 4D grid로 본다.
grid[dp, pp, cp, tp]
여기서 각 process group은 한 축만 움직이는 slice다.
TP group: grid[d, p, c, :]
CP group: grid[d, p, :, t]
PP group: grid[d, :, c, t]
DP group: grid[:, p, c, t]
이게 대규모 분산 학습 코드 읽기의 첫 번째 열쇠다. all_reduce(..., group=tp_group)를 보면 “TP 축으로 묶인 rank들끼리만 통신한다”는 뜻이다. send(..., pp_next_rank)를 보면 “pipeline 축에서 다음 stage로 activation을 보낸다”는 뜻이다.
Nanotron과 비교
Nanotron도 같은 구조를 쓴다. 다만 교육용으로 최소화한 Picotron보다 framework 상태가 더 많다.
Nanotron의 ParallelContext는 rank matrix를 다음 순서로 만든다.
ranks[ep, pp, dp, cp, tp]
그리고 tp_pg, cp_pg, dp_pg, pp_pg, ep_pg, mp_pg, dp_cp_pg를 만든다. 여기서부터는 단순히 group을 만드는 것뿐 아니라 trainer, model, checkpoint, tied weight, MoE expert sharding이 이 group들을 계속 참조한다.
즉 Nanotron에서 배울 점은 다음이다.
process group은 전역 설정이 아니라
trainer/model/checkpoint가 공유하는 execution context다
Megatron과 비교
Megatron은 이 아이디어를 더 일반화한다.
parallel_state.py에는 RankGenerator가 있고, tp-cp-ep-dp-pp 같은 order와 mask로 group을 만든다. Picotron에서는 눈으로 보이는 slicing이 Megatron에서는 “어떤 축을 mask할 것인가”라는 일반 함수가 된다.
Picotron:
grid[d, p, c, :]
Megatron:
RankGenerator(...).get_ranks("tp")
처음부터 Megatron을 읽으면 이 추상화가 과하게 느껴진다. 하지만 Picotron에서 slice 의미를 먼저 이해하면 Megatron의 RankGenerator는 “더 많은 축과 order를 처리하는 일반화된 group factory”로 보인다.
실습
이 카드의 실습은 GPU 없이 rank grid를 만든다.
cd labs/large-scale-training-parallelism
python3 rank_grid.py --dp 2 --pp 2 --cp 1 --tp 2 --ep 1 --rank 3
출력에서 rank 3의 TP group, PP group, DP group이 어떻게 달라지는지 확인한다. 실제 torch.distributed 코드는 나중에 붙인다. 먼저 중요한 것은 “통신 group은 좌표 slice다”라는 감각이다.
확인
- Picotron에서
grid[d, p, c, :]는 어떤 group인가? - Nanotron의
ParallelContext는 왜 trainer/model/checkpoint에서 계속 필요할까? - Megatron의
RankGenerator는 Picotron의 slicing을 어떤 방식으로 일반화하는가?