기술 요약: AdaptiveLoad

1. 문제 제기

세계 모델 (World Models) 을 위한 대규모 비디오 확산 트랜스포머 (DiT 및 MMDiT 등) 의 학습은 정적 이미지부터 장시간 비디오에 이르기까지 매우 이질적인 시퀀스 길이를 포함하는 혼합 모드 데이터셋을 처리할 때 두 가지 치명적인 병목 현상에 직면합니다:

• 이차 복잡도로 인한 계산 불균형: 기존 분산 학습 파이프라인은 "동일 토큰" 버킷화 전략 ($B \times S = \text{상수}$으로 제한) 에 의존합니다. 이 접근법은 셀프 어텐션 메커니즘의 $O(S^2)$ 계산 복잡도를 고려하지 못합니다. 결과적으로 긴 시퀀스를 포함하는 버킷은 "계산 지연자 (computational stragglers)"가 되어 심각한 부하 불균형을 초래합니다. 글로벌 동기화 시간은 가장 느린 GPU 에 의해 결정되므로 다른 노드에서 상당한 유휴 시간 (동기화 버블) 이 발생합니다. 실험 데이터는 단계 지연 (step latency) 과 총 토큰 수 사이에는 약한 상관관계 ($R \approx 0.35$) 가 있지만, 시퀀스 길이의 제곱과는 강한 상관관계 ($R \approx 0.92$) 가 있음을 보여줍니다.
• 메모리 벽과 연산자 오버헤드: 멀티모달 DiT(MMDiT) 아키텍처에서 적응형 레이어 정규화 (AdaLN) 는 반복당 수백 번 호출됩니다. 표준 구현은 이산적인 CUDA 커널을 사용하여 스트리밍 멀티프로세서 (SM) 와 대역폭 메모리 (HBM) 간에 중간 텐서의 빈번하고 중복된 왕복을 초래합니다. 또한, 역전파 중 변조 파라미터에 대한 기울기 감소는 비연속적인 접근 (시퀀스 차원 $N$을 따라 비연속적 접근) 으로 인한 비효율적인 스트라이드 메모리 접근 패턴으로 고통받아 하드웨어 수준의 병합 접근을 방해하고 메모리 대역폭을 심각하게 저조하게 활용합니다.

2. 방법론: AdaptiveLoad 프레임워크

이 논문은 고수준 스케줄링과 저수준 커널 융합을 통합한 풀스택 최적화 프레임워크인 AdaptiveLoad를 제안합니다.

A. 이중 제약 적응형 부하 균형

계산 불균형을 해결하기 위해 저자들은 선형 "동일 토큰" 가정을 이중 제약 스케줄링 로직으로 대체하여 다음을 동시에 강제합니다:

1. 선형 메모리 제약: $B \times S \leq M_{mem}$ (배치 크기가 GPU 메모리에 맞도록 보장).
2. 다항식 계산 제약: $B \times S^p \leq M_{comp}$ ($O(S^2)$ 복잡도 고려).

버킷의 최종 배치 크기는 이러한 제약의 교차점에 의해 결정됩니다:
$$B_{shape} = \max(1, \min(\lfloor M_{mem}/S \rfloor, \lfloor M_{comp}/S^p \rfloor))$$

• 비용 피팅 시스템: 경험적 지수 $p$와 계산 상한 $M_{comp}$를 결정하기 위해 시스템은 데이터 기반 Shape Benchmarking 워크플로우를 사용합니다. 이는 라이브 분산 환경 (예: FSDP 경로) 에서 합성 픽셀 스캔을 실행하여 배치 크기와 시퀀스 길이에 대한 단계 지연을 측정합니다. 비선형 회귀를 사용하여 비용 모델 ($\text{step time} \approx a + b \times B \times S^p$) 을 피팅하여 버킷 구성을 동적으로 보정함으로써 수동 경험적 튜닝의 필요성을 제거합니다.

B. 퓨전 LayerNorm-Modulate CUDA 커널

"메모리 벽"을 해결하기 위해 저자들은 정규화 및 변조 연산을 단일 커널 런칭으로 결합하는 퓨전 AdaLN CUDA 커널을 설계했습니다.

• 순방향 통과 (Forward Pass): 커널은 중간 통계량 (평균 및 분산) 을 저장하기 위해 레지스터와 공유 메모리를 활용하며, 2 단계 감소 (Warp 수준 셔플 및 Block 수준 감소) 를 수행합니다. 선형 변환은 특징 차원 $D$의 세 번째 탐색 동안 직접 적용되어 불필요한 HBM 읽기/쓰기를 제거합니다.
• 역방향 통과 (D-tile 병합 감소): 기울기 감소에서의 스트라이드 접근 병목 현상을 해결하기 위해 저자들은 D-tile 병합 감소 전략을 도입했습니다. 시퀀스 차원 $N$을 따라 감소시키는 대신 (이는 스트라이드 접근을 유발함), 그리드를 $(d_{tile}, n_{tile})$ 레이아웃으로 분할합니다. 스레드는 특징 인덱스 $d$에 고정되어 시퀀스 타일 내에서 수직 누적 작업을 수행합니다. 이는 Warp 내 메모리 접근이 연속적이도록 보장하여 병합된 HBM 트랜잭션을 가능하게 하고 대역폭 활용을 극대화합니다.
• 그래프 단순화: 연산자 체인을 퓨전함으로써 계산 그래프가 5~8 개의 이산 노드에서 단일 원자 노드로 축소됩니다. 중간 통계량은 레지스터에 머무르며 프레임워크의 전역 메모리 관리를 우회하여 활성화 메모리 풋프린트를 크게 줄입니다.

