使用 PyTorch 在 Apple 芯片上训练模型¶
之前，在 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+。
In [ ]:
pip install torch torchvision torchaudio
TrainingArguments 默认情况下会使用 mps 设备，如果它可用的话，这意味着你不需要显式设置设备。例如，你可以运行 run_glue.py 脚本，而无需进行任何更改即可自动启用 MPS 后端。
In [ ]:
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
对于 分布式设置的后端如 gloo 和 nccl，mps 设备不支持这些后端，这意味着你只能在单个 GPU 上使用 MPS 后端进行训练。
了解更多关于 MPS 后端的信息，可以阅读 加速 PyTorch 在 Mac 上的训练 博客文章。