NAVER Cloud가 NVIDIA Megatron Bridge로 멀티모달 사전 학습을 가속하는 방법

HyperCLOVA X는 한국어와 한국 문화에 특화된 NAVER Cloud의 foundation model 제품군으로, LLM과 멀티모달 모델을 아우릅니다. NAVER는 vision encoder에서 출발해 VLM(vision-language model)과 같은 멀티모달 LLM을 처음부터 직접 개발하고 있습니다. 이 글에서는 VLM이 대량의 멀티모달 토큰으로부터 핵심 능력을 습득하는 단계인 멀티모달 사전 학습(multimodal pretraining, MMPT)에 NVIDIA Megatron Bridge를 활용한 경험을 공유합니다. NVIDIA Megatron Bridge와 그 전신인 NeMo를 사용해 HyperCLOVA X SEED Omni 8B와 4B모델을 학습했으며, 현재는 MoE 기반 모델을 개발하고 있습니다. 이 과정에서 프로덕션 파이프라인의 여러 이슈를 해결하고 그 수정 사항을 코드베이스에 직접 기여해 왔으며, Megatron Bridge 팀과 긴밀히 협업하고 NVIDIA Korea 솔루션 아키텍트 팀의 지원을 받았습니다.

네 가지 최적화 영역

프로파일링을 통해 멀티모달 학습의 주요 병목이 되는 네 가지 영역을 식별하고, 이 영역들에 최적화 노력을 집중했습니다:

1. 스토리지와 I/O: 원시 멀티모달 데이터는 텍스트보다 훨씬 크기 때문에, 읽기만으로도 학습이 지연될 수 있습니다.
2. 샘플 처리와 sequence packing: 단순한 packing은 이미지를 시퀀스 경계에서 잘라내며, 길이가 같은 pack이라도 vision encoder 부하는 서로 다릅니다.
3. Vision Encoder: 본 설정에서는 end-to-end로 학습되므로 모든 학습 step의 critical path에 위치합니다.
4. VLM을 위한 분산 학습 최적화: 텍스트 전용 LLM을 위해 설계된 parallelism 기법이 packing된 멀티모달 시퀀스에서는 아직 올바르게 동작하지 않았습니다.

첫 번째 영역은 학습용 멀티모달 data loader인 NVIDIA Megatron Energon을 기반으로 합니다. 두 번째 영역은 Energon과 Megatron Bridge 양쪽에 걸쳐 있으며, 마지막 두 영역은 Megatron Bridge의 모델과 parallelism 레이어에 해당합니다.

그림 1에서 보듯이, 이 최적화들을 통해 30B MoE VLM 기준으로 학습 throughput(tokens/sec/GPU)이 vanilla Megatron Bridge baseline(100%) 대비 150.2%까지 향상되었습니다. 그림 1은 앞의 세 영역만을 다룹니다. 섹션 4의 작업은 속도보다는 correctness에 관한 것이었고, 이 모델 규모에서는 NVIDIA B200 클러스터에서 expert parallelism만으로 충분했기 때문입니다.

모든 실험은 NAVER Cloud에서 실제 학습에 사용하는 프로덕션 데이터로 수행했습니다. 100개가 넘는 데이터 소스에서 나온 2조 개(2T) 이상의 토큰으로 구성되어 있어, 현실적인 조건에서 데이터 확장성까지 검증할 수 있었습니다. 하나의 샘플 안에서도 이미지 개수와 해상도가 크게 달라지는데, 텍스트만 있는 샘플부터 이미지가 50장이 넘는 샘플, 8K 이상의 해상도까지 극단적인 사례를 포함합니다. 실험 모델은 Nemotron-3-Nano 스타일 아키텍처의 30B MoE로, vision encoder와 LLM을 결합한 표준 VLM 구조이며 vision encoder를 포함한 모든 파라미터를 학습 중에 업데이트합니다. 학습은 멀티 노드 B200 클러스터에서 expert parallelism(EP) 8, tensor/pipeline/context parallelism(TP/PP/CP) 1, gradient accumulation 128 설정으로 수행했습니다. vanilla Megatron Bridge baseline도 동일한 설정을 공유합니다. 즉, 같은 8K 고정 시퀀스 길이의 packing된 시퀀스, greedy bin-packing, Energon의 표준 WebDataset 포맷을 사용하며, 각 최적화는 이 baseline 위에 누적 적용됩니다.

