提示调优

这是用于复现 EMNLP 2021 论文《参数高效提示调优的规模效应》 (Lester 等, 2021) 中实验的代码。

这些模型基于 T5X 构建，其中 T5X 定义了模型和训练循环； Flaxformer 定义了实际的模型计算；Flax 定义了底层模型层；而 Jax 则提供了实际的执行。我们的实现细节可以在此处找到： 这里。

目录

• 安装
• 训练提示
  • 在 Pod 分片上训练提示
  • 自定义依赖项
• 使用提示进行推理
• 模型配置
• 提示初始化
• 发布的模型检查点
• 发布的提示
• 额外资源
• 如何引用

安装

1. 按照 T5X 安装说明 中的前三个步骤创建一个云端 TPU 虚拟机。同时按照第 5 步创建一个 Google Cloud Storage (GCS) 存储桶。我们将使用类似 gs://{bucket-name}/path/to/item/in/bucket 的 URI 格式从该存储桶读取和写入数据。我们将在其中存储缓存的数据集以及模型检查点和结果。为了便于参考，以下是一些与 TPU 虚拟机交互的常用云命令：

# 创建云端 TPU 虚拟机
$ gcloud alpha compute tpus tpu-vm create ${TPU_NAME} \
    --zone ${ZONE} \
    --accelerator-type v3-8 \
    --version v2-alpha

# SSH 到云端 TPU 虚拟机
$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE}

# 删除云端 TPU 虚拟机
$ gcloud alpha compute tpus tpu-vm delete ${TPU_NAME} --zone ${ZONE}

2. 现在您应该已经进入了 TPU 虚拟机实例的命令行界面。克隆提示调优仓库。

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

3. 安装提示调优库。

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

如果您遇到 pip 尝试安装越来越旧的依赖版本（例如 TensorFlow），直到尝试安装版本 0.0.0 并失败的错误，请在安装命令中添加 --use-deprecated=legacy-resolver。此错误与依赖项之间的所需版本有关，这种行为通常被称为回溯。如果使用该标志，可能会安装不兼容的库版本，因此您应注意安装命令输出中的版本不匹配警告。

注意： 如果您计划修改提示调优的内部实现并需要可编辑的安装（以便在运行训练时使用克隆代码中的更改），请使用 -e 标志运行 pip；如果在安装过程中出现错误，您可能需要删除 pyproject.toml 文件。

要运行测试，请使用 [test] 选项安装包（python3 -m pip install .[test] ...），然后从克隆仓库的根目录运行 python3 -m pytest。

训练提示

训练提示类似于 T5X 微调模型； 主要区别在于我们有自己的提示调优配置文件集可供使用。

我们提供了一个演示脚本 (prompt_tuning/scripts/sst2-demo.sh)，其中包含了训练提示所需的所有部分。您可以将其作为起点，或者设置 MODEL_DIR 和 TFDS_DATA_DIR 环境变量，指向您的 Google Cloud Storage 存储桶路径，直接运行此脚本。

./prompt-tuning/prompt_tuning/scripts/sst2-demo.sh

为了提高迭代速度，我们倾向于在命令行中指定更多选项，而不是将所有配置打包到单个 gin 文件中。值得注意的几个选项：

• --gin_search_paths :: 以逗号分隔的目录列表，用作 gin 文件的路径前缀。我们可以使用 prompt_tuning.scripts.find_module ${module} 来查找包含配置的库的安装位置。
• --gin_file :: 要加载的 gin 文件。我们倾向于使用相对于其安装库的路径，例如使用 prompt_tuning/configs/models/t5_1_1_base_prompt.gin 而不是 models/t5_1_1_base_prompt.gin，以避免混淆。多次使用该标志可以指定多个 gin 文件，这些文件将被合并在一起。如果在多个文件中设置了相同的配置选项，则会采用列表中最后一个文件的值。
• --gin.{PARAM}={VALUE} :: 这个通用覆盖标志会将 PARAM 设置为 VALUE。这可用于轻松设置配置选项，而无需将其作为实际的命令行参数。例如， --gin.utils.SaveCheckpointConfig.keep=20 将保存最近的 20 个检查点。

