Accelerate

Accelerate 是一個庫，旨在透過將最常用的框架（Fully Sharded Data Parallel (FSDP) 和 DeepSpeed）統一到單個介面中，從而簡化 PyTorch 在任何型別設定上的分散式訓練。Trainer 在底層由 Accelerate 提供支援，支援載入大型模型和分散式訓練。

本指南將向您展示兩種使用 Accelerate 和 Transformers 的方法，均使用 FSDP 作為後端。第一種方法演示了使用 Trainer 進行分散式訓練，第二種方法演示瞭如何調整 PyTorch 訓練迴圈。有關 Accelerate 的更多詳細資訊，請參閱文件。

pip install accelerate

首先，在命令列中執行 accelerate config，回答一系列關於您的訓練系統的問題。這將建立並儲存一個配置檔案，以幫助 Accelerate 根據您的設定正確配置訓練。

accelerate config

根據您的設定和您提供的答案，一個在具有兩個 GPU 的一臺機器上使用 FSDP 分散式訓練的示例配置檔案可能如下所示。

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

訓練器

將儲存的配置檔案的路徑傳遞給 TrainingArguments，然後將您的 TrainingArguments 傳遞給 Trainer。

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="your-model",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=2,
    fsdp_config="path/to/fsdp_config",
    fsdp="full_shard",
    weight_decay=0.01,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    processing_class=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

原生 PyTorch

Accelerate 也可以新增到任何 PyTorch 訓練迴圈中以啟用分散式訓練。Accelerator 是使您的 PyTorch 程式碼與 Accelerate 協同工作的主要入口點。它會自動檢測您的分散式訓練設定並初始化訓練所需的所有元件。您無需顯式地將模型放置在裝置上，因為 Accelerator 知道要將模型移動到哪個裝置。

from accelerate import Accelerator

accelerator = Accelerator()
device = accelerator.device

所有 PyTorch 物件（模型、最佳化器、排程器、資料載入器）現在都應傳遞給 prepare 方法。此方法將您的模型移動到適當的裝置或裝置組，使最佳化器和排程器適應使用 AcceleratedOptimizer 和 AcceleratedScheduler，並建立一個新的可分片資料載入器。

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

將訓練迴圈中的 loss.backward 替換為 Accelerate 的 backward 方法，以縮放梯度並根據您的框架（例如 DeepSpeed 或 Megatron）確定要使用的適當 backward 方法。

for epoch in range(num_epochs):
    for batch in train_dataloader:
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

將所有內容組合成一個函式，並使其可以作為指令碼呼叫。

from accelerate import Accelerator
  
def main():
  accelerator = Accelerator()

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

  for batch in training_dataloader:
      optimizer.zero_grad()
      inputs, targets = batch
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
      accelerator.backward(loss)
      optimizer.step()
      scheduler.step()

if __name__ == "__main__":
    main()

從命令列呼叫 accelerate launch 來執行您的訓練指令碼。此處也可以傳遞任何其他引數。

要在兩個 GPU 上啟動您的訓練指令碼，請新增 --num_processes 引數。

accelerate launch --num_processes=2 your_script.py

有關更多詳細資訊，請參閱啟動 Accelerate 指令碼。

< > 在 GitHub 上更新