For wandb, see wandb object_parameter, it's like following:

def wandb_hp_space(trial):
     return {
         "method": "random",
         "metric": {"name": "objective", "goal": "minimize"},
         "parameters": {
             "learning_rate": {"distribution": "uniform", "min": 1e-6, "max": 1e-4},
             "per_device_train_batch_size": {"values": [16, 32, 64, 128]},
         },
     }

Define a model_init function and pass it to the [Trainer], as an example: