为 Trainium 编译模型：TorchNeuron 的图优化实践

TorchNeuron（PyTorch 到 AWS Neuron 的图编译后端） 把 PyTorch 模型编译到 Trainium（AWS 自研 AI 训练芯片），图优化收益九成来自三件事：钉死 shape、让融合在编译期发生、把并行切分显式写进图。三步做对，单卡吞吐能翻倍；做错，编译时间会反过来吃掉迭代效率。

用大白话说：TorchNeuron（PyTorch 到 AWS Neuron 的图编译后端） 是 PyTorch 到 Trainium（AWS 自研 AI 训练芯片） 的翻译层。你丢进去一张 PyTorch 计算图，它先转成 HLO（XLA 的中间表示，High Level Operations），再由 neuronx-cc 下沉成 NeuronCore 能直接执行的 NEFF（Neuron Executable File Format）文件。算子融合、内存复用、布局转换这些重活基本都在 neuronx-cc 里做，PyTorch 侧能做的是把编译器猜不准的东西提前告诉它——shape、dtype、并行切分方式，全是猜不准的东西。

AWS 在 2026 年 10 月 26 日的周报里把 TorchNeuron（PyTorch 到 AWS Neuron 的图编译后端） for Trainium（AWS 自研 AI 训练芯片） 单独拎出来讲来源：AWS Weekly Roundup，这条编译链路已经过了试验阶段。下面是我在 trn1.32xlarge 上把一个 transformer block 编到 NeuronCore 的完整过程，带耗时和日志。

TorchNeuron 在编译链里到底扮演什么角色？#

先把职责边界划清楚，后面调参才不会乱试。

对比表：TorchNeuron 与 neuronx-cc 的分工

关键在第二行。你写的 `--optlevel=2` 不会改变模型语义，它只是允许编译器做更激进的融合和重排。这部分收益大、风险也大——我见过 `--model-type=transformer` 误用在纯 CNN 上，融合退化，NEFF 算子数不降反升。

编号步骤：环境准备

1. 建独立虚拟环境，别和系统 Python 混。Neuron 的 wheel 对 torch 版本很敏感，版本错一位就是编译期 segfault。
   python -m venv /opt/neuron-venv && source /opt/neuron-venv/bin/activate
2. 先装 CPU 版 torch（Trainium 编译阶段不需要 CUDA 依赖）。
   pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cpu
3. 装 Neuron 运行时与编译器。版本号按 2026 年 10 月的 Neuron SDK 2.21 系列写，装之前用 pip index versions torch-neuronx 核对一次。
   pip install torch-neuronx==2.7.1.2.21.0 neuronx-cc==2.21.0.0.0 'neuronx-distributed>=0.13'
4. 确认编译机能出网，后面流水线要做回调探测。
   curl -sS -o /dev/null -w "%{http_code}\n" "https://www.api.plus/doc/get"

第 4 步看着多余，实际踩过坑：编译机和训练集群的出口策略经常不是一套，编译跑完 200 秒才发现回调打不出去，只能重来。

为什么必须先解决 shape 问题？#

因为 neuronx-cc 是静态编译器。它需要在编译期就知道每个张量的完整形状，才能确定 HBM 上的内存布局和 DMA 传输块大小。你给它一个动态维度，它只能按最坏情况预留，融合策略也会变得保守。

更直接的代价是重编译。torch_neuronx.trace 产出的产物绑死 sample input 的 shape。我做过一次对照：batch 从 8 换到 16，编译好的 .pt 直接报 shape mismatch，只能重新编一次，215 秒又没了。那一晚调参换了 6 个 batch，编译时间累计到 3 小时 40 分，NeuronCore 全程闲着。

对比表：三种 shape 处理策略的取舍

我的选择是分桶。三个桶对应 512、1024、2048 三种上下文长度，超出就截断，不足就 pad 到最近的桶，attention mask 负责屏蔽。多花两次编译时间，换运行期的确定性，值。

还有一点容易忽略：dtype 也在 trace 时固化。sample input 用 bf16，编出来的图就是 bf16 通路；线上喂 fp32 会静默转换，多一次 cast 开销。这一点在 GET 示例接口 这类轻量服务的回调里看不出来，但在吞吐敏感的批量推理上会体现在 P99（99 分位延迟） 上。

