TensorFlow Models 中 XLNet 的 TPU 训练实战：从数据预处理到 IMDB 与 SQuAD 2.0 微调的完整流程

XLNet 是 TensorFlow Models 仓库 official/legacy/xlnet/ 目录实现的"广义自回归"预训练语言模型，它通过在所有分解顺序排列上最大化期望似然来学习双向上下文，并借助 Transformer-XL 的段级循环（memory）机制与相对位置编码突破固定窗口长度的限制。本文以 official/legacy/xlnet/README.md 为骨架，完整继承其环境搭建、数据预处理与两个经典任务（IMDB 情感分类、SQuAD 2.0 阅读理解）的官方命令与参数，并结合 common_flags.py、preprocess_pretrain_data.py 等源码逐项说明参数默认值、硬约束断言与相对位置编码、置换掩码的底层实现，帮助读者在 Cloud TPU 上复现预训练数据制备与微调全流程。

XLNet 的核心机制：论文结论与源码印证

README 原文给出的核心表述是：XLNet 是一个"广义自回归的 BERT 式预训练语言模型"，其两个关键设计直接继承自 Transformer-XL——

1. 排列语言建模（Permutation Language Modeling）：不固定"从左到右"的生成顺序，而是在所有分解顺序排列（factorization order）上最大化期望似然，从而像 BERT 一样看到双向上下文；
2. 段级循环（segment-level recurrence）与相对位置编码：让模型可以学习固定窗口之外的依赖关系，同时不打乱时间连贯性。

README 同时引用了原论文的结论：XLNet 在 20 个 NLP 基准任务上优于 BERT，其中 18 个任务取得当时 SOTA，覆盖问答、自然语言推理、情感分析与文档排序。

这些结论在仓库源码中都能找到对应实现：

• 相对位置编码：xlnet_modeling.py 中使用正弦/余弦构造位置向量 pos_emb = concat([sin, cos], -1)，即 Transformer-XL 风格的相对位置编码，而非 BERT 的绝对位置嵌入；
• 相对注意力：xlnet_modeling.py 的 RelativeAttention.call 把注意力得分拆成三部分：基于内容项 ac = einsum(q + r_w_bias, k_head_h)、基于相对位置项 bd = einsum(q + r_r_bias, k_head_r)（经 rel_shift 相对移位）以及基于 segment 的项 ef，最终 attn_score = (ac + bd + ef) * scale；untie_r 参数控制 r_w_bias/r_r_bias 是否在层间解绑（见 xlnet_modeling.py，untie_r=true 时按层使用 self.r_w_bias[i]）；
• 前馈层激活：xlnet_modeling.py 的 PositionwiseFF 只接受 relu 或 gelu 两种激活（否则抛出 ValueError），这解释了后面所有微调命令为何都显式指定 --ff_activation=gelu——必须与预训练检查点一致。

环境搭建（Set Up）

README 推荐的目标运行环境是 Cloud TPU。官方步骤如下：

第一步：使用 ctpu 工具创建 tf-nightly 版本的 TPU 实例：

ctpu up -name <instance name> --tf-version="nightly"

第二步：SSH 登录该 VM（或在自有机器上）后，把仓库根目录加入 PYTHONPATH，保证 official.legacy.xlnet.* 包可被导入：

export PYTHONPATH="$PYTHONPATH:/path/to/models"

第三步：安装 tf-nightly 以获得最新 TPU 支持：

pip install tf-nightly-gpu

从源码结构看，所有入口脚本（如 run_classifier.py）都通过 official.common.distribute_utils 获取分布式策略、通过 official.legacy.xlnet 下的模块组织代码，因此 PYTHONPATH 必须指向仓库根目录而非 official/ 子目录，这一点与 README 的写法一致。

数据预处理（Process Datasets）

所有预处理脚本都依赖一个 SentencePiece 模型。README 指出，可用模型存放在公开 GCS 桶 gs://cloud-tpu-checkpoints/xlnet/cased_spiece.model（注意是带大小写的 cased 版本）。同时强调：要在 Cloud TPU 上训练，数据必须存放在 GCS 桶上。统一的环境变量准备命令为：

