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

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

为了满足大规模人工智能训练不断增长的需求，我们去年在 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 TPU 上训练参数多达 128B 的 GPT-2 模型。TPU 上的 PyTorch/XLA FSDP 训练效率很高，GPT-2 的模型 FLOPS 利用率 (MFU) 高达 45.1%。

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

图 1：Google Cloud TPU v4 上 Hugging Face GPT-2 的模型 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 中，创建您可能需要的任何模型特定配置文件。这是一个具有 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.191s。使用上述公式，我们计算出 v4-8 上的 MFU 为 35.7%。有关计算 MFU 的更多详细信息，请参阅PaLM 论文。

下表展示了序列长度为 1024、参数量在 2B 到 128B 之间的 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
步长 (ms)	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，我们建议使用 FSDP 来减少模型参数所需的内存，并使用能够适应 TPU 切片 HBM 的最大全局批处理大小进行训练。

训练超大型模型（已测试至 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 包装紧接之前的一行中添加以下函数到 _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)

如果 check_fn 返回 True，函数 materialize_module 将初始化模型张量。在这种情况下，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 参数）

后续步骤和致谢

文档

教程

资源

保持联系以获取更新、活动信息和最新新闻