- 人工智能
- 大模型
- NLP
- 预训练
- 微调
【免费下载链接】text-to-text-transfer-transformer
Code for the paper "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer"
released_checkpoints.md是 T5(Text-To-Text Transfer Transformer)开源仓库中专门记录论文发表后新增实验性预训练模型检查点的文档,涵盖t5.1.1.*、t5.1.1.lm100k.*、t5.1.th.*、t5.1.n4w10.*四个系列共 15 个检查点。本文以该文档为主体,结合仓库内t5/models/gin/models/下的真实 gin 配置源码,逐一拆解每个系列的架构改进、参数规模、底层实现与适用场景,并给出在 T5 框架(Mesh TensorFlow 与 Hugging Face 两条路径)中加载、微调这些检查点的完整方法。读完本文,你将能准确区分 T5 官方论文检查点与实验检查点的差异,看懂t5.1.1.base.gin、t5.1.th.base.gin、t5.1.n4w10.base.gin等配置文件的每一行含义,并正确地把这些检查点接入自己的微调流程。
一、文档定位:论文模型之外的“实验性检查点”
1.1 与官方检查点(released checkpoints)的区别
仓库 README.md 的 “Released Model Checkpoints” 一节列出了论文中描述的 5 个官方检查点:
| 检查点 | 参数量 |
|---|---|
| T5-Small | 60M |
| T5-Base | 220M |
| T5-Large | 770M |
| T5-3B | 3B |
| T5-11B | 11B |
而released_checkpoints.md记录的是论文发表之后训练的实验模型检查点,文档开篇即明确两点:其一,这些模型“在某些场景下可以取得更好的性能”;其二,这些检查点不受官方支持,使用者需自担风险("These checkpoints are not officially supported - use at your own risk!")。这一点是选用它们前必须知晓的前提。
1.2 实验检查点家族的总体构成
文档按四个子系列组织,全部托管于 Google Cloud Storage 的gs://t5-data/pretrained_models/目录:
- t5.1.1.*:对论文模型做系统性架构改进的 5 个尺寸(small / base / large / xl / xxl);
- t5.1.1.lm100k.*:为 prompt tuning 优化的 “LM 适配”版本(5 个尺寸);
- t5.1.th.*:采用 talking-heads 注意力机制的变体(base / large);
- t5.1.n4w10.*:采用“前 4 层组窄、后 10 层组宽”非均匀层宽设计的变体(base / large / xl)。
二、t5.1.1 系列:对论文模型的核心架构改进
2.1 文档声明的五项改进
t5.1.1.*系列在论文模型(t5.1.0.*)基础上做了五处改进,released_checkpoints.md原文与仓库 gin 配置注释完全对应:
- FFN 隐藏层激活从 ReLU 换成 GEGLU(GLU 门控线性单元 + GELU 的组合,参见 GLU 论文 2002.05202);
- 预训练阶段关闭 dropout(质量收益),微调时再重新开启;
- 预训练数据仅使用 C4,不再混入下游任务;
- embedding 层与 classifier(softmax)层不再共享参数;
xl/xxl取代原先的3B/11B命名,且模型形状不同——d_model 更大、num_heads 与 d_ff 更小。
2.2 源码级验证:t5.1.1.base.gin 是改进的核心模板
仓库中 t5/models/gin/models/t5.1.1.base.gin 是这一系列所有改进的集中体现,其注释与文档逐条对应:
# Some improvements on t5.1.0.base: # - GEGLU-activated feedforward layers instead of RELU # Note that d_ff goes down by a factor of 2/3 to compensate for the extra projection. # - no dropout on training # - no sharing of embedding weights with classifier layer. include 'models/t5.1.0.base.gin' d_ff = 2048 DenseReluDense.activation = ["gelu", "linear"] dropout_rate = 0.0 Unitransformer.shared_embedding_and_softmax_weights = False逐行解读:
DenseReluDense.activation = ["gelu", "linear"]:FFN 层由原先的 “ReLU 激活 + 线性投影” 改为GEGLU结构——第一个投影使用 GELU,第二个使用线性投影,两者相乘实现门控效果。该配置项在 t5/models/gin/models/t5.1.1.base.gin 中显式写出;d_ff = 2048:相比t5.1.0.base(经 bi_v1.gin 继承 bi_bert_base.gin,d_ff 为 3072),按注释“缩减为 2/3 以补偿 GEGLU 额外的投影参数”(3072 × 2/3 ≈ 2048);dropout_rate = 0.0:预训练关闭 dropout;Unitransformer.shared_embedding_and_softmax_weights = False:关闭 embedding 与输出 softmax 层的参数共享;include 'models/t5.1.0.base.gin':以论文 base 模型为基底,仅覆盖上述差异,其余超参数继承。
2.3 五个尺寸的架构参数一览
结合t5.1.1.*.gin各文件(small、base、large、xl、xxl),各尺寸关键超参数与文档标注的参数量如下:
| 检查点 | 文档参数量 | d_model | num_layers | d_ff | num_heads | 模型并行度 |
|---|---|---|---|---|---|---|
| t5.1.1.small | ~77M | 512 | 8 | 1024 | 6 | — |
| t5.1.1.base | ~250M | 768(继承 t5.1.0.base) | 12 | 2048 | 12 | — |
| t5.1.1.large | ~800M | 1024 | 24 | 2816 | 16 | 2 |
| t5.1.1.xl | ~3B | 2048 | 24 | 5120 | 32 | 8 |
| t5.1.1.xxl | ~11B | 4096 | 24 | 10240 | 64 | 8 |
几点值得注意的细节:
- small 的 d_ff 大幅缩减:t5.1.1.small.gin 中 d_ff 直接设为 1024(相比
t5.1.0.small的 2048 明显下降),这正是 GEGLU 补偿策略的体现; - xl/xxl 的形状变化可直接与 t5.1.0 对比验证:
t5.1.0.3B(bi_v1_3B.gin)为 d_model=1024、d_ff=16384、num_heads=32;而t5.1.1.xl为 d_model=2048、d_ff=5120、num_heads=32——d_model 翻倍(更大)、d_ff 大幅缩减(更小)、num_heads 相对 d_model 的比例更低。同理t5.1.0.11B(bi_v1_11B.gin)为 d_model=1024、d_ff=65536、num_heads=128,而t5.1.1.xxl为 d_model=4096、d_ff=10240、num_heads=64,完全符合文档“larger d_model and smaller num_heads and d_ff”的描述; - large/xl/xxl 显式设置了模型并行度:large 为 2 路、xl 为 8 路、xxl 为 8 路(t5.1.1.xxl.gin 还额外设置了
tokens_per_microbatch_per_replica=1024以控制序列化微批的 token 数); - 仓库中还附带一个 t5.1.1.tiny.gin(d_model=32、num_layers=2、d_ff=64、num_heads=2),注释明确说明其用途是“主要用于测试”,并非正式发布的检查点配置。
2.4 检查点存储位置
文档给出的完整检查点路径(GCS bucket):
gs://t5-data/pretrained_models/t5.1.1.smallgs://t5-data/pretrained_models/t5.1.1.basegs://t5-data/pretrained_models/t5.1.1.largegs://t5-data/pretrained_models/t5.1.1.xlgs://t5-data/pretrained_models/t5.1.1.xxl
每个目录下包含模型权重文件与operative_config.gin(记录该检查点训练时的完整 gin 配置),后者是后续微调时的关键输入。
三、t5.1.1.lm100k 系列:面向 Prompt Tuning 的 LM 适配检查点
3.1 设计动机与训练方式
released_checkpoints.md说明:LM-Adapted 系列从 t5.1.1 检查点初始化,继续在论文中讨论的 LM(language modeling)目标上额外训练 100K 步(这也是lm100k命名的由来)。这一适配显著提升了模型用于prompt tuning(提示微调,参见 prompt tuning 论文 2104.08691)的能力——LM 目标让模型更擅长以自回归方式延续给定的前缀文本,从而在“仅优化少量提示参数、冻结模型主体”的范式下表现更好。
五个尺寸与参数量如下:
| 检查点 | 参数量 |
|---|---|
| t5.1.1.lm100k.small | ~77M |
| t5.1.1.lm100k.base | ~250M |
| t5.1.1.lm100k.large | ~800M |
| t5.1.1.lm100k.xl | ~3B |
| t5.1.1.lm100k.xxl | ~11B |
存储路径为gs://t5-data/pretrained_models/t5.1.1.lm100k.{small,base,large,xl,xxl}。
3.2 仓库中的 LM 目标实现佐证
仓库内虽未单独提供lm100k的 gin 文件,但 LM 目标的实现链路完整可查:
- 模型侧: t5/models/gin/models/lm_v1.gin 将
utils.run.model_type设为"lm",即使用单向语言模型(decoder-only 形式);其变体 lm_v1_ifa.gin 还额外开启Unitransformer.input_full_attention = True; - 数据侧: t5/models/gin/objectives/lm.gin 定义了无监督 LM 目标的完整预处理链——先
select_random_chunk(随机选取最长 65536 token 的文本块),再reduce_concat_tokens(拼接文档 token),最后split_tokens_to_targets_length(按目标长度切块),使模型把“文本块直接当作 targets 进行自回归预测”,不区分 inputs 与 targets; - 对应实现位于 t5/data/preprocessors.py(
select_random_chunk、reduce_concat_tokens、split_tokens_to_targets_length)。
理解这条链路有助于判断 lm100k 检查点的行为:它本质上是把 T5 从“双向编码器-解码器”进一步训练成更强的自回归延续器,因此特别契合 prompt tuning 这类“给定提示前缀、续写答案”的使用方式。
四、t5.1.th 系列:Talking-Heads 注意力变体
4.1 什么是 talking-heads 注意力
t5.1.th.*是t5.1.1模型引入talking-heads attention(见论文 2003.02436)的变体。普通多头注意力中,各头的 attention 权重独立作用于各自的 value 投影;talking-heads 机制则在头之间增加了一个可学习的“交换”混合层——即 softmax 之后的注意力权重先经过一个 heads×heads 的矩阵与各头的信息做混合,再与 value 相乘,相当于让不同头之间可以“对话”,从而提升注意力表达的灵活性。
4.2 配置源码解读
t5/models/gin/models/t5.1.th.base.gin 完整展现了该变体的实现方式:
include 'models/t5.1.1.base.gin' encoder/transformer.make_layer_stack.layers = [ @mesh_tensorflow.transformer.transformer_layers.TalkingHeadsSelfAttention, @mesh_tensorflow.transformer.transformer_layers.DenseReluDense, ] decoder/transformer.make_layer_stack.layers = [ @mesh_tensorflow.transformer.transformer_layers.TalkingHeadsSelfAttention, @mesh_tensorflow.transformer.transformer_layers.TalkingHeadsEncDecAttention, @mesh_tensorflow.transformer.transformer_layers.DenseReluDense, ] num_heads = 24 TalkingHeadsSelfAttention.key_heads_dims = [("key_heads", %num_heads)] TalkingHeadsSelfAttention.softmax_heads_dims = [("heads", %num_heads)] TalkingHeadsSelfAttention.value_heads_dims = [("value_heads", %num_heads)] TalkingHeadsSelfAttention.key_size = 32 TalkingHeadsSelfAttention.value_size = 32 TalkingHeadsSelfAttention.dropout_rate = %dropout_rate TalkingHeadsSelfAttention.relative_attention_type = "bias_shared"关键点:
- 编码器每层替换为
TalkingHeadsSelfAttention+DenseReluDense,解码器每层为TalkingHeadsSelfAttention+TalkingHeadsEncDecAttention+DenseReluDense; - 通过
key_heads_dims、softmax_heads_dims、value_heads_dims三个维度声明,将 key / softmax / value 的头维度分离(这正是 talking-heads 混合矩阵作用的维度空间),key_size 与 value_size 均为 32; relative_attention_type = "bias_shared"沿用 T5 一贯的相对位置注意力(bias 跨层共享,参见 bi_v1.gin 的注释说明);- base 版本使用 24 个 talking heads;t5.1.th.large.gin 将 d_model 提升到 1024、num_heads 提升到 32(32 个 talking heads),并刻意不启用模型并行——配置注释明确警告:talking-heads 配合模型并行会非常慢,且该配置“大概率只能在 TPUv3 上运行,TPUv2 因内存限制可能无法运行”。
4.3 检查点规模与位置
| 检查点 | 参数量 | 存储路径 |
|---|---|---|
| t5.1.th.base | ~250M | gs://t5-data/pretrained_models/t5.1.th.base |
| t5.1.th.large | ~800M | gs://t5-data/pretrained_models/t5.1.th.large |
五、t5.1.n4w10 系列:前窄后宽的非均匀层宽设计
5.1 架构思想:14 个层组,前 4 窄、后 10 宽
t5.1.n4w10.*同样是t5.1.1的变体,命名即架构:n4w10 = narrow first 4, wide last 10。编码器与解码器各由 14 个“层组”(layer group)组成,其中前 4 个层组较窄、后 10 个层组宽度加倍(d_ff 与 num_heads 翻倍),同时保持总体参数量与计算量与对应的t5.1.1模型相当。
这一非均匀结构带来的质量与效率权衡在文档中写得很直白:
- 对base模型:层数从
t5.1.1.base的 12 层增加到 14 层(层数变多),质量提升; - 对large / xl模型:层数从 24 层减少到 14 层(层数变少),质量下降,但大幅减少模型并行所需的通信量——因为模型并行的核心开销与层数(即跨设备切分点数量)直接相关,更少的层数意味着更少的跨设备通信。
5.2 配置源码:显式层组列表
t5/models/gin/models/t5.1.n4w10.base.gin 通过encoder/transformer.make_layer_stack.layers与decoder/transformer.make_layer_stack.layers显式列出全部 28 个组件(每 2 个组件构成 1 个层组,编码器每组为 SelfAttention + DenseReluDense,解码器每组为 SelfAttention + EncDecAttention + DenseReluDense)。以 base 编码器为例,前 4 组被显式收窄:
({"num_heads": 6}, @...SelfAttention), # 前 4 组:窄 ({"hidden_size": 1024}, @...DenseReluDense), ({"num_heads": 6}, @...SelfAttention), ({"hidden_size": 1024}, @...DenseReluDense), ...(共 4 组窄层组,共 8 个组件) @...SelfAttention, # 后 10 组:默认宽度 @...DenseReluDense, ...(共 10 组,共 20 个组件)后 10 组不再显式覆盖参数,因此继承全局默认值(base 为 num_heads=6、d_ff=2048,来自include 'models/t5.1.1.base.gin'),恰好是前 4 组(FFN 中间维度 1024,即默认 2048 的一半)的两倍宽度。
large / xl 版本在 t5.1.n4w10.large.gin 与 t5.1.n4w10.xl.gin 中遵循同一模式:
| 检查点 | d_model | 前 4 组 num_heads / FFN 维度 | 后 10 组 num_heads / FFN 维度 | 模型并行度 |
|---|---|---|---|---|
| t5.1.n4w10.base | 512 | 6 / 1024 | 6 / 2048 | — |
| t5.1.n4w10.large | 1024 | 16 / 2560 | 32 / 5120 | 2 |
| t5.1.n4w10.xl | 2048 | 32 / 5120 | 64 / 10240 | 8 |
可见 large / xl 版本中前后层组的 num_heads 与 FFN 维度均为严格的 2 倍关系,与文档描述一致。
5.3 检查点规模与位置
| 检查点 | 参数量 | 存储路径 |
|---|---|---|
| t5.1.n4w10.base | ~250M | gs://t5-data/pretrained_models/t5.1.n4w10.base |
| t5.1.n4w10.large | ~800M | gs://t5-data/pretrained_models/t5.1.n4w10.large |
| t5.1.n4w10.xl | ~3B | gs://t5-data/pretrained_models/t5.1.n4w10.xl |
六、如何加载与微调这些实验检查点
6.1 Mesh TensorFlow 路径:operative_config.gin 是加载的钥匙
按照 README.md 的微调说明,加载任意预训练检查点的核心机制是:把该检查点目录下的operative_config.gin作为 gin 配置传入训练脚本。它记录了训练时的完整架构与超参数(包括正确的预训练检查点路径),从而保证微调时模型结构与预训练权重严格一致。
以微调t5.1.1.base到 GLUE MRPC 任务为例(Mesh TensorFlow + TPU 路径):
t5_mesh_transformer \ --tpu="${TPU_NAME}" \ --gcp_project="${PROJECT}" \ --tpu_zone="${ZONE}" \ --model_dir="${MODEL_DIR}" \ --t5_tfds_data_dir="${DATA_DIR}" \ --gin_file="dataset.gin" \ --gin_param="utils.tpu_mesh_shape.model_parallelism = 1" \ --gin_param="utils.tpu_mesh_shape.tpu_topology = '${TPU_SIZE}'" \ --gin_param="MIXTURE_NAME = 'glue_mrpc_v002'" \ --gin_file="gs://t5-data/pretrained_models/t5.1.1.base/operative_config.gin"针对本文介绍的各实验检查点,把最后一行替换为对应 bucket 路径下的operative_config.gin即可,例如:
--gin_file="gs://t5-data/pretrained_models/t5.1.1.lm100k.base/operative_config.gin" --gin_file="gs://t5-data/pretrained_models/t5.1.th.base/operative_config.gin" --gin_file="gs://t5-data/pretrained_models/t5.1.n4w10.base/operative_config.gin"两点微调实践提示(同样来自 README,适用于所有检查点):
- dropout 需要在微调时重新开启:t5.1.1 系检查点在预训练阶段
dropout_rate = 0.0,文档明确要求“微调时应重新启用 dropout”。从源码结构看,可通过在命令行追加--gin_param="dropout_rate = 0.1"覆盖预训练 operative config 中的取值; - 训练步数需要显式指定:预训练检查点的 operative config 对训练步数“实际上不设上限”,且模型已训练 1,000,000 步;若想额外微调 10,000 步,需传入
--gin_param="run.train_steps = 1010000"。
6.2 用pretrained_model_dir从指定检查点续训
除了 operative config 内嵌的路径,t5_mesh_transformer二进制还提供显式的预训练目录参数: t5/models/mesh_transformer_main.py 中定义了pretrained_model_dirflag,运行时会解析该目录下最新(或指定 step)的 checkpoint 来初始化权重(对应 t5/models/hf_model.py 中pretrained_checkpoint_step=-1表示取最新检查点的逻辑)。例如从t5.1.1.xl指定步数恢复:
--pretrained_model_dir="gs://t5-data/pretrained_models/t5.1.1.xl" --init_checkpoint="gs://t5-data/pretrained_models/t5.1.1.xl/model.ckpt-1000000"6.3 Hugging Face 路径:PyTorch 单卡加载
若想在 PyTorch 单 GPU 环境快速试用,README 指出可使用实验性的HfPyTorchModel接口(t5/models/hf_model.py),它通过transformers.T5ForConditionalGeneration.from_pretrained(pretrained_model_name_or_path)加载模型。使用时把pretrained_model_name_or_path指向gs://bucket 中的检查点目录(或先同步到本地目录)即可完成权重加载,适合快速验证t5.1.1.small/t5.1.1.lm100k.small等较小模型在 prompt tuning 场景下的效果。
七、选型建议与注意事项
综合文档声明与仓库源码,可以给出如下选型视角(均为仓库内容可支持的推断,而非性能承诺):
- 追求微调质量上限:优先
t5.1.1.*系——GEGLU、无预训练 dropout、不共享 embedding/softmax 权重等改进在文档中被明确表述为质量取向;t5.1.1.lm100k.*则在“prompt tuning”场景有针对性优势; - 研究注意力机制:
t5.1.th.*是观察 talking-heads 注意力对质量影响的直接样本,但注意其 large 版本不支持模型并行、且对 TPU 硬件(TPUv3)有约束; - 受限于模型并行通信带宽:
t5.1.n4w10.*的 large/xl 用少量质量换取大幅降低的跨设备通信量,适合多卡/多片并行训练环境; - 风险提示:这些检查点全部标注为“不受官方支持,使用自担风险”;
t5.1.1.tiny仅为测试用,不应视为正式检查点;xl/xxl 级别的加载与微调对 TPU 内存与并行配置有较高要求(如 xxl 的tokens_per_microbatch_per_replica=1024等序列化微批设置)。
八、结语
released_checkpoints.md虽然篇幅不长,却浓缩了 T5 在论文之外最值得研究的四组实验模型:架构级改进(GEGLU、无 dropout 预训练、参数解绑、非均匀层宽)、训练目标适配(LM 100K 步)、注意力机制变体(talking-heads)以及模型并行友好的形状设计。仓库内的 t5/models/gin/models/ 目录为每一处文档描述提供了可逐行验证的配置源码,读者可将本文与对应 gin 文件对照阅读,从而真正理解这些检查点“是什么、为什么、怎么用”,并在自己的微调或 prompt tuning 项目中做出有依据的选型。
- 人工智能
- 大模型
- NLP
- 预训练
- 微调
【免费下载链接】text-to-text-transfer-transformer
Code for the paper "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer"
相关推荐
Megatron-LM T5 220M 分布式预训练实战指南:Slurm 启动、架构配置与训练结果解析
Megatron LM T5 220M 分布式预训练实战指南:Slurm 启动、架构配置与训练结果解析 本文基于 Megatron LM 仓库中的 T5 训练示
人工智能大模型预训练分布式训练深度学习强化学习Megatron-LM ModelOpt 进阶实战:配置体系、Slurm 集群训练与检查点管理
Megatron LM ModelOpt 进阶实战:配置体系、Slurm 集群训练与检查点管理 本文面向已经跑通 TensorRT Model Optimize
人工智能大模型强化学习AI Agent微调使用 BertGeneration 让预训练 BERT / RoBERTa 检查点适配序列生成任务(transformers 实战指南)
使用 BertGeneration 让预训练 BERT / RoBERTa 检查点适配序列生成任务(transformers 实战指南) 本文围绕 transf
人工智能大模型深度学习NLP预训练微调模型推理服务
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考