1. Update the pinned `wandb` library version. The old version raises when given new `wandb` API keys!
2. Move the `step` argument to the right place in `wandb.log` calls. The signature is `wandb.log(data: dict, step: int, commit: bool)` - previously, step counts were being included in the data dict, meaning wandb metrics incorrectly had x-axes corresponding to the number of calls to `.log` instead of the number of training steps.
3. Move `wandb.init` later in `chat_sft.py` and `base_train.py` to include config values that are calculated or read from a checkpoint.
