Architecture sweep: GPU affinity, arch search space, train.py overrides

- train.py: add --d-model, --n-layers, --n-heads, --d-ff, --lr,
--weight-decay, --warmup-steps flags to override named variants
- sweep.py: add --n-gpus (round-robin GPU pinning via CUDA_VISIBLE_DEVICES)
and --total-steps for architecture/pretrain sweeps
- pawn/sweep.py: new suggest_architecture() search space (150M-500M range),
AdapterObjective pins trials to GPUs, handles pretrain/architecture modes

pawn/sweep.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 
 import argparse
 import json
+import os
 import subprocess
 import sys
 from pathlib import Path
@@ -110,7 +111,7 @@ def suggest_tiny(trial: "optuna.Trial") -> dict:
 
 
 def suggest_pretrain(trial: "optuna.Trial") -> dict:
-    """Pretraining hyperparameters."""
+    """Pretraining hyperparameters (fixed architecture, tune training)."""
     return {
         "lr": trial.suggest_float("lr", 1e-5, 1e-3, log=True),
         "batch_size": trial.suggest_categorical("batch_size", [128, 256, 512]),
@@ -120,6 +121,29 @@ def suggest_pretrain(trial: "optuna.Trial") -> dict:
     }
 
 
+def suggest_architecture(trial: "optuna.Trial") -> dict:
+    """Architecture search space for pretraining.
+
+    Explores model size, depth/width tradeoff, and training hyperparameters.
+    Target budget: 150M-500M parameters on 80GB GPUs.
+    """
+    d_model = trial.suggest_categorical("d_model", [512, 640, 768, 896, 1024, 1280])
+    n_layers = trial.suggest_int("n_layers", 8, 24, step=2)
+    n_heads = trial.suggest_categorical("n_heads", [8, 16])
+    d_ff_mult = trial.suggest_categorical("d_ff_mult", [3, 4, 5])
+
+    return {
+        "d_model": d_model,
+        "n_layers": n_layers,
+        "n_heads": n_heads,
+        "d_ff": d_model * d_ff_mult,
+        "lr": trial.suggest_float("lr", 1e-5, 1e-3, log=True),
+        "batch_size": trial.suggest_categorical("batch_size", [128, 256]),
+        "weight_decay": trial.suggest_float("weight_decay", 0.0, 0.1),
+        "warmup_steps": trial.suggest_int("warmup_steps", 500, 3000, step=500),
+    }
+
+
 SUGGEST_FNS = {
     "lora": suggest_lora,
     "bottleneck": suggest_bottleneck,
@@ -127,6 +151,7 @@ SUGGEST_FNS = {
     "sparse": suggest_sparse,
     "hybrid": suggest_hybrid,
     "tiny": suggest_tiny,
+    "architecture": suggest_architecture,
     "pretrain": suggest_pretrain,
 }
 
@@ -138,6 +163,7 @@ ADAPTER_SCRIPTS = {
     "hybrid": "scripts/train_hybrid.py",
     "tiny": "scripts/train_tiny.py",
     "pretrain": "scripts/train.py",
+    "architecture": "scripts/train.py",
 }
 
 
@@ -208,6 +234,7 @@ class AdapterObjective:
         device: str = "cuda",
         output_base: str = "sweeps",
         epochs: int = 50,
+        n_gpus: int = 1,
         extra_args: list[str] | None = None,
     ):
         self.adapter_type = adapter_type
@@ -217,6 +244,7 @@ class AdapterObjective:
         self.output_base = Path(output_base) / adapter_type
         self.output_base.mkdir(parents=True, exist_ok=True)
         self.epochs = epochs
+        self.n_gpus = n_gpus
         self.extra_args = extra_args or []
         self.script = ADAPTER_SCRIPTS[adapter_type]
 
@@ -229,13 +257,20 @@ class AdapterObjective:
         # Build command
         cmd = [sys.executable, self.script]
 
-        # Fixed args
-        if self.adapter_type != "pretrain":
+        # Fixed args — architecture and pretrain sweeps use train.py directly
+        if self.adapter_type not in ("pretrain", "architecture"):
             cmd.extend(["--checkpoint", self.checkpoint])
             cmd.extend(["--pgn", self.pgn])
         cmd.extend(["--device", self.device])
-        cmd.extend(["--output-dir", str(trial_dir)])
-        if "epochs" not in params:
+
+        # output-dir for adapters, log-dir for pretraining
+        if self.adapter_type in ("pretrain", "architecture"):
+            cmd.extend(["--log-dir", str(trial_dir)])
+            cmd.extend(["--local-checkpoints"])
+        else:
+            cmd.extend(["--output-dir", str(trial_dir)])
+
+        if "epochs" not in params and "total_steps" not in params:
             cmd.extend(["--epochs", str(self.epochs)])
 
         # Suggested hyperparameters
@@ -244,11 +279,18 @@ class AdapterObjective:
         # Extra user-provided args
         cmd.extend(self.extra_args)
 
+        # GPU affinity: pin trial to GPU (trial.number % n_gpus)
+        env = os.environ.copy()
+        if self.n_gpus > 1:
+            gpu_id = trial.number % self.n_gpus
+            env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+
         # Run training
         result = subprocess.run(
             cmd,
             capture_output=True,
             text=True,
+            env=env,
         )
 
         if result.returncode != 0:

scripts/sweep.py

@@ -56,9 +56,13 @@ def main():
     p.add_argument("--n-trials", type=int, default=30,
                     help="Number of trials to run")
     p.add_argument("--n-jobs", type=int, default=1,
-                    help="Parallel trials (careful with GPU memory)")
+                    help="Parallel trials (match --n-gpus for multi-GPU)")
+    p.add_argument("--n-gpus", type=int, default=1,
+                    help="Number of GPUs. Trials are pinned to GPUs round-robin.")
     p.add_argument("--epochs", type=int, default=30,
-                    help="Max epochs per trial")
+                    help="Max epochs per trial (adapter sweeps)")
+    p.add_argument("--total-steps", type=int, default=20000,
+                    help="Total steps per trial (architecture/pretrain sweeps)")
     p.add_argument("--device", type=str, default="cuda")
     p.add_argument("--pruner", type=str, default="hyperband",
                     choices=["hyperband", "median", "none"])
@@ -69,7 +73,7 @@ def main():
 
     args = p.parse_args()
 
-    if args.adapter != "pretrain" and not args.pgn:
+    if args.adapter not in ("pretrain", "architecture") and not args.pgn:
         p.error("--pgn is required for adapter sweeps")
 
     study_name = args.study_name or args.adapter
@@ -77,8 +81,11 @@ def main():
 
     print(f"=== PAWN Hyperparameter Sweep ===")
     print(f"Adapter: {args.adapter}")
-    print(f"Trials: {args.n_trials} (parallel: {args.n_jobs})")
-    print(f"Epochs/trial: {args.epochs}")
+    print(f"Trials: {args.n_trials} (parallel: {args.n_jobs}, GPUs: {args.n_gpus})")
+    if args.adapter in ("pretrain", "architecture"):
+        print(f"Steps/trial: {args.total_steps}")
+    else:
+        print(f"Epochs/trial: {args.epochs}")
     print(f"Pruner: {args.pruner}")
     print(f"Storage: {db_path}")
     print(f"Dashboard: uv run optuna-dashboard {db_path}")
@@ -90,6 +97,11 @@ def main():
         pruner=args.pruner,
     )
 
+    # For architecture/pretrain sweeps, pass --total-steps via extra args
+    extra = list(args.extra_args)
+    if args.adapter in ("pretrain", "architecture"):
+        extra.extend(["--total-steps", str(args.total_steps)])
+
     objective = AdapterObjective(
         adapter_type=args.adapter,
         checkpoint=args.checkpoint,
@@ -97,7 +109,8 @@ def main():
         device=args.device,
         output_base=args.output_dir,
         epochs=args.epochs,
-        extra_args=args.extra_args,
+        n_gpus=args.n_gpus,
+        extra_args=extra,
     )
 
     study.optimize(

scripts/train.py

@@ -31,6 +31,15 @@ def parse_args():
     parser.add_argument("--discard-ply-limit", action="store_true",
                         help="Only train on games that ended naturally (no ply limit truncation)")
 
+    # Architecture overrides (for sweeps — override the named variant)
+    parser.add_argument("--d-model", type=int, default=None, help="Override d_model")
+    parser.add_argument("--n-layers", type=int, default=None, help="Override n_layers")
+    parser.add_argument("--n-heads", type=int, default=None, help="Override n_heads")
+    parser.add_argument("--d-ff", type=int, default=None, help="Override d_ff")
+    parser.add_argument("--lr", type=float, default=None, help="Override learning rate")
+    parser.add_argument("--weight-decay", type=float, default=None, help="Override weight decay")
+    parser.add_argument("--warmup-steps", type=int, default=None, help="Override warmup steps")
+
     ckpt_group = parser.add_mutually_exclusive_group(required=True)
     ckpt_group.add_argument("--hf-repo", type=str, default=None,
                             help="Push checkpoints to this HuggingFace repo (requires HF_TOKEN)")
@@ -76,6 +85,22 @@ def main():
     if args.discard_ply_limit:
         train_cfg.discard_ply_limit = True
 
+    # Architecture overrides
+    if args.d_model is not None:
+        model_cfg.d_model = args.d_model
+    if args.n_layers is not None:
+        model_cfg.n_layers = args.n_layers
+    if args.n_heads is not None:
+        model_cfg.n_heads = args.n_heads
+    if args.d_ff is not None:
+        model_cfg.d_ff = args.d_ff
+    if args.lr is not None:
+        train_cfg.lr = args.lr
+    if args.weight_decay is not None:
+        train_cfg.weight_decay = args.weight_decay
+    if args.warmup_steps is not None:
+        train_cfg.warmup_steps = args.warmup_steps
+
     print(f"Model config: {model_cfg}")
     print(f"Training config: {train_cfg}")