[Bug] Embedding 训练开启 padding_free 后，单 Token 序列可能被错误合并，导致 InfoNCE 样本错位 / zero negative 问题描述

[dog-qiuqiu]

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

[Bug] Embedding 训练开启 padding_free 后，单 Token 序列可能被错误合并，导致 InfoNCE 样本错位 / zero negative

问题描述

在使用 Qwen3-Embedding 进行训练，并开启：

--task_type embedding
--loss_type infonce
--padding_free true

时，如果某些 embedding 输入在模板编码后只有 1 个 token，padding_free 在恢复原始 sequence 边界时可能发生错误。

结果表现为：

• 实际恢复出的 sentence embedding 数量少于真实逻辑 sequence 数量；
• 轻微情况下训练不会报错，但 embedding 与 labels 可能已经错位；
• 严重情况下会触发：

RuntimeError: InfoNCE zero negative: rank=7, group=1

目前看问题出在 revert_padding_free() 对 sequence boundary 的恢复逻辑。

---

环境信息

ms-swift: 3.10.1
task_type: embedding
model: Qwen3-Embedding
loss_type: infonce
hard_negatives: 256
per_device_train_batch_size: 2
padding_free: true

每条训练样本包含：

1 anchor
1 positive
256 negatives

因此每个 local batch 理论上应该包含：

2 × (1 + 1 + 256) = 516 条 logical sequences

InfoNCE labels 数量为：

2 × (1 positive + 256 negatives) = 514

---

实际报错

出错 batch 中打印到：

========== INFONCE ZERO NEGATIVE ==========
rank=7
group_index=1
hard_negatives=256
labels_shape=(514,)
sentences_shape=(260, 2560)
raw_split_indices=[0, 257]
start=258, end=516
split_part_len=2
label_values=[(0.0, 512), (1.0, 2)]
===========================================

理论上应该是：

sentences_shape=(516, 2560)
labels_shape=(514,)

但实际变成：

sentences_shape=(260, 2560)
labels_shape=(514,)

也就是说：

labels 数量仍然正确，但 sentence embeddings 已经从 516 条错误减少到了 260 条。

---

根因分析

Embedding 训练中，_embedding_data_collator() 会把每条样本展开为独立 sequence：

anchor
positive
negative0
negative1
...
negative255

开启 padding_free=true 后，_data_collator() 会执行：

batch[:] = [self.packing_row(batch)]

将所有 sequence flatten 成一条 token stream。

packing_row() 中 position_ids 的构造逻辑为：

packed['position_ids'] = sum(
    (list(range(x)) for x in length),
    start=[]
)

例如三条 sequence 长度为：

5, 1, 3

会生成：

0 1 2 3 4 | 0 | 0 1 2

但 revert_padding_free() 当前使用：

pos = position_ids[0]
resets = torch.where(pos[1:] < pos[:-1])[0] + 1

来判断新 sequence 的边界。

如果前一条 sequence 长度为 1：

seq A: [0]
seq B: [0, 1, 2]

flatten 后：

0 | 0 1 2

真实 sequence 边界是：

0 -> 0

但当前代码判断：

0 < 0  # False

因此无法识别该边界，最终会把两个独立 sequence 错误合并。

---

实际现场证据

我在 embedding collator 中增加了 debug 日志。

某个普通训练 batch 中得到：

========== PADDING_FREE BOUNDARY MISMATCH ==========
rank=2
zero_count=516
detected_seq_count=504
position_ids_shape=(1, 1813)
total_tokens=1813
====================================================

其中 zero_count=516 很关键。

因为 packing_row() 对每一个 logical sequence 都会从：

position_id = 0

开始，所以：

zero_count=516

说明该 batch 实际包含：

516 条 logical sequences

但当前 revert_padding_free() 的边界识别逻辑只识别出：

504 条

因此这个 batch 已经静默丢失了：

516 - 504 = 12 个 sequence boundary

实际 position_ids 中也可以直接看到：

..., 0, 1, 2, 3, 4, 0, 0, 1, 2, ...
                       ^  ^

