Vision Language Model 학습 속도 개선

들어가며

Vision Language Model (이하 VLM) 학습은 생각보다 LLM 학습과 크게 다르지 않을 것이라고 믿었습니다. LLM에 Vision Encoder만 붙인 구조이고, LLM 학습 방법론도 이미 잘 정립되어 있으니까요.

하지만 막상 학습을 시작해 보니 현실은 달랐습니다. LLM에서는 크게 신경 쓰지 않던 요소들이 VLM에서는 학습 효율을 크게 좌우했고, 그 차이를 이해하기까지 적지 않은 시행착오를 겪었습니다.

이번 글에서는 Gravity-16B-A3B-Base^[1]를 기반으로 Gravity-VL-16B-A3B를 개발하며 얻은, VLM 학습 효율에 대한 경험과 교훈을 공유하고자 합니다.

VLM 학습 속도 개선을 위한 발버둥

LLM 학습은 어느 정도 최적화 레시피가 공유되어 있는 반면, VLM 학습에서는 LLM을 넘어서는 요소들을 어떻게 최적화할지에 대한 이야기가 충분히 정리되어 있지 않습니다.

물론 VLM 학습도 LLM 학습에서 사용하는 테크닉들을 기반으로 시작합니다. 하지만 LLM 쪽 최적화를 그대로 적용해두었더라도, 막상 VLM 학습을 돌려보면 Token Per Second(TPS)^[2]가 아쉽게 나오는 경우가 많습니다. 분명 같은 LLM backbone을 쓰고 있는데, 고작 400M 파라미터 규모의 Vision Transformer (이하 ViT) 하나를 붙였다는 이유만으로 학습 속도가 크게 떨어지는 것을 경험하게 됩니다.

Kimi K2.5 technical report^[3]에서도 비슷한 문제의식을 엿볼 수 있습니다. 해당 report는 K2.5가 K2의 parallel strategy를 자연스럽게 계승했으며, 그 결과 text-only training 대비 90% 수준의 multimodal training efficiency를 달성했다고 설명합니다.

K2.5 seamlessly inherits the parallel strategy of K2, achieving a multimodal training efficiency of 90% relative to text-only training.

얼핏 보면 꽤 높은 수치입니다. 하지만 다시 말하면, Vision Encoder를 붙이고 Multimodal Input을 처리하는 것만으로도 text-only 대비 약 10%의 효율 손실이 발생했다는 뜻이기도 합니다.

여기서 주목해야 할 점은 backbone LLM의 규모입니다. Kimi K2.5가 기반으로 하는 Kimi K2는 1T 파라미터 (32B active 파라미터) 규모의 MoE 모델입니다. 파라미터 수 관점에서만 보면, 400M 규모의 ViT는 1T LLM backbone의 0.04% (active 기준 1.25%)에 불과합니다. 그런데 파라미터 수는 0.04% 늘었을 뿐인데 training efficiency는 10% 줄었다고 보면, 여간 섭섭한 수치가 아닙니다. 하물며 16B LLM backbone 위에서 실험하는 저희 입장에서는 어떨까요? 16B-A3B 모델의 경우, ViT의 파라미터 비중은 2.5% (active 기준 13.3%)입니다. 시작하기도 전에 걱정이 앞설 수밖에 없습니다…

Gravity-VL-16B-A3B 학습 실험의 기본 세팅은 다음과 같았습니다.

• Hardware: B200 노드 1개
• Sequence length: 32,768
• Global batch size: 약 2.62M tokens/step
  • microbatch와 gradient accumulation은 각 실험에서 TPS가 최대가 되도록 조정
• Text : Vision-Language (이하 VL) 비율: 8 : 2
• VL 데이터 구성:
  • 90%: 이미지 해상도 2 megapixel (이하 MP) 이내 데이터
  • 10%: 최대 16MP 고해상도 데이터

안타깝게도 별 생각 없이 돌려본 첫 test run의 결과는 처참했습니다.

ViT를 붙였을 뿐인데, 학습 속도는 거의 3분의 1로 줄어들었습니다. 모델에 눈을 붙였더니 (attach eyes), 정말로 눈을 붙여버린 (take a nap) 셈입니다.

이 문제 해결 과정이 다소 고통스러웠던 이유는 tech report나 논문에서 이 부분에 관한 디테일을 접하기가 쉽지 않기 때문이기도 합니다. 간혹 VL 학습의 TPS 최적화를 위한 단편적인 테크닉들이 소개되기는 하지만, 정작 궁금한 모든 부분을 시원하게 긁어주기에는 설명이 부족한 경우가 많습니다.

