T5X

请前往 T5X ReadTheDocs 文档页面。

T5X 是一个模块化、可组合且便于研究的框架，用于在多种规模下高效、可配置、自助式的序列模型（从语言模型开始）训练、评估和推理。

它本质上是基于 Mesh TensorFlow 的 T5 代码库 在 JAX 和 Flax 上的新版改进实现。欲了解更多信息，请参阅 T5X 论文。

以下是使用 Google Cloud 上的 TPU 进行模型训练的快速入门指南。更多教程和背景信息，请参阅完整文档。

快速入门（推荐）

T5X 可以通过 XManager 在 Vertex AI 上运行。Vertex AI 是一个用于训练的平台，它可以创建 TPU 实例并在这些 TPU 上运行代码。当作业结束时，Vertex AI 还会自动关闭 TPU。这比手动管理 GCE 虚拟机和 TPU 虚拟机实例要方便得多。

1. 按照先决条件和说明安装 XManager。

2. 根据需要申请 TPU 配额。GCP 项目默认提供 8 个核心，这足以在单个 TPU 主机上运行一次训练实验。如果想要进行多主机训练或并行运行多个试验，则需要更多的配额。请导航至 Quotas。

您需要的配额如下：

• 服务：Vertex AI API
• 维度（位置）：us-central1
• 如果您想运行单主机实验：
  • Custom model training TPU V2 cores per region
  • Custom model training TPU V3 cores per region
• 如果您想运行多主机实验：
  • Custom model training TPU V2 pod cores per region
  • Custom model training TPU V3 pod cores per region

提示：您无法使用多主机配额来运行单主机实验。（即，您不能使用 TPU V2 pod 来运行 tpu_v2=8）

3. 启动位于 t5x/scripts/xm_launch.py 的 xmanager 脚本。

作为示例，我们使用 WMT14 英德翻译任务，该任务在下面的示例部分中有更详细的介绍。

export GOOGLE_CLOUD_BUCKET_NAME=...
export TFDS_DATA_DIR=gs://$GOOGLE_CLOUD_BUCKET_NAME/t5x/data
export MODEL_DIR=gs://$GOOGLE_CLOUD_BUCKET_NAME/t5x/$(date +%Y%m%d)

# 在多主机实验中预先下载数据集。
tfds build wmt_t2t_translate --data_dir=$TFDS_DATA_DIR

git clone https://github.com/google-research/t5x
cd ./t5x/

python3 ./t5x/scripts/xm_launch.py \
  --gin_file=t5x/examples/t5/t5_1_1/examples/base_wmt_from_scratch.gin \
  --model_dir=$MODEL_DIR \
  --tfds_data_dir=$TFDS_DATA_DIR

请查看 gs://$GOOGLE_CLOUD_BUCKET_NAME/t5x/ 中的输出文件，这些文件可以被 TensorBoard 读取。

GPU 使用

注意：NVIDIA 已发布此仓库的更新版本，支持 H100 FP8 并大幅提升了 GPU 性能。有关详细信息和使用说明，请访问 NVIDIA Rosetta 仓库。

T5X 可以轻松地在单节点或使用 SLURM+pyxis 集群的多节点配置中运行于 GPU 上。更多说明请参阅 t5x/contrib/gpu。t5x/contrib/gpu/scripts_gpu 文件夹包含用于在 The Pile 上预训练 T5X，以及在 SQuAD 和 MNLI 上微调的示例脚本。这些脚本及相关的 gin 配置还包含了额外的 GPU 优化，以提高吞吐量。更多示例和说明可以在由 NVIDIA 维护的 NVIDIA Rosetta 仓库中找到，该仓库支持 H100 FP8 并显著提升了 GPU 性能。

安装

请注意，除非另有说明，本文档中的所有命令都应在 TPU 虚拟机实例的命令行中执行。

1. 按照 说明 设置 Google Cloud Platform (GCP) 帐户并启用 Cloud TPU API。

   注意: T5X 也支持 GPU，如果您想使用 GPU 版本，请按照 t5x/contrib/gpu 中的说明操作。

2. 创建 Cloud TPU 虚拟机实例 按照 此说明 进行操作。我们建议您先在单个 v3-8 TPU 上开发工作流（即，--accelerator-type=v3-8），待流程准备就绪后再扩展到 Pod 切片。在本 README 中，我们重点介绍如何使用单个 v3-8 TPU。有关 TPU 架构的更多信息，请参阅 此处。

