在大语言模型训练领域，TPU 与 JAX 的组合一直以其高效能和灵活的可编程性著称。Nanocode 作为一个开源项目，首次完整展示了如何在 TPU 硬件上用纯 JAX 框架端到端训练一个类 Claude Code 的代码模型。本文将从工程实现角度，深入解析这一技术路径的核心要素，重点探讨 XLA 编译优化与 TPU 内存模型配置的具体参数与实践策略。

技术定位：为什么选择纯 JAX 加 TPU

Nanocode 的核心设计理念是用最少的成本实现可用的代码助手模型。根据项目文档，训练一个 1.3B 参数的模型（nanocode-d24）在 TPU v6e-8 上仅需约 9 小时，费用约 200 美元；而一个 477M 参数的模型（nanocode-d20）更是可以在 1.5 小时内完成训练，成本仅 34 美元。这一成本优势直接得益于 JAX 与 TPU 的深度集成：JAX 通过 XLA 编译器将 Python 代码直接编译为 TPU 指令，规避了传统框架中 Python 运行时与硬件之间的抽象层开销。

与 PyTorch 等框架相比，JAX 的函数式编程范式更契合 TPU 的执行模型。TPU 采用显式内存层次结构和脉动阵列设计，需要编译器在编译期完成数据布局与计算分片的规划。JAX 的 @jax.jit 装饰器与 shard_map 原语允许开发者精确控制这些底层细节，而 XLA 则负责将这些高级抽象转换为高度优化的 TPU 指令序列。

Nanocode 架构与 TPU 适配

Nanocode 的代码结构清晰地分离了模型定义、训练流程与硬件适配三个层面。在模型层面，nanocode/gpt.py 实现了 Transformer 架构，nanocode/generation.py 提供了带 KV 缓存的推理引擎，nanocode/tokenizer.py 则实现了 BPE 分词器。这些组件全部基于 JAX 的核心抽象构建：不使用任何 PyTorch 或 TensorFlow 依赖。

项目支持从 d3（3 层、约 4M 参数）到 d24（24 层、约 1.3B 参数）的多种模型规模。这种多规模支持不仅便于调试，也为不同资源预算的用户提供了灵活选择。值得注意的是，d24 配置在 TPU v6e-8 上可以达到 52.5% 的模型 FLOPs 利用率（MFU），这一数字在同等规模的开源项目中处于领先水平。

在 TPU 适配方面，Nanocode 充分利用了 TPU v6e 的特定优化。项目使用 --attn-impl=eager 参数来控制注意力实现：当在 CPU 或 NVIDIA GPU 上运行时，需要禁用 TPU 特有的 splash attention 内核以保证兼容性；而在 TPU 上则使用专为 TPU 优化的 splash 实现。对于多 TPU 场景，项目支持 v6e-32 等大规模配置，通过多 worker 协调实现分布式训练。

XLA 编译优化：核心机制与可调参数

JAX 对 TPU 的支持本质上是 XLA 编译器对 TPU 后端的深度集成。理解 XLA 的工作机制，是进行有效性能调优的前提。当一个 JAX 函数首次被调用时，JAX 会将其转换为中间表示（IR），随后由 XLA 执行一系列优化 passes，最终生成 TPU 可执行代码。这一过程涉及操作融合、内存布局优化、指令调度等多个维度，而开发者可以通过 XLA Flags 对这一过程进行细粒度控制。

融合策略与操作合并

操作融合（Fusion）是 XLA 在 TPU 上最重要的优化手段之一。传统编译流程中，每个数学运算都会产生独立的内核_launch，导致大量的 kernel launch 开销与中间结果内存访问。通过融合，XLA 将多个相邻操作合并为单一内核，在 TPU 的高速内存（High Bandwidth Memory）中完成整个计算流程，显著降低内存带宽压力。

对于 Nanocode 这类大语言模型训练任务，最关键的融合发生在注意力机制与前馈网络（FFN）模块。典型的 Transformer 层包含大量矩阵乘法、激活函数和 LayerNorm 操作，融合策略的优劣直接决定计算吞吐量。XLA 默认会尝试自动融合，但某些情况下需要显式引导。例如，设置 --xla_enable_fast_math=true 可以启用更激进的数学近似优化，在精度允许范围内合并更多操作。

编译缓存与持久化

大模型的首次编译时间往往很长，因为 XLA 需要完成完整的优化流程。Nanocode 充分利用了 JAX 的持久化编译缓存机制（Persistent Compilation Cache），避免重复编译相同计算图。在调试或实验不同超参数时，这一机制可以节省大量等待时间。开发者可以通过设置 jax.config.update('jax_compilation_cache_dir', '/path/to/cache') 来控制缓存目录位置，并确保在分布式训练中所有 worker 共享同一缓存目录以避免重复编译。

编译边界与分级编译

对于训练任务，另一个重要概念是编译边界（Compilation Boundary）。默认情况下，JAX 会在每个 @jax.jit 装饰的函数边界重新触发编译，这在调试时很方便，但在生产训练中可能导致不必要的重复编译开销。Nanocode 的训练脚本通过精心设计的函数层次结构，将整个前向传播、反向传播与优化器更新封装在少数几个顶层 JIT 函数中，最小化编译边界数量。

