AWS Trainium & Inferentia 文档

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

AWS Trainium & Inferentia

加入 Hugging Face 社区

并获得增强的文档体验

在模型、数据集和 Spaces 上进行协作

通过加速推理获得更快的示例

切换文档主题

开始使用

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

AWS Trainium 实例为大规模训练大型语言模型提供了强大的基础设施。一个 trn1.32xlarge 实例包含 16 个 Neuron 设备，共 32 个核心，提供 512GB 内存（每个核心 16GB）。

然而，训练大型模型带来了一个根本性挑战：默认情况下，每个 Neuron 核心作为一个独立的数据并行工作单元运行，需要将整个模型、梯度和优化器状态（大约是模型大小的 4 倍）都放入单个核心的 16GB 内存限制内，并且还需要额外的空间来存放激活值。

对于超出这些内存限制的模型，optimum-neuron 提供了先进的并行策略，可以将计算和内存分布到多个设备上，使您能够训练那些无法在单个核心上容纳的模型。

并行策略概述

1. ZeRO-1（优化器状态分片）

ZeRO-1 是一种优化器级别的优化，可以在不改变模型架构的情况下减少内存使用。

工作原理：将优化器状态（梯度、动量、方差）在数据并行等级之间进行分片，而不是在每个设备上复制一份。

内存节省：将优化器的内存使用量减少到 1/data_parellel_size。

适用场景：在使用多个设备进行训练时，无论模型大小如何，始终有益。

2. 张量并行（层内模型并行）

张量并行将单个模型层拆分到多个设备上。

工作原理：将矩阵乘法（线性层、注意力）沿行或列在设备间进行分片。每个设备计算每层的一部分，这要求在每次前向/后向传播时设备间进行通信。

内存节省：将模型参数内存减少到 1/tensor_parallel_size。

适用场景：当模型过大，即使应用了 ZeRO-1 也无法在单个设备上容纳时。

典型部署：由于通信要求高，通常在单个节点内（节点内）应用。

权衡：增加了设备间的通信开销，如果过度使用可能会减慢训练速度。

3. 序列并行（激活值分片）

序列并行是一种与张量并行协同工作的优化方法，可以进一步减少内存使用。

工作原理：在张量尚未被张量并行分片的区域，沿序列维度对激活值进行分片。

内存节省：按序列长度比例减少激活值内存，对长序列尤其有益。

适用场景：在使用张量并行时始终启用——它以最小的开销提供额外的内存节省。

要求：只能与张量并行结合使用。

4. 流水线并行（层间模型并行）

流水线并行将模型层拆分到不同的设备上。

工作原理：将模型划分为多个阶段，每个阶段包含在不同设备上运行的连续层。使用微批处理（microbatching）来保持所有设备处于繁忙状态。

内存节省：将模型参数内存减少到 1/pipeline_parallel_size。

适用场景：对于即使使用张量并行也无法容纳的非常大的模型，或者当您希望以比张量并行更低的通信开销扩展到许多设备时。

典型部署：通常跨多个节点（节点间）应用，以扩展到更多设备，同时最大限度地减少高带宽通信需求。

权衡：会引入流水线气泡（空闲时间），并需要仔细调整微批次大小。

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

optimum-neuron 仓库中的所有训练示例都通过 NeuronTrainer 使用了这些并行特性。

如何启用 ZeRO-1？

ZeRO-1 可以通过 NeuronTrainer 或直接通过 NeuronAccelerator 启用。

通过 NeuronTrainer

from optimum.neuron import NeuronTrainingArguments, NeuronTrainer

# Enable ZeRO-1 in the training arguments
training_args = NeuronTrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=1,
    zero_1=True,  # Enable ZeRO-1
    bf16=True,
    # ... other training arguments
)

