On Policy Distillation 迁移到 Trainium2（Trn2/Neuron）—— Generate() 问题、排查与解决

Created on May 27, 2026

2026   ·   migration   trn2   neuron   trainium   deep-learning   distill   ·   posts

Bug Summary

把原本在 4x L40S/CUDA 上运行的 OPSD OGLS-SD 训练脚本迁移到 AWS Trn2/Trainium2 后，问题的根源不是单个参数错位，而是 CUDA 训练栈（CUDA/NCCL/DeepSpeed/FlashAttention + Transformers 的 GPU 采样实现）与 Neuron/XLA 编译栈之间的不兼容。

原脚本假定一套在 GPU 上可用且可编译的算子（例如全词表排序、torch.isin、top-p sampling、NCCL 通信等）。在 Trn2 上，这些要么不可用、要么会触发 Neuron 编译器产生不被支持的 HLO op，从而导致编译失败或运行时错误。

主要症状

最大的坑：Transformers 的 generate() 在 Trn2 上触发不被支持的 sort 操作：

Operation sort is not supported on trn2
%sort = ... sort(... vocab_size ...)

一开始怀疑是 EOS stopping criteria 中的 torch.isin，随后又发现 top-p / sampling 的路径也会触发对全词表的排序。经过多次尝试，继续修补 generate() 不是稳定方案，最终选择在 Neuron 环境下完全绕过 model.generate()，改用更小且可编译的生成路径。

为什么 generate() 会在 Trn2 上炸

generate() 最大的问题可以拆成一句话：它不是一个简单的 forward loop，而是一个很复杂的 decoding 系统；里面很多在 GPU 上很正常的操作，在 Trn2/Neuron/XLA 上会被编译成不支持的 HLO op。

很多人直觉上以为 generate 只是这样：

outputs = model(input_ids)
next_token = outputs.logits[:, -1, :].argmax()
input_ids = torch.cat([input_ids, next_token])

也就是：

模型 forward → 取最后一个 token → 拼回去 → 重复

如果只是 greedy decoding，确实差不多。但 Transformers 的 generate() 实际上要支持很多功能：

greedy search
sampling
top-k sampling
top-p sampling
temperature
beam search
EOS stopping
repetition penalty
bad words filtering
min length / max length
logits processor
logits warper
cache management
attention mask 更新
多 batch 不同长度处理

所以它不是一个简单 loop，而是一个大而全的 generation framework。

CUDA/GPU 之所以常常没事，是因为它更偏 eager-friendly，很多 PyTorch op 可以直接丢给 GPU kernel 跑。但 Trn2 走的是：

PyTorch code
→ torch-xla 捕获计算图
→ 转成 XLA HLO
→ neuronx-cc 编译
→ Trainium 执行

也就是说，Trn2 更像是：你写的 PyTorch 逻辑必须能被编译器理解，并 lowering 成 Neuron 支持的算子。 如果中间出现一个 Neuron 不支持的 HLO，比如 sort、isin、某些 dynamic shape 操作、某些 int64 dot、某些复杂 scatter/gather、某些 multinomial/sampling 操作，整个编译就会失败。

这也解释了为什么你看到的 sort not supported on trn2 这么典型。这个报错通常来自 top-p sampling。top-p 的逻辑通常是先对整个 vocab 的概率排序，再累加到阈值 top_p 为止，所以它天然依赖 sort：

logits = outputs.logits[:, -1, :]
probs = softmax(logits)

sorted_probs, sorted_indices = torch.sort(probs, descending=True)
cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
remove_mask = cumulative_probs > top_p

当 vocab 很大时，这就变成对整条词表维度做排序。GPU 上这很常见，但在 Trn2 上会被 lowering 成 XLA HLO 的 sort，然后 Neuron 编译器直接报不支持。

一开始怀疑 torch.isin 也合理，因为 Transformers generate() 还要处理 EOS stopping：

if next_token in eos_token_ids:
   stop

如果 EOS token 有多个，内部可能会用类似 torch.isin(next_tokens, eos_token_ids) 的 tensor membership check。这个在 GPU 上也很普通，但在 XLA/Neuron 里也可能展开成不太友好的 graph。只是后来你已经确认更大的雷是 sampling / top-p / logits warper 路径，所以继续 patch generate() 基本是在打地鼠。

这也是为什么 do_sample=True 很危险。它会把 decoding 推进到 sampling path：

logits → temperature scaling → logits processors → logits warpers → top-k / top-p filtering → softmax → multinomial sample → append next token

