已提交 Parallel Cross Entropy（Triton）实现：

• 单卡 UT 21 passed（FORCE_TRITON）；2 卡词表并行 1 passed
• 默认路由：小词表走 PyTorch CE，大词表/TP 走 Triton；验收用 default_route 均值
• 详情见 PR / docs/parallel_cross_entropy/
  PR: https://gitcode.com/Ascend/TransformerEngineNPU/merge_requests/148

跟进：将 Triton CE MAX_FUSED_SIZE 降至 4096，修复大词表 UB overflow；UT 复测全绿；bench default_route mean te/pt=0.833。PR https://gitcode.com/Ascend/TransformerEngineNPU/merge_requests/148

跟进验收项：补齐特性文档（接口/支持矩阵/精度/性能/限制）与体验报告；新增显存峰值(<=5%)与重复调用泄漏烟雾测试；UT 24 passed。PR https://gitcode.com/Ascend/TransformerEngineNPU/merge_requests/148

Mass 验收已补齐：精度 case 1500 passed；性能 shape 1500，mean te/pt=0.9958，max=1.0479，单点>5%回退=0。严格门禁 mean<=1 且单点<=5% 已通过。PR https://gitcode.com/Ascend/TransformerEngineNPU/merge_requests/148

已更新 PR 描述：含任务书逐项对照、完整复现步骤，以及 mass 测试规模/范围/门禁（精度 1500/1500；性能 1500 shapes，mean te/pt=0.9958，max=1.0479，单点>5%回退=0）。
PR：https://gitcode.com/Ascend/TransformerEngineNPU/merge_requests/148

@longcat_chen
欢迎认领任务，请参考前Q3社区任务池明确该任务的：

完成的截止日期
开发进展反馈
微信答疑群
任务交付注意事项
等信息。如果您同时认领了多项任务，但无法都能进行投入，可以在部分任务中回复退出.

麻烦您加入到对应微信群，群备注名修改为"社区任务+您的gitcode账号", 后续有相关消息和问题都可以在微信群咨询答疑。 等您加入到微信群后，我这边会在社区任务池里面登记任务责任人。

【进展更新】按检视意见：实验/实践报告不再放在 PR docs/，全文归档到本 Issue。

对应实现 PR：!148 · RFC：#37

原 docs/parallel_cross_entropy.md（已从 PR docs/ 移除，归档于此）

Parallel Cross Entropy（Triton）— TransformerEngineNPU

关联 Issue：#29

1. 接口说明

from transformer_engine.pytorch import parallel_cross_entropy

loss = parallel_cross_entropy(
    inp,                      # (B, SQ, V) or (SQ, B, V), FP32 / BF16
    target,                   # (B, SQ) or (SQ, B), global vocab indices
    label_smoothing=0.0,
    reduce_loss=False,
    dist_process_group=None,  # vocabulary-parallel process group
    ignore_idx=-100,
)

语义对齐 NVTE v2.17：softmax / loss 累加在 FP32；返回 loss 为 FP32；反向梯度写回输入 dtype。

2. 支持矩阵

3. 路由策略（Ascend）

4. 精度数据（对 F.cross_entropy）

单卡参数化 UT（FORCE_TRITON）覆盖：dtype × reduce × swap_dim × label_smoothing、ignore_idx、非连续转置、异常形状。2 卡词表并行冒烟通过。

5. 性能数据（同机同卡 vs PyTorch）

性能数字写在 MR 描述 / RFC（不提交 benchmark JSON 结果文件）。本地可选：

python tests/pytorch/benchmark_parallel_cross_entropy.py --json-out /tmp/ce_perf.json

典型 shape：(2,128,4096) / (4,256,8192) / (2,512,32000) / (2,256,65536) / (1,512,128000)。

6. 显存

