이 글은 2개의 노드, 16개의 GPU로 구성된 RTX PRO 6000 Blackwell 클러스터에서 Megatron 분산 학습 (distributed training) 설정을 구축하고 테스트하는 것에 관한 독립적인 기술 보고서입니다.

목표는 최종적인 고품질 모델을 제작하는 것이 아니었습니다. 목표는 더 큰 규모의 학습을 가능하게 하는 인프라 계층(infrastructure layer)을 이해하는 것이었습니다: 모델 임포트 (model import), 분산 실행 (distributed launch), 병렬 구조 (parallelism layout), 순전파/역전파 (forward/backward) 정확성, 체크포인팅 (checkpointing), 리샤딩 (resharding), 그리고 장애 모드 (failure modes) 등이 그것입니다.

범위 (Scope)

초점은 최종 모델의 품질이 아닌 학습 인프라에 맞춰졌습니다.

품질은 데이터, 레시피 (recipes), 하이퍼파라미터 (hyperparameters), 평가 (evals), 토큰 예산 (token budget), 그리고 반복적인 어블레이션 (ablations)에 따라 달라집니다. 이번 테스트는 분산 학습 스택이 제대로 작동하는지에 초점을 맞췄습니다:

• 실행 (launch)
• 모델 구축 (model construction)
• 체크포인트 임포트 (checkpoint import)
• 체크포인트 리샤딩 (checkpoint resharding)
• 순전파/역전파 (forward/backward)
• 그래디언트 감소 (gradient reduction)
• NaN/스킵 없음 (no NaNs/skips)
• 체크포인트 저장 (checkpoint save)
• 반복 가능한 설정 기록 (repeatable config recording)

다시 말해, 질문은 다음과 같았습니다:

모델을 변환하고, GPU 전체에 분할하여, 몇 단계 동안 학습시키고, 올바르게 동기화하고, 저장한 다음, 더 긴 지속적 사전 학습 (continued-pretraining) 또는 지도 미세 조정 (supervised fine-tuning) 실행을 위한 베이스로 사용할 수 있는가?

하드웨어 / 클러스터 (Hardware / Cluster)

• 2개의 노드.
• 노드당 8개의 RTX PRO 6000 Blackwell GPU.
• 총 16개의 GPU.
• 각 노드는 데이터셋, 체크포인트 및 로그를 위해 노드 로컬 NVMe 스토리지를 사용했습니다.
• NCCL/분산 학습을 위해 프라이빗 노드 간 네트워킹 (Private internode networking)이 사용되었습니다.
• 대부분의 대규모 모델 스모크 테스트 (smoke tests)는 사용 가능한 8-GPU 노드 하나에서 실행되었습니다.
• 공개 WikiText 5D 실행은 두 노드와 16개의 GPU를 모두 사용했습니다.

스택 (Stack)

• Megatron-LM: 훈련 런타임 (training runtime), 실행 스크립트 (launch scripts), 모델 병렬 그룹 (model-parallel groups), 분산 옵티마이저 (distributed optimizer), 그리고 체크포인트 흐름 (checkpoint flow).
• Megatron-Bridge: Hugging Face에서 Megatron으로의 변환 (conversion), 모델 제공자 (model providers), 레시피 (recipes), 그리고 훈련 엔트리포인트 (training entrypoints).
• Megatron-Core: 모델에서 사용되는 텐서 (tensor), 파이프라인 (pipeline), 컨텍스트 (context), 전문가 (expert), 그리고 시퀀스 병렬 (sequence-parallel) 프리미티브 (primitives).
• Transformer Engine: 최적화된 트랜스포머 커널 (transformer kernels), FP8 경로 (FP8 paths), 어텐션 커널 (attention kernels), 그리고 통신 오버랩 (communication overlap).
• NCCL: 텐서 병렬성 (tensor parallelism), 데이터 병렬성 (data parallelism), 전문가 라우팅 (expert routing), 그리고 멀티노드 통신 (multinode communication)을 위한 GPU 컬렉티브 (GPU collectives).
• Torch distributed checkpointing: 샤딩된 체크포인트 저장/로드 (sharded checkpoint save/load) 및 서로 다른 레이아웃 간의 리샤딩 (resharding).
• Hugging Face integration: 공개 모델/토크나이저 (tokenizer) 다운로드 및 체크포인트 임포트 (import).
• BF16, FP8, 그리고 MXFP8 정밀도 경로 (precision paths).

