Bugfix for Iteration 0 checkpointing (#764)

* bugfix iteration 0 checkpointing * Update NeoXArgs docs automatically Co-authored-by: github-actions <[email protected]>
EleutherAI · Jan 6, 2023 · 375de3f · 375de3f
1 parent 008fbec
commit 375de3f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/configs/neox_arguments.md b/configs/neox_arguments.md
@@ -111,7 +111,7 @@ Logging Arguments
 
 - **git_hash**: str
 
- Default = 27e56e3
+ Default = 075a525
 
  current git hash of repository
 

diff --git a/megatron/training.py b/megatron/training.py
@@ -101,10 +101,10 @@ def pretrain(neox_args):
  timers.log(["model and optimizer", "train/valid/test data iterators"])
  print_rank_0("training ...")
 
- iteration = 0
+ iteration = neox_args.iteration
  if neox_args.do_train and neox_args.train_iters > 0:
- # edge case: save step 0 checkpoint if requested
- if neox_args.save and 0 in neox_args.save_iters:
+ # edge case: save step 0 checkpoint if requested and we're starting from step 0
+ if neox_args.save and 0 in neox_args.save_iters and iteration == 0:
  save_checkpoint(
  neox_args=neox_args,
  iteration=iteration,