align load_checkpoint args to save_checkpoint args

EleutherAI · Apr 30, 2021 · 3f42973 · 3f42973
1 parent 1d56953
commit 3f42973
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/megatron/checkpointing.py b/megatron/checkpointing.py
@@ -136,7 +136,7 @@ def save_checkpoint(neox_args, iteration, model, optimizer, lr_scheduler):
  torch.distributed.barrier()
 
 
-def load_checkpoint(model, optimizer, lr_scheduler, neox_args):
+def load_checkpoint(neox_args, model, optimizer, lr_scheduler):
  """Load a model checkpoint and return the iteration."""
 
  # Read the tracker file and set the iteration.

diff --git a/megatron/training.py b/megatron/training.py
@@ -337,7 +337,7 @@ def setup_model_and_optimizer(neox_args, inference=False, get_key_value=True):
  raise ValueError("Must be using deepspeed to run neox")
 
  if neox_args.load is not None:
- neox_args.iteration = load_checkpoint(model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, neox_args=neox_args)
+ neox_args.iteration = load_checkpoint(neox_args=neox_args, model=model, optimizer=optimizer, lr_scheduler=lr_scheduler)
  print_rank_0(f'Loading checkpoint and starting from iteration {neox_args.iteration}')
  else:
  neox_args.iteration = 0