加速

Accelerate 是一个旨在简化 PyTorch 分布式训练的库，它将最常用的框架（Fully Sharded Data Parallel (FSDP) 和 DeepSpeed）整合到一个统一的接口中。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

Trainer

将保存的配置文件的路径传递给 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_strategy="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 方法，以缩放梯度并确定要使用的适当的 backward 方法，具体取决于您的框架（例如，DeepSpeed 或 Megatron）。

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 上更新