DISTMM: Accelerating Distributed Multimodal Model Training

논문 (NSDI'24) : DistMM - https://www.usenix.org/system/files/nsdi24-huang.pdf

Multimodal model을 분산학습시키는 상황.
이 상황에서 분산학습 시스템의 시간/연산 효율성을 개선하려면?

Summary

1. ✨ 기존 분산 학습 시스템들은 이종적인(heterogeneous) 멀티모달 모델의 특성과 대규모 배치(large batch)를 요구하는 Contrastive Loss를 고려하지 않아 비효율적이었다.
2. 🛠️ DISTMM은 Modality-aware Partitioner, Data Load Balancer, Heterogeneity-aware Placement Manager를 통해 서브모듈의 이질성을 활용하고 통신 오버헤드를 줄이며, Pipeline Executor의 `batch-sync instruction` 및 DISTMM-Pipe 스케줄로 대규모 배치 요구사항을 충족한다.
3. 🚀 다양한 구조와 규모(1.1B~26B 파라미터)의 CLIP, CoCa, LiT 모델 실험에서 DISTMM은 기존 Megatron-LM보다 1.32배에서 3.27배 빠른 학습 속도를 보였다.

Problem

Multimodal model의 submodules(모델 조각들)이 모델 구조로보나, input 크기로 보나 서로 이종적(heterogeneous)이다. -> -) 연산 비효율성을 초래한다. 

기존 Multimodal Model의 특징:

• Heterogeneous submodules: 각 modality마다 특징이나 구조가 다르고, 또 multimodal task가 어떤 modality를 주력으로 하는지에 따라서도 학습 구조가 달라짐. 그리고 이 modality의 구조적 이종성이 GPU utilization에도 영향을 미침. 
  • CLIP: Vision module (ViT) 크기 >> Text module 크기
  • LiT: Vision module 크기 << Text module 크기
  • CoCa (Contrastive captioner): Vision module 크기 $\approx$ Text module 크기
  • -) 이 각 submodules을 동등하게 취급하면 sub-optimal한 시간 효율성을 가지게 됨. 
• Imbalanced input sizes: modality type에 따라 input size가 다름. 
  • e.g., CLIP: text input은 77 단어, image는 512x512 pixel 크기 제한. 
• Large batch size requirement: 기존 연구에 따르면 contrastive learning할 때 batch size 큰 게 더 성능 차원 (robust & generalized representation)이나 수렴 속도 차원에서 좋다.
  • -) 여기에 pipeline parallelism을 적용하려면 최대 마이크로 배치 크기가 $(M-\frac{M_s}{P})/M_a$​로 제한되어 batch size 키우기가 어렵다. ($M$: GPU total memory, $M_s$: static memory(weight, gradient, state)​, $M_a$: activation memory​)

•  

Multimodal Model Taxonomy:

항목	Dual Encoder	Dual-Stream Fusion	MLLM (LLM-based)
아키텍처 예시	- image → Enc →v - text → Enc →t => sim(v,t)	- image → Enc →v_tokens - text → Enc →t_tokens => cross-attn	- image→Enc→proj→LLM - text→LLM
구조 특징	완전 분리	encoder 분리 후 fusion	LLM 중심 구조
interaction 시점	없음	중간	LLM 내부
attention 구조	없음	cross-attention	causal self-attention
대표 모델	CLIP	BLIP, ALBEF	LLaVA, Qwen-VL
입력 형태	modality별 입력	tokenized separately	image→token + text
출력 형태	embedding	multimodal feature	text sequence
주요 목적	alignment	multimodal task	generation
loss function	contrastive	contrastive / ITM / LM	LM (cross entropy)
학습 방식	representation learning	hybrid learning	instruction tuning
장점	효율적	성능 균형	reasoning / generation
단점	interaction 없음	구조 복잡	매우 무거움

 

Method

⇒ 이 논문에서는 총 4가지 요소를 가지고 멀티모달 모델의 분산학습 시스템을 구성한다:

1. Modality-aware Partitioner: 전체 멀티모달 모델을 입력 모달리티(modality)에 기반하여 하위 모듈(submodules)로 분할한다. 각 하위 모듈의 신경망 아키텍처와 구성이 다르다는 점을 활용하여, 하위 모듈 크기에 따라 독립적인 병렬화 전략을 적용한다. 
   • e.g., 연산이 많은 submodule은 TP를 쓰고, 단일장치에 load 가능한 작은 크기의 submodule은 DP를 쓰도록.
   • +) 각 모달 하위 모듈의 계산 지속 시간을 균형 있게 맞춰 전체적인 높은 계산 효율성을 달성.
