使用 PyTorch/XLA FSDP 在 TPU 上大规模训练 Hugging Face Transformers

作者：Alex Wertheim, Milad Mohammadi, Jack Cao, Alex Spiridonov, Joe Spisak, Lysandre Debut, Sylvain Gugger, Sourab Mangrulkar

人工智能正通过理解和生成语言、回答问题以及提供准确推荐等先进能力，改变着许多行业。这些能力得益于不断增长的模型规模和复杂性，它们需要大量的计算能力进行训练。

为了满足日益增长的大规模 AI 训练需求，我们去年在 PyTorch/XLA 中引入了完全分片数据并行 (FSDP)。FSDP 是一种模型并行架构，能够轻松高效地将 AI 模型扩展到数千亿参数。使用 PyTorch/XLA FSDP，在分布式训练期间，每个设备可以存储特定的模型分片，并在进行前向传播时聚合完整的模型权重。嵌套 FSDP 通过仅在给定层的前向传播期间使用其完整参数来进一步优化性能。

我们很高兴地宣布，PyTorch/XLA FSDP 已登陆 Hugging Face Transformers。现在，Hugging Face 用户可以使用相同的计算能力训练参数量高达之前的 20 倍的 PyTorch 模型。

我们将 PyTorch/XLA FSDP 支持直接构建到 Hugging Face Trainer 类中，因此任何使用 Trainer 的模型都可以利用 FSDP。随着自动封装添加到 PyTorch/XLA FSDP 中，嵌套 FSDP 封装变得既灵活又易于应用。这些新特性使得在大规模下训练各种 Hugging Face 模型变得简单。在本指南中，我们将演示如何在 Google Cloud TPU 上训练参数量高达 128B 的 GPT-2 模型。TPU 上的 PyTorch/XLA FSDP 训练效率非常高，对于 GPT-2 模型可以达到高达 45.1% 的模型 FLOPS 利用率 (MFU)。

Figure 1: Model FLOPS utilization for Hugging Face GPT-2 on Google Cloud TPU v4

图 1：Hugging Face GPT-2 在 Google Cloud TPU v4 上的模型 FLOPS 利用率

在 Hugging Face Trainer 中配置 PyTorch/XLA FSDP

首先，按照您喜欢的方法创建您的 TPU 并安装 PyTorch 和 PyTorch/XLA。PyTorch 和 PyTorch/XLA 都需要 >= 2.0 版本。

    pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torc h-2.0-cp38-cp38-linux_x86_64.whl --user

    pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torc h_xla-2.0-cp38-cp38-linux_x86_64.whl

接下来，克隆并安装 Hugging Face Transformers 仓库。安装所有必需的依赖项（例如，datasets, evaluate, scikit-learn, accelerate）。

    cd $HOME
    git clone https://github.com/huggingface/transformers.git cd transformers
    git checkout v4.31-release
    pip3 install -e .
    pip3 install datasets evaluate scikit-learn
    pip3 install accelerate==0.21.0

在 $HOME/transformers 中，创建您可能需要的任何模型特定配置文件。这里是一个参数量为 2B 的 GPT-2 模型的配置文件示例，我们稍后将其称为 gpt2_config.json

{
    "activation_function": "gelu_new", 
    "architectures": [
        "GPT2LMHeadModel"
    ],
    "attn_pdrop": 0.1,
    "bos_token_id": 50256, "embd_pdrop": 0.1, "eos_token_id": 50256, "initializer_range": 0.02, "layer_norm_epsilon": 1e-05, "model_type": "gpt2",
    "n_embd": 3072,
    "n_head": 24,
    "n_layer": 18,
    "n_positions": 1024,
    "resid_pdrop": 0.1,
    "summary_activation": null,
    "summary_first_dropout": 0.1,
    "summary_proj_to_labels": true,
    "summary_type": "cls_index",
    "summary_use_proj": true,
    "task_specific_params": {
        "text-generation": {
            "do_sample": true,
            "max_length": 50
        }
    },
    "vocab_size": 50257
}

借助 PyTorch/XLA FSDP，可以在大型加速器切片上训练比这大得多的模型。我们已经使用这些技术训练了高达 128B 参数的 GPT-2 模型；有关如何复制这种规模的专家提示，请参阅附录。

在 $HOME/transformers 中，创建您的 FSDP 配置文件，这是一个 JSON 文件，其中包含您的 XLA FSDP 封装的所有可配置方面，以字典形式存储。根据Hugging Face Transformers XLA FSDP 官方文档，以下参数可供设置

xla (bool, \*可选\*, 默认为 False)：这是一个布尔值，用于确定您是否使用 XLA FSDP。请确保将其设置为 true。
xla_fsdp_settings (dict, \*可选\*)：这是一个字典，用于存储您要设置的所有 XLA FSDP 封装参数；请注意，如果您使用默认值，则无需为参数指定设置。有关设置的完整列表，请参阅此处。

对于 compute_dtype 和 buffer_dtype，请将它们输入为包含相应 torch 数据类型的字符串，例如 bfloat16。

fsdp_min_num_params (int, \*可选\*, 默认为 0)：一个整数，设置基于大小的自动封装的最小参数数量。每个参数数量至少等于 fsdp_min_num_params 的模块都将被 XLA FSDP 封装。
fsdp_transformer_layer_cls_to_wrap (List[str], \*可选\*)：一个（区分大小写的） transformer 层类名称列表，用于封装。请注意，这与 fsdp_min_num_params 互斥。示例：["GPT2Block", "GPT2MLP"]。
xla_fsdp_grad_ckpt (bool, \*可选\*, 默认为 False)：这是一个布尔值，用于确定是否在每个嵌套的 XLA FSDP 封装层上使用梯度检查点。此设置只能在 xla 标志设置为 true 且通过 fsdp_min_num_params 或 fsdp_transformer_layer_cls_to_wrap 指定了自动封装策略时使用。

注意：对于基于 transformer 的模型，在执行自动嵌套 FSDP 封装时，请使用 fsdp_transformer_layer_cls_to_wrap 而非 fsdp_min_num_params。共享权重的层不应属于单独的 FSDP 封装单元，并且基于 transformer 的模型中的输入和输出嵌入层共享权重。

对于此 GPT-2 示例，相应的 fsdp_config.json 文件如下所示

    {
        "fsdp_transformer_layer_cls_to_wrap": [
            "GPT2Block"
        ],
        "xla": true,
        "xla_fsdp_settings": {
            "compute_dtype": "bfloat16",
            "shard_param_on_dim_0": true,
            "pin_layout_in_collective_ops": true
        },
       "xla_fsdp_grad_ckpt": true
    }

现在，是时候训练您的模型了！首先，通过设置以下项确保您的 PyTorch/XLA 运行时已正确配置

    export PJRT_DEVICE=TPU

运行训练时，要传递的关键标志是

a) --fsdp "full_shard" b) --fsdp_config fsdp_config.json

其中您应该将 fsdp_config.json 替换为您自己的 FSDP 配置文件的名称。这是一个训练我们的示例 2B GPT-2 模型的样本命令，训练通过 xla_spawn.py（一个用于分布式 TPU 训练的启动脚本）启动。

    python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/language-modeling/run_clm.py \
    --num_train_epochs 1 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \ --per_device_train_batch_size 32 \ --per_device_eval_batch_size 32 \
    --do_train \
    --do_eval \
    --output_dir /tmp/test-clm \
    --overwrite_output_dir \
    --config_name gpt2_config.json \
    --cache_dir /tmp \
    --tokenizer_name gpt2 \
    --block_size 1024 \
    --optim adafactor \
    --adafactor true \
    --save_strategy no \
    --logging_strategy no \
    --fsdp "full_shard" \
    --fsdp_config fsdp_config.json