구성 요소의 결합 방식 (How The Pieces Fit Together)

스택은 세 가지 주요 레이어로 구성됩니다:

1. 모델 소스 포맷 (Model source format)

공개 모델 체크포인트는 보통 Hugging Face 스냅샷으로 시작합니다. 해당 포맷은 추론 (inference)과 생태계 호환성 측면에서 편리하지만, Megatron이 분산 학습 (distributed training)을 위해 사용하는 레이아웃과는 동일하지 않습니다.

1. 변환 및 모델 제공자 레이어 (Conversion and model provider layer)

Megatron-Bridge는 Hugging Face 모델을 Megatron의 모델 제공자 (model-provider) 및 체크포인트 포맷으로 매핑합니다. 이 단계에서는 아키텍처 특화 세부 사항이 중요하게 작용합니다: MoE 라우팅 (MoE routing), 전문가 레이아웃 (expert layout), 어텐션 변형 (attention variants), Mamba 블록 (Mamba blocks), 토크나이저 파일 (tokenizer files), 정밀도 설정 (precision config), 그리고 대상 병렬 레이아웃 (target parallelism layout).

1. 분산 학습 레이어 (Distributed training layer)

그 후 Megatron-LM/Megatron-Core는 명시적인 병렬 그룹 (parallel groups)을 사용하여 모델을 학습합니다. 모델 아키텍처와 체크포인트 레이아웃이 지원하는 한, 동일한 모델을 TP=2, PP=2, 또는 EP=4와 같은 서로 다른 런타임 레이아웃으로 로드할 수 있습니다.

실제 흐름은 다음과 같았습니다:

Hugging Face model/tokenizer
        |
        v
...

Megatron 실행 방식 (How A Megatron Run Works)

유용한 Megatron 스모크 테스트 (smoke run)는 단순히 "스크립트가 시작되었다"는 것 이상의 의미를 갖습니다. 최소한의 증거는 다음과 같습니다:

1. 모든 랭크 (rank)에 걸쳐 분산 프로세스 그룹 (distributed process group)이 초기화됩니다.
2. 요청된 TP, PP, CP, DP, EP, ETP, EDP, SP 레이아웃으로 모델이 빌드됩니다.
3. 모든 참여 노드에서 데이터셋 및 토크나이저 (tokenizer) 경로가 정상적으로 확인됩니다.
4. 백엔드 (backend) 또는 형상 (shape) 오류 없이 순전파 (forward pass)가 실행됩니다.
5. 역전파 (backward pass) 및 그래디언트 동기화 (gradient synchronization)가 완료됩니다.
6. NaN (Not a Number) 또는 건너뛴 반복 (skipped iterations)이 나타나지 않습니다.
7. 예상된 분산 레이아웃으로 체크포인트 (checkpoint)가 기록됩니다.
8. 활성화된 경우 검증/테스트 (validation/test)를 실행할 수 있습니다.

이것이 짧은 실행이 여전히 가치 있는 이유입니다. 이는 긴 학습 시간을 소비하기 전에 인프라 계약 (infrastructure contract)을 테스트합니다.

Megatron-Bridge가 사용된 방식

실험에 공개된 Hugging Face 모델이 포함되었지만 학습 대상은 Megatron이었기 때문에 Megatron-Bridge가 유용했습니다. Bridge는 중요한 중간 계층을 처리했습니다:

• 공개 모델 설정 및 토크나이저 (tokenizer) 자산 로드.
• Hugging Face 파라미터를 Megatron 파라미터 이름 및 텐서 레이아웃 (tensor layouts)으로 매핑.
• Megatron 호환 체크포인트 (checkpoint) 생성.
• 학습 레이아웃이 임포트 레이아웃과 다를 때 런타임 리샤딩 (runtime resharding) 지원.
• DeepSeek MoE 및 Nemotron 하이브리드 모델과 같은 아키텍처를 위한 모델별 레시피 (recipes) 및 프로바이더 (providers) 제공.

실질적인 교훈: 변환 (conversion)은 부수적인 작업이 아닙니다. 대규모 모델 학습에서 변환과 체크포인트 레이아웃은 학습 시스템의 일부입니다.

병렬성 멘탈 모델 (Parallelism Mental Model)

