使用 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：它是数据并行性的优化，它包括将优化器状态（通常占设备所需内存的一半）跨数据并行秩进行分片。
- 张量并行：它是一种技术，它包括将您的每个模型矩阵乘法沿着给定轴（行或列）在多个设备上进行分片。它也被称为层内模型并行。用于分片参数的设备数量称为
tensor_parallel_size。 - 序列并行：它是张量并行的优化，它在张量并行区域之外沿着序列轴对激活进行分片。它很有用，因为它可以通过分片激活来节省内存。
- 流水线并行：它包括在多个设备上对模型块层进行分片。它也被称为层间模型并行。用于分片层的设备数量称为
pipeline_parallel_size。
好消息是，可以将这些技术结合起来，而 optimum-neuron 使这变得非常容易！
optimum-neuron 仓库中提供的所有示例脚本都通过 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 目录本身。