From c762bfec18541c0a6d95f7eca2cff3c7caaab9d4 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 16:41:27 +0100 Subject: [PATCH] perf(lora): materialize select winners on the accelerator and time the reset The weight-kind schedule reset ran each winner's calc_updown on the target weight's device, which on a block-swapped denoiser is the cpu; at hundreds of layers per pass the cpu matmuls dominated every select generation. The delta now computes on the accelerator and moves back, matching the activate walk's convention. - reset and flip execution log a debug timing line (materialize, select loop, move/calc/apply split); the reset runs outside the activate walk, so its cost was invisible to the load timers --- modules/lora/lora_stack.py | 31 +++++++++++++++-- test/test-sdnq-lora-factors.py | 63 +++++++++++++++++++++++++++++++++- 2 files changed, 90 insertions(+), 4 deletions(-) diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index acd1117fa..85818c8b9 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -15,6 +15,7 @@ EST-LoRA arXiv:2508.02165 (its measured style-discrepancy estimate is exposed as an option instead of being derived from probe generations). """ +import time import weakref import hashlib @@ -305,9 +306,16 @@ def finalize(total_steps): e_den = sum(e['scores'][1] for e in state['entries'].values()) state['gamma_e'] = (e_num / e_den) if e_den > 0 else 1.0 state['flips'] = {} - if any(e['kind'] == 'weight' for e in state['entries'].values()): + stats = {'weight_n': 0, 'factor_n': 0, 'materialize': 0.0, 'select': 0.0, 'w_move': 0.0, 'w_calc': 0.0, 'w_apply': 0.0} + state['stats'] = stats + stats['weight_n'] = sum(1 for e in state['entries'].values() if e['kind'] == 'weight') + stats['factor_n'] = len(state['entries']) - stats['weight_n'] + if stats['weight_n'] > 0: + t0 = time.time() materialize_model() + stats['materialize'] = time.time() - t0 style_first = 0 + t0 = time.time() for layer_name, entry in list(state['entries'].items()): # snapshot: apply_selection drops entries whose module died flip_at = layer_flip_step(entry['scores'], state['total_steps']) initial = 1 if flip_at == 0 else 0 @@ -315,6 +323,7 @@ def finalize(total_steps): apply_selection(layer_name, entry, initial) if 0 < flip_at < state['total_steps']: state['flips'].setdefault(flip_at - 1, []).append(layer_name) # step callbacks fire after the denoise, so the flip runs one step early to be live during the crossover step's forward + stats['select'] = time.time() - t0 state['finalized'] = True if len(state['entries']) > 0: # only a built schedule can carry a flip count, so this is the line that shows selection is live rather than requested gamma = state['gamma_e'] if mode() == 'estlora' else state['gamma'] @@ -322,6 +331,8 @@ def finalize(total_steps): if report != state['reported']: # rebuilt every pass, so a batch would otherwise repeat one line per image state['reported'] = report log.info(f'Network load: type=LoRA stack={report[0]} layers={report[1]} style={report[2]} flips={report[3]} steps={report[4]} gamma={report[5]:.3f}') + # logged every pass: the reset runs outside the activate walk, so its cost is invisible to the load timers + log.debug(f'Network select: type=LoRA reset weight={stats["weight_n"]} factor={stats["factor_n"]} time={{materialize: {stats["materialize"]:.2f}, select: {stats["select"]:.2f}, move: {stats["w_move"]:.2f}, calc: {stats["w_calc"]:.2f}, apply: {stats["w_apply"]:.2f}}}') def reset(total_steps): @@ -335,10 +346,15 @@ def on_step(step): """Flip the layers whose crossover is this step; non-flip steps are a dict miss.""" if not state['finalized']: return - for layer_name in state['flips'].get(int(step), ()): + layers = state['flips'].get(int(step), ()) + if not layers: + return + t0 = time.time() + for layer_name in layers: entry = state['entries'].get(layer_name) if entry is not None: apply_selection(layer_name, entry, 1) + log.debug(f'Network select: type=LoRA flip step={int(step)} layers={len(layers)} time={time.time() - t0:.2f}') def apply_selection(layer_name, entry, winner): @@ -375,6 +391,15 @@ def weight_selection(module, entry, winner): if weight is None or weight.is_meta: warn_once('select-offloaded', 'Network stack: flip=skipped weight=offloaded') return + from modules import devices + stats = state.get('stats') or {} device = weight.device - updown = net_module.calc_updown(backup.to(device))[0] + t0 = time.time() + base = backup.to(devices.device) # a swapped-out layer keeps its weight on cpu; the delta matmul belongs on the accelerator regardless + t1 = time.time() + updown = net_module.calc_updown(base)[0].to(device) + t2 = time.time() network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it + stats['w_move'] = stats.get('w_move', 0.0) + (t1 - t0) + stats['w_calc'] = stats.get('w_calc', 0.0) + (t2 - t1) + stats['w_apply'] = stats.get('w_apply', 0.0) + (time.time() - t2) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 61627f5f8..e5c4f82fb 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -2119,6 +2119,66 @@ def test_select_weight_replay_from_cache_skips_calc(): return True +def test_select_reset_reports_timing(): + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device=DEVICE) + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.02) + lin.network_layer_name = 'lora_transformer_timed' + lin.network_current_names = () + A1, B1, _D1 = make_delta(seed=77, sigma=1e-2) + A2, B2, _D2 = make_delta(seed=78, sigma=1e-2) + n1 = make_net('t1', lin, A1, B1) + n2 = make_net('t2', lin, A2, B2) + with mock_model(lin=lin), select_mode('klora'): + activate(n1, n2) + lora_stack.reset(20) + stats = lora_stack.state.get('stats') + assert stats is not None, 'a reset must publish its timing stats' + assert stats['weight_n'] == 1 and stats['factor_n'] == 0, f'weight-kind counts wrong: {stats}' + assert stats['w_calc'] > 0.0, 'the weight-kind winner apply must account its calc time' + assert stats['select'] > 0.0 + activate() + layer = build_layer('uint4') + f1, f2, _Df1, _Df2 = select_pair(layer, seed0=79, seed1=80) + with mock_model(lin=layer), select_mode('klora'): + activate(f1, f2) + lora_stack.reset(20) + stats = lora_stack.state.get('stats') + assert stats is not None and stats['factor_n'] == 1 and stats['weight_n'] == 0, f'factor-kind counts wrong: {stats}' + assert stats['w_calc'] == 0.0, 'factor-kind resets flip segments and must not touch the weight path' + activate() + return True + + +def test_select_weight_flip_calcs_on_accelerator(): + from modules.lora import network_lora + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device='cpu') # a swapped-out layer: weight lives on cpu + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F) * 0.02) + lin.network_layer_name = 'lora_transformer_swapped' + lin.network_current_names = () + A1, B1, _D1 = make_delta(seed=81, sigma=1e-2) + A2, B2, _D2 = make_delta(seed=82, sigma=1e-2) + n1 = make_net('s1', lin, A1, B1) + n2 = make_net('s2', lin, A2, B2) + seen = [] + real = network_lora.NetworkModuleLora.calc_updown + def spy(self, target, *args, **kwargs): + seen.append(target.device.type) + return real(self, target, *args, **kwargs) + with mock_model(lin=lin), select_mode('klora'): + activate(n1, n2) + lin.to('cpu') # the offload dispatch swaps blocks back out after the walk; the reset must not follow the weight onto the cpu + network_lora.NetworkModuleLora.calc_updown = spy + try: + lora_stack.reset(20) + finally: + network_lora.NetworkModuleLora.calc_updown = real + assert seen and all(d == DEVICE.type for d in seen), f'winner materialization must calc on the accelerator, saw {seen}' + activate() + return True + + CAT_COMPILE = category('compile') @@ -2311,7 +2371,8 @@ def run_tests(): test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer, test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step, - test_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc]: + test_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc, + test_select_reset_reports_timing, test_select_weight_flip_calcs_on_accelerator]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: