Cleaned up load/save checkpoint printing

megatron/checkpointing.py

@@ -110,6 +110,11 @@ def save_checkpoint(iteration, model, optimizer, lr_scheduler):
  # Only rank zero of the data parallel writes to the disk.
  if isinstance(model, torchDDP):
  model = model.module
+
+ if torch.distributed.get_rank() == 0:
+ print('saving checkpoint at iteration {:7d} to {}'.format(
+ iteration, args.save), flush=True)
+
  if mpu.get_data_parallel_rank() == 0:
 
  # Arguments, iteration, and model.
@@ -137,14 +142,14 @@ def save_checkpoint(iteration, model, optimizer, lr_scheduler):
 
  # Save.
  checkpoint_name = get_checkpoint_name(args.save, iteration)
- print('global rank {} is saving checkpoint at iteration {:7d} to {}'.
- format(torch.distributed.get_rank(), iteration, checkpoint_name))
  ensure_directory_exists(checkpoint_name)
  torch.save(state_dict, checkpoint_name)
- print(' successfully saved {}'.format(checkpoint_name))
 
  # Wait so everyone is done (necessary)
  torch.distributed.barrier()
+ if torch.distributed.get_rank() == 0:
+ print(' successfully saved checkpoint at iteration {:7d} to {}'.format(
+ iteration, args.save), flush=True)
  # And update the latest iteration
  if torch.distributed.get_rank() == 0:
  tracker_filename = get_checkpoint_tracker_filename(args.save)
@@ -192,9 +197,9 @@ def load_checkpoint(model, optimizer, lr_scheduler, load_arg='load'):
 
  # Checkpoint.
  checkpoint_name = get_checkpoint_name(load_dir, iteration, release)
- if mpu.get_data_parallel_rank() == 0:
- print('global rank {} is loading checkpoint {}'.format(
- torch.distributed.get_rank(), checkpoint_name))
+ if torch.distributed.get_rank() == 0:
+ print(' loading checkpoint from {} at iteration {}'.format(
+ args.load, iteration), flush=True)
 
  # Load the checkpoint.
  try:
@@ -276,8 +281,9 @@ def load_checkpoint(model, optimizer, lr_scheduler, load_arg='load'):
  sys.exit()
 
  torch.distributed.barrier()
- if mpu.get_data_parallel_rank() == 0:
- print(' successfully loaded {}'.format(checkpoint_name))
+ if torch.distributed.get_rank() == 0:
+ print(' successfully loaded checkpoint from {} at iteration {}'.format(
+ args.load, iteration), flush=True)
 
  return iteration