Vision Transformer의 대규모 학습을 위한 DeepSpeed 최적화 전략

본 연구는 Vision Transformer(ViT) 모델의 확장성을 향상시키기 위해, 주로 자연어 처리 분야에서 대규모 모델 학습에 성공적으로 적용된 DeepSpeed 프레임워크를 이미지‑중심 워크로드에 적용하고 그 효과를 정량적으로 평가한다. 논문은 먼저 ViT의 구조적 특성을 설명한다. 이미지가 16×16 패치로 분할되고, 각 패치가 토큰화된 뒤 선형 임베딩을 거쳐 Transformer 인코더에 입력된다. 이 과정에서 self‑attention 연산은 입력 토큰 수에 대해 O(N²) 복잡도를 가지며, 모델 파라미터와 중간 활성값이 급격히 증가한다. 따라서 대규모 ViT(예: ViT‑b16, ViT‑L) 학습 시 GPU 메모리와 연산량이 병목이 된다. DeepSpeed는 메모리 효율적인 데이터 병렬(DDP)과 ZeRO 옵티마이저를 통해 파라미터, 옵티마이저 상태, 그래디언트를 GPU 간에 분산 저장한다. 그러나 본 연구에서는 구현 복잡성을 최소화하고 기본적인 데이터 병렬만을 사용한다. DeepSpeed의 init_distributed()와 torch.distributed.barrier()를 활용해 각 프로세스가 GPU와 통신하도록 설정하고, MPI와 NCCL을 결합해 노드 간 통신을 관리한다. 실험 환경은 세 가지 원격 클러스터(Nebula, Tesla, Vector)로 구성된다. Nebula와 Vector는 비교적 동질적인 GPU(RTX 3070, Tesla T4)를 제공하지만, Tesla 클러스터는 RTX 3070, GTX 1070, Tesla P4가 혼합돼 있다. 각 클러스터에서 ViT‑b16 모델을 CIFAR‑10, CIFAR‑100 데이터셋에 대해 학습시켰으며, ImageNet‑100은 시간 제한으로 결과를 제시하지 않는다. 평가 방법은 강 스케일링(strong scaling)과 약 스케일링(weak scaling) 두 축으로 나뉜다. 강 스케일링은 전체 데이터셋을 고정하고 GPU 수를 늘려 학습 시간 감소율을 측정한다. 약 스케일링은 전체 데이터셋을 GPU 수에 비례해 분할하고, 각 GPU가 동일한 작업량을 수행하도록 하여 epoch당 시간이 일정하게 유지되는지를 확인한다. 모든 실험은 5 epoch을 수행하고 평균값을 사용해 변동성을 감소시켰다. 핵심 실험 결과는 다음과 같다. ① Tesla 클러스터에서는 GPU 성능 차이와 작은 배치 크기(16) 때문에 동기화 비용이 급증하고, GPU 수를 늘릴수록 전체 학습 시간이 오히려 증가하는 비선형 현상이 관찰되었다. 이는 이질적인 하드웨어 환경에서 데이터 병렬이 비효율적임을 보여준다. ② Nebula 클러스터에서 배치 크기를 16→64→128으로 확대하면 통신 오버헤드가 크게 감소하고, 학습 시간은 거의 선형적으로 감소한다. 그러나 배치 256으로 확대하면 GPU 메모리 사용률은 포화되지만, CPU‑GPU 데이터 전송 비용이 새로운 병목으로 작용해 추가적인 속도 향상이 제한된다. ③ 배치 크기와 정확도 사이의 관계를 분석한 결과, 배치 64~128에서 최고 정확도(≈85 % CIFAR‑10, ≈60 % CIFAR‑100)를 달성했으며, 과도한 배치(256)에서는 과적합 및 학습 불안정성으로 정확도가 소폭 감소한다. ④ Vector 클러스터에서 1 GPU→2 GPU→4 GPU→8 GPU까지 확장했을 때, 강 스케일링에서는 1 GPU 대비 8 GPU에서 약 7배 가속을, 약 스케일링에서는 epoch당 시간이 거의 변하지 않아 이론적 선형 스케일링에 근접함을 확인했다. ⑤ 다중 노드 단일 GPU 실험(1 GPU/노드, 1 → 32 노드)에서도 intra‑node 다중 GPU와 유사한 스케일링 곡선을 보였으며, NCCL + MPI 기반 통신이 대규모 클러스터에서도 효율적으로 동작함을 입증했다. 논문의 결론은 DeepSpeed를 이용한 데이터 병렬이 ViT와 같은 이미지‑중심 모델에도 높은 확장성을 제공한다는 것이다. 특히 GPU 동질성 유지와 적절한 배치 크기(64~128) 선택이 통신 비용 최소화와 메모리 효율성 확보에 핵심적인 역할을 한다. 또한, gradient accumulation을 활용하면 메모리 제한이 있는 GPU에서도 효과적인 대규모 배치를 구현할 수 있다. 한계점으로는 현재 ZeRO‑Infinity와 같은 메모리 파티셔닝 기법을 적용하지 않아 수십억 파라미터 규모의 ViT에 대한 실험이 부족하다는 점이다. 향후 연구에서는 ZeRO‑Infinity, 모델 병렬, 파이프라인 병렬 등을 결합해 초대규모 ViT 학습을 목표로 하고, 고해상도 이미지(예: ImageNet‑1k)에서의 메모리 및 I/O 병목을 정량화할 계획이다.