6. light wam 模型中 DITBlock 模块

1. DITBlock 测试脚本

import torch
from lightwam.models.wan22.wan_video_dit import DiTBlock

def test_dit_block_forward():
print("="*50)
print("Testing DiTBlock Forward Pass")
print("="*50)

# 1. Define model hyperparameters
batch_size = 2
seq_len = 294
context_len = 129
hidden_dim = 1536
attn_head_dim = 128
num_heads = 12
ffn_dim = 8960

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

print(f"Device: {device}")
print(f"Dtype: {dtype}")

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

# Enable eval mode for deterministic testing
block.eval()
print("Initialization complete.")

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

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

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

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

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

# self_attn_mask (Optional): (seq_len, seq_len)
self_attn_mask = None

print(f" – x: {tuple(x.shape)}")
print(f" – context: {tuple(context.shape)}")
print(f" – t_mod: {tuple(t_mod.shape)}")
print(f" – freqs: {tuple(freqs.shape)}")
print(f" – context_mask: {tuple(context_mask.shape)}")

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

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

if __name__ == "__main__":
test_dit_block_forward()

• 测试结果

python test_dit_block.py
==================================================
Testing DiTBlock Forward Pass
==================================================
Device: cuda
Dtype: torch.bfloat16

Initializing DiTBlock...
Initialization complete.

Creating dummy input tensors...
– x: (2, 294, 1536)
– context: (2, 129, 1536)
– t_mod: (2, 294, 6, 1536)
– freqs: (294, 1, 64)
– context_mask: (2, 294, 129)

Running forward pass...

Output shape: (2, 294, 1536)
Test passed! Input and output shapes match.

2.前向过程

• 代码

def forward(self, x, context, t_mod, freqs, context_mask=None, self_attn_mask: Optional[torch.Tensor] = None):
if context_mask is not None and context_mask.dim() == 3:
context_mask = context_mask.unsqueeze(1) # (B, 1, seq_len, context_len), 1 for heads
has_seq = len(t_mod.shape) == 4
chunk_dim = 2 if has_seq else 1
# msa: multi-head self-attention mlp: multi-layer perceptron
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim)
if has_seq:
# means t_mod has separate modulation for each token, otherwise same modulation for all tokens in the block
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2),
shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2),
)
input_x = modulate(self.norm1(x), shift_msa, scale_msa)
x = self.gate(x, gate_msa, self.self_attn(input_x, freqs, self_attn_mask=self_attn_mask))
x = x + self.cross_attn(self.norm3(x), context, ctx_mask=context_mask)
input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
ffn_hidden = self.ffn[0](input_x)
if self.ffn_lora_in is not None:
ffn_hidden = ffn_hidden + self.ffn_lora_in(input_x)
ffn_hidden = self.ffn[1](ffn_hidden)
ffn_out = self.ffn[2](ffn_hidden)
if self.ffn_lora_out is not None:
ffn_out = ffn_out + self.ffn_lora_out(ffn_hidden)
x = self.gate(x, gate_mlp, ffn_out)
return x

• 流程图

