把原本在 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()。
自定义版本只做：
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。
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.yamlrequirements_trn2.txtopsd_train.py 改为包含 main() 并在 Neuron 下默认使用 sdpaopsd_trainer.py 在 Neuron 下绕开 transformers.generate()，使用自定义 greedy decode