기술 요약: LLM 학습을 위한 정식 보증이 포함된 온라인 동적 배칭(Online Dynamic Batching)

1. 문제 정의

현대의 대규모 언어 모델(LLM) 및 멀티모달 미세 조정(fine-tuning) 파이프라인은 학습 시점의 배칭(batching)과 관련하여 근본적인 **관측성 문제(observability problem)**를 겪고 있습니다. 표준 오프라인 배치 샘플러에서는 증강(augmentation), 채팅 템플릿 적용, 토큰화, 멀티모달 시각적 토큰 확장 등을 포함하는 복잡한 전처리 파이프라인을 거친 후에도 샘플의 실제 학습 비용(시퀀스 길이)을 알 수 없습니다.

결과적으로, 배치 구성은 패딩(padding), 메모리 사용량, GPU 포화도를 결정하는 변수를 인지하지 못한 채 "맹목적(blind)"으로 이루어집니다. 기존 솔루션들은 다음과 같은 트레이드오프에 직면해 있습니다:

• 고정 배치 샘플링(Fixed-batch sampling): 롱테일(long-tail) 데이터에서 발생할 수 있는 OOM(Out-of-Memory) 오류를 방기하기 위해 작은 배치 크기를 사용하게 되며, 이는 GPU 활용도 저하로 이어집니다. 큰 고정 배치는 과도한 패딩이나 OOM을 유발합니다.
• 오프라인 길이 캐시(Offline length caches): 길이를 미리 계산하는 방식(예: GMT/BMT 오라클)은 증강 정책, 템플릿 또는 컷오프 길이가 변경될 때마다 매번 다시 구축해야 하는 비용이 많이 드는 정적 캐시를 요구합니다.
• 시퀀스 패킹(Sequence packing): 패딩을 제거하는 데 효과적이지만, 이는 순수한 DataLoader 솔루션이라기보다 모델 레벨이나 커널 레벨의 개입(예: 가변 길이 어텐션)을 필요로 합니다.

나아가, 배치 형성을 정확한 관측이 가능한 시점(전처리 이후)으로 이동시키는 것은 **분산 그룹 정렬 문제(Distributed Group Alignment Problem, DGAP)**를 야기합니다. 분산 데이터 병렬(DDP) 환경에서 모든 랭크(rank)는 동일한 횟수의 그래디언트 감소(gradient-reduction) 단계를 실행해야 합니다. 만약 각 랭크가 로컬의 실측 길이에 따라 독립적으로 가변 크기의 그룹을 형성한다면, 그룹의 수가 랭크마다 달라져 DDP 계약을 깨뜨리고 데드락(deadlock)이나 샘플 손실을 초래하게 됩니다.

2. 방법론: 온라인 동적 배칭 (Online Dynamic Batching, ODB)

저자들은 전처리 후 실측 길이를 관찰하고, 모델, 옵티마이저 또는 어텐션 커널을 수정하지 않고도 동적으로 가변 크기 배치를 형성하는 드롭인(drop-in) 시스템인 **온라인 동적 배칭(ODB)**을 제안합니다.

아키텍처

• 배치 위치: ODB는 PyTorch DataLoader 이터레이터를 래핑(wrap)합니다. Dataset과 Model은 건드리지 않습니다.
• 워크플로우:
  1. 워커(Workers): Null 콜레이트(collate) 함수와 함께 실행되며, 단일 샘플을 전용 **콜레이트 프로세스(Collate Process)**로 전달합니다.
  2. 그룹화(Grouping): 콜레이트 프로세스는 샘플을 버퍼에 모으고, 길이에 따라 정렬한 뒤, 사용자가 지정한 토큰 예산($L_{max}$)을 충족하도록 탐욕적(greedy)으로 가변 크기 배치를 형성합니다.
  3. 정렬(Alignment): 전용 Gloo 그룹(메인 NCCL 그룹과 격리됨)이 모든 랭크 간의 그룹 수를 동기화합니다.
  4. 출력: 정렬된 그룹이 트레이너로 전달됩니다. 채워지지 않은 슬롯은 IDLE_DATA 센티넬(sentinel)로 패딩되며, 메인 프로세스는 이를 건너뛰어 스텝 정렬을 보장합니다.