如何用 torch_neuronx.trace 把融合和并行写进图里？#

核心就一个函数调用，参数怎么填决定编译结果。先看代码。

可执行脚本：编译一个 FFN block

示例说明：以下代码仅用于逻辑验证，生产环境需补充异常处理、超时与重试、日志与监控告警及降级策略。

import time
import torch
import torch_neuronx

class FFNBlock(torch.nn.Module):
    def __init__(self, d=1024):
        super().__init__()
        self.fc1 = torch.nn.Linear(d, 4 * d, bias=False)
        self.act = torch.nn.GELU()
        self.fc2 = torch.nn.Linear(4 * d, d, bias=False)
        self.norm = torch.nn.LayerNorm(d)

    def forward(self, x):
        return self.norm(x + self.fc2(self.act(self.fc1(x))))

model = FFNBlock().eval().to(torch.bfloat16)
sample = torch.randn(8, 512, 1024, dtype=torch.bfloat16)

t0 = time.time()
neuron_model = torch_neuronx.trace(
    model,
    sample,
    compiler_args=[
        "--optlevel=2",
        "--model-type=transformer",
        "--enable-mixed-precision-accumulation",
    ],
    compiler_workdir="/var/log/neuron/workdir",
    inline_weights_to_neff=True,
)
print("compile wall time: %.1fs" % (time.time() - t0))
torch.jit.save(neuron_model, "/opt/neuron/artifacts/ffn_b8_s512.pt")

参数含义逐条说：

• model：必须 eval()。训练态的 dropout 留在图里会阻断融合，编译器无法把相邻算子合成一个 kernel。
• sample：决定静态 shape 和 dtype，同时是 check_trace 的对照输入。batch=8、seq=512、hidden=1024 是 trn1 单核上比较舒服的配置，HBM 占用约 2.1GB。
• --optlevel=2：融合等级。默认 1 只做基础融合，2 会放开算子重排和跨层融合。
• --model-type=transformer：告诉编译器图结构特征，它会启用针对 attention / FFN 的专用融合模板。这个参数对结构判断错了会反效果。
• --enable-mixed-precision-accumulation：矩阵乘用 bf16 累加器，省 HBM 带宽，数值误差在 LayerNorm 之前可控。
• compiler_workdir：中间产物落地目录。编译失败时唯一的排查入口，别省。
• inline_weights_to_neff=True：权重内联进 NEFF，加载快，代价是产物文件变大（本例从 38MB 涨到 156MB）。

编译完成后翻日志确认融合效果。

可执行命令：从编译日志提取算子数与耗时