里面最危险的点就是 top-p 依赖全词表 sort，multinomial 自身也可能不好编译，EOS stopping / logits warper 也可能触发复杂 mask / scatter / dynamic checks。所以在 Trn2 上，先把 OPSD_GENERATION_DO_SAMPLE=0、TOP_P=1.0 这些参数压住，是为了尽量避开复杂路径；但最终最稳的做法仍然是直接绕开 generate()。

为什么自定义 greedy decode 更稳

自定义版本只做：

for _ in range(max_new_tokens):
   outputs = model(
      input_ids=next_input_ids,
      attention_mask=current_attention_mask,
      past_key_values=past_key_values,
      use_cache=True,
      return_dict=True,
   )
   next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)
   generated = torch.cat([generated, next_token], dim=1)
   current_attention_mask = torch.cat(
      [current_attention_mask, torch.ones_like(next_token)], dim=1
   )
   next_input_ids = next_token
   past_key_values = outputs.past_key_values

这个路径很干净，只有 forward、argmax、cat、attention_mask 更新和 KV cache 更新，没有 sort、top-p、multinomial、beam search、复杂 EOS stopping、复杂 logits processor，所以更容易通过 Neuron 编译。这就是为什么最终解决方案不是继续 patch generate()，而是完全绕开 Transformers 的 generate()。

对 OPSD 来说，这一点尤其关键。普通 SFT 只需要 forward/backward，但 OPSD / OPD / GRPO / on-policy distillation 这类训练，核心还包括 rollout generation：

student 先生成 rollout
→ teacher/student 对 rollout 算 logprob 或 KL
→ 再训练 student

也就是说，生成阶段不是一个可有可无的 eval，而是训练主路径的一部分。CUDA 上它只是普通推理步骤，但在 Trn2 上，generation 也必须能被 XLA/Neuron 编译。因此更准确地说：SFT forward/backward 能跑，不代表 OPSD 的 rollout generation 能跑。

所以这个 bug 的关键不是“generate() 参数没设对”，而是：Transformers 的 model.generate() 不是一个简单 forward loop，而是一个很复杂的 decoding 系统；它里面很多 GPU 上很正常的操作，在 Trn2/Neuron/XLA 上会被编译成不支持的 HLO op。 你遇到的 sort not supported on trn2 就是最典型的例子。

最终解决方案（核心）

在 Neuron/Trn2 上不使用 Transformers 的 model.generate()，改用自定义的 greedy decode（仅使用 argmax，避免排序、top-p、采样等）：

for _ in range(max_new_tokens):
   outputs = model(
      input_ids=next_input_ids,
      attention_mask=current_attention_mask,
      past_key_values=past_key_values,
      use_cache=True,
      return_dict=True,
   )
   next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)
   generated = torch.cat([generated, next_token], dim=1)
   current_attention_mask = torch.cat(
      [current_attention_mask, torch.ones_like(next_token)], dim=1
   )
   next_input_ids = next_token
   past_key_values = outputs.past_key_values

通过环境变量控制行为：

export PJRT_DEVICE=NEURON
export OPSD_NEURON_CUSTOM_GREEDY=1
export OPSD_GENERATION_DO_SAMPLE=0
export TOP_P=1.0

在该模式下，生成阶段只走最小可编译路径（greedy + argmax），避免触发 Neuron 编译器不支持的 HLO。

Takeaway（迁移经验）

Trn2/Trainium 的迁移不是把 CUDA 改成 Neuron 就完成。训练的前向/反向（forward/backward）路径虽然多数情况下能通过编译，但在 on-policy RL、self-distillation 或任何在训练中调用模型生成（generate()）的任务里，生成逻辑里常见的 GPU 优化（如全词表排序、top-p、sampling warper、torch.isin）极易踩到 XLA/Neuron 的编译限制。

稳妥的迁移策略：

先用最小可编译的 generation path 跑通（greedy + argmax + no top-p + no sampling），
再逐步恢复复杂采样功能并验证每一步是否可编译。

项目侧实际改动（摘要）

  • 新增 accelerate_trn2_ogls.yaml
  • 新增 requirements_trn2.txt
  • launcher 改为通过 Neuron/XLA 的环境变量启动
  • 将 opsd_train.py 改为包含 main() 并在 Neuron 下默认使用 sdpa
  • opsd_trainer.py 在 Neuron 下绕开 transformers.generate()，使用自定义 greedy decode
  • WandB 支持 online/offline 的显式控制（环境变量控制）