Scaling Up Models and Data with t5x and seqio
\emailLead Authors \nameAdam Roberts \nameHyung Won Chung \nameAnselm Levskaya \nameGaurav Mishra \nameJames Bradbury \emailTechnical Contributors \nameDaniel Andor, Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorne, Aitor Lewkowycz, Alex Salcianu, Marc van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bulian, Xavier Garcia Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee Dan Garrette, James Lee-Thorp \emailTechnical Advisors \nameColin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Maitin-Shepard \emailLeadership \nameNoah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, Andrea Gesmundo \email*Authors are ordered by impact within groups
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: t5x simplifies the process of building and training large language models at scale while maintaining ease of use, and 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.
中文速览
训练超大规模语言模型时，研究者面临两大痛点——如何在超级计算机集群上高效切分模型与数据，以及如何保证训练数据管线快速、可复现。为此，谷歌推出了两个配套开源库：t5x 基于 JAX 框架，通过对 XLA GSPMD 分区器的高层封装，让研究者能以极少的配置改动实现数据并行、模型并行等多种并行策略，从而把拥有数千亿参数的 Transformer 模型训练起来；seqio 则提供以"任务（Task）"为核心的数据管线 API，支持确定性（deterministic）数据流，能从任意断点恢复训练、全局随机打乱、跨多机分片读取，有效避免数据重复带来的过拟合和训练不稳定问题。两者结合后，谷歌内部已用其训练了多个千亿参数级模型，并复现了 T5、mT5、ByT5 等经典模型。这套工具链的意义在于，它将此前需要大量手工调优的大模型训练基础设施打包成研究友好的开源工具，大幅降低了学术界和工业界复现与探索超大模型的门槛。
原文 arXiv:2203.17189;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2203.17189v1