Integrate wandb sweeps

[ashleve]

[emilemathieu]

Would also love that! Would there be a way to specific a sweep_id so that a sweep launched with hydra via -m would appear in the wandb view?

[ashleve]

I think it should be possible to just add a 'wandb_sweep_id' parameter to the main config (then it would be easily overridable from command line) and use it to connect to existing sweep server (or create a new one) using wandb api (https://docs.wandb.ai/sweeps/python-api). Maybe I will add example of this when I have a time.

It would probably be the best if W&B released official sweeper for Hydra. Then such a sweep could automatically override chosen parameters from hydra config, without you needing to worry about it. I just made an issue about it here: https://github.com/wandb/client/issues/1856

Currently I use Optuna for hyperparameter search, since it has a custom sweeper for hydra, so launching it doesn't require coding any extra logic into train.py (I will add some example to repo soon).

[thomas-happify]

@ashleve possible work around:

1. create configs/hparams_search/mnist_wandb.yaml

program: run.py
args_no_hyphens: True
method: bayes
metric:
  name: val/acc
  goal: maximize
parameters:
  model.lr:
    min: 0.0001
    max: 0.1
  model.lin1_size:
    values: [32, 64, 128, 256, 512]
  model.lin2_size:
    values: [32, 64, 128, 256, 512]
  model.lin3_size:
    values: [32, 64, 128, 256, 512]
  datamodule.batch_size:
    values: [32, 64, 128]

early_terminate:
  type: hyperband
  s: 2
  eta: 3
  max_iter: 27

2. Remove "--" + in wandb/wandb_agent.py line 373 had to do this otherwise wamdb sweep will call something like python run.py --model.lr=0.0001 ...

wandb sweep configs/hparams_search/mnist_wandb.yaml
wandb agent your-sweep-id

[ashleve]

@thomas-happify thanks! I did not realize this integrates with hydra so effortlessly.

You actually don't even need the second step, just add the command section and ${args_no_hyphens} to config:

program: run.py

project: lightning-hydra-template

description: "MNIST wandb sweep"

method: bayes

metric:
  name: val/acc
  goal: maximize

parameters:
  model.lr:
    distribution: uniform
    min: 0.0001
    max: 0.1
  model.lin1_size:
    values: [32, 64, 128, 256, 512]
  model.lin2_size:
    values: [32, 64, 128, 256, 512]
  model.lin3_size:
    values: [32, 64, 128, 256, 512]
  datamodule.batch_size:
    values: [32, 64, 128]

early_terminate:
  type: hyperband
  s: 2
  eta: 3
  max_iter: 2

command:
  - ${interpreter}
  - ${program}
  - ${args_no_hyphens}
  - logger=wandb

It's hidden in the docs here: https://docs.wandb.ai/guides/sweeps/configuration#command

The cool thing is you can also hardcode the logger=wandb arg like I did above.

[thomas-happify]

@ashleve
awesome!

[ashleve]

@thomas-happify The only problem is parameters overriden in wandb config get doubled:

Lightning saves them with / while wandb sweeper uses . I'm not sure if this affects the performance of bayesian optimization

[tesfaldet]

There's this hydra-wandb-sweeper repo in case anyone's willing to revive it. I gave up after a single PR because wandb's code isn't open source. The author of the repo, at the time I was working on my PR for it, was working for W&B but it wasn't an "official" wandb Hydra sweeper afaik. The author even made a W&B blog post about having Hydra sweeps work with W&B sweeps (albeit nowhere near as integrated as my PR which enabled native wandb sweep support within Hydra but in the most janky way possible). It's been almost a year since my PR and I've moved to Aim instead for my experiment tracking and Orion for HPO.

EDIT: What I mean by my comment of wandb not being open source is that although the client codebase is open source, various important parts of the server-side code is closed source. Since I only had an understanding of how communication works from W&B client -> W&B server but not W&B server -> W&B client, it made it extremely difficult and sometimes impossible to extend the client so that it can access important information from the sweep controller (which is running server-side) and that would enable tighter Hydra integration. This is why you see a lot of monkeypatching in my PR.