이 글은 제가 스스로 가려웠던 부분을 긁어보기 위해 작성하는, 말하자면 효자손 같은 글이기도 합니다. 구체적으로는 다음 두 가지를 정리해보려 합니다.

• VLM 학습의 TPS를 낮추는 요인들
• 그 TPS를 개선하기 위해 저희가 시도한 여러 방법과 결론

VLM 학습은 어떤 식으로 되는가 (LLaVA 방식)

본론에 들어가기 전에, 이 글에서 다루는 VLM 구조를 빠르게 짚고 가겠습니다. 가장 보편적인 LLaVA^[4] 계열 구조를 기준으로 설명합니다.

LLaVA 스타일의 VLM 구조는 크게 세 덩어리로 나눌 수 있습니다.

1. Vision Encoder
   • 이미지를 patch 단위로 쪼갠 뒤, 각 patch를 visual feature로 인코딩합니다. 보통 ViT 기반 모델을 사용합니다.
2. Projector (Connector)
   • visual feature를 LLM의 embedding 공간으로 매핑합니다. 일반적으로 MLP가 많이 사용됩니다.
3. LLM
   • projected visual token과 text token을 이어 붙인 시퀀스를 입력으로 받아 autoregressive 방식으로 학습합니다.

학습 시 forward는 대략 다음과 같은 흐름으로 진행됩니다.

즉, 이미지 한 장은 ViT를 거쳐 수백–수천 개의 visual token이 되고, 이 visual token들이 text token 사이에 합류해 하나의 긴 시퀀스를 이룹니다. 여기서부터 이미 두 가지 부담이 보입니다. 이미지가 많아질수록 LLM 토큰 시퀀스가 길어지고, 거기에 ViT 연산까지 추가됩니다.

무엇이 VLM 학습을 느리게 하는가

이를 실제 학습 파이프라인 관점에서 조금 더 풀어보면, 병목은 크게 세 군데에서 드러납니다. 이미지를 읽고 전처리하는 단계, vision encoder를 통과시키는 단계, 그리고 분산 학습에서 이 부하가 GPU마다 불균형하게 분산되는 단계입니다.

1. Image Read / Decoding의 병목

이건 GPU 밖, 즉 CPU·I/O 쪽 문제입니다. 이미지 데이터는 텍스트와 다르게 다음 과정을 추가로 거쳐야만 비로소 텐서가 됩니다:

• 스토리지, 보통 NFS나 object storage에서 읽기
• JPEG/PNG 디코딩
• Resize, normalize 같은 전처리

이 과정에서 병목이 생기면, GPU가 다음 배치를 기다리며 놀게 됩니다. 특히 이미지 한 장이라도 이런저런 이유로 읽는 데 오래 걸리면 해당 GPU 프로세스의 학습에 제동이 걸리게 되고, 그렇게 한 GPU에 제동이 걸리면 학습 전체가 함께 멈추게 됩니다.

2. 덜 최적화된 Vision Encoder 연산

LLM backbone은 보통 수 B–수십 B 파라미터 규모임에도 여러 최적화 덕분에 꽤 효율적으로 학습됩니다. 반면 vision encoder는 상대적으로 작은 모델인데도 의외로 step time을 많이 잡아먹습니다. 이유는 단순합니다.

• 이미지 한 장이 수백–수천 개의 patch로 펼쳐지고, 고해상도 이미지의 경우 patch 수는 몇 만 개까지 갑니다.
• ViT의 attention 연산량은 patch 수에 대해 $O(N^2)$입니다.
• LLM 쪽 최적화, 예를 들어 FlashAttention이나 fused kernel 등은 잘 챙기면서도, 정작 ViT는 reference 구현에 가깝게 남겨두는 경우가 많습니다.

즉, 상대적으로 덜 최적화된 모듈에 의외로 많은 patch가 흘러들어간다는 것이 핵심입니다.

3. 이미지가 각 프로세스에 골고루 분산되지 않으면 생기는 문제

이 부분은 VLM 학습에서 가장 음흉한 병목 중 하나입니다. Data parallel 환경에서 각 GPU는 매 step마다 gradient를 동기화합니다. 그런데 ViT forward 시간은 대체로 해당 GPU에 들어온 이미지, 더 정확히는 patch 양에 비례합니다.

