분산 LLM 병렬화 실행 흐름 시각화

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

대규모 언어 모델의 분산 학습에서는 FSDP, Tensor Parallel, Expert Parallel, Context parallel처럼 서로 다른 병렬화 방식과 장치 간 통신의 실행 순서를 이해해야 합니다. 이 프로젝트는 실제 torchtitan profile을 바탕으로 compute kernel과 collective를 drag and drop해 DDP, TP, FSDP, EP, CP trace를 구성하도록 지원합니다. 사용자는 계산과 통신의 겹침을 시각적으로 비교하면서 각 병렬화 방식의 실행 흐름과 잠재적 병목을 직관적으로 파악할 수 있습니다.

섹션별 상세

대규모 언어 모델 학습은 단순히 GPU 수를 늘리는 문제가 아니라, 여러 장치에 모델과 데이터를 나누는 샤딩 방식과 시스템 실행 순서를 함께 설계해야 하는 과제입니다. 기존에 LLM 확장 자료는 많지만, 병렬화 방식이 실제 학습 실행에서 어떻게 겹치고 어떤 순서로 진행되는지 직관적으로 파악할 수 있는 시각화 자료에는 공백이 있었습니다. 이 프로젝트는 그 공백을 분산 학습 trace를 직접 조작하는 방식으로 줄이려는 목적을 가집니다.

FSDP, Tensor Parallel, Expert Parallel, Context parallel의 실행 흐름을 시각적으로 구성할 수 있도록 compute kernel과 collective를 drag and drop 방식으로 배치합니다. 사용자는 각 연산과 통신 요소의 위치를 바꾸면서 DDP, TP, FSDP, EP, CP trace를 만들고 병렬화 방식별 실행 순서와 겹침을 비교할 수 있습니다. trace의 기반은 실제 torchtitan profile이므로 추상적인 도식보다 실제 분산 학습 실행 데이터에 가까운 형태로 병렬 처리 흐름을 살필 수 있습니다.

이 도구의 핵심 가치는 특정 병렬화 기법 하나의 사용법보다, 계산 커널과 장치 간 통신이 언제 시작되고 서로 얼마나 겹치는지를 이해하게 만드는 데 있습니다. 사용자가 trace를 직접 조립하면 샤딩 방식에 따른 통신 지점과 실행 순서를 확인하고 분산 학습 병목이 발생하는 위치를 추적할 수 있습니다. 따라서 LLM 학습 시스템을 설계하거나 성능을 해석할 때 병렬화 전략의 차이를 시각적 실행 trace로 연결하는 학습 자료가 됩니다.

용어 해설

FSDP: — Fully Sharded Data Parallel의 약자로, 모델 파라미터·그래디언트·옵티마이저 상태를 여러 GPU에 분산 저장하는 병렬화 방식입니다. 필요한 계산 시점에 데이터를 모으고 계산이 끝나면 다시 샤딩해 메모리 사용량을 줄입니다.
Tensor Parallelism: — 하나의 신경망 레이어에서 행렬 연산을 여러 장치로 나누는 방식입니다. 각 장치가 텐서 일부를 계산한 뒤 collective communication으로 중간 결과를 교환하며, 대형 모델의 단일 장치 메모리 한계를 완화합니다.
Expert Parallelism: — Mixture-of-Experts 구조에서 서로 다른 Expert 네트워크를 여러 장치에 나누어 배치하는 병렬화 방식입니다. 토큰을 선택된 Expert로 라우팅하고 결과를 모으는 과정의 실행 순서와 통신 겹침이 성능에 영향을 줍니다.
Context Parallelism: — 긴 입력 시퀀스의 Context를 여러 장치에 분할해 처리하는 병렬화 방식입니다. 각 장치가 일부 토큰 구간을 계산하고 필요한 정보를 교환하므로, 시퀀스 길이에 따른 메모리 부담과 통신 흐름을 함께 고려해야 합니다.
DDP: — Distributed Data Parallel의 약자로, 각 장치가 서로 다른 데이터 배치를 같은 모델로 처리한 뒤 그래디언트를 동기화하는 방식입니다. 전체 장치가 동일한 파라미터를 유지하면서 데이터 병렬 학습 처리량을 높이는 데 사용됩니다.
Collective Communication(collectives): — 여러 장치가 함께 수행하는 통신 연산으로, 데이터 교환과 결과 집계에 사용됩니다. 분산 학습에서는 collective의 시작 시점과 계산 커널의 실행 순서를 겹치게 배치하는지가 전체 실행 시간과 자원 활용률을 좌우합니다.

기술

FSDP
Tensor Parallel
Expert Parallel
Context parallel
DDP
TP
EP
CP
torchtitan

활용 사례

분산 LLM 학습 trace 시각화
병렬화 방식별 실행 순서 비교
compute kernel과 collective communication의 겹침 파악
torchtitan profile 기반 학습 시스템 직관 형성

언급된 리소스

튜토리얼https://huggingface.co/spaces/nanotron/ultrascale-playbook

문서https://jax-ml.github.io/scaling-book/