3. 使用 Cloud TPU 虚拟机时，您可以直接通过 SSH 登录到 TPU 虚拟机的宿主机。您可以在宿主机上安装软件包、运行代码等。一旦 TPU 实例创建完毕，即可通过以下命令登录：

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}
   

   其中 TPU_NAME 和 ZONE 是第 2 步中使用的名称和区域。

4. 安装 T5X 和其依赖项。

   git clone --branch=main https://github.com/google-research/t5x
   cd t5x
   
   python3 -m pip install -e '.[tpu]' -f \
     https://storage.googleapis.com/jax-releases/libtpu_releases.html
   

5. 创建 Google Cloud Storage (GCS) 存储桶，用于存储数据集和模型检查点。有关创建 GCS 存储桶的说明，请参阅这些 instructions。

6. （可选）如果您更喜欢 Jupyter/Colab 风格的环境，可以按照 t5x/notebooks 中的步骤设置自定义 Colab 运行环境。

示例：英语到德语的翻译

作为一个示例，我们使用 WMT14 英-德翻译数据集。原始数据集可以在 TensorFlow Datasets 中找到，其名称为 "wmt_t2t_translate"。

T5 将诸如以下的翻译任务

{'en': 'That is good.', 'de': 'Das ist gut.'}

转换为“文本到文本”的形式：

{'inputs': 'translate English to German: That is good.', 'targets': 'Das ist gut.'}

这种表述方式使得多种不同的语言任务能够以统一的方式表达，并且单个编码器-解码器架构就可以在无需任何特定于任务的参数的情况下处理这些任务。更多细节请参阅 T5 论文（Raffel 等人，2019）。

为了构建可扩展的数据管道和评估框架，我们使用了 SeqIO，它是从 T5 库 中独立出来的。一个 seqio.Task 将原始数据集、词汇表、预处理（如分词）以及评估指标（如 BLEU [aclanthology.org/P02-1040.pdf]）封装在一起，并提供一个 tf.data 实例。

T5 库 提供了许多在 T5 论文 中使用过的 seqio.Task。在这个示例中，我们使用 wmt_t2t_ende_v003。

