69 lines
2.9 KiB
Python
69 lines
2.9 KiB
Python
"""Capture Qwen decoder layer-0 intermediates for direct parity debugging."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
|
|
source = path.read_text(encoding="utf-8")
|
|
if "import os\n" not in source:
|
|
source = "import os\n" + source
|
|
old = (
|
|
" # Self Attention\n"
|
|
" residual = x\n"
|
|
" x = self.input_layernorm(x)\n"
|
|
" x, present_key_value = self.self_attn(\n"
|
|
" hidden_states=x,\n"
|
|
" attention_mask=attention_mask,\n"
|
|
" freqs_cis=freqs_cis,\n"
|
|
" optimized_attention=optimized_attention,\n"
|
|
" past_key_value=past_key_value,\n"
|
|
" )\n"
|
|
" x = residual + x\n"
|
|
"\n"
|
|
" # MLP\n"
|
|
" residual = x\n"
|
|
" x = self.post_attention_layernorm(x)\n"
|
|
" x = self.mlp(x)\n"
|
|
" x = residual + x\n"
|
|
)
|
|
new = (
|
|
" # Self Attention\n"
|
|
" residual = x\n"
|
|
" x = self.input_layernorm(x)\n"
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
|
" if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"qwen0_norm1.pt\"))\n"
|
|
" attention, present_key_value = self.self_attn(\n"
|
|
" hidden_states=x,\n"
|
|
" attention_mask=attention_mask,\n"
|
|
" freqs_cis=freqs_cis,\n"
|
|
" optimized_attention=optimized_attention,\n"
|
|
" past_key_value=past_key_value,\n"
|
|
" )\n"
|
|
" if capture_dir: torch.save(attention.detach().cpu(), os.path.join(capture_dir, \"qwen0_attention.pt\"))\n"
|
|
" x = residual + attention\n"
|
|
" if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"qwen0_post_attention.pt\"))\n"
|
|
"\n"
|
|
" # MLP\n"
|
|
" residual = x\n"
|
|
" x = self.post_attention_layernorm(x)\n"
|
|
" if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"qwen0_norm2.pt\"))\n"
|
|
" mlp = self.mlp(x)\n"
|
|
" if capture_dir: torch.save(mlp.detach().cpu(), os.path.join(capture_dir, \"qwen0_mlp.pt\"))\n"
|
|
" x = residual + mlp\n"
|
|
" if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"qwen0_output.pt\"))\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate Qwen TransformerBlock.forward.")
|
|
source = source.replace(old, new)
|
|
|
|
old = " for i, layer in enumerate(self.layers):\n"
|
|
new = (
|
|
" for i, layer in enumerate(self.layers):\n"
|
|
" layer._h3_trace_index = i\n"
|
|
" layer.self_attn._h3_trace_index = i\n"
|
|
" layer.mlp._h3_trace_index = i\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate Qwen decoder layer loop.")
|
|
path.write_text(source.replace(old, new), encoding="utf-8")
|
|
print("Applied Qwen layer-0 sublayer capture patch.")
|