export SPIECE_DIR=~/cased_spiece/
export SPIECE_MODEL=${SPIECE_DIR}/cased_spiece.model
export DATASETS_DIR=gs://some_bucket/datasets
mkdir -p ${SPIECE_DIR}
gsutil cp gs://cloud-tpu-checkpoints/xlnet/cased_spiece.model ${SPIECE_DIR}

预训练数据制备

预训练语料通过 preprocess_pretrain_data.py 转换为 TFRecord。输入要求是纯文本文件（或纯文本文件 glob），每行一句话。README 给出的官方命令：

export INPUT_GLOB='path/to/wiki_cased/*.txt'

python3 preprocess_pretrain_data.py --bsz_per_host=32 --num_core_per_host=16 \
--seq_len=512 --reuse_len=256 --input_glob='path/to/wiki_cased/*.txt' \
--save_dir=${DATASETS_DIR}/pretrain --bi_data=True --sp_path=${SPIECE_MODEL} \
--mask_alpha=6 --mask_beta=1 --num_predict=85

README 特别警告：为了让 memory（段级循环复用）机制正确工作，bsz_per_host 和 num_core_per_host 在制备 TFRecord 时被严格固定，训练时必须使用相同的 TPU 配置。源码印证了这一约束——preprocess_pretrain_data.py 中 create_data 首先断言 bsz_per_host % num_core_per_host == 0；开启 bi_data 时 create_tfrecords 还要求 bsz_per_host % (2 * num_core_per_host) == 0（因为批量的一半是反向数据流）。

关键参数说明（默认值取自 preprocess_pretrain_data.py 的 define_flags）：

参数	默认值	含义
bsz_per_host	32	每 host 的总 batch size
num_core_per_host	8	每 host 的 TPU core 数
seq_len	512	序列长度
reuse_len	256	下一 batch 复用的 token 数（即 memory 长度），通常为 seq_len 的一半
bi_data	True	是否生成双向数据流（正序+反序）
mask_alpha	6	每组掩码 span 的 token 分组数
mask_beta	1	每组中掩码的 token 数
num_predict	85	每个样本进行部分预测（partial prediction）的 token 数
input_glob	—	输入纯文本文件 glob
sp_path	—	SentencePiece 模型路径
save_dir	—	TFRecord 输出目录
split	train	数据划分（train/dev/test）
pass_id / num_task / task	0 / 1 / 0	多 worker 并行制备时切分文件用

部分预测（span masking）的实现：_sample_mask（preprocess_pretrain_data.py）按 1/n（n 取 1~5，归一化）的概率随机选 n-gram 长度，依据 ctx_size = (n * mask_alpha) // mask_beta 决定 span 两侧保留多少可见上下文，因此 mask_alpha=6, mask_beta=1 时掩码 span 两侧最多保留约 6n 个可见 token。当 --num_predict=85 指定时，每个样本恰好掩码 85 个 token（不足时随机补齐），这与命令参数一一对应。

样本结构：每个 TFRecord Example 包含 5 个特征键（见 preprocess_pretrain_data.py）：

• input：[reuse 段 inp] + [A 段] + [SEP] + [B 段] + [SEP] + [CLS] 拼接后的 id 序列；
• target：对应的预测目标（右移一位）；
• seg_id：segment 编号（A 段为 0，B 段为 1，CLS 为 2）；
• label：下一句关系标签（A/B 来自相邻句时为 1，否则为 0，见 _split_a_and_b）；
• is_masked：被选中做部分预测的 token 掩码。

词表固定为 32000（VOCAB_SIZE = 32000），特殊符号含 <eod>（文档结尾）、<eop> 等（preprocess_pretrain_data.py）。输出 TFRecord 的命名规则由 format_filename 生成，形如 train-0-0.bsz-32.seqlen-512.reuse-256.bi.alpha-6.beta-1.fnp-85.tfrecords，训练端按此规则 glob 匹配（preprocess_pretrain_data.py），所以制备与训练参数不一致时会直接找不到数据。

分类数据制备（IMDB）

按 README 官方步骤下载并解压 aclImdb 数据集：

export IMDB_DIR=~/imdb
mkdir -p ${IMDB_DIR}

