Transformers 文档

在苹果硅芯片上进行PyTorch训练

在Apple芯片上进行PyTorch训练

以前,在Mac上训练模型仅限于CPU。随着PyTorch v1.12的发布,你可以利用Apple的硅GPU来训练模型,从而显著提高性能和训练速度。这是通过在PyTorch中集成Apple的Metal Performance Shaders (MPS)作为后端来实现的。MPS后端将PyTorch操作实现为自定义的Metal着色器,并将这些模块放置在mps设备上。

一些PyTorch操作尚未在MPS中实现,并且会抛出错误。为了避免这种情况,您应该设置环境变量PYTORCH_ENABLE_MPS_FALLBACK=1以使用CPU内核代替(您仍然会看到UserWarning)。


如果您遇到任何其他错误,请在PyTorch仓库中提出问题,因为Trainer仅集成了MPS后端。

设置mps设备后,您可以:

  • 在本地训练更大的网络或批量大小
  • 减少数据检索延迟,因为GPU的统一内存架构允许直接访问整个内存存储
  • 降低成本,因为您不需要在基于云的GPU上进行训练或添加额外的本地GPU

首先确保您已安装PyTorch。MPS加速支持macOS 12.3及以上版本。

pip install torch torchvision torchaudio

TrainingArguments 默认使用 mps 设备(如果可用),这意味着您不需要显式设置设备。例如,您可以运行 run_glue.py 脚本,MPS 后端会自动启用,无需进行任何更改。

export TASK_NAME=mrpc

python examples/pytorch/text-classification/run_glue.py \
  --model_name_or_path google-bert/bert-base-cased \
  --task_name $TASK_NAME \
- --use_mps_device \
  --do_train \
  --do_eval \
  --max_seq_length 128 \
  --per_device_train_batch_size 32 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir /tmp/$TASK_NAME/ \
  --overwrite_output_dir

不支持像 gloonccl 这样的分布式设置的后端,这意味着你只能在单个 GPU 上使用 MPS 后端进行训练。

您可以在Introducing Accelerated PyTorch Training on Mac博客文章中了解更多关于MPS后端的信息。

< > Update on GitHub