[论文解读] Scaling Up Models and Data with $ exttt{t5x}$ and $ exttt{seqio}$

Adam Roberts, Hyung Won Chung+41 more|arXiv (Cornell University)|Mar 31, 2022

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.

研究动机与目标

• 提出对大规模、可重复训练的语言模型的需求动机。
• 介绍 t5x，基于 JAX 的库，用于大规模构建、训练、评估和推断 Transformer 模型。
• 介绍 seqio，作为基于任务的数据管道 API，用于高效、确定性和可重复的数据处理。
• 展示这些库如何支持编码器-解码器和解码器前置架构，并与现有框架集成。

提出的方法

• 描述 t5x 的模块化架构，以及它如何将 jax.pjit 包装到 XLA GSPMD，以实现模型、数据和激活的分区。
• 解释分区选项（数据分区与模型分区，1D/2D 参数和激活分区），并将其与 ZeRO-3 和 Megatron 等已知方案联系起来。
• 详细说明如何使用 Flax 和 Gin-config 实现模型与配置，包括与传统 T5 与 Mesh TensorFlow 模型的兼容性。
• 介绍 seqio 作为建立在 tensorflow.data 之上的基于任务的 API，用于可扩展、确定性的数据管道和多任务混合。
• 概述通过 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 内部和外部研究者中的采用情况表明，该工具在大规模语言建模方面具有可用性和研究友好性。