cd ${IMDB_DIR}
wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
tar zxvf aclImdb_v1.tar.gz -C ${IMDB_DIR}
rm aclImdb_v1.tar.gz

再调用 preprocess_classification_data.py 转成 TFRecord：

export TASK_NAME=imdb

python3 preprocess_classification_data.py --max_seq_length=512 --spiece_model_file=${SPIECE_MODEL} --output_dir=${DATASETS_DIR}/${TASK_NAME} --data_dir=${IMDB_DIR}/aclImdb --task_name=${TASK_NAME}

README 明确指出：要在 IMDB 上取得 SOTA，必须使用序列长度 512（该脚本 max_seq_length 默认值仅为 128，见 preprocess_classification_data.py）。该脚本还支持 --num_passes（无损耗切批，供 TPU 使用）、--eval_split、--uncased、--is_regression 等参数。

SQuAD 2.0 数据制备

先下载 SQuAD v2.0 的 train/dev json（SQuAD 官方网站提供了数据集与评测的完整说明）：

export SQUAD_DIR=~/squad

mkdir -p ${SQUAD_DIR} && cd ${SQUAD_DIR}
wget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json
wget https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v2.0.json

再执行转换（README 官方命令）：

python3 preprocess_squad_data.py --spiece_model_file=${SPIECE_MODEL} --train_file=${SQUAD_DIR}/train-v2.0.json --predict_file=${SQUAD_DIR}/dev-v2.0.json --output_dir=${DATASETS_DIR}/squad --uncased=False --max_seq_length=512 --num_proc=1 --proc_id=0

gsutil cp ${SQUAD_DIR}/dev-v2.0.json ${DATASETS_DIR}/squad

最后一条 gsutil cp 是必需的：评测时要用原始 dev json 做答案抽取与 F1 计算，因此必须随 TFRecord 一起上传到 GCS。preprocess_squad_data.py 的参数默认值为：max_seq_length=512、max_query_length=64、doc_stride=128；输出的训练文件命名为 {spiece 模型名}.{proc_id}.slen-{max_seq_length}.qlen-{max_query_length}.train.tf_record，这正是后文训练命令中 squad_cased 目录下的文件来源。--num_proc/--proc_id 用于把训练样本按 train_examples[proc_id::num_proc] 切给多进程并行制备。

微调（Fine-tuning with XLNet）

预训练检查点

解压后的预训练模型存放在 GCS 目录 gs://cloud-tpu-checkpoints/xlnet/keras_xlnet：

export XLNET_DIR=gs://cloud-tpu-checkpoints/xlnet/keras_xlnet
export MODEL_DIR=gs://some_bucket/my_output_dir

（README 原文示例写作 gs:/cloud-tpu-checkpoints/xlnet/keras_xlnet，注意 GCS 规范路径应为 gs:// 双斜杠。）

微调时的模型结构超参必须与检查点匹配。从 xlnet_config.py 的 XLNetConfig 可以看到，构成模型"指纹"的参数共 8 个：n_layer、d_model、n_head、d_head、d_inner、ff_activation、untie_r、n_token，源码注释明确说明"这些超参在预训练与微调间必须相同"。官方 LARGE 检查点 cased_L-24_H-1024_A-16 即对应 n_layer=24, d_model=d_embed=1024, n_head=16。

分类任务：IMDB 微调

README 给出的官方预期：在 v3-8 TPU 上约 11 分钟得到前 500 步结果，全程约 1 小时，最终精度介于 96.15% 与 96.33%。完整命令（v3-8）：

export TPU_NAME=my-tpu

python3 run_classifier.py \
--strategy_type=tpu \
--tpu=${TPU_NAME} \
--init_checkpoint=${XLNET_DIR}/xlnet_model.ckpt \
--model_dir=${MODEL_DIR} \
--test_data_size=25024 \
--train_tfrecord_path=${DATASETS_DIR}/imdb/cased_spiece.model.len-512.train.tf_record \
--test_tfrecord_path=${DATASETS_DIR}/imdb/cased_spiece.model.len-512.dev.eval.tf_record \
--train_batch_size=32 \
--seq_len=512 \
--n_layer=24 \
--d_model=1024 \
--d_embed=1024 \
--n_head=16 \
--d_head=64 \
--d_inner=4096 \
--untie_r=true \
--n_class=2 \
--ff_activation=gelu \
--learning_rate=2e-5 \
--train_steps=4000 \
--warmup_steps=500 \
--iterations=500 \
--bi_data=false \
--summary_type=last

