Reduce attribution GPU synchronization

evaluate_single_disease_mortality_attribution.py

@@ -555,7 +555,7 @@ def main() -> None:
 
     batch_size = int(cfg_get(args, cfg, "batch_size", 128))
     attribution_batch_size = int(
-        cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 8, batch_size))
+        cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096))
     )
     if attribution_batch_size <= 0:
         raise ValueError("attribution_batch_size must be positive")
@@ -613,15 +613,20 @@ def main() -> None:
     row_base_cache: dict[int, dict[str, Any]] = {}
     pending_batch_chunks: list[Dict[str, torch.Tensor]] = []
     pending_meta_chunks: list[list[dict[str, Any]]] = []
+    pending_orig_risk_chunks: list[torch.Tensor] = []
+    pending_orig_hazard_chunks: list[torch.Tensor] = []
     pending_n = 0
 
     def flush_pending() -> None:
-        nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks, pending_n
+        nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks
+        nonlocal pending_orig_risk_chunks, pending_orig_hazard_chunks, pending_n
         if pending_n == 0:
             return
 
         ablated_batch = concat_padded_tensor_batches(pending_batch_chunks)
         meta_rows = [row for chunk in pending_meta_chunks for row in chunk]
+        orig_risk = torch.cat(pending_orig_risk_chunks, dim=0).to(device)
+        orig_hazard = torch.cat(pending_orig_hazard_chunks, dim=0).to(device)
         with torch.no_grad():
             ablated_risk = death_risk_for_batch(
                 model=model,
@@ -634,35 +639,33 @@ def main() -> None:
                 tau=tau,
             )
         ablated_hazard = mortality_hazard_from_risk(ablated_risk)
-        orig_risk = torch.as_tensor(
-            [row.pop("_death_risk") for row in meta_rows],
-            dtype=ablated_risk.dtype,
-            device=ablated_risk.device,
-        )
-        orig_hazard = torch.as_tensor(
-            [row.pop("_death_hazard") for row in meta_rows],
-            dtype=ablated_hazard.dtype,
-            device=ablated_hazard.device,
-        )
-
         attr_prob = orig_risk - ablated_risk
         attr_hazard = orig_hazard - ablated_hazard
         ratio_prob = safe_ratio(orig_risk, ablated_risk, eps=float(args.ratio_eps))
         ratio_hazard = safe_ratio(orig_hazard, ablated_hazard, eps=float(args.ratio_eps))
+        value_block = torch.stack(
+            [
+                orig_risk,
+                orig_hazard,
+                ablated_risk,
+                ablated_hazard,
+                attr_prob,
+                attr_hazard,
+                ratio_prob,
+                ratio_hazard,
+            ],
+            dim=1,
+        ).detach().cpu().numpy()
 
         for i, row in enumerate(meta_rows):
-            row["death_risk"] = float(orig_risk[i].detach().cpu())
+            row["death_risk"] = float(value_block[i, 0])
-            row["death_hazard"] = float(orig_hazard[i].detach().cpu())
+            row["death_hazard"] = float(value_block[i, 1])
-            row["ablated_death_risk"] = float(ablated_risk[i].detach().cpu())
+            row["ablated_death_risk"] = float(value_block[i, 2])
-            row["ablated_death_hazard"] = float(ablated_hazard[i].detach().cpu())
+            row["ablated_death_hazard"] = float(value_block[i, 3])
-            row["mortality_attribution_probability"] = float(attr_prob[i].detach().cpu())
+            row["mortality_attribution_probability"] = float(value_block[i, 4])
-            row["mortality_attribution_hazard"] = float(attr_hazard[i].detach().cpu())
+            row["mortality_attribution_hazard"] = float(value_block[i, 5])
-            row["mortality_attribution_probability_ratio"] = float(
+            row["mortality_attribution_probability_ratio"] = float(value_block[i, 6])
-                ratio_prob[i].detach().cpu()
+            row["mortality_attribution_hazard_ratio"] = float(value_block[i, 7])
-            )
-            row["mortality_attribution_hazard_ratio"] = float(
-                ratio_hazard[i].detach().cpu()
-            )
 
         table = pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)
         update_summary_accumulator(summary_accumulator, table)
@@ -674,6 +677,8 @@ def main() -> None:
         written_rows += len(meta_rows)
         pending_batch_chunks = []
         pending_meta_chunks = []
+        pending_orig_risk_chunks = []
+        pending_orig_hazard_chunks = []
         pending_n = 0
 
     def get_row_base(row_idx: int) -> dict[str, Any]:
@@ -779,7 +784,6 @@ def main() -> None:
                         f"{disease_token} for row {row_idx}, but cached history has count 0"
                     )
 
-                orig_row = int(local_rows[local_pos].detach().cpu().item())
                 meta_chunk.append(
                     {
                         "patient_id": row_base["patient_id"],
@@ -803,14 +807,14 @@ def main() -> None:
                         "history_count__selected_organ_system": int(
                             group_counts.get(str(disease_meta.get("organ_system", "")), 0)
                         ),
-                        "_death_risk": float(death_risk_tensor[orig_row].detach().cpu()),
-                        "_death_hazard": float(death_hazard_tensor[orig_row].detach().cpu()),
                     }
                 )
 
             if meta_chunk:
                 pending_batch_chunks.append(ablated_chunk)
                 pending_meta_chunks.append(meta_chunk)
+                pending_orig_risk_chunks.append(death_risk_tensor[local_rows].detach())
+                pending_orig_hazard_chunks.append(death_hazard_tensor[local_rows].detach())
                 pending_n += len(meta_chunk)
             pair_offset = pair_stop