각 영역을 살펴보기 전에, 그림 2는 데이터가 NVIDIA Megatron Energon을 거쳐 Megatron Bridge까지 흐르는 과정을 보여줍니다. 서로 다른 flavor의 데이터셋이 하나의 샘플 스트림으로 블렌딩되고, 각 샘플은 TaskEncoder를 통과하면서 text token과 이미지 텐서로 변환되어 학습 시퀀스로 packing됩니다. TaskEncoder는 Megatron Bridge의 핵심 커스터마이징 지점으로, 이를 상속받아 워크로드에 맞게 수정했습니다. 섹션 1은 이 흐름의 저장과 로딩 단계를, 섹션 2는 TaskEncoder 단계와 그 직후에 적용되는 reordering을, 섹션 3과 4는 학습 step에 해당하는 vision encoder와 분산 학습을 다룹니다.

1. 스토리지와 I/O

효율적인 데이터 로딩: JSONL offset indexing과 prefetching

NVIDIA Megatron Energon은 안정적인 WebDataset 기반 shard 포맷에 더해 최근 JSONL 지원을 추가했습니다. 파이프라인은 이미 JSONL로 데이터를 관리하고 있었고, 100개 이상의 소스에 걸친 2조 개 이상의 토큰을 shard로 변환하면 상당한 스토리지·파이프라인 오버헤드가 발생하기에 JSONL 경로를 유지했습니다. 다만 새로 추가된 경로인 만큼 WebDataset 로딩 수준까지는 아직 최적화되어 있지 않았고, zip 패키징, offset indexing, 비동기 prefetching으로 그 격차를 메웠습니다. 그림 3은 이 레이아웃을 WebDataset 포맷과 비교해 보여줍니다.

zip 패키징은 수백만 건의 작은 파일 읽기를 피하게 해 주지만, 파일을 통합하면서 개별 이미지를 찾는 속도가 느려졌습니다. 그래서 각 이미지의 정확한 offset과 크기를 JSONL 메타데이터에 기록해, 그림 4처럼 전체 스캔 없이 파일에 직접 접근할 수 있게 했습니다. 여기에 비동기 이미지 prefetching이 남은 읽기 지연을 연산 뒤로 숨겨 오버헤드를 감춰 줍니다.

기타 I/O 개선

데이터셋 포맷 외에도 CPU 메모리 누수 수정, 샘플당 최대 이미지 수 제한 같은 작은 변경들이 안정성과 throughput을 더했습니다. CPU 메모리가 줄면서 data loader worker를 더 많이 띄울 수 있었습니다. 이 스토리지·I/O 변경들을 합치면 end-to-end throughput이 2.5% 개선됩니다.

2. 샘플 처리와 Sequence Packing

샘플이 로드되면 TaskEncoder가 각 샘플을 text token과 이미지 텐서로 변환하고 학습 시퀀스로 packing합니다. 멀티모달 packing에는 텍스트 전용 학습에는 없는 제약이 있습니다. 각 이미지가 하나의 시퀀스 안에 온전히 들어가야 한다는 점입니다. 아래에서는 Energon의 packing hook 위에 구현되어 baseline에 포함된 greedy bin-packing과, 그 위에 구축한 최적화인 global batch 단위 pack reordering(vision encoder 연산 균형화)을 차례로 설명합니다. 그림 5는 packing 방식과 그것이 남기는 불균형을, 그림 6은 reordering을 보여줍니다.

Greedy bin-packing으로 LLM 연산 균형 맞추기

샘플 길이는 자릿수가 달라질 만큼 큰 폭으로 차이가 나므로 모든 샘플을 패딩하면 연산이 낭비되고, LLM에서 흔히 쓰는 이어 붙인 뒤 시퀀스 길이에서 자르는 방식은 이미지를 중간에서 잘라낼 수 있습니다. greedy bin-packing은 각 고정 길이 시퀀스를 통째 샘플로 채워 패딩을 최소화하며, 샘플 단위로 packing하기 때문에 이미지가 시퀀스 경계에서 잘리는 일이 없습니다. 8K 고정 시퀀스 길이 기준으로 worker별 LLM 부하가 균형을 이루며, packing 효율(전체 토큰 대비 유효 토큰 비율)은 99.4%에 달합니다.