문제는 데이터를 단순하게 나눠 담으면 이미지 부하가 특정 GPU에 쏠릴 수 있다는 점입니다. 어떤 GPU에는 이미지가 10장 들어가고, 어떤 GPU에는 이미지가 하나도 들어가지 않는 식입니다. 혹은 특정 GPU에 고해상도 이미지가 집중될 수도 있습니다.

이렇게 되면 매 step의 all_reduce는 가장 느린 GPU를 기다리게 됩니다. 결국 이미지가 몰린 GPU 하나 때문에 전체 학습 속도가 그 GPU에 맞춰지는 셈입니다. GPU 64장을 붙여놔도 그중 한 장이 이미지 폭탄을 맞으면, 나머지 63장은 멀뚱멀뚱 기다리게 됩니다. 자원 낭비가 꽤 심해지는 지점입니다.

LLM Optimization과 다른 부분

요약하면, VLM 학습의 병목은 LLM 학습에서 챙기던 것과 결이 다릅니다.

그래서 LLM 학습 레시피를 그대로 가져온다고 해도, 위 문제들이 자동으로 해결되지는 않습니다.

VLM 학습 속도 개선을 위한 여러 시도

(1) Dataloader for Fast Image Decoding

가장 먼저 손봐야 하는 건 보통 dataloader입니다. 비싼 GPU를 I/O 병목 때문에 못 쓰는 것만큼 아까운 것이 없지요.

이를 해결하기 위해 저희는 두 가지 방법을 적용했습니다. Node-local NVMe Staging과 Async Prefetch가 바로 그것입니다.

Storage Locality: NFS → Node-local NVMe Staging^[5].

처음에는 이미지를 공용 디스크(NFS)에서 직접 로딩하도록 구성했습니다. 그런데 이 방식은 단순히 느린 것을 넘어, 학습 속도를 꽤 불안정하게 만들었습니다.

분산 학습에서는 수많은 GPU와 data loader worker가 동시에 작은 이미지 파일을 열고, metadata를 조회하고, 읽는 작업을 반복합니다. 즉, many-small-files, metadata-heavy, many-client 접근 패턴이 만들어집니다. 이런 패턴은 NFS 같은 원격 파일시스템에서 특히 불리합니다.

• 높은 per-request latency: 파일을 열고, metadata를 확인하고, 데이터를 읽는 과정마다 네트워크 왕복 비용이 붙습니다.
• 동시성 하에서의 throughput 저하와 tail latency 증가: 여러 GPU가 동시에 작은 파일 요청을 쏟아내면 서버와 네트워크가 쉽게 경합 지점이 되고, 일부 read가 유난히 오래 걸리는 tail latency가 커집니다.
• 간헐적 stall 또는 wedged read: 네트워크 순간 지연, 서버 부하, 마운트 상태 등의 이유로 특정 read가 정상적으로 반환되지 않거나 매우 오래 걸리는 경우가 생길 수 있습니다. 이때 해당 이미지를 기다리던 GPU 전체가 멈추고, 결국 매 step의 gradient synchronization에서 나머지 GPU까지 함께 기다리게 됩니다. 이를 방어하는 로직으로, 특정 이미지 로딩이 지정해놓은 timeout을 넘어가면 skip하게끔 할 수 있습니다.

결국 이미지 read 단계가 파이프라인 전체에서 가장 느리고, 가장 예측하기 어려운 병목이 됩니다. 아무리 decode나 prefetch를 최적화해도, 그 앞단에서 데이터가 제때 흘러나오지 않으면 GPU는 다음 batch를 기다릴 수밖에 없습니다.

저희가 선택한 해결책은 데이터를 연산 가까이로 옮기는 것이었습니다. 학습 시작 전에 필요한 이미지 데이터를 각 노드의 local NVMe cache로 staging하고, 학습 중에는 NFS가 아니라 로컬 디스크에서 이미지를 읽도록 바꿨습니다.

그 결과 다음과 같은 효과가 있었습니다.

• read 요청이 네트워크를 타지 않고 로컬 고IOPS·저지연 스토리지에 직접 닿습니다.
• NFS의 metadata / network 경합이 training hot path에서 빠지면서 throughput이 더 높고 안정적으로 바뀝니다.
• tail latency와 stall이 줄어들어, timeout 같은 방어 로직에 의존해야 하는 상황 자체가 줄어듭니다.

Async Prefetch.

NVMe staging이 “데이터가 어디에 있는가”를 해결했다면, async prefetch는 “데이터를 언제 준비하는가”를 해결합니다.