grep -hE "NEFF op count|Compile time|Compiler status" /var/log/neuron/workdir/*/log/neuroncc*.log | tail -5

我这台的输出片段：

2026-10-29T09:12:41Z INFO 1021 NEFF op count: 41
2026-10-29T09:12:41Z INFO 1021 Compile time: 214.6 sec
2026-10-29T09:12:41Z INFO 1021 Compiler status: SUCCESS

不加 --optlevel 时 NEFF op count 是 287，加上之后 41。下降的不是计算量，是 kernel 启动次数和中间结果落 HBM 的次数——Linear 加 GELU 加残差被合成一个大 kernel，张量不再来回搬运。

对比表：编译耗时分解（同一次编译）

注意这段编译跑在 host CPU 上，NeuronCore 全程空闲。所以流水线里编译任务和推理任务要分开调度，别让编译卡住推理队列。

编译耗时与吞吐怎么量化？#

别凭感觉说"变快了"。我把编译当作一次可以打点的构建任务，至少记录四类数据：编译墙钟时间、NEFF 算子数、产物大小、HBM 峰值。这四个数够判断一次编译是不是正常。

流水线里我做一次出网与接口连通性探测，避免编译完发现回调打不出去。用的是 API.PLUS 的 GET 示例接口，参数简单、返回快，适合当探针。

可执行命令：编译前连通性探测

curl -sS --max-time 5 -o /dev/null \
  -w "http=%{http_code} dns=%{time_namelookup}s total=%{time_total}s\n" \
  "https://www.api.plus/api/get?apikey=YOUR_KEY&ip=8.8.8.8"

参数解释：apikey 是你在 API.PLUS 控制台注册后拿到的密钥，直接拼在查询串里；ip 是待查询的目标地址，这里用 8.8.8.8 做固定值便于对比历史数据；-w 输出 HTTP 状态码和各阶段耗时，-o /dev/null 丢弃响应体只看时间。

本机输出：

http=200 dns=0.021s total=0.183s

如果 total 超过 1 秒，说明编译机到公网的链路有问题，流水线直接中止，别浪费 215 秒编译时间。这个探针我放在编译任务的第一步，跑了两周拦下三次网络策略变更导致的故障。

需要看完整返回结构时去掉 -o /dev/null，返回是标准 JSON，字段含义见接口文档。类似的 POST 示例接口 我用在上报编译指标这一步——把耗时、NEFF 算子数、产物大小打成一个 JSON body，POST 给内部的看板服务。

编译完成后还有一项验证容易被跳过：用随机输入跑一遍编译产物，和 host 侧 PyTorch 输出比对。bf16 在 1024 维 hidden 上相对误差通常在 1e-2 量级，超过 5e-2 就要怀疑 --enable-mixed-precision-accumulation 用得过激。这一步失败过一次，原因是 LayerNorm 的 epsilon 默认值和训练侧不一致，查了两小时。

监控与告警清单：编译上线之后盯什么#

指标清单：编译与运行期

• 编译失败率：阈值 5%，超过就告警。动作是保留 compiler_workdir 并按失败阶段分类，别急着重试。
• 单次编译耗时 P95：阈值 400 秒。超了先查 shape 是否漂移、optlevel 是否被人改动。
• 每小时重编译次数：阈值 2 次。超了几乎可以确定是动态 shape 混进来了，查 cache key 的组成。
• 单模型 NEFF 数量：阈值 3 个。超过说明 shape 分桶没收敛，需要合并桶。
• HBM 峰值：阈值单核容量的 60%。接近上限时 next-token 生成会因内存碎片抖动。
• 推理 P99（99 分位延迟） 延迟：基线 1.3 倍，动作是回滚到上一个编译产物并对比 NEFF 算子数。
• NeuronCore 利用率：低于 40% 且 QPS（Queries Per Second，每秒请求数） 正常，说明融合后 kernel 太小、启动开销占比高，考虑加大 batch。

这些指标我用 沐小悠统计（t.api.plus） 做统一看板，编译耗时和命中率变更通过事件接口打点，省得自己维护一套时序库。站点流量侧的访客分析也在同一处，运维和业务数据放一起看反而方便。

总结#

Trainium 上的模型编译不是"跑一条命令就完事"，它是一次和静态编译器的协商。你必须把 shape 说死、把图结构特征说清楚、把并行方式说清楚，编译器才知道能融合到哪一步。我的经验值：--optlevel=2 加正确 --model-type 能让 NEFF 算子数降到原来的七分之一，代价是编译时间多 18%；这两个数在迭代频率高的阶段完全值得。

如果要把这套流程接进自己的服务，先注册 API.PLUS 账号拿一个免费 API Key，用前面那条 curl 命令跑通连通性探测，再逐步改成 POST 上报编译指标。编译日志解析和产物比对脚本可以放到流水线里做成一个独立步骤，参考资料在 API 文章中心 有同类实践，遇到接口层面的问题直接走 技术支持。

适用范围

• 使用 trn1 / trn2 实例，且模型已经能在 PyTorch 上稳定跑通前向的推理场景。
• 输入 shape 可枚举或可截断分桶，batch 不超过单核 16 的常规 transformer 类模型。
• 需要把编译纳入 CI 流水线、要求编译时间和失败率可度量的团队。

不适用范围

• 训练场景的完整反向传播编译，本文只覆盖推理侧 trace 流程。
• 维度和序列长度完全随机、无法分桶的流式输入。
• 依赖大量自定义算子的模型，Neuron 不支持时会 fallback 到 host，融合收益归零。