tpu-starter
tpu-starter 是一个面向开发者和研究人员的开源指南，旨在帮助用户快速上手 Google Cloud TPU（张量处理单元）。它系统性地解答了初学者常见的疑问，比如如何免费申请 TPU 资源、TPU VM 与 TPU Pod 的区别、如何配置开发环境、远程连接调试等，并提供了详细的实操步骤。项目还涵盖了 JAX 框架的最佳实践，包括随机数管理、数据类型转换、优化器使用等关键技巧，帮助用户高效发挥 TPU 的并行计算能力。特别值得一提的是，tpu-starter 不仅讲解基础用法，还整合了 TRC（TPU Research Cloud）免费计划的申请流程和多机共享存储、资源监控等进阶配置，大幅降低了 TPU 的使用门槛。无论是刚接触 TPU 的学生，还是希望规模化训练模型的研究人员，都能从中获得清晰、实用的指导。
使用场景
某高校研究团队正在复现一篇基于 JAX 的大规模 Transformer 论文，需要在 Google Cloud TPU 上训练模型，但团队成员此前只用过 GPU，对 TPU 架构和配置流程完全陌生。
没有 tpu-starter 时
- 团队花费数天查阅零散的官方文档，仍不清楚如何申请免费 TRC（TPU Research Cloud）资源，多次申请被拒。
- 成功创建 TPU VM 后，不确定是否正确识别到 TPU 设备，反复调试环境依赖却无法运行 JAX 示例代码。
- 配置 VSCode Remote-SSH 和 Jupyter Notebook 时遇到权限和网络问题，开发效率极低。
- 尝试使用 TPU Pod 进行分布式训练时，因不熟悉 NFS 和 SSH 密钥同步机制，集群始终无法协同工作。
- 缺乏 JAX 最佳实践指导，在随机数生成、数据类型转换等细节上频繁出错，导致训练结果不稳定。
使用 tpu-starter 后
- 直接参考项目中“如何申请 TRC”的清晰步骤，一次性成功获得免费 TPU 配额。
- 按照“验证 TPU 是否可用”和“设置开发环境”章节操作，10 分钟内跑通 JAX 基础示例。
- 利用内置的 VSCode Remote-SSH 和 Jupyter 配置指南，快速搭建熟悉的本地开发体验。
- 通过 TPU Pod 部署手册逐步完成子网、NFS 和 podrun 配置，顺利启动多芯片并行训练。
- 遵循 JAX 最佳实践（如随机键管理、Optax 优化器使用），显著提升代码健壮性和训练稳定性。
tpu-starter 将 TPU 从“高门槛专用硬件”转变为“开箱即用的研究工具”，大幅降低 AI 研究者的基础设施学习成本。
运行环境要求
- Linux
未说明
335 GiB（TPU VM实例配置）
快速开始
TPU 入门指南
关于 Google Cloud TPU 的一切你想知道的内容
本项目受 Cloud Run FAQ 启发，后者是关于另一款 Google Cloud 产品的社区维护知识库。
1. 社区
Google 的官方 Discord 服务器已设立 #tpu-research-cloud 频道。
2. TPU 简介
2.1. 为什么选择 TPU？
TL;DR（太长不看）：TPU 之于 GPU，正如 GPU 之于 CPU。
TPU 是专为机器学习设计的硬件。有关性能对比，请参阅 Hugging Face Transformers 中的 Performance Comparison：
此外，Google 的 TRC 计划向研究人员提供免费的 TPU 资源。如果你曾思考过该用什么计算资源来训练模型，你应该尝试 TRC 计划，因为这是我所知的最佳选择。下文将提供更多关于 TRC 计划的信息。
2.2. 如何免费获得 TPU 使用权限？
研究人员可以申请 TRC 计划以获得免费的 TPU 资源。
2.3. 如果 TPU（张量处理单元）这么好，为什么我很少看到别人使用它？
如果你打算使用 PyTorch，TPU 可能并不适合你。PyTorch 对 TPU 的支持很差。在我过去的一个使用 PyTorch 的实验中，一个 batch 在 CPU 上只需 14 秒，但在 TPU 上却需要 4 小时。Twitter 用户 @mauricetpunkt 也认为 PyTorch 在 TPU 上的性能很差。
总之，如果你想用 TPU 进行深度学习，你应该选择 JAX 作为你的深度学习框架。事实上，许多流行的深度学习库都支持 JAX。例如：
此外，JAX 的设计非常简洁，广受好评。例如，JAX 是我最喜欢的开源项目。我曾发推文讨论过 JAX 为何优于 PyTorch。
2.4. 我现在知道 TPU 很棒了。我能摸到一块 TPU 吗？
不幸的是，我们通常无法实际触摸到真实的 TPU。TPU 是通过 Google Cloud 服务访问的。
在一些展览中，TPU 会 被展示出来供人观看，这可能是你能最接近“亲手触摸”TPU 的方式了。
也许只有成为 Google Cloud 基础设施工程师，才能真正感受到 TPU 的“触感”。
2.5. 创建 TPU 实例是什么意思？我实际上得到了什么？
在 Google Cloud Platform 上创建一个 TPU v3-8 实例后，你会获得一台运行 Ubuntu 系统的云服务器，拥有 sudo 权限、96 个 CPU 核心、335 GiB 内存，以及一个包含 8 个核心的 TPU 设备（总计 128 GiB TPU 内存）。
事实上，这与我们使用 GPU 的方式类似。通常，当我们使用 GPU 时，实际上是使用一台连接了 GPU 的 Linux 服务器。同样地，当我们使用 TPU 时，也是在使用一台连接了 TPU 的服务器。
3. TRC 计划介绍
3.1. 如何申请 TRC 计划？
除了 TRC 计划的 主页 外，Shawn 在 google/jax#2108 上撰写了一篇关于 TRC 计划的精彩文章。任何对 TPU 感兴趣的人都应立即阅读。
3.2. 真的免费吗？
在注册 Google Cloud 时，前三个月由于赠送的免费试用额度，TRC 计划完全免费。三个月后，我每月大约花费 13.95 港币（约合 1.78 美元）。这笔费用仅用于 TPU 服务器的网络流量，而 TPU 设备本身由 TRC 计划免费提供。
4. 使用 TPU VM
4.1. 创建 TPU VM
打开 Google Cloud Platform，进入 TPU 管理页面。
点击右上角的控制台按钮以激活 Cloud Shell。
在 Cloud Shell 中，输入以下命令以创建一个 Cloud TPU v3-8 VM：
until gcloud alpha compute tpus tpu-vm create node-1 --project tpu-develop --zone europe-west4-a --accelerator-type v3-8 --version tpu-vm-base ; do : ; done
其中，node-1 是你要创建的 TPU VM 的名称，--project 是你的 Google Cloud 项目名称。
上述命令会不断尝试创建 TPU VM，直到成功为止。
4.2. 向 Google Cloud 添加 SSH 公钥
对于 Google Cloud 的服务器，如果你想通过 SSH 登录，使用 ssh-copy-id 是错误的做法。正确的方法是：
首先，在 Google Cloud 网页的搜索框中输入 “SSH keys”，进入相关页面，点击编辑，然后添加你电脑的 SSH 公钥。
查看你电脑的 SSH 公钥：
cat ~/.ssh/id_rsa.pub
如果你尚未创建 SSH 密钥对，可使用以下命令创建，然后再执行上述命令查看：
ssh-keygen -t rsa -f ~/.ssh/id_rsa -N ""
向 Google Cloud 添加 SSH 公钥时，需特别注意用户名的值。在 SSH 公钥字符串末尾 @ 符号前的部分即为用户名。添加到 Google Cloud 后，它会在当前项目的所有服务器上创建一个该名称的用户。例如，对于字符串 ayaka@instance-1，Google Cloud 会在服务器上创建名为 ayaka 的用户。如果你希望 Google Cloud 创建不同的用户名，可以手动修改该字符串。例如，将字符串改为 nixie@instance-1，Google Cloud 就会创建名为 nixie 的用户。而且，这种修改不会影响 SSH 密钥的功能。
4.3. SSH 登录 TPU VM
在你本地电脑上创建或编辑 ~/.ssh/config 文件：
nano ~/.ssh/config
添加以下内容：
Host tpuv3-8-1
User nixie
Hostname 34.141.220.156
其中，tpuv3-8-1 是任意名称，User 是上一步在 Google Cloud 中创建的用户名，Hostname 是 TPU VM 的 IP 地址。
然后，在你本地电脑上使用以下命令 SSH 登录 TPU VM：
ssh tpuv3-8-1
其中 tpuv3-8-1 是在 ~/.ssh/config 中设置的名称。
4.4. 验证 TPU VM 是否拥有 TPU
ls /dev/accel*
如果出现如下输出：
/dev/accel0 /dev/accel1 /dev/accel2 /dev/accel3
说明该 TPU VM 确实配备了 TPU。
4.5. 在 TPU VM 中设置开发环境
更新软件包：
sudo apt-get update -y -qq
sudo apt-get upgrade -y -qq
sudo apt-get install -y -qq golang neofetch zsh byobu
安装最新的 Python 3.12：
sudo apt-get install -y -qq software-properties-common
sudo add-apt-repository -y ppa:deadsnakes/ppa
sudo apt-get install -y -qq python3.12-full python3.12-dev
安装 Oh My Zsh：
sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended
sudo chsh $USER -s /usr/bin/zsh
创建虚拟环境（venv）：
python3.12 -m venv ~/venv
激活 venv：
. ~/venv/bin/activate
在 venv 中安装 JAX：
pip install -U pip
pip install -U wheel
pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
4.6. 验证 JAX 是否正常工作
激活 venv 后，使用以下命令验证 JAX 是否正常工作：
python -c 'import jax; print(jax.devices())'
如果输出中包含 TpuDevice，说明 JAX 已按预期正常工作。
4.7. 使用 Byobu 确保程序持续运行
许多教程使用在命令末尾附加 & 的方法，使其在后台运行，这样即使退出 SSH 后程序仍会继续执行。然而，这是一种基础方法。正确的方式是使用像 Byobu 这样的终端窗口管理器（terminal window manager）。
要运行 Byobu，只需使用 byobu 命令。然后在打开的窗口中执行命令即可。若要关闭窗口，可以直接强制关闭你本地计算机上的当前窗口。Byobu 会在服务器上继续运行。下次连接到服务器时，再次使用 byobu 命令即可恢复之前的窗口。
Byobu 拥有许多高级功能。你可以通过观看官方视频 边听莫扎特边学习 Byobu 来了解这些功能。
4.8. 配置 VSCode Remote-SSH
打开 VSCode，在左侧访问扩展（Extensions）面板，搜索并安装 Remote - SSH。
按下 F1 打开命令面板，输入 ssh，点击 “Remote-SSH: Connect to Host...”，然后点击你在 ~/.ssh/config 中设置的服务器名称（例如 tpuv3-8-1）。一旦 VSCode 在服务器上完成设置，你就可以直接在服务器上使用 VSCode 进行开发。
在你的本地计算机上，可以使用以下命令快速打开服务器上的某个目录：
code --remote ssh-remote+tpuv3-8-1 /home/ayaka/tpu-starter
该命令将使用 VSCode 打开 tpuv3-8-1 上的 /home/ayaka/tpu-starter 目录。
4.9. 在 TPU VM 上使用 Jupyter Notebook
在配置好 VSCode 的 Remote-SSH 后，你可以在 VSCode 内使用 Jupyter Notebook。效果如下所示：
这里需要注意两点：第一，在 Jupyter Notebook 界面的右上角，应选择来自 venv 的 Kernel，即我们在前面步骤中创建的 ~/venv/bin/python。第二，首次运行时，系统会提示你为 VSCode 安装 Jupyter 扩展，并在 venv 中安装 ipykernel。你需要确认这些操作。
5. 使用 TPU Pod
5.1. 创建子网
要创建 TPU Pod，首先需要创建一个新的 VPC 网络（Virtual Private Cloud network），然后在该网络对应的区域（例如 europe-west4-a）中创建一个子网（subnet）。
TODO: Purpose?
5.2. 禁用 Cloud Logging
TODO: Reason? Steps?
5.3. 创建 TPU Pod
使用之前创建 TPU VM 时描述的方法打开 Cloud Shell，并使用以下命令创建一个 TPU v3-32 Pod：
until gcloud alpha compute tpus tpu-vm create node-1 --project tpu-advanced-research --zone europe-west4-a --accelerator-type v3-32 --version v2-alpha-pod --network advanced --subnetwork advanced-subnet-for-europe-west4 ; do : ; done
其中 node-1 是你希望为 TPU VM 设置的名称，--project 是你的 Google Cloud 项目名称，--network 和 --subnetwork 是上一步中创建的网络和子网的名称。
5.4. SSH 登录 TPU Pod
由于 TPU Pod 由多个主机组成，我们需要选择其中一个主机作为 Host 0，然后通过 SSH 登录 Host 0 来执行命令。由于在 Google Cloud 网页上添加的 SSH 公钥会被自动分发到所有主机，因此每个主机都可以通过该 SSH 密钥直接连接，我们可以任意指定一个主机作为 Host 0。登录 Host 0 的方法与前述 TPU VM 相同。
5.5. 修改 Host 0 上的 SSH 配置文件
SSH 登录 Host 0 后，需要进行如下配置：
nano ~/.ssh/config
添加以下内容：
Host 172.21.12.* 127.0.0.1
StrictHostKeyChecking no
UserKnownHostsFile /dev/null
LogLevel ERROR
这里的 172.21.12.* 取决于上一步创建的子网 IP 地址范围。我们使用 172.21.12.* 是因为在创建子网时指定的 IP 地址范围为 172.21.12.0/24。
之所以这样做，是因为 SSH 的 known_hosts 文件用于防止中间人攻击（man-in-the-middle attacks）。而此处我们使用的是内网环境，无需防范此类攻击，也就不需要该文件，因此将其指向 /dev/null。此外，保留 known_hosts 文件会导致首次连接时需要手动确认服务器指纹，这在内网环境中是不必要的，也不利于自动化操作。
然后，运行以下命令修改该配置文件的权限。如果不修改权限，配置文件将不会生效：
chmod 600 ~/.ssh/config
5.6. 将 Host 0 的 SSH 公钥添加到所有主机
在 Host 0 上生成密钥对：
ssh-keygen -t rsa -f ~/.ssh/id_rsa -N ""
查看生成的 SSH 公钥：
cat ~/.ssh/id_rsa.pub
将此公钥添加到 Google Cloud 的 SSH 密钥中。该密钥将自动分发到所有主机。
5.7. 配置 podrun 命令
podrun 是一个正在开发中的工具。当在 Host 0 上执行时，它可以通过 SSH 在所有主机上运行命令。
下载 podrun：
wget https://raw.githubusercontent.com/ayaka14732/llama-2-jax/18e9625f7316271e4c0ad9dea233cfe23c400c9b/podrun
chmod +x podrun
使用以下命令编辑 ~/podips.txt：
nano ~/podips.txt
将其他主机的内网 IP 地址逐行保存到 ~/podips.txt 中。例如：
172.21.12.86
172.21.12.87
172.21.12.83
一个 TPU v3-32 包含 4 个主机。除去 Host 0，还有 3 个主机。因此，TPU v3-32 的 ~/podips.txt 应包含 3 个 IP 地址。
使用系统自带的 pip3 安装 Fabric：
pip3 install fabric
使用 podrun 让所有主机像小猫一样“喵喵”叫：
./podrun -iw -- echo meow
5.8. 配置 NFS
安装 NFS 服务器和客户端：
./podrun -i -- sudo apt-get update -y -qq
./podrun -i -- sudo apt-get upgrade -y -qq
./podrun -- sudo apt-get install -y -qq nfs-common
sudo apt-get install -y -qq nfs-kernel-server
sudo mkdir -p /nfs_share
sudo chown -R nobody:nogroup /nfs_share
sudo chmod 777 /nfs_share
修改 /etc/exports：
sudo nano /etc/exports
添加：
/nfs_share 172.21.12.0/24(rw,sync,no_subtree_check)
执行：
sudo exportfs -a
sudo systemctl restart nfs-kernel-server
./podrun -- sudo mkdir -p /nfs_share
./podrun -- sudo mount 172.21.12.2:/nfs_share /nfs_share
./podrun -i -- ln -sf /nfs_share ~/nfs_share
touch ~/nfs_share/meow
./podrun -i -- ls -la ~/nfs_share/meow
请将 172.21.12.2 替换为 Host 0 实际的内网 IP 地址。
5.9. 在 TPU Pod 中设置开发环境
将以下内容保存为 ~/nfs_share/setup.sh：
#!/bin/bash
export DEBIAN_FRONTEND=noninteractive
sudo apt-get update -y -qq
sudo apt-get upgrade -y -qq
sudo apt-get install -y -qq golang neofetch zsh byobu
sudo apt-get install -y -qq software-properties-common
sudo add-apt-repository -y ppa:deadsnakes/ppa
sudo apt-get install -y -qq python3.12-full python3.12-dev
sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended
sudo chsh $USER -s /usr/bin/zsh
python3.12 -m venv ~/venv
. ~/venv/bin/activate
pip install -U pip
pip install -U wheel
pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
然后执行：
chmod +x ~/nfs_share/setup.sh
./podrun -i ~/nfs_share/setup.sh
5.10. 验证 JAX 是否正常工作
./podrun -ic -- ~/venv/bin/python -c 'import jax; jax.distributed.initialize(); jax.process_index() == 0 and print(jax.devices())'
如果输出中包含 TpuDevice，则表示 JAX 已按预期工作。
6. TPU 最佳实践
6.1. 优先使用 Google Cloud Platform 而非 Google Colab
Google Colab 仅提供 TPU v2-8 设备，而在 Google Cloud Platform 上你可以选择 TPU v2-8 和 TPU v3-8。
此外，在 Google Colab 中你只能通过 Jupyter Notebook 界面使用 TPU。即使你通过 SSH 登录到 Colab 服务器，那也是一个 Docker 镜像，并且你没有 root 权限。而在 Google Cloud Platform 上，你可以完全访问 TPU VM。
如果你确实想在 Google Colab 上使用 TPU，则需要运行以下脚本来设置 TPU：
import jax
from jax.tools.colab_tpu import setup_tpu
setup_tpu()
devices = jax.devices()
print(devices) # 应该打印出 TpuDevice
6.2. 优先使用 TPU VM 而非 TPU node
当你创建 TPU 实例时，需要在 TPU VM 和 TPU node 之间进行选择。始终优先选择 TPU VM，因为这是新架构，其中 TPU 设备直接连接到宿主机 VM，这将使 TPU 设备的设置更加简单。
7. JAX 最佳实践
7.1. 导入约定
你可能会看到两种不同的导入约定。一种是将 jax.numpy 导入为 np，并将原始 NumPy 导入为 onp；另一种是将 jax.numpy 导入为 jnp，并将原始 NumPy 保留为 np。
2019 年 1 月 16 日，Colin Raffel 在一篇博客文章中提到，当时的约定是将原始 NumPy 导入为 onp。
2020 年 11 月 5 日，Niru Maheswaranathan 在一条推文中表示，他认为当时的约定是将 jax.numpy 导入为 jnp，并将原始 NumPy 保留为 np。
我们可以得出结论：新的约定是将 jax.numpy 导入为 jnp。
7.2. 在 JAX 中管理随机密钥（random keys）
常规做法如下：
key, *subkey = rand.split(key, num=4)
print(subkey[0])
print(subkey[1])
print(subkey[2])
7.3. NumPy 数组与 JAX 数组之间的转换
使用 np.asarray 和 onp.asarray。
import jax.numpy as np
import numpy as onp
a = np.array([1, 2, 3]) # JAX 数组
b = onp.asarray(a) # 转换为 NumPy 数组
c = onp.array([1, 2, 3]) # NumPy 数组
d = np.asarray(c) # 转换为 JAX 数组
7.4. PyTorch 张量与 JAX 数组之间的转换
将 PyTorch 张量转换为 JAX 数组：
import jax.numpy as np
import torch
a = torch.rand(2, 2) # PyTorch 张量
b = np.asarray(a.numpy()) # JAX 数组
将 JAX 数组转换为 PyTorch 张量：
import jax.numpy as np
import numpy as onp
import torch
a = np.zeros((2, 2)) # JAX 数组
b = torch.from_numpy(onp.asarray(a)) # PyTorch 张量
这将产生一个警告：
UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:178.)
如果你需要可写的张量，可以使用 onp.array 而不是 onp.asarray，以复制原始数组。
7.5. 获取嵌套字典中所有参数的形状
jax.tree_map(lambda x: x.shape, params)
7.6. 在 CPU 上生成随机数的正确方式
使用 jax.default_device() 上下文管理器：
import jax
import jax.random as rand
device_cpu = jax.devices('cpu')[0]
with jax.default_device(device_cpu):
key = rand.PRNGKey(42)
a = rand.poisson(key, 3, shape=(1000,))
print(a.device()) # TFRT_CPU_0
参见 https://github.com/google/jax/discussions/9691#discussioncomment-3650311。
7.7. 使用 Optax 中的优化器
7.8. 使用 Optax 中的交叉熵损失实现
optax.softmax_cross_entropy_with_integer_labels
8. 如何...
8.1. 在多个 TPU VM 实例之间共享文件
同一区域（zone）中的 TPU VM 实例通过内部 IP 相互连接，因此你可以使用 NFS 创建共享文件系统。
8.2. 监控 TPU 使用情况
8.3. 在 TPU VM 上启动服务器
示例：TensorBoard
尽管每个 TPU VM 都分配了一个公网 IP，但在大多数情况下你不应将服务器直接暴露在互联网上，因为这不安全。
通过 SSH 进行端口转发：
ssh -C -N -L 127.0.0.1:6006:127.0.0.1:6006 tpu1
8.4. 在不同的 TPU 核心上运行独立进程
https://gist.github.com/skye/f82ba45d2445bb19d53545538754f9a3
9. 常见陷阱
9.1. TPU VM 会偶尔重启
截至 2022 年 10 月 24 日，如果发生维护事件，TPU VM 会偶尔重启。
以下情况会发生：
- 所有正在运行的进程将被终止
- 外部 IP 地址将发生变化
我们可以定期保存模型参数、优化器状态以及其他有用的数据，以便在终止后轻松恢复模型训练。
我们应该使用 gcloud 命令而不是直接通过 SSH 连接。如果我们必须使用 SSH（例如，如果我们想使用 VSCode，SSH 是唯一的选择），则需要手动更改目标 IP 地址。
9.2. 一个 TPU 核心（TPU core）同一时间只能被一个进程使用
另见：§10.5。
与 GPU 不同，如果你同时在 TPU 上运行两个进程，将会收到错误：
I0000 00:00:1648534265.148743 625905 tpu_initializer_helper.cc:94] libtpu.so already in use by another process. Run "$ sudo lsof -w /dev/accel0" to figure out which process is using the TPU. Not attempting to load libtpu.so in this process.
9.3. TCMalloc 会导致多个程序异常
TCMalloc 是 Google 定制的内存分配库。在 TPU VM 上，默认通过设置 LD_PRELOAD 来使用 TCMalloc：
$ echo LD_PRELOAD
/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4
然而，以这种方式使用 TCMalloc 可能会导致多个程序（如 gsutil）出错：
$ gsutil --help
/snap/google-cloud-sdk/232/platform/bundledpythonunix/bin/python3: /snap/google-cloud-sdk/232/platform/bundledpythonunix/bin/../../../lib/x86_64-linux-gnu/libm.so.6: version `GLIBC_2.29' not found (required by /usr/lib/x86_64-linux-gnu/libtcmalloc.so.4)
TCMalloc 的主页 也指出，使用 LD_PRELOAD 的方式较为棘手，不推荐这样使用。
如果你遇到与 TCMalloc 相关的问题，可以在当前 shell 中通过以下命令禁用它：
unset LD_PRELOAD
9.4. libtpu.so 已被另一个进程占用
if ! pgrep -a -u $USER python ; then
killall -q -w -s SIGKILL ~/.venv311/bin/python
fi
rm -rf /tmp/libtpu_lockfile /tmp/tpu_logs
另见 https://github.com/google/jax/issues/9220#issuecomment-1015940320。
9.5. JAX 不支持 multiprocessing 的 fork 启动策略（start method）
请使用 spawn 或 forkserver 启动策略。
参见 https://github.com/google/jax/issues/1805#issuecomment-561244991。
版本历史
v2.02023/10/08v1.02022/10/08常见问题
相似工具推荐
stable-diffusion-webui
stable-diffusion-webui 是一个基于 Gradio 构建的网页版操作界面，旨在让用户能够轻松地在本地运行和使用强大的 Stable Diffusion 图像生成模型。它解决了原始模型依赖命令行、操作门槛高且功能分散的痛点，将复杂的 AI 绘图流程整合进一个直观易用的图形化平台。 无论是希望快速上手的普通创作者、需要精细控制画面细节的设计师，还是想要深入探索模型潜力的开发者与研究人员，都能从中获益。其核心亮点在于极高的功能丰富度：不仅支持文生图、图生图、局部重绘（Inpainting）和外绘（Outpainting）等基础模式，还独创了注意力机制调整、提示词矩阵、负向提示词以及“高清修复”等高级功能。此外，它内置了 GFPGAN 和 CodeFormer 等人脸修复工具，支持多种神经网络放大算法，并允许用户通过插件系统无限扩展能力。即使是显存有限的设备，stable-diffusion-webui 也提供了相应的优化选项，让高质量的 AI 艺术创作变得触手可及。
everything-claude-code
everything-claude-code 是一套专为 AI 编程助手（如 Claude Code、Codex、Cursor 等）打造的高性能优化系统。它不仅仅是一组配置文件，而是一个经过长期实战打磨的完整框架，旨在解决 AI 代理在实际开发中面临的效率低下、记忆丢失、安全隐患及缺乏持续学习能力等核心痛点。 通过引入技能模块化、直觉增强、记忆持久化机制以及内置的安全扫描功能，everything-claude-code 能显著提升 AI 在复杂任务中的表现，帮助开发者构建更稳定、更智能的生产级 AI 代理。其独特的“研究优先”开发理念和针对 Token 消耗的优化策略，使得模型响应更快、成本更低，同时有效防御潜在的攻击向量。 这套工具特别适合软件开发者、AI 研究人员以及希望深度定制 AI 工作流的技术团队使用。无论您是在构建大型代码库，还是需要 AI 协助进行安全审计与自动化测试，everything-claude-code 都能提供强大的底层支持。作为一个曾荣获 Anthropic 黑客大奖的开源项目，它融合了多语言支持与丰富的实战钩子（hooks），让 AI 真正成长为懂上
ComfyUI
ComfyUI 是一款功能强大且高度模块化的视觉 AI 引擎，专为设计和执行复杂的 Stable Diffusion 图像生成流程而打造。它摒弃了传统的代码编写模式，采用直观的节点式流程图界面，让用户通过连接不同的功能模块即可构建个性化的生成管线。 这一设计巧妙解决了高级 AI 绘图工作流配置复杂、灵活性不足的痛点。用户无需具备编程背景，也能自由组合模型、调整参数并实时预览效果，轻松实现从基础文生图到多步骤高清修复等各类复杂任务。ComfyUI 拥有极佳的兼容性，不仅支持 Windows、macOS 和 Linux 全平台，还广泛适配 NVIDIA、AMD、Intel 及苹果 Silicon 等多种硬件架构，并率先支持 SDXL、Flux、SD3 等前沿模型。 无论是希望深入探索算法潜力的研究人员和开发者，还是追求极致创作自由度的设计师与资深 AI 绘画爱好者，ComfyUI 都能提供强大的支持。其独特的模块化架构允许社区不断扩展新功能，使其成为当前最灵活、生态最丰富的开源扩散模型工具之一，帮助用户将创意高效转化为现实。
NextChat
NextChat 是一款轻量且极速的 AI 助手，旨在为用户提供流畅、跨平台的大模型交互体验。它完美解决了用户在多设备间切换时难以保持对话连续性，以及面对众多 AI 模型不知如何统一管理的痛点。无论是日常办公、学习辅助还是创意激发，NextChat 都能让用户随时随地通过网页、iOS、Android、Windows、MacOS 或 Linux 端无缝接入智能服务。 这款工具非常适合普通用户、学生、职场人士以及需要私有化部署的企业团队使用。对于开发者而言，它也提供了便捷的自托管方案，支持一键部署到 Vercel 或 Zeabur 等平台。 NextChat 的核心亮点在于其广泛的模型兼容性，原生支持 Claude、DeepSeek、GPT-4 及 Gemini Pro 等主流大模型，让用户在一个界面即可自由切换不同 AI 能力。此外，它还率先支持 MCP（Model Context Protocol）协议，增强了上下文处理能力。针对企业用户，NextChat 提供专业版解决方案，具备品牌定制、细粒度权限控制、内部知识库整合及安全审计等功能，满足公司对数据隐私和个性化管理的高标准要求。
ML-For-Beginners
ML-For-Beginners 是由微软推出的一套系统化机器学习入门课程，旨在帮助零基础用户轻松掌握经典机器学习知识。这套课程将学习路径规划为 12 周，包含 26 节精炼课程和 52 道配套测验，内容涵盖从基础概念到实际应用的完整流程，有效解决了初学者面对庞大知识体系时无从下手、缺乏结构化指导的痛点。 无论是希望转型的开发者、需要补充算法背景的研究人员，还是对人工智能充满好奇的普通爱好者，都能从中受益。课程不仅提供了清晰的理论讲解，还强调动手实践，让用户在循序渐进中建立扎实的技能基础。其独特的亮点在于强大的多语言支持，通过自动化机制提供了包括简体中文在内的 50 多种语言版本，极大地降低了全球不同背景用户的学习门槛。此外，项目采用开源协作模式，社区活跃且内容持续更新，确保学习者能获取前沿且准确的技术资讯。如果你正寻找一条清晰、友好且专业的机器学习入门之路，ML-For-Beginners 将是理想的起点。
ragflow
RAGFlow 是一款领先的开源检索增强生成（RAG）引擎，旨在为大语言模型构建更精准、可靠的上下文层。它巧妙地将前沿的 RAG 技术与智能体（Agent）能力相结合，不仅支持从各类文档中高效提取知识，还能让模型基于这些知识进行逻辑推理和任务执行。 在大模型应用中，幻觉问题和知识滞后是常见痛点。RAGFlow 通过深度解析复杂文档结构（如表格、图表及混合排版），显著提升了信息检索的准确度，从而有效减少模型“胡编乱造”的现象，确保回答既有据可依又具备时效性。其内置的智能体机制更进一步，使系统不仅能回答问题，还能自主规划步骤解决复杂问题。 这款工具特别适合开发者、企业技术团队以及 AI 研究人员使用。无论是希望快速搭建私有知识库问答系统，还是致力于探索大模型在垂直领域落地的创新者，都能从中受益。RAGFlow 提供了可视化的工作流编排界面和灵活的 API 接口，既降低了非算法背景用户的上手门槛，也满足了专业开发者对系统深度定制的需求。作为基于 Apache 2.0 协议开源的项目，它正成为连接通用大模型与行业专有知识之间的重要桥梁。