变更摘要

该 PR 主要为 Flux DanceGRPO 训练添加了 A5 服务器的适配支持，包括新增 A5 性能测试脚本、引入 NPU 版本检测机制以自动适配 HCCL 环境变量、将图片保存改为异步执行以减少训练阻塞，并在训练循环中增加了 rollout time、SFT time 和 step time 等性能统计信息的打印。

主要改动

• 新增 A5 性能测试脚本 posttrain_flux_dancegrpo_A5.sh：提供面向 A5 NPU 服务器的完整训练启动脚本，包含 NPU 相关环境变量配置、多卡分布式参数、以及训练结束后自动计算并输出平均单步耗时和每秒样本数的性能统计逻辑。
• 新增 NPU 版本检测工具模块 device.py：定义 NPUVersion 枚举和 check_npu_version()、get_npu_version() 函数，通过 torch_npu.npu.get_device_name() 自动识别 NPU 型号（A2/A3/A5），用于在代码中按硬件版本做出适配决策。
• A5 服务器 HCCL 配置适配：在 SoraGRPOTrainer 初始化中，当检测到 NPU 版本为 A5 及以上时跳过设置 HCCL_OP_EXPANSION_MODE=AIV，避免对 A5 服务器产生不必要的环境变量覆盖。
• 图片保存改为异步执行：在 FluxGRPOTrainer 中引入 ThreadPoolExecutor，将原本同步的 image.save() 调用改为通过 _save_image_async 方法异步提交保存任务；同时新增 --save_images 布尔参数（默认 True），允许在纯性能测试时禁用图片保存。
• 训练时间统计增强：在 SoraGRPOTrainer 训练循环中加入 rollout time（sample_reference 耗时）、SFT time（每次 grpo_one_step + loss backward 耗时）和 step time 的测量与打印，并在 FluxGRPOTrainer 中增加 step_counter 计数及按步归档图片的目录结构。

代码审查

All four files have been thoroughly reviewed. Let me now provide the closing summary.

审查总结

审查覆盖的 4 个文件：

按优先级统计：

• P0：0 个
• P1：0 个
• P2：1 个（total_loss / total_loss_addnum 除零风险）
• P3：3 个（异常静默吞没、循环内重复 os.makedirs、未使用的 buffer 变量）

整体风险评估：低风险。 此改动主要为性能测试脚本和训练日志增强，不修改核心训练逻辑。新增的 device.py 模块逻辑正确，NPU 版本检测的保守回退策略合理。P2 的除零问题在常规配置下极难触发，但建议防御性修复。三个 P3 问题为代码质量改进项，不影响功能正确性。

💬 仅评论

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here

PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

💡 Tip:

• Committer can comment /approve or /lgtm
• Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

MoCuishle-M, thanks for your pull request. All authors of the commits have signed the CLA. 👍

/approve

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.