在 TPU Pod 分片上训练提示

随着模型规模的增大，例如 XL 和 XXL 模型，它们已经无法容纳在单个 TPU 虚拟机所配备的 8 个 TPU 上。在这种情况下，我们需要使用 TPU Pod 的一个分片（有关 TPU 架构和可用配置的更多信息，请参阅 此处）。在单个 TPU 虚拟机上与在 Pod 分片上训练提示的主要区别在于，我们现在有多个 TPU 虚拟机，并且将在每个虚拟机上运行相同的 SPMD JAX 程序。本页面提供了关于 多主机 JAX 程序的更多信息。这篇指南简要介绍了如何在 TPU Pod 分片上运行 JAX 程序，但我们在此处将重点概述主要步骤。

1. 创建 TPU Pod 分片。 此页面列出了哪些地区提供哪种类型的加速器。这与上述创建 TPU 虚拟机的过程相同，只是我们请求的是 32 个 TPU 而不是 8 个。

$ gcloud alpha compute tpus tpu-vm create ${TPU_NAME} \
    --zone ${ZONE} \
    --accelerator-type v3-32 \
    --version v2-alpha

2. 安装 Prompt Tuning 库。由于我们现在有 4 个 TPU 虚拟机，每个虚拟机拥有 8 个 TPU，因此我们不希望通过 SSH 直接登录到每个虚拟机来执行操作。相反，Google Cloud 的 SSH 命令允许我们使用 --command= 标志指定要在所有虚拟机（称为工作节点）上运行的命令，并通过 --worker=all 参数实现这一点。

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone ${ZONE} \
  --worker=all \
  --command="git clone --branch=main https://github.com/google-research/prompt-tuning && cd prompt-tuning && "
python3 -m pip install . -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

3. 编写用于训练提示的脚本。我们提供了一个演示脚本 (/prompt_tuning/scripts/sst2-xxl-demo.sh)，该脚本使用 T5 1.1 lm100k XXL 模型来训练提示以解决 SST2 数据集问题。您可以将其用作起点，或者只需填写指向您 Google Cloud Storage 存储桶的路径，以指定保存结果的位置 (MODEL_DIR) 和缓存 TFDS 数据的位置 (TFDS_DATA_DIR)，也可以将这些路径设置为环境变量。

4. 将您的训练脚本复制到每个工作节点。如果您是第一次使用 scp，可能会遇到错误。请按照错误消息中的说明运行 ssh-add /.../.ssh/google_compute_engine 命令，然后重试。

$ gcloud alpha compute tpus tpu-vm scp sst2-xxl-demo.sh ${TPU_NAME}: \
  --zone=${ZONE}
  --worker=all

5. 执行您的训练脚本。

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone ${ZONE} \
  --worker=all \
  --command="./sst2-xxl-demo.sh"

如果其中一个工作节点在训练过程中出现错误，其他工作节点上的 TPU 将被占用而无法释放。这将阻止您重新启动作业，直到这些进程被终止并释放 TPU。以下命令应能终止所有这些进程。您可能会从最初出错的工作节点返回 kill 命令的手册页。

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone ${ZONE} \
  --worker=all \
  --command="sudo lsof -t /dev/accel0 | xargs kill -9"

自定义依赖项

要使用自定义组件（如您自己的数据集）训练提示，请遵循 T5X 关于自定义组件的说明。

如果您将自己的代码打包成可使用 pip 安装的 Python 包，则不必局限于单个目录，并且可以使用 python3 -m prompt_tuning.scripts.find_module {your_module} 来帮助设置 gin_search_paths，以便找到您库中捆绑的 gin 配置文件。注意： 如果您计划将 gin 配置文件捆绑到可安装的包中，请确保包含配置文件的目录具有 __init__.py 文件，因为 gin 要求文件必须位于 Python 包中。