3. 주요 기여

1. 부하 불일치의 정량화 및 완화: 이 논문은 비디오 확산 학습의 부하 불일치를 엄격하게 정량화하고 이중 제약 적응형 버킷 스케줄링 알고리즘을 제안합니다. 이 접근법은 긴 시퀀스의 $O(S^2)$ 복잡도를 능동적으로 관리함으로써 부하 불균형률을 최대 50% 까지 감소시킵니다.
2. 데이터 기반 비용 모델링: 저자들은 휴리스틱 기반에서 데이터 기반으로 분산 학습 구성을 전환하는 엔드투엔드 시스템을 설계했습니다. Shape Benchmarking 을 통한 자동화된 비용 모델 피팅을 통해 시스템은 수동 개입 없이 클러스터 전체의 자원 활용을 최적화합니다.
3. 고성능 퓨전 커널: MMDiT 백본을 위한 전용 퓨전 AdaLN CUDA 커널이 개발되었습니다. D-tile 병합 감소를 도입함으로써 역전파의 스트라이드 접근 병목 현상을 해결하여 훈련 처리량을 25% 이상 증가시키면서 동시에 활성화 메모리 사용량을 줄입니다.

4. 실험 결과

실험은 8-GPU 및 16-GPU 클러스터에서 WebDataset 과 Koala-36m 를 포함한 1000 만 개의 샘플로 구성된 혼합 코퍼스를 사용하여 Wan 2.1 세계 모델에서 수행되었습니다.

• 처리량 향상:
  • 8-GPU 클러스터: 평균 훈련 처리량이 초당 14,383 에서 18,069 토큰으로 증가 (+25.6%).
  • 16-GPU 클러스터: 평균 처리량이 초당 30,170 에서 38,372 토큰으로 상승 (+27.18%).
  • AdaptiveLoad 곡선은 최악의 시퀀스 분포에서도 베이스라인보다 더 높은 처리량 바닥을 유지했습니다.
• 부하 균형 효율성:
  • 계산 불균형률 (계산 부하의 변동 계수로 측정) 이 **39% 에서 18.9%**로 감소했습니다.
  • 동기화 오버헤드 (단계 시간의 변동 계수) 가 크게 감소하여, 16-GPU 구성의 경우 AdaptiveLoad 하에서 베이스라인 평균 18.7% 에서 10.4% 로 감소했습니다.
• 메모리 및 커널 성능:
  • 최대 VRAM 활용도: 22.7% 향상.
  • 커널 속도 향상: 퓨전 순방향 통과는 3.21 배에서 3.39 배의 속도 향상을 달성했습니다. 역방향 통과 (D-tile 감소) 는 64k 시퀀스 길이에서 최대 1.42 배의 속도 향상을 달성했습니다.
  • 메모리 절감: 퓨전 커널은 활성화 메모리를 약 61.9% 감소시켰습니다 (예: $N=64,000$에서 약 4.8 GB 절감).
  • 부하 용량: 최적화된 시스템은 최대 메모리 사용량을 증가시키지 않고 52,800 토큰의 최대 시퀀스 길이를 지원했습니다 (베이스라인 48,000 토큰 제한 대비 10% 증가).
• 학습 안정성: AdaptiveLoad 의 손실 곡선은 베이스라인과 매우 일치하여 최적화가 수치적 충실도와 수렴 품질을 유지했음을 입증했습니다.

5. 의의

이 논문은 AdaptiveLoad 가 초장기 시퀀스 비디오 생성 모델의 확장성을 위한 견고한 기반을 제공한다고 주장합니다. 이중 제약 스케줄링을 통해 메모리 안전성과 계산 균형을 조화시키고, 하드웨어 인식 연산자 퓨전을 통해 "메모리 벽"을 해결함으로써, 이 프레임워크는 이질적인 비디오 데이터셋에 내재된 동기화 병목 현상을 효과적으로 제거합니다. 결과는 비디오 생성 작업에서 대규모 GPU 클러스터의 효율성을 극대화하기 위해 선형 토큰 카운팅 휴리스틱을 넘어 이차 복잡도 인식 스케줄링으로 전환하는 것이 필수적임을 보여줍니다.