jlens_v3: chunk VJP probes (16) to fit K2200 4GB; proxy chunk fix; e2e verified (last-layer cos-sim=1.0)

1 files 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)])