trainer = NeuronTrainer(
    model=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/training/qwen3/finetune_qwen3.py \
    --model_name_or_path Qwen/Qwen2.5-0.5B \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
    --per_device_train_batch_size 1 \
    --block_size 1024 \
    --bf16 \
    --zero_1 \
    --tensor_parallel_size 2 \
    --output_dir my_training/

通过 NeuronAccelerator

当直接使用 NeuronAccelerator 时，您需要创建一个 TrainingNeuronConfig 并单独启用 ZeRO-1。

from torch.optim import AdamW
from optimum.neuron import NeuronAccelerator
from optimum.neuron.models.training.config import TrainingNeuronConfig

# Create the training configuration
trn_config = TrainingNeuronConfig()

# Create accelerator with ZeRO-1 enabled
accelerator = NeuronAccelerator(
    trn_config=trn_config,
    zero_1=True,  # Enable ZeRO-1
    mixed_precision="bf16",
)

model = ...  # Your model instance
optimizer = AdamW(model.parameters(), lr=5e-5)

# Prepare model and optimizer
model, optimizer = accelerator.prepare(model, optimizer)

如何启用张量并行？

张量并行可以与 NeuronTrainer 或 NeuronAccelerator 一起使用。

重要提示：张量并行要求模型在 optimum.neuron.models.training 中有自定义的模型实现。

在使用张量并行时，您有几个重要的设置。

tensor_parallel_size：理想情况下，它应该是能将模型装入内存的最小值。
是否启用序列并行：序列并行在张量并行区域之外沿序列轴对激活值进行分片，通过分片激活值来节省内存。

使用分布式训练时，训练脚本由 torchrun 调用，它会将任务分派给工作单元，每个核心一个工作单元。每个工作单元将加载分片后的模型，并自动在核心之间分派参数。tensor_parallel_size 是用于分片模型参数的工作单元数量。

通过 NeuronTrainer

from optimum.neuron import NeuronTrainingArguments, NeuronTrainer

# Configure tensor parallelism in training arguments
training_args = NeuronTrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=1,
    bf16=True,
    tensor_parallel_size=8,
    # ... other training arguments
)

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

trainer.train()

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

例如：

torchrun --nproc_per_node=8 examples/training/qwen3/finetune_qwen3.py \
    --model_name_or_path Qwen/Qwen2.5-0.5B \
    --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 8 \
    --output_dir my_training/

通过 NeuronAccelerator

当直接使用 NeuronAccelerator 时，您可以通过 TrainingNeuronConfig 来配置张量并行。

from torch.optim import AdamW
from optimum.neuron import NeuronAccelerator
from optimum.neuron.models.training.config import TrainingNeuronConfig

# Configure tensor parallelism
trn_config = TrainingNeuronConfig(
    tensor_parallel_size=8,
    sequence_parallel_enabled=True,
    checkpoint_dir=None,  # Can be specified when resuming from checkpoint
)

accelerator = NeuronAccelerator(
    trn_config=trn_config,
    mixed_precision="bf16",
)

model = ...  # Your model instance
optimizer = AdamW(model.parameters(), lr=5e-5)

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

如何启用流水线并行？

流水线并行允许您将模型层拆分到多个设备上，从而能够训练那些无法在单个设备甚至单个节点上容纳的非常大的模型。

重要提示：流水线并行要求模型在 optimum.neuron.models.training 中有自定义的模型实现，并声明 SUPPORTS_PIPELINE_PARALLELISM = True。

配置选项

流水线并行有几个配置参数。

pipeline_parallel_size：流水线阶段的数量（用于拆分层的设备数量）
pipeline_parallel_num_microbatches：用于流水线调度的微批次数量
当启用流水线并行时，ZeRO-1 可以自动应用于流水线并行优化器。

通过 NeuronTrainer

from optimum.neuron import NeuronTrainingArguments, NeuronTrainer
from optimum.neuron.models.training import LlamaForCausalLM  # Custom model implementation

# Configure pipeline parallelism in training arguments
training_args = NeuronTrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=4,  # Will be split into microbatches
    bf16=True,
    tensor_parallel_size=2,
    pipeline_parallel_size=4,                    # Split model across 4 pipeline stages
    pipeline_parallel_num_microbatches=4,        # Number of microbatches
    zero_1=True,                                 # Enable ZeRO-1 with pipeline parallelism
    # ... other training arguments
)

# Load model using custom implementation - must be done with the model class directly
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B",
    trn_config=training_args.trn_config  # Pass the auto-generated trn_config
)

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

trainer.train()

通过 NeuronAccelerator

from optimum.neuron import NeuronAccelerator
from optimum.neuron.models.training.config import TrainingNeuronConfig
from optimum.neuron.models.training import LlamaForCausalLM
from torch.optim import AdamW

