diff --git a/cli/lora-quant-fidelity.py b/cli/lora-quant-fidelity.py index 1dc41bfe5..8aba5675c 100644 --- a/cli/lora-quant-fidelity.py +++ b/cli/lora-quant-fidelity.py @@ -252,7 +252,7 @@ class Bf16Repo: return f.get_tensor(key) -def analyze_module(W_dq, deq_params, mods, calib_rms=None): +def analyze_module(W_dq, deq_params, mods, calib_rms=None, step_live=None): """Return fidelity metrics for one quantized module and the adapters targeting it. Deltas come from each module's production calc_updown and sum the way the @@ -260,6 +260,9 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None): is measured as applied. A module is factor-path eligible only when every contribution is a plain additive lora. With ``calib_rms``, hosting mirrors the calibrated production path and its rho is scored in the weighted norm. + With ``step_live`` (the layer's own pre-add scale), the production routing + rule applies: a delta fat against the grid whose truncation capture is low + reports the requantize path, the way the loader would route it. """ D = None for mod in mods: @@ -316,17 +319,23 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None): with torch.random.fork_rng(devices=[D.device] if D.device.type == 'cuda' else []): torch.manual_seed(0) U, S, V = torch.svd_lowrank(Dw, q=min(q + 64, *D.shape), niter=8) - Dk = (U[:, :q] * S[:q]) @ V[:, :q].t() - if rms is not None: - Dk = Dk / rms - base16 = W_dq.to(torch.bfloat16).float() - realized = (W_dq.to(torch.bfloat16) + Dk.to(torch.bfloat16)).float() - base16 - if rms is not None: # weighted norm: the diagonal-covariance output-error proxy the calibrated truncation optimizes - Dr = D * rms - applied_rho = float((realized * rms).flatten() @ Dr.flatten() / Dr.square().sum()) - else: - applied_rho = float(realized.flatten() @ D.flatten() / nD.square()) - hosted = True + energy = float(S[:q].square().sum() / Dw.square().sum().clamp(min=1e-30)) + routed = False + if step_live is not None and not deq_params.get('use_svd', False): + sr = float(D.square().mean().sqrt() / step_live.float().mean()) + routed = sr > lora_sdnq.REQUANT_RATIO and energy < lora_sdnq.REQUANT_ENERGY + if not routed: # the loader routes fat, genuinely-truncated deltas back to requantize + Dk = (U[:, :q] * S[:q]) @ V[:, :q].t() + if rms is not None: + Dk = Dk / rms + base16 = W_dq.to(torch.bfloat16).float() + realized = (W_dq.to(torch.bfloat16) + Dk.to(torch.bfloat16)).float() - base16 + if rms is not None: # weighted norm: the diagonal-covariance output-error proxy the calibrated truncation optimizes + Dr = D * rms + applied_rho = float((realized * rms).flatten() @ Dr.flatten() / Dr.square().sum()) + else: + applied_rho = float(realized.flatten() @ D.flatten() / nD.square()) + hosted = True return dict(rank=getattr(mods[0], 'dim', None), rms_delta=float(D.pow(2).mean().sqrt()), rms_weight=float(W_dq.pow(2).mean().sqrt()), step_ratio=step_ratio, crossers=crossers, requant_rho=rho, requant_resid=resid, factor_eligible=factor_eligible, hosted=hosted, applied_rho=applied_rho, @@ -398,6 +407,7 @@ def main(): skip_quantized_matmul=deq.use_quantized_matmul, dtype=torch.float32, skip_compile=True).to(device) params = dict(weights_dtype=deq.weights_dtype, group_size=deq.group_size, hadamard_group_size=deq.hadamard_group_size, use_hadamard=deq.use_hadamard, use_svd=layer.svd_up is not None, svd_rank=deq.svd_rank, svd_steps=deq.svd_steps) + step_live = layer.scale.detach().to(device) sd_module = layer else: W = bf16_repo.get(f'{path}.weight') @@ -412,16 +422,18 @@ def main(): if args.dtype == 'bf16': W_dq = W.to(device, torch.bfloat16).float() params = dict(weights_dtype='bf16', group_size=0, hadamard_group_size=0, use_hadamard=False) + step_live = None else: deq0, data0 = sdnq_quantize_layer_weight(W.to(device, torch.float32), layer_class_name='Linear', weights_dtype=args.dtype, group_size=args.group, hadamard_group_size=args.hadamard_group, use_hadamard=args.hadamard_group > 0, use_svd=False, use_quantized_matmul=False, dequantize_fp32=False, torch_dtype=torch.bfloat16) W_dq = deq0(data0['weight'], data0['scale'], zero_point=data0['zero_point'], svd_up=None, svd_down=None, dtype=torch.float32, skip_compile=True) params = dict(weights_dtype=args.dtype, group_size=deq0.group_size, hadamard_group_size=deq0.hadamard_group_size, use_hadamard=deq0.use_hadamard) + step_live = data0['scale'].detach() sd_module = make_stub(W.shape) try: mods = [build_module(fam, path, w, net, sd_module) for fam, w in entries] - row = analyze_module(W_dq, params, mods, calib_rms=calib_stats.get(lname)) + row = analyze_module(W_dq, params, mods, calib_rms=calib_stats.get(lname), step_live=step_live) except Exception as e: # a family the tool cannot rebuild must not read as a clean module failed.append(f'{path}: {type(e).__name__}: {e}') del W_dq diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 3f9b9a0a0..5a341c137 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -33,6 +33,13 @@ magnitude. When activation statistics for the checkpoint exist (see error instead of weight error. At 8 bits and above requantize retains most of the delta, so hosting is skipped there and the requantize path remains. + +A small tail of deltas inverts the tradeoff: when the delta is large +against the grid step AND the truncation genuinely cuts it, requantize +retains more than hosting drops, and the layer routes back to the +requantize path (``REQUANT_RATIO``/``REQUANT_ENERGY``). Both terms must +agree: a thin delta rounds away on the grid however low its capture, and +a low-rank delta hosts exactly however fat it is. """ import torch @@ -45,6 +52,10 @@ from modules.logger import log fallback_layers: list[str] = [] hosted_layers: list[tuple[str, float, bool]] = [] +routed_layers: list[str] = [] + +REQUANT_RATIO = 0.30 # delta rms over mean grid step above which requantize can retain the delta +REQUANT_ENERGY = 0.90 # sketch capture below which truncation genuinely loses part of it def rank_bucket(r): """Fixed rank ladder for compiled-graph reuse: powers of two up to 256, multiples of 64 above (hosted rank plus exact members).""" @@ -243,8 +254,8 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): directions minimize output error rather than weight error. Computed factors are disk-cached per configuration (``lora_factor_cache``) and replayed bit-identically on later applies. Returns None when the delta - cannot ride the channel (wrong shape); the caller falls back to - requantize. + cannot ride the channel (wrong shape) or when the routing rule prefers + the grid for it; the caller falls back to requantize. """ from sdnq.quant_utils import rotate_hadamard @@ -255,19 +266,33 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): if updown is None or updown.ndim != 2 or tuple(updown.shape) != tuple(deq.original_shape): return None dtype = deq.result_dtype - lora_factor_cache.begin_pass(wanted_names) - cached = lora_factor_cache.fetch(network_layer_name) - D = None if cached is not None else updown.detach().to(devices.device, torch.float32) - ups, downs = [], [] + members = [] for net in l.loaded_networks: module = net.modules.get(network_layer_name, None) if module is None: continue factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) - if factors is None: - continue - up_eff, down = factors + if factors is not None: + members.append(factors) + + # requantize keeps a delta the grid can resolve and that truncation would genuinely + # cut: both terms must agree, since a thin delta rounds away on the grid however + # low its capture, and a low-rank delta hosts exactly however fat it is. Scoped to + # sets the side-channel would otherwise carry whole: factorable members ride + # exactly. + maybe_requant = len(members) == 0 and self.svd_up is None + if maybe_requant: + step = float(self.scale.detach().float().mean()) + rms = float(updown.detach().float().square().mean().sqrt()) + maybe_requant = step > 0 and rms / step > REQUANT_RATIO + + lora_factor_cache.begin_pass(wanted_names) + cached = lora_factor_cache.fetch(network_layer_name) + D = None if cached is not None else updown.detach().to(devices.device, torch.float32) + + ups, downs = [], [] + for up_eff, down in members: if D is not None: D = D.sub_(up_eff.to(torch.float32) @ down.to(torch.float32)) # factorable members ride exactly; host only the remainder if deq.use_hadamard: @@ -277,10 +302,33 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): if cached is not None: up_h, down_h, energy, calibrated = cached + if maybe_requant and energy < REQUANT_ENERGY: + routed_layers.append(network_layer_name) + return None append_factors(self, ups + [up_h.to(device=devices.device, dtype=dtype)], downs + [down_h.to(device=devices.device, dtype=dtype)]) hosted_layers.append((network_layer_name, energy, calibrated)) return True + up_h, down_h, energy, calibrated = truncate_delta(self, D, dtype) + up_h, down_h = lora_factor_cache.store(network_layer_name, up_h, down_h, energy, calibrated) + if maybe_requant and energy < REQUANT_ENERGY: + routed_layers.append(network_layer_name) # the stored entry memoizes the routing; replays skip the sketch + return None + append_factors(self, ups + [up_h], downs + [down_h]) + hosted_layers.append((network_layer_name, energy, calibrated)) + return True + + +def truncate_delta(self, D, dtype): + """Truncate one dense fp32 delta to hosted factors in the layer's channel layout; consumes ``D``. + + Calibration-weighted when statistics exist; the sketch is oversampled past + the kept rank so the truncation sits within noise of exact svd. Returns + ``(up_h, down_h, energy, calibrated)`` with the down factor rotated into the + layer's hadamard domain. + """ + from sdnq.quant_utils import rotate_hadamard + deq = self.sdnq_dequantizer cap = int(shared.opts.lora_sdnq_host_rank) q = min(cap, *D.shape) rms = lora_calib.rms_for(self) @@ -304,15 +352,12 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): if deq.use_hadamard: down_h = rotate_hadamard(down_h, group_size=deq.hadamard_group_size) down_h = down_h.to(dtype=dtype) - up_h, down_h = lora_factor_cache.store(network_layer_name, up_h, down_h, energy, rms is not None) - append_factors(self, ups + [up_h], downs + [down_h]) - hosted_layers.append((network_layer_name, energy, rms is not None)) - return True + return up_h, down_h, energy, rms is not None def note_fallback(self, network_layer_name): - """Record a quantized layer taking the lossy requantize path (summary-logged per pass).""" - if getattr(self, 'sdnq_dequantizer', None) is not None: + """Record a quantized layer taking the requantize path (summary-logged per pass); layers the routing rule sent there are counted apart.""" + if getattr(self, 'sdnq_dequantizer', None) is not None and network_layer_name not in routed_layers: fallback_layers.append(network_layer_name) @@ -328,6 +373,11 @@ def report_fallbacks(): if l.debug: log.debug(f'Network load: type=LoRA quant=sdnq hosted={[(n, round(e, 3)) for n, e, _c in hosted_layers[:8]]}{"..." if len(hosted_layers) > 8 else ""}') hosted_layers.clear() + if len(routed_layers) > 0: + log.info(f'Network load: type=LoRA quant=sdnq apply=requantize layers={len(routed_layers)} routed=fat-delta') + if l.debug: + log.debug(f'Network load: type=LoRA quant=sdnq routed={routed_layers[:8]}{"..." if len(routed_layers) > 8 else ""}') + routed_layers.clear() if len(fallback_layers) > 0: if enabled(): log.warning(f'Network load: type=LoRA quant=sdnq layers={len(fallback_layers)} non-factorable networks requantized in place (reduced fidelity on quantized weights)') diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 53872f021..da9ffa12b 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -753,6 +753,129 @@ def test_hosted_transitions_and_rng_isolation(): return True +@contextmanager +def requant_rule(ratio, energy): + old_r, old_e = lora_sdnq.REQUANT_RATIO, lora_sdnq.REQUANT_ENERGY + lora_sdnq.REQUANT_RATIO, lora_sdnq.REQUANT_ENERGY = ratio, energy + try: + yield + finally: + lora_sdnq.REQUANT_RATIO, lora_sdnq.REQUANT_ENERGY = old_r, old_e + + +def test_route_fat_dense_delta_requantizes(): + layer = build_layer('uint4') + torch.manual_seed(21) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 # full-rank and well above the grid step: the grid retains it, truncation would cut it + net = make_dense_net('fatnet', layer, D) + with host_rank(256), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'a fat full-rank delta must route to requantize' + assert isinstance(getattr(layer, 'network_weights_backup', None), torch.Tensor), 'the routed layer takes the requantize backup' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.7, f'the grid must retain the routed delta: rho={rho:.3f}' + activate() + assert torch.equal(dq(layer), Wdq0), 'restore from backup must be bit-exact' + return True + + +def test_route_rule_terms_gate_both_ways(): + layer = build_layer('uint4') + torch.manual_seed(23) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 # sr about 0.8, capture about 0.8 at cap 256: each term alone can hold it hosted + net = make_dense_net('gatenet', layer, D) + with host_rank(256), mock_model(lin=layer): + with requant_rule(ratio=10.0, energy=0.90): + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'sr below the ratio must host regardless of capture' + activate() + with requant_rule(ratio=0.30, energy=0.0): + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'capture above the energy floor must host regardless of sr' + activate() + with requant_rule(ratio=0.30, energy=0.90): + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'both terms crossed must requantize' + activate() + return True + + +def test_route_low_rank_fat_delta_stays_hosted(): + layer = build_layer('uint4') + _A, _B, D = make_delta(seed=22, sigma=3e-3) # rank-8: fat against the grid, exact under the cap + net = make_dense_net('fatlow', layer, D) + with host_rank(64), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'a low-rank delta hosts exactly at any magnitude' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.95, f'rho={rho:.4f}' + activate() + return True + + +def test_route_mixed_set_keeps_hosting(): + layer = build_layer('uint4') + A, B, _D1 = make_delta(seed=24) + torch.manual_seed(25) + D2 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + net1 = make_net('mixp', layer, A, B) + net2 = make_dense_net('mixf', layer, D2) + with host_rank(256), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net1, net2) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'a set with factorable members keeps the side-channel' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_route_svd_checkpoint_keeps_hosting(): + layer = build_layer('uint4', use_svd=True) + torch.manual_seed(26) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + net = make_dense_net('svdfat', layer, D) + with host_rank(256), mock_model(lin=layer): + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'svd checkpoints keep hosting; the rule is not grounded there' + activate() + return True + + +def test_route_dense_stack_keeps_hosting(): + layer = build_layer('uint4') + torch.manual_seed(27) + D1 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + D2 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + net1, net2 = make_dense_net('df1', layer, D1), make_dense_net('df2', layer, D2) + with host_rank(256), stack_mode('ties'), mock_model(lin=layer): + activate(net1, net2) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'dense-combined deltas host at any magnitude' + activate() + return True + + +def test_route_replay_from_cache(): + import tempfile + layer = build_layer('uint4') + with tempfile.TemporaryDirectory() as tmp: + with host_rank(256), host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=layer): + net, _D = cache_fixture(tmp, layer, name='fatcache', sigma=1e-2, seed=28) + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'fat delta must route on the fresh-sketch path' + activate() + real_svd = torch.svd_lowrank + torch.svd_lowrank = raise_no_svd + try: + activate(net) # the stored entry memoizes the routing: same decision, no sketch + finally: + torch.svd_lowrank = real_svd + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'cache replay must route the same way' + activate() + return True + + CAT_CALIB = category('calibration') @@ -1294,7 +1417,10 @@ def run_tests(): run_test(CAT_TRANS, fn) log.warning('=== Hosting ===') for fn in [test_hosted_low_rank_delta_is_kept, test_hosted_dense_delta_beats_requant, test_hosted_skips_int8, - test_hosted_disabled_by_option, test_hosted_transitions_and_rng_isolation]: + test_hosted_disabled_by_option, test_hosted_transitions_and_rng_isolation, + test_route_fat_dense_delta_requantizes, test_route_rule_terms_gate_both_ways, test_route_low_rank_fat_delta_stays_hosted, + test_route_mixed_set_keeps_hosting, test_route_svd_checkpoint_keeps_hosting, test_route_dense_stack_keeps_hosting, + test_route_replay_from_cache]: run_test(CAT_HOST, fn) log.warning('=== Calibration ===') for fn in [test_calibrated_hosting_beats_plain, test_calibrated_low_rank_delta_survives, test_calib_option_off_matches_plain,