【Bug已解决】How to run Pytorch on Macbook pro (M1) GPU? 解决方案问题描述随着 Apple SiliconM1/M2/M3/M4芯片的普及越来越多的开发者在 MacBook Pro 上进行深度学习开发。然而PyTorch 最初是为 NVIDIA CUDA 设计的在 Apple Silicon GPU 上运行 PyTorch 需要使用特殊的 MPSMetal Performance Shaders后端。开发者常遇到的问题包括torch.cuda.is_available()返回False因为 M1 没有 CUDA安装标准 PyTorch 后无法利用 GPU 加速MPS 后端的某些算子不支持导致 fallback 到 CPU内存管理差异导致 OOM 错误混合精度训练在 MPS 上行为不一致某些模型架构在 MPS 上出现 NaN 或精度问题Apple Silicon 的 GPU 使用统一内存架构Unified MemoryCPU 和 GPU 共享同一块物理内存。这与 NVIDIA 的独立显存架构有根本性差异导致内存管理策略需要调整。错误复现以下代码演示了在 MacBook Pro M1 上使用 PyTorch 时常见的问题import torch # 标准方式检测 GPU对 M1 无效 print(fPyTorch 版本: {torch.__version__}) print(fCUDA 是否可用: {torch.cuda.is_available()}) # False # 尝试使用 CUDA try: x torch.randn(3, 3).cuda() except Exception as e: print(fCUDA 错误: {e}) # RuntimeError: CUDA is not available # 检查 MPS 是否可用 if hasattr(torch.backends, mps): print(fMPS 是否可用: {torch.backends.mps.is_available()}) print(fMPS 是否已构建: {torch.backends.mps.is_built()}) else: print(当前 PyTorch 版本不支持 MPS需要 PyTorch 1.12) # 尝试使用 MPS try: device torch.device(mps) x torch.randn(3, 3, devicedevice) print(fMPS 张量创建成功: {x.device}) except Exception as e: print(fMPS 错误: {e}) # 可能的错误: # RuntimeError: Placeholder storage has not been allocated on MPS device # NotImplementedError: ... not implemented for MPS # 某些算子在 MPS 上不支持 try: x torch.randn(3, 3, devicemps) # 某些操作可能不支持 result torch.fft.fft(x) # 可能在 MPS 上不支持 except NotImplementedError as e: print(f算子不支持: {e}) # NotImplementedError: The function fft is not currently implemented for the MPS device.根因分析1. Apple Silicon GPU 架构Apple M1/M2/M3 芯片集成了 GPU使用 Apple 自研的 Metal 图形 API 进行编程。PyTorch 通过 MPSMetal Performance Shaders后端来利用 Apple GPU。MPS 是 Apple 提供的高性能 GPU 计算框架类似于 NVIDIA 的 cuDNN。2. MPS 后端的实现状态PyTorch 从 1.12 版本开始实验性支持 MPS 后端到 2.0 版本后逐渐成熟。但并非所有 PyTorch 算子都在 MPS 上实现了。当遇到未实现的算子时PyTorch 会抛出NotImplementedError需要手动将张量移到 CPU 执行后再移回 MPS。3. 统一内存架构的影响Apple Silicon 的统一内存架构意味着 CPU 和 GPU 共享同一块物理内存。这带来了以下影响无需数据拷贝CPU 和 GPU 之间的数据传输几乎是零成本的内存总量受限GPU 可用的内存就是系统总内存如 M1 16GB没有独立的显存内存竞争CPU 和 GPU 的内存使用会互相影响4. FP16/BF16 精度问题MPS 后端对 FP16半精度浮点数的支持与 CUDA 有差异。某些算子在 FP16 模式下可能产生 NaN 或精度损失。此外MPS 对 BF16Brain Float 16的支持也有限。5. 随机数生成器差异MPS 后端使用与 CPU/CUDA 不同的随机数生成器实现这意味着相同的种子在 MPS 和 CPU 上会产生不同的随机数序列。这对于需要可重复性的实验是一个重要问题。6. 内存管理差异在 CUDA 上torch.cuda.empty_cache()可以释放缓存的 GPU 内存。在 MPS 上内存管理机制不同因为使用的是统一内存。MPS 的内存释放行为与 CUDA 不完全一致可能导致内存碎片或 OOM。解决方案方案一安装支持 MPS 的 PyTorch确保安装了支持 MPS 的 PyTorch 版本1.12# 方法1使用 pip 安装推荐 # 对于 Apple Silicon Macpip 会自动安装 ARM64 版本 pip install torch torchvision torchaudio # 方法2使用 conda 安装 conda install pytorch torchvision torchaudio -c pytorch # 验证 MPS 支持 python -c import torch; print(torch.backends.mps.is_available()) # 应输出: True方案二正确使用 MPS 设备import torch def get_device(): 获取最佳可用设备 if torch.backends.mps.is_available(): return torch.device(mps) elif torch.cuda.is_available(): return torch.device(cuda) else: return torch.device(cpu) device get_device() print(f使用设备: {device}) # 创建 MPS 张量 x torch.randn(3, 3, devicedevice) print(f张量设备: {x.device}) # MPS 上的基本操作 y torch.randn(3, 3, devicedevice) z torch.mm(x, y) print(f矩阵乘法结果:\n{z})方案三处理不支持的算子当遇到 MPS 不支持的算子时将张量临时移到 CPU 执行import torch class MPSCompatWrapper: MPS 兼容性包装器。 自动处理 MPS 不支持的算子。 # MPS 上已知不支持的算子列表 UNSUPPORTED_OPS { fft, ifft, stft, istft, linalg.svd, linalg.eig, linalg.eigh, complex, view_as_complex, view_as_real, } staticmethod def run_on_cpu(tensor, op, *args, **kwargs): 将张量移到 CPU 执行操作再移回原设备。 original_device tensor.device tensor_cpu tensor.cpu() # 处理参数中的张量 cpu_args [] for arg in args: if isinstance(arg, torch.Tensor): cpu_args.append(arg.cpu()) else: cpu_args.append(arg) cpu_kwargs {} for k, v in kwargs.items(): if isinstance(v, torch.Tensor): cpu_kwargs[k] v.cpu() else: cpu_kwargs[k] v # 在 CPU 上执行 result op(tensor_cpu, *cpu_args, **cpu_kwargs) # 移回原设备 if isinstance(result, torch.Tensor): return result.to(original_device) elif isinstance(result, (tuple, list)): return [r.to(original_device) if isinstance(r, torch.Tensor) else r for r in result] else: return result staticmethod def safe_fft(x, *args, **kwargs): 安全的 FFT 操作 return MPSCompatWrapper.run_on_cpu(x, torch.fft.fft, *args, **kwargs) staticmethod def safe_svd(x, *args, **kwargs): 安全的 SVD 操作 return MPSCompatWrapper.run_on_cpu(x, torch.linalg.svd, *args, **kwargs) # 使用示例 device torch.device(mps) x torch.randn(4, 4, devicedevice) # FFT 在 MPS 上可能不支持使用兼容包装器 fft_result MPSCompatWrapper.safe_fft(x) print(fFFT 结果设备: {fft_result.device}) # SVD 在 MPS 上可能不支持 U, S, Vh MPSCompatWrapper.safe_svd(x) print(fSVD 结果设备: {U.device})方案四内存管理优化针对 Apple Silicon 统一内存架构的内存管理import torch import gc class MPSMemoryManager: Apple Silicon MPS 内存管理器。 针对统一内存架构优化。 staticmethod def get_memory_info(): 获取内存使用信息 if torch.backends.mps.is_available(): # MPS 使用统一内存通过系统级 API 获取 import psutil mem psutil.virtual_memory() print(f系统总内存: {mem.total / 1024**3:.1f} GB) print(f已使用: {mem.used / 1024**3:.1f} GB) print(f可用: {mem.available / 1024**3:.1f} GB) print(f使用率: {mem.percent:.1f}%) return mem return None staticmethod def clear_mps_cache(): 清理 MPS 缓存 # MPS 没有像 CUDA 那样的 empty_cache # 但可以通过以下方式释放内存 gc.collect() if hasattr(torch.mps, empty_cache): torch.mps.empty_cache() print(MPS 缓存已清理) staticmethod def monitor_memory(func): 内存监控装饰器 def wrapper(*args, **kwargs): import psutil mem_before psutil.virtual_memory().used / 1024**3 result func(*args, **kwargs) gc.collect() if hasattr(torch.mps, empty_cache): torch.mps.empty_cache() mem_after psutil.virtual_memory().used / 1024**3 print(f内存变化: {mem_before:.2f} GB - {mem_after:.2f} GB f(差值: {mem_after - mem_before:.2f} GB)) return result return wrapper # 使用示例 MPSMemoryManager.get_memory_info()完整修复代码以下是一个完整的、在 Apple Silicon M1 上使用 PyTorch MPS 后端进行训练的代码import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader import numpy as np import time import gc import warnings class MPSDeviceManager: Apple Silicon MPS 设备管理器。 处理设备检测、内存管理和兼容性问题。 def __init__(self): self.device self._detect_device() self.fallback_device torch.device(cpu) self.unsupported_ops set()  self._setup() def _detect_device(self) - torch.device: 检测最佳可用设备 if hasattr(torch.backends, mps) and torch.backends.mps.is_available(): print([OK] MPS (Metal Performance Shaders) 可用) return torch.device(mps) elif torch.cuda.is_available(): print([OK] CUDA 可用) return torch.device(cuda) else: print([INFO] 使用 CPU) return torch.device(cpu) def _setup(self): 初始化设备设置 if self.device.type mps: # 设置 MPS 相关的环境变量 # PYTORCH_ENABLE_MPS_FALLBACK1 可以让不支持的算子自动 fallback 到 CPU import os os.environ.setdefault(PYTORCH_ENABLE_MPS_FALLBACK, 1) print([OK] 已启用 MPS fallback 机制) def to_device(self, data, deviceNone): 将数据移到指定设备 target_device device or self.device if isinstance(data, torch.Tensor): return data.to(target_device) elif isinstance(data, (list, tuple)): return type(data)(self.to_device(d, target_device) for d in data) elif isinstance(data, dict): return {k: self.to_device(v, target_device) for k, v in data.items()} else: return data def safe_op(self, op, *args, fallbackTrue, **kwargs): 安全执行操作如果 MPS 不支持则 fallback 到 CPU。 Args: op: 要执行的函数 fallback: 是否在失败时 fallback 到 CPU try: return op(*args, **kwargs) except (NotImplementedError, RuntimeError) as e: if not fallback or self.device.type ! mps: raise op_name getattr(op, __name__, str(op)) if op_name not in self.unsupported_ops: print(f[WARN] {op_name} 在 MPS 上不支持fallback 到 CPU: {e}) self.unsupported_ops.add(op_name) # 将参数移到 CPU cpu_args [] for arg in args: if isinstance(arg, torch.Tensor): cpu_args.append(arg.cpu()) else: cpu_args.append(arg) cpu_kwargs {} for k, v in kwargs.items(): if isinstance(v, torch.Tensor): cpu_kwargs[k] v.cpu() else: cpu_kwargs[k] v # 在 CPU 上执行 result op(*cpu_args, **cpu_kwargs) # 将结果移回 MPS if isinstance(result, torch.Tensor): return result.to(self.device) elif isinstance(result, (tuple, list)): return type(result)( r.to(self.device) if isinstance(r, torch.Tensor) else r for r in result ) return result def clear_cache(self): 清理设备缓存 gc.collect() if self.device.type mps: if hasattr(torch.mps, empty_cache): torch.mps.empty_cache() elif self.device.type cuda: torch.cuda.empty_cache() def get_memory_usage(self) - dict: 获取内存使用情况 import psutil mem psutil.virtual_memory() return { total_gb: mem.total / 1024**3, used_gb: mem.used / 1024**3, available_gb: mem.available / 1024**3, percent: mem.percent, } class MPSTrainingManager: Apple Silicon 上的训练管理器。 处理 MPS 设备上的模型训练、内存管理和兼容性。 def __init__(self, device_manager: MPSDeviceManager): self.dm device_manager self.device device_manager.device self.scaler None # 混合精度 scaler def prepare_model(self, model: nn.Module) - nn.Module: 准备模型移到 MPS 设备 model model.to(self.device) # 在 MPS 上启用混合精度如果支持 if self.device.type mps: # MPS 支持 autocast但行为可能与 CUDA 不同 print([INFO] MPS 设备上使用 autocast 混合精度) return model def prepare_dataloader(self, dataset, batch_size32, shuffleTrue): 创建 DataLoader # 在 MPS 上pin_memory 无效统一内存架构 pin_memory (self.device.type cuda) dataloader DataLoader( dataset, batch_sizebatch_size, shuffleshuffle, pin_memorypin_memory, num_workers0, # MPS 上建议使用 0 避免多进程问题 ) return dataloader def train_epoch(self, model, dataloader, optimizer, criterion, epoch): 训练一个 epoch model.train() total_loss 0 correct 0 total 0 start_time time.time() for batch_idx, (data, target) in enumerate(dataloader): # 移到设备 data self.dm.to_device(data) target self.dm.to_device(target) optimizer.zero_grad() # 使用 autocast 进行混合精度训练 with torch.autocast( enabled(self.device.type in [mps, cuda]), device_typeself.device.type, ): output model(data) loss criterion(output, target) loss.backward() optimizer.step() total_loss loss.item() pred output.argmax(dim1, keepdimTrue) correct pred.eq(target.view_as(pred)).sum().item() total target.size(0) # 每 100 批次打印一次 if batch_idx % 100 0: mem self.dm.get_memory_usage() print(f Epoch {epoch} | Batch {batch_idx}/{len(dataloader)} | fLoss: {loss.item():.4f} | fMem: {mem[used_gb]:.1f}/{mem[total_gb]:.1f} GB) elapsed time.time() - start_time avg_loss total_loss / len(dataloader) accuracy 100. * correct / total return { loss: avg_loss, accuracy: accuracy, time: elapsed, } def validate(self, model, dataloader, criterion): 验证模型 model.eval() total_loss 0 correct 0 total 0 with torch.no_grad(): for data, target in dataloader: data self.dm.to_device(data) target self.dm.to_device(target) with torch.autocast( enabled(self.device.type in [mps, cuda]), device_typeself.device.type, ): output model(data) loss criterion(output, target) total_loss loss.item() pred output.argmax(dim1, keepdimTrue) correct pred.eq(target.view_as(pred)).sum().item() total target.size(0) avg_loss total_loss / len(dataloader) accuracy 100. * correct / total return { loss: avg_loss, accuracy: accuracy, } def train(self, model, train_loader, val_loader, epochs10, lr1e-3): 完整训练流程 model self.prepare_model(model) optimizer optim.Adam(model.parameters(), lrlr) criterion nn.CrossEntropyLoss() print(f\n{*60}) print(f开始训练 | 设备: {self.device} | 轮数: {epochs}) print(f{*60}) best_val_acc 0 history {train: [], val: []} for epoch in range(1, epochs 1): # 训练 train_metrics self.train_epoch(model, train_loader, optimizer, criterion, epoch) # 验证 val_metrics self.validate(model, val_loader, criterion) # 记录历史 history[train].append(train_metrics) history[val].append(val_metrics) print(fEpoch {epoch}/{epochs} | fTrain Loss: {train_metrics[loss]:.4f}, Acc: {train_metrics[accuracy]:.2f}% | fVal Loss: {val_metrics[loss]:.4f}, Acc: {val_metrics[accuracy]:.2f}% | fTime: {train_metrics[time]:.1f}s) # 保存最佳模型 if val_metrics[accuracy] best_val_acc: best_val_acc val_metrics[accuracy] torch.save(model.state_dict(), best_model_mps.pth) print(f - 最佳模型已保存 (Val Acc: {best_val_acc:.2f}%)) # 清理缓存 self.dm.clear_cache() print(f\n训练完成最佳验证准确率: {best_val_acc:.2f}%) return model, history # 示例模型和数据 class ConvNet(nn.Module): 卷积神经网络示例 def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(), nn.AdaptiveAvgPool2d(1), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(128, 256), nn.ReLU(), nn.Dropout(0.5), nn.Linear(256, num_classes), ) def forward(self, x): x self.features(x) x self.classifier(x) return x class SyntheticDataset(Dataset): 合成数据集用于测试 def __init__(self, num_samples1000, image_size32, num_classes10): self.num_samples num_samples self.image_size image_size self.num_classes num_classes # 预生成数据 self.images torch.randn(num_samples, 3, image_size, image_size) self.labels torch.randint(0, num_classes, (num_samples,)) def __len__(self): return self.num_samples def __getitem__(self, idx): return self.images[idx], self.labels[idx] def benchmark_mps_vs_cpu(): MPS vs CPU 性能基准测试 print(\n * 60) print(MPS vs CPU 性能基准测试) print( * 60) device_manager MPSDeviceManager() sizes [512, 1024, 2048, 4096] for size in sizes: # CPU 基准 x_cpu torch.randn(size, size) y_cpu torch.randn(size, size) start time.time() for _ in range(5): z_cpu torch.mm(x_cpu, y_cpu) cpu_time (time.time() - start) / 5 # MPS 基准 if device_manager.device.type mps: x_mps x_cpu.to(mps) y_mps y_cpu.to(mps) # 预热 torch.mm(x_mps, y_mps) start time.time() for _ in range(5): z_mps torch.mm(x_mps, y_mps) mps_time (time.time() - start) / 5 speedup cpu_time / mps_time print(f {size}x{size}: CPU{cpu_time*1000:.1f}ms, fMPS{mps_time*1000:.1f}ms, 加速比{speedup:.1f}x) else: print(f {size}x{size}: CPU{cpu_time*1000:.1f}ms (MPS 不可用)) if __name__ __main__: # 设备检测 device_manager MPSDeviceManager() # 性能基准 benchmark_mps_vs_cpu() # 训练示例 print(\n * 60) print(MPS 训练示例) print( * 60) # 创建数据集 train_dataset SyntheticDataset(num_samples2000, image_size32) val_dataset SyntheticDataset(num_samples400, image_size32) # 创建训练管理器 trainer MPSTrainingManager(device_manager) train_loader trainer.prepare_dataloader(train_dataset, batch_size64) val_loader trainer.prepare_dataloader(val_dataset, batch_size64) # 创建模型 model ConvNet(num_classes10) # 训练 trained_model, history trainer.train( model, train_loader, val_loader, epochs5, lr1e-3 ) print(\n训练完成) # 内存信息 mem device_manager.get_memory_usage() print(f\n最终内存使用: {mem[used_gb]:.1f}/{mem[total_gb]:.1f} GB ({mem[percent]:.1f}%))常见陷阱与注意事项1.PYTORCH_ENABLE_MPS_FALLBACK环境变量设置PYTORCH_ENABLE_MPS_FALLBACK1可以让不支持的算子自动 fallback 到 CPU但这会带来性能损失。建议在开发阶段启用生产阶段针对不支持的算子手动处理。2. BatchNorm 在 MPS 上的行为早期版本的 PyTorch 中BatchNorm 在 MPS 上的实现可能有数值精度问题。如果遇到训练不稳定尝试使用nn.GroupNorm或nn.LayerNorm替代。3. 随机数一致性MPS 的随机数生成器与 CPU/CUDA 不同。如果需要可重复性在 CPU 上生成随机数据再移到 MPS# 确保可重复性 torch.manual_seed(42) data torch.randn(100, 3, 32, 32) # 在 CPU 上生成 data data.to(mps) # 移到 MPS4. 内存碎片问题长时间训练可能导致 MPS 内存碎片。定期调用torch.mps.empty_cache()和gc.collect()可以缓解。5. DataLoader 的num_workers在 macOS 上使用 MPS 时num_workers 0可能导致问题特别是使用 spawn 模式时。建议在 MPS 训练时使用num_workers0。6. 模型保存和加载MPS 上的模型保存和加载需要注意设备映射# 保存自动移到 CPU torch.save(model.state_dict(), model.pth) # 加载时指定设备 state_dict torch.load(model.pth, map_locationcpu) model.load_state_dict(state_dict) model model.to(mps)7. MPS 上的梯度累积MPS 上梯度累积时需要注意torch.autocast的作用域。确保scaler.scale(loss).backward()在正确的上下文中执行。总结在 MacBook Pro M1 上运行 PyTorch 的核心是正确使用 MPS 后端。关键要点如下安装 PyTorch 1.12MPS 支持从 PyTorch 1.12 开始2.0 版本更加稳定。使用torch.device(mps)替代torch.device(cuda)API 基本兼容。处理不支持的算子设置PYTORCH_ENABLE_MPS_FALLBACK1或手动 fallback 到 CPU。统一内存管理Apple Silicon 的统一内存架构意味着无需数据拷贝但需要注意内存总量限制。混合精度训练使用torch.autocast(device_typemps)进行混合精度训练。避免多进程 DataLoader在 MPS 上使用num_workers0避免兼容性问题。定期清理缓存使用torch.mps.empty_cache()和gc.collect()管理内存。Apple Silicon 的 MPS 后端正在快速成熟对于个人开发、原型验证和小规模训练已经完全可用。对于大规模生产训练仍然建议使用 NVIDIA CUDA 集群。