# Configure combined parallelism strategies
trn_config = TrainingNeuronConfig(
    tensor_parallel_size=2,
    pipeline_parallel_size=4,
    pipeline_parallel_num_microbatches=4,
    sequence_parallel_enabled=True,
)

accelerator = NeuronAccelerator(
    trn_config=trn_config,
    zero_1=True,  # Can combine with ZeRO-1
    mixed_precision="bf16",
)

# Load model with custom implementation
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B",
    trn_config=trn_config
)

optimizer = AdamW(model.parameters(), lr=5e-5)
model, optimizer = accelerator.prepare(model, optimizer)

使用流水线并行时，总进程数应至少为 tensor_parallel_size * pipeline_parallel_size。例如，当 tensor_parallel_size=2 和 pipeline_parallel_size=4 时，您总共需要 8 个进程。

组合并行策略

您可以组合多种并行策略以实现最大的内存效率和性能。以下是一个结合了所有策略的示例。

通过 NeuronTrainer

from optimum.neuron import NeuronTrainingArguments, NeuronTrainer
from optimum.neuron.models.training import LlamaForCausalLM

# Example: Combine all parallelism strategies
training_args = NeuronTrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=32,
    bf16=True,
    gradient_checkpointing=True,
    
    # ZeRO-1
    zero_1=True,
    
    # Tensor parallelism
    tensor_parallel_size=4,
    disable_sequence_parallel=False,     # Enable sequence parallelism
    
    # Pipeline parallelism
    pipeline_parallel_size=2,
    pipeline_parallel_num_microbatches=8,
    
    # Additional optimizations
    fuse_qkv=True,                      # Fuse QKV projections for efficiency
    kv_size_multiplier=None,            # Auto-calculate optimal KV multiplier
)

# Load model using custom implementation
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B",
    trn_config=training_args.trn_config
)

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

trainer.train()

此配置总共使用 4 * 2 = 8 个进程。

每个张量并行组有 4 个进程。
每个流水线阶段在一个张量并行组上运行。

然后我们可以在拥有 32 个 Neuron 核心的 trn1.32xlarge 实例上运行训练脚本，得到以下配置：dp=4, tp=4, pp=2，这意味着 4 个数据并行组，每个组有 4 个张量并行设备，以及 2 个流水线阶段。

检查点合并

由于分布式训练使用跨不同工作单元的分片检查点，您需要将它们合并以创建一个标准的模型检查点，以便在特定训练配置之外共享和使用。

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
│   ├── shards/
│   ├── tokenizer.json
│   ├── tokenizer.model
│   ├── tokenizer_config.json
│   ├── trainer_state.json
│   └── training_args.bin
├── config.json
├── special_tokens_map.json
├── 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_00_pp_rank_01
│   ├── tp_rank_01_pp_rank_01
│   ├── tp_rank_02_pp_rank_01
│   └── tp_rank_03_pp_rank_01
├── tokenizer.json
├── tokenizer.model
├── tokenizer_config.json
├── train_results.json
├── trainer_state.json
├── training_args.bin
└── trn_config.json

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

optimum-cli neuron consolidate my_training my_training_consolidated_checkpoint

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

最佳实践

选择并行策略

从张量并行开始：使用能将模型装入内存的最小 tensor_parallel_size。
添加流水线并行：对于非常大的模型，与流水线并行结合使用。
启用序列并行：在使用张量并行时始终启用以节省内存（设置 disable_sequence_parallel=False）。
使用 ZeRO-1：与任何并行策略结合使用以节省优化器内存。

内存优化

为大型模型启用 gradient_checkpointing
为流水线并行设置合适的 pipeline_parallel_num_microbatches

问题排查

常见问题

内存不足：减小批次大小、增加并行度或启用梯度检查点
模型不受支持：确保您正在使用来自 optimum.neuron.models.training 的模型
流水线并行失败：检查模型是否支持流水线并行
进程数不正确：确保 nproc_per_node 与您的并行配置匹配

调试技巧

从较小的模型和并行规模开始
检查所有进程是否能正常通信
验证检查点目录和权限
监控 Neuron 设备利用率

←Neuron 模型缓存将模型导出到 Inferentia→