perf(sub_zero): kill AtP per-pair thrash (autograd.grad + GPU reduction)

Root cause of 51s/pair: the backward hook did g.detach().float().cpu() on
every scored weight, forcing up to ~96 CUDA syncs + ~10-22GB PCIe transfer
per pair and serializing the backward.

Fix:
- torch.autograd.grad over exactly the scored weights (no hooks, one pass)
- grad->SV reduction on GPU; only a tiny [rank] vector crosses to CPU
- drop gradient_checkpointing_enable() (scoped grad fixed the OOM; checkpointing
was doubling backward compute for no memory benefit on 48GB/80GB)

Expected: ~0.2-2s/pair on A100 vs 51s/pair. Resume/OOM-skip/fingerprint tests pass.

Co-Authored-By: Claude <noreply@anthropic.com>

deploy_to_hf_space.py

@@ -41,7 +41,8 @@ def main() -> None:
         repo_type="space",
         folder_path=str(args.local_dir),
         path_in_repo="",
-        ignore_patterns=["outputs/*", "atlas/*", ".venv/*", "__pycache__/*", "*.pyc", ".DS_Store"],
+        ignore_patterns=["outputs/*", "atlas/*", ".venv/*", "__pycache__/*", "*.pyc", ".DS_Store",
+                          "*.tmp", "*.npz", "*.npy", "*.sqlite", "sub_zero_ckpt/*", "l*_census_raw*"],
     )
     print("[hf] upload complete")
 

qwip_atlas/sub_zero_surgery.py

@@ -589,40 +589,44 @@ def _capture_atp_gradients(
                     "data": atp_accum}, tmp)
         tmp.replace(ckpt_path)  # atomic; a kill mid-write can't corrupt it
 
-    for idx in tqdm(range(start_idx, n_pairs), desc="AtP gradient", unit="pair",
-                    initial=start_idx, total=n_pairs):
-        param_grads: Dict[Tuple[int, str], torch.Tensor] = {}
-        hooks = []
+    # Build the flat list of scored (li, pname, weight) ONCE. Using
+    # torch.autograd.grad over exactly these leaves (instead of a full
+    # loss.backward() + a register_hook per weight) is the key perf fix: the
+    # old hook did g.detach().float().cpu() on every scored weight, which
+    # forced up to ~96 individual GPU->CPU syncs per pair and serialised the
+    # whole backward (~51 s/pair on a 6000 Ada). Now the grads stay on GPU and
+    # only one tiny [rank]-sized vector crosses to CPU per projection.
+    scored_params: List[Tuple[int, str, torch.nn.Parameter]] = []
+    for li, layer in enumerate(layers):
+        if li not in proj_svd:
+            continue
+        for pname, pmod in _get_projection_map(layer).items():
+            if pname in proj_svd[li] and hasattr(pmod, "weight") and pmod.weight.requires_grad:
+                scored_params.append((li, pname, pmod.weight))
+    scored_weight_list = [w for _, _, w in scored_params]
 
-        for li, layer in enumerate(layers):
-            if li not in proj_svd:
-                continue
-            for pname, pmod in _get_projection_map(layer).items():
-                if pname not in proj_svd[li]:
-                    continue
-                def _mk(li=li, pname=pname, w=pmod.weight):
-                    def _hook(g):
-                        param_grads[(li, pname)] = g.detach().float().cpu()
-                    hooks.append(w.register_hook(_hook))
-                _mk()
+    # Pre-stage the SVD right singular vectors on the model device once, so the
+    # per-pair grad->SV reduction is a single small GPU matmul (not a CPU one).
+    vh_gpu: Dict[Tuple[int, str], torch.Tensor] = {}
+    for li, pname, _ in scored_params:
+        _, _, vh = proj_svd[li][pname]
+        vh_gpu[(li, pname)] = vh.to(device=device, dtype=torch.float32, non_blocking=True)
 
+    for idx in tqdm(range(start_idx, n_pairs), desc="AtP gradient", unit="pair",
+                    initial=start_idx, total=n_pairs):
         enc = tokenizer(corp_prompts[idx], return_tensors="pt", truncation=True, max_length=max_length)
-        enc = {k: v.to(device) for k, v in enc.items()}
-        model.zero_grad()
+        enc = {k: v.to(device, non_blocking=True) for k, v in enc.items()}
         try:
             out = model(**enc, use_cache=False)
             logits = out.logits[0, :-1, :]
             targets = enc["input_ids"][0, 1:]
             loss = F.cross_entropy(logits, targets)