#mermaid-svg-3pQewxA1NIw4fmep{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-3pQewxA1NIw4fmep .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-3pQewxA1NIw4fmep .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-3pQewxA1NIw4fmep .error-icon{fill:#552222;}#mermaid-svg-3pQewxA1NIw4fmep .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-3pQewxA1NIw4fmep .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-3pQewxA1NIw4fmep .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-3pQewxA1NIw4fmep .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-3pQewxA1NIw4fmep .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-3pQewxA1NIw4fmep .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-3pQewxA1NIw4fmep .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-3pQewxA1NIw4fmep .marker{fill:#333333;stroke:#333333;}#mermaid-svg-3pQewxA1NIw4fmep .marker.cross{stroke:#333333;}#mermaid-svg-3pQewxA1NIw4fmep svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-3pQewxA1NIw4fmep p{margin:0;}#mermaid-svg-3pQewxA1NIw4fmep .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-3pQewxA1NIw4fmep .cluster-label text{fill:#333;}#mermaid-svg-3pQewxA1NIw4fmep .cluster-label span{color:#333;}#mermaid-svg-3pQewxA1NIw4fmep .cluster-label span p{background-color:transparent;}#mermaid-svg-3pQewxA1NIw4fmep .label text,#mermaid-svg-3pQewxA1NIw4fmep span{fill:#333;color:#333;}#mermaid-svg-3pQewxA1NIw4fmep .node rect,#mermaid-svg-3pQewxA1NIw4fmep .node circle,#mermaid-svg-3pQewxA1NIw4fmep .node ellipse,#mermaid-svg-3pQewxA1NIw4fmep .node polygon,#mermaid-svg-3pQewxA1NIw4fmep .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-3pQewxA1NIw4fmep .rough-node .label text,#mermaid-svg-3pQewxA1NIw4fmep .node .label text,#mermaid-svg-3pQewxA1NIw4fmep .image-shape .label,#mermaid-svg-3pQewxA1NIw4fmep .icon-shape .label{text-anchor:middle;}#mermaid-svg-3pQewxA1NIw4fmep .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-3pQewxA1NIw4fmep .rough-node .label,#mermaid-svg-3pQewxA1NIw4fmep .node .label,#mermaid-svg-3pQewxA1NIw4fmep .image-shape .label,#mermaid-svg-3pQewxA1NIw4fmep .icon-shape .label{text-align:center;}#mermaid-svg-3pQewxA1NIw4fmep .node.clickable{cursor:pointer;}#mermaid-svg-3pQewxA1NIw4fmep .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-3pQewxA1NIw4fmep .arrowheadPath{fill:#333333;}#mermaid-svg-3pQewxA1NIw4fmep .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-3pQewxA1NIw4fmep .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-3pQewxA1NIw4fmep .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-3pQewxA1NIw4fmep .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-3pQewxA1NIw4fmep .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-3pQewxA1NIw4fmep .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-3pQewxA1NIw4fmep .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-3pQewxA1NIw4fmep .cluster text{fill:#333;}#mermaid-svg-3pQewxA1NIw4fmep .cluster span{color:#333;}#mermaid-svg-3pQewxA1NIw4fmep div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-3pQewxA1NIw4fmep .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-3pQewxA1NIw4fmep rect.text{fill:none;stroke-width:0;}#mermaid-svg-3pQewxA1NIw4fmep .icon-shape,#mermaid-svg-3pQewxA1NIw4fmep .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-3pQewxA1NIw4fmep .icon-shape p,#mermaid-svg-3pQewxA1NIw4fmep .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-3pQewxA1NIw4fmep .icon-shape .label rect,#mermaid-svg-3pQewxA1NIw4fmep .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-3pQewxA1NIw4fmep .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-3pQewxA1NIw4fmep .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-3pQewxA1NIw4fmep :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

Feed Forward Network (MLP)

Cross Attention

Self Attention (MSA)

Preparation

Inputs to DiTBlock

Unsqueeze(1) for heads

x(B, seq_len, hidden_dim)

context(B, context_len, hidden_dim)

t_mod(B, seq_len, 6, hidden_dim)

freqs(RoPE position encodings)

context_mask

self_attn_mask

self.modulation + t_modChunk into 6 parts

shift_msa

scale_msa

gate_msa

shift_mlp

scale_mlp

gate_mlp

context_mask (B, 1, seq_len, ctx_len)

self.norm1(x)

modulate(…, shift_msa, scale_msa)

self.self_attn(…)

self.gate(x, gate_msa, MSA_out)

Add (Residual)

self.norm3(x)

self.cross_attn(…)

Add (Residual)

self.norm2(x)

modulate(…, shift_mlp, scale_mlp)

self.ffn[0](Linear)

self.ffn_lora_in (Optional)

Add

self.ffn[1](GELU)

self.ffn[2](Linear)

self.ffn_lora_out (Optional)

Add

self.gate(x, gate_mlp, FFN_out)

Add (Residual)

Output x

• 流程解释：