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