diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index 220877c04..a5ef4041a 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -8,13 +8,17 @@ statistics, so they are cached on disk keyed by exactly that identity and replayed bit-identically on the next apply of the same configuration. One safetensors file per configuration under ``models/lora-factor-cache``, -holding every hosted layer's post-rotation factor pair. Files are named by -the model and network set with an identity-hash suffix, and the exact -signature is embedded in the file metadata. The -``lora_sdnq_host_cache`` option is the size budget in GB (0 disables); -least-recently-used entries are evicted past the budget. Any doubt about -identity (unknown checkpoint, unreadable lora file, signature mismatch) -disables caching for the pass rather than risking a stale hit. +holding every hosted layer's post-rotation factor pair as rowwise int8 +with fp32 scales (measured fidelity-free in output space, half the bytes +of bf16). Files are named by the model and network set with an +identity-hash suffix, and the exact signature is embedded in the file +metadata. Factors are quantized before first use: ``store`` returns the +dequantized round-trip for the caller to apply, so a fresh compute and a +later cache hit attach bit-identical tensors. The ``lora_sdnq_host_cache`` +option is the size budget in GB (0 disables); least-recently-used entries +are evicted past the budget. Any doubt about identity (unknown checkpoint, +unreadable lora file, signature mismatch) disables caching for the pass +rather than risking a stale hit. """ import os @@ -86,46 +90,73 @@ def begin_pass(wanted_names): sig = json.dumps(parts, sort_keys=True) key = hashlib.sha256(sig.encode()).hexdigest()[:24] path = os.path.join(cache_root, f'{label(parts)}-{key}.safetensors') - store = {} + entries = {} if os.path.isfile(path): try: from safetensors import safe_open with safe_open(path, framework='pt', device='cpu') as f: - if (f.metadata() or {}).get('sig') == sig: + meta = f.metadata() or {} + if meta.get('sig') == sig and meta.get('fmt') == '2': for k in f.keys(): - store[k] = f.get_tensor(k) + entries[k] = f.get_tensor(k) os.utime(path, None) # freshness for LRU eviction except Exception as e: log.debug(f'Network cache: read failed path="{path}" {e}') - store = {} + entries = {} state.update(sig=sig, path=path) - state['store'] = store - log.debug(f'Network cache: entry="{path}" keys={len(store)}') + state['store'] = entries + log.debug(f'Network cache: entry="{path}" keys={len(entries)}') + + +def quantize_rowwise(t): + t32 = t.detach().to(torch.float32) + scale = t32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12) / 127.0 + q = (t32 / scale).round().clamp(-127, 127).to(torch.int8) + return q, scale + + +def dequantize_rowwise(q, scale): + # int8 * fp32 with a single fp32 rounding: identical on any device, so hit and miss replay the same values + return q.to(torch.float32) * scale def fetch(network_layer_name): - """Cached (up, down, energy, calibrated) for a layer, or None.""" + """Cached (up, down, energy, calibrated) for a layer, or None; factors return as fp32.""" if state['sig'] is None: return None - up = state['store'].get(f'{network_layer_name}.up') - down = state['store'].get(f'{network_layer_name}.down') - energy = state['store'].get(f'{network_layer_name}.energy') - calib = state['store'].get(f'{network_layer_name}.calib') - if up is None or down is None or energy is None or calib is None: + st = state['store'] + up_q, up_s = st.get(f'{network_layer_name}.up_q'), st.get(f'{network_layer_name}.up_s') + down_q, down_s = st.get(f'{network_layer_name}.down_q'), st.get(f'{network_layer_name}.down_s') + energy = st.get(f'{network_layer_name}.energy') + calib = st.get(f'{network_layer_name}.calib') + if up_q is None or up_s is None or down_q is None or down_s is None or energy is None or calib is None: state['misses'] += 1 return None state['hits'] += 1 - return up, down, float(energy), bool(calib) + return dequantize_rowwise(up_q, up_s), dequantize_rowwise(down_q, down_s), float(energy), bool(calib) -def put(network_layer_name, up, down, energy, calibrated): +def store(network_layer_name, up, down, energy, calibrated): + """Quantize-before-use: returns the pair the caller must apply. + + With caching inactive the inputs pass through untouched. Otherwise the + factors are stored as rowwise int8 and the dequantized round-trip comes + back, so the factors applied now and the factors a later hit replays are + the same tensors. + """ if state['sig'] is None: - return - state['store'][f'{network_layer_name}.up'] = up.detach().to('cpu').contiguous() - state['store'][f'{network_layer_name}.down'] = down.detach().to('cpu').contiguous() - state['store'][f'{network_layer_name}.energy'] = torch.tensor(float(energy)) - state['store'][f'{network_layer_name}.calib'] = torch.tensor(1 if calibrated else 0, dtype=torch.uint8) + return up, down + up_q, up_s = quantize_rowwise(up) + down_q, down_s = quantize_rowwise(down) + st = state['store'] + st[f'{network_layer_name}.up_q'] = up_q.to('cpu').contiguous() + st[f'{network_layer_name}.up_s'] = up_s.to('cpu').contiguous() + st[f'{network_layer_name}.down_q'] = down_q.to('cpu').contiguous() + st[f'{network_layer_name}.down_s'] = down_s.to('cpu').contiguous() + st[f'{network_layer_name}.energy'] = torch.tensor(float(energy)) + st[f'{network_layer_name}.calib'] = torch.tensor(1 if calibrated else 0, dtype=torch.uint8) state['dirty'] = True + return dequantize_rowwise(up_q, up_s).to(up.dtype), dequantize_rowwise(down_q, down_s).to(down.dtype) def evict(): @@ -159,7 +190,7 @@ def flush(): from safetensors.torch import save_file os.makedirs(cache_root, exist_ok=True) tmp = state['path'] + '.tmp' - save_file(state['store'], tmp, metadata={'sig': state['sig']}) + save_file(state['store'], tmp, metadata={'sig': state['sig'], 'fmt': '2'}) os.replace(tmp, state['path']) evict() except Exception as e: diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 51f91044b..8d5330e09 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -283,7 +283,7 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa if deq.use_hadamard: down_h = rotate_hadamard(down_h, group_size=deq.hadamard_group_size) down_h = down_h.to(dtype=dtype) - lora_factor_cache.put(network_layer_name, up_h, down_h, energy, rms is not None) + 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 diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 95c6ddb35..58028801c 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -952,6 +952,9 @@ def test_factor_cache_roundtrip_bitexact(): activate() # pass end flushed the entry; unload restores the base files = os.listdir(os.path.join(tmp, 'cache')) assert len(files) == 1, f'one cache entry expected, got {files}' + bf16_bytes = (first_up.numel() + first_down.numel()) * 2 + entry_bytes = os.path.getsize(os.path.join(tmp, 'cache', files[0])) + assert entry_bytes < bf16_bytes * 0.62 + 8192, f'int8 entry must be about half the bf16 factor bytes: {entry_bytes} vs {bf16_bytes}' real_svd = torch.svd_lowrank torch.svd_lowrank = raise_no_svd try: @@ -984,6 +987,20 @@ def test_factor_cache_invalidates_on_multiplier(): return True +def test_factor_cache_int8_quantization(): + from modules.lora import lora_factor_cache as fc + torch.manual_seed(71) + t = torch.randn(64, 128, device=DEVICE) * torch.logspace(-3, 0, 64, device=DEVICE)[:, None] # rows spanning magnitudes + q, s = fc.quantize_rowwise(t) + assert q.dtype == torch.int8 + dq = fc.dequantize_rowwise(q, s) + err = (dq - t).abs().max(dim=1).values + assert bool((err <= s.squeeze(1) * 0.51).all()), 'rowwise int8 error must stay within half a step' + cos = torch.nn.functional.cosine_similarity(dq.flatten(), t.flatten(), dim=0) + assert float(cos) > 0.99995, f'int8 roundtrip cosine {float(cos):.6f}' + return True + + def test_factor_cache_disabled_at_zero(): import tempfile layer = build_layer('uint4') @@ -1070,7 +1087,8 @@ def run_tests(): test_calib_capture_persist_roundtrip, test_calib_capture_gates]: run_test(CAT_CALIB, fn) log.warning('=== Factor cache ===') - for fn in [test_factor_cache_roundtrip_bitexact, test_factor_cache_invalidates_on_multiplier, test_factor_cache_disabled_at_zero]: + for fn in [test_factor_cache_roundtrip_bitexact, test_factor_cache_invalidates_on_multiplier, + test_factor_cache_int8_quantization, test_factor_cache_disabled_at_zero]: run_test(CAT_FCACHE, fn) log.warning('=== Robustness ===') for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: