使用 optimum-neuron 进行分布式训练

AWS Trainium 实例非常适合训练模型。它们最多可以包含 16 个 Neuron 设备，每个设备包含 2 个 Neuron 核心，并具有 32GB 的内存（每个核心 16GB）。例如，一个 trn1.32xlarge 实例具有 32 x 16 = 512GB 的内存。

但需要注意的是：默认情况下，每个 Neuron 核心都是一个独立的数据并行工作器。这意味着模型、梯度状态和优化器状态（总计约为模型大小的 4 倍）必须适合每个 Neuron 核心（16GB），才能进行训练。如果满足此条件，则激活也必须适合剩余内存。

为了缓解这种情况，optimum-neuron 支持并行功能，使您能够充分利用 Trainium 实例的强大功能

ZeRO-1：这是一种数据并行性的优化，它将优化器状态（通常占设备所需内存的一半）分片到数据并行 ranks 上。
张量并行：这是一种技术，它沿着给定轴（行或列）在多个设备上分片模型的每个矩阵乘法。它也称为层内模型并行性。用于分片参数的设备数量称为 tensor_parallel_size。
序列并行：这是对张量并行性的一种优化，它在张量并行区域之外的序列轴上分片激活。它很有用，因为它通过分片激活来节省内存。
流水线并行：它包括在多个设备上分片模型块层。它也称为层间模型并行性。用于分片层的设备数量称为 pipeline_parallel_size。

好消息是，可以组合这些技术，而 optimum-neuron 使其非常容易！

optimum-neuron repo 中提供的所有示例脚本都通过 NeuronTrainer 实现了这些功能。

如何启用 ZeRO-1？

无论您使用 NeuronTrainer 还是决定使用自己的训练脚本（使用 NeuronAccelerator），启用 ZeRO-1 优化都非常容易。

通过 NeuronTrainer

from optimum.neuron import NeuronTrainingArguments, NeuronTrainer

# To enable ZeRO-1, set the `zero_1` argument to `True` in the training arguments.
training_args = NeuronTrainingArguments(
    ...,
    zero_1=True,
)

trainer = NeuronTrainer(
    model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()

由于示例脚本使用 NeuronTrainer，您可以在使用它们时通过向命令行添加 --zero_1 标志来启用 ZeRO-1。

例如

torchrun --nproc_per_node=2 examples/language-modeling/run_clm.py \
    --model_name_or_path TinyLlama/TinyLlama-1.1B-Chat-v0.6 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
    --per_device_train_batch_size 1 \
    --block_size 1024 \
    --bf16 \
    --zero_1 \
    --output_dir my_training/

通过 NeuronAccelerator

当不使用 NeuronTrainer 时，需要做更多的工作

（可选）包装优化器类使其变为惰性加载。启用 ZeRO-1 后，原始优化器将被覆盖以使用其分片版本。因此，可以惰性加载原始优化器，以便在实际分片之前不会物化优化器状态。

from torch.optim import AdamW
from optimum.neuron.distributed import make_optimizer_constructor_lazy

lazy_adamw = make_optimizer_constructor_lazy(AdamW)

在实例化 NeuronAccelerator 时，将 zero_1 参数设置为 True。

accelerator = NeuronAccelerator(
    ...
    zero_1=True,
)

model = ...
lazy_optimizer = lazy_adamw(...) # Actually instantiate the optimizer.


model, optimizer = accelerator.prepare(model, lazy_optimizer)

如何启用张量并行？

与 ZeRO-1 一样，可以使用 NeuronTrainer 或 NeuronAccelerator 应用张量并行。

在进行张量并行时，您有不同的设置

tensor_parallel_size。理想情况下，它应该是模型可以容纳的最小值。
是否应启用序列并行。序列并行在张量并行区域之外的序列轴上分片激活。它很有用，因为它通过分片激活来节省内存。
是否应完成嵌入层的并行化。默认情况下会这样做，因为它提供了多项好处

并行化嵌入层可以节省内存，从而能够适应更大的批次大小和/或序列长度。
对于语言模型，其中嵌入层权重和语言建模头权重通常是绑定的，语言建模头最终是并行的，并且不需要 all-gather 其输出，因为它被馈送到与并行性兼容的交叉熵损失，从而节省了昂贵的通信。

最重要的是，务必确保以高效的方式加载原始模型：训练脚本将由 torchrun 调用，torchrun 会将其分派给工作器，每个核心一个工作器。如果每个工作器（在 trn1.32xlarge 实例中有 32 个）加载完整的模型权重，则会花费大量时间并很快耗尽内存。

optimum-neuron 提供了一个上下文管理器 distributed.lazy_load_for_parallelism()，它惰性加载模型以防止这种情况，只有相应模型分片的参数才会在每个工作器中物化。

通过 NeuronTrainer

from optimum.neuron import NeuronTrainingArguments, NeuronTrainer
from optimum.neuron.distributed import lazy_load_for_parallelism

# Specify the `tensor_parallel_size` in the training arguments.
training_args = NeuronTrainingArguments(
    ...,
    tensor_parallel_size=8,
    disable_embedding_parallelization=False, # It is `False` by default.
    disable_sequence_parallel=False, # It is `False` by default.
)

with lazy_load_for_parallelism(tensor_parallel_size=training_args.tensor_parallel_size):
    model = ...


trainer = NeuronTrainer(
    model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()

由于示例脚本使用 NeuronTrainer，您可以在使用它们时通过指定 --tensor_parallel_size 参数以及可选的 disable_embedding_parallelization 和 disable_sequence_parallel 标志来启用张量并行。

例如

torchrun --nproc_per_node=2 examples/language-modeling/run_clm.py \
    --model_name_or_path TinyLlama/TinyLlama-1.1B-Chat-v0.6 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
    --per_device_train_batch_size 1 \
    --block_size 1024 \
    --bf16 \
    --tensor_parallel_size 2 \
    --output_dir my_training/

通过 NeuronAccelerator

与 ZeRO-1 一样，可以包装优化器类使其变为惰性加载。由于模型参数将被分片，因此无需在模型并行化之前物化优化器状态：包装器确保它保持未物化状态。

from torch.optim import AdamW
from optimum.neuron import NeuronAccelerator
from optimum.neuron.accelerate.utils import ModelParallelismPlugin
from optimum.neuron.distributed import lazy_load_for_parallelism

tensor_parallel_size = 8
mp_plugin = ModelParallelismPlugin(
    tensor_parallel_size,
    parallelize_embeddings=True,
    sequence_parallel_enabled=True,
    checkpoint_dir=None, # Can be specified when resuming from checkpoint.
)

accelerator = NeuronAccelerator(
    ...
    mp_plugin=mp_plugin,
)

with lazy_load_for_parallelism(tensor_parallel_size=tensor_parallel_size):
    model = ...

lazy_adamw = make_optimizer_constructor_lazy(AdamW)
lazy_optimizer = lazy_adamw(...) # Actually instantiate the optimizer.

model, optimizer = accelerator.prepare(model, lazy_optimizer)

检查点合并

由于张量并行包括在不同的工作器之间分片模型权重，因此训练期间只会保存分片检查点。有必要合并分片检查点，以便能够在创建它们时使用的特定训练配置之外共享和使用它们。

Optimum CLI 提供了一种通过 optimum neuron consolidate 命令非常容易地执行此操作的方法

optimum-cli neuron consolidate --help

usage: optimum-cli neuron consolidate [-h] [-f {pytorch,safetensors}] checkpoint_dir output_dir

positional arguments:
  checkpoint_dir        The path to the directory containing the checkpoints.
  output_dir            The path to the output directory containing the consolidated checkpoint.

optional arguments:
  -h, --help            show this help message and exit
  -f {pytorch,safetensors}, --format {pytorch,safetensors}
                        The format used to save the consolidated checkpoint.

您只需指定分片检查点目录和将包含合并检查点的输出目录，该命令就会处理其余的工作。还可以指定合并检查点的输出格式，默认情况下，它会将其导出为 safetensors 格式，这是建议使用的格式。

示例

使用张量并行进行的训练刚刚完成，输出目录名为 my_training。该目录如下所示

my_training/
├── README.md
├── all_results.json 
├── checkpoint-10 
│   ├── config.json
│   ├── scheduler.pt
│   ├── special_tokens_map.json
│   ├── tensor_parallel_shards
│   ├── tokenizer.json
│   ├── tokenizer.model
│   ├── tokenizer_config.json
│   ├── trainer_state.json
│   └── training_args.bin
├── config.json
├── special_tokens_map.json
├── tensor_parallel_shards
│   ├── tp_rank_00_pp_rank_00
│   ├── tp_rank_01_pp_rank_00
│   ├── tp_rank_02_pp_rank_00
│   ├── tp_rank_03_pp_rank_00
│   ├── tp_rank_04_pp_rank_00
│   ├── tp_rank_05_pp_rank_00
│   ├── tp_rank_06_pp_rank_00
│   └── tp_rank_07_pp_rank_00
├── tokenizer.json
├── tokenizer.model
├── tokenizer_config.json
├── train_results.json
├── trainer_state.json
└── training_args.bin

可以通过运行以下命令来合并 my_training/tensor_parallel_shards 中的分片检查点，这对应于在训练结束时保存的分片检查点

optimum-cli neuron consolidate my_training my_training_consolidated_checkpoint

分片检查点保存在名为 tensor_parallel_shards 的目录下。optimum-cli neuron consolidate 命令接受包含 tensor_parallel_shards 目录的目录或 tensor_parallel_shards 目录本身作为输入。

AWS Trainium & Inferentia

使用 optimum-neuron 进行分布式训练

如何启用 ZeRO-1？

通过 NeuronTrainer

通过 NeuronAccelerator

如何启用张量并行？

通过 NeuronTrainer

通过 NeuronAccelerator

检查点合并