Megatron은 여러 병렬성 축 (parallelism axes)을 직접 노출하기 때문에 유용합니다:

축 (Axis)	분할 대상	중요한 이유
TP	레이어 내부의 밀집 텐서 연산 (Dense tensor operations)	대규모 행렬 연산에 대한 GPU당 메모리 및 연산량 감소
...

중요한 점은 이러한 축들이 단순히 곱하기 관계가 아니라는 것입니다. 수학적으로 유효한 월드 사이즈 (world size)라 할지라도 레이어 수, 어텐션 백엔드 (attention backend) 지원, 커널 (kernel) 가용성, 메모리 압박 또는 체크포인트 레이아웃으로 인해 실패할 수 있습니다.

예를 들어:

• PP=2는 DeepSeek-V2-Lite에서 불균형한 파이프라인 분할 (pipeline splitting)을 사용하기 전까지 실패했습니다. 이는 모델이 27개의 레이어를 가지고 있기 때문입니다.
• CP=2는 커스텀 MoE GPT 및 Nemotron 경로에서는 작동했지만, DeepSeek-V2-Lite의 MLA 어텐션 (attention) 경로는 백엔드 (backend) 문제를 일으켰습니다.
• SP는 TP와 함께 메모리에 도움을 주었지만, 새로운 랭크 차원 (rank dimension)을 추가하지는 않았습니다.

체크포인팅 (Checkpointing) 및 리샤딩 (Resharding)

체크포인팅 (Checkpointing)은 인프라 테스트의 일부입니다. 순전파 (forward) 및 역전파 (backward)는 완료되었으나 저장할 수 없는 실행은 완전한 분산 학습 (distributed training) 증명이라고 할 수 없습니다.

두 가지 체크포인트 아이디어가 중요했습니다:

• 체크포인트 가져오기 (Import checkpoint): 공개된 Hugging Face 모델로부터 생성된 Megatron 형식의 체크포인트.
• 런타임 체크포인트 (Runtime checkpoint): 특정 병렬성 레이아웃 (parallelism layout)을 가진 분산 학습 후 기록된 체크포인트.

Megatron은 레이아웃 간에 리샤딩 (resharding)을 수행할 수 있지만, 주변 스토리지와 메타데이터 (metadata)는 여전히 깨끗해야 합니다. 16-GPU WikiText 공개 실행은 새로운 학습/저장/평가 (train/save/eval)를 입증했습니다. 임시 노드 로컬 (node-local) 체크포인트 레이아웃으로부터의 재개 (Resume)는 메타데이터 로드 후 타임아웃이 발생했기 때문에 여전히 정리가 필요합니다.

용어 사전 (Glossary)

용어	의미
CPT	계속 사전 학습 (Continued pretraining): 텍스트 데이터에 대한 다음 토큰 학습 (next-token training)
...

이 작업에 Megatron을 사용하는 이유

FSDP는 많은 밀집 모델 (dense-model) 및 미세 조정 (fine-tuning) 작업에 유용합니다. 하지만 이번 작업의 목표는 명시적인 토폴로지 제어 (topology control)였기 때문에 Megatron이 더 나은 대상이었습니다:

• TP: 텐서 병렬성 (tensor parallelism), 밀집 텐서 연산을 분할.
• PP: 파이프라인 병렬성 (pipeline parallelism), 레이어를 스테이지 (stages) 간에 분할.
• CP: 컨텍스트 병렬성 (context parallelism), 긴 시퀀스/컨텍스트 토큰을 분할.
• DP: 데이터 병렬성 (data parallelism), 데이터 배치 (data batches) 전반에 걸쳐 모델 그룹을 복제.
• EP: 전문가 병렬성 (expert parallelism), MoE 전문가 (experts)를 GPU 간에 분산.
• ETP: 전문가 텐서 병렬성 (expert tensor parallelism), 각 전문가 내부의 가중치 (weights)를 샤딩 (sharding).
• EDP: 전문가 데이터 병렬성 (expert data parallelism), 전문가 병렬 그룹을 복제.
• SP: 시퀀스 병렬성 (sequence parallelism), TP 그룹 내부의 활성화 메모리 (activation memory)를 감소.