TPU 内存模型配置：从理论到实践

TPU 的内存架构与 GPU 有本质区别，这是进行有效配置的前提。TPU v6e 每个核心拥有独立的 High Bandwidth Memory（HBM），通过互连网络形成统一的内存空间。XLA 在编译时需要决定数据的分片（Sharding）策略、内存分配时机与数据移动模式，这些决策直接影响训练效率。

内存池配置

XLA 使用内存池（Memory Pool）来管理设备端内存分配。默认情况下，XLA 会动态管理内存池大小，但在训练大模型时，显式配置内存池参数可以避免运行时分配抖动。关键的环境变量包括 XLA_PYTHON_CLIENT_MEM_FRACTION，用于设置 XLA 可使用的设备内存比例；XLA_PYTHON_CLIENT_PREALLOCATE，控制是否在初始化时预分配全部内存。

对于 Nanocode 的 1.3B 参数模型，建议将内存池比例设置为 0.9 左右，留出约 10% 的内存用于编译器临时缓冲区、碎片整理与异常处理。这一配置在 v6e-8 上经过验证，可以稳定运行 4096 序列长度的训练。

分片与网格布局

JAX 的 shard_map 与 jax.experimental.mesh_utils 提供了在 TPU 上表达数据与计算分片的原语。Nanocode 使用这些工具将模型参数与激活值分布在多个 TPU 核心上。对于多 Slice 配置（如 v6e-32），需要正确设置 jax.distributed.initialize 以协调所有 worker 的初始化，并使用正确的设备网格（Mesh）拓扑。

分片策略的选择涉及计算效率与通信开销的权衡。最常见的策略是数据并行（每个核心处理不同的训练样本）与模型并行（将模型参数分片到不同核心）的组合。对于 1.3B 参数的模型，在 v6e-8（8 核心）上，典型的配置是将批维度分布在 8 个核心上，每个核心持有完整的模型参数副本（数据并行）。对于更大的模型或更多核心，则需要引入张量并行或流水线并行。

流水线与通信重叠

TPU 的互连网络支持高速的集体通信操作（All-Reduce、All-Gather 等），但这些操作仍然会引入同步开销。XLA 提供了流水线（Pipelining）机制，通过重叠计算与通信来隐藏延迟。在训练大型模型时，启用流水线可以在前一个 micro-batch 反向传播的同时进行下一个 micro-batch 的前向传播，显著提升硬件利用率。

相关的 XLA Flags 包括 --xla_enable_async_collective_operations=true 与 --xla_pipeline_budget_factor。这些参数允许 XLA 在编译时插入异步通信操作，并在运行时自动调度计算与通信的重叠执行。

实践参数清单：可落地的配置建议

基于上述分析，以下是在 TPU 上运行 Nanocode 的推荐配置参数。这些参数经过项目实际验证，可以作为起步配置使用。

环境变量配置：在启动训练脚本前，建议设置以下环境变量。XLA_FLAGS 应在导入 JAX 之前设置，因为 XLA 后端在首次使用时初始化，之后修改将无效。具体配置为：export XLA_FLAGS='--xla_enable_fast_math=true --xla_enable_async_collective_operations=true --xla_mlir_emit_hlo_as_text=false'。其中，--xla_enable_fast_math 启用激进的数学优化，--xla_enable_async_collective_operations 开启异步集体通信。

JAX 运行时配置：在 Python 代码或启动脚本中，通过 jax.config 设置运行时选项。建议的配置包括：jax.config.update('jax_enable_mits', True) 启用内存跟踪以便调试；jax.config.update('jax_compilation_cache_dir', '/tmp/jax_cache') 设置编译缓存目录；jax.config.update('jax_pmap_shmap_merge', True) 优化 pmap 与 shard_map 的交互。

TPU 内存配置：对于 v6e-8，建议设置 XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 与 XLA_PYTHON_CLIENT_PREALLOCATE=true。如果遇到 OOM，可以逐步降低至 0.8 或 0.7。

训练超参数：根据目标模型规模选择合适的配置。d24（1.3B 参数）推荐使用批量大小 256、序列长度 4096、学习率 1e-4、权重衰减 0.1。d20（477M 参数）可使用批量大小 512、序列长度 2048、学习率 1.5e-4。

监控与调试：使用 jax.profiler 捕获 Profile 数据，分析 TPU 利用率与内存使用情况。如果发现计算与通信未充分重叠，可以调整 --xla_pipeline_budget_factor 参数，增大流水线深度。

小结

Nanocode 项目展示了纯 JAX 加 TPU 技术栈在大语言模型训练领域的工程可行性。通过深入理解 XLA 编译优化机制与 TPU 内存模型，开发者可以在这一平台上实现高效且经济的大模型训练。上述配置参数为实际落地提供了可操作的起点，建议读者结合自身硬件环境与模型规模进行微调。

参考资料