[Bug]: Qwen3.5-9B/35B训练使用thd格式时，存在精度问题-MindSpeed-Bridge-AtomGit

已关闭

[Bug]: Qwen3.5-9B/35B训练使用thd格式时，存在精度问题 #41

lixionglong创建于 7月25日关闭于 4 天前

lixionglong

7月25日创建

在提交新问题之前，请确保您已经在社区中搜索过相关问题，并使用了社区中提供的资源/工具后，仍未找到满意的解决方式。

⚠️ 安全信息提醒：请仔细检查提供的文本内容，确保其不包含敏感数据信息，包括但不限于：

API 令牌或密钥
密码或身份验证凭证
私有网址或接口地址
个人或机密数据
...

在分享配置信息或代码示例时，请将敏感信息脱敏处理，或使用 <TOKEN> 等占位符替代原有内容。

环境信息

CANN：9.0.0
torch：2.10.0
torch-npu：2.10.0
triton-ascend：3.2.1
megatron：core_v0.16.1
mindspeed：core_r0.16.0
megatron-bridge：v0.3.1
mindSpeed-bridge：3655c07cbcc9

🐛 问题描述

问题概述

使用 qkv_format="thd"训练 Qwen3.5-9B 时，loss 高达 21.37，grad_norm 达 570+。而 qkv_format="bshd" 格式下 loss 约 0.37，正常。

根因：Qwen3VLModel.forward 中创建一个 dtype=torch.int32 的 attention_mask，传入 preprocess_packed_seqs 后，input_ids[i, attention_mask[i]] 将 int 值当作行号索引而非 bool 掩码，导致 combined_embeddings（embedding 输出）和 position_ids 全部被破坏。

复现步骤

使用 --qkv-format thd在 Qwen3.5 模型上通过 mindspeed-bridge 启动训练
运行 SFT 训练
观察 train/loss ≈ 21.37，train/grad_norm ≈ 571，而非正常值（~0.3–1.0）

BSHD 格式（--qkv-format bshd）不受影响，因为此时 packed_seq_params is None，Qwen3VLModel.forward 中整个 if packed_seq_params is not None 分支被跳过。

根因分析

model.py 中 attention_mask 的创建方式：

# 约第503行和第562行，共两处
if attention_mask is None:
    attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)

这个 mask 随后被传入 3 处 preprocess_packed_seqs 调用：

处理 input_ids（token 序列）
处理 combined_embeddings（embedding 输出 / 隐藏状态）
处理 position_ids（RoPE 位置编码）

preprocess_packed_seqs 内部（utils.py:724）：

input_ids_rmpad[start : start + seqlen] = input_ids[i, attention_mask[i]]

当 attention_mask[i] 是 torch.int32 且值全为 1 时：

# int32 mask: [1, 1, 1, ..., 1]
input_ids[i, attention_mask[i]]
# → input_ids[0, [1, 1, ..., 1]]  ← PyTorch 整型索引!
# → 取第 1 行，重复 2048 次！
# → 所有 token 都被替换为 position=1 的 embedding/position_id!

这导致进入 transformer 层的隐藏状态完全错误，后续所有计算产生垃圾输出。

具体例子

假设单条数据长度 2000 tokens，TP=4，CP=1：

TND 格式:
  input_ids            = [1, 2048]           (pad 到 TP*128=512 的倍数)
  attention_mask       = [[1,1,...,1]]       shape [1, 2048], dtype=int32
  combined_embeddings  = [1, 2048, 4096]     (embedding 输出)

preprocess_packed_seqs(combined_embeddings, attention_mask, ...):
  seqlens_in_batch  = [2048]           (全1 mask 求和)
  cu_seqlens_padded = [0, 2048]

  循环 i=0:
    input_ids_rmpad[0:2048] = combined_embeddings[0, attention_mask[0]]

  Bug: attention_mask[0] = [1, 1, 1, ..., 1]  (int32, 2048 个值为 1)
    → combined_embeddings[0, [1,1,...,1]]      ← 整型索引!
    → 取出 token position=1 的 embedding，重复 2048 次
    → 结果: 2048 个 token 全部拿到同一个错误的 embedding!

  期望行为 (bool): attention_mask[0] = [True, True, ..., True]
    → combined_embeddings[0, [True]*2048]       ← bool 掩码!
    → 原样取出全部 2048 行
    → 结果: embedding 值正确保持不变

修复方案

将两处 dtype 从 torch.int32 改为 torch.bool：

 if attention_mask is None:
-    attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
+    attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device)

这使 input_ids[i, attention_mask[i]] 从整型索引（按行号选取特定行）变为 bool 掩码（保留 True 对应的行），即预期行为。