Scaling Vision Transformers: Evaluating DeepSpeed for Image-Centric Workloads

논문 요약: DeepSpeed 를 활용한 비전 트랜스포머 (ViT) 의 확장성 평가

1. 문제 제기 (Problem Statement)

• 비전 트랜스포머 (ViT) 의 한계: ViT 는 자기 주의 (self-attention) 메커니즘을 통해 이미지 내의 전역적 관계를 포착하여 이미지 처리 분야에서 뛰어난 성능을 보이지만, 모델 크기와 데이터 규모가 커질수록 계산 비용과 메모리 요구량이 급격히 증가합니다.
• 확장성 장벽: 기존에 언어 모델 (LLM) 에 성공적으로 적용되어 온 DeepSpeed 와 같은 분산 학습 프레임워크가 이미지 중심 작업 (ViT) 에 어떻게 적용될 수 있는지, 그리고 그 확장성 (Scalability) 이 어떤지 체계적으로 평가된 연구가 부족합니다.
• 핵심 질문: DeepSpeed 의 데이터 병렬성 (Data Parallelism) 을 ViT 에 적용할 때, 학습 속도, 통신 오버헤드, 그리고 하드웨어 구성 (동일/이질적 GPU) 이 시스템의 확장성에 어떤 영향을 미치는가?

2. 방법론 (Methodology)

• 실험 환경:
  • 클러스터: 워터루 대학교의 Nebula, Tesla, Vector 등 3 개의 원격 클러스터를 활용.
  • 하드웨어 구성:
    • Nebula: 동종 GPU 환경 (내부 노드 및 노드 간 학습 테스트).
    • Tesla: 이질적 GPU 환경 (RTX 3070 과 GTX 1070/Tesla P4 혼재).
    • Vector: 54 노드, 8 개의 Tesla T4 GPU 를 보유한 대규모 환경.
  • 소프트웨어 스택: DeepSpeed, NCCL, MPI (OpenMPI), PyTorch 분산 백엔드.
• 학습 전략:
  • 모델: ViT-Base-16 (ViT b 16) 아키텍처 사용.
  • 데이터셋: CIFAR-10, CIFAR-100 (ImageNet-100 도 계획되었으나 시간 제약으로 미완료).
  • 확장성 평가:
    • 강한 확장성 (Strong Scaling): 전체 데이터셋 크기를 고정하고 GPU 수를 증가시켜 학습 시간 단축 효과 측정.
    • 약한 확장성 (Weak Scaling): GPU 수에 비례하여 데이터셋 크기를 증가시켜 각 GPU 의 작업 부하를 일정하게 유지하며 성능 측정.
  • 변수 조절: 배치 크기 (Batch Size), 그라디언트 누적 (Gradient Accumulation) 등 소프트웨어 파라미터를 변경하여 성능 영향 분석.

3. 주요 기여 (Key Contributions)

1. DeepSpeed 의 ViT 적용: 언어 모델에 특화된 DeepSpeed 를 Vision Transformer 에 적용하여 내부 (Intra-node) 및 노드 간 (Inter-node) 학습 파이프라인을 구축하고 확장성 추세를 분석함.
2. 분산 학습 효율성 정량화: GPU 수 증가에 따른 학습 속도와 통신 오버헤드를 정밀하게 측정하여 데이터 병렬성의 한계와 가능성을 규명함.
3. 소프트웨어 파라미터 최적화: 배치 크기와 그라디언트 누적 단계가 분산 학습 성능에 미치는 영향을 체계적으로 평가하여 최적의 설정을 도출함.

4. 실험 결과 (Results)

• 하드웨어 동질성의 중요성:
  • Tesla 클러스터 (이질적 GPU): 약한 성능의 GPU 가 포함될 경우, 동기화 비용이 급증하여 학습 시간이 오히려 증가하는 등 확장성이 저하됨. 이는 GPU 동질성이 효율적인 확장에 필수적임을 시사.
  • Vector 클러스터 (동질적 GPU): 32 노드까지 확장 가능한 양호한 확장성 확인.
• 배치 크기와 통신 오버헤드:
  • 소규모 배치 (예: 16): 동기화 비용이 비례하여 급증하여 확장성이 떨어짐.
  • 최적 배치 크기: 64 또는 128이 동기화 비용과 메모리 사용량 간의 최적 균형을 제공.
  • 과도한 배치 크기 (256 이상): GPU 메모리 로딩 오버헤드 (CPU to GPU) 로 인해 성능 향상이 정체됨 (Plateau).
• 확장성 패턴:
  • 강한 확장성: GPU 수 증가에 따라 학습 시간이 선형적으로 감소 (1 GPU 에서 2 GPU 로 증가 시 약 50% 시간 단축).
  • 약한 확장성: GPU 수 증가에 따라 학습 시간이 일정하게 유지됨.
  • 노드 간 vs 내부: Vector 클러스터 실험 결과, 단일 노드 내 다중 GPU 와 다중 노드 단일 GPU 간의 성능 차이가 크지 않음을 확인.
• 정확도: 배치 크기가 너무 커지면 과적합 (Overfitting) 으로 인해 정확도가 하락하는 경향 관찰.

5. 의의 및 향후 작업 (Significance & Future Work)

• 의의:
  • DeepSpeed 가 언어 모델뿐만 아니라 이미지 처리 작업 (ViT) 에도 효과적으로 적용될 수 있음을 입증.
  • 대규모 ViT 학습을 위한 하드웨어 구성 (동질적 GPU) 과 소프트웨어 파라미터 (배치 크기 64~128) 에 대한 실용적인 가이드라인 제시.
• 향후 작업:
  • 메모리 최적화: ZeRO (Zero Redundancy Optimizer) 의 각 단계 (Stage 1~3) 를 평가하여 메모리 절감 효과와 오버헤드 분석.
  • 최적화 알고리즘 비교: SGD, LAMB 등 다양한 옵티마이저 성능 비교.
  • 고해상도 및 긴 시퀀스 처리: DeepSpeed-Ulysses(시퀀스 병렬성) 기술을 ViT 에 적용하여 고해상도 이미지 패치를 시퀀스 차원에서 분할하는 연구 진행.
  • 응용 분야 확장: 의료 영상 (fastMRI), 로봇 공학 (CoSTAR), 유전체학 (GTDB) 등 고해상도 및 다채널 이미지 데이터셋으로 연구 범위 확대.

결론

본 연구는 DeepSpeed 를 활용한 ViT 의 분산 학습이 대규모 이미지 처리 작업에서 확장 가능함을 입증했습니다. 특히, 하드웨어의 동질성과 적절한 배치 크기 (64~128) 설정이 통신 오버헤드를 최소화하고 학습 효율을 극대화하는 핵심 요소임을 규명했습니다. 이는 향후 고해상도 이미지 및 멀티모달 모델의 대규모 학습 파이프라인 구축에 중요한 기초를 제공합니다.