feat(probe): per-probe per-layer sweep pipeline with parallel coordinator

- configs: add layer 0 to probe_layers ([0, 10, 20, 30, 39]) for both models
- run_probe.py: add --layer arg to scope sweep/final to a single layer;
  sweep mode loops over probe_layers when --layer is not set
- src/probe.py: run_sweep uses probe_layers[0] as sweep target (no middle-layer
  heuristic); create_sweep names sweep with layer suffix; run_sweep/run_final
  tag and group W&B runs by layer; run_final merges into existing results.pt
  so per-layer jobs accumulate correctly
- probe_sweep_coordinator.sh: redesigned as a pure orchestrator — creates one
  W&B sweep, submits N parallel worker jobs, submits final training with
  --dependency=afterok on all workers; --count is total trials (divided by
  --n-agents internally); coordinator exits immediately after submitting
- probe_sweep.sh / probe_final.sh: bump --mem to 256G for large layer caches
- extract_swebench_thin.sh / extract_swebench.sh: add --chunk-size 8192 to
  handle long trajectories without OOM on thin nodes
- README: rewrite to document the current agentic SWE-bench pipeline only

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

Author: andre15silva

README.md

@@ -1,3 +1,3 @@
 model_id: "poolside/Laguna-XS.2"
-probe_layers: [10, 20, 30, 39]
+probe_layers: [0, 10, 20, 30, 39]
 adapter: laguna

configs/models/laguna_xs2.yaml

@@ -1,3 +1,3 @@
 model_id: "Qwen/Qwen3.6-35B-A3B"
-probe_layers: [10, 20, 30, 39]
+probe_layers: [0, 10, 20, 30, 39]
 adapter: qwen

configs/models/qwen36_35b_a3b.yaml

@@ -62,6 +62,9 @@ def main():
     eval_p.add_argument("--output-run-id", required=True,
                         help="Run ID under which to save results.pt")
 
+    parser.add_argument("--layer", type=int, default=None,
+                        help="Run sweep/final for this single layer only (overrides model config loop)")
+
     args = parser.parse_args()
     if args.n_eval_bins is not None and args.n_eval_bins != args.n_bins and args.n_bins != 1:
         parser.error("--n-eval-bins requires --n-bins 1 or --n-eval-bins equal to --n-bins")
@@ -73,26 +76,29 @@ def main():
             run_id=args.run_id,
             probe_name=args.probe,
             fixed_params=args.fixed,
+            layer=args.layer,
         )
     elif args.mode == "sweep":
-        run_sweep(
-            run_id=args.run_id,
-            probe_name=args.probe,
-            probe_layers=model_cfg.probe_layers,
-            seed=args.seed,
-            sweep_id=args.sweep_id,
-            count=args.count,
-            cache_dir=args.cache_dir,
-            cache_run_id=cache_run_id,
-            n_bins=args.n_bins,
-            probe_arch=args.probe_arch,
-            n_eval_bins=args.n_eval_bins,
-            eval_bin_axis=args.eval_bin_axis,
-            then_final=args.then_final,
-            results_dir=args.results_dir,
-            fixed_params=args.fixed,
-            shuffle_labels=args.shuffle_labels,
-        )
+        layers = [args.layer] if args.layer is not None else model_cfg.probe_layers
+        for layer in layers:
+            run_sweep(
+                run_id=args.run_id,
+                probe_name=args.probe,
+                probe_layers=[layer],
+                seed=args.seed,
+                sweep_id=args.sweep_id,
+                count=args.count,
+                cache_dir=args.cache_dir,
+                cache_run_id=cache_run_id,
+                n_bins=args.n_bins,
+                probe_arch=args.probe_arch,
+                n_eval_bins=args.n_eval_bins,
+                eval_bin_axis=args.eval_bin_axis,
+                then_final=args.then_final,
+                results_dir=args.results_dir,
+                fixed_params=args.fixed,
+                shuffle_labels=args.shuffle_labels,
+            )
     elif args.mode == "eval":
         run_eval(
             weights_run_id=args.weights_run_id or args.run_id,
@@ -111,10 +117,11 @@ def main():
             tool_nll_run_id=args.tool_nll_run_id,
         )
     else:
