Fix diagnostic performance and report timing
Performance: hooks now only capture when sigma in [0.85, 0.95] - one step per scene instead of every step. Eliminates the 2x overhead and 5x slowdown on step 6. Report timing: fires in sigma tracker when sigma < 0.5 AND scene is snapped AND ≥2 scenes captured. Covers the last sigma=0.421875 step which was below the old hook_remover threshold of 0.05. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
21fb111956
commit
bf6cb80f9f
+27
-13
@@ -672,7 +672,7 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
|
||||
inj_gate = (0.2, 0.5) if is_anchor else (0.35, 0.80)
|
||||
lazy = lazy_injects.get(idx)
|
||||
|
||||
def _make_hook(block_idx, buf, lazy_ref, strength, gate, s_state):
|
||||
def _make_hook(block_idx, buf, lazy_ref, strength, gate, s_state, scene_cur=_scene_cur_buf):
|
||||
def _extract_tensor(out):
|
||||
if isinstance(out, dict):
|
||||
return out.get("img") or out.get("hidden_states") or next(iter(out.values()), None), out
|
||||
@@ -681,13 +681,16 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
|
||||
return out, out
|
||||
|
||||
def _hook(module, inp, out):
|
||||
# Capture
|
||||
try:
|
||||
t, _ = _extract_tensor(out)
|
||||
if isinstance(t, torch.Tensor):
|
||||
buf[block_idx] = t.detach().cpu()
|
||||
except Exception:
|
||||
pass
|
||||
# Capture only at mid-sigma (avoid per-step tensor copies)
|
||||
if 0.85 <= s_state[0] <= 0.95:
|
||||
try:
|
||||
t, _ = _extract_tensor(out)
|
||||
if isinstance(t, torch.Tensor):
|
||||
buf[block_idx] = t.detach().cpu()
|
||||
if scene_cur is not None:
|
||||
scene_cur[0][block_idx] = buf[block_idx]
|
||||
except Exception:
|
||||
pass
|
||||
# Inject (only when lazy has a tensor)
|
||||
if lazy_ref is None or lazy_ref.tensor is None:
|
||||
return None
|
||||
@@ -717,11 +720,15 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
|
||||
_make_hook(idx, capture_buf, lazy, inj_strength, inj_gate, sigma_state)
|
||||
)
|
||||
hook_handles.append(h)
|
||||
# Register capture-only diagnostic hooks on extra candidate blocks
|
||||
# Register capture-only diagnostic hooks on extra candidate blocks.
|
||||
# Gate by sigma to avoid capturing (and copying tensors) on every step.
|
||||
for idx in _diag_blocks - blocks_to_hook:
|
||||
if idx < len(block_list):
|
||||
def _make_diag_hook(block_idx, scene_cur):
|
||||
def _make_diag_hook(block_idx, scene_cur, s_state):
|
||||
def _hook(module, inp, out):
|
||||
# Only capture in the mid-sigma window once per scene
|
||||
if not (0.85 <= s_state[0] <= 0.95):
|
||||
return None
|
||||
try:
|
||||
if isinstance(out, dict):
|
||||
t = out.get("img") or out.get("hidden_states") or next(iter(out.values()), None)
|
||||
@@ -733,8 +740,9 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
|
||||
scene_cur[0][block_idx] = t.detach().cpu()
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
return _hook
|
||||
h = block_list[idx].register_forward_hook(_make_diag_hook(idx, _scene_cur_buf))
|
||||
h = block_list[idx].register_forward_hook(_make_diag_hook(idx, _scene_cur_buf, sigma_state))
|
||||
hook_handles.append(h)
|
||||
|
||||
# Also redirect the main capture hooks to write into scene_cur_buf too
|
||||
@@ -833,8 +841,14 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
|
||||
|
||||
# Per-scene diagnostic snapshot at mid-sigma
|
||||
if scene_cur is not None and not _scene_snapped[0] and sigma < 0.95:
|
||||
_scene_snapped[0] = True
|
||||
scene_cur[0] = {k: v for k, v in scene_cur[0].items()} # snapshot copy
|
||||
_scene_snapped[0] = True # hooks already wrote into scene_cur[0] at this sigma
|
||||
|
||||
# Fire similarity report at end of each scene (sigma low) when ≥2 scenes captured
|
||||
if scene_captures is not None and sigma < 0.5 and _scene_snapped[0]:
|
||||
if scene_cur is not None and scene_cur[0] and (not scene_captures or scene_captures[-1] is not scene_cur[0]):
|
||||
scene_captures.append(dict(scene_cur[0]))
|
||||
if len(scene_captures) >= 2:
|
||||
_report_block_similarity(scene_captures)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
Reference in New Issue
Block a user