ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

MPS 训练完全指南:在 Apple Silicon GPU 上使用 PyTorch Lightning

MPS 训练完全指南:在 Apple Silicon GPU 上使用 PyTorch Lightning MPS 训练完全指南在 Apple Silicon GPU 上使用 PyTorch Lightning【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000 GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning核心主题本文基于 PyTorch Lightning 官方文档中的 MPS 基础指南docs/source-pytorch/accelerators/mps_basic.rst系统讲解如何在 Apple SiliconM 系列芯片的 GPU 上使用 Lightning 进行训练包括硬件原理、一行代码切换加速器、MPS 后端特性、常见故障排查以及配套源码实现细节。适用读者想在 MacBook / Mac Mini / Mac Studio / iMac 等 Apple Silicon 设备上用 GPU 训练 PyTorch 模型的开发者。读完本文后你将掌握如何用Trainer(acceleratormps)启用 MPS 加速、MPS 加速器在 Lightning 内部的工作原理与设备限制、遇到MPSAccelerator can not run on your system等异常时的排查方法。Apple Silicon 是什么Apple silicon 芯片是 Apple 基于 ARM 架构设计的一体化片上系统system on a chipSoC。与传统的分离式硬件不同它在同一块芯片上集成了多种计算单元CPU 核心负责通用计算与串行逻辑GPU 核心负责大规模并行计算是深度学习训练的主要加速单元神经引擎Neural Engine针对机器学习推理等特定负载优化的专用硬件统一内存shared memory以上所有单元共享同一块内存数据无需在 CPU 内存与显存之间复制这也是 M 系列芯片能高效跑模型的关键硬件基础。它只是一颗 CPU 吗不是。Apple silicon 芯片内部确实包含 CPU 核心但如果只用CPUAccelerator在 CPU 上运行你无法发挥这颗芯片的全部硬件加速潜力——因为芯片上还同时存在 GPU 与神经引擎它们具备 CPU 不具备的并行算力。PyTorch 通过Metal Performance ShadersMPS后端把这些 GPU 核心暴露给深度学习框架。在 Lightning 中对应的支持实现是MPSAccelerator。它在两个子包中均有实现且lightning.pytorch的版本直接复用lightning.fabric的底层检测逻辑src/lightning/pytorch/accelerators/mps.pyPyTorch Trainer 使用的MPSAccelerator继承自Accelerator基类src/lightning/fabric/accelerators/mps.pyFabric 与 PyTorch 共用的底层实现包含is_available()等核心检测逻辑。在 Apple silicon GPU 上运行MPS 设备启用方式非常简单在Trainer中把accelerator设为mps并指定设备数量即可from lightning.pytorch import Trainer trainer Trainer(acceleratormps, devices1)你也可以直接传入加速器对象from lightning.pytorch.accelerators import MPSAccelerator trainer Trainer(acceleratorMPSAccelerator(), devices1)注意MPSAccelerator一次只支持 1 个设备。目前市面上还没有搭载多块 MPS 能力 GPU 的机器M 系列芯片的单颗 SoC 只暴露一个 MPS 设备。从源码看这一限制是硬编码的MPSAccelerator.auto_device_count()返回固定值1_get_all_available_mps_gpus()在加速器可用时只返回[0]。因此如果你试图请求超过 1 个设备例如devices2或devices[0, 2]设备解析逻辑会抛出MisconfigurationException: ... But your machine only has: [0]这一点由测试 tests/tests_fabric/accelerators/test_mps.py 专门验证。MPS 在 Lightning 中的设备解析流程当你在Trainer(acceleratormps, devices1)中传入devices后Lightning 会调用MPSAccelerator.parse_devices()其内部最终走向 src/lightning/fabric/utilities/device_parser.py 的_parse_gpu_ids(gpus, include_mpsTrue)devices1→ 归一化为[0]devices[0]、devices0,→ 直接解析为[0]devices-1或-1表示“使用所有可用设备”→ 由于 MPS 机器只有 1 个设备仍解析为[0]请求2个及以上设备 → 抛出MisconfigurationException。随后get_parallel_devices()将解析结果映射为torch.device(mps, 0)训练便会在 MPS 设备上执行。测试 tests/tests_pytorch/accelerators/test_mps.py 中覆盖了devices取1、[0]、-1三种写法都能正常训练一个完整的BoringModel。自动选择加速器如果不显式指定acceleratorLightning 会根据硬件可用性自动选择。在 device_parser.py 的_select_auto_accelerator()中检测顺序为TPU →MPS→ CUDA → CPU。也就是说在一台没有 XLA/TPU 的 Apple Silicon Mac 上Trainer()会自动选中 MPS 作为加速器而在 NVIDIA GPU 机器上则自动选中 CUDA。把 batch 数据搬到 MPS 设备Lightning 的策略层负责把每个 batch 转移到目标设备。在 MPS 场景下batch_to_device支持几乎所有常见数据结构测试 tests/tests_pytorch/accelerators/test_mps.py 覆盖了纯 TensorTensor 列表、列表的列表Tensor 字典、字典列表元组混合嵌套、namedtuple定义了.to()方法的自定义对象。这些结构都会被递归地迁移到torch.mps.FloatTensor。简单类型None、{}、[]、1.0、字符串等则原样返回不会报错。在 Lightning Fabric 中使用 MPS如果你在使用更轻量的 Lightning Fabric写法等价from lightning.fabric import Fabric fabric Fabric(acceleratormps, devices1)命令行方式同样支持--accelerator mps见 docs/source-fabric/fundamentals/launch.rst。Fabric 的MPSAccelerator与 PyTorch 版本共享is_available()、parse_devices()等核心逻辑区别仅在于setup_device()校验失败时抛出的异常类型Fabric 抛ValueErrorPyTorch 抛MisconfigurationException。MPS 代表什么MPS 是Metal Performance Shaders的缩写这是 Apple 的Metal图形与计算框架中用于 GPU 通信与计算的技术Metal 是 Apple 面向 GPU 编程的底层 API。PyTorch 的 MPS 后端通过它把张量运算映射到 Apple silicon 的 GPU 核心上执行从而获得比 CPU 高得多的并行吞吐。当前状态提醒MPS 加速器与 PyTorch 的 MPS 后端目前仍处于实验阶段。并非所有算子都已被支持。不过随着 PyTorch 团队的持续开发可用算子的覆盖范围正在不断扩大。不支持算子的回退方案如果在训练中遇到某个算子尚未被 MPS 后端支持可以通过设置环境变量让 PyTorch 自动回退到 CPU 执行该算子PYTORCH_ENABLE_MPS_FALLBACK1 python your_script.py这条命令的作用是当 MPS 后端遇到不支持的算子时回退到 CPU 实现而不是直接报错从而让实验性功能尽可能跑通。禁用 MPS 的开关Lightning 还提供了一个显式禁用 MPS 的环境变量DISABLE_MPS。在 src/lightning/fabric/accelerators/mps.py 的is_available()实现中mps_disabled os.getenv(DISABLE_MPS, 0) 1 return not mps_disabled and torch.backends.mps.is_available() and platform.processor() in (arm, arm64)它同时满足三个条件才算“MPS 可用”未设置DISABLE_MPS1torch.backends.mps.is_available()返回True即你安装的 PyTorch 版本编译并启用了 MPS 后端CPU 架构是arm或arm64即真实运行在 Apple Silicon 上而不是被转译的 x86 环境。结果会被lru_cache缓存避免重复检测的开销。故障排查Troubleshooting报错MPSAccelerator can not run on your system since the accelerator is not available如果 Lightning 无法在你的系统上检测到 Apple silicon 硬件Trainer(acceleratormps)会抛出如下异常MisconfigurationException: MPSAccelerator can not run on your system since the accelerator is not available.最可能的原因你的 Python 正在被模拟emulated运行导致系统认为你身处 Intel CPU 环境。当你在一台 ARM 版 Mac 上安装了为 Intel 编译的 Python 解释器时macOS 的 Rosetta 转译层会把它模拟成 x86 进程。此时platform.processor()返回的不是arm64MPSAccelerator.is_available()的三个条件无法全部满足Lightning 便认为 MPS 不可用——尽管你的硬件完全支持。解决方法重新安装原生 ARMApple siliconM1/M2 及更新版本的 Python。具体来说卸载当前被模拟的 Python 可执行文件如果你使用 conda / miniconda 等环境管理器需要连同环境管理器一起重新安装ARM 版本conda 官方提供 Apple M1/M2 版安装包注意选择 Apple silicon 版本而非 Intel 版本不要选 x86_64 的安装包重新创建虚拟环境并安装 PyTorch需使用支持 MPS 的版本PyTorch 2.0 起 macOS 官方构建默认包含 MPS 支持。安装完成后可以通过一个最小脚本快速验证import torch from lightning.pytorch.accelerators import MPSAccelerator print(PyTorch MPS backend available:, torch.backends.mps.is_available()) print(Lightning MPSAccelerator available:, MPSAccelerator.is_available())两个输出都为True时即可正常使用Trainer(acceleratormps, devices1)。其他常见疑问“GPU available but not used”警告如果在 MPS 机器上显式使用acceleratorcpuLightning 会发出GPU available but not used的用户警告见测试 tests/tests_pytorch/accelerators/test_mps.py。这是提示你有 MPS 设备可用但当前配置没有使用它。如何确认训练确实跑在 MPS 上训练开始时的日志中Lightning 会打印类似GPU available: True (mps)的信息也可通过trainer.strategy.batch_to_device(batch, torch.device(mps))后检查张量类型是否为torch.mps.FloatTensor来确认。进阶监控 MPS 设备状态除训练外Lightning 还提供了 MPS 设备状态监控能力。MPSAccelerator.get_device_stats()返回当前机器的 CPU GPU 整体状态M 系列为统一内存架构故统计的是整机指标字段包括M1_vm_percent虚拟内存virtual memory使用百分比M1_percentCPU 使用百分比M1_swap_percent交换内存swap使用百分比。该功能依赖psutil包未安装时会抛出ModuleNotFoundError提示。可通过 tests/tests_pytorch/accelerators/test_mps.py 中展示的方式获取from lightning.pytorch.accelerators import MPSAccelerator import torch stats MPSAccelerator().get_device_stats(torch.device(mps)) print(stats)对于长时间训练可以利用这些指标结合 Lightning 的 DeviceStatsMonitor 回调定期记录帮助定位显存/内存压力问题。小结在 Apple silicon 上使用 PyTorch Lightning 训练模型核心步骤只有一句话把Trainer的accelerator设为mpsdevices设为1。其余工作——设备检测、batch 迁移、单设备策略选择——都由 Lightning 的MPSAccelerator自动完成。需要记住的三件事MPS 目前仅支持单设备MPSAccelerator.auto_device_count()固定返回 1多设备请求会抛出MisconfigurationException实验性功能PyTorch 的 MPS 后端尚在完善中遇到不支持的算子可用PYTORCH_ENABLE_MPS_FALLBACK1回退到 CPU报错先查 Python 架构MPSAccelerator can not run on your system最常由 Rosetta 模拟的 x86 Python 导致重新安装 Apple silicon 原生构建即可解决。【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000 GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表