이미지를 로컬 NVMe에서 빠르게 읽어오더라도, decode, resize, tokenization, packing 같은 준비 과정은 여전히 CPU에서 수행됩니다. 이 과정을 학습 loop 안에서 동기적으로 처리하면, GPU는 매 step마다 다음 batch가 준비될 때까지 기다리게 됩니다. 결국 step time이 compute + data preparation에 가까워집니다.

Async prefetch는 이 준비 과정을 GPU 연산과 겹치게 만듭니다. GPU가 현재 step을 계산하는 동안 background worker가 다음 batch를 미리 준비해 buffer에 채워두고, 학습 loop는 준비된 batch를 꺼내 쓰기만 합니다. 잘 동작하면 data preparation 시간이 compute 뒤에 숨어, step time은 compute + data preparation이 아니라 max(compute, data preparation)에 가까워집니다.

구체적으로는 이미지 decode를 multi-thread로 병렬화하고, prefetch buffer에 여러 batch를 미리 쌓아 tail latency를 흡수합니다. 또한 timeout 방어를 넣어 특정 이미지의 read/decode가 오래 걸리면 skip하거나 fallback 처리해, 문제 이미지 한 장이 전체 학습을 멈추지 않도록 했습니다.

결국 NVMe staging이 read latency를 줄이고, async prefetch가 decode와 batch 준비를 compute와 overlap시키면서, 데이터 파이프라인이 학습 hot path에서 차지하는 비중을 크게 줄일 수 있었습니다.

(2) Flash Attention & Vision Encoder Compile & Activation Checkpointing

이번에는 ViT 자체를 최적화하는 작업입니다. LLM 학습 시 흔히 적용하는 최적화들이지만, ViT에도 잊지 않고 적용해볼 만하다는 의미에서 다시 정리해보겠습니다.

Flash Attention^[6].

ViT attention backend는 생각보다 큰 차이를 만듭니다. 기본 eager attention은 QKᵀ로 만들어지는 [heads, L, L] attention score matrix를 HBM에 저장한 뒤, softmax와 V multiplication을 이어서 수행합니다. 이 과정에서 L² 크기의 intermediate를 계속 읽고 쓰게 되므로, sequence length가 커질수록 memory traffic과 메모리 사용량이 빠르게 증가합니다.

FlashAttention은 기본적으로 이 score matrix를 HBM에 materialize하지 않습니다. 대신 Q, K, V를 tile 단위로 on-chip SRAM에서 처리하고, online softmax를 사용해 attention을 하나의 fused kernel 안에서 계산합니다. 덕분에 attention의 memory traffic을 크게 줄이고, bandwidth-bound에 가까운 연산을 더 compute-bound하게 만들 수 있습니다.

이 차이는 ViT에서 특히 중요합니다. ViT는 이미지 patch들 사이에 full bidirectional attention을 사용하고, patch 수 L은 해상도에 따라 빠르게 늘어납니다. 따라서 max pixel이 커질수록 eager attention은 memory bandwidth와 device memory를 강하게 압박합니다.

저희는 ViT attention backend를 eager에서 FlashAttention, 구체적으로는 Blackwell/B200 세대의 FA4 kernel로 교체했습니다. 여기에 바로 이어서 설명할 vision encoder compile까지 함께 적용한 결과, ViT를 포함한 학습 throughput이 11,800 tokens/s/GPU에서 22,200 tokens/s/GPU로 증가했습니다. 약 +88% 개선으로, 저희 세팅에서 가장 큰 폭의 최적화였습니다.

Compile^[7].

torch.compile은 PyTorch 코드를 graph로 잡아 최적화하고, 일부 연산을 fuse해 Python overhead와 kernel launch overhead를 줄이는 방식입니다. Vision encoder에 적용하면 ViT block 내부의 반복적인 연산을 더 효율적으로 실행할 수 있어 TPS 개선을 기대할 수 있습니다.

다만 이미지 input shape이 다양하면 shape이 바뀔 때마다 recompile이 발생할 수 있습니다. 이 경우 compile 자체가 병목이 될 수 있으므로 dynamic=True로 설정했습니다. dynamic=True는 일부 dimension을 symbolic하게 다루어, shape 변화가 있어도 같은 compiled graph를 최대한 재사용하려는 옵션입니다.

Activation Checkpointing^[8].

Activation checkpointing은 forward 중간 activation을 저장하는 대신, backward 시 필요한 부분을 다시 계산하는 방식입니다. 즉, 메모리를 아끼는 대신 compute를 더 쓰는 memory-compute trade-off입니다.