-            loss.backward()
+            grads = torch.autograd.grad(loss, scored_weight_list, allow_unused=True)
         except torch.cuda.OutOfMemoryError:
             # A single pathological pair must not wall the whole stage. Free
             # what we can, drop this pair's partial state, advance, and flush
             # so a re-run does not retry the same OOMing pair forever.
-            for h in hooks:
-                h.remove()
-            param_grads.clear()
-            enc.clear()
+            del enc
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
             gc.collect()
@@ -630,33 +634,25 @@ def _capture_atp_gradients(
             _flush(idx + 1)
             continue
 
-        for h in hooks:
-            h.remove()
-
-        for li in range(len(layers)):
-            if li not in proj_svd:
+        for (li, pname, _w), g_w in zip(scored_params, grads):
+            if g_w is None:
+                continue
+            _, _, vh = proj_svd[li][pname]
+            c_act = corp_proj_acts[li].get(pname)
+            a_act = auth_proj_acts[li].get(pname)
+            if c_act is None or a_act is None or c_act.shape[-1] != vh.shape[1]:
                 continue
-            for pname in proj_svd[li]:
-                g_w = param_grads.get((li, pname))
-                if g_w is None:
-                    continue
-
-                u, s, vh = proj_svd[li][pname]
-                c_act = corp_proj_acts[li].get(pname)
-                a_act = auth_proj_acts[li].get(pname)
-
-                if c_act is None or a_act is None or c_act.shape[-1] != vh.shape[1]:
-                    continue
 
-                c_sv = (c_act @ vh.T).mean(0)
-                a_sv = (a_act @ vh.T).mean(0)
-                diff = c_sv - a_sv
+            c_sv = (c_act @ vh.T).mean(0)        # CPU, small ([rank])
+            a_sv = (a_act @ vh.T).mean(0)        # CPU, small
+            diff = c_sv - a_sv                    # CPU, small
 
-                g_sv = (g_w @ vh.T).norm(dim=0)
-                atp_accum[li][pname].append(diff * g_sv)
+            # grad->SV reduction on GPU; one tiny CPU transfer at the end.
+            g_sv = (g_w.float() @ vh_gpu[(li, pname)].T).norm(dim=0).cpu()
+            atp_accum[li][pname].append(diff * g_sv)
 
         # Explicit cleanup to prevent CUDA memory fragmentation across pairs.
-        del loss, out, logits, targets, enc, param_grads
+        del loss, out, logits, targets, enc, grads
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         if idx % 50 == 0:
@@ -1467,8 +1463,11 @@ def build_brain_atlas(
     # requires_grad on all ~8B params allocates a .grad buffer for every
     # parameter during backward, which OOMs a 48GB card already holding a 16GB
     # model; scoped grad keeps peak memory to the scored MLP projections.
-    # Gradient checkpointing trades a little compute for not retaining every
-    # layer's forward activations through backward.
+    # No gradient checkpointing: with scoped grad the model + ~11GB of grad
+    # buffers + batch-1 activations fit in 48GB (and trivially in 80GB), so
+    # checkpointing would just double the backward compute by recomputing the
+    # forward. The per-pair OOM-skip guard inside _capture_atp_gradients is the
+    # safety net if a pathological pair ever exceeds headroom.
     scored_weights = []
     for li, layer in enumerate(layers):
         if li not in proj_svd:
@@ -1478,13 +1477,6 @@ def build_brain_atlas(
                 pmod.weight.requires_grad_(True)
                 scored_weights.append(pmod.weight)
 
-    model.train()
-    try:
-        model.gradient_checkpointing_enable()
-    except Exception as _e:
-        print(f"[sub_zero_surgery] gradient checkpointing unavailable "
-              f"({_e.__class__.__name__}); proceeding without")
-
     try:
         atp_scores = _capture_atp_gradients(
             model, tokenizer, corp[:32], auth[:32], layers,
@@ -1492,13 +1484,8 @@ def build_brain_atlas(
             ckpt_dir=ckpt_dir, fingerprint=fp_atp, ckpt_name="stage3_atp",
         )
     finally:
-        try:
-            model.gradient_checkpointing_disable()
-        except Exception:
-            pass
         for w in scored_weights:
             w.requires_grad_(False)
-        model.eval()
         for p in model.parameters():
             p.requires_grad_(False)
     print("  AtP complete")