测量 GPT-2 的模型 FLOPS 利用率 (MFU)

模型 FLOPS 是执行一次前向和后向传播所需的浮点运算。模型 FLOPS 与硬件和实现无关，仅取决于底层模型。在每个步骤中，FLOPS 数量通过以下公式计算

tokens_per_batch = global_batch_size \* seq_len

FLOPS_per_step = 6 \* tokens_per_batch \* num_params

其中 seq_len 是序列长度，num_params 是模型中的参数数量。我们注意到，此估计假设输入维度远大于输入序列长度（d_model >> seq_len）。如果违反此假设，自注意力 FLOPS 开始变得足够重要，此表达式将低估真实的 MFU。

基于步长时间和硬件细节（芯片数量和每个芯片的峰值 FLOPS），我们可以计算模型 FLOPS 利用率 (MFU)，它衡量了我们的实现在使用底层硬件方面的效率。达到 100% MFU 意味着硬件被模型完美利用。我们使用以下公式计算 MFU

model_FLOPS_utilization = FLOPS_per_step / step_time(s) / chip_count / FLOPS_per_chip

在使用上述 XLA FSDP 配置文件在 Cloud TPU v4-8 上训练参数量为 2B 的 GPT-2 模型时，我们测得步长时间为 4.191 秒。使用上述公式，我们在 v4-8 上计算出 35.7% 的 MFU。有关 MFU 计算的更多详细信息，请参阅PaLM 论文。

下表列出了参数量介于 2B 和 128B 之间、序列长度为 1024 的 GPT-2 模型的 MFU。

TPU 核数	v4-8	v4-64	v4-128	v4-128	v4-256	v4-512
每批次 Token 数	131,072	524,288	524,288	524,288	1,048,576	1,048,576
参数数量	2B	16B	20B	32B	64B	128B
步长时间 (毫秒)	4,191	14,592	7,824	12,970	25,653	30,460
PFLOPS / 步	1.65	50	62	101	404	809
MFU	35.7%	38.8%	45.1%	44.4%	44.7%	37.7%

表 1：GPT-2 模型 FLOPS 利用率计算详情

在这些配置中，20B 参数模型在 v4-128 上的 MFU 峰值为 45.1%。这一结果与例如22B Megatron-like 模型的 41.5% MFU 相比，表现良好。

从这些实验中可以得出两个可操作的见解

首先，简单地增加芯片数量而不增加批次大小通常意味着 FLOPS 利用率较低，因为更多时间花费在共享模型分片上。FSDP 使用 All-reduce 通信集合，它们是非异步的，这意味着芯片间的通信不能与计算重叠。随着芯片数量的增加，需要通信的模型分片数量也增加，因此我们可以预期步长时间中花费在通信上的部分会随着芯片数量的增加而增加。

其次，增加批次大小通常意味着更好的 FLOPS 利用率。随着芯片数量的增加，模型的内存占用减少，这通常会释放高带宽内存 (HBM) 来扩大全局批次大小。全局批次大小越大，每个步骤处理的 token 数量就越多，因此每个步骤的 FLOPS 也越高。只要步长时间不成比例地增加，我们预计更大的全局批次大小会提高 MFU。

因此，为了最大化 MFU，我们建议使用能够在 TPU 切片的 HBM 中容纳的尽可能大的全局批次大小进行训练，同时利用 FSDP 减少模型参数所需的内存。

训练超大型模型（已测试高达 128B 参数）

使用 PyTorch/XLA 时，张量必须先在 CPU 上初始化，然后才能移动到 XLA 设备。这意味着如果模型足够大，即使分片后模型可以容纳在设备 HBM 中，也可能会遇到主机端内存不足错误。为避免这种情况，我们必须推迟每个子模块的初始化，直到它被 FSDP 封装，这确保了子模块在填充其值后立即进行分片，从而避免了主机端限制。