Vision encoder에 적용하면 고해상도 이미지나 큰 ViT를 사용할 때 activation memory를 줄이는 데 도움이 됩니다. patch 수가 많아질수록 ViT의 activation memory도 커지기 때문에, checkpointing을 통해 더 큰 batch size나 더 높은 resolution을 사용할 여지를 만들 수 있습니다.

적용 방식은 크게 두 가지입니다.

• Full checkpointing: 대부분의 block에 적용합니다. 메모리 절감 효과는 크지만, recomputation이 많아져 TPS가 떨어질 수 있습니다.
• Selective checkpointing: 메모리 사용량이 큰 일부 block에만 적용합니다. 메모리 절감 폭은 작지만, 속도 손실을 줄이면서 타협점을 찾기 좋습니다. 어떤 operation에 activation checkpointing을 적용할지 직접 조절할 수 있기 때문에, TPS를 최대한 지키면서 메모리를 줄여야 한다면 시간을 들여 최적의 checkpointing 범위를 찾아볼 만합니다.

물론 checkpointing은 기본적으로 속도를 일부 희생해 메모리를 아끼는 기법입니다. 하지만 절약한 메모리 덕분에 microbatch 크기를 키울 수 있다면 이야기가 달라집니다. recomputation overhead가 추가되더라도, 더 큰 microbatch로 GPU utilization이 좋아지거나 gradient accumulation step을 줄일 수 있다면 end-to-end TPS는 오히려 상승할 수 있습니다.

(3) Batchwise Modality Segregation

text-only 데이터와 VL 데이터를 함께 학습할 때, 이 둘을 같은 microbatch에 섞으면 효율이 떨어집니다. 이유는 앞서 본 이미지 쏠림 문제와 같은 메커니즘입니다.

batch 단위 gradient backprop은 그 배치에서 가장 느린 샘플이 끝나기를 기다립니다. text-only 데이터는 처리가 빠른데, 같은 배치에 VL 샘플이 끼어 있으면 빠른 text 데이터가 느린 VL 데이터를 하염없이 기다리게 됩니다.

해결책은 단순합니다. text만 들어가는 batch와 VL만 들어가는 batch를 따로 구성하는 것입니다. DeepSeek-VL^[10]에서도 이와 비슷한 시도를 했습니다. DeepSeek-VL에 따르면, 이렇게 batch를 구분해서 학습을 하면 속도는 빨라지고 성능에는 차이가 없었다고 합니다.

그러나 여기서 한발 더 나아가서 gradient accumulation을 잘 조정하면, global batch 관점에서는 mixed batch 학습과 완전히 동일한 학습을 재현할 수 있습니다.

예를 들어 text-only : VL = 4 : 1 비율을 원한다면 다음과 같이 batch를 구성할 수 있습니다.

grad_accum = 2
text-only batch size = 4
VL-only batch size = 1

이렇게 하면 modality를 batch 단위로 분리하면서도, global batch 기준으로는 원하는 배치 조합을 그대로 만들 수 있습니다. 빠른 text 배치는 빠르게, 느린 VL 배치는 그것대로 — 서로 발목 잡지 않도록 하는 것이죠.

이런 식으로 text-only batch를 앞에 배치해두면 위에서 설명한 이미지 async prefetch로부터 더 큰 이득을 볼 수 있기도 합니다. text-only batch가 학습을 하는 동안 이미지 batch를 구성함으로써, 이미지 로딩에 걸리는 시간을 학습 시간과 겹쳐 효과적으로 숨길 수 있게 됩니다.

(4) max pixel Curriculum

마지막은 데이터 자체의 해상도를 다루는 curriculum 전략입니다^[11].

실제 VLM 학습 데이터에는 고해상도 이미지가 얼마나 포함되어 있을까요?

Figure 8은 대표적인 공개 VLM 학습셋인 FineVision^[12]과 LLaVA-OneVision^[13] 데이터셋에서 각각 10K 샘플을 뽑아, 이미지 해상도 분포를 정리한 것입니다. 정말 다행히도, 대부분의 이미지는 2MP보다 작습니다. 저희는 정보 밀도가 높은 고해상도 이미지까지 모델이 잘 이해하기를 원했기 때문에, 그런 데이터를 의도적으로 신경 써서 준비했습니다. 그럼에도 불구하고 학습 데이터의 대부분은 상대적으로 작은 max_pixel 안에서 충분히 처리할 수 있다고 볼 수 있습니다.