如果您的自定义组件的部分内容可以通过 gin 进行配置，则需要在 gin 文件中显式导入；如果它们是在解析 gin 文件之后才被导入，将会导致错误。如果您所有的依赖项都不包含可配置的 gin 设置，则可以通过传递 --gin.MIXTURE_OR_TASK_MODULE="'path.to.your.module'" 来避免编写 gin 文件。这将自动导入您的模块，非常适合仅需更换数据集的情况。

使用提示进行推理

我们建议的使用提示进行推理的方法是加载用于初始化模型的原始检查点，并从文件中加载提示。正如本节关于 部分加载 中所解释的那样，T5X 支持在初始化时加载部分模型参数，而其他参数则从头开始初始化。我们将此功能与 from_array 提示初始化器结合使用，以从原始检查点和提示文件中重新加载冻结的参数。configs/runs/prompt_eval.gin 文件为您设置了此配置；您只需提供 PROMPT_FILE 即可。如果您的模型是使用任何 prompts/ 配置文件训练的，您可以将其从评估脚本的参数中移除。

随附的 sst2-demo-eval.sh 脚本展示了以这种方式进行评估的示例。您只需设置 EVAL_DIR 和 TFDS_DATA_DIR 环境变量，分别指向存储评估输出和 TensorFlow 数据集缓存的路径即可。

在 T5X 中，评估脚本假定您的数据集包含标签，并根据数据集的指标函数输出最终结果。而推理脚本则不需要标签，而是直接输出模型的预测结果。我们还提供了一个类似的 prompt_infer.gin 文件，供与推理脚本一起使用。

如果您希望使用由提示微调训练生成的 T5X 检查点进行推理或评估，可以直接使用 T5X 中的 (eval|infer).gin 配置文件。不过，您需要更新 utils.RestoreChekcpointConfig。应将 path 设置为新的检查点，assignment_map=()，并设置 fallback_to_scratch=False。

模型配置

所有关于模型、训练、评估、保存、恢复等的配置均通过 gin 完成。有关 gin 的通用介绍，请参阅 gin-config 仓库，以及 这篇入门指南。

我们遵循 T5X 的配置布局：

• runs/ :: 包含用于实际训练模型的配置。数据集和评估相关的配置通常放在这里。
• architectures/ :: 包含模型架构相关的配置。例如，编码器-解码器架构与纯解码器架构的选择，以及嵌入共享等设置。
• models/ :: 包含设置模型特定参数的配置，如层数或嵌入表的大小。此外，还配置所使用的 T5X 模型封装器等。
  • models/decoding/ :: 提供易于使用的配置，用于在推理过程中切换文本生成方式，包括束搜索和核采样等配置。
  • models/sizes/ :: 包含用于创建不同规模模型的各种设置。这些设置会与默认版本结合，形成特定规模的模型，例如，t5_1_1_prompt.gin + sizes/large.gin 可以创建一个 T5 1.1 Large 模型。一些常用组合已作为包含正确引入语句的 gin 文件提供（如上例中的 t5_1_1_large_prompt.gin）。注意： 这些规模文件必须__位于__主模型文件之后。
• prompts/ :: 我们的额外目录包含用于设置 PROMPT gin 变量的配置，从而可以通过添加不同的提示文件作为 --gin_file 参数来轻松切换提示初始化方式（该文件需置于 models/ gin 文件之后）。

gin 配置文件的顺序

在命令行中指定 --gin_file 参数时，文件的顺序非常重要。gin 文件的一般指定顺序如下：