Reordering으로 vision encoder 연산 균형 맞추기

greedy bin-packing은 LLM 연산은 균형 있게 만들지만 vision encoder는 그렇지 못합니다. 길이가 같은 pack이라도 담고 있는 image token 수가 다르기 때문입니다. pack이 도착하는 순서대로 data parallel(DP) worker에 배정되면 매 step이 가장 이미지가 많은 pack을 쥔 worker에 발목 잡히고, 이 straggler 효과는 data parallel 폭이 커질수록 심해집니다.

이를 없애기 위해 pack을 image token 수 기준으로 DP worker에 재분배합니다(그림 6). 이 reordering은 통째 pack을 배치 수준에서 다루며, data loader와 forward step 사이의 얇은 iterator wrapper로 구현되어 하위의 학습 루프는 전혀 건드리지 않습니다. reordering 풀은 정확히 하나의 global batch(DP × gradient accumulation 개의 pack)입니다. pack은 microbatch와 worker 경계는 넘나들지만 global batch 경계는 절대 넘지 않으므로, 각 optimizer step이 보는 샘플 집합은 동일하게 유지됩니다.

배정은 모든 worker가 한 번의 메타데이터 all-gather로부터 동일하게 계산합니다. 교환되는 것은 pack 자체가 아니라 image token 수와 텐서 shape 정보로 킬로바이트 수준에 불과하며, encoder 연산 비용 모델을 기준으로 worker 간 균형을 맞추고 전송량을 줄이기 위한 추가 기법도 배정 로직에 담았습니다. 그 후 pack은 한 번의 all-to-all로 교환되며, 이미 올바른 worker에 있는 pack은 통신을 완전히 건너뜁니다.

핵심은 이 reordering이 드러나지 않게 설계되었다는 점입니다. 다음 global batch의 메타데이터 교환과 배정 계산은 background thread에서 돌아가고, pack 교환은 전용 CUDA stream에서 forward·backward 연산과 겹쳐 수행되며, 각 worker는 이미 로컬에 있는 pack으로 첫 microbatch를 시작합니다. Comunication 오버랩의 효과를 정량화하기 위해 교환을 기다리는 동기 방식도 측정했는데, 균형화만으로는 1.7%의 이득에 그친 반면 비동기 방식은 이 섹션의 총 이득을 13.3%까지 끌어올렸습니다.

3. Vision Encoder

일반적인 VLM에서 모든 이미지는 vision encoder를 거쳐 image token이 되고, text token과 함께 LLM에 입력됩니다. 그림 7은 NAVER의 프로덕션 모델 중 하나인 HyperCLOVA X 32B Think에서 이 구조를 보여줍니다. 이 글의 실험은 30B MoE VLM으로 진행했지만 전체 구조는 동일합니다. encoder는 LLM과 함께 end-to-end로 학습되므로 encoder의 효율이 step 시간에 직접 영향을 줍니다. 개선은 두 단계로 진행했습니다. 먼저 Hugging Face encoder에 activation recomputation을 적용하고, 이어서 그림 8처럼 네이티브 Megatron으로 재구현했습니다.

Activation recomputation을 적용한 Huggingface encoder

처음에는 사내 vision encoder를 Hugging Face 구현 그대로 연결했습니다. Megatron Bridge가 전용 Megatron 모델 없이도 이를 지원하기 때문에, encoder 아키텍처가 아직 진화 중이던 시기에 가장 빠른 학습 경로였습니다. 다만 이 방식에서는 encoder가 Megatron parallelism 바깥에 놓입니다. 고해상도 이미지 활성화의 전체 시퀀스가 GPU 한 장에 들어가야 해서 학습 가능한 시퀀스 길이에 상한이 생겨 activation recomputation을 이용하여 그 상한을 풀어 주었습니다. forward 연산이 추가되지만 확보된 메모리 덕분에 시퀀스 길이를 8K에서 16K로 두 배 늘릴 수 있었고, packing이 길어진 시퀀스를 빽빽하게 채우면서 최종적으로 5.4%의 throughput 이득을 얻었습니다.

