253 lines
12 KiB
Python
253 lines
12 KiB
Python
"""Localize the padding-embedding divergence by comparing comfy (torch) and MLX
|
|
intermediates through layer 0, at chosen token positions.
|
|
|
|
It runs the SAME padded prompt through both encoders and compares, stage by stage:
|
|
embed_in -> input to layer 0 (token embedding)
|
|
ln1 -> input_layernorm(x)
|
|
attn_out -> self_attn(...) output (o_proj, before residual)
|
|
layer_out -> output of decoder layer 0
|
|
The FIRST stage that diverges (low cosine / high max-abs-diff) is the culprit:
|
|
embed_in bad -> tokenization/embedding; ln1 bad -> layernorm;
|
|
attn_out bad -> attention (mask/rope/qk-norm/sdpa); layer_out bad -> mlp/residual.
|
|
|
|
Run on your Mac (needs comfy + torch + mlx):
|
|
cd ComfyUI
|
|
python custom_nodes/ComfyUI-SeFi-MLX/sefi_mlx/dump_layer0.py \
|
|
models/text_encoders/sefi-image-5b-turbo_qwen3vl.safetensors \
|
|
"a blue ceramic mug on a white desk"
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
|
|
|
|
def _cmp(name, a, b):
|
|
import numpy as np
|
|
a = np.asarray(a, dtype=np.float64).reshape(-1)
|
|
b = np.asarray(b, dtype=np.float64).reshape(-1)
|
|
n = min(a.shape[0], b.shape[0])
|
|
a, b = a[:n], b[:n]
|
|
cos = float((a * b).sum() / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
|
|
md = float(np.abs(a - b).max())
|
|
na, nb = float(np.linalg.norm(a)), float(np.linalg.norm(b))
|
|
print(f" {name:10s} cos={cos:.4f} max|d|={md:.4e} |comfy|={na:.3f} |mlx|={nb:.3f}")
|
|
|
|
|
|
def main():
|
|
here = os.path.abspath(__file__)
|
|
cn = os.path.dirname(os.path.dirname(os.path.dirname(here)))
|
|
comfy_root = os.path.dirname(cn)
|
|
sys.path.insert(0, comfy_root)
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(here)))
|
|
|
|
import numpy as np
|
|
import torch
|
|
import mlx.core as mx
|
|
|
|
ckpt = sys.argv[1] if len(sys.argv) > 1 else os.path.join(
|
|
comfy_root, "models", "text_encoders", "sefi-image-5b-turbo_qwen3vl.safetensors")
|
|
prompt = sys.argv[2] if len(sys.argv) > 2 else "a blue ceramic mug on a white desk"
|
|
pad_len = 512
|
|
positions = [16, 17, 18, 19, 20, 21, 64, 256] # think tokens + padding
|
|
|
|
# ---------------- comfy side (hooks on layer 0) ----------------
|
|
import comfy.sd
|
|
import folder_paths
|
|
ctype = getattr(comfy.sd.CLIPType, "FLUX2", None) or comfy.sd.CLIPType.FLUX
|
|
clip = comfy.sd.load_clip(ckpt_paths=[ckpt], clip_type=ctype,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
for _n, sub in vars(clip.tokenizer).items():
|
|
if hasattr(sub, "min_length"):
|
|
sub.min_length = pad_len
|
|
|
|
# find layer 0 (first module with input_layernorm + self_attn) on the torch model
|
|
root = clip.cond_stage_model
|
|
layer0 = None
|
|
for name, mod in sorted(root.named_modules(), key=lambda x: x[0]):
|
|
if hasattr(mod, "input_layernorm") and hasattr(mod, "self_attn"):
|
|
layer0 = mod
|
|
print(f"comfy layer0 module: {name}")
|
|
break
|
|
if layer0 is None:
|
|
print("could not find layer0 in comfy model"); return
|
|
|
|
caps = {}
|
|
|
|
def grab(key):
|
|
def hook(_m, inp, out):
|
|
t = out[0] if isinstance(out, tuple) else out
|
|
caps[key] = t.detach().float().cpu().numpy()
|
|
return hook
|
|
|
|
def grab_in(key):
|
|
def pre(_m, args, kwargs):
|
|
t = args[0] if args else kwargs.get("x")
|
|
if t is None:
|
|
t = kwargs.get("hidden_states")
|
|
caps[key] = t.detach().float().cpu().numpy()
|
|
return pre
|
|
|
|
hs = [layer0.register_forward_pre_hook(grab_in("embed_in"), with_kwargs=True),
|
|
layer0.input_layernorm.register_forward_hook(grab("ln1")),
|
|
layer0.self_attn.q_proj.register_forward_hook(grab("q_proj")),
|
|
layer0.self_attn.o_proj.register_forward_hook(grab("attn_out")),
|
|
layer0.register_forward_hook(grab("layer_out"))]
|
|
if getattr(layer0.self_attn, "q_norm", None) is not None:
|
|
hs.append(layer0.self_attn.q_norm.register_forward_hook(grab("q_norm")))
|
|
|
|
# wrap comfy's apply_rope to capture layer-0 post-rope q/k
|
|
import comfy.text_encoders.llama as _cl
|
|
_orig_rope = _cl.apply_rope
|
|
roped = {}
|
|
|
|
# wrap precompute_freqs_cis to capture the EXACT args comfy uses
|
|
_orig_pfc = _cl.precompute_freqs_cis
|
|
|
|
def _patched_pfc(head_dim, position_ids, theta, *a, **kw):
|
|
if "pfc" not in roped:
|
|
try:
|
|
import numpy as _np
|
|
roped["pfc"] = {
|
|
"head_dim": head_dim, "theta": theta,
|
|
"pos_shape": tuple(position_ids.shape),
|
|
"pos_sample": position_ids.detach().float().cpu().numpy()[:, :6].tolist(),
|
|
"args": [repr(x) for x in a], "kwargs": {k: repr(v) for k, v in kw.items()},
|
|
}
|
|
except Exception as e:
|
|
roped["pfc_err"] = str(e)
|
|
return _orig_pfc(head_dim, position_ids, theta, *a, **kw)
|
|
_cl.precompute_freqs_cis = _patched_pfc
|
|
|
|
def _patched_rope(xq, xk, freqs_cis):
|
|
out = _orig_rope(xq, xk, freqs_cis)
|
|
if "q" not in roped: # first call == layer 0
|
|
roped["q"] = out[0].detach().float().cpu().numpy() # (B,H,L,D)
|
|
roped["k"] = out[1].detach().float().cpu().numpy()
|
|
try:
|
|
roped["cos"] = freqs_cis[0].detach().float().cpu().numpy()
|
|
roped["cos_shape"] = tuple(freqs_cis[0].shape)
|
|
except Exception as e:
|
|
roped["cos_err"] = str(e)
|
|
return out
|
|
_cl.apply_rope = _patched_rope
|
|
|
|
tokens = clip.tokenize(prompt)
|
|
clip.encode_from_tokens_scheduled(tokens)
|
|
_cl.apply_rope = _orig_rope
|
|
_cl.precompute_freqs_cis = _orig_pfc
|
|
for h in hs:
|
|
h.remove()
|
|
print("comfy precompute_freqs_cis call:", roped.get("pfc"), roped.get("pfc_err"))
|
|
comfy_ids = [t[0] for t in tokens[next(iter(tokens))][0]]
|
|
print("comfy seq len:", len(comfy_ids))
|
|
|
|
# ---------------- MLX side (manual layer 0, fp32 to match) ----------------
|
|
from sefi_mlx.encoder_mlx import (load_sefi_encoder_mlx, SEFI_TEMPLATE,
|
|
sefi_rope_tables, ENCODER_ROPE_THETA)
|
|
rope_tables = lambda L, hd, _theta: sefi_rope_tables(L, hd, ENCODER_ROPE_THETA)
|
|
cond = load_sefi_encoder_mlx(ckpt, max_length=pad_len, weight_dtype="fp32")
|
|
enc = cond.tokenizer(SEFI_TEMPLATE.format(prompt), truncation=True,
|
|
padding="max_length", max_length=pad_len, return_tensors="np")
|
|
ids = mx.array(enc["input_ids"].astype("int32"))
|
|
amask_bool = mx.array(enc["attention_mask"].astype(bool))
|
|
m = cond.model
|
|
cfg = m.config
|
|
x = m.embed_tokens(ids)
|
|
L = x.shape[1]
|
|
cos_t, sin_t = rope_tables(L, cfg.head_dim, cfg.rope_theta)
|
|
cos_t, sin_t = cos_t.astype(x.dtype), sin_t.astype(x.dtype)
|
|
causal = mx.tril(mx.ones((L, L), dtype=mx.float32)) > 0
|
|
keep = causal[None, None, :, :] & amask_bool[:, None, None, :]
|
|
amask = mx.where(keep, mx.array(0.0, dtype=x.dtype), mx.array(-1e9, dtype=x.dtype))
|
|
layer = m.layers[0]
|
|
ln1 = layer.input_layernorm(x)
|
|
attn = layer.self_attn
|
|
B = x.shape[0]
|
|
hd, nq, nkv = cfg.head_dim, attn.n_heads, attn.n_kv
|
|
qp = attn.q_proj(ln1) # (B, L, H*D)
|
|
qn = attn.q_norm(qp.reshape(B, L, nq, hd)) # (B,L,H,D)
|
|
attn_out = layer.self_attn(ln1, cos_t, sin_t, amask) # fast-SDPA path
|
|
layer_out = layer(x, cos_t, sin_t, amask)
|
|
|
|
# ---- manual fp32 attention (no flash kernel) ----
|
|
from krea2_mlx.encoder_mlx import _apply_rope
|
|
q = attn.q_norm(qp.reshape(B, L, nq, hd)).transpose(0, 2, 1, 3) # (B,nq,L,D)
|
|
k = attn.k_norm(attn.k_proj(ln1).reshape(B, L, nkv, hd)).transpose(0, 2, 1, 3)
|
|
v = attn.v_proj(ln1).reshape(B, L, nkv, hd).transpose(0, 2, 1, 3)
|
|
q_pre = q
|
|
q, k = _apply_rope(q, k, cos_t, sin_t)
|
|
q_roped = np.array(q.astype(mx.float32)) # (B,nq,L,D) post-rope, for comparison
|
|
rep = nq // nkv
|
|
k = mx.repeat(k, rep, axis=1)
|
|
v = mx.repeat(v, rep, axis=1)
|
|
scale = hd ** -0.5
|
|
scores = (q.astype(mx.float32) @ k.astype(mx.float32).transpose(0, 1, 3, 2)) * scale
|
|
scores = scores + amask.astype(mx.float32)
|
|
w = mx.softmax(scores, axis=-1)
|
|
o = (w @ v.astype(mx.float32)).transpose(0, 2, 1, 3).reshape(B, L, nq * hd)
|
|
attn_manual = attn.o_proj(o.astype(x.dtype))
|
|
|
|
mx.eval(x, ln1, qp, qn, attn_out, attn_manual, layer_out)
|
|
mlx_caps = {"embed_in": np.array(x.astype(mx.float32)),
|
|
"ln1": np.array(ln1.astype(mx.float32)),
|
|
"q_proj": np.array(qp.astype(mx.float32)),
|
|
"q_norm": np.array(qn.astype(mx.float32)), # (B,L,H,D)
|
|
"attn_out": np.array(attn_out.astype(mx.float32)),
|
|
"attn_manual": np.array(attn_manual.astype(mx.float32)),
|
|
"layer_out": np.array(layer_out.astype(mx.float32))}
|
|
mlx_ids = enc["input_ids"][0].tolist()
|
|
print("ids match:", sum(1 for i in range(min(len(comfy_ids), len(mlx_ids)))
|
|
if comfy_ids[i] == mlx_ids[i]), "/", len(mlx_ids))
|
|
|
|
# ---- compare comfy's rope cos table to mine, to see the freq mapping ----
|
|
print("comfy rope cos shape:", roped.get("cos_shape"), "err:", roped.get("cos_err"))
|
|
if "cos" in roped:
|
|
cc = roped["cos"]
|
|
while cc.ndim > 2: # squeeze to (L, head_dim)
|
|
cc = cc[0] if cc.shape[0] == 1 else cc.reshape(-1, cc.shape[-1])
|
|
mine = np.array(cos_t.astype(mx.float32)) # (L, head_dim)
|
|
for p in (16, 256):
|
|
print(f" rope-cos @pos{p}: comfy[:8]={np.round(cc[p,:8],4)}")
|
|
print(f" mlx [:8]={np.round(mine[p,:8],4)}")
|
|
d = np.abs(cc[p] - mine[p])
|
|
print(f" max|d|={d.max():.4f} at dim {int(d.argmax())}")
|
|
# ---- reverse-engineer comfy's dim -> inv_freq-index permutation ----
|
|
half = cfg.head_dim // 2
|
|
inv = 1.0 / (cfg.rope_theta ** (np.arange(0, cfg.head_dim, 2) / cfg.head_dim)) # (half,)
|
|
perm = []
|
|
for j in range(half):
|
|
err = (np.abs(cc[16, j] - np.cos(16 * inv)) +
|
|
np.abs(cc[256, j] - np.cos(256 * inv)) +
|
|
np.abs(cc[64, j] - np.cos(64 * inv)))
|
|
perm.append(int(err.argmin()))
|
|
perm = np.array(perm)
|
|
print(f" dim->invfreq permutation (first 32): {perm[:32].tolist()}")
|
|
print(f" is identity: {bool((perm == np.arange(half)).all())}")
|
|
# show it as 3 interleaved groups if applicable
|
|
print(f" perm[::3][:12]={perm[0::3][:12].tolist()}")
|
|
print(f" perm[1::3][:12]={perm[1::3][:12].tolist()}")
|
|
print(f" perm[2::3][:12]={perm[2::3][:12].tolist()}")
|
|
|
|
# ---------------- compare per position, per stage ----------------
|
|
for p in positions:
|
|
print(f"\n[pos {p}] id comfy={comfy_ids[p]} mlx={mlx_ids[p]}")
|
|
# compare both the fast-SDPA and manual-fp32 attention against comfy's attn_out
|
|
for stage in ("embed_in", "ln1", "q_proj", "attn_out", "attn_manual", "layer_out"):
|
|
ca = caps.get("attn_out" if stage == "attn_manual" else stage)
|
|
ma = mlx_caps.get(stage)
|
|
if ca is None or ma is None:
|
|
print(f" {stage}: missing"); continue
|
|
_cmp(stage, ca[0, p], ma[0, p])
|
|
# q_norm: comfy (B,H,L,D) vs mlx (B,L,H,D) -> compare (H,D) at position p
|
|
if "q_norm" in caps and "q_norm" in mlx_caps:
|
|
_cmp("q_norm", caps["q_norm"][0, :, p, :], mlx_caps["q_norm"][0, p, :, :])
|
|
# post-rope q: comfy roped (B,H,L,D) vs mlx q_roped (B,H,L,D) at position p
|
|
if "q" in roped:
|
|
_cmp("q_roped", roped["q"][0, :, p, :], q_roped[0, :, p, :])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|