33 lines
1.6 KiB
Python
33 lines
1.6 KiB
Python
|
"""Capture Comfy Qwen layer-0 projection boundaries for direct parity checks."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
|
||
|
|
source = path.read_text(encoding="utf-8")
|
||
|
|
old = " xq = self.q_proj(hidden_states)\n xk = self.k_proj(hidden_states)\n xv = self.v_proj(hidden_states)\n"
|
||
|
|
new = old + (
|
||
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
||
|
|
" if capture_dir:\n"
|
||
|
|
" torch.save({\"q\": xq.detach().cpu(), \"k\": xk.detach().cpu(), \"v\": xv.detach().cpu()}, os.path.join(capture_dir, \"qwen0_qkv.pt\"))\n"
|
||
|
|
)
|
||
|
|
if source.count(old) != 1:
|
||
|
|
raise RuntimeError("Unable to locate Qwen QKV projections.")
|
||
|
|
source = source.replace(old, new)
|
||
|
|
|
||
|
|
old = " return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
|
||
|
|
new = (
|
||
|
|
" gate = self.gate_proj(x)\n"
|
||
|
|
" up = self.up_proj(x)\n"
|
||
|
|
" activated = self.activation(gate) * up\n"
|
||
|
|
" output = self.down_proj(activated)\n"
|
||
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
|
||
|
|
" if capture_dir:\n"
|
||
|
|
" torch.save({\"gate\": gate.detach().cpu(), \"up\": up.detach().cpu(), \"activated\": activated.detach().cpu(), \"down\": output.detach().cpu()}, os.path.join(capture_dir, \"qwen0_mlp_projections.pt\"))\n"
|
||
|
|
" return output\n"
|
||
|
|
)
|
||
|
|
if source.count(old) != 1:
|
||
|
|
raise RuntimeError("Unable to locate Qwen MLP projections.")
|
||
|
|
path.write_text(source.replace(old, new), encoding="utf-8")
|
||
|
|
print("Applied Qwen layer-0 projection capture patch.")
|