逐组参数解读（默认值取自 common_flags.py 与 run_classifier.py）：

• 运行与检查点：--strategy_type=tpu + --tpu 指定 TPU 策略（strategy_type 默认就是 tpu）；--init_checkpoint 指向预训练 xlnet_model.ckpt；--model_dir 为输出目录。
• 数据：--test_data_size=25024 即 IMDB dev 集样本数，run_classifier.py 用 test_data_size / test_batch_size 计算评估步数；训练/测试 tfrecord 路径对应前面 IMDB 制备的输出文件名。
• 模型结构：n_layer/d_model/d_embed/n_head/d_head/d_inner/untie_r/ff_activation 八元组锁定 LARGE 检查点（d_inner=4096 为 4 倍 d_model）。
• 任务头：--n_class=2（IMDB 二分类）；--summary_type=last 用序列最后一个位置的表示作为文档向量（run_classifier.py 中的 summary_type flag，默认 last）。
• 优化：--learning_rate=2e-5、--train_steps=4000、--warmup_steps=500、--iterations=500（每 500 步触发一次评估循环）。从 run_classifier.py 可见，优化器由 optimization.create_optimizer 按 learning_rate、总步数与 warmup 步数构造（默认带线性 warmup 与余弦/线性衰减）。
• --bi_data=false：与预训练（bi_data=True）不同，微调时不用双向数据流；bi_data 只在预训练配置中生效（见 xlnet_config.py 的 create_run_config：is_finetune=True 时不注入 bi_data/reuse_len/mem_len 等字段）。

评估的实现细节值得注意：验证集中包含为对齐 batch 而补入的假样本，run_classifier.py 因此先用 is_real_example 掩码过滤，再把各副本的 logits/labels 收集到本地用 NumPy 统计 acc = correct/total，并写入 eval_acc 的 TensorBoard 摘要。

SQuAD 2.0 任务

SQuAD（Stanford Question Answering Dataset）是经典的阅读理解基准。README 以 XLNet-LARGE（cased_L-24_H-1024_A-16）在 v3-8 上为例，预期：读 pickle 约 5 分钟、前 1000 步约 18 分钟、全程约 2 小时，best_f1 落在 88.30 ~ 88.80 区间。完整命令：

export TPU_NAME=my-tpu

python3 run_squad.py \
  --strategy_type=tpu \
  --tpu=${TPU_NAME} \
  --init_checkpoint=${XLNET_DIR}/xlnet_model.ckpt \
  --model_dir=${MODEL_DIR} \
  --train_tfrecord_path=${DATASETS_DIR}/squad/squad_cased \
  --test_tfrecord_path=${DATASETS_DIR}/squad/squad_cased/12048.eval.tf_record \
  --test_feature_path=${DATASETS_DIR}/squad/spiece.model.slen-512.qlen-64.eval.features.pkl \
  --predict_dir=${MODEL_DIR} \
  --predict_file=${DATASETS_DIR}/squad/dev-v2.0.json \
  --train_batch_size=48 \
  --seq_len=512 \
  --reuse_len=256 \
  --mem_len=0 \
  --n_layer=24 \
  --d_model=1024 \
  --d_embed=1024 \
  --n_head=16 \
  --d_head=64 \
  --d_inner=4096 \
  --untie_r=true \
  --ff_activation=gelu \
  --learning_rate=.00003 \
  --train_steps=8000 \
  --warmup_steps=1000 \
  --iterations=1000 \
  --bi_data=false \
  --query_len=64 \
  --adam_epsilon=.000001 \
  --lr_layer_decay_rate=0.75

相对分类任务，SQuAD 特有的参数及其作用（默认值取自 run_squad.py 与 common_flags.py）：

参数	取值	说明
--train_batch_size	48	跨所有 host 的总训练 batch（默认 16）
--seq_len / --reuse_len	512 / 256	与数据制备时的 slen-512 一致；reuse_len 是 segment 复用的 token 数
--mem_len	0	额外缓存的历史步数（默认 0），微调时通常关闭 memory
--query_len	64	问题最大长度，对应数据文件的 qlen-64
--learning_rate	3e-5	略高于分类任务的 2e-5
--train_steps / --warmup_steps	8000 / 1000	总步数与线性 warmup 步数
--iterations	1000	每 1000 步做一次评估（与预期"前 1000 步出结果"对应）
--adam_epsilon	1e-6	Adam 的 epsilon（默认 1e-8）
--lr_layer_decay_rate	0.75	逐层学习率衰减：越深的层学习率越高，lr[l-1] = lr[l] * 0.75（见 common_flags.py 的 flag 说明），这是 Transformer-XL 系列经验配置

其余答案抽取相关 flag 保持默认：start_n_top=5、end_n_top=5（start/end span 的 beam 宽度）、n_best_size=5、max_answer_length=64。评测时模型对每个 dev 样本输出候选 span 概率，再由 run_squad.py 的 run_evaluation 结合 test_feature_path 中的 token 映射与 predict_file（dev json）做答案抽取并计算 F1；注意 README 提到的"读 pickle 约 5 分钟"即指加载 spiece.model.slen-512.qlen-64.eval.features.pkl 这一步。

源码级补充：置换顺序如何实现双向性

预训练数据中每个样本只固定了输入序列，"以什么顺序预测"是在训练管线中在线采样的。preprocess_pretrain_data.py 的 _local_perm 函数即 XLNet 论文"permutation factorization order"的 TF 实现：

1. 对位置索引做 random_shuffle，把序列切成若干置换块（块长 perm_size 不得超过 reuse_len，否则会数据泄漏，见 get_dataset 中的断言）；
2. 未被掩码且非功能 token（非 SEP/CLS）的位置索引设为 -1，使其"可见于所有位置、且不能看到被掩码位置"，从而杜绝信息泄漏；
3. 生成 perm_mask（注意矩阵）与 target_mask（哪些位置参与损失），并区分 LM 目标（next token）与 PLM 目标（current token）。

由于解析是 .cache().map(parser).repeat()（preprocess_pretrain_data.py），同一输入每个 epoch 都会重新采样置换顺序与掩码，这正是"在所有排列上最大化期望似然"在工程上的落地方式；源码注释也说明了为什么不缓存解析结果（在线预处理结果每次不同，缓存会导致容器 OOM）。

适用前提与注意事项

• 环境：全部流程面向 TF2 + TPU 生态（tf-nightly、gsutil、ctpu），GPU/CPU 下可运行但命令中的 --strategy_type=tpu 与 GCS 路径需相应调整；预处理脚本内部仍使用 tensorflow.compat.v1 的 TFRecord 写入接口（见 preprocess_pretrain_data.py）。
• 参数一致性铁律：预训练 TFRecord 的 bsz_per_host/num_core_per_host 必须与训练 TPU 配置一致；微调的 8 项 XLNetConfig 超参必须与检查点一致；seq_len/reuse_len/max_seq_length 等要贯穿数据制备与训练两端，否则文件名 glob 匹配失败或张量形状断言失败。
• 该模块位于 official/legacy/ 目录：仓库将其归入 legacy 代码，接口基于较旧的 TF 2.0 早期 API 风格；新实验建议评估仓库中非 legacy 的 NLP 模块，本文内容以该目录当前代码为准。
• README 中的精度/时长（IMDB 96.15~96.33%、SQuAD 2.0 best_f1 88.30~88.80）均为官方文档在 v3-8 TPU 上的参考区间，实际结果随硬件与随机种子波动，本文不对其作进一步外推。

通过本文流程，读者可以完整走通：环境准备 → SentencePiece 与 GCS 数据布局 → 预训练 TFRecord 制备（含部分预测与段级循环约束）→ IMDB 与 SQuAD 2.0 微调，并能对照 xlnet_modeling.py、training_utils.py、optimization.py 深入理解 XLNet 的相对位置编码、置换掩码与分布式训练实现。