[논문 리뷰] Scaling Up Models and Data with $ exttt{t5x}$ and $ exttt{seqio}$

Adam Roberts, Hyung Won Chung외 41명|arXiv (Cornell University)|2022. 03. 31.

Topic Modeling인용 수 48

한 줄 요약

이 논문은 두 개의 오픈 소스 라이브러리, t5x와 seqio를 제시하여 Transformer 모델 및 데이터 파이프라인의 확장을 단순화하고, TPU 및 기타 하드웨어에서 대규모 학습 및 재현 가능한 평가를 가능하게 한다.

ABSTRACT

Recent neural network-based language models have benefited greatly from scaling up the size of training datasets and the number of parameters in the models themselves. Scaling can be complicated due to various factors including the need to distribute computation on supercomputer clusters (e.g., TPUs), prevent bottlenecks when infeeding data, and ensure reproducible results. In this work, we present two software libraries that ease these issues: $ exttt{t5x}$ simplifies the process of building and training large language models at scale while maintaining ease of use, and $ exttt{seqio}$ provides a task-based API for simple creation of fast and reproducible training data and evaluation pipelines. These open-source libraries have been used to train models with hundreds of billions of parameters on datasets with multiple terabytes of training data. Along with the libraries, we release configurations and instructions for T5-like encoder-decoder models as well as GPT-like decoder-only architectures. $ exttt{t5x}$ and $ exttt{seqio}$ are open source and available at https://github.com/google-research/t5x and https://github.com/google/seqio, respectively.

연구 동기 및 목표

• 대형 언어 모델의 확장 가능하고 재현 가능한 학습의 필요성에 대한 동기 부여.
• JAX 기반 라이브러리로서 대규모로 Transformer 모델을 구축, 학습, 평가, 추론하기 위한 t5x를 소개한다.
• 효율적이고 결정적이며 재현 가능한 데이터 처리를 위한 작업 기반 데이터 파이프라인 API로서 seqio를 소개한다.
• 이 라이브러리들이 인코더-디코더 및 디코더 전용 아키텍처를 어떻게 지원하고 기존 프레임워크와의 통합을 보여준다.

제안 방법

• jax.pjit를 XLA GSPMD로 래핑하여 모델, 데이터, 활성화 분할을 위한 모듈식 아키텍처와 t5x의 구성을 설명한다.
• 데이터 대 모델 분할, 1D/2D 파라미터 및 활성화 분할과 같은 분할 옵션을 설명하고 이를 ZeRO-3 및 Megatron과 같은 알려진 스킴과 연결한다.
• Flax와 Gin-config가 모델 및 구성을 구현하는 데 어떻게 사용되는지, 레거시 T5 및 Mesh TensorFlow 모델과의 호환성을 포함하여 상세히 다룬다.
• tensorflow.data 위에 구축된 작업 기반 API인 seqio를 확장 받아 확장 가능하고 결정론적 데이터 파이프라인과 다작업 혼합을 제공한다.
• 재현성을 위한 샤딩 및 전역 셔플링을 포함해 Apache Beam을 통해 구현된 결정론적 파이프라인 기능을 개요한다.

실험 결과

연구 질문

• RQ1 고수준 추상화를 사용하여 데이터 및 모델 병렬성 축에서 Transformer 모델의 확장을 어떻게 단순화할 수 있는가?
• RQ2 공정한 비교 및 대규모 디버깅을 가능하게 하는 재현 가능하고 결정론적 데이터 파이프라인을 어떻게 보장할 수 있는가?
• RQ3 t5x와 seqio 내에서 인코더-디코더 대 디코더 전용 아키텍처에 대한 실용적인 구성 및 작업 흐름은 무엇인가?
• RQ4 t5x와 seqio가 기존 모델 구현체 및 학습 생태계(Flax, TensorFlow, PyTorch)와 어떻게 통합되는가?

주요 결과

• t5x는 JAX/XLA GSPMD에 대한 고수준 인터페이스를 제공하여 대규모 Transformer 모델에 대해 데이터, 매개변수 및 활성화 분할을 유연하게 가능하게 한다.
• seqio는 결정론적 파이프라인, 재현성, 복구 가능성 및 효율적인 분산 읽기를 갖춘 작업 기반 데이터 파이프라인을 가능하게 하여 대규모 학습 및 평가에 도움을 준다.
• 라이브러리들은 인코더-디코더 및 디코더 전용 모델 구성에 대한 호환성 및 레거시 T5 코드베이스와 Mesh TensorFlow 모델에서의 전환 경로를 제공한다.
• T5-like 및 GPT-like 아키텍처를 위한 구성과 가이드가 포함된 오픈 소스 릴리스로 TPU 환경 전반에서 빠른 실험과 확장을 촉진한다.
• Google 내 채택 및 외부 연구자들의 활용은 대규모 언어 모델링에 대한 사용성 및 연구 친화성을 보여준다.