MoE (Mixture-of-Experts) 및 하이브리드 모델의 경우, 병목 현상은 파라미터 메모리(parameter memory)에만 국한되지 않습니다. 이는 밀집 가중치 (dense weights), 라우팅된 전문가 (routed experts), 시퀀스 길이 (sequence length), 옵티마이저 상태 (optimizer state), 체크포인트 레이아웃 (checkpoint layout) 또는 통신 (communication) 문제일 수 있습니다. Megatron은 이러한 축들을 직접적으로 노출합니다.

실행된 실험 (Experiments Run)

모델 / 작업	실행 이름	설정 (Config)	정밀도 (Precision)	반복 횟수 (Iters)	시퀀스 (Seq)	결과
Public WikiText 5D 연속 사전 학습	Public WikiText 5D	dense TP=2 PP=2 CP=2 DP=2, expert ETP=1 EP=4 EDP=2, PP=2	BF16 + TE + FlashAttention	20	—	16개 GPU에서 학습/저장/평가 통과
...

학습 로그를 통한 런타임 관찰 (Runtime Observations From Training Logs)

다음은 학습 로그에서 추출한 Rank-0 메모리 관찰 결과입니다. 이는 전체 성능 벤치마크는 아니지만, 왜 서로 다른 병렬성 레이아웃 (parallelism layouts)이 중요한지를 보여줍니다.

모델 / 레이아웃	기록된 Rank-0 메모리	시사점
DeepSeek-V2-Lite EP=8 BF16	약 42.6-43.2 GB 예약됨	EP 전용 베이스라인은 8-GPU 노드 하나에 여유롭게 적재됨
...

가장 강력한 실질적 교훈은 PP가 단순히 처리량 (throughput)을 조절하는 노브 (knob)가 아니었다는 점입니다. 하이브리드 Nemotron 모델의 경우, PP=2는 주요한 메모리 관리 도구였습니다. 이는 일부 레이아웃을 순수 TP+EP 방식보다 훨씬 더 안정적으로 만들어 주었습니다.

모델 참고 사항 및 학습 내용 (Model Notes And Learnings)

Public WikiText 16-GPU 5D 실행

• 데이터셋 (Dataset): Hugging Face Salesforce/wikitext, 설정 (config) wikitext-2-raw-v1.
• 토크나이저 (Tokenizer) 파일: Qwen/Qwen3-30B-A3B에서 가져온 로컬 Qwen3 토크나이저 파일.
• 노드 (Nodes): 2개.
• 총 GPU 수: 16개.
• 밀집 차원 (Dense dimensions): TP=2, PP=2, CP=2, DP=2.
• 전문가 차원 (Expert dimensions): ETP=1, EP=4, EDP=2, PP=2 적용.
• 전문가 (Experts): 총 128개, EP 샤드(shard)당 32개의 로컬 전문가.
• 제어 모델 (Control model): 4개 레이어, 은닉 크기 (hidden size) 256, FFN 512, 8개 어텐션 헤드 (attention heads), 2개 쿼리 그룹 (query groups).
• 정밀도/경로 (Precision/path): BF16, Transformer Engine, FlashAttention.
• 결과:
  • 1회차부터 20회차까지 학습 (trained iterations) 수행
  • 10회차 및 20회차에서 체크포인트 (checkpoints) 저장
  • 검증/테스트 (validation/test) 실행
  • 메모리 사용량은 낮게 유지되었으며, 표시된 랭크 (ranks)에서 예약된 메모리는 약 0.9 GB 미만임
• 중요 주의 사항: 신규 학습/저장/평가 (train/save/eval)는 정상 작동했으나, 노드 로컬 체크포인트 레이아웃 (node-local checkpoint layout)으로부터 재개 (resume)할 경우 로드 후 NCCL 타임아웃 (timeout)이 발생하며 이는 아직 해결되지 않은 상태임.

DeepSeek-V2-Lite

• HF 모델: deepseek-ai/DeepSeek-V2-Lite
• 가져온 체크포인트 (Imported checkpoint): 공개된 Hugging Face 스냅샷 (snapshot)으로부터 생성된 내부 Megatron 형식 체크포인트.
• 스냅샷/체크포인트 크기: 약 30G
• 공개 규모 (Public scale):
  • 총 파라미터 (Total parameters): 16B
  • 토큰당 활성화된 파라미터 (Activated parameters per token): 2.4B
• 중요 수정 사항: 모델이 27개 레이어를 가지고 있으므로 PP=2 설정 시 불균등한 레이어 분할이 필요함.
  • FIRST_PP_LAYERS=14
  • LAST_PP_LAYERS=13