Native Megatron vision encoder

다음으로 동일한 vision encoder를 Megatron 블록으로, Megatron Bridge의 비전 트랜스포머 레퍼런스를 참고해 네이티브로 재구현했습니다. 이로써 Hugging Face 경로가 포기했던 것들, 즉 Megatron의 최적화된 kernel과 LLM과 동일한 방식의 vision encoder 병렬화를 되찾았습니다. 네이티브 구현은 Megatron의 parallelism를 그대로 상속받는 반면, Hugging Face encoder를 분산하려면 별도의 추가 작업이 필요합니다.

표 1은 두 설정 모두에서의 이득을 보여줍니다. 메인 설정인 TP=1, EP=8에서는 네이티브 encoder가 throughput을 23.9% 개선하며(그림 1에서 가장 큰 단일 개선 항목), GPU 메모리는 사실상 변화가 없습니다. TP=1에서는 두 backend 모두 encoder를 복제하므로 이득이 Megatron의 최적화된 kernel에서 나오기 때문입니다. TP=4, EP=2에서는 네이티브 encoder가 weight와 activation을 TP rank에 추가로 shard하면서 이득이 50.3%로 커지고 메모리도 뚜렷하게 줄어듭니다.

4. VLM을 위한 분산 학습 최적화

Megatron의 분산 학습 전략은 텍스트 전용 LLM을 중심으로 설계되었지만, VLM은 매 step에 연산 특성이 크게 다른 vision encoder를 더합니다. 메인 설정은 EP만 사용(TP/PP/CP=1)하므로 이 섹션의 결과는 앞의 대표 수치에는 포함되지 않지만, 시퀀스 길이와 모델 규모가 커질수록 CP 같은 parallelism 전략이 필수가 되기에 이를 필요로 하는 팀들을 위해 경험을 공유합니다.

packing된 멀티모달 워크로드에 CP를 적용하는 과정에서, Vision DP when CP가 packing된 시퀀스를 아직 올바르게 처리하지 못하고 gradient backpropagation에서 오류가 발생하는 것을 발견했습니다. NVIDIA Megatron Bridge 팀과 긴밀히 협업해 문제를 수정하고 PR #4784로 업스트림에 반영했습니다.

Context Parallelism with Sequence Packing과 Vision DP when CP

CP로 확장하면 GPU당 throughput은 필연적으로 줄어듭니다. CP 그룹이 커질수록 data parallel 폭이 좁아지고 통신이 늘어나기 때문입니다. 측정에서는 CP=1(16K)에서 CP=2(24K)로 확장할 때 완화 조치 이전 기준 38%의 비용이 발생했습니다. sequence packing이 늘어난 컨텍스트를 빽빽하게 채워 이를 완화하며, 그 위에 Vision DP when CP를 활성화했습니다. 이는 이미지 처리를 CP rank들에 분할해 각 이미지가 CP 그룹당 한 번만 인코딩되도록 하는 Megatron Bridge의 기존 기능으로, 그림 9에 나타나 있습니다. 표 2에서 보듯 이 기능만으로 CP=2 실행이 20% 빨라져 전환 비용의 상당 부분을 회복합니다.

요약 및 향후 계획

이 글에서는 프로덕션 워크로드 프로파일링으로 식별한 네 가지 영역, 즉 스토리지와 I/O, sequence packing, vision encoder, 분산 학습에 걸쳐 NVIDIA Megatron Bridge 기반 멀티모달 사전 학습을 최적화한 과정을 공유했습니다. 이 최적화들을 통해 30B MoE VLM에서 학습 throughput이 vanilla Megatron Bridge 대비 150.2%로 향상되었습니다. NAVER는 이 기반을 Omni 스타일 멀티모달 학습으로 확장하고 있으며, 추가 결과도 공유할 예정입니다. 긴밀히 협업해 주신 NVIDIA Megatron Bridge 팀과 NVIDIA Korea 솔루션 아키텍트 팀에 감사드립니다.

참고자료

NAVER Tech Blog: 속도와 비용의 열쇠: 멀티모달 AI 학습을 갉아먹는 ‘이미지 토큰 불균형’ 이야기