在开始训练或微调之前，您需要先下载 ["wmt_t2t_translate"] 数据集 (https://www.tensorflow.org/datasets/catalog/wmt_t2t_translate)。

# 用于保存处理后数据集的数据目录，格式为 "gs://data_dir"。
TFDS_DATA_DIR="..."

# 确保数据集包是最新的。
python3 -m pip install --upgrade tfds-nightly

# 预先下载数据集。
tfds build wmt_t2t_translate ${TFDS_DATA_DIR}

训练

要运行训练任务，我们使用 t5x/train.py 脚本。

# 用于保存日志、检查点等的模型目录，格式为 "gs://model_dir"。
MODEL_DIR="..."
T5X_DIR="..."  # T5X 仓库被克隆的目录。
TFDS_DATA_DIR="..."

python3 ${T5X_DIR}/t5x/train.py \
  --gin_file="t5x/examples/t5/t5_1_1/examples/base_wmt_from_scratch.gin" \
  --gin.MODEL_DIR=\"${MODEL_DIR}\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

本次训练的配置定义在 Gin 文件 base_wmt_from_scratch.gin 中。Gin-config 是一个基于依赖注入来管理配置的库。Gin 的诸多优势之一是允许用户向 T5X 库传递自定义组件（例如自定义模型），而无需修改核心库。关于如何实现这一点，请参阅“自定义组件”部分。

尽管核心库与 Gin 无关，但它是我们提供的示例的核心。因此，我们在 T5X 的上下文中提供了一个简短的 Gin 入门介绍。所有配置都写入 MODEL_DIR 中的一个文件 config.gin，这使得调试和实验复现变得更加容易。

除了 config.json 外，model-info.txt 文件还总结了模型参数（形状、轴名称、分区信息）以及优化器状态。

TensorBoard

要在 TensorBoard 中监控训练过程，由于身份验证问题，最好在您自己的机器上启动 TensorBoard，而不是在 TPU 虚拟机中启动。因此，在您通过 SSH 登录到 TPU 虚拟机的命令行中，启动 TensorBoard 并将 logdir 指向 MODEL_DIR。

# 注意：请在您的机器上运行，而不是在 TPU 虚拟机上！
MODEL_DIR="..."  # 从 TPU 虚拟机复制过来。
tensorboard --logdir=${MODEL_DIR}

或者，您也可以在 Colab 中启动 TensorBoard。在 Colab 单元格中运行

from google.colab import auth
auth.authenticate_user()

以授权 Colab 访问 GCS 存储桶，然后启动 TensorBoard。

%load_ext tensorboard
model_dir = "..."  # 从 TPU 虚拟机复制过来。
%tensorboard --logdir=model_dir

微调

我们可以利用自监督预训练的优势，从我们预训练好的模型之一进行初始化。这里我们使用 T5.1.1 Base 检查点。

# 用于保存日志、检查点等的模型目录，格式为 "gs://model_dir"。
MODEL_DIR="..."

# 用于保存处理后数据集的数据目录，格式为 "gs://data_dir"。
TFDS_DATA_DIR="..."
T5X_DIR="..."  # T5X 仓库被克隆的目录。

python3 ${T5X_DIR}/t5x/train.py \
  --gin_file="t5x/examples/t5/t5_1_1/examples/base_wmt_finetune.gin" \
  --gin.MODEL_DIR=\"${MODEL_DIR}\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

注意： 当通过标志传递字符串、字典、列表、元组值或 Bash 变量时，必须将其用引号括起来。对于字符串，需要使用转义引号 (\"<string>\")。例如： --gin.utils.DatasetConfig.split=\"validation\" 或 --gin.MODEL_DIR=\"${MODEL_DIR}\"。

Gin 使更改许多配置变得非常容易。例如，您可以更改 partitioning.PjitPartitioner.num_partitions（覆盖 base_wmt_from_scratch.gin 中的值），以调整并行策略，并将其作为命令行参数传递。

--gin.partitioning.PjitPartitioner.num_partitions=8

评估

要进行离线（即不涉及训练的）评估，可以使用 t5x/eval.py 脚本。

EVAL_OUTPUT_DIR="..."  # 用于写入评估结果的目录
T5X_DIR="..."  # t5x 被克隆的目录，例如 `${HOME}"/t5x"`.
TFDS_DATA_DIR="..."
CHECKPOINT_PATH="..."

python3 ${T5X_DIR}/t5x/eval.py \
  --gin_file="t5x/examples/t5/t5_1_1/examples/base_wmt_eval.gin" \
  --gin.CHECKPOINT_PATH=\"${CHECKPOINT_PATH}\" \
  --gin.EVAL_OUTPUT_DIR=\"${EVAL_OUTPUT_DIR}\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

推理

要进行推理，可以使用 t5x/infer.py 脚本。在这里，我们使用相同的 seqio.Task，但在推理过程中，我们不会使用目标特征，而是将它们与预测结果一起记录在一个 JSON 文件中。

INFER_OUTPUT_DIR="..."  # 用于写入推理结果的目录
T5X_DIR="..."  # t5x 被克隆的目录，例如 `${HOME}"/t5x"`.
TFDS_DATA_DIR="..."
CHECKPOINT_PATH="..."

python3 ${T5X_DIR}/t5x/infer.py \
  --gin_file="t5x/examples/t5/t5_1_1/examples/base_wmt_infer.gin" \
  --gin.CHECKPOINT_PATH=\"${CHECKPOINT_PATH}\" \
  --gin.INFER_OUTPUT_DIR=\"${INFER_OUTPUT_DIR}\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

导出为 TensorFlow Saved Model

预训练模型可以导出为 TensorFlow Saved Model，并使用 [优化的 TensorFlow 运行时] (https://cloud.google.com/vertex-ai/docs/predictions/optimized-tensorflow-runtime) 部署到 Vertex AI Prediction 服务。请注意，导出的模型无法在基于开源的 TensorFlow Model Server 上运行。

T5X_DIR="..."  # t5x 被克隆的目录，例如 ${HOME}"/t5x"。
CHECKPOINT_PATH="..."

BATCH_SIZE=None
BEAM_SIZE=1

# 如果计划在 NVIDIA A100 或更新的 GPU 上运行导出的模型，请使用 'bfloat16'；
# 对于其他 GPU，则使用 'float32'。
ACTIVATION_DTYPE=bfloat16

# 版本号必须是数字。我们根据当前日期和时间生成一个版本号。
VERSION=$(date +%Y%m%d%H%M%S)

NAME=t5x_base_${ACTIVATION_DTYPE}  # 模型名称。

# 模型导出路径。请注意，导出脚本会在模型名称后添加 _cpu 后缀。
OUTPUT=${CHECKPOINT_PATH}/saved_model.${NAME}/${VERSION}

declare -a ARGS=(
--gin_file=t5x/examples/t5/t5_1_1/base.gin
--gin_file=t5x/t5x/configs/runs/export.gin
--gin.TASK_FEATURE_LENGTHS="{'inputs': 256, 'targets': 256}"
--gin.CHECKPOINT_PATH=\"${CHECKPOINT_PATH}\"
--gin.MODEL_NAME=\"/ml/${USER}/t5x_base\"
--gin.MODEL_OUTPUT_DIR=\"${OUTPUT}\"
--gin.BEAM_SIZE=${BEAM_SIZE}
--gin.BATCH_SIZE=${BATCH_SIZE}
--gin.export_lib.save.partitioner=None
--gin.export_lib.save.warmup_examples="['hello world']"
--gin.export_lib.ExportableModule.use_batch_function=False
--gin.export_lib.ExportableModule.use_gpu=False
--gin.export_lib.ExportableModule.jit_compile=False
--gin.ACTIVATION_DTYPE=\"${ACTIVATION_DTYPE}\"
--gin.network.T5Config.dtype=\"${ACTIVATION_DTYPE}\"
--gin.utils.RestoreCheckpointConfig.dtype=\"${ACTIVATION_DTYPE}\"
--gin.DROPOUT_RATE=0.0
)

(python3 ${T5X_DIR}/t5x/export.py "${ARGS[@]}")

有关详细参数定义，请参阅 [export.gin] (t5x/configs/runs/export.gin)。

您可以在 NVIDIA A100 40GB 上运行 XL 及更小的模型，在 NVIDIA A100 80GB 上运行 XXL 模型。

自定义组件

[翻译示例] (#example-english-to-german-translation) 使用了 T5X 提供的编码器-解码器模型以及 T5 库中的数据集。本节将展示如何使用您自己的数据集和模型，并通过 Gin 配置文件进行传递。

示例：用户目录中的自定义数据集

对于此示例，我们具有以下目录结构，其中 ${HOME}/dir1/user_dir 表示包含自定义组件的用户目录。

${HOME}
└── dir1
    └── user_dir
        ├── t5_1_1_base_de_en.gin
        └── tasks.py

作为示例，让我们定义一个新的数据集。这里我们使用相同的 Translation 数据集，但我们将翻译任务的方向设置为相反方向，即从德语到英语，而不是从英语到德语。我们在 tasks.py 中定义此任务：

# ${HOME}/dir1/user_dir/tasks.py

import functools
import seqio
import tensorflow_datasets as tfds
from t5.evaluation import metrics
from t5.data import preprocessors

vocabulary = seqio.SentencePieceVocabulary(
    'gs://t5-data/vocabs/cc_all.32000/sentencepiece.model', extra_ids=100)
output_features = {
    'inputs': seqio.Feature(vocabulary=vocabulary),
    'targets': seqio.Feature(vocabulary=vocabulary)
}

seqio.TaskRegistry.add(
    'wmt_t2t_de_en_v003',
    source=seqio.TfdsDataSource(tfds_name='wmt_t2t_translate/de-en:1.0.0'),
    preprocessors=[
        functools.partial(
            preprocessors.translate,
            source_language='de', target_language='en'),
        seqio.preprocessors.tokenize,
        seqio.CacheDatasetPlaceholder(),
        seqio.preprocessors.append_eos_after_trim,
    ],
    metric_fns=[metrics.bleu],
    output_features=output_features)

在 Gin 文件中，大多数设置与 [英德翻译示例] (#example-english-to-german-translation) 中使用的设置相同。因此，我们直接引用该示例中的 Gin 文件。为了使用我们刚刚定义的“wmt_t2t_de_en_v003”任务，我们需要导入任务模块“tasks.py”。请注意，我们使用相对于用户目录的相对路径。这将作为标志指定。

# ${HOME}/dir1/user_dir/t5_1_1_base_de_en.gin
from __gin__ import dynamic_registration
import tasks  # 导入 dir1/user_dir/tasks.py 中定义的任务。

include "t5x-tmp/t5x/examples/t5/t5_1_1/examples/base_wmt_from_scratch.gin"
MIXTURE_OR_TASK_NAME = "wmt_t2t_de_en_v003"

最后，我们通过将用户目录作为 gin_search_paths 标志来启动训练，以便 Gin 文件和 Python 模块可以使用相对路径进行指定。

PROJECT_DIR=${HOME}"/dir1/user_dir"
T5X_DIR="..."  # t5x 被克隆的目录。
TFDS_DATA_DIR="..."
MODEL_DIR="..."
export PYTHONPATH=${PROJECT_DIR}

python3 ${T5X_DIR}/t5x/train.py \
  --gin_search_paths=${PROJECT_DIR} \
  --gin_file="t5_1_1_base_de_en.gin" \
  --gin.MODEL_DIR=\"${MODEL_DIR}\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

检查点

原生检查点

我们已发布许多原始 T5 模型及其变体的原生 T5X 格式检查点，以实现最高效率。请参阅 [完整列表] (https://github.com/google-research/t5x/blob/main/docs/models.md)，其中包括匹配的 Gin 配置文件。

这些检查点是从公开的 [Mesh TensorFlow 检查点] (https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511) 转换而来的。

与 Mesh TensorFlow 检查点的兼容性

使用 T5 库 训练的 Mesh TensorFlow 检查点可以直接加载到 T5X 中。例如，我们可以通过更改 INIT_CHECKPOINT Gin 宏，从 MTF 检查点初始化并重新运行微调示例。

# 用于保存日志、检查点等的模型目录，格式为“gs://model_dir”。
MODEL_DIR="..."

# 用于保存处理后的数据集的目录，格式为“gs://data_dir”。
TFDS_DATA_DIR="..."
T5X_DIR="..."  # t5x 代码库被克隆的目录。

python3 ${T5X_DIR}/t5x/train.py \
  --gin_file="t5x/examples/t5/t5_1_1/examples/base_wmt19_ende_train.gin" \
  --gin.MODEL_DIR=\"${MODEL_DIR}\" \
  --gin.MIXTURE_OR_TASK_NAME=\"wmt_t2t_ende_v003\" \
  --gin.INIT_CHECKPOINT=\"gs://t5-data/pretrained_models/t5.1.1.base/model.ckpt-1000000\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

请注意，如果大型模型使用了高并行度，直接从 Mesh TensorFlow 检查点恢复可能会效率较低。这是因为每个主机首先会加载整个模型副本，然后仅保留由模型并行度规范决定的相关切片。如果您经常运行 Mesh TensorFlow 检查点，我们建议使用 [convert_tf_checkpoint 脚本] (t5x/scripts/convert_tf_checkpoint.py) 将检查点转换为 T5X 原生格式。

引用 T5X

请使用以下 BibTeX 条目来引用 T5X。

@article{roberts2022t5x,
  url = {https://arxiv.org/abs/2203.17189},
  author = {Roberts, Adam 和 Chung, Hyung Won 和 Levskaya, Anselm 和 Mishra, Gaurav 和 Bradbury, James 和 Andor, Daniel 和 Narang, Sharan 和 Lester, Brian 和 Gaffney, Colin 和 Mohiuddin, Afroz 和 Hawthorne, Curtis 和 Lewkowycz, Aitor 和 Salcianu, Alex 和 van Zee, Marc 和 Austin, Jacob 和 Goodman, Sebastian 和 Soares, Livio Baldini 和 Hu, Haitang 和 Tsvyashchenko, Sasha 和 Chowdhery, Aakanksha 和 Bastings, Jasmijn 和 Bulian, Jannis 和 Garcia, Xavier 和 Ni, Jianmo 和 Chen, Andrew 和 Kenealy, Kathleen 和 Clark, Jonathan H. 和 Lee, Stephan 和 Garrette, Dan 和 Lee-Thorp, James 和 Raffel, Colin 和 Shazeer, Noam 和 Ritter, Marvin 和 Bosma, Maarten 和 Passos, Alexandre 和 Maitin-Shepard, Jeremy 和 Fiedel, Noah 和 Omernick, Mark 和 Saeta, Brennan 和 Sepassi, Ryan 和 Spiridonov, Alexander 和 Newlan, Joshua 和 Gesmundo, Andrea},
  title = {使用 $\texttt{t5x}$ 和 $\texttt{seqio}$ 扩展模型和数据},
  journal={arXiv 预印本 arXiv:2203.17189},
  year = {2022},
}

注意

这并非 Google 官方支持的产品

T5X 快速上手指南

T5X 是一个基于 JAX 和 Flax 构建的模块化、高性能序列模型框架，适用于大规模语言模型的训练、评估和推理。它是原始 T5 代码库（基于 Mesh TensorFlow）的现代化重构版本。

环境准备

系统要求

• 硬件: 推荐使用 Google Cloud TPU (v3-8 或更高) 或 NVIDIA GPU (需参考 NVIDIA Rosetta 仓库以获取 H100 FP8 支持)。
• 操作系统: Linux (TPU VM 或 GPU 服务器)。
• 云平台: 若使用 TPU，需拥有 Google Cloud Platform (GCP) 账号并启用 Cloud TPU API。
• 存储: 需要 Google Cloud Storage (GCS) Bucket 用于存放数据集和模型检查点。

前置依赖

• Python 3.8+
• Git
• GCP SDK (仅 TPU 用户需要)
• XManager (可选，推荐用于 Vertex AI 自动化部署)

注意: 国内开发者若无法直接访问 GCP 或 HuggingFace，请自行配置网络代理或使用相应的镜像加速服务。

安装步骤

以下命令需在 TPU VM 实例或配置好的 Linux 环境中执行。

1. 克隆仓库

   git clone --branch=main https://github.com/google-research/t5x
   cd t5x
   

2. 安装依赖包

   • TPU 环境:

     python3 -m pip install -e '.[tpu]' -f \
       https://storage.googleapis.com/jax-releases/libtpu_releases.html
     

   • GPU 环境: 请参考 t5x/contrib/gpu/README.md 或 NVIDIA Rosetta 仓库进行特定安装。

3. 准备数据存储 创建一个 GCS Bucket 用于存储数据和模型输出：

   # 替换为您的 bucket 名称
   export GOOGLE_CLOUD_BUCKET_NAME="your-bucket-name"
   export TFDS_DATA_DIR=gs://$GOOGLE_CLOUD_BUCKET_NAME/t5x/data
   export MODEL_DIR=gs://$GOOGLE_CLOUD_BUCKET_NAME/t5x/models/$(date +%Y%m%d)
   

4. 预处理数据集 (示例：WMT14 英德翻译)

   # 确保 tfds 包为最新版本
   python3 -m pip install --upgrade tfds-nightly
   
   # 下载并构建数据集
   tfds build wmt_t2t_translate --data_dir=$TFDS_DATA_DIR
   

基本使用

T5X 使用 Gin 配置文件来管理实验参数。以下是最简单的从头开始训练示例。

1. 启动训练任务

使用 t5x/train.py 脚本启动训练。以下示例使用 WMT14 数据集从头训练一个基础模型：

python3 ./t5x/train.py \
  --gin_file=t5x/examples/t5/t5_1_1/examples/base_wmt_from_scratch.gin \
  --gin.MODEL_DIR=\"${MODEL_DIR}\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

• --gin_file: 指定配置文件路径，定义了模型架构、超参数等。
• --gin.MODEL_DIR: 指定模型日志和检查点的保存路径（注意字符串需转义引号）。
• --tfds_data_dir: 指定预处理后的数据路径。

2. 监控训练进度 (TensorBoard)

建议在本地机器（而非 TPU VM）上运行 TensorBoard 以避免认证问题，前提是本地已配置好 GCP 凭证。

# 在本地终端执行，将 MODEL_DIR 替换为实际的 GCS 路径
tensorboard --logdir=${MODEL_DIR}

或者在 Colab 中运行：

from google.colab import auth
auth.authenticate_user()

%load_ext tensorboard
model_dir = "gs://your-bucket-name/t5x/models/..." # 替换为实际路径
%tensorboard --logdir=model_dir

3. 微调预训练模型 (可选)

利用自监督预训练优势，加载现有检查点进行微调：

python3 ./t5x/train.py \
  --gin_file=t5x/examples/t5/t5_1_1/examples/base_wmt_finetune.gin \
  --gin.MODEL_DIR=\"${MODEL_DIR}\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

4. 离线评估

训练完成后，使用 eval.py 对模型进行评估：

CHECKPOINT_PATH="gs://path/to/checkpoint"
EVAL_OUTPUT_DIR="gs://path/to/output"

python3 ./t5x/eval.py \
  --gin_file=t5x/examples/t5/t5_1_1/examples/base_wmt_eval.gin \
  --gin.CHECKPOINT_PATH=\"${CHECKPOINT_PATH}\" \
  --gin.EVAL_OUTPUT_DIR=\"${EVAL_OUTPUT_DIR}\" \
  --tfds_data_dir=${TFDS_DATA_DIR}

提示: 若需修改并行策略或批次大小等参数，可通过命令行直接覆盖 Gin 配置，例如： --gin.partitioning.PjitPartitioner.num_partitions=8