| [fix] use packing positions for chat boundaries Co-authored-by: mystri<hanboyou@huawei.com> # message auto-generated for no-merge-commit merge: !481 merge fix/chat-packing-position-boundaries into master [fix] use packing positions for chat boundaries Created-by: mystri Commit-by: mystri Merged-by: cann-robot Description: ## 描述 修复 greedy packing 场景下 block-causal attention 文档边界识别错误的问题。 Qwen3 ChatML 会在同一段对话的不同消息之间插入 EOS，而 greedy packing 后的样本连接位置不一定保留 EOS。因此，根据 EOS 推断文档边 界会： - 将同一段多轮对话错误拆分； - 漏掉两个 packed samples 之间的真实边界。 本 PR 改为使用 dataloader positions 中的归零位置识别 packed-document 边界，并统一用于： - 构造 dense block-causal SDPA mask； - 构造 NPU Varlen Attention 所需的 CPU int64 metadata； - 在 Trainer post-dataloading 阶段通过标准 attention_masks 参数传递结果。 同时保留原有基于 EOS 的 SDPA 路径，兼容没有提供预计算 mask 的直接调用。 ## 类型 - [x] Bug 修复 - [ ] 新功能 - [ ] 重构（即不是新增功能，也不是修改bug的代码变动） - [ ] 构建过程或辅助工具的变动 - [ ] 文档内容更新 ## Checklist: - [ ] 我的代码遵循这个项目的代码风格 - [ ] 我已经自己测试过我的代码 - [ ] 我已经更新了相应的文档 - [ ] 我已经在标题中正确使用了类型标签（例如：feat, fix, refactor, docs, test） ## 如何测试 ### 新增测试用例 1. test_position_boundaries_ignore_message_eos_and_split_packed_samples 验证 dense block-causal SDPA mask 使用 positions 归零识别边界： - ChatML 消息内部的 EOS 不会错误拆分同一段对话； - packed samples 连接处即使没有 EOS，也会被 positions == 0 正确隔离； - 每个 sample 内仍保持 causal attention。 2. test_position_resets_create_varlen_metadata_for_packed_samples 验证多 batch、多个 packed samples 的 position resets 能正确生成 Varlen metadata，包括： - cu_seq_q 和 cu_seq_k 的累计边界； - max_q 和 max_k 的最大序列长度； - metadata 位于 CPU 且使用 int64。 3. test_qwen3_tnd_decoder_consumes_precomputed_metadata 验证启用 NPU Varlen Attention 后，Qwen3 Decoder 能直接消费 Trainer 根据 positions 预计算的 metadata，并从 extra_inputs 中取走 attention_masks。 ### 修改测试用例 1. TestVarlenKeepsPacking 在原有“Varlen Attention 保持 greedy packing”检查基础上，新增对 dataloader positions 的断言，确认每个 packed sample 及 padding 区域起点都会归零。 2. test_qwen3_block_causal_sdpa_get_attention_masks_returns_dense_mask 从验证 EOS 推导 mask 改为验证 Decoder 消费基于 positions 预计算的 dense mask，同时确认 attention_masks 不会残留在 extra_inputs。 3. test_sdpa_forwards_dense_mask_to_functional_sdpa 改为使用 position resets 构造 dense mask，验证该 mask 被正确传递给 PyTorch functional SDPA，并关闭额外的 is_causal 路径。 4. test_trainer_post_dataloading_process_requests_sdpa_mask 验证 Trainer post-dataloading 流程： - 使用原始 positions 构造 block-causal SDPA mask； - 保留 positions 供模型使用； - 将 attention_masks 从 extra_inputs 移至支持 Pipeline Parallel 转发的 extra_kwargs。 ## 其他信息 在这里可以添加任何与这个 Pull Request 相关的其他说明。 See merge request: cann/torchtitan-npu!481 | 5 天前 |