1. models/*.gin
2. prompts/*.gin
3. models/sizes/*.gin
4. models/decoding/*.gin
5. runs/*.gin

必填字段

T5X 有一些必填字段，如 MIXTURE_OR_TASK_NAME 或 TASK_FEATURE_LENGTHS。我们又增加了两个必填字段：

• PROMPT_LENGTH :: 我们所使用的提示长度。由于该值会在多个地方使用，因此我们将其定义为一个 gin 宏，以便在多处引用并确保各处数值一致。
• PROMPT :: 这是实际提示模块的配置，将在 Flaxformer 的 PromptX 子类中使用。

注： 目前，提示调优不支持示例打包。这意味着我们的最大目标长度只需足够容纳每个示例的目标即可。因此，在 TASK_FEATURE_LENGTHS 映射中，targets 键的值可以显著缩短，例如对于许多 SuperGLUE (Wang et al., 2019) 任务，可以设为约 4，而 P5X 的默认值则是 62。

提示初始化

提示参数的初始化有多种选项。我们支持我们在论文 3.2 节 中介绍的各种方法，同时也支持从文件中进行初始化。后者允许用户例如从 MNLI 上学到的提示开始，在 BoolQ 数据集上进行训练。

所有初始化器都遵循 Flax 初始化器 API，即它们是带参数的函数，返回一个关于初始化函数的闭包。实际的初始化函数总是具有以下签名：

def initializer(rng: Array, shape: Sequence[int]) -> Array:
  ...

我们将每种初始化方案都提供为 configs/prompts 目录下的 gin 配置文件。可以通过使用 --gin_file=path/to/configs/prompts/scheme.gin 来包含 gin 文件来使用这些配置。该文件需要位于主模型文件之后，否则默认的（随机均匀）方法会覆盖您选择的方法。其中一些初始化方法可能需要您通过覆盖标志或在其中一个 gin 文件中设置额外的 gin 值。

随机均匀

这是一种标准的随机初始化，类似于人们用于嵌入初始化的方式。这是默认设置，无需任何 gin 文件。随机值的尺度可以通过覆盖 prompt_init/linen.initializers.uniform.scale=N 来调整。

采样词汇表

使用 from_sample_of_embeddings 初始化器，为每个提示位置采样一个标记嵌入作为初始化。您可以使用 prompt_init/prompts.from_samples_of_embeddings.population_size 参数将采样限制为前 n 个嵌入。

这可以与 --gin_file=prompt_tuning/configs/prompts/from_sampled_vocab.gin 一起使用。这种方法使用从初始模型检查点中提取的嵌入表。您也可以通过 --gin_file=prompt_tuning/configs/prompts/from_sampled_vocab_numpy.gin 提供您自己的嵌入文件。此方法要求您提供一个名为 EMBEDDING_FILE 的值，该值是一个包含模型嵌入表的 numpy 数组。这可以通过使用 prompt_tuning.scripts.extract_variable 从模型检查点中提取。

类标签

我们支持通过 from_embedded_list 初始化器，使用类标签（又称“表述词”）的嵌入来初始化提示时间步。用户可以提供一个要使用的单词列表（类标签）。每个单词由提供的词汇表进行分词；使用提供的词汇表嵌入；如有必要，跨子标记进行聚合；然后用于初始化提示的一个时间步。如果提供的标记未能覆盖整个提示长度，则缺失的标记将使用提供的后备初始化器进行初始化。

为了与论文中的做法保持一致——未填充的提示标记通过从嵌入表中采样来填充——我们可以将此初始化方法与上述方法组合使用。它可以与 --gin_file=prompt_tuning/configs/prompts/from_class_labels.gin 一起使用。这需要设置 CLASS_LABELS，即您希望嵌入作为提示初始化的单词列表。您还可以通过 --gin_file=prompt_tuning/configs/prompts/from_class_labels_numpy.gin 提供您自己的嵌入文件（与上述相同）。此外，还需要设置 EMBEDDING_FILE。

来自字符串

我们还支持使用某个字符串的嵌入来初始化提示，通常用于从离散提示或任务描述开始。这使用 from_embedded_string 初始化器。字符串由提供的词汇表进行分词，每个标记在提供的嵌入表中查找，最终得到的字符串嵌入表示被用作提示初始化。如果提供的标记未能覆盖整个提示长度，则缺失的标记将使用提供的后备初始化器进行初始化。

注意： 词汇表只是将字符串转换为一系列 ID，您需要确保字符串与您的 SeqIO 任务所执行的任何文本格式化（如标点符号周围的空格等）相匹配。

从文件

您还可以使用 from_array 初始化器从文件中加载提示，以实现跨任务的迁移。这可以通过 --gin_file=prompt_tuning/configs/prompts/from_file.gin 来完成。这需要设置 PROMPT_FILE，指定包含要加载提示的 Numpy 文件的路径。默认情况下，训练时会输出提示的 Numpy 版本，但也可以使用上述脚本提取提示。

已发布的模型检查点

我们发布了经过 10 万步语言模型适应的 T5 1.1 检查点的 T5X 原生检查点。

• t5_1_1_lm100k_small（约 7700 万个参数）： gs://t5-data/pretrained_models/t5x/t5_1_1_lm100k_small/checkpoint_1100000
• t5_1_1_lm100k_base（约 2.5 亿个参数）： gs://t5-data/pretrained_models/t5x/t5_1_1_lm100k_base/checkpoint_1100000
• t5_1_1_lm100k_large（约 8 亿个参数）： gs://t5-data/pretrained_models/t5x/t5_1_1_lm100k_large/checkpoint_1100000
• t5_1_1_lm100k_xl（约 30 亿个参数）： gs://t5-data/pretrained_models/t5x/t5_1_1_lm100k_xl/checkpoint_1100000
• t5_1_1_lm100k_xxl（约 110 亿个参数）： gs://t5-data/pretrained_models/t5x/t5_1_1_lm100k_xxl/checkpoint_1100000

这些检查点是从公开的 Mesh TensorFlow 检查点 转换而来的。

已发布的提示词

我们已在多种任务上发布了预训练的提示词，并计划随着时间的推移继续增加更多内容。

提示词位于 pretrained_prompts 目录中。该目录下的每个子目录按其对应的预训练模型进行分组。引用随库一起打包的这些提示词最简单的方式是：

  --PROMPT_FILE=`python3 -m prompt_tuning.scripts.find_module prompt_tuning`/pretrained_prompts/{MODEL_SIZE}/{PROMPT}.npy

由于并行计算固有的随机性，在训练和评估过程中需要确保若干设置完全一致，才能得到完全相同的数值结果。每个模型子目录都包含一个 README.md 文件，其中详细说明了这些必须匹配的设置。最重要的几个设置包括：批量大小、TPU 拓扑结构以及模型并行化的划分方式。表格中列出了在使用这些提示词运行 t5x.eval 时应预期获得的分数。

其他资源

这里收集了一些关于提示词微调的额外资源。

• 演示文稿：
  • EMNLP： 视频 幻灯片 海报

如何引用

如果您将本工作作为起点，请引用以下文献：

@inproceedings{lester-etal-2021-power,
    title = "参数高效的提示词微调中的规模效应",
    author = "Lester, Brian  and
      Al-Rfou, Rami  and
      Constant, Noah",
    booktitle = "2021年自然语言处理经验方法会议论文集",
    month = nov,
    year = "2021",
    address = "线上及多米尼加共和国蓬塔卡纳",
    publisher = "计算语言学协会",
    url = "https://aclanthology.org/2021.emnlp-main.243",
    doi = "10.18653/v1/2021.emnlp-main.243",
    pages = "3045--3059",
}

注意事项

本项目并非 Google 官方支持的产品。

Prompt Tuning 快速上手指南

本指南基于 Google Research 开源的 prompt-tuning 项目，帮助开发者快速复现 EMNLP 2021 论文中的参数高效提示微调实验。该项目主要依赖 TPU 环境运行。

环境准备

• 硬件要求：推荐使用 Google Cloud TPU VM（至少 v3-8 配置）。对于超大模型（如 XXL），需要 TPU Pod Slice（多主机环境）。
• 软件依赖：
  • Python 3
  • Google Cloud SDK (gcloud)
  • JAX, Flax, Flaxformer, T5X
• 存储要求：需要创建一个 Google Cloud Storage (GCS) Bucket，用于存放数据集缓存、模型检查点和结果。路径格式为 gs://{bucket-name}/path/to/item。

安装步骤

1. 创建并连接 TPU VM

使用 gcloud 命令创建 TPU 虚拟机并 SSH 登录：

# 创建 Cloud TPU VM (以 v3-8 为例)
$ gcloud alpha compute tpus tpu-vm create ${TPU_NAME} \
    --zone ${ZONE} \
    --accelerator-type v3-8 \
    --version v2-alpha

# SSH 登录到 TPU VM
$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE}

2. 克隆代码库

在 TPU VM 命令行中克隆仓库：

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

3. 安装依赖库

安装 Prompt Tuning 库及其 TPU 依赖。 注意：若遇到依赖版本回溯错误（backtracking error），请添加 --use-deprecated=legacy-resolver 标志。

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

(可选) 如果需要开发模式安装（修改代码即时生效）：

# 可能需要先删除 pyproject.toml 以避免安装错误
rm pyproject.toml
python3 -m pip install -e .[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

基本使用

训练一个 Prompt

训练流程与 T5X 微调类似，但需使用专用的配置文件。项目提供了针对 SST-2 数据集的演示脚本。

前置设置： 确保已设置环境变量指向你的 GCS Bucket：

• MODEL_DIR: 保存模型检查点和结果的目录。
• TFDS_DATA_DIR: 缓存 TensorFlow Datasets 的目录。

运行演示脚本：

./prompt-tuning/prompt_tuning/scripts/sst2-demo.sh

关键参数说明：

• --gin_search_paths: 指定 gin 配置文件的搜索路径前缀。
• --gin_file: 加载具体的 gin 配置文件（支持合并多个文件）。
• --gin.{PARAM}={VALUE}: 直接覆盖配置参数。例如 --gin.utils.SaveCheckpointConfig.keep=20 表示保留最近 20 个检查点。

多主机训练 (TPU Pod Slice)

对于无法放入单台 VM (8 TPU) 的大模型（如 XL, XXL），需使用 TPU Pod Slice。

1. 创建 Pod Slice (例如 32 TPU):

   $ gcloud alpha compute tpus tpu-vm create ${TPU_NAME} \
       --zone ${ZONE} \
       --accelerator-type v3-32 \
       --version v2-alpha
   

2. 在所有 Worker 上安装代码:

   $ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone ${ZONE} \
     --worker=all \
     --command="git clone --branch=main https://github.com/google-research/prompt-tuning && cd prompt-tuning && python3 -m pip install . -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
   

3. 分发训练脚本:

   $ gcloud alpha compute tpus tpu-vm scp sst2-xxl-demo.sh ${TPU_NAME}: \
     --zone=${ZONE} \
     --worker=all
   

4. 执行训练:

   $ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone ${ZONE} \
     --worker=all \
     --command="./sst2-xxl-demo.sh"
   

(故障处理) 若训练中断导致 TPU 被占用，可使用以下命令清理进程：

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone ${ZONE} \
  --worker=all \
  --command="sudo lsof -t /dev/accel0 | xargs kill -9"

推理与评估

评估 (Evaluation): 加载原始模型检查点和训练好的 Prompt 文件进行评估。需设置 EVAL_DIR 和 TFDS_DATA_DIR。

# 示例脚本参考
./prompt_tuning/scripts/sst2-demo-eval.sh

配置核心在于使用 configs/runs/prompt_eval.gin 并指定 PROMPT_FILE。

纯推理 (Inference): 如果不需标签，仅需模型预测输出，使用 configs/runs/prompt_infer.gin 配置。

直接使用 T5X 检查点: 若想对训练生成的完整 T5X 检查点进行推理，可直接使用 T5X 原生的 (eval|infer).gin，但需更新恢复配置：

• path: 指向新检查点路径
• assignment_map=()
• fallback_to_scratch=False