Optuna integration for hyperparameter optimization

[forklady42]

Optuna setup for initial coarse hyperparameter search. WandB runs are under betsy-mp-hpo-tier1.

Best trial over 15 epochs:

=== Best Trial ===
Trial number: 71
Value (val_loss): 0.033743

Best hyperparameters:
  n_channels: 32
  n_residual_blocks: 1
  kernel_size: 5
  lr: 0.003580430023049488

[forklady42]

@claude review this PR

[claude]

Claude finished @forklady42's task in 2m 4s —— View job

---

PR Review: Optuna Integration for Hyperparameter Optimization

Tasks:

• Gather context and read changed files
• Review hpo.py entrypoint
• Review YAML config files
• Review pyproject.toml changes
• Post comprehensive review feedback

---

Overall this is a solid and well-structured HPO integration. The design choices (ddp_spawn over ddp, per-trial W&B runs, optional dependency group) are all sensible. A few issues worth addressing before merge:

Bugs / Correctness

val_loss metric name ambiguity

lightning.py logs val_loss with on_step=True, on_epoch=True:

self.log("val_loss", loss, on_step=True, on_epoch=True, sync_dist=True)

In PyTorch Lightning, this creates both val_loss_step and val_loss_epoch keys during training. The plain val_loss key in callback_metrics should resolve to the epoch aggregate at epoch end, but this depends on Lightning version behavior. The ModelCheckpoint and EarlyStopping monitors, and the final trainer.callback_metrics.get("val_loss") at hpo.py:210, all rely on this. Worth verifying the logged key name matches exactly — if it resolves to val_loss_epoch instead, all three will silently fail (checkpoint won't save, early stopping won't trigger, and val_loss will be None, causing all trials to prune).

Fix this →

---

Issues

train_workers/val_workers commented out in tier1 config (config_hpo_tier1.yaml:49-50)

  # train_workers: 8
  # val_workers: 2

If RhoRead defaults to 0 workers, data loading will be single-threaded and a significant bottleneck across 50+ trials. The full config has these set — the tier1 config should too unless there's a deliberate reason to drop them.

Fix this →

---

OOM heuristic ignores batch size (hpo.py:131-140)

The memory check guards against large channel counts but doesn't account for batch_size, which is also in the full search space in config_hpo.yaml. A combination like batch_size=4, depth=4, n_channels=32 could OOM without triggering the bottleneck check. Consider incorporating batch size into the heuristic, or moving the check after batch_size is applied.

---

No installation hint in docstring (hpo.py:1-8)

optuna and optuna-integration are optional extras. The module-level imports on lines 19 and 26 will raise ImportError if the user runs the script without uv sync --extra hpo. The docstring at the top should note this requirement.

Fix this →

---

Nits

WandB reinit=True is deprecated (hpo.py:162)

In WandB SDK ≥0.18, reinit=True is deprecated in favor of reinit="finish_previous". Worth updating to avoid future breakage, though it works for now.

Checkpoint accumulation (hpo.py:169-172)

With n_trials=50, this will create 50 subdirectories under hpo_checkpoints/, each with a saved checkpoint. For a coarse search that's probably fine, but in longer runs this could consume significant disk space. Consider adding a note or cleanup step.

gradient_clip_value: 20.0 is a very high upper bound (config_hpo.yaml:51)

The search space includes [1.0, 5.0, 20.0] for gradient clipping. A value of 20.0 is effectively no clipping for most training scenarios and could destabilize learning. Consider capping at 10.0 or using [0.5, 1.0, 5.0].

sqlite:/// storage for parallel trials (config_hpo.yaml:9, config_hpo_tier1.yaml:8)

SQLite is fine for sequential trials but can have contention with n_jobs > 1. If you ever parallelize the outer study (multiple HPO processes), consider switching to PostgreSQL or RDB storage.

---

[forklady42]

We're not currently experimenting with batch_size. It causing an OOM isn't currently a concern.

@claude re-review this PR

[claude]

Claude finished @forklady42's task in 1m 56s —— View job

---

PR Re-Review: Optuna Integration for Hyperparameter Optimization

Tasks:

• Read changed files
• Re-review with batch_size OOM concern dropped
• Post updated review

