如何让JAX在M1 Mac上启用GPU加速？配置后未生效求排查

我在M1设备上尝试运行JAX并启用GPU加速，但未成功。已执行以下步骤：

• 卸载并清除所有Anaconda实例
• 重新安装ARM版本的Anaconda
• 创建名为serotonin-gpu的conda环境，所有安装操作均在此环境内完成
• 安装PyTorch nightly版本以获取M1 GPU支持
• 安装适配Mac的MoltenVK
• 克隆JAX仓库的HEAD版本
• 编译并安装jaxlib的wheel包
• 安装jax
• 安装iree的最新快照
• 修改jax/jax/_src/iree.py以传递额外参数

之后在该环境中运行代码，发现GPU未被利用，针对我的疑问给出以下解答：

疑问1：笔记本最初在未配置GPU加速的普通JAX环境创建，是否会有影响？ #

不会直接影响。但要确保运行笔记本时处于serotonin-gpu环境内，且代码中没有硬编码指定CPU设备（比如显式调用jax.devices("cpu")）。可在笔记本开头添加代码验证当前可用设备：

import jax
print(jax.devices())

若输出包含Vulkan相关设备，说明环境配置生效，只是代码未触发GPU执行；若仅显示CPU，说明环境配置仍存在问题。

疑问2：不知道在哪里输入这段与dalle-playground相关的命令 #

# if you want to try the known-good CPU-only mode, then remove these env vars
JAX_PLATFORMS=iree JAX_IREE_BACKEND=vulkan python app.py 8080

这段命令是在终端运行dalle-playground应用时设置环境变量的方式。如果是运行自己的代码，可按以下方式配置：

• 终端运行脚本：在命令前拼接环境变量，示例：JAX_PLATFORMS=iree JAX_IREE_BACKEND=vulkan python your_script.py
• Jupyter Notebook：在笔记本开头添加代码设置环境变量，设置后重启内核生效：

import os
os.environ["JAX_PLATFORMS"] = "iree"
os.environ["JAX_IREE_BACKEND"] = "vulkan"

疑问3：是否有更优质的通用资源指导M1上配置JAX GPU加速？ #

目前JAX对M1 GPU的支持仍处于实验阶段，官方专门指导较少，可参考以下方向：

• 关注JAX官方仓库的Issues和Discussions，筛选M1/Vulkan相关的最新内容，社区用户的实操经验更及时
• 查看IREE官方文档中Mac/Vulkan的配置说明（JAX通过IREE间接调用M1 GPU）
• 搜索社区技术博客中M1用户分享的JAX GPU配置实操教程，这类内容通常包含踩坑细节

疑问4：是否可以联系JAX或Apple团队推动兼容性优化？ #

可以。

• 针对JAX团队：在官方仓库提交Issue，详细描述环境配置（M1型号、系统版本、安装步骤、错误日志等）和GPU未被利用的现象，帮助团队定位问题；也可在Discussions发起话题寻求支持
• 针对Apple团队：通过Apple开发者反馈平台提交兼容性反馈，说明JAX在M1 GPU上的适配需求，但这类反馈响应周期通常较长

内容的提问来源于stack exchange，提问作者seeker_after_truth