분산 LLM 학습 19 - Megatron-LM과 DeepSpeed를 어떻게 읽을 것인가

프레임워크를 기능 목록으로 외우면 남는 게 적다

Megatron-LM, DeepSpeed, FSDP 기반 스택은 각자 제공하는 옵션이 많다. 하지만 실무에서 중요한 것은 "이 프레임워크가 무엇을 추상화하고 무엇은 여전히 사용자가 책임져야 하는가"를 읽는 것이다.

Megatron-LM을 볼 때

Megatron 계열은 transformer 구조에 맞춘 tensor parallel, pipeline parallel, sequence parallel 설계를 비교적 노골적으로 드러낸다. 따라서 "연산을 어떻게 나누는가"를 배우기 좋은 코드다.

볼 포인트:

• layer 내부 parallel group 구성
• pipeline stage 배치
• micro-batch schedule
• optimizer / checkpoint가 병렬화 구조와 어떻게 결합되는가

DeepSpeed를 볼 때

DeepSpeed는 더 넓은 실무 기능을 제공하는 쪽에 가깝다. ZeRO, optimizer offload, engine abstraction 같은 구성 요소를 통해 메모리와 운영 문제를 다룬다.

볼 포인트:

• stage별 state sharding
• runtime engine이 step을 어떻게 감싼는가
• checkpoint와 optimizer 상태를 어떻게 관리하는가

중요한 것은 "어떤 병목을 겨냥했는가"

프레임워크를 비교할 때는 이렇게 보는 편이 좋다.

• 모델이 너무 커서 메모리가 문제인가
• intra-layer 분산이 필요한가
• multi-node 환경에서 scheduler가 중요한가
• checkpoint와 운영 도구가 더 중요한가

즉 프레임워크는 추상화의 묶음이지, 만능 정답이 아니다.

코드 읽기의 실제 기준

• parallel group 정의가 어디 있는가
• collectives가 어느 층에서 발생하는가
• state dict와 checkpoint 구조가 어떻게 생겼는가
• optimizer step과 grad sync가 어느 abstraction 아래 묶여 있는가

이 기준이 있으면 거대한 코드베이스도 조금 덜 압도적으로 느껴진다.

다음 글에서는 이 시리즈를 마무리하며, 실제 LLM 학습 스택을 설계할 때 어떤 순서로 병렬화 전략을 선택하고 검증할지 정리한다.