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-memoryconda 环境需验证 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 50smoke,看 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 的五步全部可行,已对现码核实;三个注意点:
- 行号已过时:§2d 说的"trainer init ~984"现对应
training_modules/wan_training_module.py(镜像use_moc的传参与pipe.*赋值位置);eager 建头必须在train.py:620的 AdamW 构造之前。 - 只接
context/condition两分支即可(two-chunk 全走 context);standard 分支变量名不同,不必动。 - 头的 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 定义一致。