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

Hugging Face는 Vision-Language Model (VLM) 학습에서 발생하는 주요 병목 현상인 GPU 활용도 저하와 과도한 패딩 문제를 해결하기 위해 효율적인 멀티모달 데이터 파이프라인을 개발했습니다. 배치를 포장 문제로 간주함으로써, 팀은 배치당 데이터 밀도를 최대화하고 GPU 간 이미지 분포를 균형 있게 유지하는 균형 잡힌 배낭(knapsack) 전략을 구현했습니다.

문제점: GPU 활용도 저하와 패딩 낭비

비효율적인 데이터 파이프라인은 종종 "idle GPU" 현상을 초래합니다. 이는 하드웨어가 데이터를 기다리는 동안 충분히 활용되지 않는 상황을 의미합니다. 주요 원인 중 하나는 "패딩 지옥"으로, 배치 내 가장 긴 시퀀스에 맞추기 위해 불필요한 패딩 토큰이 과다하게 삽입되는 것입니다. nanoVLM 프로젝트 초기 테스트에서 Hugging Face는 단순 패딩이 배치의 약 60%를 빈 토큰으로 낭비한다는 사실을 발견했으며, 이는 큰 연산 낭비와 비용 증가로 이어졌습니다.

데이터 파이프라인 전략의 진화

Hugging Face는 다섯 단계에 걸친 반복적인 개선을 통해 파이프라인을 최적화했으며, 단순 접근 방식에서 정교한 포장 시스템으로 전환했습니다.

단계 1 & 2: 시각화와 단순 패딩

초기 작업은 데이터셋(이미지, 텍스트 프롬프트, 응답)을 시각화하고 단순 패딩을 적용하는 데 초점을 맞췄습니다. 이 방식에서는 배치 내 모든 시퀀스가 가장 긴 시퀀스에 맞춰 패딩됩니다. 결과적으로 GPU가 많은 양의 빈 토큰을 처리하게 되어 낭비가 크게 발생합니다.

단계 3: 제한된 패딩

낭비를 줄이기 위해 전역 최대 길이를 설정했습니다. 이 길이를 초과하는 샘플은 삭제되었습니다. 일부 패딩이 감소했지만, 실제 내용과 관계없이 모든 시퀀스를 고정 길이로 패딩해야 했으며, 필터링으로 인해 학습 샘플이 손실되는 문제가 있었습니다.

단계 4: 텍스트용 배낭 포장

패딩을 없애기 위해 Hugging Face는 배치를 "배낭 문제" 로 모델링했습니다. 목표는 최대 토큰 제한(max_length)을 초과하지 않으면서 가능한 많은 시퀀스(아이템)를 배치(배낭) 안에 넣는 것입니다.

두 가지 전략을 장난감 데이터로 테스트했습니다:

  • Greedy Packing: 시퀀스를 순차적으로 추가해 배치가 가득 찰 때까지 진행합니다. 빠르지만 이후 배치에 빈 공간이 많이 남을 수 있습니다.
  • Bin-Packing (First Fit Decreasing): 시퀀스를 길이 순(긴 것부터)으로 정렬한 뒤, 첫 번째로 들어갈 수 있는 공간에 배치합니다. 훨씬 더 촘촘한 배치를 만들며 낭비 공간이 최소화됩니다.

이 동적 배치를 지원하기 위해 팀은 map‑style 데이터셋에서 IterableDataset 으로 전환하고, Python 큐를 활용한 생산자‑소비자 패턴을 구현해 포장 과정이 GPU를 병목시키지 않도록 했습니다.

단계 5: 멀티모달 데이터를 위한 균형 잡힌 배낭

최종 단계에서는 토큰 제한과 이미지 예산을 모두 고려한 균형 잡힌 포장 개념을 멀티모달 데이터에 적용했습니다. 이를 통해 샘플당 이미지 수가 GPU 간에 균형을 이루어, 어느 한 GPU가 과도한 이미지 데이터를 처리하는 상황을 방지합니다.

ConstantLengthDataset 클래스는 다음 과정을 관리합니다:

  1. 이미지와 텍스트를 읽음.
  2. 토큰 또는 이미지 제한을 초과하는 샘플을 필터링.
  3. 균형 잡힌 탐욕적 배낭 전략으로 샘플을 포장.
  4. 최종 배치만 고정 길이로 패딩하여 전체 패딩을 최소화.

포장 전략의 기술적 비교

Concept Stage 4 (Toy Data) Stage 5 (Multimodal Data)
Item Integer (sequence length) Full sample (image, prompt, response)
Weight The integer itself Number of tokens (len(input_ids))
Knapsack Batch of integers ≤ max_length Batch of samples ≤ seq_length and image limit
Packing Strategy Greedy or Binpack Greedy packing with token and image constraints
Output List of integers Dict with input_ids, labels, attention_mask, images

결론 및 출처

데이터 파이프라인을 위한 균형 잡힌 배낭 전략은 NVIDIA의 연구 논문 Eagle 2: Building Post-Training Data Strategies from Scratch for Frontier Vision-Language Models 에서 영감을 얻었습니다. 단순 패딩에서 제한적이고 균형 잡힌 포장 방식으로 전환함으로써 개발자는 GPU를 최대한 활용하고 학습 비용을 절감할 수 있습니다.

Sources