下面，我们将解释如何修改 Hugging Face transformers 仓库的本地副本，以使用此技术训练高达 128B 参数的 GPT-2 模型。

首先，使用以下命令安装 torchdistX，这是一个包含实验性 PyTorch 分布式功能的库。这是延迟初始化的引擎，允许您创建不需要立即存储并可在以后实例化的张量。您还需要安装利用此包的特定 PyTorch/XLA 2.0 版本；请注意，如果您之前安装了 PyTorch 和 PyTorch/XLA，则必须先卸载它们。

pip3 install torch==2.0 --index-url [https://download.pytorch.org/whl/test/cpu](https://download.pytorch.org/whl/test/cpu) --user
pip3 install torch_xla[torchdistx] -f https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/experimen tal/torch_xla-2.0-cp38-cp38-linux_x86_64.whl

接下来，对您的 Hugging Face Transformers 本地副本应用以下更改

在 src/transformers/trainer.py 中，在 PyTorch/XLA FSDP 封装紧前的一行中添加以下函数

from torchdistx import deferred_init

def _init_with_torchdistX(module):
    def check_fn(k):
        return not isinstance(k, FSDP)
    deferred_init.materialize_module(module, check_fn=check_fn)

函数 materialize_module 如果 check_fn 返回 True，将初始化模型张量。在这种情况下，check_fn 检查模块是否已进行 FSDP 封装。

在 _wrap_model 中，修改您的 FSDP 封装以接受附加参数 param_init_fn=_init_with_torchdistX

self.model = model = FSDP(
        model,
        auto_wrap_policy=auto_wrap_policy,
        auto_wrapper_callable=auto_wrapper_callable,
        param_init_fn=_init_with_torchdistX,
        \*\*fsdp_kwargs,
    )

在 examples/pytorch/language-modeling/run_clm.py 中，在文件开头添加以下导入语句

from torchdistx import deferred_init

编辑模型初始化，以便通过替换以下行来使用 deferred_init.deferred_init 封装模型

model = AutoModelForCausalLM.from_config(config)

替换为

model = deferred_init.deferred_init(AutoModelForCausalLM.from_config, config)

请注意，这假设您提供了自己的模型配置文件。否则，您应相应地修改您的模型初始化语句。

您还应该注释掉紧随其后的这两行

n_params = sum({p.data_ptr(): p.numel() for p in model.parameters()}.values()) logger.info(f"Training new model from scratch - Total size={n_params/2\*\*20:.2f}M params")

如果不对其进行修改，它们将导致错误，因为执行这些行时模型张量实际上没有存储空间。

进行这些更改后，只要加速器大小足够大，您现在就可以运行参数量多达 128B 的 GPT-2 模型了。

后续步骤与致谢

要了解更多信息，请参阅此处的文档。如果您在使用 PyTorch/XLA 中的 FSDP 时遇到任何问题，或者只是想告诉我们您是如何使用它的，我们很乐意听取您的意见。

我们对 PyTorch/XLA 的未来感到欣喜，并邀请社区加入我们。PyTorch/XLA 完全以开源形式开发。因此，请在 GitHub 上提交问题、发送拉取请求和 RFC，以便我们公开协作。

我们要感谢 Meta AI 的 Ronghang Hu 和 Ross Girshick，以及 Lysandre Debut, Sourab Mangrulkar, Sylvain Gugger 和 Arthur Zucker 的所有支持与协作。我们还要感谢 Jiewen Tan, Liyang Lu, Will Cromar, Vaibhav Singh 和 Chandra Devarakonda 在准备此文中的协助。

此致

Google 的 PyTorch/XLA 团队

使用 PyTorch/XLA FSDP 在 TPU 上大规模训练 Hugging Face Transformers

在 Hugging Face Trainer 中配置 PyTorch/XLA FSDP

测量 GPT-2 的模型 FLOPS 利用率 (MFU)

训练超大型模型（已测试高达 128B 参数）

后续步骤与致谢

文档

教程

资源