Optuna doesn't converge at all (TPE)

[neel04]

I think the problem is pretty self-explanatory - optuna doesn't converge, even with a large number of trials (~100) and this is using a median-pruner as well.

From what I can gather, there's probably an incorrect feedback metric provided to the TPE sampler. I assumed that I was trial.report-ing the wrong value, but it seems ok. This is how my code is structured:

study = optuna.create_study(
    direction="minimize",
    load_if_exists=True,
    sampler=optuna.samplers.TPESampler(
        seed=69,
        consider_magic_clip=True,
        consider_endpoints=True,
        n_startup_trials=5,
    ),
    pruner=optuna.pruners.MedianPruner(
        n_startup_trials=5, n_warmup_steps=200, n_min_trials=10
    ),
)

  wandb_kwargs = {
      "project": "****",
      "config": args,
      "anonymous": "allow",
      "entity": "****",
      "magic": True,
  }

  trainer_kwargs = {
      "args": args,
      "loaders": (trainloader, valloader),
      "decode_fn": train_dataset.tok.decode,
      "key": key
  }

  study.optimize(
      lambda trial: kickoff_optuna(trial=trial, **trainer_kwargs),
      n_trials=100,
      callbacks=[wandbc]
  )

Where the kickoff function is defined as:

def kickoff_optuna(trial, **trainer_kwargs):
    args = trainer_kwargs['args']

    args.epochs = 1
    args.lr = trial.suggest_float('lr', 1e-4, 1e-3, step=1e-4)
    args.drop_rate = trial.suggest_float('drop_rate', 0.0, 0.1, step=0.02)
    args.weight_decay = trial.suggest_float('weight_decay', 1e-5, 1e-3, step=2e-4)
    args.warmup_steps = trial.suggest_int('warmup_steps', 0, 500, step=100)

    args = trainer_kwargs['args']
    ...
    trainer = Trainer(**trainer_kwargs)

    with jax.spmd_mode('allow_all'):
        loss = trainer.train(trial)
    
    return loss

Effectively I just override the arguments to those suggested by the sampler to keep things compatible with my codebase.
After that, is just reporting the trial metric:

for step, batch in tqdm(enumerate(self.trainloader), total=len(self.trainloader), desc=f'Epoch {epoch}'):

    make_step(.....)
    
    ...

    if step % 100 == 0:
        accuracy, loss, perplexity = self.compute_metrics(...)

        train_acc.append(accuracy)
        train_loss.append(loss)
        train_ppl.append(perplexity.mean())

        loss = loss.item()
        
        if trial is not None:
            loss = jnp.nan_to_num(loss, nan=9999.0) # handle NaNs
            trial.report(loss, step)
            
            if trial.should_prune() or jnp.isnan(loss):
                raise optuna.exceptions.TrialPruned()

I've tried doing more trials for longer, and changing other hyperparams to no avail.

[neel04]

@nzw0301 Is that alright? do you need some other details/information?

[neel04]

My run failed cuz I didn't have plotly and thus it couldn't viz it, but I changed something - the trial.report just reports a single metric now once - the loss at the end of the training, instead of every n steps.

This looks like its doing significantly better, but I'm not 100% sure

[neel04]

Also, this my optuna~y plot I made in wandb, if it helps:

Here, it converges to a run with relatively high dropout, which I don't really think is optimal. The optimal would probably be 0 regularization, high LR, medium WD

[neel04]

For context, the viz for the older run I posted (at the top) is this: (It reported loss at every n steps) @nzw0301

[nzw0301]

Here, it converges to a run with relatively high dropout, which I don't really think is optimal. The optimal would probably be 0 regularization, high LR, medium WD

If you think so, you can give such hyper-param config to your study by using https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.enqueue_trial. Unfortunately, in my understanding, TPESampler cannot guarantee researching the global solution.

[neel04]

That's a great feature - I'll try that out and see if my params actually outperform. Thanks!