효율적인 멀티모달 데이터 파이프라인

알겠습니다? 우리는 그 상황을 겪었습니다. nanoVLM 프로젝트의 탐정 작업 끝에, 실수인은 모델이나 하드웨어가 아니라, 매우 낭비적인 데이터 파이프라인임을 발견했습니다.

우리가 찾은 것:

방치된 GPU: 모델이 데이터를 기다리느라字字 literally 대기하고 있었습니다.패딩 지옥 (Padding hell): 각 배치마다 무용지부에 기여하지 않는 불필요한 패딩 토큰으로 가득 차 있었습니다.

이 글에서는 5 단계로 효율적인 파이프라인을 구축합니다. 각 단계에서 이전 단계에 추가하거나 제거하며, 무엇이 잘되었고 무엇을 하지 않았는지에 대해 논의합니다.

• 단계 0: 사전 요구사항 (Pre Requisites)
• 단계 1: 데이터셋 시각화 (Visualising the Dataset)
• 단계 2: 무난한 패딩 (Naive Padding)
• 단계 3: 제한된 패딩 (Constrained Padding)
• 단계 4: Knapsack을 활용한 더 스마트한 패킹 (Packing Smarter with Knapsacks)
• 단계 5: 멀티모달 데이터용 Knapsack (Knapsack for Multimodal Data)
• 결론

데이터 준비 작업을 따라가기 쉽게 하기 위해, 데이터 파이프라인에만 초점을 맞춘 별도의 저장소를 만들었습니다. nanoVLM 저장소와 통합된 코드를 읽는 것보다 훨씬 이해하기 쉬울 것입니다. 또한, 다른 데이터 파이프라인을 부팅 (bootstrap)하는 데에도 유용할 수 있습니다.

저장소: https://github.com/ariG23498/mmdp

따라하려면, 모든 것이 필요한 것은 저장소를 클론하는 것입니다. 최종 데이터 준비 작업을 포함하지만, 각 단계를 보여주기 위해 설계되었습니다.

$ git clone https://github.com/ariG23498/mmdp.git

무언가를 최적화하기 전에, 우리가 작업하고 있는 것을 이해해야 합니다. 우리의 멀티모달 데이터셋에는 이미지, 텍스트 프롬프트, 응답이 있습니다.

$ uv run 01_check_dataset.py

훈련 데이터를 잘 아는 것은 성공에 필수적입니다. 이전 스크립트는 실행할 때마다 랜덤 샘플을 보여줍니다. 데이터에 대한 느낌을 얻기 위해 스니펫을 노트북으로 복사하여 여러 번 실행하는 것을 원할 수 있습니다.

우리의 첫 번째 훈련 시도는 명백한 (또는 매우 빈번한) 접근법을 사용했습니다:

• 모든 것을 토큰화 (Tokenize)
• 각 배치에서 가장 긴 시퀀스 찾기
• 나머지를 모두 패딩에 맞추기

$ uv run 02_naive_pad_dataloader.py

결과가 고통스러웠습니다. 이 시각화를 보세요:

그 모든 회색 (gray)을 보셨나요? 그건 패딩입니다. GPU 는 아무것도 처리하지 않고 계산 시간을 지불합니다. 우리는 빈 토큰으로 배치의 약 60% 를 낭비했습니다.

다음 단계는 간단했습니다. 전역 최대 길이를 설정하고 이를 유지하세요. 샘플이 너무 길다면, 그냥 제거하세요.

배치에 한 샘플이 더 적다는 것을您可能 noticed 했을 것입니다. 이는 필터링 프로세스 때문입니다. 도움이 되었지만, 우리는 여전히 실제 내용에 관계없이 모든 것을 동일한 고정 길이에 패딩했습니다. 이전보다 좋았지만 여전히 낭비적이었습니다.

이제 우리는 배치를 완전히 재고할 준비가 되었습니다. 패딩은 적대적이며, 각 배치에 들어갈 수 있는 데이터를 최대화하면서 패딩을 최소화하는 전략이 필요합니다. **Knapsack 문제 (knapsack problem)**라는 컴퓨터 과학의 고전적인 문제가 여기에 적합합니다.

산책용 가방을 채우는 것을 상상해 보세요. 그것은 일정한 무게만 지탱할 수 있으며, 최대한 많은 유용한 물건을 넣으려는 것입니다. 우리 경우:

• 백팩은 최대 토큰 제한 (max_length)이 있는 훈련 배치입니다.
• 각 **물건 (item)**은 시퀀스 (토큰화된 프롬프트-응답 쌍) 이며, 그 **무게 (weight)**는 토큰 수입니다.
• 목표는 토큰 제한을 초과하지 않으면서 최대한 많은 시퀀스를 배치에 넣는 것입니다. 낭비된 공간을 최소화합니다.

이 아이디어를 테스트하기 위해, 우리는 작은 데이터셋으로 시작합니다: 1 에서 25 까지의 숫자 목록만이며, 각각은 시퀀스 길이를 나타냅니다. 이미지와 텍스트의 복잡성 없이 실험할 수 있습니다.

대부분의 PyTorch 데이터셋은 map 스타일 (데이터셋을 dataset[i] 로 접근) 입니다. 하지만 동적 배치 (dynamic batching) 를 위해 더 유연한 것이 필요합니다. 따라서 우리는 torch.utils.data.IterableDataset 을 상속받아 iterable 스타일 데이터를 구축했습니다. 이를 통해 실시간으로 배치를 생성하고 여러 작업자 (workers) 간 데이터 분할 (sharding) 같은 트릭을 처리할 수 있습니다:

def _get_data_range(self):
worker_info = get_worker_info()
if worker_info is None: # 단일 작업자의 경우, 전체 데이터셋 반환
...

시퀀스 (sequence) 를 패킹하는 것은 특히 정렬이나 섞기를 할 때 느릴 수 있습니다. 이를 유지하기 위해 Python 큐를 사용하여 프로듀서 - 컨슈머 패턴을 사용합니다:

def _producer(self, data_iter, queue, stop_signal):
if self.strategy == "greedy":
for pack in self._greedy_packing(data_iter):
...

프로듀스 (producer) 스레드는 배치를 패킹하고 큐에 넣으며, 메인 스레드는 필요할 때 꺼냅니다. 이 중첩은 파이프라인을 원활하게 흐르게 합니다.

먼저 간단한 greedy packing 전략을 시도해 보겠습니다:

def _greedy_packing(self, iterator):
pack, pack_sum = [], 0
for item in iterator:
...

이것은 데이터를 순차적으로 걸으며, 패치가 가득 차면 새 패치로 시작합니다. 빠르지만 완벽하지는 않습니다. 배치가 어떻게 보이는지 확인해 보겠습니다:

=== Strategy: GREEDY ===
[tensor([1]), tensor([2]), tensor([3]), tensor([4]), tensor([5]), tensor([6]), tensor([7]), tensor([8]), tensor([9]), tensor([10]), tensor([11]), tensor([12]), tensor([13])]
[tensor([14]), tensor([15]), tensor([16]), tensor([17]), tensor([18]), tensor([19])]
...

후속 배치들이 희소 (sparse) 해진 것을 주목해 보십시오. 우리는 공백을 남기게 됩니다.

더 똑똑한 접근법을 시도해 보겠습니다: bin-packing (구체적으로, First Fit Decreasing):

def _bin_packing(self, buffer: List[int]):
buffer = sorted(buffer, reverse=True)
knapsacks = []
...

이는 시퀀스를 길이로 정렬 (가장 긴 것부터) 한 후, 공간이 있는 첫 번째 패치에 넣으려 시도합니다. 어떤 것도 맞지 않으면 새 패치를 시작합니다. 결과는 다음과 같습니다:

=== Strategy: BINPACK ===
[tensor([24]), tensor([23]), tensor([22]), tensor([21]), tensor([10])]
[tensor([20]), tensor([19]), tensor([18]), tensor([17]), tensor([16]), tensor([9]), tensor([1])]
...

이 배치는 훨씬 더 조밀하며, 낭비되는 공간이 적습니다. 데이터와 테트리스 (Tetris) 를 하는 것처럼, 조각을 조밀하게 맞춰 넣는 것입니다.

이제 우리의 멀티모달 (multimodal) 데이터셋에 knapsack packing 을 적용하는 실제적인 부분입니다.

우리는 이미지, 프롬프트 (prompts), 응답으로 돌아가며, 토크 제한과 이미지 예산을 모두 존중하면서 효율적으로 패킹해야 합니다. 이미지 예산은 샘플당 이미지 수를 균형을 맞추기 위해 수행됩니다. 우리는 하나의 GPU 가 다른 GPU 보다 훨씬 더 많은 이미지를 처리해야 하는 경우를 피하고 싶습니다.

새로운 ConstantLengthDataset 클래스가 중량을 담당합니다. 이는 Stage 4 와 비교하여 다음과 같이 작동합니다:

Concept	Stage 4 (Toy Data)	Stage 5 (Multimodal Data)	Function(s)
Item	Integer (sequence length)	Full sample (image, prompt, response)	VQADataset.__getitem__
Weight	The integer itself	Number of tokens (len(input_ids) )	—
Knapsack	Batch of integers ≤ max_length	Batch of samples ≤ seq_length and image limit	_balanced_greedy_knapsack
Packing Strategy	Greedy or Binpack	Greedy packing with token and image constraints	_balanced_greedy_knapsack
Producer-Consumer	Producer fills queue	Same as the toy example, but with multimodal samples	_producer , __iter__
Sample Filtering	Skip integers > max_length	Skip samples with too many tokens or images	_producer
Sharding	Split integer range	Shard dataset indices	make_base_iterator()
Batching	Group integers	Concatenate and align tokens/images	_pack_one_group
Output	List of integers	Dict with input_ids , labels , attention_mask , images
yield from __iter__

The ConstantLengthDataset

do es todo:

• Reads samples (images and text).
• Filters out samples that are too long or have too many images.
• Packs samples into batches using a greedy knapsack strategy, balancing token count and image count.
• Pads the final batches to a fixed length, but with
  way lesspadding than before.

Here's the result:

Look at that! The gray (padding) is minimal, and the batches are dense with useful data. It's like packing a suitcase so well you can still zip it up without sitting on it.

The image might seem unintuive at the first glance, but let us look at the image side by side with constrained padding.

Here you will notice that the samples in knapsack are more evenly distributed. We also do not run into the issue of having less samples in the batch due to filtering.

What started as a simple "why is training so slow?" investigation led to a complete rethink of how we handle multimodal data.

The balanced knapsack strategy for data pipeline comes from the Eagle 2: Building Post-Training Data Strategies from Scratch for Frontier Vision-Language Models paper from NVIDIA.

The key lessons:

• Padding everything to the longest sequences is a good first approach (but wasteful)
• Think of batching as a packing problem
• Consider all your constraints (text length, image memory, etc.)
• Test with toy data first to validate your approach

Want to dig deeper? Check out:

Happy training (and may your GPUs stay busy)!