From a4665ab7119ca62ee2c89785d59317d1d82b6db8 Mon Sep 17 00:00:00 2001 From: Void Agent Date: Fri, 31 Jul 2026 16:26:13 +0100 Subject: jlens_v3: chunk VJP probes (16) to fit K2200 4GB; proxy chunk fix; e2e verified (last-layer cos-sim=1.0) --- src/jlens_v3.py | 45 +++++++++++++++++++++++++++++---------------- 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/src/jlens_v3.py b/src/jlens_v3.py index 7112cb1..7316fe8 100644 --- a/src/jlens_v3.py +++ b/src/jlens_v3.py @@ -100,22 +100,29 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32): # sum over all future positions t' (causal: t' < t contributes zero grad) h_final_sum = h_final.sum(dim=1) # (B, d) - # grad_outputs[v, b, :] = W_U[v] -> batched VJPs for all V tokens - grad_outputs = W_U.to(device).unsqueeze(1).expand(V, B, d).contiguous() - grads = torch.autograd.grad( - h_final_sum, h_l, grad_outputs=grad_outputs, - is_grads_batched=True)[0] # (V, B, T, d) - accum += grads.detach().cpu().reshape(V, -1, d).sum(dim=1) + # grad_outputs[v, b, :] = W_U[v] -> batched VJPs, chunked over vocab + # (is_grads_batched with all V=65 probes at once OOMs the 4GB K2200) + for v0 in range(0, V, chunk): + C = min(chunk, V - v0) + grad_outputs = (W_U[v0:v0 + C].to(device).unsqueeze(1) + .expand(C, B, d).contiguous()) + grads = torch.autograd.grad( + h_final_sum, h_l, grad_outputs=grad_outputs, + is_grads_batched=True, + retain_graph=(v0 + C < V))[0] # (C, B, T, d) + accum[v0:v0 + C] += grads.detach().cpu().reshape(C, -1, d).sum(dim=1) + del grads, grad_outputs + torch.cuda.empty_cache() n_pairs += B * T * (T + 1) // 2 # valid (t, t' >= t) pairs - del logits, loss, h_l, h_final, h_final_sum, grads, grad_outputs + del logits, loss, h_l, h_final, h_final_sum torch.cuda.empty_cache() accum /= n_pairs return accum # (V, d): faithful J-lens vectors, rows of W_U * J_l -def compute_proxy_norms(model, layer_idx, batches, device, vocab_size): +def compute_proxy_norms(model, layer_idx, batches, device, vocab_size, chunk=16): """ Old proxy (jlens_v2 quantity): || E_x[ grad log_softmax(logits)[:, :, k] ] || averaged over positions — norm of the mean gradient, per token. @@ -134,15 +141,20 @@ def compute_proxy_norms(model, layer_idx, batches, device, vocab_size): logits, loss = model(x, y) handle.remove() log_probs = torch.nn.functional.log_softmax(logits, dim=-1) - # sum over positions -> (V,); batched VJP over all vocab at once + # sum over positions -> (V,); batched VJP over vocab, chunked for GPU memory lp_sum = log_probs.sum(dim=(0, 1)) - grad_outputs = torch.eye(vocab_size, device=device) - grad = torch.autograd.grad( - lp_sum, resid_captured['val'], grad_outputs=grad_outputs, - is_grads_batched=True)[0] # (V, B, T, d) - accum += grad.detach().cpu().reshape(vocab_size, -1, d).sum(dim=1) + for v0 in range(0, vocab_size, chunk): + C = min(chunk, vocab_size - v0) + grad_outputs = torch.eye(vocab_size, device=device)[v0:v0 + C] + grad = torch.autograd.grad( + lp_sum, resid_captured['val'], grad_outputs=grad_outputs, + is_grads_batched=True, + retain_graph=(v0 + C < vocab_size))[0] # (C, B, T, d) + accum[v0:v0 + C] += grad.detach().cpu().reshape(C, -1, d).sum(dim=1) + del grad, grad_outputs + torch.cuda.empty_cache() counts += B * T - del logits, loss, log_probs, lp_sum, grad + del logits, loss, log_probs, lp_sum torch.cuda.empty_cache() norms = {} for token_id in range(vocab_size): @@ -203,7 +215,8 @@ def main(): f"{sims.mean().item():.4f} (expect ~1.0 if J=identity)") print(" Computing old proxy norms...") - proxy_norms = compute_proxy_norms(model, layer_idx, batches, device, V) + proxy_norms = compute_proxy_norms(model, layer_idx, batches, device, V, + args.chunk) # correlations with frequency f_arr = np.array([freq[k] for k in range(V)]) -- cgit v1.2.3