LR scheduler fix no longer breaks inference (#1060)

* Add lr_scheduler check for inference. Signed-off-by: Dashiell Stander <[email protected]> * Update NeoXArgs docs automatically --------- Signed-off-by: Dashiell Stander <[email protected]> Co-authored-by: github-actions <[email protected]>
EleutherAI · Oct 17, 2023 · e001a04 · e001a04
1 parent f6ac04d
commit e001a04
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git 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
@@ -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