+        layers = [args.layer] if args.layer is not None else model_cfg.probe_layers
         run_final(
             run_id=args.run_id,
             probe_name=args.probe,
-            probe_layers=model_cfg.probe_layers,
+            probe_layers=layers,
             lr=args.lr,
             weight_decay=args.weight_decay,
             batch_size=args.batch_size,

run_probe.py

@@ -42,4 +42,5 @@ uv run python run_extract_swebench.py \
   --shard-rank "$RANK" \
   --num-shards "$NUM_SHARDS" \
   --extraction-batch-size 1 \
+  --chunk-size 8192 \
   "$@"

slurm/extract_swebench.sh

@@ -33,4 +33,5 @@ uv run python run_extract_swebench.py \
   --shard-rank "$RANK" \
   --num-shards "$NUM_SHARDS" \
   --extraction-batch-size 1 \
+  --chunk-size 8192 \
   "$@"

slurm/extract_swebench_thin.sh

@@ -2,7 +2,7 @@
 #SBATCH -J pp-probe-final
 #SBATCH -p berzelius-cpu
 #SBATCH -n 8
-#SBATCH --mem=150G
+#SBATCH --mem=256G
 #SBATCH -t 02:00:00
 #SBATCH -o logs/probe_final_%j.out
 #SBATCH -e logs/probe_final_%j.err

slurm/probe_final.sh

@@ -2,7 +2,7 @@
 #SBATCH -J pp-probe-sweep
 #SBATCH -p berzelius-cpu
 #SBATCH -n 8
-#SBATCH --mem=150G
+#SBATCH --mem=256G
 #SBATCH -t 04:00:00
 #SBATCH -o logs/probe_sweep_%j.out
 #SBATCH -e logs/probe_sweep_%j.err

slurm/probe_sweep.sh

@@ -1,52 +1,74 @@
 #!/bin/bash
+# Coordinator for a single (probe, layer) sweep.
+# Creates one W&B sweep, submits N parallel worker jobs, then submits final
+# training as a dependent job that runs after all workers complete.
+#
+# Usage — one call per (probe, layer):
+#   sbatch slurm/probe_sweep_coordinator.sh \
+#     --model-config configs/models/laguna_xs2.yaml \
+#     --layer 0 \
+#     --probe currently_correct \
+#     --probe-arch linear \
+#     --run-id laguna_xs2_full_pooled \
+#     --cache-dir cache/swebench \
+#     --cache-run-id laguna_xs2_full \
+#     --results-dir results/swebench \
+#     --n-bins 1 \
+#     --n-agents 4 \
+#     --count 5
+#
 #SBATCH -J pp-probe-coord
 #SBATCH -p berzelius-cpu
-#SBATCH -n 8
-#SBATCH --mem=150G
-#SBATCH -t 04:00:00
+#SBATCH -n 2
+#SBATCH --mem=8G
+#SBATCH -t 00:10:00
 #SBATCH -o logs/probe_sweep_coordinator_%j.out
 #SBATCH -e logs/probe_sweep_coordinator_%j.err
 
 set -euo pipefail
 mkdir -p logs
 
 # ---------------------------------------------------------------------------
-# Parse --n-agents and --count; collect everything else as probe_args
+# Parse --n-agents and --count; pass everything else through to run_probe.py
 # ---------------------------------------------------------------------------
-N_AGENTS=2
-COUNT=10
+N_AGENTS=4
+COUNT=20
 PROBE_ARGS=()
 
 while [[ $# -gt 0 ]]; do
     case "$1" in
-        --n-agents)  N_AGENTS="$2"; shift 2 ;;
-        --count)     COUNT="$2";    shift 2 ;;
-        *)           PROBE_ARGS+=("$1"); shift ;;
+        --n-agents) N_AGENTS="$2"; shift 2 ;;
+        --count)    COUNT="$2";    shift 2 ;;
+        *)          PROBE_ARGS+=("$1"); shift ;;
     esac
 done
 
+COUNT_PER_AGENT=$(( COUNT / N_AGENTS ))
+
 # ---------------------------------------------------------------------------
-# 1. Create the sweep and capture only the ID (printed to stdout by create_sweep)
+# 1. Create the sweep — layer name baked in via --layer in PROBE_ARGS
 # ---------------------------------------------------------------------------
-echo "[coordinator] Creating sweep (n_agents=${N_AGENTS}, count_per_agent=${COUNT})"
+echo "[coordinator] Creating sweep (n_agents=${N_AGENTS}, count=${COUNT}, count_per_agent=${COUNT_PER_AGENT})"
 SWEEP_ID=$(uv run python run_probe.py "${PROBE_ARGS[@]}" create-sweep)
 echo "[coordinator] Sweep ID: ${SWEEP_ID}"
 
 # ---------------------------------------------------------------------------
-# 2. Submit N-1 additional agent jobs
+# 2. Submit N parallel worker jobs
 # ---------------------------------------------------------------------------
-AGENT_SLURM_TIME="${SLURM_TIMELIMIT:-04:00:00}"
-
-for i in $(seq 2 "${N_AGENTS}"); do
-    JOB_ID=$(sbatch --parsable \
-        --time="${AGENT_SLURM_TIME}" \
+WORKER_IDS=()
+for i in $(seq 1 "${N_AGENTS}"); do
+    JID=$(sbatch --parsable \
         slurm/probe_sweep.sh \
-        "${PROBE_ARGS[@]}" sweep --sweep-id "${SWEEP_ID}" --count "${COUNT}")
-    echo "[coordinator] Submitted agent ${i} as SLURM job ${JOB_ID}"
+        "${PROBE_ARGS[@]}" sweep --sweep-id "${SWEEP_ID}" --count "${COUNT_PER_AGENT}")
+    WORKER_IDS+=("${JID}")
+    echo "[coordinator] Worker ${i} → job ${JID}"
 done
 
 # ---------------------------------------------------------------------------