• 峰值：相对 PyTorch CE 增幅目标 ≤5%（见 test_memory_peak_vs_pytorch）。
• 连续调用：无泄漏（见 test_memory_no_leak_repeated）。
• 实现上对输入做 in-place 梯度暂存，额外缓冲主要为 loss[N] 与（TP 时）m_d_Xy gather。

7. 已知限制

• Ascend Triton MAX_FUSED_SIZE=4096（更大 tile 易触发 910B3 UB overflow）。
• n_non_ignore 在 Host 侧统计（避免 atomic_add 兼容问题）。
• is_cg_capturable 仅为 API 兼容，NPU 无 CUDA Graph 路径。
• 上游暂无独立 2.17 分支；合入 main，待官方开分支后再 cherry-pick。

8. 文件清单

• transformer_engine/pytorch/cross_entropy.py
• transformer_engine/pytorch/triton/cross_entropy.py
• transformer_engine/pytorch/triton/cross_entropy_kernels.py
• tests/pytorch/test_parallel_cross_entropy.py
• tests/pytorch/test_parallel_cross_entropy_dist.py
• tests/pytorch/benchmark_parallel_cross_entropy.py
• docs/parallel_cross_entropy.md
• docs/parallel_cross_entropy/experience_report.md

原 docs/parallel_cross_entropy/experience_report.md（已从 PR docs/ 移除，归档于此）

昇腾社区开发体验报告 — TransformerEngineNPU #29 Parallel CE

任务

基于 TransformerEngineNPU 实现 Triton Parallel Cross Entropy，语义对齐 NVTE v2.17。

环境

• 硬件：Ascend 910B3，卡 0–3
• CANN：9.0.0（CANN toolkit）
• 沙箱：NPU 沙箱（隔离 HOME/TMP/Triton cache；复用 te_npu_30 工具链只读）

开发过程

1. 对照 NVTE parallel_cross_entropy / Triton online-softmax 两阶段内核移植。
2. NPU 适配：Host 统计 n_non_ignore；减小 BLOCK/num_warps；保持 FP32 归约。
3. 单卡精度/反向对齐 PyTorch；2 卡词表并行 HCCL 冒烟。
4. 性能脚本对比同机同卡 PyTorch CE；补充显存峰值与泄漏烟雾测试。

问题与解决

• Triton atomic_add 在 NPU 上不稳定 → 改为 Host sum。
• 非连续 (B,SQ,V) 转置输入 → 与 NVTE 一致做 contiguous。
• MAX_FUSED_SIZE=32768 触发 910B3 UB overflow → 降至 4096。
• 小词表 Triton 启动开销高于 ACL CE → 默认 V>=16384 才走 Triton（TE_PARALLEL_CE_FORCE_TRITON=1 可强制）。

验收（910B3 / CANN 9.0.0）

• UT：单卡 + 多卡全绿（含 FORCE_TRITON 精度与显存烟雾）
• 性能 default_route mean ours/pt ≈ 0.833（无均值回退、单点 ≤5%）
• PR：https://gitcode.com/Ascend/TransformerEngineNPU/merge_requests/148

建议

• 上游补充 Ascend Triton 对 atomic / warp / UB 的文档约束。
• 官方若开 2.17 分支，请再 cherry-pick 本 PR。

已给您申请A2 8卡 资源用于后续长跑使用，长跑具体要求预计相关同事今日会给出。
issue 30、31复用此回复。

集成环境 框架选择：Megatron core_r0.18.0，MegatronAdaptor core_r0.18.0， TransformerEngineNPU main&2.17，具体快速上手参考：https://gitcode.com/clc2025/MegatronAdaptor.git
基础脚本，可以参考MegatronAdaptor 的 ST脚本，脚本8卡，带TP2PP2EP2就行。数据集enwiki或者deepseek都行。基线为megatron配置config.cross_entropy_fusion_impl == 'native'，实验为config.cross_entropy_fusion_impl == 'te'。
验收标准：长跑2000步，精度在MAPE 2% 误差以内。性能不低于基线。
验收材料：结果日志，测试脚本，环境信息