59 lines
4.1 KiB
Python
59 lines
4.1 KiB
Python
|
"""Capture loaded Comfy Qwen QuantizedTensor fields and dispatch variants."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
|
||
|
|
source = path.read_text(encoding="utf-8")
|
||
|
|
helper = (
|
||
|
|
"def _h3_owned_cpu(tensor):\n"
|
||
|
|
" return tensor.detach().to(device=\"cpu\", copy=True).contiguous().clone()\n\n"
|
||
|
|
)
|
||
|
|
if "def _h3_owned_cpu(tensor):\n" not in source:
|
||
|
|
source = helper + source
|
||
|
|
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"
|
||
|
|
" def probe(module):\n"
|
||
|
|
" stored = module.weight\n"
|
||
|
|
" qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n"
|
||
|
|
" weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, hidden_states, offloadable=True)\n"
|
||
|
|
" try:\n"
|
||
|
|
" output = torch.nn.functional.linear(hidden_states, weight, bias)\n"
|
||
|
|
" return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n"
|
||
|
|
" finally:\n"
|
||
|
|
" comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n"
|
||
|
|
" torch.save({\"input\": hidden_states.detach().cpu(), \"modules\": {name: probe(module) for name, module in modules.items()}}, os.path.join(capture_dir, \"qwen0_quantized_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"
|
||
|
|
" def probe(module, value):\n"
|
||
|
|
" stored = module.weight\n"
|
||
|
|
" qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n"
|
||
|
|
" weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, value, offloadable=True)\n"
|
||
|
|
" try:\n"
|
||
|
|
" output = torch.nn.functional.linear(value, weight, bias)\n"
|
||
|
|
" return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n"
|
||
|
|
" finally:\n"
|
||
|
|
" comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n"
|
||
|
|
" torch.save({\"input\": x.detach().cpu(), \"activated\": activated.detach().cpu(), \"modules\": {\"gate\": probe(self.gate_proj, x), \"up\": probe(self.up_proj, x), \"down\": probe(self.down_proj, activated)}}, os.path.join(capture_dir, \"qwen0_quantized_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 Qwen QuantizedTensor dispatch probe.")
|