2. Data Load Balancer: Modality-aware Partitioner에 의해 분할된 하위 모듈과 클러스터 구성을 입력으로 받아, 각 하위 모듈 파티션에 할당할 장치 수와 데이터 배치 크기를 결정한다. (Dynamic programming으로 최적화)
   • +) 각 모달 하위 모듈의 계산 지속 시간을 균형 있게 맞춰 전체적인 높은 계산 효율성을 달성.
3. Heterogeneity-aware Placement Manager: 각 submodule의 locality (=어떤 submodule을 어떤 node의 어떤 GPU에 배치할지)를 최적화했다.
   • 노드 내 NVLink와 노드 간 Ethernet의 bandwidth 차이를 고려.
   • e.g., 동일 modality 내 submodules -> 서로 가깝게 배치하여 고대역폭 링크를 활용 / 다른 modality의 submodules -> 별도의 노드에 배치하여 저대역폭 링크에서의 통신량(communication volume)을 줄임.
   • +) submodule 간 통신량을 고려하여 배치함으로써 통신량, 통신 시간 최적화. 
4. Pipeline Executor: 기존의 pipeline execution schedule의 batch size 한계를 극복하고자 새로운 pipeline executor를 제안.

DistMM Overview

그림을 보면,

1. Modality-aware Partitioner가 image module (파랑)과 text module (노랑)을 각각 4등분, 16등분한다. 아마도 parameter size 같은 모델의 meta 정보 가지고 4와 16이라는 값을 결정한듯. 이 4등분/16등분 값은 다시 말해 TP degree x PP degree를 의미한다. (e.g., 파랑 모듈: TP degree x PP degree = 4가 되어야 함)
2. 그 다음, Data Load Balancer가 전체 batch size (Gbs=72)에 맞춰서, DP degree와 각 DP group 별 batch size를 결정. 그래서 파랑 모듈은 4개의 replica가 필요하고 (dp degree=4, 각 replica 당 batchsize (Rbs)=18), 노랑 모듈은 3개의 replica가 필요하다 (dp degree=3, Rbs=24). 
3. Heterogeneity-aware Placement Manager가 이제 마지막으로 실제로 GPU를 할당해주고, 이에 따라 TP degree와 PP degree도 결정된다. 이 예시에서는 전체 GPU cluster가 8개 GPU 짜리 node 8개인 상황에서, GPU 16개짜리 노랑 모듈 replica 3개와 GPU 4개짜리 파랑 모듈 replica 4개를 통신 효율적으로 배치해야 한다.
   • 아무래도 DP replica들은 sync를 해줘야 하니까 replica들끼리는 TP degree / PP degree가 통일되는 게 연산 시간을 맞추는데 도움이 되겠지.
   • 노랑이 먼저 생각해보면 TP deg x PP deg = 16이어야 하는데, TP가 통신량이 많이 heavy하니까 TP degree를 낮추고 TP group끼리는 같은 node에 배치해주어야 함. 그러면 TP deg = 2이거나 4 정도가 적당. 근데 TP deg =2인 경우는 PP deg=8이 되어 pipeline bubble 현상이 심화되는 문제도 발생할 수 있고, 파랑 모듈 replica들을 모두 서로 다른 node에 배치시켜야 해서 파랑 모듈의 DP grad sync 통신량도 상당해진다. 뭐 이러저러한 계산 결과로.. (실제 결정할 때는 이렇게 heuristic하게 결정하지는 않겠지. 아님 말고) 
   • 결과적으로 ... 최종 degree (DP deg, PP deg, TP deg)는 노랑 모듈: (3, 4, 4), 파랑 모듈: (4, 2, 2) 로 결정이 난다. 
4.  마지막, Pipeline Executor가 이 PP deg를 고려해서 interleaved 1f1b schedule을 만들었다. 이 때 각 GPU의 메모리 사정을 고려했을 때 microbatch size (Mbs)는 노랑 모듈이 3, 파랑 모듈이 2가 되고, 기존에 Rbs가 각각 18, 24이었으니까 microbatch 수는 각각 6, 12가 된다. 
   • 이 때 노랑 모듈과 파랑 모듈 feature vector를 짬뽕해서 similarity 계산하는 modality interactive module의 경우 언제 연산해주어야 하느냐? (이 모듈을 학습할 때는 노랑 pipeline과 파랑 pipeline 간 sync가 필요한 상황.) 얘는 모듈 간 통신이 필요하다보니까 매 microbatch마다 연산하지 않고, 그림에서 Interactive Computations로 적힌 빨간 선 부분에서 한번에 통합해서 연산함. 

Modality-aware Partitioner