这里对应：

... | 一个 1-token sequence | 下一个 sequence ...

中间的：

0 -> 0

无法被：

pos[1:] < pos[:-1]

检测到。

---

为什么这个问题可能比 crash 更严重

这个问题并不一定每次都会报错。

例如：

logical sequences: 516
recovered sequences: 504

此时 InfoNCE 中仍然可能存在大量 negative，因此训练可以继续进行。

但这意味着：

sentence embeddings 数量

已经和：

labels / group boundary

不再严格对应。

也就是说可能出现：

训练 loss 正常、训练继续执行，但实际 positive / negative 的 embedding 分组已经发生错位。

只有在极端 batch 中，丢失的 sequence 足够多，才会最终触发：

InfoNCE zero negative

本次出错 batch 为：

expected: 516
actual:   260
difference: 256

而：

hard_negatives=256

两者完全一致。

---

最小复现

无需完整 embedding 数据集即可复现这个边界问题。

例如：

position_ids = torch.tensor([
    [0, 1, 2, 0, 0, 1, 2]
])

它实际表示三条 sequence：

seq0 = [0, 1, 2]
seq1 = [0]
seq2 = [0, 1, 2]

但当前逻辑：

pos = position_ids[0]
resets = torch.where(pos[1:] < pos[:-1])[0] + 1

detected_seq_count = 1 + len(resets)
print(detected_seq_count)

输出：

2

而真实 sequence 数应该是：

3

---

期望行为

padding_free 应该完整保留 logical sequence boundary，而不应受到 sequence 长度影响。

对于 embedding 训练：

516 条 logical sequences

无论其中是否存在 1-token sequence，最终都应该恢复为：

516 条 sentence embeddings

而不是：

504
260
...

---

修复建议

建议不要仅通过：

pos[1:] < pos[:-1]

从 position_ids 反推 sequence boundary。

因为对于：

... | [0] | [0, 1, 2] | ...

这种情况，仅通过 position ID 无法利用严格下降条件恢复真实边界。

更可靠的方式是直接使用 packing / FlashAttention 阶段已经存在的显式 sequence boundary 信息，例如：

cu_seq_lens_q

或在 packing_row() 阶段显式保留每条 sequence 的长度 / cumulative offsets，在 revert_padding_free() 中直接使用。

不建议简单把：

<

改成：

<=

因为还需要确认：

• multimodal position ids
• mRoPE
• sequence parallel
• 其他可能出现重复 position id 的场景

是否会受到影响。

---

临时规避方案

关闭：

--padding_free false

即可绕过该问题。

关闭后，每个：

anchor / positive / negative

都会保持为独立 batch row，不需要通过 position_ids 再恢复 sequence 边界。

---

相关代码

Embedding 展开：

def _embedding_data_collator(...):
    ...
    indexes = ['anchor_', 'positive_']
    if max_neg is not None:
        for i in range(0, max_neg):
            indexes.append(f'negative{i}_')

    for prefix in indexes:
        new_batch += self._fetch_inputs_startswith([b], prefix)

    res = self._data_collator(new_batch, padding_to=padding_to)

Padding-free flatten：

if self.padding_free:
    batch[:] = [self.packing_row(batch)]

Position IDs 构造：

packed['position_ids'] = sum(
    (list(range(x)) for x in length),
    start=[]
)

当前 sequence 恢复逻辑：

pos = position_ids[0]
resets = torch.where(pos[1:] < pos[:-1])[0] + 1

当一个 1-token sequence 后面紧跟下一条 sequence 时，会出现：

0 -> 0

从而导致 sequence boundary 丢失。

How to Reproduce / 如何复现

参考问题描述

Additional Information / 补充信息

No response

[tastelikefeet]

已经修复。把判断条件改为 < 改成 <=，额外加了 >= 0 来处理 sequence parallel 的 -1 padding，另外最后一段的长度也改成按有效 token 数算。没有换成 cu_seq_lens，因为那条分支本来也是先展开成 position_ids 再处理。可以试一下main 分支 #9924

[dog-qiuqiu]

thanks