switch_sweep.yaml

# program: train.py
method: random
metric:
  goal: minimize
  name: val_loss
parameters:
  LR:
    values: [0.0005, 0.001, 0.003, 0.005]
  batch_size:
    values: [16, 32, 64, 128, 256, 512]
  optimizer:
    values: ['adamw', 'sgd']
  activation:
    values: ['GELU', 'GEGLU', 'SwiGLU']
  n_experts:
    values: [2, 4, 8, 16, 32, 64]
  capacity_factor:
    values: [0.5, 0.75, 1, 1.25, 1.5, 1.75, 2]
  aux_loss_coef:
    values: [0.005, 0.01, 0.05, 0.1, 0.15]
  norm_first:
    values: [True, False]
  switch_first:
    values: [True, False]
  every_n_switch:
    values: [1, 2, 3, 4]
  mlp_dropout:
    values: [0.1, 0.2, 0.3, 0.4]
  expert_dropout:
    values: [0.1, 0.2, 0.3, 0.4]
  rope_scale:
    values: [0.25, 0.5, 0.75, 1]
    
early_terminate:
  type: hyperband
  s: 2
  eta: 3 # reduction factor, e.g. 1/eta of trials left after each check (i.e. (eta-1)/eta are discarded) 
  max_iter: 27
  ## e.g. brackets set at floor(max_iter/eta), floor(max_iter/eta**2), etc.
  ## here they are at 9, 3
  
# sweep using specified distribution
# parameters_dict.update({
#     'learning_rate': {
#         # a flat distribution between 0 and 0.1
#         'distribution': 'uniform',
#         'min': 0,
#         'max': 0.1
#       }})