【Bug已解决】How to move PyTorch model to GPU on Apple M1 chips? 解决方案问题描述Apple SiliconM1/M2/M3/M4芯片的推出为机器学习开发带来了新的硬件选择。这些芯片内置了强大的 GPU通过统一内存架构但它们不使用 NVIDIA 的 CUDA 平台而是使用 Apple 自己的 Metal Performance ShadersMPS框架。这导致了一个常见的困惑如何在 Apple M1 芯片上利用 GPU 加速 PyTorch 模型训练具体问题表现为torch.cuda.is_available()返回False因为 M1 没有 CUDAmodel.to(cuda)报错CUDA 不可用开发者不知道 MPS 后端的存在模型只能在 CPU 上运行无法利用 M1 的 GPU从 CUDA 代码迁移到 MPS 时遇到各种兼容性问题某些操作在 MPS 上不支持需要回退到 CPU错误复现场景一尝试使用 CUDAimport torch import torch.nn as nn # 检查 CUDA在 M1 上返回 False print(fCUDA available: {torch.cuda.is_available()}) # False # 尝试使用 CUDA model nn.Linear(10, 5) model model.to(cuda) # 报错报错AssertionError: Torch not compiled with CUDA enabled场景二不知道 MPS 后端import torch # 开发者可能只检查 CUDA if torch.cuda.is_available(): device cuda else: device cpu # M1 上会回退到 CPU浪费 GPU 资源 print(fUsing device: {device}) # cpu场景三MPS 上的操作不支持import torch # 正确使用 MPS device torch.device(mps) x torch.randn(10, 10, devicedevice) # 某些操作在 MPS 上不支持 result torch.fft.fft(x) # 可能报错报错RuntimeError: Placeholder only has a single piece of memory, but multiple views were given.或NotImplementedError: The following operators are not implemented for MPS: ...场景四MPS 内存不足import torch device torch.device(mps) # 创建大型张量 x torch.randn(10000, 10000, 10000, devicedevice) # 可能导致内存错误报错RuntimeError: MPS exceeded available system memory根因分析1. Apple Silicon 不支持 CUDACUDA 是 NVIDIA 专有的 GPU 计算平台。Apple Silicon 使用 ARM 架构和 Apple 自研的 GPU不兼容 CUDA。PyTorch 的 CUDA 后端无法在 Apple Silicon 上使用。2. MPS 后端的引入PyTorch 从 1.12 版本开始引入 MPSMetal Performance Shaders后端为 Apple Silicon 提供 GPU 加速能力。MPS 使用 Apple 的 Metal 框架直接与 GPU 通信。3. MPS 后端的成熟度与 CUDA 相比MPS 后端相对较新存在以下限制不是所有 PyTorch 操作都支持 MPS某些操作的 MPS 实现可能与 CUDA/CPU 有数值精度差异MPS 的性能在某些操作上不如 CUDA调试工具不如 CUDA 丰富4. 统一内存架构Apple Silicon 使用统一内存架构Unified Memory ArchitectureCPU 和 GPU 共享同一块物理内存。这与 CUDA 的离散内存模型不同带来了不同的内存管理挑战不需要 CPU-GPU 数据拷贝但 PyTorch 的 API 仍然模拟了这种拷贝内存上限受系统总内存限制而不是独立的 GPU 显存大型模型可能占用过多系统内存5. PyTorch 版本要求MPS 后端需要 PyTorch 1.12 和 macOS 12.3。旧版本的 PyTorch 不支持 MPS。解决方案方案一使用 MPS 后端import torch # 检查 MPS 是否可用 if torch.backends.mps.is_available(): device torch.device(mps) elif torch.cuda.is_available(): device torch.device(cuda) else: device torch.device(cpu) print(fUsing device: {device}) # 使用设备 model model.to(device) x x.to(device)方案二升级 PyTorch 到支持 MPS 的版本# 安装最新版 PyTorch支持 MPS pip install torch torchvision torchaudio # 或指定版本 pip install torch1.12 torchvision torchaudio方案三处理 MPS 不支持的操作import torch device torch.device(mps) def safe_to_mps(tensor, devicedevice): 安全地将张量移到 MPS不支持时回退到 CPU try: return tensor.to(device) except RuntimeError: return tensor def safe_op_mps(op, *args, **kwargs): 安全地在 MPS 上执行操作不支持时回退到 CPU try: return op(*args, **kwargs) except (RuntimeError, NotImplementedError): # 将所有参数移到 CPU cpu_args [a.cpu() if isinstance(a, torch.Tensor) else a for a in args] cpu_kwargs {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()} result op(*cpu_args, **cpu_kwargs) # 将结果移回 MPS if isinstance(result, torch.Tensor): return result.to(device) return result完整修复代码以下是一个完整的 Apple Silicon PyTorch GPU 训练方案 Apple Silicon (M1/M2/M3/M4) PyTorch GPU 训练完整方案 使用 MPS (Metal Performance Shaders) 后端 import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset import time import warnings def get_device(): 自动选择最佳设备 优先级: CUDA MPS CPU if torch.cuda.is_available(): device torch.device(cuda) device_name torch.cuda.get_device_name(0) elif torch.backends.mps.is_available(): device torch.device(mps) device_name Apple Silicon GPU (MPS) else: device torch.device(cpu) device_name CPU print(fUsing device: {device} ({device_name})) return device def check_mps_availability(): 检查 MPS 后端的可用性 print( * 60) print(MPS 后端可用性检查) print( * 60) print(fPyTorch version: {torch.__version__}) print(fMPS available: {torch.backends.mps.is_available()}) print(fMPS built: {torch.backends.mps.is_built()}) if torch.backends.mps.is_available(): print(\n[OK] MPS 后端可用) print(可以在 Apple Silicon GPU 上运行 PyTorch。) else: print(\n[WARNING] MPS 后端不可用。) print(可能原因:) print( 1. PyTorch 版本 1.12) print( 2. macOS 版本 12.3) print( 3. 不是 Apple Silicon 芯片) print(\n升级命令: pip install --upgrade torch) class MPSModel(nn.Module): 示例模型 def __init__(self, input_dim784, hidden_dim256, output_dim10): super().__init__() self.fc1 nn.Linear(input_dim, hidden_dim) self.bn1 nn.BatchNorm1d(hidden_dim) self.fc2 nn.Linear(hidden_dim, hidden_dim) self.bn2 nn.BatchNorm1d(hidden_dim) self.fc3 nn.Linear(hidden_dim, output_dim) self.dropout nn.Dropout(0.2) def forward(self, x): x torch.relu(self.bn1(self.fc1(x))) x self.dropout(x) x torch.relu(self.bn2(self.fc2(x))) x self.dropout(x) return self.fc3(x) def train_on_mps(num_epochs10, batch_size128, lr0.001): 在 MPS 上训练模型 print(\n * 60) print(MPS GPU 训练) print( * 60) device get_device() # 生成数据 num_samples 20000 X torch.randn(num_samples, 784) y torch.randint(0, 10, (num_samples,)) dataset TensorDataset(X, y) # DataLoader 配置 # MPS 支持 pin_memory在 Apple Silicon 上效果不同于 CUDA loader DataLoader( dataset, batch_sizebatch_size, shuffleTrue, num_workers0, # MPS 上多进程可能有兼容性问题 pin_memory(device.type cuda) # MPS 不需要 pin_memory ) # 创建模型 model MPSModel().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lrlr) scheduler optim.lr_scheduler.StepLR(optimizer, step_size3, gamma0.5) # 训练循环 start_time time.time() for epoch in range(num_epochs): model.train() epoch_loss 0.0 correct 0 total 0 for batch_X, batch_y in loader: batch_X batch_X.to(device) batch_y batch_y.to(device) optimizer.zero_grad() outputs model(batch_X) loss criterion(outputs, batch_y) loss.backward() optimizer.step() epoch_loss loss.item() _, predicted outputs.max(1) correct predicted.eq(batch_y).sum().item() total batch_y.size(0) scheduler.step() acc 100. * correct / total print(fEpoch {epoch1}/{num_epochs} | Loss: {epoch_loss/len(loader):.4f} | Acc: {acc:.2f}%) elapsed time.time() - start_time print(f\nTotal training time: {elapsed:.2f}s) return model def benchmark_cpu_vs_mps(): CPU vs MPS 性能基准测试 print(\n * 60) print(CPU vs MPS 性能基准测试) print( * 60) if not torch.backends.mps.is_available(): print(MPS 不可用跳过基准测试) return # 测试参数 sizes [(100, 784), (1000, 784), (5000, 784)] hidden_dim 512 for input_size, input_dim in sizes: print(f\n输入形状: ({input_size}, {input_dim})) x torch.randn(input_size, input_dim) # CPU 测试 model_cpu nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 10) ) start time.time() for _ in range(100): with torch.no_grad(): _ model_cpu(x) cpu_time time.time() - start # MPS 测试 model_mps copy_model_to_mps(model_cpu) x_mps x.to(mps) # 预热 with torch.no_grad(): _ model_mps(x_mps) start time.time() for _ in range(100): with torch.no_grad(): _ model_mps(x_mps) mps_time time.time() - start speedup cpu_time / mps_time if mps_time 0 else 0 print(f CPU: {cpu_time:.4f}s | MPS: {mps_time:.4f}s | 加速比: {speedup:.2f}x) def copy_model_to_mps(model): 将模型复制到 MPS 设备 return model.to(mps) def demo_mps_fallback(): 演示 MPS 不支持操作的回退机制 print(\n * 60) print(MPS 操作回退演示) print( * 60) if not torch.backends.mps.is_available(): print(MPS 不可用) return device torch.device(mps) # 一些可能在 MPS 上不支持的操作 operations [ (矩阵乘法, lambda: torch.mm( torch.randn(100, 100, devicedevice), torch.randn(100, 100, devicedevice) )), (卷积, lambda: nn.functional.conv2d( torch.randn(1, 3, 32, 32, devicedevice), torch.randn(16, 3, 3, 3, devicedevice) )), (BatchNorm, lambda: nn.BatchNorm1d(100).to(device)( torch.randn(32, 100, devicedevice) )), (注意力机制, lambda: nn.MultiheadAttention(128, 4).to(device)( torch.randn(10, 32, 128, devicedevice), torch.randn(10, 32, 128, devicedevice), torch.randn(10, 32, 128, devicedevice) )), ] for name, op in operations: try: result op() print(f [OK] {name}: {type(result)}, shape{result.shape if hasattr(result, shape) else N/A}) except (RuntimeError, NotImplementedError) as e: print(f [FAIL] {name}: {str(e)[:80]}...) # 回退到 CPU print(f - 回退到 CPU) def demo_memory_management(): 演示 MPS 内存管理 print(\n * 60) print(MPS 内存管理) print( * 60) if not torch.backends.mps.is_available(): print(MPS 不可用) return device torch.device(mps) # 清空 MPS 缓存 if hasattr(torch.mps, empty_cache): torch.mps.empty_cache() print(已清空 MPS 缓存) # 检查内存使用 if hasattr(torch.mps, current_allocated_memory): mem_before torch.mps.current_allocated_memory() / 1024**2 print(f当前 MPS 内存使用: {mem_before:.2f} MB) # 分配一些张量 tensors [] for i in range(5): t torch.randn(1000, 1000, devicedevice) tensors.append(t) if hasattr(torch.mps, current_allocated_memory): mem torch.mps.current_allocated_memory() / 1024**2 print(f 分配张量 {i1} 后: {mem:.2f} MB) # 释放 del tensors if hasattr(torch.mps, empty_cache): torch.mps.empty_cache() if hasattr(torch.mps, current_allocated_memory): mem_after torch.mps.current_allocated_memory() / 1024**2 print(f释放后 MPS 内存使用: {mem_after:.2f} MB) def main(): 主函数 print(\n * 60) print( Apple Silicon PyTorch GPU 训练指南) print( * 60) check_mps_availability() if torch.backends.mps.is_available(): train_on_mps(num_epochs5) benchmark_cpu_vs_mps() demo_mps_fallback() demo_memory_management() else: print(\nMPS 不可用使用 CPU 演示) train_on_mps(num_epochs3) print(\n * 60) print( 演示完成) print( * 60) if __name__ __main__: main()常见陷阱与注意事项1. MPS 的数值精度差异MPS 后端使用 Float16 或修改过的 Float32 计算可能与 CPU/CUDA 有微小的数值差异。在需要高精度的场景中如科学计算需要注意这种差异。# CPU 和 MPS 的结果可能有微小差异 x_cpu torch.randn(1000, 1000) x_mps x_cpu.to(mps) result_cpu torch.mm(x_cpu, x_cpu) result_mps torch.mm(x_mps, x_mps) diff (result_cpu - result_mps.cpu()).abs().max() print(f最大差异: {diff}) # 通常在 1e-5 量级2. 不是所有操作都支持 MPS某些 PyTorch 操作在 MPS 上没有实现。遇到NotImplementedError时需要将相关计算移到 CPUtry: result some_operation(x_mps) except NotImplementedError: result some_operation(x_mps.cpu()).to(mps)3. MPS 的num_workers兼容性在 DataLoader 中使用num_workers 0时MPS 可能出现多进程兼容性问题。建议在 MPS 上使用num_workers0或测试多进程是否正常工作。4. MPS 内存是共享内存Apple Silicon 的统一内存架构意味着 GPU 和 CPU 共享同一块物理内存。大型模型可能占用过多系统内存导致系统变慢或其他应用崩溃。使用torch.mps.empty_cache()释放未使用的内存。5. MPS 不支持pin_memoryCUDA 使用pin_memory来加速 CPU-GPU 数据传输。MPS 由于统一内存架构不需要pin_memory设置它不会有任何效果。6. 模型保存和加载在 MPS 上训练的模型可以正常保存和加载。保存的是模型参数与设备无关加载时可以指定任意设备# 保存在 MPS 上训练的模型 torch.save(model.state_dict(), model.pth) # 加载到任意设备 model.load_state_dict(torch.load(model.pth, map_locationcpu)) model model.to(cuda) # 或 mps 或 cpu7. MPS 的性能特点MPS 在某些操作上比 CPU 快很多如大型矩阵乘法、卷积但在某些操作上可能比 CPU 慢如小型张量操作、复杂的索引操作。建议对不同操作进行基准测试选择最优的执行设备。8. macOS 版本和 PyTorch 版本确保 macOS 12.3 和 PyTorch 1.12。较新的 macOS 和 PyTorch 版本通常有更好的 MPS 支持和性能。9. 调试 MPS 问题MPS 的调试工具不如 CUDA 丰富。遇到问题时可以使用PYTORCH_ENABLE_MPS_FALLBACK1环境变量自动回退不支持的操作到 CPU检查 PyTorch 的 MPS 兼容性列表在 PyTorch GitHub 上搜索相关问题# 设置环境变量在运行 Python 之前 # export PYTORCH_ENABLE_MPS_FALLBACK110. 从 CUDA 代码迁移从 CUDA 代码迁移到 MPS 通常只需要将cuda替换为mps# CUDA 代码 device torch.device(cuda if torch.cuda.is_available() else cpu) # 兼容 CUDA 和 MPS 的代码 if torch.cuda.is_available(): device torch.device(cuda) elif torch.backends.mps.is_available(): device torch.device(mps) else: device torch.device(cpu)总结在 Apple M1 芯片上使用 PyTorch GPU 加速的关键是使用 MPS 后端而非 CUDA检查 MPS 可用性torch.backends.mps.is_available()使用 MPS 设备device torch.device(mps)处理不支持的操作回退到 CPU 执行注意内存管理使用torch.mps.empty_cache()释放内存注意数值精度MPS 可能与 CPU/CUDA 有微小差异通过本文提供的完整代码可以在 Apple Silicon 上实现 GPU 加速的 PyTorch 训练。虽然 MPS 的性能和功能不如 CUDA 成熟但对于开发、原型验证和中小规模训练来说M1 的 GPU 已经提供了显著的加速效果。对于大规模生产训练仍然推荐使用 NVIDIA GPU CUDA 环境。
)
