feat: precision verification framework + PP dump fix — diagnostic, comparison, and baseline tools for ON/OFF prefix-sharing - #40
Merged
Conversation
- 删除 test_logprob_extended.py：core.logprob 已删、函数无调用，测试成孤儿导致收集中断 - 移除 test_prefix_last_in_chain_reuse_resolves_to_ancestor_with_packed_slot 及其 helper： 手工拼的 row2->row1(缺 slot 的 reuser) 链不可能由 trie 检测产生，且断言(解析到祖先) 与 _resolve_provider_for_position 的 keep_start-1 扩展正确行为相反 - 新增 test_chain_reuse_resolves_position_to_provider_with_matching_label： 用 [1,2,3,4,5]/[1,2,3,7,8]/[1,2,3,7,9] 自然链验证 seq2 位置2(token 3) 解析到 直接 provider seq1(label 7) 而非 seq0(label 4)，prefix-last 落在 seq1 真实 packed slot
…ec 数据结构 - verl_mcore/context/vocab_logprobs/observability: restore 由逐元素 2D 注入改为整段切片 + prefix-last 重算 - planner: 仅生成 prefix-last spec，删除 interior spec 与 is_shared_prefix_interior 标记 - 删除死字段 reuse_1d_pos、valid_indices 链、provider_predict_pos、reuse_first_suffix_label_pos - 同步单测与 overview 文档；更新精度诊断工具
restore循环：i→reuser_idx、provider→provider_idx、idx→prefix_last_spec、prefix_last_by_row→prefix_last_spec_by_reuser、saved_key→saved_logits_key、provider_logits→saved_provider_logits；unfold/fold循环：row→seq_idx，lp_/ent_/logp_等隐晦缩写展开为log_probs_/entropy_；vocab_logprobs patch docstring收敛为prefix-last单语义。纯命名/注释调整，零行为变化。
last-attn/rope_freqs的ON dump失败提示从prefix_log.warning改为print，与同区块其它diag输出一致（直接走stdout）。
新增 flash_atten_npu_test.py，验证 sparse_mode=3 (rightDownCausal) 在 TND varlen 下对 reuser(Q<Skv) 的右对齐 causal。 - Probe A-E: provider-only mode2 sanity / 单reuser mode3(核心) / 全batch mode3单次调用 / 反向(128-tile) / seq>2048 - ground truth = TorchReferenceBackend；NPU-only，非NPU 干净 skip；__main__ 自带决策树 - 配套 plan: note 目录 npu-flash-attn-mode3-tnd-plan.md
新后端 flash_atten_npu_tnd (NpuFlashAttentionTndBackend)：TND varlen + sparse_mode=3 (rightDownCausal)，一次调用覆盖全 batch，reuser Q 仍 suffix-only。 - mode 3 原生支持 reuser(Q<Skv) 右对齐 causal（prefix 全见 + suffix causal），用现成压缩 [2048,2048] mask，不自建、不拆调用 - provider 段(Q==KV) 退化成标准 causal，与 baseline 一致 - 经 flash_atten_npu_test.py 5 probe 实测验证（A-E 全过，含反向 + seq>2048） - 老 BSH 后端 flash_atten_npu 保留作回退；config 里 backend='flash_atten_npu_tnd' 切换 - 注册到 factory/config/__init__，加 test_factory_flash_atten_npu_tnd
attention() 里多写了一个 _，导致 ValueError: not enough values to unpack。 _prepare_flash_inputs 返回 8 元组 (q,k,v,cu_seqlens_q,cu_seqlens_kv,max_seqlen_q,max_seqlen_kv,pad_layout)。 实验文件走自己的 _varlen_tnd_call 没踩到，真实 backend 路径才暴露。
- 新增 --token POS（单个 int，默认 0），比较 packed[POS] 而非写死 packed[0]
- cmp_first_token → cmp_packed_token(dir_on, dir_off, pos, layer)；_logits_first_token → _logits_at_pos；_first_token_metrics → _vec_metrics；_print_first_token → _print_packed_token
- 输出名/section 头不再有 first_token 字样：[packed_token] attn_L{layer}_pos{pos} / logits_pos{pos}
- --layer 同时控制 packed_token 的 attn 层（默认最后一层）；logits 永远最后一层
- logits top-K 表按 val 排序、去掉 REL_ERR 列；attn 表保持原样
之前 packed[pos] 直接比 ON[pos] vs OFF[pos]，但 ON 是裁剪后的 suffix-only packed、OFF 是完整 packed，两者 token 不对应——比较的是不同 token，结论无意义。 修复：新增 _aligned_vec_at_pos，先用 align_mask（OFF cu_seqlens + ON prefix_lens）把 OFF 的 suffix 段抽出来与 ON 对齐，再取 [pos]。pos 索引对齐后的 suffix-packed 空间（ON/OFF 一致，指向同一 token）。attn 和 logits 都走这条对齐路径，top-K 表同样先对齐。
abs 排序会把大负值 logit（softmax 后概率极低、不会被选中的 token）顶到表头，掩盖真正的高 logit 候选。改为 max(on, off) 带符号值排序，让候选 token 排前面。 仅改 1D _print_topk_vec 的 val 分支（logits/attn 表）；2D logp/entropy 表的 val 仍用 abs（logp 本身为负，abs 有意义）。
新增环境变量 PREFIX_SHARING_FORCE_ZERO_PREFIX： - prefix_detector.detect() 命中即早返回 0-prefix（所有 provider、无复用） - verl_mcore build 旁路 has_sharing 早返回，让 ON pipeline 仍跑全路径： trim(0-prefix→不裁)/build_kv(全 provider→无注入)/attention(mode 3 全序列)/restore(no-op) - 用途：跑 ON(0-prefix, mode 3, 全序列, 无裁剪/注入/位置补偿) vs OFF(mode 2, 全序列) ≈ → kernel mode 不是根因；偏差大 → kernel mode 是根因。
…nd-truth + skip-restore S1.1: RoPE 外推自动诊断日志（step_min/max/mean/std + step_is_uniform） - 不需要 env var，每次外推触发时在 megatron_runtime.py 自动打印 - step_is_uniform=False → 非线性 RoPE 调度，外推可能引入误差 S1.2: PREFIX_SHARING_DIAG_ROPE_GROUND_TRUTH=1 - 从 inv_freq 直接计算完整频率表，跳过线性外推 - 启用后 ON vs OFF 一致 → RoPE 外推是根因 S4.3: PREFIX_SHARING_DIAG_SKIP_RESTORE=1 - 跳过 logprobs restore，直接返回 forward 原始输出 - suffix 一致 → restore 是根因；不一致 → 偏差在 forward Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
新增 step01_vs_step12_maxdiff / is_linear 字段： - 计算 pos_emb[0→1] 和 pos_emb[1→2] 的 step 差值 - maxdiff ≈ 0 → pos_emb[p] 对 p 严格线性 → 外推数学正确 - maxdiff > 0 → pos_emb[p] 对 p 非线性 → 外推引入误差 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
ON path: dump_rope_emb_layer() called in _apply_positioned_rope after RoPE
OFF path: extract QKV, reconstruct standard 0..seg_len-1 position IDs,
apply_rotary_pos_emb manually, dump via dump_rope_emb_layer()
Both paths write to rope_emb.pt for cmp_diag comparison.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…position-aligned comparison 1. rope_emb.pt 新增 positions 字段（per-token absolute position IDs） 2. cmp_rope_emb 重写: - first packed token Q/K 直接对比（row 0, pos 0，provider 必相等） - first row Q/K 按 position ID 对齐后逐位对比（0..L-1，两路径同） - reuser row 不直接对比（ON absolute pos vs OFF relative pos，设计差异） 3. ON 调用点传入 packed_position_ids 4. OFF 调用点传入重建的 segment-relative positions Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
dump: diagnostic_dump_verl080.py
- dump_rope_emb_verl080(layer, query, key, num_layers)
- 格式: rope_emb.pt = {layer: {"query": [T,H,D], "key": [T,H,D]}}
- ON suffix-only / OFF 完整 packed，不需要 positions（suffix 对齐用已有元数据）
cmp: cmp_diag_verl080.py
- cmp_rope_emb_layer(): per-layer Q/K cosine（suffix 对齐，同 cmp_attn_layer）
- cmp_rope_emb_token(): 指定位置 Q/K 向量对比（同 cmp_packed_token）
- _print_rope_emb_per_layer(): 每层 Q_cos + K_cos 报告
调用点:
- ON: megatron_runtime.py → dump_rope_emb_verl080
- OFF: attention.py → dump_rope_emb_verl080
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ge/cat) OFF 路径 diag dump 用 torch.arange/torch.cat 构造每 segment 的 position IDs， 但 attention.py 顶部未导入 torch（与 forward_step.py:170 惯例一致）， 导致 OFF 路径报 NameError: name 'torch' is not defined。 在 diag block 内局部 import torch，保持顶部零 torch 导入且非 diag 路径零开销。
OFF 路径 post-RoPE Q/K 之前用 get_query_key_value_tensors + 重建 positions + index_select + 重算 apply_rotary_pos_emb，既脆又有 bug（_k_pos_emb or _q_pos_emb 对 tensor 求 bool 崩溃）。post-RoPE Q/K 是 original_forward 内部中间变量，唯一 拿到真实张量的办法是 hook。 - diagnostic_dump_verl080: 新增 capture_post_rope_qk() context manager， monkey-patch megatron.core.transformer.attention.apply_rotary_pos_emb， 捕获 mcore THD prefill 每层 Q→K 两次调用的返回值，finally 还原。 ON 不受影响（apply_rotary_pos_emb 在 megatron_runtime 自有命名空间，不走 attention 模块全局）。 - attention.py OFF 分支：original_forward 外套 capture_post_rope_qk， 直接取 captures[0]/[1] dump，删掉重算逻辑。 - 修潜伏 bug：dump_rope_emb_verl080 原手写 os.path.join 但未 import os， NameError 被 except 静默吞掉导致 rope_emb.pt 从未写盘；改用 _save_tensor。 → ON/OFF 都需用本提交重跑一次。
- _print_shapes 加入 rope_emb.pt；_shape_of 支持嵌套 dict（每层 {query,key,
positions}），显示 (dict,NL,Q(T,H,D))，一眼看出 ON(suffix-only)/OFF(full)
token 数差异。
- cmp_rope_emb_token 的 pos 失败从笼统「无法对齐或越界」改为 _diag_rope_pos_fail
三分：缺失 / 对齐失败(带 n_on,n_off,align_mask len,sum) / pos 越界(带对齐后
token 数)，定位 pos16 失败根因。
dump_rope_emb_verl080 的 flush 之前用 _save_tensor，但 _save_tensor 内部对入参 做 .detach().cpu().clone()，dict 没 .detach() → AttributeError 被其 except 静默吞掉， rope_emb.pt 永不写盘（ON/OFF 都没写，cmp 显示 missing）。 da39661 原本想修 os 未导入的 bug，换成 _save_tensor 反而引入这个 dict 回归。 现仿 _flush_attn_buffer（写 attn_outputs.pt 的方式）：rank0 + 直接 torch.save(dict) + 本地 import os，失败 print 不再静默。entries 插入时已 cpu clone，无需再搬。 注：用户目录里的 rope_emb_on.pt / rope_emb_off.pt 是旧版本残留，当前代码不产生 （grep 零命中），真正名字是 rope_emb.pt（dump 与 cmp 一致），需清掉残留重跑。
dump 侧（rope_preqk.pt，旋转前 Q/K）：
- diagnostic_dump_verl080: 抽 _flush_dict_buffer 公用写盘 helper（避开 _save_tensor 对
dict 失败的 bug）；新增 dump_rope_preqk_verl080；capture_post_rope_qk → capture_rope_qk，
每次 apply_rotary_pos_emb 同时存 {"pre":输入, "post":返回}。
- attention.py(OFF): hook 解包 pre/post，分别 dump rope_preqk.pt / rope_emb.pt。
- megatron_runtime.py(ON): _apply_positioned_rope 入口 dump 旋转前 Q/K。
cmp 侧（cmp_diag_verl080）：
- _ROPE_STAGES=[pre,post]；cmp_rope_emb_layer/token 改 stage 参数，支持单 stage 调用。
- rope_freqs 加 layer 过滤（cmp_rope_freqs(layer=...)）；新增 cmp_rope_freqs_token
（指定 pos 的角度向量对比，max_abs 应为 0）。抽 _align_rope_freqs_layer helper。
- shapes 表加 rope_preqk.pt / rope_emb.pt；_shape_of 支持嵌套 dict。
- main 按 RoPE 计算顺序重排：per-layer 与 token 都是 pre Q/K → rope_freqs → post Q/K，
定位分歧出现在哪一步（pre 偏=上游；freqs 偏=角度表；post 才偏=旋转应用）。
- pos 失败诊断细化（_diag_rope_pos_fail 三分：缺失/对齐失败/越界）。
注：layer 1-indexed（megatron self.layer_number 从 1 起），第一层用 --layer 1。
1) rope_emb → rope_postqk 改名（全仓 6 文件） rope_emb.pt 这名字误导（存的是旋转后 Q/K，非 embedding 表）。统一为 rope_postqk.pt， 与 rope_preqk.pt（旋转前）配成一对。函数/变量/buffer/label 一并改： dump_rope_emb_verl080→dump_rope_postqk_verl080, cmp_rope_emb_*→cmp_rope_postqk_*, _ROPE_EMB_BUFFER→_ROPE_POSTQK_BUFFER, 结果名/section 头同步。 判断 pre/post 用 "preqk" in name（postqk 不含 preqk，仍正确分流）。 2) rope_freqs: ON/OFF 统一 per-token，合并 dump 函数 原 rope_freqs_on.pt(per-token) / rope_freqs_off.pt(raw 表) 内容不对称、文件名冗余 （目录已区分 ON/OFF）。统一：两边都存 per-token 角度 rope_freqs.pt [T,1,1,D]。 - dump_rope_freqs_on/off 合并为 dump_rope_freqs（ON/OFF 共用）。 - OFF 调用方 index_select raw 表成 per-token 再 dump。 - cmp 去掉「从 raw 表按 cu_seqlens 重建」逻辑，_align_rope_freqs_layer 简化为直接对齐。 - shapes 表加 rope_freqs.pt；docstring 更新。 注：v070 cmp_diag.py 的 rope_freqs_on/off 引用保留（v070 dump 不写 rope_freqs，dead code）。
--topk 时 rope 三个 stage 的 dim 不再各自独立排序，而是以 rope_postqk 的 sort-err top-K dim 为基准，rope_preqk 显示同样 dim，rope_freqs 显示 dim%D（角度按 head_dim 共享），便于逐 dim 追溯误差来源： - preqk[dim] 已偏 → 上游（投影/hidden_states） - freqs[dim%D] 偏 → 角度/位置 ID - 两者都 ~0 但 postqk[dim] 偏 → 旋转应用本身 实现： - _print_topk_vec 返回选中的 dim 列表。 - 新增 _print_vec_at_dims（按指定 dim 打印，不排序）。 - 新增 _load_rope_freqs_vec_at_pos（对齐后取 [pos] 角度向量 [D]）。 - main rope topk 段：postqk Q/K sort-err top-K → preqk 同 dim、freqs dim%D。 - postqk 排序用 --sort-err（默认 abs），原先硬编码 val。
重构 OFF diag 块把 _off_positions 上提时漏抄了 device=（原 device=_off_q.device）， arange 默认建在 CPU，而 _q_pos_emb 在 GPU → _q_pos_emb.index_select(0, _off_positions) 抛 RuntimeError(device 不匹配)。该 diag 块无 try/except，异常直接崩 OFF forward， 导致 rope_freqs/rope_postqk/rope_preqk 三个文件都没写盘。 修：arange 显式 device=_cu.device（cu_seqlens 在 GPU），并用 int() 把 GPU 标量转 Python int 再给 arange。
定位「postqk 对但 attention_output 偏」的根因：dump attention 实际用的完整 KV， 逐元素比 ON(expanded) vs OFF(full)。 dump 侧： - ON(megatron_runtime): build_kv 后 dump expanded_key/value → expanded_kv.pt （prefix 复用 + suffix，attention 实际用的完整 KV）。 - OFF(attention): dump 完整 K(hook 截的 post-RoPE K) + 完整 V(get_qkv，V 不旋转 故忠实) → full_kv.pt。 cmp 侧(cmp_diag_verl080)： - cmp_attn_kv：ON expanded_kv.pt vs OFF full_kv.pt，逐层 K/V 分别 max_diff + cos， 两者 reshape [T,-1]，全量不需 align_mask。阈值极严（expanded 应精确==full）。 - _print_attn_kv 表插在 postqk 与 attn_output 之间。 - K/V 都 OK → attention 输入一致，偏差必来自 attention 计算/mask； K 或 V DIFF → bug 在 build_kv 的 prefix 复用。 - shapes 表加 expanded_kv.pt / full_kv.pt（per-side 文件，对侧显示 missing 正常）。
…call 之前 OFF 的 V 用 get_query_key_value_tensors 事后重算（re-call），而 K 走 apply_rotary hook 是 in-context。数据显示 L1「K 精确(0)、V 偏 1.2e-4」的不对称——K 和 V 同出 get_qkv， 理应同对/同偏，怀疑是 re-call 时 hidden_states 已被 forward 的 in-place 操作改写导致 V 失真。 改：capture_rope_qk(attention_module) 额外 hook 实例方法 get_query_key_value_tensors， 把返回的 V（第 3 元素）in-context 截下。yield (qk_captures, v_captures)。 attention.py OFF 块传 self、解包、用 _v_caps[-1] 作为 full_kv 的 V（删掉 re-call）。 重跑后看 attn_kv 的 V_MAXDIFF： - 若 V 也归零（和 K 一致）→ 之前是 dump artifact，expanded_kv 其实没问题； - 若 V 仍偏 → ON 侧 V 真有问题，深挖 build_kv/store。
_orig_get_qkv 之前取成绑定方法（attention_module.get_query_key_value_tensors，self 已绑）， _capturing_get_qkv 又经 MethodType 绑定把 attention_module 作为 self_ 传入，原方法收到 两次 self，位置参数错位 → 'got multiple values for argument output_gate'。 修：_orig_get_qkv 改取未绑定的类方法函数 type(attention_module).get_query_key_value_tensors， _capturing_get_qkv 显式接收 self_ 并传给未绑定原函数（调用方经 MethodType 自动注入 self_）。
postqk 对、build_kv 出来 V 偏，需区分 V 是 build_kv 之前就偏（get_qkv/hidden_states） 还是 build_kv 引入。在 megatron_runtime build_kv 调用前 dump 进去的 value（raw V， get_qkv 出来、build_kv 之前），存 build_kv_input_v.pt。 cmp_build_kv_input_v：ON build_kv_input_v vs OFF full_kv V（suffix 对齐）。 输出 ON_T / OFF_T —— 若 ON_T < OFF_T 说明 ON 把 hidden_states 裁成 suffix-only。 判读：PASS（build_kv 前 V 一致）→ 偏由 build_kv 引入；FAIL（已偏）→ 根因在 get_qkv/hidden_states。
_rope_postqk_cos_for_layer 额外算 Q_max_diff / K_max_diff（对齐后整体 abs max）。 printer 表格改显 max_diff（替换 cos_min 列；cos_min 仍在 metrics 用于 WARN 判定）。 目的：验证 Q/K 的 max_diff 是否和 V 的 ~6e-5 一致。若一致 → Q/K/V 同源同偏（裁剪 matmul tiling）；若 Q/K≈0 而 V~6e-5 → V 走了不同路径，需深查模型 attention 实现。
… shape mismatch - 2D: rel_max/rel_mean/abs_mean from actual per-pair diffs (not hardcoded) - 2D: pearson from paired rows (fix shape mismatch 1024 vs 4096) - 2D: _print_topk_2d for logp/entropy positions - within_batch: same fixes Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…t entries Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…cal wrappers Added _print_plain_baseline and _print_rope_baseline wrappers with SNG_T/STK_T or neutral column labels instead of ON_T/OFF_T. within_batch uses within-batch pairwise headers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2D logp/entropy now compares row i vs row i+k*num_seq (same sequence stack copies) instead of all C(N,2) pairs across different sequences. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…opy) Was comparing row 0 vs row 1 (different sequences → huge false diffs). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e quality) - sd/md → single_data/multi_data - cu_s/cu_m → cu_seqlens_single/cu_seqlens_multi - st/mt → single_tensor/multi_tensor - B → num_sequences, T_i → tokens_in_seq - lyr → layer_index, n → stack_count - i/k → seq_index/copy_index, j → multi_offset - worst_md → worst_max_diff, layer_md → layer_max_diff - on_f/off_f → single_flat/multi_flat - q_all/k_all → first_cos_list/second_cos_list - Consolidated comparison loops into DRY for-loops over file lists. - Removed all single-letter variable names in logic functions. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ag] ... ##### - attention.py (patch ON): added missing END marker - megatron_runtime.py: unified ######### → ##### [PS-diag] format - Megatron attention.py (invasive): added ##### wrappers to both blocks - All 12 dump blocks now consistently wrapped for easy identification Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…mble tool + variable readability refactor ## New features - dump 侧 PP-aware routing: 替换 _rank0_only() 为 scope 模型 (global/tp_vocab/pp_last/pp_stage) - 新增 _stage_last_layer() 解决 PP 下 buffer flush 永不触发的问题 - 新增 assemble_dump.py: 将多卡分散 dump 文件合并为单卡兼容格式 - 修复 forward_step.py PP 崩溃: restore_via_2d_unfold_verl080 加 isinstance guard ## Refactored - 所有 dump 代码段变量名可读性优化 (_lp→log_probs_nested, _ol→original_lengths 等) - cmp_diag.py / cmp_diag_verl080.py 变量名规范化 (lo/lf→logits_on/logits_off, lyr→layer_idx 等) - exception 变量统一为 exc ## Files: 10 files (9 modified + 1 new)
…verl080.py Key renames: - fp→filepath, fa/fb→filepath_on/off, da/db→attn_dict_on/off - ma/mb→meta_on/off, la/lb→layers_on/off - d→attn_dict/rope_dict/kv_dict (context-specific) - s→shard_tensor, t→tp_rank, pl_fp→prefix_lens_filepath - _q→query_tensor, _name→result_name, _D→head_dim - _sec→section_label, _stage→stage_label, _crop→crop_note - _on_row1_end→on_row1_end, _cmp_len→compare_length, etc. - Various underscore-prefixed locals→descriptive names
…assembly Root cause: assemble_dump.py copied parallel_info.json verbatim from raw dump, leaving tp_size/pp_size > 1 and scopes like 'tp_vocab'/'pp_stage'. cmp_diag then tried to read logits_tp0.pt instead of logits.pt → showed (missing). Fix: after merging all shards, update parallel_info.json to tp_size=1, pp_size=1, all scopes='global', matching the single-card assembled directory.
…semble manifest rewrite cmp_diag_verl080.py: - _load_logits: remove TP vocab concat logic → just load logits.pt - _logits_shape: deleted - _load_manifest: deleted - _print_topology: deleted - _print_shapes: remove manifest params, treat logits like any other file - main(): remove manifest loading and topology printing Rationale: multi-rank assembly is done by assemble_dump.py before cmp runs; cmp should work on flat single-card data only — no parallel awareness. assemble_dump.py: revert the manifest rewrite (no longer needed since cmp doesn't read manifest).
boundless-future
force-pushed
the
parallel-precision-verify
branch
3 times, most recently
from
July 1, 2026 13:30
ce3b3c9 to
8b8a042
Compare
boundless-future
force-pushed
the
parallel-precision-verify
branch
from
July 14, 2026 08:22
50310e4 to
5cab79d
Compare
…U autograd compat On Ascend NPU, torch.nested.nested_tensor is intercepted by torch_npu and throws. verl/__init__.py unwraps it (via __wrapped__), but the unwrapped version does not support autograd backward on NPU tensors -- gradients are silently detached. In PP training, this causes the last stage's loss backward to never reach the model, so no send_backward P2P is issued, and all prior PP ranks hang forever at recv_backward. TP doesn't hang because there is no P2P dependency, even though gradients are also broken. Fix: replace all torch.nested.nested_tensor(rows, layout=torch.jagged) with torch.nested.as_nested_tensor(rows, layout=torch.jagged) in verl_mcore.py. The latter preserves the autograd graph correctly on NPU after unwrapping. Affected functions: - _fold_2d_to_nested (root cause of PP hang) - _trim_nested_batch (forward batch trimming) - _trim_plain_batch_thd (forward batch trimming) Co-Authored-By: Claude <noreply@anthropic.com>
boundless-future
force-pushed
the
parallel-precision-verify
branch
from
July 14, 2026 12:33
b9904c0 to
b7fe400
Compare
…rify # Conflicts: # prefix-sharing/prefix_sharing/integrations/megatron_runtime.py # prefix-sharing/prefix_sharing/integrations/verl_mcore.py # prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py # prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py # prefix-sharing/prefix_sharing/tools/diagnostic_dump.py
- Add missing _load_manifest function to cmp_diag_verl080.py (F821) - Remove unused imports: Sequence, ensure_global_packed_token_lengths, inspect, torch, _cached_parallel_info, json, os (F401) - Fix f-strings without placeholders to plain strings (F541) - Remove unused variable device assignment (F841) All flagged by ruff --select F,E9 (pyflakes + syntax errors).
📊 Coverage reportChanged files in this PR |
boundless-future
changed the title
boundless-future
changed the title
Jackie2049
self-requested a review
July 16, 2026 07:28
Jackie2049
reviewed
Jul 16, 2026
Jackie2049
left a comment
Jackie2049
left a comment
Owner
There was a problem hiding this comment.
File-level review comments (1/2)
每条 comment 罗列该文件的开发目标及对应修改。
| @@ -1069,6 +1069,23 @@ def forward( | |||
| value = value.squeeze(1) | |||
Owner
There was a problem hiding this comment.
开发目标 A：OFF 路径插桩 pre-RoPE Q/K/V + hidden_states dump（post-squeeze、pre-rotary）
dump_rope_preqk_verl080(layer_number, query, key, num_layers)dump_build_kv_input_v_on(layer_number, value, num_layers)dump_hidden_states_on(layer_number, hidden_states, num_layers)
开发目标 B：OFF 路径插桩 post-RoPE Q/K + full KV dump（rotary 后、core attention 前）
dump_rope_postqk_verl080(layer_number, query, key, num_layers)dump_full_kv_off(layer_number, key, value, num_layers)
两者由 PREFIX_SHARING_DIAG_DUMP env var 控制，未设时仅一次 os.environ.get 检查。try/except 保护 dump 失败不崩训练。
| @@ -1354,8 +1354,26 @@ def fit(self): | |||
| max_prompt_length=self.config.data.max_prompt_length, | |||
Owner
There was a problem hiding this comment.
开发目标：在 fit() 训练循环前注入 synthetic baseline 数据
- 第 1354-1374 行：读
PREFIX_SHARING_BASELINE_SYNTHETIC/NUM_SEQ/STACK/SEEDenv var - 调用
inject_baseline_synthetic.patch_baseline_synthetic()替换 dataloader - 用于无真实数据加载器的 CI 环境
| from prefix_sharing.backends.base import BackendCapabilities, PrefixAttentionBackend, PrefixDeltanetBackend | ||
| from prefix_sharing.backends.block_causal_mask import build_block_causal_mask | ||
| from prefix_sharing.backends.factory import get_backend_instance | ||
| from prefix_sharing.backends.flash_atten_base import FlashAttentionMixin, FlashBackendValidationError |
Owner
There was a problem hiding this comment.
开发目标：导出新的 NPU TND varlen + sparse_mode=3 backend
- 导入
NpuFlashAttentionTndBackend（第 89 行） - 加入
__all__（第 97 行）
| @@ -0,0 +1,257 @@ | |||
| """Ascend NPU Flash Attention backend — TND varlen + sparse_mode=3 (rightDownCausal). | |||
Owner
There was a problem hiding this comment.
开发目标：实现 NPU TND varlen + sparse_mode=3 attention 后端
替代旧 BSH + per-sample B1SS mask + sparse_mode=1 方案（reuser 上结果错误）。利用 TND varlen + sparse_mode=3 rightDownCausal 的"Q<KV 右对齐"语义，一次调用覆盖 provider（标准 causal）和 reuser（prefix 全可见 + suffix causal）。
_compressed_causal_mask()：全局缓存的 [2048,2048] 下三角 maskapply_rope(query, key, plan, ...)：按 position_offsets 逐 token 应用 RoPEbuild_kv(key, value, store, plan, ...)：与 TorchRef 相同的 expanded KV 语义attention(query, ek, ev, plan, ...)：input_layout="TND"、sparse_mode=3、actual_seq_qlen/kvlen=cu_seqlens
| @@ -0,0 +1,468 @@ | |||
| """TND (varlen) + sparse_mode=3 (rightDownCausal) 可行性实验（NPU only）。 | |||
Owner
There was a problem hiding this comment.
开发目标：5-probe 验证矩阵系统性验证 TND varlen + sparse_mode=3 的可行性
- Probe A (provider-only mode 2)：baseline sanity——接线正确验证
- Probe B (single reuser mode 3)：核心验证——mode 3 rightDownCausal 在 varlen 下对 Q<KV 真的右对齐
- Probe C (full batch mode 3)：单次 TND 调用同时处理 provider+reuser（vs 退路：拆两次调用）
- Probe D (full batch backward)：128-tile 约束下 gradient 对齐
- Probe E (seq>2048)：边界验证——压缩 mask [2048,2048] 下超长序列兼容性
各 probe 输出 {name, pass, detail, out_diff, grad_diff} + 决策树建议。
| """ | ||
|
|
||
| if not self.enable_prefix_sharing: | ||
| return |
Owner
There was a problem hiding this comment.
开发目标：将 flash_atten_npu_tnd 加入支持的 backends 列表
_SUPPORTED_BACKENDS新增"flash_atten_npu_tnd"（与 factory 注册同步）
| @@ -106,6 +106,15 @@ def prefix_attention( | |||
| f"built expanded kv: expanded_key_shape={tuple(expanded_key.shape)}, expanded_value_shape={tuple(expanded_value.shape)}" | |||
Owner
There was a problem hiding this comment.
开发目标 A：修复 PP dump 中 attention output 永不 flush 的 bug（核心 bug 修复）
旧逻辑用 layer_number == num_layers 判断最后层——PP stage 最大 layer_number < num_layers，永远不 flush。
- 新增
_last_layer_for_stage(layer_number, num_layers, pp_size)：按 PP partition 计算当前 stage 的最大层号 - flush 条件改为
layer_number == _last_layer_for_stage(…) - pp_size=1 时退化为旧逻辑，完全兼容
开发目标 B：新增 SKIP_RESTORE env var 诊断 probe
- restore 前检查
PREFIX_SHARING_SKIP_RESTORE：设值时跳过 restore，用于隔离"分歧在 restore 前/时" - 添加
[PS-diag]打印警告
| @@ -35,7 +35,6 @@ | |||
| from prefix_sharing.integrations.parallel_info import get_megatron_parallel_info | |||
Owner
There was a problem hiding this comment.
开发目标 A：迁移 NestedTensor API 实现 NPU autograd 兼容
nested_tensor(seqs, layout=torch.jagged)→as_nested_tensor(rows, layout=torch.jagged)（~6 处）- NPU ascend 不支持
nested_tensor的 autograd
开发目标 B：新增 SKIP_RESTORE 诊断 probe
- restore 路径前加 env var 检测分支
Jackie2049
reviewed
Jul 16, 2026
Jackie2049
left a comment
Jackie2049
left a comment
Owner
There was a problem hiding this comment.
File-level review comments (2/2)
每条 comment 罗列该文件的开发目标及对应修改。
| @@ -36,7 +36,11 @@ def patched_forward( | |||
| ctx = current_prefix_sharing_context() | |||
Owner
There was a problem hiding this comment.
开发目标 A：解耦 per-layer diagnostic dump 检测，支持 PP 感知的 flush
- 新增
_compute_stage_last_layer辅助函数 - 提取
_dump_attn_output薄委托层从_FSDP_ATTN_BUFFER分离 - 新增 expanded_kv / full_kv / attn_inputs 的 per-layer dump 委托点
开发目标 B：添加 layer_index/config 诊断日志
- patch 启动时打印
[PS-diag] num_hidden_layers={n}, layer_idx={idx}
| @@ -122,57 +122,58 @@ def patched_forward_step( | |||
| # ON: prefix_lens / original_lengths 取自 plan； | |||
Owner
There was a problem hiding this comment.
开发目标：在 forward_step ON/OFF 路径衔接处新增 per-layer diagnostic dump 接线
- dump_meta_verl080() / dump_attention_mask_verl080() / dump_label_mask_verl080() 于 build_micro_batch 前后
- dump_fsdp_model_output_2d_verl080() 于 restore 前后
- dump_logprobs_2d_verl080() / dump_entropy_2d_verl080() 于 loss 前后
- import 路径同步更新
| @@ -26,9 +26,9 @@ def patched_fn(logits, labels): | |||
| # ##### [PS-diag] dump logits（ON/OFF 都 dump，必须在 original_fn 之前） ##### | |||
Owner
There was a problem hiding this comment.
开发目标：同步诊断 dump 的 logits 处理路径
- logits 前后嵌入
dump_logits_verl080(logits)/dump_logprobs_verl080(logits, labels) - 对称包装 ON/OFF 两端 logits dump
| @@ -0,0 +1,196 @@ | |||
| """Multi-rank diagnostic dump assembler — merges PP/TP-sharded files into a | |||
Owner
There was a problem hiding this comment.
开发目标：多卡（TP/PP）诊断 dump → 单卡格式离线拼装
- TP 拼装：按
_tp{r}.pt后缀识别 shard，vocab 维拼接 logits - PP 拼装：按
_pp{p}.pt后缀合并层 dict - 单卡 fast path：TP/PP=1 直接 copy
| @@ -0,0 +1,679 @@ | |||
| """GEMM Precision Baseline — Cross-Batch-Size Comparison. | |||
Owner
There was a problem hiding this comment.
开发目标："跨 batch"精度基线——测量同一配置下两次独立运行之间的自然噪声底
GPU GEMM 非确定性导致相同输入多次 forward 有微小差异。此差值就是 ON/OFF 对比时应排除的噪声底。
_compare_attn_cross_batch()：per-layer Q/K cosine + max_diff_compare_logits_cross_batch()：packed logits 按 cu_seqlens 切片对比_compare_2d_cross_batch()：2D tensor 逐元素对比 + pearson- CLI：
--dir-single --dir-stacked
| @@ -7,7 +7,7 @@ | |||
|
|
|||
Owner
There was a problem hiding this comment.
开发目标：v080 精度对比引擎全面增强——3 大维度
RoPE 对比体系：cmp_rope_postqk_layer/token()（pre/post 两阶段）、cmp_rope_freqs()、_rope_postqk_vec_at_pos()、_diag_rope_pos_fail()
KV 对比体系：cmp_attn_kv()（expanded vs full）、cmp_build_kv_input_v()、_print_attn_kv()
Suffix 对齐统一：删除 2D attention_mask 对齐（_build_alignment_mask_from_2d、_load_attention_mask_2d），全改为 prefix_lens + cu_seqlens
代码清理：_first_token_metrics→_vec_metrics 重命名、变量名规范化、删除旧 logging
| @@ -10,7 +10,7 @@ | |||
|
|
|||
Owner
There was a problem hiding this comment.
开发目标：新增 OFF 侧诊断 dump 函数 + 统一 attention output buffer 管理
- 新增
dump_fsdp_attn_output(output, module)：_FSDP_ATTN_BUFFER buffer（layer1 清空 / last layer flush） - 新增
dump_rope_pre_post_qk(rope_dict_on, rope_dict_off, ...) - 新增
_rank0_only()门控：多 rank 下仅 rank 0 存盘 - 新增
_get_dump_dir()：从 env var 统一获取 dump 目录
| @@ -22,6 +22,7 @@ | |||
|
|
|||
Owner
There was a problem hiding this comment.
开发目标：新增 6 个 OFF 侧侵入式 dump 函数
dump_rope_preqk_verl080(layer_number, query, key, num_layers)→ rope_preqk.ptdump_rope_postqk_verl080(layer_number, query, key, num_layers)→ rope_postqk.ptdump_full_kv_off(layer_number, key, value, num_layers)→ full_kv.ptdump_build_kv_input_v_on(layer_number, value, num_layers)→ build_kv_input_v.ptdump_hidden_states_on(layer_number, hidden_states, num_layers)→ hidden_states.pt
每函数带 try/except 保护、_rank0_only() + _FSDP_ATTN_BUFFER 协议
| @@ -0,0 +1,115 @@ | |||
| """Inject synthetic data for GEMM precision baseline — reuse nested-prefix | |||
Owner
There was a problem hiding this comment.
开发目标：生成合成 baseline 数据，供 CI 环境验证精度比较工具链
patch_baseline_synthetic(trainer, json_path, max_prompt_length, max_response_length, num_seq, stack, seed)：- 构造填充后 input_ids / attention_mask / sequences 等 tensor
stack参数控制数据规模（相同数据重复 K 次）- 挂载到
trainer.train_dataloader替换真实加载器
|
|
||
| import pytest | ||
|
|
||
| from prefix_sharing.backends.factory import get_backend_instance |
Owner
There was a problem hiding this comment.
开发目标：为 flash_atten_npu_tnd 后端添加 factory 单元测试
- 新增
test_factory_flash_atten_npu_tnd()：验证后端创建及其 build_kv/attention 签名 - skip 门控：仅在 flash_atten_npu_tnd importable 时运行
…est backends Delete flash_atten_npu_tnd.py and flash_atten_npu_test.py, and clean up all references in __init__.py, factory.py, config.py, and tests. Co-Authored-By: Claude <noreply@anthropic.com>
📊 Coverage reportChanged files in this PR |
Jackie2049
approved these changes
Jul 16, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Build numerical precision verification infrastructure for prefix-sharing Phase 1, covering per-tensor to per-layer, single-card to TP/PP ON/OFF comparison. Fix PP dump silent loss in attention hooks.
Key Changes
nested_tensor→as_nested_tensorfor NPU autograd compat; add SKIP_RESTORE diagnostic probeTest Results
Remaining Issues
assemble_dump.py), not yet integrated with training pipeline