앞서 설명했듯이, 해상도가 높아질수록 ViT에 걸리는 부담은 선형을 넘어 초선형적으로 증가합니다. ViT는 이미지 내부 patch들 사이에 full attention을 적용하기 때문에, patch 수를 p라고 하면 MLP 계열 연산은 대체로 $O(p)$로 증가하지만 attention 연산은 $O(p^2)$로 증가합니다. 따라서 16MP 이미지 한 장이 들어간 GPU가 저해상도 이미지 수십~수백 장을 받은 GPU보다 더 느릴 수 있습니다.

고해상도 이미지가 일부 끼게 된다면 결국 매 step의 all_reduce는 이 “이미지 폭탄”을 맞은 GPU 하나를 기다리게 되고, 나머지 GPU들은 그동안 멀뚱멀뚱 대기하게 됩니다. 즉, 1% 전후의 고해상도 이미지들만으로도 전체 학습 속도에 큰 악영향을 줄 수 있습니다.

이를 가장 쉽게 해결하는 방법은 특정 해상도를 넘는 이미지를 모두 resize하는 것입니다. 실제로 굉장히 실용적인 해결책입니다. 하지만 분명한 한계도 있습니다. 촘촘한 정보가 집약된 고해상도 이미지의 경우, 해상도를 줄였을 때 정보 손실이 너무 커질 수 있기 때문입니다. 예를 들어 B4 사이즈 종이에 글씨가 빼곡히 적힌 이미지를 색종이 크기로 줄인다고 생각해보겠습니다. 이 경우 글씨가 뭉개져, 모델은커녕 사람도 내용을 읽기 어려워질 수 있습니다.

따라서 저희가 선택한 방법은 max_pixel, 즉 허용 최대 해상도를 학습 진행에 따라 점진적으로 키우는 curriculum입니다.

• 초반: 낮은 max_pixel → patch 수 감소 → 빠른 step, 더 많은 step 소화
• 후반: 높은 max_pixel → 디테일한 시각 정보 학습

사실 max pixel curriculum은 선택이라기보다 어느 정도는 강제되는 측면도 있습니다. LLM 학습에서도 TPS를 최대한 챙기기 위해 context length를 4K에서 시작해 256K 혹은 그 이상으로 점진적으로 늘리는 경우가 많습니다. VLM에서도 비슷합니다. 16MP 이미지는 이미지 토큰만으로도 16K token을 차지하기 때문에 — patch size 16 기준 4096² 이미지는 65,536 patch가 되고, 2×2 merge를 거쳐 LLM 입력 기준 약 16K token이 됩니다 (Figure 2 참고) — 애초에 4K sequence length 학습 단계에는 포함시키기 어렵습니다.

따라서 sequence length curriculum에 맞춰 max_pixel을 함께 늘려가는 것은 경제적인 선택일 뿐 아니라, 현실적으로도 불가피한 선택에 가깝습니다.

지금까지의 최적화를 그대로 둔 채, 허용 해상도 상한만 max_pixel=1M에서 max_pixel=16M으로 올리면 TPS는 다음과 같이 변합니다.

고해상도 이미지가 시퀀스에 포함되기 시작하면서, 앞서 본 이미지 쏠림 문제가 그만큼 커진 것입니다. 다음 (5)의 bin packing은 바로 이 지점을 되찾아오는 작업입니다.

(5) Bin Packing for Even Image Workload Distribution

앞서 말한 이미지 쏠림 문제를 정면으로 푸는 방법입니다^[9]. 핵심 아이디어는 단순합니다.

packing 단계에서부터 image workload가 특정 GPU에 몰리지 않도록 고르게 분배하자.

일반적인 sample packing은 주로 sequence token 수만 맞춥니다. 하지만 VLM에서는 text token 수가 비슷하더라도 이미지 개수나 해상도에 따라 ViT workload가 크게 달라질 수 있습니다. 따라서 저희는 packing 단계에서 image workload까지 함께 고려하도록 bin packing을 수정했습니다.

v1: image token 수로 균형 맞추기

첫 번째 버전은 image token 수를 기준으로 bin 간 균형을 맞추는 방식입니다. 각 bin은 결국 한 DP GPU의 batch slot이 되므로, bin마다 들어가는 image token 수가 비슷해지도록 문서를 배치합니다.

휴리스틱은 간단합니다.

• 이미지가 있는 sample은 현재 image token 수가 가장 적은 bin에 넣습니다.
• 텍스트만 있는 sample은 일반적인 packing처럼 token 수가 가장 적은 bin에 넣습니다.
• 더 이상 들어갈 bin이 없으면 가장 꽉 찬 bin을 pack으로 내보내고 새 bin으로 교체합니다.

