| # CGLA Wan 移植审计(cgla_audit.md,2026-07-22) |
| |
| 背景:Echo-Memory 的 CGLA / PRoPE / UCPE 三个消融行,是 DFOT 侧原型(`/data1/dfot_new`,RealEstate10k 上正在训 `uvit_cgla_poserope` / `uvit_cgla_gatemod`)向 Wan2.1-1.3B two-chunk 框架的移植。本审计对照了三方代码:`diffsynth/models/memory/u_vit_cgla_blocks.py`(移植)、`flash-linear-attention/fla/layers/sse.py`(共享内核)、`dfot_new/algorithms/dfot/backbones/u_vit/u_vit_cgla_blocks.py`(原型)。 |
|
|
| **总结论:移植当前无法启动训练(3 处必崩),另有 2 处静默失效、3 处相对原型的语义/结构偏离。** |
|
|
| ## 1. 必崩(按触发顺序) |
|
|
| | # | 位置 | 问题 | 修法 | |
| |---|---|---|---| |
| | 1 | `train/memory_baselines_basic/run_ablation_prope_memory_two_chunk.sh:21` | 不传 `--add_action_attn`,撞上 `u_vit_cgla_blocks.py:200` 的 `assert add_action_attn` → AssertionError | 启动器补 flag(cgla/ucpe 两个启动器都传了),或放宽 assert | |
| | 2 | `u_vit_cgla_blocks.py:204` | `self.self_attn_with_action.o_proj.weight` — Wan `SelfAttention` 的输出投影叫 `.o`(`wan_video_dit.py:134`),无 `.o_proj` → AttributeError,cgla/ucpe 行构造 block 即崩 | 改成 `.o.weight`,并对照 `training_modules/action_blocks.py:58-59` 顺带把 `.o.bias` 也置零 | |
| | 3 | `src/model_training/train.py:547-548` | CGLA 冻结循环的行继续后少了 `or`:`("self_attn" in name) \ ("cam_encoder" in name)` 解析为"布尔值调用" → TypeError | 补上 `or` | |
|
|
| ## 2. 静默失效(跑起来也不对) |
|
|
| | # | 位置 | 问题 | 修法 | |
| |---|---|---|---| |
| | 4 | `train.py:189`(只 add_argument) | `--cgla_aux_loss_weight 0.01` **从未赋到 pipe**;消费端 `wan_video_new.py:629` 和 `training_loop.py:373` 都是 `getattr(pipe, "cgla_aux_loss_weight", 0.0)` → SSE 负载均衡辅助损失恒为 0(DFOT 侧同名权重 0.01 生效中) | train.py 构造 model 后 `model.pipe.cgla_aux_loss_weight = float(_arg('cgla_aux_loss_weight', 0.0))` | |
| | 5 | `u_vit_cgla_blocks.py:210-228` | `--cgla_gate_logit_normalizer/--cgla_gate_low_rank_dim` 是 **SSEGLA 的参数**,传给 SSEGDN 被 `**kwargs` 吞掉,无任何效果;`--cgla_pose_dim` 亦被 `self.pose_dim = int(dim)`(:187)覆盖失效(pose 实为 action_mlp 输出的 D 维嵌入,自洽但与 flag 语义不符) | 见 #6 的家族决策;pose_dim 覆盖行为写进 flag help 或删 flag | |
| |
| ## 3. 相对 DFOT 原型的关键偏离(设计审查) |
| |
| | # | 主题 | DFOT 原型 | Wan 移植 | 建议 | |
| |---|---|---|---|---| |
| | 6 | 注意力家族 | 三种 block 全部 `SSEGLA`(`u_vit_cgla_blocks.py:15,189,307-308`),正在训的两个 run 也是 | `SSEGDN`(delta-rule 家族,`:210`),注释却写 SSEGLA | 受控移植应换回 SSEGLA;若有意升级 GDN,需文档言明并同步 DFOT 侧对照 | |
| | 7 | write_gate 输入 | `noise_write_gate(emb)` = **噪声等级嵌入**(:209-234;零权重 bias+5 → sigmoid≈0.993),语义:噪声大时少写入记忆 | 输入换成了 token 特征 `x`(:236-238,303),gate 无法按噪声水平调控 | Wan 侧有 `t`(timestep embedding)可用:把 write_gate 输入换成 t 的投影(broadcast 到 token) | |
| | 8 | step-0 恒等性 | 全参训练 + 零初始化注入 ≈ 恒等起步 | 注释承诺 `tanh(cgla_gate)` 零门(:139-142)**但 cgla_gate 根本不存在**,实际走 AdaLN 的 gate_msa(非零)→ step-0 输出大幅偏离预训练 Wan,与"冻结骨干只训新模块"的配方自相矛盾 | 给 CGLA 残差加真正的零初始化标量门(即把注释里的设计实现出来)——这是训练稳定性的关键 | |
| | 9 | 双向 | `BiJoint` 用两套独立权重 `attn_fwd/attn_bwd`(:307-308) | 单模块对翻转序列跑两遍取均值 | 容量减半 + 每步计算×2;可接受但应写明,预算按前向×2 计 | |
| | 10 | 替换范围 | 只替换 bottleneck mid 层,且 `softmax_mid_layers=[0,5,10,15]` 保留 softmax(混合) | `cgla_every_n_blocks=4` 全深度均匀替换(prope/ucpe 启动器还是 every_n=1 全替换) | 模式不同本身可以,但消融叙事要自洽;建议至少保一行 every_n=4 与 DFOT 混合模式对齐 | |
| |
| ## 4. 次级 / 健壮性 |
| |
| - `u_vit_cgla_blocks.py` 缺 `import torch.nn.functional as F` → `MLP_Action`(use_cam_pose=False 路径)NameError;现启动器都传 `--train_cam_pose` 暂不触发,但属地雷。 |
| - forward 对 `actions=None` 无守卫(:334 直接 `.to`)→ 无位姿调用即崩(例如不带 action_path 的推理);`DiTBlock_w_Action` 有守卫,应对齐。 |
| - `_pose_emb_from_actions`(:272-297)是死代码;`pose_bottleneck` 被注释(:223);文档写 SSEGLA 实现是 SSEGDN,名实不符。 |
| - Wan `self_attn.{q,k,v,o}` 的 **bias 无法载入**(SSEGDN 的 `*_proj` 均 bias=False)→ remap 后成为 unexpected 静默丢弃;`norm_q/norm_k` 同。对 1.3B 影响可控,但应知情。 |
| - 环境:chunk kernel 依赖 triton(vendored fla 在仓库根,`u_vit_cgla_blocks.py:49-54` 自动挂 sys.path);`echo-memory` conda 环境需验证 triton 可用。 |
| - 推理侧 `env/loop_utils.py` 的 CGLA 检测/重建/remap(:234-270,:365-388)是通的;训练侧修完后两侧键名兼容。 |
| |
| ## 5. 建议修复顺序(最小 patch 集) |
| |
| - **P0(恢复可启动,约 10 行)**:#1 #2 #3 + #4(aux 接线)。 |
| - **P1(训练质量)**:#8 零初始化门(稳定性关键)→ #7 write_gate 换 t 输入。 |
| - **P2(方法论)**:#6 家族决策(建议换回 SSEGLA 对齐正在训的 DFOT 原型)+ #12 None 守卫 + F import。 |
| - 验证路径:`tests/test_cgla_wan.py`(权重覆盖率,GPU)→ GPU5 上 `--max_train_steps 50` smoke,看 step-0 loss 是否与 no_memory 行同量级(#8 修好后应当接近)。 |
| |
| ## 6. 附:DFOT 侧变体矩阵(移植选型参考) |
| |
| 8 个配置仅切换 backbone 变体;公共超参全走代码默认(`num_sparse_partition=4, num_writer=num_reader=1, gate_low_rank_dim=16, gate_logit_normalizer=16, pose_bottleneck=64, head_dim=bottleneck_ch//num_heads`),`aux_loss_weight=0.01`,位姿条件 = ray_encoding 相对首帧: |
| |
| | 配置 | 块 | pose_rope | pose_gate_mod | |
| |---|---|---|---| |
| | cgla | Joint | – | – | |
| | bidi | BiJoint | – | – | |
| | poserope(在训) | Joint | ✓ | – | |
| | bidi_poserope | BiJoint | ✓ | – | |
| | gatemod(在训) | Joint | – | ✓ | |
| | full / bidi_full | Joint / BiJoint | ✓ | ✓ | |
| | fact | 帧内 softmax + 帧间 SSE(pose_rope/gatemod 在此变体**静默无效**) | – | – | |
| |
| 要点:即便 base cgla,位姿也**始终注入**稀疏流(q2/k2/g2/η,零初始化上投影);poserope 是叠加的相对位姿 RoPE(两条流都转);gatemod 是遗忘门乘性调制(`1+tanh(Linear(pose))`,step-0 起就有梯度)。DFOT 侧**不冻结任何参数**。 |
| |
| ## 7. retrieve.md(不确定性损失)接线核验 |
| |
| §2 的五步全部可行,已对现码核实;三个注意点: |
| |
| 1. 行号已过时:§2d 说的"trainer __init__ ~984"现对应 `training_modules/wan_training_module.py`(镜像 `use_moc` 的传参与 `pipe.*` 赋值位置);eager 建头必须在 `train.py:620` 的 AdamW 构造**之前**。 |
| 2. 只接 `context`/`condition` 两分支即可(two-chunk 全走 context);standard 分支变量名不同,不必动。 |
| 3. 头的 ckpt 键带 `pipe.uncertainty_head.` 前缀:`trainable_modules()/export_trainable_state_dict`(`diffsynth/trainers/utils.py:1214,1233`)已确认会收进优化器并存进 ckpt;推理加载时 unexpected,无害。 |
| |
| `uncertainty.py` 模块本身质量合格:零初始化恒等、um clamp 防 bf16 溢出、`x0 = x_t − σ·v` 与 FlowMatchScheduler 定义一致。 |
| |
