# PR #5904 完整报告

- 仓库：`verl-project/verl`
- 标题：[megatron] fix: Adjust the attention mask shape for VLM with Megatron on NPU
- 合并时间：2026-04-08 19:47
- 原文链接：http://prhub.com.cn/verl-project/verl/pull/5904

---

# 执行摘要

- 一句话：修复 VLM+Megatron 在 NPU 上的注意力掩码形状问题，适配 NPU 环境。
- 推荐动作：该 PR 值得精读，尤其是 NPU 环境下掩码处理的特殊设计（模型前向传 None，后处理保留 2D 掩码），这是硬件适配的典型模式。关注 `build_vlm_attn_mask_*` 函数的实现和 `is_npu_available` 的条件分支。

# 功能与动机

根据 PR body 引用的 issue #5878，VLM+Megatron 流水线在 NPU 上使用较少，之前未适配 NPU 环境。PR 旨在解决 NPU 上 VLM 模型因注意力掩码形状不兼容导致的训练失败问题。

# 实现拆解

1. 在 `verl/models/mcore/model_forward.py` 中，将 VLM 模型的注意力掩码生成逻辑替换为调用新的工具函数 `build_vlm_attn_mask_thd` 和 `build_vlm_attn_mask_bshd`。
2. 在 `verl/models/mcore/util.py` 中，新增上述两个工具函数，它们会在 NPU 环境下返回 `None`（因为 NPU 的 `npu_fusion_attention` 后端期望 `None` 掩码），非 NPU 环境下生成标准的 2D 布尔掩码。
3. 移除 `_build_npu_attn_mask` 函数，并在 `preprocess_bshd_engine` 中不再将 2D 掩码转换为 4D 掩码（改为保留 2D 掩码供后处理使用）。
4. 调整 `postprocess_bshd_engine` 中的 NPU 特殊处理逻辑，使其适应掩码为 `None` 或 2D 的情况。

关键文件：
- `verl/models/mcore/model_forward.py`（模块 model）: 核心模型前向函数，修改了 VLM 注意力掩码的生成方式，直接调用新工具函数。
- `verl/models/mcore/util.py`（模块 model）: 新增 VLM 注意力掩码生成工具函数，并调整 NPU 相关的预处理和后处理逻辑。

关键符号：gptmodel_forward_model_engine, build_vlm_attn_mask_thd, build_vlm_attn_mask_bshd, preprocess_bshd_engine, postprocess_bshd_engine


# 评论区精华

reviewer gemini-code-assist[bot] 指出两处关键问题：1) 在 `preprocess_bshd_engine` 中将 `attention_mask` 设为 `None` 会导致 `postprocess_bshd_engine` 崩溃，因为后处理需要 2D 掩码来过滤填充 token。建议保留 2D 掩码供后处理使用，仅在模型前向调用时传 `None`。2) `postprocess_bshd_engine` 中针对 NPU 的 `.diagonal()` 转换逻辑已失效，因为 4D 掩码不再生成。建议通过保留 2D 掩码来让标准断言逻辑正常工作。最终提交采纳了这些建议，调整了实现。

- NPU 注意力掩码处理逻辑 (correctness): 采纳建议，保留 2D 掩码供后处理使用，仅在模型前向调用时传 None。
- postprocess_bshd_engine 中 NPU 分支逻辑错误 (correctness): 调整逻辑，通过保留 2D 掩码让标准断言工作，避免崩溃。

# 风险与影响

- 风险：1. 回归风险：修改了核心模型前向路径（`gptmodel_forward_model_engine`），如果新工具函数在非 NPU 环境下生成掩码有误，可能影响所有 VLM+Megatron 训练。
2. 兼容性风险：NPU 环境下掩码处理逻辑变更（从生成 4D 掩码改为传 `None`），需确保 `npu_fusion_attention` 后端确实支持 `None` 掩码。
3. 代码健壮性：`postprocess_bshd_engine` 中 NPU 分支的逻辑调整可能引入边界条件错误，如掩码为 `None` 时的处理。
- 影响：1. 对用户：修复后，用户可以在 NPU 上正常运行 VLM+Megatron 训练，扩展了硬件支持范围。
2. 对系统：核心模型前向路径的修改影响所有使用 Megatron 引擎的 VLM 模型训练，但变更集中在掩码生成逻辑，对非 VLM 模型无影响。
3. 对团队：提供了统一的 VLM 掩码生成工具函数，简化了后续维护。
- 风险标记：核心路径变更 , 硬件适配风险 , 缺少测试覆盖

# 关联脉络

- PR #5680 [trainer] feat: add mindspeedllm backend engine support on NPU.: 同为 NPU 适配相关 PR，涉及硬件后端支持。
- PR #5884 [megatron] fix: enable_routing_replay fails with MLATransformerConfig…: 同为 Megatron 引擎修复，涉及模型配置和训练问题。