Add W&B lr logging and project naming

[morganmcg1]

Adds:

• logging of learning rate to Weights & Biases in device_train.py
• ability to change the W&B project name in the config if needed (for device_train and train.py)
• wandb_project to the configs
• name and wandb_project description to doc string
• updates wandb version to >=11.2

[kingoflolz]

@morganmcg1 Thanks for making this PR, agree with @nostalgebraist about the use of optax step count. Happy to merge once that's updated and merge conflict is resolved

[morganmcg1]

Thanks @kingoflolz, @nostalgebraist for the comments, conflict addressed. Also, wandb_project was moved from the arguments to the config since your last review.

The PR now adds:

• logging of learning rate to Weights & Biases in device_train.py
• ability to change the W&B project name in the config (for device_train and train.py)
• wandb_project to the existing config files
• name and wandb_project description to doc string
• updates wandb version to >=0.11.2

[kingoflolz]

Thanks!