분산 그룹 정렬 문제 (DGAP)

ODB는 동기화 요구 사항을 DGAP로 공식화합니다. 이는 모든 랭크가 데드락이나 샘플 손실 없이 동일한 횟수의 AllReduce 연산을 수행하도록 하는 Max 기반 양방향 그룹 정렬 프로토콜을 도입합니다.

• 대상 계산: 대상 그룹 수($T_{grp}$)는 활성 랭크들 중 최소 양수 출력 용량과 최소 양수 버퍼 샘플 수 중 최댓값으로 계산됩니다.
• 조정:
  • 분할(Split): 랭크의 그룹 수가 $T_{grp}$보다 적은 경우, 더 큰 그룹을 싱글톤(singleton)으로 분할합니다.
  • 오버플로(Overflow): 랭크의 그룹 수가 더 많은 경우, 상위 $T_{grp}$개의 그룹만 유지하고 나머지 샘플은 버퍼로 재순환시킵니다.
• 종료 모드:
  • 기본 조인 모드(Default Join Mode): 글로벌 완료 전 대기 중인 샘플러 뷰를 모두 소진하여 엄격한 항등성 커버리지(identity coverage)를 보장합니다.
  • 옵트인 비조인 모드(Opt-in Non-Join Mode): 엄격한 반복당 항등성 보장 대신 샘플 할당량 폐쇄(cumulative count)를 보장하며, 이는 제한된 런타임에 적합합니다.

손실 스케일링 (Loss Scaling)

ODB 배치는 랭크당 토큰 수가 다르므로, 단순한 DDP 평균은 편향된 손실 추정치를 산출할 수 있습니다. ODB는 각 랭크의 손실에 전체 토큰 대비 해당 랭크의 비율($w_r = t_r / T_{tok}$)을 가중치로 부여하는 토큰 레벨 손실 스케일링을 구현하여 비트 정밀도(bit-precise)의 토큰당 참조 손실을 복구합니다.

3. 주요 기여

1. 시스템 설계: 전처리 파이프라인 이후의 길이를 관찰하고, 모델 재작성이나 사전 계산된 길이 캐시 없이 DDP 안전한 가변 크기 배치를 형성하는 DataLoader 측 시스템인 ODB를 도입했습니다.
2. 정식 보증: DGAP를 공식화하고 데드락 없는 유계 종료(deadlock-free bounded termination), 엄격한 항등성 커버리지(조인 모드 시), 샘플 할당량 폐쇄(비조인 모드 시)를 증명했습니다.
3. 실험적 성능: Qwen3-VL 2B/8B 모델을 사용하여 UltraChat, LLaVA, ShareGPT4o에 대해 평가하였으며, 표준 방식과 유사한 품질을 유지하면서 상당한 처리량(throughput) 향상을 입증했습니다.
4. 오픈 소스: 경량화된 트레이너 어댑터를 포함한 online-dynamic-batching 패키지를 공개했습니다.

4. 실험 결과

실험은 8×H20 노드(DeepSpeed ZeRO-2, bf16)에서 Qwen3-VL-2B/8B 모델을 사용하여 수행되었습니다.

처리량 이득 (Throughput Gains)

ODB는 고정 배치 베이스라인(Standard, Sorted)을 크게 능가하며, 오프라인 오라클 방식(GMT/BMT)의 캐시 오버헤드 없이도 그 성능에 근접합니다:

• 단일 노드 (Full FT/LoRA): Standard 대비 1.58배 – 2.51배 개선.
• 2개 노드 (Full FT): Standard 대비 1.71배 – 3.78배 개선.
• 고-이질성 (ShareGPT4o, CV=1.00): ODB는 2.46배(8B) 및 2.47배(2B)의 속도 향상을 달성한 반면, Sorted는 롱테일 제약으로 인해 작은 배치 크기에 묶여 약 1.03배에 그쳤습니다.
• 프로덕션 케이스 (MM-Mix): 짧은 샘플 밀도가 높은 프로덕션 멀티모달 믹스에서 ODB는 Standard 대비 4.43배의 속도 향상을 달성했습니다.

품질 지표 (Quality Metrics)

• 검증 손실 및 벤치마크: ODB는 "Standard와 유사한 범위" 내의 품질을 유지합니다. 예를 들어, 8B UltraChat에서 ODB의 MMLU(74.75%)는 Packing(75.18%) 및 GMT(75.14%)와 대등합니다.
• Sorted와의 비교: Sorted는 Standard보다 높은 처리량을 달리는 경우가 많지만, 검증 손실 퇴보 및 답변 형식 저하(예: LLaVA에서 Sorted의 생성 답변 MMMU 점수가 Standard의 22.30% 대비 5.52%로 급락)를 빈번히 겪습니다. ODB는 정렬(sorting) 대신 그룹화(grouping)를 수행함으로써 이를 방지합니다.
• 오라클과의 비교: 오프라인 오라클(GMT/BMT)은 약간 더 높은 원시 처리량을 달성하기도 하지만, 정책 변경 시 재구축해야 하는 값비싼 정적 캐시를 요구합니다. ODB는 "온라인/드롭인" 영역에서 작동하여, 캐시 사전 계산 없이도 대등한 품질을 제공합니다.

절제 연구 (Ablation Studies)

• 토큰 예산 ($L_{max}$): 메모리 압박이나 스텝 시간이 지배하기 전까지 $L_{max}$가 증가함에 따라 처리량이 증가합니다. 최적의 $L_{max}$는 데이터셋마다 다릅니다 (예: UltraChat 12288, ShareGPT4o 14336).
• 대기 깊이 (Outstanding Depth, $D$): 파이프라인 중첩이 포화되면(통상 >95%), $D$를 늘려도 수익 체감(diminishing returns)이 발생합니다.
• 짧은 샘플 활용도: 속도 향상은 변동 계수(CV)와 상관관계가 있지만, 짧은 샘플의 비율($f_s$)에 의해 증폭됩니다. MM-Mix(낮은 CV, 높은 $f_s$)는 ShareGPT4o(높은 CV, 낮은 $f_s$)보다 더 큰 이득을 보았습니다.

5. 의의 및 주장

본 논문은 ODB가 고-이질성 LLM 미세 조정의 온라인/드롭인 영역을 성공적으로 점유한다고 주장합니다. 그 의의는 다음과 같습니다:

• 격차 해소: 고정 배치 학습과 더 강력한 오프라인/모델 측 방식(Packing 또는 Oracle Batching) 사이의 처리량 격차를 줄이면서도, 스칼라 길이 캐시나 모델 측 커널 재작성의 필요성을 피합니다.
• 형식적 안전성: 기존의 정적 시스템들이 피해왔던 영역인, DDP에서의 런타임 가변 크기 배칭에 대해 최초의 정식 보증(데드락 프리, 항등성 커버리지, 할당량 폐쇄)을 제공합니다.
• 배포 가능성: 순수하게 DataLoader 경계에서 작동함으로써, 기존의 훈련 스택(DeepSpeed, LLaMA-Factory)과 호환되며, 오프라인 캐시를 무효화하는 증강 또는 템플릿 정책의 동적 변화에도 견고합니다.

저자들은 이득이 데이터에 따라 달라질 수 있음(고-CV/고-짧은 샘플 데이터셋에서 최적)을 명시하며 겸손한 태도를 유지합니다. 또한 ODB가 모델 측 패킹(packing)이 사용 가능한 경우 이를 대체하기보다는 보완하는 관계임을 밝힙니다. ODB는 모든 배칭 비효효율성을 해결한다고 주장하는 것이 아니라, 현대의 멀티모달 및 증강 LLM 학습에서 발생하는 특정 관측성 병목 현상에 대한 견고하고 정식으로 보증된 솔루션을 제공하는 것입니다.