From 744d2cef84d16ec2b28d0562b4f82f1cb2f79e0c Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Mon, 18 May 2026 20:26:18 +0300 Subject: [PATCH] Log learning_rate to wandb during training Recreate the LR schedule in the training loop and log its value alongside the existing metrics, so the learning rate is visible in wandb without an external launcher patch. --- scripts/train.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/scripts/train.py b/scripts/train.py index 5d28941..ed2ba8a 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -255,6 +255,9 @@ def main(config: _config.TrainConfig): dynamic_ncols=True, ) + # Recreate the LR schedule (it is otherwise hidden inside the optimizer) so we can log it. + lr_schedule = config.lr_schedule.create() + infos = [] for step in pbar: with sharding.set_mesh(mesh): @@ -263,6 +266,7 @@ def main(config: _config.TrainConfig): if step % config.log_interval == 0: stacked_infos = common_utils.stack_forest(infos) reduced_info = jax.device_get(jax.tree.map(jnp.mean, stacked_infos)) + reduced_info["learning_rate"] = float(lr_schedule(step)) info_str = ", ".join(f"{k}={v:.4f}" for k, v in reduced_info.items()) pbar.write(f"Step {step}: {info_str}") wandb.log(reduced_info, step=step)