{"id":91537,"date":"2026-08-07T23:16:23","date_gmt":"2026-08-07T15:16:23","guid":{"rendered":"https:\/\/www.wsisp.com\/helps\/91537.html"},"modified":"2026-08-07T23:16:23","modified_gmt":"2026-08-07T15:16:23","slug":"6-light-wam-%e6%a8%a1%e5%9e%8b%e4%b8%ad-ditblock-%e6%a8%a1%e5%9d%97","status":"publish","type":"post","link":"https:\/\/www.wsisp.com\/helps\/91537.html","title":{"rendered":"6. light wam \u6a21\u578b\u4e2d DITBlock \u6a21\u5757"},"content":{"rendered":"

1. DITBlock \u6d4b\u8bd5\u811a\u672c<\/h3>\n

import<\/span> torch
\nfrom<\/span> lightwam.<\/span>models.<\/span>wan22.<\/span>wan_video_dit import<\/span> DiTBlock<\/p>\n

def<\/span> test_dit_block_forward<\/span>(<\/span>)<\/span>:<\/span>
\n print<\/span>(<\/span>"="<\/span>*<\/span>50<\/span>)<\/span>
\n print<\/span>(<\/span>"Testing DiTBlock Forward Pass"<\/span>)<\/span>
\n print<\/span>(<\/span>"="<\/span>*<\/span>50<\/span>)<\/span><\/p>\n

# 1. Define model hyperparameters<\/span>
\n batch_size =<\/span> 2<\/span>
\n seq_len =<\/span> 294<\/span>
\n context_len =<\/span> 129<\/span>
\n hidden_dim =<\/span> 1536<\/span>
\n attn_head_dim =<\/span> 128<\/span>
\n num_heads =<\/span> 12<\/span>
\n ffn_dim =<\/span> 8960<\/span><\/p>\n

device =<\/span> torch.<\/span>device(<\/span>"cuda"<\/span> if<\/span> torch.<\/span>cuda.<\/span>is_available(<\/span>)<\/span> else<\/span> "cpu"<\/span>)<\/span>
\n dtype =<\/span> torch.<\/span>bfloat16 if<\/span> torch.<\/span>cuda.<\/span>is_available(<\/span>)<\/span> else<\/span> torch.<\/span>float32<\/p>\n

print<\/span>(<\/span>f"Device: <\/span>{<\/span>device}<\/span><\/span>"<\/span><\/span>)<\/span>
\n print<\/span>(<\/span>f"Dtype: <\/span>{<\/span>dtype}<\/span><\/span>"<\/span><\/span>)<\/span><\/p>\n

# 2. Instantiate the DiTBlock<\/span>
\n print<\/span>(<\/span>"\\\\nInitializing DiTBlock…"<\/span>)<\/span>
\n block =<\/span> DiTBlock(<\/span>
\n hidden_dim=<\/span>hidden_dim,<\/span>
\n attn_head_dim=<\/span>attn_head_dim,<\/span>
\n num_heads=<\/span>num_heads,<\/span>
\n ffn_dim=<\/span>ffn_dim,<\/span>
\n eps=<\/span>1e-6<\/span>
\n )<\/span>.<\/span>to(<\/span>device=<\/span>device,<\/span> dtype=<\/span>dtype)<\/span><\/p>\n

# Enable eval mode for deterministic testing<\/span>
\n block.<\/span>eval<\/span>(<\/span>)<\/span>
\n print<\/span>(<\/span>"Initialization complete."<\/span>)<\/span><\/p>\n

# 3. Create dummy input tensors matching the training flow shapes<\/span>
\n print<\/span>(<\/span>"\\\\nCreating dummy input tensors…"<\/span>)<\/span>
\n # x_tokens: (B, seq_len, hidden_dim) -> e.g., (2, 294, 1536)<\/span>
\n x =<\/span> torch.<\/span>randn(<\/span>batch_size,<\/span> seq_len,<\/span> hidden_dim,<\/span> device=<\/span>device,<\/span> dtype=<\/span>dtype)<\/span><\/p>\n

# context_emb: (B, context_len, hidden_dim) -> e.g., (2, 129, 1536)<\/span>
\n context =<\/span> torch.<\/span>randn(<\/span>batch_size,<\/span> context_len,<\/span> hidden_dim,<\/span> device=<\/span>device,<\/span> dtype=<\/span>dtype)<\/span><\/p>\n

# t_mod: (B, seq_len, 6, hidden_dim) -> e.g., (2, 294, 6, 1536)<\/span>
\n t_mod =<\/span> torch.<\/span>randn(<\/span>batch_size,<\/span> seq_len,<\/span> 6<\/span>,<\/span> hidden_dim,<\/span> device=<\/span>device,<\/span> dtype=<\/span>dtype)<\/span><\/p>\n

# freqs (RoPE): (seq_len, 1, 64) -> e.g., (294, 1, 64) <\/span>
\n # (Assuming attn_head_dim=128, freq dim usually is head_dim \/\/ 2 = 64)<\/span>
\n freq_dim =<\/span> attn_head_dim \/\/<\/span> 2<\/span>
\n freqs =<\/span> torch.<\/span>randn(<\/span>seq_len,<\/span> 1<\/span>,<\/span> freq_dim,<\/span> device=<\/span>device,<\/span> dtype=<\/span>dtype)<\/span><\/p>\n

# context_mask: (B, seq_len, context_len) -> e.g., (2, 294, 129)<\/span>
\n context_mask =<\/span> torch.<\/span>ones(<\/span>batch_size,<\/span> seq_len,<\/span> context_len,<\/span> device=<\/span>device,<\/span> dtype=<\/span>torch.<\/span>bool<\/span>)<\/span><\/p>\n

# self_attn_mask (Optional): (seq_len, seq_len)<\/span>
\n self_attn_mask =<\/span> None<\/span> <\/p>\n

print<\/span>(<\/span>f" – x: <\/span>{<\/span>tuple<\/span>(<\/span>x.<\/span>shape)<\/span>}<\/span><\/span>"<\/span><\/span>)<\/span>
\n print<\/span>(<\/span>f" – context: <\/span>{<\/span>tuple<\/span>(<\/span>context.<\/span>shape)<\/span>}<\/span><\/span>"<\/span><\/span>)<\/span>
\n print<\/span>(<\/span>f" – t_mod: <\/span>{<\/span>tuple<\/span>(<\/span>t_mod.<\/span>shape)<\/span>}<\/span><\/span>"<\/span><\/span>)<\/span>
\n print<\/span>(<\/span>f" – freqs: <\/span>{<\/span>tuple<\/span>(<\/span>freqs.<\/span>shape)<\/span>}<\/span><\/span>"<\/span><\/span>)<\/span>
\n print<\/span>(<\/span>f" – context_mask: <\/span>{<\/span>tuple<\/span>(<\/span>context_mask.<\/span>shape)<\/span>}<\/span><\/span>"<\/span><\/span>)<\/span><\/p>\n

# 4. Run Forward Pass<\/span>
\n print<\/span>(<\/span>"\\\\nRunning forward pass…"<\/span>)<\/span>
\n with<\/span> torch.<\/span>no_grad(<\/span>)<\/span>:<\/span>
\n out =<\/span> block(<\/span>
\n x=<\/span>x,<\/span>
\n context=<\/span>context,<\/span>
\n t_mod=<\/span>t_mod,<\/span>
\n freqs=<\/span>freqs,<\/span>
\n context_mask=<\/span>context_mask,<\/span>
\n self_attn_mask=<\/span>self_attn_mask
\n )<\/span><\/p>\n

# 5. Verify Output<\/span>
\n print<\/span>(<\/span>f"\\\\nOutput shape: <\/span>{<\/span>tuple<\/span>(<\/span>out.<\/span>shape)<\/span>}<\/span><\/span>"<\/span><\/span>)<\/span>
\n assert<\/span> out.<\/span>shape ==<\/span> x.<\/span>shape,<\/span> f"Shape mismatch! Expected <\/span>{<\/span>x.<\/span>shape}<\/span><\/span>, got <\/span>{<\/span>out.<\/span>shape}<\/span><\/span>"<\/span><\/span>
\n print<\/span>(<\/span>"Test passed! Input and output shapes match."<\/span>)<\/span><\/p>\n

if<\/span> __name__ ==<\/span> "__main__"<\/span>:<\/span>
\n test_dit_block_forward(<\/span>)<\/span><\/p>\n