36 lines
3 KiB
Python
36 lines
3 KiB
Python
"""Capture actual loaded Comfy Qwen layer-0 projection module behavior."""
|
|
|
|
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"
|
|
" modules = {\"q\": self.q_proj, \"k\": self.k_proj, \"v\": self.v_proj}\n"
|
|
" metadata = {name: {\"class\": f\"{type(module).__module__}.{type(module).__qualname__}\", \"weight_type\": f\"{type(module.weight).__module__}.{type(module.weight).__qualname__}\", \"weight_dtype\": str(module.weight.dtype), \"weight_shape\": tuple(module.weight.shape), \"layout_type\": getattr(module, \"layout_type\", None), \"full_precision\": getattr(module, \"_full_precision_mm\", None), \"pre_quant_scale\": getattr(module, \"pre_quant_scale\", None) is not None} for name, module in modules.items()}\n"
|
|
" torch.save({\"input\": hidden_states.detach().cpu(), \"output\": {\"q\": xq.detach().cpu(), \"k\": xk.detach().cpu(), \"v\": xv.detach().cpu()}, \"metadata\": metadata}, os.path.join(capture_dir, \"qwen0_loaded_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"
|
|
" modules = {\"gate\": self.gate_proj, \"up\": self.up_proj, \"down\": self.down_proj}\n"
|
|
" metadata = {name: {\"class\": f\"{type(module).__module__}.{type(module).__qualname__}\", \"weight_type\": f\"{type(module.weight).__module__}.{type(module.weight).__qualname__}\", \"weight_dtype\": str(module.weight.dtype), \"weight_shape\": tuple(module.weight.shape), \"layout_type\": getattr(module, \"layout_type\", None), \"full_precision\": getattr(module, \"_full_precision_mm\", None), \"pre_quant_scale\": getattr(module, \"pre_quant_scale\", None) is not None} for name, module in modules.items()}\n"
|
|
" torch.save({\"input\": x.detach().cpu(), \"output\": {\"gate\": gate.detach().cpu(), \"up\": up.detach().cpu(), \"activated\": activated.detach().cpu(), \"down\": output.detach().cpu()}, \"metadata\": metadata}, os.path.join(capture_dir, \"qwen0_loaded_mlp.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 loaded Qwen layer-0 projection probe.")
|