• 검증된 내용: 체크포인트 가져오기 (checkpoint import), BF16, FP8 지연 스케일링 (delayed scaling), 전문가 병렬성 (expert parallelism), 텐서 병렬성 (tensor parallelism), 시퀀스 병렬성 (sequence parallelism), 지도 미세 조정 (supervised fine-tuning) 흐름, 그리고 지속적 사전 학습 (continued-pretraining) 흐름.

Nemotron-3-Nano

• 모델/레시피 (Model/recipe): nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
• 아키텍처 (Architecture): Mamba2-Transformer Hybrid Mixture of Experts (MoE).
• HF 토크나이저/소형 파일 (HF tokenizer/small files): 공개된 Hugging Face 모델 자산의 로컬 복사본.
• 공개 규모 (Public scale):
  • 총 파라미터 (Total parameters): 30B
  • 토큰당 활성화 파라미터 (Activated parameters per token): 약 3B
• causal-conv1d가 설치되지 않아 Mamba fast path를 비활성화했습니다.
• Transformer Engine CUDA 멀티캐스트 오버랩 (multicast overlap) 문제로 인해 TP=2 설정 시 UB_SKIPMC=1이 필요했습니다.
• 검증 내용: EP, TP, PP, CP 및 ETP 조합을 통한 Mamba2-Transformer Hybrid MoE 아키텍처에서의 Megatron 학습 검증.

DeepSeek-V4-Flash

• 대상 (Target): deepseek-ai/DeepSeek-V4-Flash
• 테스트 체크포인트 (Probe checkpoint): 내부 Megatron 형식의 개발용 체크포인트.
• 결과: 분산 Megatron 설정에는 도달했으나, 단일 8x96 GB 노드에서 전체 학습을 수행하기에는 용량이 부족했습니다.
• 검증 내용: 예상된 메모리 벽 (memory wall)에 부딪히기 전, 실행, 설정 구성, 분산 설정 및 MXFP8 경로 탐색을 검증했습니다.

직면한 문제 및 해결 방법 (Issues Faced And Fixes)

영역 (Area)	문제 (Issue)	해결 / 상태 (Fix / Status)
선택적 임포트 (Optional imports)	Bridge 레시피 임포트 시 diffusers, nvidia_resiliency_ext, megatron.energon과 같은 선택적 패키지를 불러옴	안전한 경우 더 좁은 범위의 레시피 로딩 및 작은 스텁 (stub) 사용
...

증명된 사항 (What Was Proven)

축 (Axis)	상태 (Status)
TP	DeepSeek-V2-Lite 및 Nemotron-3-Nano에서 증명됨
...

시사점 (Takeaways)

Megatron은 학습 토폴로지 (topology)를 명시적으로 만들어 주기 때문에 가치가 있습니다. 이번 작업은 "모델을 실행할 수 있는가?"에서 "대규모 모델이 분할, 학습, 체크포인트 저장 및 디버깅되는 방식을 제어할 수 있는가?"로 나아갔습니다.

주요 실무적 교훈은 다음과 같습니다:

• Megatron-Bridge 변환(conversion) 및 체크포인트 레이아웃(checkpoint layout)은 학습 시스템의 일부입니다.
• 분산 학습(distributed training)의 증명에는 순전파/역전파(forward/backward)뿐만 아니라 체크포인트 저장(checkpoint save)이 포함되어야 합니다.
• MoE(Mixture-of-Experts) 모델은 밀집 가중치(dense weights), 라우팅된 전문가(routed experts), 전문가 라우팅(expert routing), 그리고 옵티마이저 상태(optimizer state)에 대한 별도의 추론이 필요합니다.
• 컨텍스트 병렬성(Context parallelism)은 모델 경로(model-path)에 의존적입니다. 즉, 한 아키텍처에서는 작동하지만 다른 백엔드 경로에서는 실패할 수 있습니다.
• 파이프라인 병렬성(Pipeline parallelism)은 단순히 처리량(throughput)만을 위한 것이 아닙니다. 메모리 여유 공간(memory headroom)을 확보하는 가장 깔끔한 방법이 될 수 있습니다.
• 짧은 스모크 테스트(Short smoke runs)는 전체 인프라 경로를 검증하고 다음 병목 지점(bottleneck)을 드러낼 때 유용합니다.