이 휴리스틱이 실제로 어떻게 동작하는지는 아래 데모에서 직접 돌려볼 수 있습니다. bin 수 N을 1(greedy packing)로 뒀을 때와 4 이상으로 키웠을 때, pack 간 image token 표준편차가 어떻게 달라지는지 비교해보세요.

이 방식만으로도 low-to-mid resolution 세팅에서는 효과가 꽤 있었습니다. 예를 들어 max pixel = 1MP 세팅에서는 pack_bin_size를 6 정도로만 키워도 시퀀스 간 image token 수의 표준편차가 크게 줄었습니다. bin이 많을수록 이미지를 분산시킬 여유가 생기기 때문에, 비교적 작은 수정으로 rank 간 visual workload를 더 고르게 만들 수 있었습니다.

아래 표에서 N은 packing 시 동시에 열어두는 bin의 개수, 즉 pack_bin_size입니다. N=1은 bin을 하나만 두고 차례로 채우는 일반적인 greedy packing과 같고, N이 클수록 각 이미지 sample을 어느 bin에 넣을지 고를 선택지가 많아집니다.

하지만 max pixel이 커지고 고해상도 이미지가 들어오기 시작하면, image token 수만으로는 부족했습니다. 이유는 ViT 비용이 patch 수에 대해 선형이 아니라 초선형적으로 증가하기 때문입니다. 즉, image token 수는 LLM sequence 길이를 설명하는 데는 좋지만, ViT forward 비용을 설명하는 proxy로는 부족할 수 있습니다.

극단적으로 보면, 초고해상도 이미지 한 장과 저해상도 이미지 여러 장이 같은 patch 수를 가질 수 있습니다. token 수 기준으로는 두 bin이 균형 잡힌 것처럼 보이지만, ViT attention 비용은 큰 이미지 한 장 쪽이 훨씬 비쌀 수 있습니다. 결국 token 기준으로는 예쁘게 나눈 것처럼 보여도, 실제 step time에서는 여전히 특정 GPU가 병목이 됩니다.

v2: ViT cost 기준으로 균형 맞추기

그래서 두 번째 버전에서는 균형 지표를 image token 수에서 ViT cost로 바꿨습니다. patch 수가 적을 때는 거의 token count처럼 동작하되, patch 수가 커질수록 attention 비용을 반영해 더 큰 penalty를 주는 방식입니다.

def vit_patch_cost(patches, p_ref):
    return patches * (1 + patches / p_ref)

여기서 앞의 patches 항은 MLP나 projection처럼 patch 수에 비례하는 비용을, 뒤의 patches² / p_ref 항은 attention처럼 patch 수의 제곱에 비례하는 비용을 근사합니다. p_ref는 두 항이 비슷해지는 기준점으로, 이 값보다 patch 수가 커질수록 고해상도 이미지에 더 큰 비용이 부여됩니다.

정리하면 v1은 image token 수를 고르게 나누는 방식이고, v2는 ViT가 실제로 느려지는 정도까지 반영해 나누는 방식입니다. max pixel이 낮은 구간에서는 v1만으로도 충분히 효과가 있지만, 고해상도 이미지가 섞이기 시작하면 v2처럼 cost-aware하게 packing하는 편이 더 안정적입니다.

위의 테이블을 보시면, bin packing의 효과를 확실하게 보실 수 있습니다. 다만, cost-aware bin packing이 image-token aware bin packing보다 월등한 속도 향상을 봤다고 하기에는 애매할 수 있는데요, 이는 고해상도 이미지의 비율이 높아질수록, 그리고 학습 sequence length가 길어질수록 차이가 벌어질 것으로 예상됩니다.

(6) Chunked ViT Forward / Max Image Capping

bin packing으로 분산을 시켜도, 어쩌다 한 배치에 patch가 많이 몰리면 ViT forward에서 peak memory가 튑니다. OOM 한 번이면 학습 전체가 죽으니 이에 대한 방어 로직이 필요합니다.

Chunked ViT Forward. 배치 내 모든 이미지를 한 번에 ViT에 태우지 않고, 지정된 최대 patch 수 이하로 그룹화해서 순차적으로 ViT를 돌린 뒤 출력을 다시 합칩니다. throughput을 약간 희생하는 대신 peak memory를 억제합니다.

