From e001a04f536076544228316ad653ff70029735c6 Mon Sep 17 00:00:00 2001 From: Dashiell Stander Date: Tue, 17 Oct 2023 14:43:09 -0400 Subject: [PATCH] LR scheduler fix no longer breaks inference (#1060) * Add lr_scheduler check for inference. Signed-off-by: Dashiell Stander * Update NeoXArgs docs automatically --------- Signed-off-by: Dashiell Stander Co-authored-by: github-actions --- configs/neox_arguments.md | 2 +- megatron/training.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/configs/neox_arguments.md b/configs/neox_arguments.md index 3b4b253f7..badc95e46 100644 --- a/configs/neox_arguments.md +++ b/configs/neox_arguments.md @@ -111,7 +111,7 @@ Logging Arguments - **git_hash**: str - Default = ec71f71 + Default = a97bd1f current git hash of repository diff --git a/megatron/training.py b/megatron/training.py index 0f9f7a0e0..548f81cb0 100644 --- a/megatron/training.py +++ b/megatron/training.py @@ -686,7 +686,9 @@ def setup_model_and_optimizer(neox_args, use_cache=False, iteration=None): neox_args.iteration = 0 # need this for correct lr scheduling resume from ckpt - lr_scheduler.optimizer = model.optimizer + # but it will not exist if this is being called for inference + if lr_scheduler is not None: + lr_scheduler.optimizer = model.optimizer return model, optimizer, lr_scheduler