Large Scale Training of Hugging Face Transformers on TPUs With PyTorch/XLA FSDP

由亚历克斯·沃特海姆、米尔达德·莫哈马迪、杰克·曹、亚历克斯·斯皮里多诺夫、乔·斯皮萨克、利桑德·德布特、西尔万·古格、索拉布·曼格鲁卡尔

人工智能通过理解生成语言、回答问题和提供准确推荐等高级功能正在改变许多行业。这些功能由不断增长和复杂的人工智能模型驱动，这些模型需要大量的计算能力来训练。

为了满足大规模人工智能训练不断增长的需求，去年我们在 PyTorch/XLA 中引入了完全分片数据并行（FSDP）。FSDP 是一种模型并行架构，它解锁了轻松高效地将人工智能模型扩展到数百亿参数的能力。在 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 TPUs 上使用 GPT-2 模型进行训练，参数高达 128B。PyTorch/XLA FSDP 在 TPUs 上的训练效率极高，GPT-2 的模型 FLOPS 利用率（MFU）高达 45.1%：

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（s），并安装 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 中，创建您可能需要的任何特定模型配置文件。以下是一个 GPT-2 模型（2B 参数）的配置文件示例，我们稍后将称之为 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 配置文件，一个包含所有可配置的 XLA FSDP 包装方面的 JSON 文件，存储为字典。遵循官方的 Hugging Face Transformers XLA FSDP 文档，以下参数可供设置：

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

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

fsdp_min_num_params (int, \*optional\*, defaults to 0) : 设置基于大小的自动换行所需的最小参数数量。具有至少与 fsdp_min_num_params 一样多参数的每个模块都将被 XLA FSDP 包装。
fsdp_transformer_layer_cls_to_wrap (List[str], \*optional\*) : 要包装的（区分大小写）转换层类名列表。请注意，这与 fsdp_min_num_params 互斥。例如： ["GPT2Block", "GPT2MLP"] 。
xla_fsdp_grad_ckpt (bool, \*optional\*, defaults to False) : 这是一个布尔值，用于确定是否在每个嵌套的 XLA FSDP 包装层上使用梯度检查点。此设置只能在将 xla 标志设置为 true 时使用，并且通过 fsdp_min_num_params 或 fsdp_transformer_layer_cls_to_wrap 指定自动包装策略。

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

对于这个 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 上训练具有 20 亿参数的 GPT-2 模型时，我们测量到步长时间为 4.191 秒。使用上述公式，我们在 v4-8 上计算出 35.7%的 MFU。有关计算 MFU 的更多详细信息，请参阅 PaLM 论文。

下表展示了大小在 2B 到 128B 之间的 GPT-2 模型的 MFU，序列长度为 1024。

TPU NumCores	v4-8	v4-64	v4-128	v4-128	v4-256	v4-512
每个批次/批次的令牌数	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 使用全量减少通信集体，这些集体不是异步的，这意味着芯片间的通信不能与计算重叠。随着芯片数量的增加，必须通信的模型分片数量也增加，因此我们预计通信所占的步骤时间比例会随着芯片数量的增加而增加。

其次，增加批次大小通常意味着更好的 FLOPS 利用率。随着芯片数量的增加，模型的内存占用减少，这通常可以释放出高带宽内存（HBM），以扩大全局批次大小。随着全局批次大小的增加，每步处理的标记数量增加，因此每步的 FLOPS 也增加。只要步骤时间不按比例增加，我们预计更大的全局批次大小可以提高 MFU。

因此，为了最大化 MFU，我们建议使用尽可能大的全局批次大小进行训练，该批次大小可以适应 TPU 切片的 HBM，并使用 FSDP 来减少模型参数所需的内存。

训练非常大的模型（已测试至 128B 参数）

当使用 PyTorch/XLA 时，必须在将张量移动到 XLA 设备之前在 CPU 上初始化它们。这意味着如果模型足够大，可能会遇到主机端内存不足错误，即使模型在分片后可以适应设备 HBM。为了避免这种情况，我们必须将每个子模块的初始化推迟到它被 FSDP 包装时，这确保了子模块在它们的值被填充后立即分片，避免了主机端限制。

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

首先，使用以下命令安装 torchdistX 库，这是一个包含实验性 PyTorch 分布式功能的库。这是延迟初始化的引擎，允许您创建不需要立即存储的 tensor，可以在稍后实现。您还需要安装一个特定的 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 处，在 _wrap_model 之前的一行添加以下函数：

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 时初始化模型 tensor。在这种情况下， 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 遇到任何问题，或者想告诉我们您是如何使用它的，我们非常乐意听取您的意见。

我们对 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 在准备这篇帖子中的帮助。

喝彩！

谷歌的 PyTorch/XLA 团队

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

在 Hugging Face Trainer 中配置 PyTorch/XLA FSDP

测量 GPT-2 模型 FLOPS 利用率（MFU）

训练非常大的模型（已测试至 128B 参数）

下一步 & 致谢

文档

教程

资源