Add Modal GPU benchmark workflow and script

- `modal/benchmark.py`: Configurable benchmark using `electrai-data`
  Volume (dataset_4, 2,885 samples). Filters by file size (≤100MB
  for L4), subsamples, uses temp data root with symlinks. Gradient
  checkpointing enabled. Handles `RhoData` trailing-newline bug.
- `.github/workflows/gpu-benchmark-modal.yml`: GHA workflow wrapping
  `modal run modal/benchmark.py` with configurable inputs.
- Default config (8ch/2blk) fits on L4; use `--gpu A100 --channels 32
  --residual-blocks 16` for production-size benchmarks.

Verified: 50 samples, 5 epochs, L4 → green.

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

Author: ryan-williams

.github/workflows/gpu-benchmark-modal.yml

@@ -0,0 +1,70 @@
+name: GPU Benchmark (Modal)
+
+on:
+  push:
+    branches: [modal-benchmark]
+  workflow_dispatch:
+    inputs:
+      gpu:
+        description: 'GPU type (L4, A10G, A100, H100)'
+        default: 'L4'
+        type: string
+      epochs:
+        description: 'Number of training epochs'
+        default: '5'
+        type: string
+      channels:
+        description: 'Number of model channels (8=tiny/L4, 32=prod/A100, 64=large)'
+        default: '8'
+        type: string
+      residual_blocks:
+        description: 'Number of residual blocks (2=tiny/L4, 16=prod/A100)'
+        default: '2'
+        type: string
+      samples:
+        description: 'Number of samples from Volume (0=all eligible)'
+        default: '50'
+        type: string
+      max_file_size:
+        description: 'Skip samples larger than N MB (0=no limit, 100=safe for L4)'
+        default: '100'
+        type: string
+      wandb_project:
+        description: 'WandB project name'
+        default: 'elf-net-ci'
+        type: string
+
+jobs:
+  benchmark:
+    runs-on: ubuntu-latest
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install Modal
+        run: uv pip install --system modal
+
+      - name: Run benchmark
+        env:
+          GPU: ${{ inputs.gpu || 'L4' }}
+          EPOCHS: ${{ inputs.epochs || '5' }}
+          CHANNELS: ${{ inputs.channels || '8' }}
+          RESIDUAL_BLOCKS: ${{ inputs.residual_blocks || '2' }}
+          SAMPLES: ${{ inputs.samples || '50' }}
+          MAX_FILE_SIZE: ${{ inputs.max_file_size || '100' }}
+          WANDB_PROJECT: ${{ inputs.wandb_project || 'elf-net-ci' }}
+        run: |
+          modal run modal/benchmark.py \
+            --gpu "$GPU" \
+            --epochs "$EPOCHS" \
+            --channels "$CHANNELS" \
+            --residual-blocks "$RESIDUAL_BLOCKS" \
+            --samples "$SAMPLES" \
+            --max-file-size "$MAX_FILE_SIZE" \
+            --wandb-project "$WANDB_PROJECT"

modal/benchmark.py

@@ -4,11 +4,11 @@
 electrai-data Volume, logs metrics to WandB, and reports wall-clock time.
 
 Usage:
-    # Default: 50 samples, 5 epochs, 32 channels, L4
+    # Default: 50 samples, 5 epochs, 8 channels, 2 blocks, L4
     modal run modal/benchmark.py
 
-    # Production model on A100
-    modal run modal/benchmark.py --gpu A100 --channels 64 --residual-blocks 16
+    # Production model on A100 (32ch/16blk needs >=40GB VRAM)
+    modal run modal/benchmark.py --gpu A100 --channels 32 --residual-blocks 16
 
     # All 2,885 Volume samples
     modal run modal/benchmark.py --samples 0
@@ -59,10 +59,10 @@
 )
 def run_benchmark(
     epochs: int = 5,
-    channels: int = 32,
-    residual_blocks: int = 16,
+    channels: int = 8,
+    residual_blocks: int = 2,
     samples: int = 50,
-    max_file_size: float = 25,
+    max_file_size: float = 100,
     seed: int = 42,
     wandb_project: str = "elf-net-ci",
     gpu_type: str = "L4",
@@ -77,22 +77,52 @@ def run_benchmark(
     sys.path.insert(0, "/root/electrai/scripts")
     from e2e_train import run_training
 
-    # Subsample from volume filelist if requested
+    # Filter and subsample from volume filelist
     filelist = Path(DATA_ROOT) / "mp_filelist.txt"
-    all_ids = filelist.read_text().strip().splitlines()
-    if samples > 0 and samples < len(all_ids):
+    all_ids = [s for s in filelist.read_text().strip().splitlines() if s]
+
+    # Filter by file size first (avoid OOM on large grids)
+    if max_file_size > 0:
+        max_bytes = int(max_file_size * 1024 * 1024)
+        eligible = [
+            sid
+            for sid in all_ids
+            if (Path(DATA_ROOT) / "data" / f"{sid}.CHGCAR").stat().st_size <= max_bytes
+        ]
+        log.info(
+            "File size filter: %d/%d eligible (<=%.0fMB)",
+            len(eligible),
+            len(all_ids),
+            max_file_size,
+        )
+    else:
+        eligible = all_ids
+
+    # Subsample from eligible
+    if 0 < samples < len(eligible):
         import random
 
         random.seed(seed)
-        subset = sorted(random.sample(all_ids, samples))
-        sub_filelist = Path(DATA_ROOT) / "mp_filelist_benchmark.txt"
-        sub_filelist.write_text("\n".join(subset) + "\n")
-        data_root = DATA_ROOT
-        log.info("Subsampled %d/%d samples", samples, len(all_ids))
+        subset = sorted(random.sample(eligible, samples))
+        log.info("Subsampled %d/%d eligible samples", samples, len(eligible))
     else:
-        data_root = DATA_ROOT
-        samples = len(all_ids)
-        log.info("Using all %d samples", samples)
+        subset = eligible
+        samples = len(subset)
+        log.info("Using all %d eligible samples", samples)
+
+    if not subset:
+        raise ValueError(
+            f"No eligible samples (total={len(all_ids)}, max_file_size={max_file_size}MB)"
+        )
+
+    # Create temp data root with symlinks + clean filelist (no trailing newline)
+    data_root = "/tmp/benchmark_data"
+    Path(data_root).mkdir(parents=True, exist_ok=True)
+    for subdir in ["data", "label"]:
+        link = Path(data_root) / subdir
+        if not link.exists():
+            link.symlink_to(Path(DATA_ROOT) / subdir)
+    Path(data_root, "mp_filelist.txt").write_text("\n".join(subset))
 
     log.info(
         "Benchmark: gpu=%s, epochs=%d, channels=%d, blocks=%d, samples=%d",
@@ -109,8 +139,9 @@ def run_benchmark(
         epochs=epochs,
         seed=seed,
         gpu=True,
+        gradient_checkpoint=True,  # save VRAM for large grids
         data_root=data_root,
-        max_file_size=max_file_size,
+        max_file_size=0,  # already filtered above
         wandb_project=wandb_project,
         verbose=True,
     )
@@ -127,10 +158,10 @@ def run_benchmark(
 def main(
     gpu: str = "L4",
     epochs: int = 5,
-    channels: int = 32,
-    residual_blocks: int = 16,
+    channels: int = 8,
+    residual_blocks: int = 2,
     samples: int = 50,
-    max_file_size: float = 25,
+    max_file_size: float = 100,
     seed: int = 42,
     wandb_project: str = "elf-net-ci",
 ):