[논문 리뷰] TorchScale: Transformers at Scale

Shuming Ma, Hongyu Wang외 9명|arXiv (Cornell University)|2022. 11. 23.

Topic Modeling인용 수 6

한 줄 요약

TorchScale는 Magneto 아키텍처, DeepNorm 초기화 및 X-MoE를 포함한 고급 모델링 기법과 함께 새로운 그래디언트 클리핑 방법(SparseClip)을 통합함으로써, 대규모 Transformer 모델을 효율적이고 안정적으로 확장할 수 있도록 해주는 오픈소스 PyTorch 툴킷입니다. 이는 언어 모델링 및 신경 기계 번역 작업 전반에서 안정적인 훈련과 향상된 성능을 달성하며, 96층의 깊이와 256개의 전문가를 가진 MoE 구성에서도 성능을 유지를 합니다.

ABSTRACT

Large Transformers have achieved state-of-the-art performance across many tasks. Most open-source libraries on scaling Transformers focus on improving training or inference with better parallelization. In this work, we present TorchScale, an open-source toolkit that allows researchers and developers to scale up Transformers efficiently and effectively. TorchScale has the implementation of several modeling techniques, which can improve modeling generality and capability, as well as training stability and efficiency. Experimental results on language modeling and neural machine translation demonstrate that TorchScale can successfully scale Transformers to different sizes without tears. The library is available at https://aka.ms/torchscale.

연구 동기 및 목표

• 표준 크기 이상으로 확장할 때 발생하는 불안정성과 비효율성 문제를 해결하기 위해.
• 특히 희소 MoE 아키텍처를 사용할 경우, 깊고 넓은 Transformer 확장 과정에서의 훈련 안정성을 향상시키기 위해.
• 수렴성이나 성능을 희생시키지 않고 모델의 깊이, 너비 및 전문가 수를 효과적으로 확장할 수 있도록 하기 위해.
• 언어 모델링, 기계 번역, 다중모달 프리트레인과 같은 다양한 작업을 지원하는 통합적이고 일반적인 프레임워크를 제공하기 위해.
• 희소 MoE 모델에 특화된 새로운 그래디언트 클리핑 방법(SparseClip)을 도입하여 그래디언트 지배 문제를 완화하고 훈련 안정성을 향상시키기 위해.

제안 방법

• 일반적인 모델링을 위해 서브-LN 초기화와 통합된 인코더-디코더 아키텍처를 사용하는 Magneto 아키텍처를 기본 백본으로 채택합니다.
• DeepNet 이론에 기반한 Sub-LN 초기화를 구현하여 깊은 Transformer에서 최적화 안정성을 향상시킵니다.
• 희소성에 기반한 파라미터 효율적 확장을 가능하게 하기 위해, 상위 1개 및 상위 2개 라우팅 방식을 사용하는 X-MoE(eXtended Mixture of Experts)를 통합합니다.
• 모델의 밀집 및 MoE 구성 요소 간 그래디언트 노름을 균형 잡기 위해 $ \kappa = 1/\sqrt{E} $ 비율로 MoE 그래디언트를 스케일링하는 재가중 그래디언트 클리핑 방법인 SparseClip을 제안합니다.
• 수정된 노름 계산 방식을 사용하는 L2 노름 기반 그래디언트 클리핑: $ ||g||_2 = ||(g_d, \kappa g_s)||_2 $, 여기서 $ E $는 전문가 수입니다.
• 일관된 훈련 및 추론 파이프라인을 갖춘 인코더, 디코더, 인코더-디코더 아키텍처에서 밀집 모델과 희소 MoE 모델을 모두 지원합니다.

실험 결과

연구 질문

• RQ1 훈련 발산 없이 96층과 256개의 전문가를 가진 Transformer를 안정적으로 확장할 수 있는가?
• RQ2 기본 그래디언트 클리핑 대비 희소 MoE 모델에서 제안된 SparseClip 방법이 훈련 안정성과 수렴성에 어떻게 기여하는가?
• RQ3 상위 2개 라우팅 방식을 사용하는 X-MoE 아키텍처는 성능 손실 없이 효율적인 확장을 얼마나 잘 달성하는가?
• RQ4 Sub-LN 초기화를 적용한 Magneto 아키텍처가 언어 모델링 및 기계 번역과 같은 다양한 작업에서 더 나은 일반화와 훈련 안정성을 제공하는가?
• RQ5 TorchScale의 고급 모델링 기법 조합이 다양한 모델 크기와 작업에서 일관된 성능 향상을 달성할 수 있는가?

주요 결과

• TorchScale는 언어 모델링 작업에서 96층 디코더 모델을 안정적인 수렴과 함께 훈련시키며, 깊이가 증가할수록 검증 퍼플렉서티(PPL)가 향상됨을 확인했습니다.
• 상위 2개 라우팅을 사용하는 256개 전문가를 가진 X-MoE 모델은 파라미터 수가 100배 이상 증가했음에도 불구하고 더 낮은 PPL을 기록하여 효과적인 확장을 입증했습니다.
• SparseClip는 희소 MoE 모델의 훈련 안정성과 수렴성을 크게 향상시켰으며, 기본 그래디언트 클리핑 대비 PPL을 유의미하게 감소시켰습니다.
• SparseClip를 사용할 경우 시간이 지남에 따라 그래디언트 클리핑 트리거 수가 감소하여 보다 일관되고 정확한 그래디언트 업데이트가 이루어짐을 나타냅니다.
• SparseClip 하에서 라우팅 엔트로피가 낮아져 훈련 중 전문가 활용도 향상과 더 안정적인 라우팅 행동이 이루어짐을 시사합니다.
• 전문가 수가 증가함에 따라 훈련 FLOPs가 거의 일정하게 유지되어, TorchScale에서 희소 MoE 접근 방식의 효율성을 확인합니다.