在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
)。
设置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
不支持像 gloo
和 nccl
这样的分布式设置的后端,这意味着你只能在单个 GPU 上使用 MPS 后端进行训练。
您可以在Introducing Accelerated PyTorch Training on Mac博客文章中了解更多关于MPS后端的信息。
< > Update on GitHub