-# 3. Run as agent 1
+# 3. Submit final training, dependent on all workers completing successfully
 # ---------------------------------------------------------------------------
-echo "[coordinator] Running as agent 1"
-uv run python run_probe.py "${PROBE_ARGS[@]}" sweep --sweep-id "${SWEEP_ID}" --count "${COUNT}"
+DEP="afterok:$(IFS=:; echo "${WORKER_IDS[*]}")"
+FINAL_JID=$(sbatch --parsable --dependency="${DEP}" \
+    slurm/probe_final.sh \
+    "${PROBE_ARGS[@]}" final --from-sweep "${SWEEP_ID}")
+echo "[coordinator] Final training → job ${FINAL_JID} (depends on ${WORKER_IDS[*]})"

slurm/probe_sweep_coordinator.sh

@@ -360,10 +360,11 @@ def train_probe_layer(
     return results, weights
 
 
-def create_sweep(run_id: str, probe_name: str, fixed_params: dict | None = None) -> str:
+def create_sweep(run_id: str, probe_name: str, fixed_params: dict | None = None, layer: int | None = None) -> str:
     import wandb
+    layer_suffix = f"-layer{layer}" if layer is not None else ""
     sweep_config = {
-        "name": f"sweep-{run_id}-{probe_name}",
+        "name": f"sweep-{run_id}-{probe_name}{layer_suffix}",
         "method": "random",
         "metric": {"name": "mean_val_auc", "goal": "maximize"},
         "parameters": {
@@ -427,16 +428,16 @@ def run_sweep(
     import wandb
     cache_base = Path(cache_dir) / (cache_run_id or run_id) / probe_name
 
-    middle_layer = probe_layers[len(probe_layers) // 2]
-    middle_cache = str(cache_base / f"layer_{middle_layer}.pt")
+    sweep_layer = probe_layers[0]
+    sweep_cache = str(cache_base / f"layer_{sweep_layer}.pt")
 
     if sweep_id is None:
-        sweep_id = create_sweep(run_id, probe_name, fixed_params=fixed_params)
+        sweep_id = create_sweep(run_id, probe_name, fixed_params=fixed_params, layer=sweep_layer)
 
     def sweep_fn():
         with wandb.init(
-            group=f"{run_id}/{probe_name}",
-            tags=[run_id, probe_name, "sweep"],
+            group=f"{run_id}/{probe_name}/layer_{sweep_layer}",
+            tags=[run_id, probe_name, f"layer_{sweep_layer}", "sweep"],
         ) as run:
             cfg = run.config
 
@@ -451,7 +452,7 @@ def log_fn(metrics: dict) -> None:
                 })
 
             results, _ = train_probe_layer(
-                middle_cache, middle_layer,
+                sweep_cache, sweep_layer,
                 lr=cfg.lr, weight_decay=cfg.weight_decay,
                 batch_size=cfg.batch_size, patience=cfg.patience,
                 seed=seed, n_bins=n_bins, probe_arch=probe_arch,
@@ -543,12 +544,14 @@ def run_final(
     all_results: dict[int, list[ProbeResult]] = {}
     all_weights: dict[int, dict] = {}
 
+    layer_suffix = f"/layer_{probe_layers[0]}" if len(probe_layers) == 1 else ""
     with wandb.init(
         project="program-probes",
         job_type="final",
-        name=f"final-{run_id}-{probe_name}-{probe_arch}",
-        group=f"{run_id}/{probe_name}",
-        tags=[run_id, probe_name, "final", probe_arch],
+        name=f"final-{run_id}-{probe_name}-{probe_arch}{layer_suffix}",
+        group=f"{run_id}/{probe_name}{layer_suffix}",
+        tags=[run_id, probe_name, "final", probe_arch]
+              + ([f"layer_{probe_layers[0]}"] if len(probe_layers) == 1 else []),
     ) as run:
         for layer_idx in probe_layers:
             cache_path = str(cache_base / f"layer_{layer_idx}.pt")
@@ -597,8 +600,21 @@ def log_fn(metrics: dict) -> None:
 
     out = Path(results_dir) / run_id / probe_name
     out.mkdir(parents=True, exist_ok=True)
-    torch.save(all_results, out / "results.pt")
-    torch.save(all_weights, out / "weights.pt")
+
+    results_path = out / "results.pt"
+    if results_path.exists():
+        existing = torch.load(results_path, weights_only=False)
+        existing.update(all_results)
+        all_results = existing
+
+    weights_path = out / "weights.pt"
+    if weights_path.exists():
+        existing_w = torch.load(weights_path, weights_only=False)
+        existing_w.update(all_weights)
+        all_weights = existing_w
+
+    torch.save(all_results, results_path)
+    torch.save(all_weights, weights_path)
     return {layer: [vars(r) for r in rs] for layer, rs in all_results.items()}