기존 분산학습 시스템에서는 submodule을 다 통합 고려해서 전체 모델 대상으로 모델을 split했다면, 이 Modality-aware Partitioner에서는 각 submodule을 독립적으로 고려해서 split한다. 

즉, 이 단계에서는 GPU에다가 모델 쪼개는 방식을 최적화.

(a) 기존 Colocated solution 대비 (b) DistMM의 Modality-aware Partitioner

1. modality interactive submodule 같은 경우는 연산이 heavy하지 않고 통신량도 적기 때문에 각 device에 균등하게 나눠담는다.
2. 그리고 각 modal module의 경우, efficiency를 최대화하는 방향으로 adaptive하게, 그리고 독립적으로 split해서 device에 나눠담는다. 
   • 이 때 parallelism strategy라던가 parallelism degree 같은 것들은 common practice를 따른다 (model size 고려해서 DP를 할 수도 있고 TP를 할 수도 있고 PP를 할 수도 있고).

Data Load Balancer

위에서 모델을 잘 쪼개줬다면, Data Load Balancer는 cluster setup(node 몇 개고 node 당 device 몇 개인지)과 training config(Modality-aware partitioner에 의해 쪼개진 model partition들 정보와 global batch size)를 input으로 받아서, 각 model partition이 몇 개의 device에 할당될지 (e.g., DP의 경우 여러 대 GPU에 model partition replica를 load해야 함)를 결정("resource assignment plan")한다.

목표는 가장 느린 model partition의 연산 시간을 최소화하는 것.

(a) 기존 Colocated solution 대비 (b) DistMM의 Data load balancer

Heterogeneity-aware Placement Manager

이제 각 model partition 별 GPU 수가 정해졌으니, 이를 바탕으로 Heterogeneity-aware Placement Manager에서는 bandwidth를 고려해 communication-efficient한 방식으로 실제 device에 배치해준다. 

위에서 모델 어떻게 쪼개서 몇 개 GPU 쓰라고 결정했다면 여기서는 "그래서 어떤 GPU에 어떤 model partition을 배정할지"를 정하는 단계. 

(a) 기존 Colocated solution 대비 (b) DistMM의 Heterogeneity-aware Placement

 

위 그림 예시를 보면, 원래 colocated solution에서는 파랑색 모달과 주황색 모달 모두 4개 GPU에 쪼갰다보니 4개 GPU가 다 통합적으로 gradient AllReduce를 해야 했는데, 이제는 파랑색 모달은 GPU 하나에 담겨서 allreduce가 필요 없어지고 주황색 모달은 3개 GPU끼리만 allreduce하면 돼서 확실히 comm. efficient해졌다. 

Pipeline Executor

여기서 PP를 적용할 때 가장 포인트는, 학습에서 요구되는 batch size (Required batch size, Rbs)와 실제 GPU의 메모리가 감당할 수 있는 micro batch size (Mbs)의 격차를 고려하는 것이다. 

K=Rbs/Mbs라고 할 때, DistMM-Pipe는 Mbs/2 크기의 2K개 microbatch로 쪼개서 ...???

Pipeline Executor는 실행 스케줄 명령하는 Batch-sync instruction 부분과 

원래 PP에서는 각 stage에서 gradient까지Contrastive Learning할 때 similarity를 계산하는데, 이걸 위해서는 ...???

Batch-sync instruction은 총 4단계로 구성됨:

1. Memory movement: 이전 K개 microbatch의 forward 결과 feature vector를 continuous feature vector로 합치는 작업. 
2. Forward pass of modal interactive submodule
3. Backward pass of modal interactive submodule: modality interactive submodule의 backward 실행 -> continuous feature vector에 대한 gradient 생성. 
4. Memory dispatching: continuous gradients를 각 microbatch의 feature vector에 대한 개별적인 gradient (총 K개)로 분리하는 작업.

-) 근데 솔직히 .. pipeline executor 부분 이해 못했음. 그림에서 overhead reduction은 그냥 GPipe 쓸 거 1f1b 쓰고 microbatch 두 배로 쪼개서 줄어든 거 아니야? 이 논문의 novelty에 의한 감소는 아닌 거 같은데.. 그리고 논문에서 Mbs/2 크기의 2K개 microbatch로 나눴다는데, 왜 굳이 2야?

Strength

+) VLM의 모달리티 간 독립성을 고려한 분산 시스템 디자인이라는 점. 

Weakness

-) 그리고 전반적으로 dual encoder (e.g., CLIP) 구조에 맞춰서 설계된 거 같다. 요즘 많이 쓰는 MLLM 구조랑은 잘 안 맞는 듯. 

-) 그리고 transformer 구조들은 Sequence Parallelism도 많이 쓰는데 여기서는 고려가 안 되어 있다.