Max Image Capping. packing 단계에서 per-pack image cap(img_cap)을 둬서, 애초에 한 pack에 들어가는 image token 수의 상한을 겁니다. 위 (5)의 bin packing pseudocode에 등장하는 img_cap 부분이 바로 이것입니다 (Figure 10).

더 효율적인 VLM을 향해: 아직 안 해본 것들

위 작업들은 대부분 “주어진 LLaVA-style 구조를 그대로 두고, 학습 인프라/엔지니어링으로 쥐어짜는” 접근이었습니다. 그런데 한 발 물러서서 보면, 더 근본적인 방향도 있습니다. 아직 저희가 시도해보진 않았지만 해볼 만한 것들을 적어둡니다.

Efficient Architecture of Vision Encoder

병목의 상당 부분이 ViT의 $O(N^2)$ attention과 많은 patch 수에서 온다면, vision encoder 자체를 더 효율적인 구조로 바꾸는 게 정공법일 수 있습니다. low precision 학습, attention 효율화 (linear attention, sparse attention…) 등의 연산량 자체를 줄이는 방향입니다.

Vision Encoder-Free VLM

더 과감하게는, 별도의 ViT vision encoder를 아예 없애는 방향이 있습니다. 예를 들어 Gemma 4 12B^[14]는 550M ViT encoder를 35M짜리 단일 matmul projection으로 대체하고, 48×48 RGB patch를 곧장 LLM에 먹입니다. 이렇게 단일 backbone이 modality를 모두 처리하게 하면 — ViT라는 병목 모듈 자체가 사라집니다. 물론 성능 trade-off는 별도로 검증해야 합니다.

마치며

여러 가지 시도를 통해 text-only LLM 학습 대비 78% TPS(max_pixel=1M 세팅 기준)라는 지점까지 도달했습니다. 이게 상한은 아니겠지만, 처음 반도 되지 않던 TPS를 두 배 이상으로 이끌어낸 것은 고무적인 결과라고 할 수 있겠습니다.

tl;dr

• Dataloader부터 최적화합시다. GPU가 I/O를 기다리며 노는 것만큼 아까운 것이 없습니다.
• LLM 최적화에 쓰이는 기술 대부분은 ViT에 그대로 적용하는 것만으로도 큰 이득을 볼 수 있습니다.
• Modality는 batch 단위로 분리하는 것을 추천합니다.
• GPU 간 image workload를 균등하게 만드는 것도 중요한 관건입니다.
• 처음부터 굳이 높은 해상도로 학습할 필요는 없습니다.

References

[1] Trillion Labs. Gravity-16B-A3B-Base. huggingface.co/trillionlabs/Gravity-16B-A3B-Base

[2] Tokens per second. grokipedia.com/page/Tokens_per_second

[3] Kimi Team. Kimi K2.5 Technical Report. arXiv:2602.02276. “K2.5 seamlessly inherits the parallel strategy of K2, achieving a multimodal training efficiency of 90% relative to text-only training.”

[4] Haotian Liu et al. Visual Instruction Tuning (LLaVA). arXiv:2304.08485

[5] NVIDIA. DGX Systems Documentation: Storage. docs.nvidia.com/dgx/bp-dgx/storage.html

[6] Ted Zadouri et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. arXiv:2603.05451

[7] PyTorch. torch.compile documentation. docs.pytorch.org. dynamic=True는 up-front로 최대한 dynamic한 kernel을 생성해 size 변화 시 recompile을 피하려는 옵션입니다.

[8] PyTorch Blog. Activation Checkpointing Techniques. pytorch.org/blog/activation-checkpointing-techniques

[9] Bin packing problem. en.wikipedia.org/wiki/Bin_packing_problem

[10] Haoyu Lu et al. DeepSeek-VL: Towards Real-World Vision-Language Understanding. arXiv:2403.05525. “we observe that this method does not compromise the model’s performance while enhancing the model’s training efficiency by 20%.”

[11] arXiv:2504.00595. low-to-high resolution training이라는 거의 동일한 접근을 보고합니다.

[12] HuggingFaceM4. FineVision. huggingface.co/datasets/HuggingFaceM4/FineVision

[13] LLaVA-OneVision-1.5-Mid-Training-85M. huggingface.co/datasets/tz1002

[14] Gemma Team. Gemma 4 Technical Report. arXiv:2607.02770. Gemma 4 12B는 별도의 vision encoder 없이 48×48×3 RGB patch를 lightweight projection (35M parameters)으로 직접 받는 encoder-free 구조를 사용합니다.