remove block steering + NVFP4 loader - both failed live validation
Block steering: run-averaged block activity is architecture-determined (varies in the 4th decimal run-to-run), so no scoring scheme - pole EMAs or reward correlation - could extract a rating signal from it. The idea needs within-run per-step measurement to have a chance; that is a research project, not a feature. Removed: module, sampler wiring, the two trailing widgets (positional reference mapping unaffected), the Refiner pairing hook, editor knobs, tests. Stale block_snapshot/ block_profile sidecars are inert and covered by the existing delete_refinement_key sweep. FunPackNVFP4ModelLoader: the FP4 GEMM path was real (3.35x at 16k tokens) but end-to-end the user's runs are dominated by everything the loader cannot touch (attention, prep, encode - the sampler node is ~23s of a ~60s run) while dynamic input quantization added a fixed per-layer-call toll. Net observed: overhead. Users are better served by dedicated NVFP4 loaders. Removed: module, registration, tests. Kept (independently useful): the chain sampler's phase timing report + quantized-layer census, the nested-AV x0 snapshot fix, and the sm_120 masked-attention work. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Fable 5
Claude Fable 5
parent
e9aee731e7
commit
b7ed4f9fb6
@@ -20,7 +20,6 @@ if __package__:
|
||||
FunPackVideoStitch,
|
||||
)
|
||||
from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader
|
||||
from .nvfp4_patch import FunPackNVFP4ModelLoader
|
||||
from .samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler
|
||||
from .templates import FunPackRefinementKeyLoader
|
||||
try:
|
||||
@@ -63,10 +62,6 @@ else:
|
||||
except Exception:
|
||||
FunPackApplyLoraWeights = None
|
||||
FunPackLoraLoader = None
|
||||
try:
|
||||
from nvfp4_patch import FunPackNVFP4ModelLoader
|
||||
except Exception:
|
||||
FunPackNVFP4ModelLoader = None
|
||||
try:
|
||||
from samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler
|
||||
except Exception:
|
||||
@@ -101,7 +96,6 @@ NODE_CLASS_MAPPINGS = {
|
||||
"FunPackLTXAVSceneChainSampler": FunPackLTXAVSceneChainSampler,
|
||||
"FunPackApplyLoraWeights": FunPackApplyLoraWeights,
|
||||
"FunPackLoraLoader": FunPackLoraLoader,
|
||||
"FunPackNVFP4ModelLoader": FunPackNVFP4ModelLoader,
|
||||
"FunPackRefinementKeyLoader": FunPackRefinementKeyLoader,
|
||||
}
|
||||
NODE_CLASS_MAPPINGS = {name: cls for name, cls in NODE_CLASS_MAPPINGS.items() if cls is not None}
|
||||
@@ -127,7 +121,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FunPackLTXAVSceneChainSampler": "FunPack LTXAV Scene Chain Sampler",
|
||||
"FunPackApplyLoraWeights": "FunPack Apply LoRA Weights",
|
||||
"FunPackLoraLoader": "FunPack LoRA Loader",
|
||||
"FunPackNVFP4ModelLoader": "FunPack NVFP4 Model Loader (Experimental)",
|
||||
"FunPackRefinementKeyLoader": "FunPack Refinement Key Loader",
|
||||
}
|
||||
|
||||
|
||||
-395
@@ -1,395 +0,0 @@
|
||||
"""Block Steering (EXPERIMENTAL): rating-learned per-block residual gains for LTX-AV.
|
||||
|
||||
RLHF at the MODEL level instead of the conditioning level. The rating system already
|
||||
attributes quality to sampling settings and conditioning; this adds the missing
|
||||
coordinate system - which REGIONS OF THE MODEL were doing the work when a generation
|
||||
came out good or bad - and then steers per-block application accordingly.
|
||||
|
||||
Three stages, all riding existing FunPack infra:
|
||||
|
||||
1. MEASURE - forward hooks on the transformer blocks record, per model call, each
|
||||
block's relative residual delta on the VIDEO stream: ||vx_out - vx_in|| / ||vx_in||.
|
||||
A block that barely moves the hidden state is coasting; one with a large relative
|
||||
delta is doing the work for this prompt at this noise level. Hooks are tagged with
|
||||
ltx_enhancements' _FUNPACK_HOOK_TAG so the existing leak-strip covers them.
|
||||
|
||||
2. ATTRIBUTE - two signals, folded into a per-key `block_profile.pt` sidecar:
|
||||
* run-level: the end-of-run mean activity fingerprint is staged (`block_snapshot.pt`,
|
||||
same stage-then-pair convention as x0_snapshot) and paired with the user's rating
|
||||
reward on the next Refiner cycle - liked / disliked / all-runs EMA profiles.
|
||||
* step-level: when output_guidance's value function is trained, a model-function
|
||||
wrapper scores each step's x0_hat (one tiny-MLP forward, no extra model pass) and
|
||||
credits blocks by (value delta this step) x (their activity this step, centered).
|
||||
Converges in a handful of runs - matters because keys are retrained per instance.
|
||||
|
||||
3. STEER - per-block gains on the video residual only: vx_out' = vx_in + g_b*(vx_out -
|
||||
vx_in), g_b = 1 + strength * score_b, scores zero-mean/unit-max so emphasis is
|
||||
REDISTRIBUTED rather than globally inflated, clamped to +/-MAX_GAIN_DELTA. The audio
|
||||
stream is never touched (residual-gain on the video half is the intervention family
|
||||
that works on LTX-AV; hidden-state injection through joint attention is the one that
|
||||
corrupts audio). Exact no-op until the profile has enough paired data.
|
||||
|
||||
Everything is per-refinement-key and rebuilt from ratings in a few generations - a
|
||||
mechanism upgrade, not a data store to migrate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
from .ltx_enhancements import _FUNPACK_HOOK_TAG, _funpack_locate_blocks
|
||||
except ImportError:
|
||||
from ltx_enhancements import _FUNPACK_HOOK_TAG, _funpack_locate_blocks
|
||||
|
||||
# Steering never moves a block's residual contribution by more than this, regardless of
|
||||
# strength widget or how lopsided the learned scores are.
|
||||
MAX_GAIN_DELTA = 0.10
|
||||
# Profile readiness: at least this many rating-paired runs (with 2, per-block correlation
|
||||
# is degenerate — every block correlates ±1 trivially)…
|
||||
MIN_RATED_RUNS = 3
|
||||
# …with at least this much reward spread. Attribution is the per-block CORRELATION between
|
||||
# reward and activity across runs — if every run is rated the same there is nothing to
|
||||
# attribute, no matter how many runs were paired (the v1 liked/disliked-EMA design failed
|
||||
# exactly there: 16 same-pole ratings produced zero contrast and looked "not ready").
|
||||
# Correlation, not raw covariance: run-mean activity is nearly deterministic (averaged
|
||||
# over ~1200 steps it varies in the 4th decimal run-to-run), so raw covariance is
|
||||
# numerically dead ("flat activity across runs") — but a tiny variation that CONSISTENTLY
|
||||
# tracks reward is exactly the signal we want, and correlation is scale-free.
|
||||
MIN_REWARD_STD = 0.05
|
||||
# Confidence ramp: scores scale by min(1, (runs - 2) / 6) so a young profile steers
|
||||
# gently (correlation over 3 runs is noisy) and reaches full strength around 8 runs.
|
||||
CONFIDENCE_FULL_AT = 8
|
||||
# …or this many value-credit runs, before scores are considered meaningful.
|
||||
MIN_CREDIT_RUNS = 2
|
||||
# Rolling history of (activity fingerprint, reward) pairs kept per key.
|
||||
HISTORY_MAX = 32
|
||||
# Blend weights when both attribution signals are available.
|
||||
RATING_WEIGHT = 0.6
|
||||
CREDIT_WEIGHT = 0.4
|
||||
_EPS = 1e-8
|
||||
|
||||
|
||||
def _video_stream(x):
|
||||
"""The video tensor of a block input/output: (vx, ax) tuple on LTX-AV, plain tensor
|
||||
on video-only LTXV."""
|
||||
if isinstance(x, (tuple, list)) and len(x) >= 1:
|
||||
return x[0]
|
||||
return x if isinstance(x, torch.Tensor) else None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stage 1 - measurement
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class BlockActivityRecorder:
|
||||
"""Accumulates per-block relative residual deltas (video stream) across a run.
|
||||
|
||||
`cur` always holds the most recent model call's per-block activity so the value-
|
||||
credit wrapper (which runs after all blocks in the same call) can read it.
|
||||
"""
|
||||
|
||||
def __init__(self, n_blocks: int):
|
||||
self.n_blocks = int(n_blocks)
|
||||
self.sums = torch.zeros(self.n_blocks, dtype=torch.float32)
|
||||
self.counts = torch.zeros(self.n_blocks, dtype=torch.float32)
|
||||
self.cur = torch.zeros(self.n_blocks, dtype=torch.float32)
|
||||
# value-credit accumulation (filled by the credit wrapper, if installed)
|
||||
self.credit = torch.zeros(self.n_blocks, dtype=torch.float32)
|
||||
self.credit_steps = 0
|
||||
self._prev_score = None
|
||||
|
||||
def record(self, idx: int, rel_delta: float):
|
||||
if 0 <= idx < self.n_blocks:
|
||||
self.sums[idx] += rel_delta
|
||||
self.counts[idx] += 1.0
|
||||
self.cur[idx] = rel_delta
|
||||
|
||||
def fingerprint(self):
|
||||
"""Mean per-block activity over the run, or None if nothing was recorded."""
|
||||
if float(self.counts.max()) <= 0:
|
||||
return None
|
||||
return self.sums / self.counts.clamp(min=1.0)
|
||||
|
||||
def credit_step(self, score: float):
|
||||
"""Fold one model call's value score into per-block credit: blocks that were
|
||||
MORE active than average during a step where predicted quality rose (fell) get
|
||||
positive (negative) credit."""
|
||||
if self._prev_score is not None:
|
||||
delta = float(score) - self._prev_score
|
||||
centered = self.cur - self.cur.mean()
|
||||
self.credit += delta * centered
|
||||
self.credit_steps += 1
|
||||
self._prev_score = float(score)
|
||||
|
||||
def credit_vector(self):
|
||||
if self.credit_steps <= 0:
|
||||
return None
|
||||
return self.credit / float(self.credit_steps)
|
||||
|
||||
|
||||
def install_recorder(model, recorder: BlockActivityRecorder):
|
||||
"""Tagged forward hooks on every transformer block; returns handles (possibly [])."""
|
||||
blocks = _funpack_locate_blocks(model)
|
||||
if blocks is None:
|
||||
return []
|
||||
handles = []
|
||||
for i, blk in enumerate(blocks):
|
||||
def _hook(_module, args, output, _idx=i, _rec=recorder):
|
||||
try:
|
||||
v_in = _video_stream(args[0] if args else None)
|
||||
v_out = _video_stream(output)
|
||||
if v_in is None or v_out is None or v_in.shape != v_out.shape:
|
||||
return output
|
||||
with torch.no_grad():
|
||||
denom = float(v_in.norm()) + _EPS
|
||||
_rec.record(_idx, float((v_out - v_in).norm()) / denom)
|
||||
except Exception:
|
||||
pass
|
||||
return output
|
||||
setattr(_hook, _FUNPACK_HOOK_TAG, True)
|
||||
handles.append(blk.register_forward_hook(_hook))
|
||||
return handles
|
||||
|
||||
|
||||
def make_credit_wrapper(old_wrapper, recorder: BlockActivityRecorder, value_fn):
|
||||
"""model_function_wrapper that scores each call's x0_hat with the (already trained)
|
||||
output-space value function and folds the step-to-step delta into block credit.
|
||||
Installed INNERMOST so it scores the base model's own prediction, before
|
||||
embed_guidance / output_guidance corrections layer on top."""
|
||||
|
||||
def _call(apply_fn, a):
|
||||
if old_wrapper is not None:
|
||||
return old_wrapper(apply_fn, a)
|
||||
return apply_fn(a["input"], a["timestep"], **a.get("c", {}))
|
||||
|
||||
def _wrapper(apply_fn, args, _rec=recorder, _vf=value_fn):
|
||||
denoised = _call(apply_fn, args)
|
||||
try:
|
||||
with torch.inference_mode(False), torch.no_grad():
|
||||
c = _vf.compress(denoised.detach().float())
|
||||
score = float(_vf(c).mean())
|
||||
_rec.credit_step(score)
|
||||
except Exception:
|
||||
pass
|
||||
return denoised
|
||||
|
||||
return _wrapper
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stage 3 - steering
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def install_steer(model, gains):
|
||||
"""Tagged forward hooks scaling each block's VIDEO residual by its gain:
|
||||
vx' = vx_in + g*(vx_out - vx_in). Audio passes through untouched. Returns handles;
|
||||
blocks whose gain is ~1.0 get no hook at all (byte-identical native path)."""
|
||||
blocks = _funpack_locate_blocks(model)
|
||||
if blocks is None:
|
||||
return []
|
||||
handles = []
|
||||
for i, blk in enumerate(blocks):
|
||||
g = float(gains[i]) if i < len(gains) else 1.0
|
||||
if abs(g - 1.0) < 1e-4:
|
||||
continue
|
||||
def _hook(_module, args, output, _g=g):
|
||||
try:
|
||||
v_in = _video_stream(args[0] if args else None)
|
||||
if isinstance(output, tuple) and len(output) == 2:
|
||||
v_out, a_out = output
|
||||
if v_in is None or v_in.shape != v_out.shape:
|
||||
return output
|
||||
return (v_in + _g * (v_out - v_in), a_out)
|
||||
if isinstance(output, torch.Tensor) and v_in is not None and v_in.shape == output.shape:
|
||||
return v_in + _g * (output - v_in)
|
||||
except Exception:
|
||||
pass
|
||||
return output
|
||||
setattr(_hook, _FUNPACK_HOOK_TAG, True)
|
||||
handles.append(blk.register_forward_hook(_hook))
|
||||
return handles
|
||||
|
||||
|
||||
def remove_handles(handles):
|
||||
for h in handles or []:
|
||||
try:
|
||||
h.remove()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stage 2 - persistence + attribution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _state_path(refinement_key, name):
|
||||
try:
|
||||
from .conditioning import refinement_state_path
|
||||
except ImportError:
|
||||
from conditioning import refinement_state_path
|
||||
return refinement_state_path(refinement_key, name, prefix="refine_v2", extension="pt")
|
||||
|
||||
|
||||
def save_run_snapshot(refinement_key, recorder: BlockActivityRecorder):
|
||||
"""End of run: stage this run's fingerprint (+ optional value credit) so the NEXT
|
||||
rating cycle can pair it with a reward. Atomic write - a rating read never sees a
|
||||
partial file. Same convention as the sampler's x0_snapshot."""
|
||||
if not refinement_key:
|
||||
return False
|
||||
fp = recorder.fingerprint()
|
||||
if fp is None:
|
||||
return False
|
||||
# Store RELATIVE emphasis (mean-normalized): run-global scale differences (step
|
||||
# counts, schedulers, scene counts) must not masquerade as per-block signal.
|
||||
fp = fp / (fp.mean() + _EPS)
|
||||
try:
|
||||
path = _state_path(refinement_key, "block_snapshot")
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
tmp = path + ".tmp"
|
||||
torch.save({
|
||||
"activity": fp.cpu(),
|
||||
"credit": (recorder.credit_vector().cpu() if recorder.credit_vector() is not None else None),
|
||||
"n_blocks": recorder.n_blocks,
|
||||
}, tmp)
|
||||
os.replace(tmp, path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"[FunPackBlockSteer] snapshot save failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def _ema(old, x, n):
|
||||
a = max(0.2, 1.0 / (float(n) + 1.0))
|
||||
return x if old is None else old * (1.0 - a) + x * a
|
||||
|
||||
|
||||
def update_profile_with_rating(refinement_key, reward):
|
||||
"""Pair the staged snapshot with a rating reward (called by the Refiner at the same
|
||||
site that trains the output value function). Appends (activity fingerprint, reward)
|
||||
to a rolling per-key history — EVERY admissible reward contributes with its actual
|
||||
value (no liked/disliked pole thresholds, no dead band) — and folds the run's value
|
||||
credit (rating-independent, self-supervised) into the credit EMA. Consumes the
|
||||
snapshot so one run is never paired twice. Returns the number of rating-paired runs
|
||||
in the history, or None when there was nothing to pair."""
|
||||
if not refinement_key:
|
||||
return None
|
||||
try:
|
||||
snap_path = _state_path(refinement_key, "block_snapshot")
|
||||
if not os.path.exists(snap_path):
|
||||
return None
|
||||
snap = torch.load(snap_path, map_location="cpu", weights_only=False)
|
||||
activity = snap.get("activity")
|
||||
if not isinstance(activity, torch.Tensor):
|
||||
return None
|
||||
prof_path = _state_path(refinement_key, "block_profile")
|
||||
prof = {}
|
||||
if os.path.exists(prof_path):
|
||||
prof = torch.load(prof_path, map_location="cpu", weights_only=False)
|
||||
# Start fresh on: different model depth, a v1 pole-EMA profile, or a v2
|
||||
# profile whose fingerprints predate mean-normalization (raw scales would
|
||||
# poison the correlation).
|
||||
if (int(prof.get("n_blocks", -1)) != activity.numel()
|
||||
or ("history_act" not in prof and "liked_n" in prof)
|
||||
or int(prof.get("fp_version", 1)) < 3):
|
||||
prof = {}
|
||||
prof["n_blocks"] = activity.numel()
|
||||
prof["fp_version"] = 3
|
||||
|
||||
row = activity.float().reshape(1, -1)
|
||||
rw = torch.tensor([float(reward)], dtype=torch.float32)
|
||||
ha, hr = prof.get("history_act"), prof.get("history_rw")
|
||||
prof["history_act"] = row if ha is None else torch.cat([ha, row], dim=0)[-HISTORY_MAX:]
|
||||
prof["history_rw"] = rw if hr is None else torch.cat([hr, rw], dim=0)[-HISTORY_MAX:]
|
||||
|
||||
credit = snap.get("credit")
|
||||
if isinstance(credit, torch.Tensor) and credit.numel() == activity.numel():
|
||||
prof["credit"] = _ema(prof.get("credit"), credit, prof.get("credit_n", 0))
|
||||
prof["credit_n"] = int(prof.get("credit_n", 0)) + 1
|
||||
|
||||
tmp = prof_path + ".tmp"
|
||||
torch.save(prof, tmp)
|
||||
os.replace(tmp, prof_path)
|
||||
os.remove(snap_path) # consumed - never pair one run with two ratings
|
||||
return int(prof["history_rw"].numel())
|
||||
except Exception as e:
|
||||
print(f"[FunPackBlockSteer] profile update failed: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _normalize(v):
|
||||
v = v - v.mean()
|
||||
m = float(v.abs().max())
|
||||
return v / m if m > _EPS else None
|
||||
|
||||
|
||||
def block_scores_with_status(refinement_key):
|
||||
"""(scores, status): combined normalized per-block scores in [-1, 1] (zero-mean,
|
||||
unit max-abs), or (None, why-not). Rating attribution is the per-block covariance
|
||||
between reward and activity across the run history; value credit blends in when
|
||||
trained. `status` is a user-facing sentence for the sampler log."""
|
||||
if not refinement_key:
|
||||
return None, "no refinement key wired"
|
||||
try:
|
||||
prof_path = _state_path(refinement_key, "block_profile")
|
||||
if not os.path.exists(prof_path):
|
||||
return None, "no profile yet — rate a generation made with block_steer on"
|
||||
prof = torch.load(prof_path, map_location="cpu", weights_only=False)
|
||||
ha, hr = prof.get("history_act"), prof.get("history_rw")
|
||||
credit, credit_n = prof.get("credit"), int(prof.get("credit_n", 0))
|
||||
|
||||
rating_part = None
|
||||
peak_corr = 0.0
|
||||
n = int(hr.numel()) if isinstance(hr, torch.Tensor) else 0
|
||||
conf = min(1.0, max(0.0, (n - (MIN_RATED_RUNS - 1)) / float(CONFIDENCE_FULL_AT - (MIN_RATED_RUNS - 1))))
|
||||
if not isinstance(ha, torch.Tensor) or n < MIN_RATED_RUNS:
|
||||
rating_status = f"needs {MIN_RATED_RUNS}+ rated runs for this key (have {n})"
|
||||
elif float(hr.std()) < MIN_REWARD_STD:
|
||||
rating_status = (f"{n} rated runs but all rated alike — rate both good AND bad "
|
||||
"runs so blocks have something to contrast")
|
||||
else:
|
||||
# Per-block Pearson correlation between reward and activity. Scale-free:
|
||||
# run-mean activity varies only in the 4th decimal run-to-run, so raw
|
||||
# covariance is numerically dead — but consistency with reward isn't.
|
||||
rc = hr - hr.mean()
|
||||
ac = ha - ha.mean(dim=0)
|
||||
corr = (rc.unsqueeze(1) * ac).sum(dim=0) / (
|
||||
(rc.square().sum().sqrt() * ac.square().sum(dim=0).sqrt()) + _EPS)
|
||||
corr = corr.clamp(-1.0, 1.0)
|
||||
peak_corr = float(corr.abs().max())
|
||||
rating_part = _normalize(corr)
|
||||
rating_status = ("no block's activity varies at all across runs"
|
||||
if rating_part is None else "ready")
|
||||
|
||||
credit_part = _normalize(credit) if (credit is not None and credit_n >= MIN_CREDIT_RUNS) else None
|
||||
|
||||
if rating_part is not None and credit_part is not None:
|
||||
combined = RATING_WEIGHT * rating_part + CREDIT_WEIGHT * credit_part
|
||||
else:
|
||||
combined = rating_part if rating_part is not None else credit_part
|
||||
if combined is None:
|
||||
return None, rating_status
|
||||
scores = _normalize(combined)
|
||||
if scores is None:
|
||||
return None, "flat combined scores"
|
||||
# Confidence ramp: correlation over few runs is noisy — steer gently until the
|
||||
# history has substance. Applied AFTER normalization so it survives into gains.
|
||||
scores = scores * conf
|
||||
return scores, (f"ready ({n} runs, peak |corr| {peak_corr:.2f}, "
|
||||
f"confidence {conf:.2f})")
|
||||
except Exception as e:
|
||||
return None, f"score load failed: {e}"
|
||||
|
||||
|
||||
def load_block_scores(refinement_key):
|
||||
"""Scores only (None while not ready) — see block_scores_with_status."""
|
||||
return block_scores_with_status(refinement_key)[0]
|
||||
|
||||
|
||||
def gains_from_scores(scores, strength):
|
||||
"""Per-block gains 1 + strength*score, hard-clamped to 1 +/- MAX_GAIN_DELTA."""
|
||||
s = max(0.0, float(strength))
|
||||
g = 1.0 + (s * scores).clamp(-MAX_GAIN_DELTA, MAX_GAIN_DELTA)
|
||||
return [float(x) for x in g]
|
||||
@@ -11910,21 +11910,6 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
|
||||
n_out = self._v2_train_output_value_function(refinement_key, float(learning_profile.get("reward", 0.0)))
|
||||
if n_out is not None:
|
||||
print(f"[FunPackRefiner] Output value function updated — {n_out} samples")
|
||||
# Block Steering: the same reward pairs with the sampler's staged per-block
|
||||
# activity fingerprint (refinement_key.block_snapshot.pt) — feeds the
|
||||
# liked/disliked/credit block profile the Chain Sampler steers from. No-op
|
||||
# when the sampler ran without block_steer (no snapshot staged).
|
||||
try:
|
||||
try:
|
||||
from . import block_steer as _bsmod
|
||||
except ImportError:
|
||||
import block_steer as _bsmod
|
||||
n_blocks = _bsmod.update_profile_with_rating(
|
||||
refinement_key, float(learning_profile.get("reward", 0.0)))
|
||||
if n_blocks is not None:
|
||||
print(f"[FunPackRefiner] Block profile updated — {n_blocks} rated run(s)")
|
||||
except Exception as _bse:
|
||||
print(f"[FunPackRefiner] Block profile update failed: {_bse}")
|
||||
# Absolute store: the same rating also feeds the keyless, prompt-agnostic taste prior.
|
||||
# Runs even with no refinement_key (Absolute is global), so standalone runs still build it.
|
||||
# Skipped for Wrong-* repair ratings (skip_value_function): Absolute reads reward as pure
|
||||
|
||||
@@ -254,8 +254,6 @@
|
||||
{ name: "score_slider_strength", label: "Slider strength (eta)", kind: "float", default: 1.0, min: 0, max: 3, step: 0.25, dependsOn: "score_slider" },
|
||||
{ name: "output_guidance", label: "Output guidance", kind: "bool", default: false },
|
||||
{ name: "output_guidance_strength", label: "Output guidance strength", kind: "float", default: 0.02, min: 0.005, max: 0.1, step: 0.005, dependsOn: "output_guidance" },
|
||||
{ name: "block_steer", label: "Block steering", kind: "bool", default: false },
|
||||
{ name: "block_steer_strength", label: "Block steer strength", kind: "float", default: 0.05, min: 0.01, max: 0.15, step: 0.005, dependsOn: "block_steer" },
|
||||
{ name: "decode_noise_scale", label: "Decode noise scale", kind: "float", default: 0.0, min: 0, max: 1, step: 0.01 },
|
||||
{ name: "decode_timestep", label: "Decode timestep", kind: "float", default: 0.05, min: 0, max: 1, step: 0.01 },
|
||||
{ name: "decode_tile_size", label: "Decode tile size", kind: "int", default: 0, min: 0, max: 4096, step: 64 },
|
||||
@@ -408,7 +406,7 @@
|
||||
{ id: "chain_auto", title: "Auto continuity", defaultOpen: true, knobs: [] },
|
||||
{ id: "chain_timing", title: "Timing", defaultOpen: false, knobs: ["frame_overlap", "transition_duration", "use_same_seed"] },
|
||||
{ id: "chain_cont", title: "Manual continuity", defaultOpen: false, knobs: ["carry_i2v_guides"] },
|
||||
{ id: "chain_guid", title: "Guidance", defaultOpen: false, knobs: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "output_guidance", "output_guidance_strength", "block_steer", "block_steer_strength"] },
|
||||
{ id: "chain_guid", title: "Guidance", defaultOpen: false, knobs: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "output_guidance", "output_guidance_strength"] },
|
||||
{ id: "chain_dec", title: "Decode", defaultOpen: false, knobs: ["decode_noise_scale", "decode_timestep", "decode_tile_size"] },
|
||||
{ id: "chain_exp", title: "Experimental", defaultOpen: false, knobs: ["mid_scene_guide", "mid_scene_guide_strength", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "bounded_attention_enabled"] },
|
||||
];
|
||||
|
||||
-527
@@ -1,527 +0,0 @@
|
||||
"""FunPack NVFP4 on-the-fly model patching (EXPERIMENTAL).
|
||||
|
||||
Blackwell's 5th-gen tensor cores run NVFP4 (e2m1 with per-16-element FP8-E4M3 block
|
||||
scales + a per-tensor FP32 scale) at roughly 2x FP8 / 4x BF16 GEMM throughput. ComfyUI
|
||||
core already ships the whole runtime for it (comfy.quant_ops layouts + MixedPrecisionOps
|
||||
dispatch, dynamic activation quantization, LoRA fallback, dequant emulation on non-FP4
|
||||
GPUs) - but only for checkpoints that were quantized ahead of time. Stock loaders offer
|
||||
fp8 at best, and on sm_120 the fp8 path is mostly storage, not tensor-core speedup.
|
||||
|
||||
This module closes that gap: it quantizes an existing bf16/fp16 LTX-AV checkpoint to
|
||||
NVFP4 *at load time* as a pure state-dict transform - selected 2D Linear weights are
|
||||
re-serialized in core's exact quantized-checkpoint format (weight/weight_scale/
|
||||
weight_scale_2 + per-layer `comfy_quant` metadata), then handed to the normal
|
||||
`load_diffusion_model_state_dict` path. Core's own detection routes the model through
|
||||
MixedPrecisionOps; every layer we did NOT tag simply stays at compute dtype. No custom
|
||||
forward, no monkeypatching, no new inference code paths to maintain.
|
||||
|
||||
Zone-aware precision map (the part generic quantizers can't do): LTX-AV is a joint
|
||||
audio-video transformer and its zones are not equally robust. By default we quantize
|
||||
only the video-branch Linears (attn1/attn2/ff) of the middle blocks, keeping at full
|
||||
precision:
|
||||
- the audio branch (audio_attn1/audio_attn2/audio_ff) - joint attention is the part
|
||||
that corrupts first (see the self-consistency post-block injection failure),
|
||||
- the cross-modal bridges (audio_to_video_attn / video_to_audio_attn),
|
||||
- the first/last N blocks (standard quantization practice; they carry the most
|
||||
input/output-adjacent signal),
|
||||
- everything that is not a plain 2D matmul (norms, ada modulation, embeddings).
|
||||
|
||||
Scopes widen from there for experimentation. W4A4 without calibration WILL cost some
|
||||
fidelity - this is an experimental speed/VRAM lever, off unless you wire the node in.
|
||||
|
||||
Requires: ComfyUI with comfy.quant_ops (comfy_kitchen), torch cu13+ on Blackwell
|
||||
(sm_120+) for real FP4 GEMMs. On anything older the model still loads and runs through
|
||||
core's dequant emulation (no speedup, same quality) - useful for A/B sanity checks.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
|
||||
import torch
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Static input-scale calibration
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core quantizes each layer's INPUT dynamically (amax reduction + quantize) on every
|
||||
# call unless the layer has a static `input_scale` parameter. Live probe on sm_120:
|
||||
# the dynamic path costs ~0.30 ms/layer-call and a provided scale cuts it to ~0.11 ms
|
||||
# (amax dominates) - across ~800 quantized layers that is a top-line cost of every
|
||||
# model call. Calibration is capture-then-freeze: a "calibrate" run records each
|
||||
# quantized layer's input amax via forward-pre-hooks into a per-checkpoint sidecar;
|
||||
# later loads inject `input_scale = amax * MARGIN / (448 * 6)` into the state dict,
|
||||
# where core's loader picks it up natively. MARGIN leaves headroom because NVFP4
|
||||
# clips activations above the scale; per the disposable-keys philosophy the sidecar
|
||||
# is per-machine and rebuilt in one run whenever absent.
|
||||
CALIBRATION_MARGIN = 1.25
|
||||
_F8F4_MAX = 448.0 * 6.0 # F8_E4M3_MAX * F4_E2M1_MAX
|
||||
INPUT_SCALE_MODES = ["dynamic", "auto-calibrate"]
|
||||
|
||||
|
||||
def _scales_path(checkpoint_path):
|
||||
base = os.path.join(os.path.dirname(os.path.abspath(__file__)), "refinements", "nvfp4_scales")
|
||||
return os.path.join(base, os.path.basename(checkpoint_path) + ".json")
|
||||
|
||||
|
||||
def load_calibration(checkpoint_path):
|
||||
"""{module_path: input_amax} recorded by a previous calibrate run, or None."""
|
||||
try:
|
||||
p = _scales_path(checkpoint_path)
|
||||
if not os.path.exists(p):
|
||||
return None
|
||||
with open(p, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
amax = data.get("amax") or {}
|
||||
return {str(k): float(v) for k, v in amax.items()} or None
|
||||
except Exception as e: # noqa: BLE001
|
||||
logging.warning(f"[FunPack NVFP4] could not read calibration sidecar: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def inject_input_scales(sd, target_weight_keys, calibration, margin=CALIBRATION_MARGIN):
|
||||
"""Add `<layer>.input_scale` next to each quantized weight whose module path has a
|
||||
calibrated amax. Keys are matched by suffix so any checkpoint prefix
|
||||
('model.diffusion_model.' or none) works. Returns the number injected."""
|
||||
injected = 0
|
||||
for k in target_weight_keys:
|
||||
prefix = k[: -len("weight")] # '...transformer_blocks.3.attn1.to_q.'
|
||||
module_path = prefix.rstrip(".")
|
||||
for cal_key, amax in calibration.items():
|
||||
if module_path == cal_key or module_path.endswith("." + cal_key):
|
||||
sd[prefix + "input_scale"] = torch.tensor(
|
||||
float(amax) * margin / _F8F4_MAX, dtype=torch.float32)
|
||||
injected += 1
|
||||
break
|
||||
return injected
|
||||
|
||||
|
||||
def install_calibration_capture(model, checkpoint_path):
|
||||
"""Forward-pre-hooks on every quantized Linear recording the running input amax.
|
||||
The sidecar is written on the first full sweep and refreshed periodically; hooks
|
||||
live for the model instance's lifetime (a calibration run is expected to be
|
||||
slightly SLOWER - it pays the amax it is trying to eliminate, once)."""
|
||||
store = {"amax": {}, "writes": 0, "calls": 0}
|
||||
path = _scales_path(checkpoint_path)
|
||||
|
||||
def _flush():
|
||||
try:
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
tmp = path + ".tmp"
|
||||
with open(tmp, "w", encoding="utf-8") as f:
|
||||
json.dump({"amax": store["amax"], "margin": CALIBRATION_MARGIN}, f)
|
||||
os.replace(tmp, path)
|
||||
store["writes"] += 1
|
||||
except Exception as e: # noqa: BLE001
|
||||
logging.warning(f"[FunPack NVFP4] calibration flush failed: {e}")
|
||||
|
||||
hooked = 0
|
||||
try:
|
||||
for name, m in model.model.diffusion_model.named_modules():
|
||||
if getattr(m, "quant_format", None) != "nvfp4":
|
||||
continue
|
||||
|
||||
def _hook(_module, args, _name=name):
|
||||
try:
|
||||
x = args[0] if args else None
|
||||
if isinstance(x, torch.Tensor):
|
||||
a = float(x.detach().abs().amax())
|
||||
if a > store["amax"].get(_name, 0.0):
|
||||
store["amax"][_name] = a
|
||||
store["calls"] += 1
|
||||
if store["calls"] % 4096 == 0:
|
||||
_flush()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
m.register_forward_pre_hook(_hook)
|
||||
hooked += 1
|
||||
except Exception as e: # noqa: BLE001
|
||||
logging.warning(f"[FunPack NVFP4] calibration capture install failed: {e}")
|
||||
return 0
|
||||
|
||||
if hooked:
|
||||
# First flush after one full model call so a single generation is enough.
|
||||
def _first_flush_hook(_module, _args):
|
||||
if store["writes"] == 0 and len(store["amax"]) >= hooked:
|
||||
_flush()
|
||||
try:
|
||||
last = [m for _, m in model.model.diffusion_model.named_modules()
|
||||
if getattr(m, "quant_format", None) == "nvfp4"][-1]
|
||||
last.register_forward_pre_hook(_first_flush_hook)
|
||||
except Exception:
|
||||
pass
|
||||
return hooked
|
||||
|
||||
QUANTIZE_SCOPES = [
|
||||
"video blocks", # video-branch Linears in middle blocks (default, safest)
|
||||
"video + cross-modal", # + audio_to_video_attn / video_to_audio_attn bridges
|
||||
"all blocks", # every Linear inside transformer_blocks (incl. audio branch)
|
||||
"all 2D layers", # + eligible Linears outside the blocks (most aggressive)
|
||||
]
|
||||
|
||||
# Video-branch submodules inside a BasicAVTransformerBlock (comfy/ldm/lightricks/av_model.py).
|
||||
_VIDEO_MODULES = ("attn1.", "attn2.", "ff.")
|
||||
_CROSS_MODULES = ("audio_to_video_attn.", "video_to_audio_attn.")
|
||||
|
||||
# Matches "...transformer_blocks.<idx>.<rest>" with any leading prefix
|
||||
# ("model.diffusion_model." or none - load_diffusion_model_state_dict strips prefixes
|
||||
# from ALL keys uniformly, so our metadata keys ride along).
|
||||
_BLOCK_RE = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.(.*)$")
|
||||
|
||||
# Outside the blocks ("all 2D layers" scope) never touch input/output-adjacent or
|
||||
# modulation weights even when they are 2D Linears.
|
||||
_GENERIC_SKIP = ("proj_out", "adaln", "embed", "emb.", "norm", "scale_shift", "patchify")
|
||||
|
||||
# Below this, quantization overhead outweighs any GEMM win and relative error is worst.
|
||||
_MIN_DIM = 512
|
||||
|
||||
|
||||
def _block_ids(sd) -> list[int]:
|
||||
ids = set()
|
||||
for k in sd:
|
||||
m = _BLOCK_RE.search(k)
|
||||
if m:
|
||||
ids.add(int(m.group(1)))
|
||||
return sorted(ids)
|
||||
|
||||
|
||||
def _selected(key_body: str, scope: str, keep: set[int]) -> bool:
|
||||
"""Decide whether the Linear at `key_body` (state-dict key minus '.weight') is
|
||||
quantized under `scope`, given the high-precision `keep` block set."""
|
||||
m = _BLOCK_RE.search(key_body + ".")
|
||||
if not m:
|
||||
if scope != "all 2D layers":
|
||||
return False
|
||||
low = key_body.lower()
|
||||
return not any(s in low for s in _GENERIC_SKIP)
|
||||
idx, rest = int(m.group(1)), m.group(2)
|
||||
if idx in keep:
|
||||
return False
|
||||
if scope in ("all blocks", "all 2D layers"):
|
||||
return True
|
||||
if any(rest.startswith(p) for p in _VIDEO_MODULES):
|
||||
return True
|
||||
if scope == "video + cross-modal" and any(rest.startswith(p) for p in _CROSS_MODULES):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _eligible(key: str, w) -> bool:
|
||||
return (
|
||||
key.endswith(".weight")
|
||||
and isinstance(w, torch.Tensor)
|
||||
and w.dim() == 2
|
||||
and w.shape[0] >= _MIN_DIM
|
||||
and w.shape[1] >= _MIN_DIM
|
||||
and w.shape[1] % 16 == 0 # NVFP4 blocks run along in_features
|
||||
)
|
||||
|
||||
|
||||
def quantize_state_dict_nvfp4(sd, scope="video blocks", keep_first_blocks=2,
|
||||
keep_last_blocks=2, device=None, progress_callback=None):
|
||||
"""Return (new_sd, quantized_layer_names, kept_block_ids).
|
||||
|
||||
Pure transform: selected 2D weights are replaced by NVFP4 QuantizedTensor
|
||||
serialization (weight uint8-packed, weight_scale fp8-e4m3 block scales,
|
||||
weight_scale_2 fp32 tensor scale) plus a per-layer `comfy_quant` metadata key -
|
||||
byte-for-byte the format `comfy.ops._load_quantized_module` consumes. Unselected
|
||||
keys pass through untouched.
|
||||
"""
|
||||
import comfy.quant_ops as quant_ops
|
||||
|
||||
if not quant_ops._CK_AVAILABLE:
|
||||
raise RuntimeError(
|
||||
"comfy_kitchen is unavailable in this ComfyUI - NVFP4 quantization needs it. "
|
||||
"Update ComfyUI / reinstall requirements."
|
||||
)
|
||||
if scope not in QUANTIZE_SCOPES:
|
||||
raise ValueError(f"Unknown quantize scope: {scope}")
|
||||
for k in sd:
|
||||
if k.endswith(".comfy_quant") or k.endswith("scaled_fp8"):
|
||||
raise ValueError(
|
||||
"This checkpoint is already quantized (comfy_quant / scaled_fp8 metadata "
|
||||
"found). NVFP4 patching needs the original bf16/fp16 weights."
|
||||
)
|
||||
|
||||
ids = _block_ids(sd)
|
||||
keep: set[int] = set()
|
||||
if ids:
|
||||
if keep_first_blocks > 0:
|
||||
keep |= set(ids[:keep_first_blocks])
|
||||
if keep_last_blocks > 0:
|
||||
keep |= set(ids[-keep_last_blocks:])
|
||||
|
||||
targets = [k for k, w in sd.items() if _eligible(k, w) and _selected(k[: -len(".weight")], scope, keep)]
|
||||
|
||||
meta = torch.tensor(list(json.dumps({"format": "nvfp4"}).encode("utf-8")), dtype=torch.uint8)
|
||||
out = {}
|
||||
done = 0
|
||||
for k, w in sd.items():
|
||||
if k not in targets:
|
||||
out[k] = w
|
||||
continue
|
||||
wq = w
|
||||
if wq.dtype not in (torch.float32, torch.bfloat16):
|
||||
# fp16 and fp8-stored checkpoints alike: quantize from bf16.
|
||||
wq = wq.to(torch.bfloat16)
|
||||
if device is not None:
|
||||
wq = wq.to(device)
|
||||
qt = quant_ops.QuantizedTensor.from_float(wq, "TensorCoreNVFP4Layout", scale="recalculate")
|
||||
for kk, vv in qt.state_dict(k).items():
|
||||
out[kk] = vv.to("cpu")
|
||||
out[k[: -len("weight")] + "comfy_quant"] = meta.clone()
|
||||
del qt, wq
|
||||
done += 1
|
||||
if progress_callback is not None:
|
||||
progress_callback(done, len(targets))
|
||||
if done % 100 == 0:
|
||||
logging.info(f"[FunPack NVFP4] quantized {done}/{len(targets)} layers…")
|
||||
|
||||
return out, targets, keep
|
||||
|
||||
|
||||
def log_nvfp4_diagnostics(model, device=None):
|
||||
"""Print everything needed to tell WHY NVFP4 is (or is not) fast on this machine.
|
||||
|
||||
'Loaded fine' and 'runs fast' are different claims. The three ways the fast path
|
||||
silently degrades, in the order they bite:
|
||||
1. comfy_kitchen's CUDA backend is disabled when torch is built for CUDA < 13 -
|
||||
quantized matmuls then fall through to the pure-PyTorch eager backend, which
|
||||
is CORRECT but no faster than bf16 (the exact "speed didn't change" symptom).
|
||||
2. The GPU lacks FP4 tensor cores (pre-Blackwell) - core flags every nvfp4 layer
|
||||
_full_precision_mm and dequantizes (visible in the layer census below).
|
||||
3. Everything is native but the quantized scope is too small a slice of the
|
||||
total step time to move the needle.
|
||||
A timed 4096x4096 GEMM probe settles it empirically: on a native FP4 path the
|
||||
quantized matmul must beat bf16 clearly; parity or worse means emulation."""
|
||||
lines = []
|
||||
try:
|
||||
import comfy_kitchen as ck
|
||||
for name, info in ck.list_backends().items():
|
||||
if info.get("available"):
|
||||
caps = info.get("capabilities") or []
|
||||
lines.append(f" ck backend '{name}': AVAILABLE, scaled_mm_nvfp4="
|
||||
f"{'yes' if 'scaled_mm_nvfp4' in caps else 'NO'}")
|
||||
else:
|
||||
lines.append(f" ck backend '{name}': unavailable ({info.get('unavailable_reason')})")
|
||||
except Exception as e: # noqa: BLE001
|
||||
lines.append(f" comfy_kitchen introspection failed: {e}")
|
||||
cap = torch.cuda.get_device_capability(device) if torch.cuda.is_available() else None
|
||||
lines.append(f" torch {torch.__version__} | cuda build {torch.version.cuda} | device capability {cap}")
|
||||
if torch.version.cuda is not None:
|
||||
try:
|
||||
if tuple(map(int, str(torch.version.cuda).split(".")))[0] < 13:
|
||||
lines.append(" ⚠ torch is built for CUDA < 13 -> comfy_kitchen DISABLES its CUDA "
|
||||
"backend -> nvfp4 matmuls run on the eager (pure PyTorch) backend: "
|
||||
"correct but NO speedup. Install a cu13x torch build to unlock FP4 GEMMs.")
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
n_q = n_emul = 0
|
||||
try:
|
||||
for m in model.model.diffusion_model.modules():
|
||||
if getattr(m, "quant_format", None) == "nvfp4":
|
||||
n_q += 1
|
||||
if getattr(m, "_full_precision_mm", False):
|
||||
n_emul += 1
|
||||
lines.append(f" nvfp4 layers in loaded model: {n_q}"
|
||||
+ (f" — {n_emul} flagged FULL-PRECISION EMULATION (no FP4 compute on this GPU)"
|
||||
if n_emul else " — all on the quantized matmul path"))
|
||||
except Exception as e: # noqa: BLE001
|
||||
lines.append(f" layer census failed: {e}")
|
||||
|
||||
if torch.cuda.is_available():
|
||||
try:
|
||||
import time
|
||||
import comfy.quant_ops as quant_ops
|
||||
d = torch.device(device) if device is not None else torch.device("cuda")
|
||||
w = torch.randn(4096, 4096, dtype=torch.bfloat16, device=d)
|
||||
qw = quant_ops.QuantizedTensor.from_float(w, "TensorCoreNVFP4Layout", scale="recalculate")
|
||||
|
||||
def _timed(fn, n=30):
|
||||
for _ in range(5):
|
||||
fn()
|
||||
torch.cuda.synchronize(d)
|
||||
t0 = time.perf_counter()
|
||||
for _ in range(n):
|
||||
fn()
|
||||
torch.cuda.synchronize(d)
|
||||
return (time.perf_counter() - t0) / n
|
||||
|
||||
# Decomposed probe at two token counts: M=1024 (small scene chunk) and
|
||||
# M=16384 (realistic video token count). Times the input quantization and
|
||||
# the FP4 GEMM SEPARATELY — "no speedup" has two very different fixes
|
||||
# depending on which one is at fault:
|
||||
# * GEMM alone beats bf16, quantization dominates -> per-call overhead
|
||||
# problem (x1232 layers/step); needs static input scales or fusion.
|
||||
# * GEMM alone loses to bf16 -> the FP4 kernel isn't engaging at all.
|
||||
results = {}
|
||||
for m in (1024, 16384):
|
||||
x = torch.randn(m, 4096, dtype=torch.bfloat16, device=d)
|
||||
t_bf16 = _timed(lambda: torch.nn.functional.linear(x, w))
|
||||
t_quant = _timed(lambda: quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout"))
|
||||
# Same quantization with a PREcomputed static scale: isolates the amax
|
||||
# reduction (which per-layer input_scale calibration could remove) from
|
||||
# tensor-construction/kernel-launch overhead (which it can't).
|
||||
s = (x.abs().amax() / (448.0 * 6.0)).float()
|
||||
t_quant_s = _timed(lambda: quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout", scale=s))
|
||||
qx = quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout")
|
||||
t_gemm = _timed(lambda: torch.nn.functional.linear(qx, qw))
|
||||
total = t_quant + t_gemm
|
||||
results[m] = (t_bf16, t_quant, t_quant_s, t_gemm)
|
||||
lines.append(
|
||||
f" probe 4096x4096 @ M={m}: bf16 {t_bf16 * 1e3:.3f} ms | "
|
||||
f"input-quant {t_quant * 1e3:.3f} ms (static-scale {t_quant_s * 1e3:.3f} ms) "
|
||||
f"+ fp4-gemm {t_gemm * 1e3:.3f} ms = {total * 1e3:.3f} ms "
|
||||
f"({t_bf16 / max(total, 1e-9):.2f}x net, {t_bf16 / max(t_gemm, 1e-9):.2f}x gemm-only)")
|
||||
del x, qx
|
||||
|
||||
# Verdict from the LARGE-M run (realistic video token counts); small-M numbers
|
||||
# are launch-overhead dominated by nature and must not drive the conclusion.
|
||||
t_bf16, t_quant, t_quant_s, t_gemm = results[16384]
|
||||
total = t_quant + t_gemm
|
||||
if t_gemm < t_bf16 * 0.8 and total < t_bf16 * 0.8:
|
||||
lines.append(" verdict: FAST PATH ACTIVE — FP4 GEMMs win at real token counts. "
|
||||
"If end-to-end speed still looks flat, the input-quant overhead "
|
||||
f"(~{t_quant * 1e3:.2f} ms/layer-call, largely size-independent) is "
|
||||
"cancelling the win on SMALL layers: prefer scope 'video blocks' "
|
||||
"(skips the audio branch, whose matmuls are too small to profit).")
|
||||
if t_quant_s < t_quant * 0.6:
|
||||
lines.append(" note: static input scales would reclaim most of the input-quant "
|
||||
"cost (amax dominates) — calibration support is worth building.")
|
||||
else:
|
||||
lines.append(" note: static input scales would NOT help much "
|
||||
"(overhead is construction/launch, not amax).")
|
||||
elif t_gemm < t_bf16 * 0.8:
|
||||
lines.append(" verdict: FP4 GEMM fast but INPUT QUANTIZATION dominates even at "
|
||||
"large M — overhead-bound; scope down to the largest layers only.")
|
||||
else:
|
||||
lines.append(" verdict: FP4 GEMM not faster than bf16 even at M=16384 — the "
|
||||
"kernel is not engaging; report these numbers upstream.")
|
||||
except Exception as e: # noqa: BLE001
|
||||
lines.append(f" GEMM probe failed: {e}")
|
||||
else:
|
||||
lines.append(" GEMM probe skipped (no CUDA device)")
|
||||
|
||||
print("[FunPack NVFP4] diagnostics:\n" + "\n".join(lines))
|
||||
return {"nvfp4_layers": n_q, "emulated_layers": n_emul}
|
||||
|
||||
|
||||
class FunPackNVFP4ModelLoader:
|
||||
"""EXPERIMENTAL: load a diffusion model with on-the-fly NVFP4 quantization.
|
||||
|
||||
Drop-in replacement for 'Load Diffusion Model' aimed at LTX-AV on Blackwell
|
||||
(sm_120+): selected transformer Linears run real FP4 tensor-core GEMMs (~2x fp8
|
||||
throughput) and weight VRAM drops to ~1/4 of bf16. Audio branch, cross-modal
|
||||
bridges, first/last blocks, and all non-matmul weights stay at full precision
|
||||
by default. On GPUs without FP4 compute the model still works via emulation.
|
||||
"""
|
||||
|
||||
CATEGORY = "FunPack/Experimental"
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "load"
|
||||
DESCRIPTION = __doc__
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
import folder_paths
|
||||
return {
|
||||
"required": {
|
||||
"unet_name": (folder_paths.get_filename_list("diffusion_models"),
|
||||
{"tooltip": "Original bf16/fp16 checkpoint - NOT a pre-quantized one."}),
|
||||
"quantize_scope": (QUANTIZE_SCOPES, {
|
||||
"default": "video blocks",
|
||||
"tooltip": "Which Linears go NVFP4. 'video blocks' keeps the whole audio "
|
||||
"path + cross-modal bridges at full precision (safest). Widen "
|
||||
"only if quality holds.",
|
||||
}),
|
||||
"keep_first_blocks": ("INT", {"default": 2, "min": 0, "max": 16,
|
||||
"tooltip": "First N transformer blocks stay full precision."}),
|
||||
"keep_last_blocks": ("INT", {"default": 2, "min": 0, "max": 16,
|
||||
"tooltip": "Last N transformer blocks stay full precision."}),
|
||||
"input_scales": (INPUT_SCALE_MODES, {
|
||||
"default": "dynamic",
|
||||
"tooltip": "auto-calibrate: first run records each layer's input range "
|
||||
"(slightly slower) into a per-checkpoint sidecar; every later "
|
||||
"load injects static input scales, removing the per-call amax "
|
||||
"cost (~0.2 ms x every quantized layer x every step). Reload "
|
||||
"the model after the calibration run to activate.",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
def load(self, unet_name, quantize_scope, keep_first_blocks, keep_last_blocks,
|
||||
input_scales="dynamic"):
|
||||
import comfy.model_management as mm
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
|
||||
path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)
|
||||
# return_metadata is NOT optional here: LTX-AV variants carry their architecture
|
||||
# config in the safetensors metadata, and detection without it builds the WRONG
|
||||
# model (6 vs 9 ada params, connector dims) -> size-mismatch wall on load.
|
||||
sd, metadata = comfy.utils.load_torch_file(path, return_metadata=True)
|
||||
|
||||
if mm.supports_nvfp4_compute(mm.get_torch_device()):
|
||||
logging.info("[FunPack NVFP4] FP4 tensor-core compute available - native GEMM path.")
|
||||
else:
|
||||
logging.warning(
|
||||
"[FunPack NVFP4] This GPU has no FP4 compute (needs Blackwell, sm_100+). "
|
||||
"The model will load and run through dequant emulation: correct output, "
|
||||
"no speedup. Use this only for A/B quality checks."
|
||||
)
|
||||
|
||||
device = mm.get_torch_device() if torch.cuda.is_available() else None
|
||||
pbar = comfy.utils.ProgressBar(1)
|
||||
|
||||
def _tick(done, total):
|
||||
if total:
|
||||
pbar.total = total
|
||||
pbar.update_absolute(done)
|
||||
|
||||
sd, targets, keep = quantize_state_dict_nvfp4(
|
||||
sd, scope=quantize_scope,
|
||||
keep_first_blocks=keep_first_blocks, keep_last_blocks=keep_last_blocks,
|
||||
device=device, progress_callback=_tick,
|
||||
)
|
||||
if not targets:
|
||||
raise RuntimeError(
|
||||
"No layers matched the NVFP4 quantization scope - is this an LTX-AV "
|
||||
"diffusion model checkpoint? (Expected transformer_blocks.* keys.)"
|
||||
)
|
||||
logging.info(
|
||||
f"[FunPack NVFP4] quantized {len(targets)} Linear layers "
|
||||
f"(scope='{quantize_scope}', full-precision blocks: {sorted(keep) or 'none'})."
|
||||
)
|
||||
|
||||
calibration = load_calibration(path) if input_scales == "auto-calibrate" else None
|
||||
if calibration:
|
||||
injected = inject_input_scales(sd, targets, calibration)
|
||||
logging.info(
|
||||
f"[FunPack NVFP4] static input scales ACTIVE on {injected}/{len(targets)} layers "
|
||||
f"(calibrated {len(calibration)}, margin {CALIBRATION_MARGIN}x) — per-call input "
|
||||
"amax eliminated. Delete the sidecar under refinements/nvfp4_scales/ to recalibrate."
|
||||
)
|
||||
|
||||
model = comfy.sd.load_diffusion_model_state_dict(sd, metadata=metadata)
|
||||
if model is None:
|
||||
raise RuntimeError("ComfyUI could not detect the model type after NVFP4 patching.")
|
||||
|
||||
if input_scales == "auto-calibrate" and not calibration:
|
||||
hooked = install_calibration_capture(model, path)
|
||||
logging.info(
|
||||
f"[FunPack NVFP4] CALIBRATING input scales on {hooked} layers this session "
|
||||
"(slightly slower). Run one full generation, then reload the model (change any "
|
||||
"widget or restart) — static scales activate on the next load."
|
||||
)
|
||||
|
||||
try:
|
||||
log_nvfp4_diagnostics(model, device)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logging.warning(f"[FunPack NVFP4] diagnostics failed: {e}")
|
||||
return (model,)
|
||||
-91
@@ -2281,14 +2281,6 @@ class FunPackLTXAVSceneChainSampler:
|
||||
"default": 0.02, "min": 0.005, "max": 0.1, "step": 0.005,
|
||||
"tooltip": "Per-step correction strength applied to the model's predicted output. Same scale/units as embed_guidance_strength — start there and adjust.",
|
||||
}),
|
||||
"block_steer": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "EXPERIMENTAL: RLHF at the model level. Measures which transformer blocks do the work each run (per-block video-residual activity, ~free), pairs those fingerprints with your ratings (plus per-step value-function credit when output_guidance's VF is trained), and on later runs redistributes emphasis: blocks correlated with your Perfects get boosted, blocks correlated with Awfuls damped (video residual only — audio untouched, hard ±10% cap). Learns automatically while on; exact no-op until 2+ rated runs (or 2+ value-credit runs) for this key. Requires refinement_key_input.",
|
||||
}),
|
||||
"block_steer_strength": ("FLOAT", {
|
||||
"default": 0.05, "min": 0.01, "max": 0.15, "step": 0.005,
|
||||
"tooltip": "Max per-block residual gain deviation (scores are normalized to ±1, so 0.05 = up to ±5% per block; hard-capped at ±10% regardless). Start low — this reshapes the model's internal emphasis, not the conditioning.",
|
||||
}),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
@@ -3930,7 +3922,6 @@ class FunPackLTXAVSceneChainSampler:
|
||||
alg_blur_guides=False,
|
||||
bounded_attention_enabled=False,
|
||||
output_guidance=False, output_guidance_strength=0.02,
|
||||
block_steer=False, block_steer_strength=0.05,
|
||||
unique_id=None, prompt=None):
|
||||
if not isinstance(positive, list) or not positive:
|
||||
raise ValueError("positive conditioning must contain at least one scene entry.")
|
||||
@@ -4036,65 +4027,6 @@ class FunPackLTXAVSceneChainSampler:
|
||||
print(f"[FunPackSceneChain] output_guidance: active ({_output_value_fn.n_trained} samples), "
|
||||
f"strength={output_guidance_strength}")
|
||||
|
||||
# Block Steering (see block_steer.py): run-level machinery, installed once for the
|
||||
# whole chain. The recorder + steer hooks are tagged module hooks (the defensive
|
||||
# strip above cleans any leak); the value-credit wrapper goes in INNERMOST so all
|
||||
# per-scene wrappers (embed/score/output guidance, temporal) stack on top and see
|
||||
# it as their old_wrapper — it scores the base model's own prediction.
|
||||
_bs_recorder = None
|
||||
_bs_handles = []
|
||||
_bs_wrapper_old = None
|
||||
_bs_wrapper_installed = False
|
||||
if block_steer:
|
||||
try:
|
||||
try:
|
||||
from . import block_steer as _bsmod
|
||||
except ImportError:
|
||||
import block_steer as _bsmod
|
||||
if not refinement_key_input:
|
||||
print("[FunPackSceneChain] block_steer: no refinement_key_input wired — "
|
||||
"nothing to learn into or steer from, skipping.")
|
||||
else:
|
||||
# Recover from a credit wrapper leaked by a crashed previous run.
|
||||
_leaked = model.model_options.get("model_function_wrapper")
|
||||
if getattr(_leaked, "_funpack_block_steer", False):
|
||||
_old = getattr(_leaked, "_funpack_old", None)
|
||||
if _old is not None:
|
||||
model.model_options["model_function_wrapper"] = _old
|
||||
else:
|
||||
del model.model_options["model_function_wrapper"]
|
||||
print("[FunPackSceneChain] block_steer: stripped a leaked credit wrapper from a previous run")
|
||||
_bs_blocks = _bsmod._funpack_locate_blocks(model)
|
||||
if _bs_blocks is None:
|
||||
print("[FunPackSceneChain] block_steer: couldn't locate transformer blocks — skipping.")
|
||||
else:
|
||||
_bs_recorder = _bsmod.BlockActivityRecorder(len(_bs_blocks))
|
||||
_bs_handles = _bsmod.install_recorder(model, _bs_recorder)
|
||||
_bs_scores, _bs_status = _bsmod.block_scores_with_status(refinement_key_input)
|
||||
if _bs_scores is not None and _bs_scores.numel() != len(_bs_blocks):
|
||||
_bs_scores, _bs_status = None, (
|
||||
f"profile was learned on a {_bs_scores.numel()}-block model, this one has "
|
||||
f"{len(_bs_blocks)} — relearning")
|
||||
if _bs_scores is not None:
|
||||
_gains = _bsmod.gains_from_scores(_bs_scores, block_steer_strength)
|
||||
_bs_handles += _bsmod.install_steer(model, _gains)
|
||||
_dev = max(abs(g - 1.0) for g in _gains) if _gains else 0.0
|
||||
print(f"[FunPackSceneChain] block_steer: steering active "
|
||||
f"({len(_bs_blocks)} blocks, max gain deviation ±{_dev:.3f}, {_bs_status})")
|
||||
else:
|
||||
print(f"[FunPackSceneChain] block_steer: learning only — {_bs_status}; "
|
||||
"steering is a no-op this run.")
|
||||
_bs_vf = _output_value_fn or self._load_output_value_function(refinement_key_input)
|
||||
if _bs_vf is not None:
|
||||
_bs_wrapper_old = model.model_options.get("model_function_wrapper")
|
||||
_w = _bsmod.make_credit_wrapper(_bs_wrapper_old, _bs_recorder, _bs_vf)
|
||||
_w._funpack_block_steer = True
|
||||
_w._funpack_old = _bs_wrapper_old
|
||||
model.model_options["model_function_wrapper"] = _w
|
||||
_bs_wrapper_installed = True
|
||||
except Exception as e:
|
||||
print(f"[FunPackSceneChain] block_steer setup failed (continuing without): {e}")
|
||||
|
||||
# Phase timing + quantized-path census: "the matmuls got faster" and "the video
|
||||
# arrives sooner" are different claims — the report separates sampling, decode,
|
||||
# and everything else so an optimization's real ceiling is visible. The census
|
||||
@@ -4450,29 +4382,6 @@ class FunPackLTXAVSceneChainSampler:
|
||||
except Exception as e:
|
||||
print(f"[FunPackLTXAVSceneChainSampler] Failed to write sampler context: {e}")
|
||||
|
||||
# Block Steering teardown: drop hooks, restore the wrapper chain, and stage this
|
||||
# run's activity fingerprint for the NEXT rating cycle (same stage-then-pair
|
||||
# convention as the x0 snapshot below).
|
||||
if _bs_recorder is not None:
|
||||
try:
|
||||
try:
|
||||
from . import block_steer as _bsmod
|
||||
except ImportError:
|
||||
import block_steer as _bsmod
|
||||
_bsmod.remove_handles(_bs_handles)
|
||||
if _bs_wrapper_installed:
|
||||
if _bs_wrapper_old is not None:
|
||||
model.model_options["model_function_wrapper"] = _bs_wrapper_old
|
||||
elif "model_function_wrapper" in model.model_options:
|
||||
del model.model_options["model_function_wrapper"]
|
||||
if _bsmod.save_run_snapshot(refinement_key_input, _bs_recorder):
|
||||
_cs = _bs_recorder.credit_steps
|
||||
print(f"[FunPackSceneChain] block_steer: staged activity fingerprint "
|
||||
f"({int(_bs_recorder.counts.max())} calls"
|
||||
+ (f", value credit from {_cs} steps" if _cs else "") + ")")
|
||||
except Exception as e:
|
||||
print(f"[FunPackSceneChain] block_steer teardown failed: {e}")
|
||||
|
||||
# Snapshot the final output for output_guidance's NEXT training cycle (the rating that
|
||||
# scores THIS run pairs with this snapshot, not with denoised — see
|
||||
# _save_output_value_snapshot). Independent of the output_guidance toggle itself: a
|
||||
|
||||
@@ -1,239 +0,0 @@
|
||||
"""Block Steering: recorder/steer hooks on a toy AV model + profile/score/gain math.
|
||||
|
||||
The toy model mirrors the LTX-AV calling convention the hooks rely on: blocks held in
|
||||
`model.model.diffusion_model.transformer_blocks` (a ModuleList of >= 28 entries, which
|
||||
is what _funpack_locate_blocks requires), each called as block((vx, ax), ...) and
|
||||
returning (vx_out, ax_out). Persistence tests point the refinement store at a tmp dir.
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
import block_steer as bs # noqa: E402
|
||||
|
||||
N_BLOCKS = 28 # _funpack_locate_blocks' minimum believable transformer depth
|
||||
DIM = 8
|
||||
|
||||
|
||||
class _ToyAVBlock(torch.nn.Module):
|
||||
def __init__(self, v_gain, a_gain=0.5):
|
||||
super().__init__()
|
||||
self.v_gain = v_gain
|
||||
self.a_gain = a_gain
|
||||
|
||||
def forward(self, x, **kwargs):
|
||||
vx, ax = x
|
||||
return vx + self.v_gain * torch.ones_like(vx), ax + self.a_gain * torch.ones_like(ax)
|
||||
|
||||
|
||||
def _toy_model(v_gains):
|
||||
diff = torch.nn.Module()
|
||||
diff.transformer_blocks = torch.nn.ModuleList([_ToyAVBlock(g) for g in v_gains])
|
||||
inner = types.SimpleNamespace(diffusion_model=diff)
|
||||
return types.SimpleNamespace(model=inner, model_options={})
|
||||
|
||||
|
||||
def _run(model, vx=None, ax=None):
|
||||
vx = vx if vx is not None else torch.ones(1, 4, DIM)
|
||||
ax = ax if ax is not None else torch.ones(1, 2, DIM)
|
||||
for blk in model.model.diffusion_model.transformer_blocks:
|
||||
vx, ax = blk((vx, ax))
|
||||
return vx, ax
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path, monkeypatch):
|
||||
"""Redirect refinement sidecars into a tmp dir (no conditioning.py import needed)."""
|
||||
def _path(key, name, prefix="refine_v2", extension="json"):
|
||||
return str(tmp_path / f"{key}.{name}.{extension}")
|
||||
monkeypatch.setattr(bs, "_state_path", lambda key, name: _path(key, name, extension="pt"))
|
||||
return tmp_path
|
||||
|
||||
|
||||
def test_recorder_ranks_blocks_by_video_activity():
|
||||
gains = [0.01] * N_BLOCKS
|
||||
gains[5], gains[20] = 2.0, 1.0 # blocks 5 and 20 move the video stream hardest
|
||||
model = _toy_model(gains)
|
||||
rec = bs.BlockActivityRecorder(N_BLOCKS)
|
||||
handles = bs.install_recorder(model, rec)
|
||||
assert len(handles) == N_BLOCKS
|
||||
_run(model)
|
||||
_run(model)
|
||||
bs.remove_handles(handles)
|
||||
fp = rec.fingerprint()
|
||||
assert fp is not None and int(rec.counts[0]) == 2
|
||||
assert int(fp.argmax()) == 5
|
||||
assert fp[20] > fp[3]
|
||||
# hooks removed - a third run must not record
|
||||
_run(model)
|
||||
assert int(rec.counts[0]) == 2
|
||||
|
||||
|
||||
def test_hooks_carry_leak_strip_tag():
|
||||
"""Recorder AND steer hooks must be tagged so the sampler's defensive
|
||||
strip_funpack_block_hooks cleans them if a run crashes mid-sample."""
|
||||
model = _toy_model([1.0] * N_BLOCKS)
|
||||
rec = bs.BlockActivityRecorder(N_BLOCKS)
|
||||
h1 = bs.install_recorder(model, rec)
|
||||
h2 = bs.install_steer(model, [1.05] * N_BLOCKS)
|
||||
try:
|
||||
for blk in model.model.diffusion_model.transformer_blocks:
|
||||
fns = list(blk._forward_hooks.values())
|
||||
assert fns and all(getattr(f, bs._FUNPACK_HOOK_TAG, False) for f in fns)
|
||||
finally:
|
||||
bs.remove_handles(h1 + h2)
|
||||
|
||||
|
||||
def test_steer_scales_video_residual_only():
|
||||
model = _toy_model([1.0] * N_BLOCKS)
|
||||
v_ref, a_ref = _run(model)
|
||||
gains = [1.0] * N_BLOCKS
|
||||
gains[3] = 1.10 # boost block 3's video contribution by 10%
|
||||
gains[7] = 0.90 # damp block 7's
|
||||
handles = bs.install_steer(model, gains)
|
||||
assert len(handles) == 2 # gain==1.0 blocks get NO hook (native path)
|
||||
v_out, a_out = _run(model)
|
||||
bs.remove_handles(handles)
|
||||
# video: each block adds +1; two blocks changed by +/-0.1 -> net exactly equal
|
||||
assert torch.allclose(v_out, v_ref, atol=1e-5) # +0.1 and -0.1 cancel
|
||||
gains[7] = 1.0
|
||||
handles = bs.install_steer(model, gains)
|
||||
v_out2, _ = _run(model)
|
||||
bs.remove_handles(handles)
|
||||
assert torch.allclose(v_out2 - v_ref, torch.full_like(v_out2, 0.10), atol=1e-5)
|
||||
# audio stream byte-identical under steering
|
||||
assert torch.equal(a_out, a_ref)
|
||||
|
||||
|
||||
def test_credit_wrapper_scores_and_credits():
|
||||
model = _toy_model([1.0] * N_BLOCKS)
|
||||
rec = bs.BlockActivityRecorder(N_BLOCKS)
|
||||
handles = bs.install_recorder(model, rec)
|
||||
|
||||
class _VF(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.scores = iter([0.1, 0.5, 0.2])
|
||||
def compress(self, x):
|
||||
return x.mean().unsqueeze(0)
|
||||
def forward(self, c):
|
||||
return torch.tensor([next(self.scores)])
|
||||
|
||||
def apply_fn(x, t, **c):
|
||||
_run(model) # blocks fire -> recorder.cur populated
|
||||
return x * 0.9
|
||||
|
||||
wrapper = bs.make_credit_wrapper(None, rec, _VF())
|
||||
args = {"input": torch.ones(1, 4, DIM), "timestep": torch.tensor([0.9]), "c": {}}
|
||||
for _ in range(3):
|
||||
wrapper(apply_fn, args)
|
||||
bs.remove_handles(handles)
|
||||
assert rec.credit_steps == 2 # first call only sets the baseline
|
||||
cv = rec.credit_vector()
|
||||
assert cv is not None and cv.shape == (N_BLOCKS,)
|
||||
# In the toy model vx grows by +1 per block, so relative deltas DECREASE with depth:
|
||||
# early blocks are the most active. Net score delta across the run is positive
|
||||
# (0.1 -> 0.5 -> 0.2), so credit must follow the centered activity gradient -
|
||||
# positive for the most active (early) blocks, negative for the least active (late).
|
||||
assert float(cv[0]) > 0 > float(cv[-1])
|
||||
assert int(cv.argmax()) == 0
|
||||
|
||||
|
||||
def _rec_with(activity):
|
||||
r = bs.BlockActivityRecorder(activity.numel())
|
||||
r.sums = activity.clone().float(); r.counts = torch.ones(activity.numel())
|
||||
return r
|
||||
|
||||
|
||||
def test_snapshot_profile_scores_gains_roundtrip(store):
|
||||
key = "toykey"
|
||||
base = torch.full((N_BLOCKS,), 0.5)
|
||||
hot = base.clone(); hot[10] = 1.5 # block 10 fires hard on the GOOD run
|
||||
cold = base.clone(); cold[2] = 1.5 # block 2 fires hard on the BAD run
|
||||
|
||||
assert bs.save_run_snapshot(key, _rec_with(hot))
|
||||
assert bs.update_profile_with_rating(key, 1.0) == 1 # Perfect
|
||||
assert bs.update_profile_with_rating(key, 1.0) is None # snapshot consumed - no double pair
|
||||
scores, status = bs.block_scores_with_status(key)
|
||||
assert scores is None and "have 1" in status # 1 rated run < MIN_RATED_RUNS
|
||||
|
||||
assert bs.save_run_snapshot(key, _rec_with(cold))
|
||||
assert bs.update_profile_with_rating(key, -0.9) == 2 # Awful
|
||||
scores, status = bs.block_scores_with_status(key)
|
||||
assert scores is None and "have 2" in status # correlation needs 3+
|
||||
|
||||
assert bs.save_run_snapshot(key, _rec_with(base))
|
||||
assert bs.update_profile_with_rating(key, 0.2) == 3 # neutral third run
|
||||
scores, status = bs.block_scores_with_status(key)
|
||||
assert status.startswith("ready")
|
||||
assert "corr" in status and "confidence" in status
|
||||
assert scores is not None and scores.shape == (N_BLOCKS,)
|
||||
assert int(scores.argmax()) == 10 and int(scores.argmin()) == 2
|
||||
assert abs(float(scores.mean())) < 1e-5 # zero-mean: redistributes, no drift
|
||||
# confidence ramp: 3 runs -> young profile steers gently, not at unit strength
|
||||
conf = (3 - 2) / (bs.CONFIDENCE_FULL_AT - 2)
|
||||
assert float(scores.abs().max()) == pytest.approx(conf, abs=1e-4)
|
||||
|
||||
gains = bs.gains_from_scores(scores, 0.05)
|
||||
# gains follow the (asymmetric, zero-mean) scores exactly; the largest deviation
|
||||
# on either side is bounded by strength * confidence
|
||||
assert max(gains) == pytest.approx(1.0 + 0.05 * float(scores.max()), abs=1e-4)
|
||||
assert min(gains) == pytest.approx(1.0 + 0.05 * float(scores.min()), abs=1e-4)
|
||||
assert max(abs(g - 1.0) for g in gains) == pytest.approx(0.05 * conf, abs=1e-4)
|
||||
# hard cap regardless of strength widget
|
||||
gains_hi = bs.gains_from_scores(scores, 50.0)
|
||||
assert max(gains_hi) <= 1.0 + bs.MAX_GAIN_DELTA + 1e-6
|
||||
assert min(gains_hi) >= 1.0 - bs.MAX_GAIN_DELTA - 1e-6
|
||||
|
||||
|
||||
def test_same_pole_ratings_need_contrast_not_more_runs(store):
|
||||
"""The v1 regression: MANY rated runs, all rated alike -> the profile must say WHY
|
||||
it can't steer (no reward contrast), not pretend it needs more runs."""
|
||||
key = "samepole"
|
||||
for _ in range(16):
|
||||
assert bs.save_run_snapshot(key, _rec_with(torch.rand(N_BLOCKS)))
|
||||
assert bs.update_profile_with_rating(key, 1.0) is not None
|
||||
scores, status = bs.block_scores_with_status(key)
|
||||
assert scores is None
|
||||
assert "rated alike" in status and "16" in status
|
||||
# one contrasting rating unlocks attribution
|
||||
assert bs.save_run_snapshot(key, _rec_with(torch.rand(N_BLOCKS)))
|
||||
bs.update_profile_with_rating(key, -0.9)
|
||||
scores, status = bs.block_scores_with_status(key)
|
||||
assert scores is not None and status.startswith("ready")
|
||||
|
||||
|
||||
def test_every_reward_value_contributes(store):
|
||||
"""No dead band: mild rewards like Missing action (0.05) / Wrong details (0.20)
|
||||
count with their actual value instead of being dropped."""
|
||||
key = "mild"
|
||||
base = torch.full((N_BLOCKS,), 0.5)
|
||||
hi = base.clone(); hi[7] = 1.2
|
||||
lo = base.clone(); lo[21] = 1.2
|
||||
bs.save_run_snapshot(key, _rec_with(hi)); bs.update_profile_with_rating(key, 0.35)
|
||||
bs.save_run_snapshot(key, _rec_with(lo)); bs.update_profile_with_rating(key, 0.05)
|
||||
bs.save_run_snapshot(key, _rec_with(base)); bs.update_profile_with_rating(key, 0.20)
|
||||
scores, status = bs.block_scores_with_status(key)
|
||||
assert status.startswith("ready")
|
||||
assert int(scores.argmax()) == 7 and int(scores.argmin()) == 21
|
||||
|
||||
|
||||
def test_depth_mismatch_and_v1_profile_reset(store):
|
||||
key = "depth"
|
||||
bs.save_run_snapshot(key, _rec_with(torch.rand(N_BLOCKS)))
|
||||
bs.update_profile_with_rating(key, 1.0)
|
||||
r2 = torch.rand(N_BLOCKS + 4) # different model depth
|
||||
bs.save_run_snapshot(key, _rec_with(r2))
|
||||
assert bs.update_profile_with_rating(key, 1.0) == 1 # fresh profile, not a mixed one
|
||||
prof = torch.load(bs._state_path(key, "block_profile"), weights_only=False)
|
||||
assert prof["n_blocks"] == N_BLOCKS + 4
|
||||
# a v1 pole-EMA profile on disk (liked_n, no history) is discarded, not mixed in
|
||||
torch.save({"n_blocks": N_BLOCKS, "liked": torch.rand(N_BLOCKS), "liked_n": 16},
|
||||
bs._state_path(key, "block_profile"))
|
||||
bs.save_run_snapshot(key, _rec_with(torch.rand(N_BLOCKS)))
|
||||
assert bs.update_profile_with_rating(key, 1.0) == 1
|
||||
@@ -1,241 +0,0 @@
|
||||
"""NVFP4 on-the-fly patching: transform format + scope selection + core-load roundtrip.
|
||||
|
||||
Uses the REAL comfy.quant_ops / comfy.ops from the local ComfyUI install (comfy_kitchen's
|
||||
eager backend quantizes and matmuls NVFP4 on CPU), so these tests prove byte-for-byte
|
||||
compatibility with the format `comfy.ops._load_quantized_module` consumes - not just our
|
||||
own invariants. Skipped wholesale if ComfyUI or comfy_kitchen is unavailable.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
_COMFY_ROOT = os.environ.get("COMFYUI_ROOT", str(Path.home() / "Documents" / "ComfyUI"))
|
||||
if _COMFY_ROOT not in sys.path and os.path.isdir(_COMFY_ROOT):
|
||||
sys.path.insert(0, _COMFY_ROOT)
|
||||
|
||||
quant_ops = pytest.importorskip("comfy.quant_ops")
|
||||
if not quant_ops._CK_AVAILABLE:
|
||||
pytest.skip("comfy_kitchen unavailable", allow_module_level=True)
|
||||
|
||||
import comfy.ops # noqa: E402
|
||||
|
||||
from nvfp4_patch import ( # noqa: E402
|
||||
QUANTIZE_SCOPES,
|
||||
_selected,
|
||||
quantize_state_dict_nvfp4,
|
||||
)
|
||||
|
||||
D = 512 # >= _MIN_DIM and a multiple of 16
|
||||
|
||||
|
||||
def _w(out_f=D, in_f=D):
|
||||
return torch.randn(out_f, in_f, dtype=torch.bfloat16) * 0.02
|
||||
|
||||
|
||||
def _fake_ltxav_sd(n_blocks=4, prefix=""):
|
||||
"""Minimal LTX-AV-shaped state dict: video attn/ff, audio branch, cross-modal
|
||||
bridges, plus non-block layers that must never be touched by block scopes."""
|
||||
sd = {}
|
||||
for i in range(n_blocks):
|
||||
b = f"{prefix}transformer_blocks.{i}."
|
||||
for mod in ("attn1", "attn2", "audio_attn1", "audio_to_video_attn", "video_to_audio_attn"):
|
||||
for lin in ("to_q", "to_k", "to_v", "to_out.0"):
|
||||
sd[f"{b}{mod}.{lin}.weight"] = _w()
|
||||
sd[f"{b}ff.net.0.proj.weight"] = _w(2 * D, D)
|
||||
sd[f"{b}ff.net.2.weight"] = _w(D, 2 * D)
|
||||
sd[f"{b}audio_ff.net.0.proj.weight"] = _w(2 * D, D)
|
||||
sd[f"{b}scale_shift_table"] = torch.randn(6, D, dtype=torch.bfloat16) # 2D but not .weight
|
||||
sd[f"{b}attn1.norm_q.weight"] = torch.randn(D, dtype=torch.bfloat16) # 1D
|
||||
sd[f"{prefix}patchify_proj.weight"] = _w()
|
||||
sd[f"{prefix}proj_out.weight"] = _w()
|
||||
sd[f"{prefix}adaln_single.linear.weight"] = _w(6 * D, D)
|
||||
return sd
|
||||
|
||||
|
||||
def _quant_layers(out_sd):
|
||||
return {k[: -len(".comfy_quant")] for k in out_sd if k.endswith(".comfy_quant")}
|
||||
|
||||
|
||||
def test_video_scope_protects_audio_cross_edges_and_nonblock():
|
||||
sd = _fake_ltxav_sd()
|
||||
out, targets, keep = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1)
|
||||
q = _quant_layers(out)
|
||||
assert keep == {0, 3}
|
||||
# middle-block video linears quantized
|
||||
assert "transformer_blocks.1.attn1.to_q" in q
|
||||
assert "transformer_blocks.2.ff.net.2" in q
|
||||
# first/last blocks untouched
|
||||
assert not any(l.startswith("transformer_blocks.0.") or l.startswith("transformer_blocks.3.") for l in q)
|
||||
# audio branch, cross-modal bridges, non-block layers untouched
|
||||
assert not any("audio" in l for l in q)
|
||||
assert not any(l.split(".")[-2:] == ["patchify_proj"] or "proj_out" in l or "adaln" in l for l in q)
|
||||
# untouched keys pass through identically
|
||||
assert torch.equal(out["transformer_blocks.0.attn1.to_q.weight"],
|
||||
sd["transformer_blocks.0.attn1.to_q.weight"])
|
||||
assert len(targets) == len(q)
|
||||
|
||||
|
||||
def test_scopes_widen_monotonically():
|
||||
sd = _fake_ltxav_sd()
|
||||
counts = []
|
||||
for scope in QUANTIZE_SCOPES:
|
||||
_, targets, _ = quantize_state_dict_nvfp4(sd, scope, 1, 1)
|
||||
counts.append(len(targets))
|
||||
assert counts == sorted(counts) and counts[0] < counts[-1]
|
||||
# cross-modal scope adds exactly the two bridges of the two middle blocks (4 linears each)
|
||||
assert counts[1] - counts[0] == 2 * 2 * 4
|
||||
|
||||
|
||||
def test_all_2d_scope_still_skips_sensitive_generic_layers():
|
||||
sd = _fake_ltxav_sd()
|
||||
out, _, _ = quantize_state_dict_nvfp4(sd, "all 2D layers", 0, 0)
|
||||
q = _quant_layers(out)
|
||||
for name in ("proj_out", "adaln_single.linear", "patchify_proj"):
|
||||
assert not any(name in l for l in q)
|
||||
assert any("audio_attn1" in l for l in q) # audio IS included in the widest scopes
|
||||
|
||||
|
||||
def test_serialized_format_matches_core_loader_contract():
|
||||
sd = _fake_ltxav_sd(n_blocks=3)
|
||||
out, _, _ = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1)
|
||||
layer = "transformer_blocks.1.attn1.to_q"
|
||||
assert out[f"{layer}.weight"].dtype == torch.uint8 # packed fp4 pairs
|
||||
assert out[f"{layer}.weight"].shape[1] == D // 2
|
||||
assert out[f"{layer}.weight_scale"].dtype == torch.float8_e4m3fn
|
||||
assert out[f"{layer}.weight_scale_2"].dtype == torch.float32
|
||||
import json
|
||||
conf = json.loads(bytes(out[f"{layer}.comfy_quant"].tolist()))
|
||||
assert conf == {"format": "nvfp4"}
|
||||
|
||||
|
||||
def test_prefix_and_refusals():
|
||||
sd = _fake_ltxav_sd(prefix="model.diffusion_model.")
|
||||
out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1)
|
||||
assert any(k.startswith("model.diffusion_model.") and k.endswith(".comfy_quant") for k in out)
|
||||
assert targets
|
||||
# already-quantized checkpoints are refused
|
||||
with pytest.raises(ValueError):
|
||||
quantize_state_dict_nvfp4(out, "video blocks", 1, 1)
|
||||
with pytest.raises(ValueError):
|
||||
quantize_state_dict_nvfp4({"scaled_fp8": torch.zeros(2)}, "video blocks", 1, 1)
|
||||
|
||||
|
||||
def test_selector_never_matches_bias_or_small():
|
||||
sd = {"transformer_blocks.1.attn1.to_q.bias": torch.randn(D, dtype=torch.bfloat16),
|
||||
"transformer_blocks.1.attn1.small.weight": _w(64, 64)}
|
||||
_, targets, _ = quantize_state_dict_nvfp4(sd, "all blocks", 0, 0)
|
||||
assert targets == []
|
||||
assert _selected("transformer_blocks.1.attn1.to_q", "video blocks", set())
|
||||
assert not _selected("transformer_blocks.1.audio_attn1.to_q", "video blocks", set())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disabled", [[], ["nvfp4"]])
|
||||
def test_core_mixed_precision_linear_loads_and_forwards(disabled):
|
||||
"""The emitted keys load through comfy.ops.mixed_precision_ops Linear and produce
|
||||
output close to the bf16 reference - both on the eager NVFP4 matmul path and on
|
||||
the dequant-emulation path used by GPUs without FP4 compute."""
|
||||
torch.manual_seed(7)
|
||||
w = _w()
|
||||
sd = {"transformer_blocks.1.attn1.to_q.weight": w,
|
||||
"transformer_blocks.1.attn1.to_q.bias": torch.zeros(D, dtype=torch.bfloat16)}
|
||||
out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 0, 0)
|
||||
assert targets == ["transformer_blocks.1.attn1.to_q.weight"]
|
||||
|
||||
ops = comfy.ops.mixed_precision_ops({}, torch.bfloat16, disabled=disabled)
|
||||
lin = ops.Linear(D, D, bias=True, device="cpu", dtype=torch.bfloat16)
|
||||
layer_sd = {k.replace("transformer_blocks.1.attn1.to_q.", ""): v for k, v in out.items()}
|
||||
missing, unexpected = lin.load_state_dict(layer_sd, strict=False)
|
||||
assert not missing and not unexpected
|
||||
assert isinstance(lin.weight, quant_ops.QuantizedTensor)
|
||||
assert lin.quant_format == "nvfp4"
|
||||
|
||||
x = torch.randn(2, 8, D, dtype=torch.bfloat16) * 0.5
|
||||
with torch.no_grad():
|
||||
y = lin.forward(x)
|
||||
y_ref = torch.nn.functional.linear(x, w)
|
||||
assert y.shape == y_ref.shape
|
||||
rel = (y.float() - y_ref.float()).abs().mean() / (y_ref.float().abs().mean() + 1e-8)
|
||||
assert rel < 0.25, f"relative error too high: {rel:.3f} (disabled={disabled})"
|
||||
|
||||
|
||||
def test_diagnostics_reports_layer_census(capsys):
|
||||
"""log_nvfp4_diagnostics counts nvfp4 layers + emulation flags and never raises
|
||||
(GEMM probe skipped without CUDA)."""
|
||||
import types
|
||||
from nvfp4_patch import log_nvfp4_diagnostics
|
||||
|
||||
class _Q:
|
||||
quant_format = "nvfp4"
|
||||
_full_precision_mm = False
|
||||
|
||||
class _QEmul(_Q):
|
||||
_full_precision_mm = True
|
||||
|
||||
class _Diff(torch.nn.Module):
|
||||
def modules(self):
|
||||
return [self, _Q(), _Q(), _QEmul()]
|
||||
|
||||
model = types.SimpleNamespace(model=types.SimpleNamespace(diffusion_model=_Diff()))
|
||||
out = log_nvfp4_diagnostics(model, None)
|
||||
assert out == {"nvfp4_layers": 3, "emulated_layers": 1}
|
||||
text = capsys.readouterr().out
|
||||
assert "nvfp4 layers in loaded model: 3" in text
|
||||
assert "EMULATION" in text
|
||||
assert "GEMM probe skipped" in text or "GEMM probe" in text
|
||||
|
||||
|
||||
def test_input_scale_injection_and_core_load(tmp_path, monkeypatch):
|
||||
"""Calibration sidecar -> inject_input_scales matches by suffix across checkpoint
|
||||
prefixes, and core's MixedPrecisionOps Linear registers the injected input_scale
|
||||
and still forwards correctly."""
|
||||
from nvfp4_patch import inject_input_scales, _scales_path
|
||||
import nvfp4_patch as np_mod
|
||||
|
||||
ckpt = str(tmp_path / "model.safetensors")
|
||||
sidecar = _scales_path(ckpt)
|
||||
assert sidecar.endswith("model.safetensors.json")
|
||||
|
||||
cal = {"transformer_blocks.1.attn1.to_q": 12.5}
|
||||
w = _w()
|
||||
sd = {"model.diffusion_model.transformer_blocks.1.attn1.to_q.weight": w,
|
||||
"model.diffusion_model.transformer_blocks.1.attn1.to_q.bias": torch.zeros(D, dtype=torch.bfloat16)}
|
||||
out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 0, 0)
|
||||
n = inject_input_scales(out, targets, cal)
|
||||
assert n == 1
|
||||
key = "model.diffusion_model.transformer_blocks.1.attn1.to_q.input_scale"
|
||||
assert key in out and out[key].dtype == torch.float32
|
||||
assert float(out[key]) == pytest.approx(12.5 * np_mod.CALIBRATION_MARGIN / (448.0 * 6.0))
|
||||
|
||||
# core load consumes it: input_scale becomes a registered parameter
|
||||
import comfy.ops
|
||||
ops = comfy.ops.mixed_precision_ops({}, torch.bfloat16, disabled=[])
|
||||
lin = ops.Linear(D, D, bias=True, device="cpu", dtype=torch.bfloat16)
|
||||
layer_sd = {k.replace("model.diffusion_model.transformer_blocks.1.attn1.to_q.", ""): v
|
||||
for k, v in out.items()}
|
||||
missing, unexpected = lin.load_state_dict(layer_sd, strict=False)
|
||||
assert not missing and not unexpected
|
||||
assert isinstance(getattr(lin, "input_scale", None), torch.nn.Parameter)
|
||||
x = torch.randn(2, 4, D, dtype=torch.bfloat16) * 0.5
|
||||
with torch.no_grad():
|
||||
y = lin.forward(x)
|
||||
y_ref = torch.nn.functional.linear(x, w)
|
||||
rel = (y.float() - y_ref.float()).abs().mean() / (y_ref.float().abs().mean() + 1e-8)
|
||||
assert y.shape == y_ref.shape and rel < 0.3
|
||||
|
||||
# no calibration entry for a layer -> nothing injected for it
|
||||
sd2 = {"transformer_blocks.2.attn1.to_q.weight": _w()}
|
||||
out2, targets2, _ = quantize_state_dict_nvfp4(sd2, "video blocks", 0, 0)
|
||||
assert inject_input_scales(out2, targets2, cal) == 0
|
||||
|
||||
|
||||
def test_load_calibration_reads_sidecar(tmp_path, monkeypatch):
|
||||
import nvfp4_patch as np_mod
|
||||
monkeypatch.setattr(np_mod, "_scales_path", lambda p: str(tmp_path / "s.json"))
|
||||
assert np_mod.load_calibration("whatever") is None
|
||||
(tmp_path / "s.json").write_text('{"amax": {"a.b": 3.0}, "margin": 1.25}')
|
||||
assert np_mod.load_calibration("whatever") == {"a.b": 3.0}
|
||||
Reference in New Issue
Block a user