---

Good progress — the previous round of comments has been well addressed. One new bug surfaced during this pass.

Addressed Since Previous Review

• val_loss metric ambiguity — Fixed. validation_step now logs with on_step=False, on_epoch=True, so val_loss in callback_metrics unambiguously resolves to the epoch aggregate. All three monitors (ModelCheckpoint, EarlyStopping, callback_metrics.get) will work correctly.
• train_workers/val_workers commented out — Fixed. Both are now set in config_hpo_tier1.yaml.
• No installation hint in docstring — Fixed. Docstring now includes the uv sync --extra hpo note.
• WandB reinit=True deprecated — Fixed. Updated to reinit="finish_previous".
• OOM heuristic / batch_size — Not a concern; noted and dropped.

---

Bug: src.electrai imports in HPO files

The parent commit 8c4418a explicitly fixed all from src.electrai. imports and _target_: src.electrai. references across the codebase. The new files in this PR reintroduce the old pattern:

hpo.py:31:

from src.electrai.lightning import LightningGenerator  # broken

All other entrypoints use from electrai.lightning import LightningGenerator. The src. prefix only works if the project root is on sys.path, which isn't guaranteed when the package is installed.

config_hpo.yaml:63,78 and config_hpo_tier1.yaml:44,58:

_target_: src.electrai.dataloader.dataset.RhoRead   # broken
_target_: src.electrai.model.resunet.ResUNet3D       # broken

Compare with the existing config_resunet.yaml which correctly uses:

_target_: electrai.dataloader.dataset.RhoRead
_target_: electrai.model.resunet.ResUNet3D

Fix this →

---

Remaining Nits

Checkpoint accumulation — With n_trials=50, hpo_checkpoints/ and hpo_tier1_checkpoints/ will each accumulate 50 subdirectories. Fine for a coarse search, but worth a comment or a cleanup note for longer runs.

gradient_clip_value: 20.0 in config_hpo.yaml (config_hpo.yaml:51) — choices: [1.0, 5.0, 20.0] still includes 20.0, which is effectively no clipping. Minor, but could destabilize some trials.

SQLite storage for parallel outer-loop runs — No concern if n_jobs=1 (the default), but worth a note if you ever parallelize across multiple HPO processes.

---

[forklady42]

Rebased and fixed the src.electrai references. Ready for review

[hanaol]

Thanks for looking into this Betsy. Did you happen to check larger kernel sizes like 7?

[forklady42]

@hanaol I initially tested 3 and 5 for kernel size. Since there's architectural changes still in flight, I've been limiting some of the sweeps until those land.

To be clear, this PR focuses on the Optuna integration for running hyperparameter sweeps, rather than the output of the sweep itself.

[hanaol]

I understand. Since you reported the best trial, I was curious which kernel sizes you had explored. I can now see the [3,5] case in your code.

[forklady42]

@hanaol anything you'd like addressed on this branch or should I go ahead and merge so we can keep iterating on hyperparameter sweep code?

[forklady42]

@claude review once more

[claude]

Claude encountered an error —— View job

---

I'll analyze this and get back to you.

[hanaol]

@forklady42 I had missed this one, reviewing now...

[hanaol]

If a new parameter is added to search_space in the YAML but not added here, it will be silently dropped. This could produce a study with no error or warning.

There could be a guard at the top:

known_params = {"depth", "n_channels", "n_residual_blocks", "kernel_size",
                "lr", "weight_decay", "warmup_length", "gradient_clip_value",
                "batch_size", "augmentation"}
unknown = set(params) - known_params
if unknown:
    raise ValueError(f"Unknown hyperparameters with no config mapping: {unknown}")

[hanaol]

Curious how you arrived at the heuristic. Do you think the spatial dimensions of the input grid and precision should also factor in?

[hanaol]

Makes sense for HPO since on_step=True creates two metrics (val_loss and val_loss_epoch) which can confuse the monitoring callbacks. That said, per-step validation tracking is still useful during early model development to catch instabilities mid-epoch. Would it be worth making this configurable (e.g. a flag in the config) so we can track it closely while we're still iterating on the architecture?

[hanaol]

Minor comments. Feel free to merge.