diff --git a/training/colossalai/run_clm.py b/training/colossalai/run_clm.py index 72f5c8a..609cccc 100644 --- a/training/colossalai/run_clm.py +++ b/training/colossalai/run_clm.py @@ -600,6 +600,7 @@ def main(): train_perplexity = float("inf") writer.add_scalar("Train/Perplexity (Step)", train_perplexity, global_step) writer.add_scalar("Train/Loss (Step)", loss, global_step) + writer.add_scalar("Train/Learning Rate (Step)", lr_scheduler.get_last_lr()[-1], global_step) if args.output_dir is not None and args.checkpointing_steps is not None: if args.checkpointing_steps != "epoch" and completed_steps % int(args.checkpointing_steps) == 0: