Hugging Face 的 LOGO

AWS Trainium & Inferentia 文件

使用 optimum-neuron 進行分散式訓練

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

使用 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 中有自定義的模型實現。

在使用張量並行時，您有幾個重要的設定。

  1. tensor_parallel_size：理想情況下，它應該是能將模型裝入記憶體的最小值。
  2. 是否啟用序列並行：序列並行在張量並行區域之外沿序列軸對啟用值進行分片，透過分片啟用值來節省記憶體。

使用分散式訓練時，訓練指令碼由 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 目錄本身作為輸入。

最佳實踐

選擇並行策略

  1. 從張量並行開始：使用能將模型裝入記憶體的最小 tensor_parallel_size。
  2. 新增流水線並行：對於非常大的模型，與流水線並行結合使用。
  3. 啟用序列並行：在使用張量並行時始終啟用以節省記憶體（設定 disable_sequence_parallel=False）。
  4. 使用 ZeRO-1：與任何並行策略結合使用以節省最佳化器記憶體。

記憶體最佳化

  • 為大型模型啟用 gradient_checkpointing
  • 為流水線並行設定合適的 pipeline_parallel_num_microbatches

問題排查

常見問題

  1. 記憶體不足：減小批次大小、增加並行度或啟用梯度檢查點
  2. 模型不受支援：確保您正在使用來自 optimum.neuron.models.training 的模型
  3. 流水線並行失敗：檢查模型是否支援流水線並行
  4. 程序數不正確：確保 nproc_per_node 與您的並行配置匹配

除錯技巧

  • 從較小的模型和並行規模開始
  • 檢查所有程序是否能正常通訊
  • 驗證檢查點目錄和許可權
  • 監控 Neuron 裝置利用率

←Neuron 模型快取 將模型匯出到 Inferentia→
© . This site is unofficial and not affiliated with Hugging Face, Inc.