[BUG] Checkpoint error in Text Generation for GPT

Hey,

I am getting the following error when I am trying to run the generate_sample_gpt.py script.

[2022-01-17 17:22:04,627] [INFO] [runner.py:327:main] Using IP address of 130.20.68.20 for node 127.0.0.1
[2022-01-17 17:22:04,627] [INFO] [runner.py:398:main] cmd = /people/shar703/anaconda3/envs/deepspeed/bin/python -u -m deepspeed.launcher.launch --world_info=eyIxMjcuMC4wLjEiOiBbMSwgMl19 --master_addr=130.20.68.20 --master_port=29500 tools/generate_samples_gpt.py --num-layers 24 --hidden-size 1024 --num-attention-heads 16 --seq-length 1024 --max-position-embeddings 1024 --micro-batch-size 4 --global-batch-size 8 --lr 0.00015 --train-iters 5000 --lr-decay-iters 320 --lr-decay-style cosine --vocab-file ../deepspeed_megatron/gpt_files/gpt2-vocab.json --merge-file ../deepspeed_megatron/gpt_files/gpt2-merges.txt --lr-warmup-fraction .01 --fp16 --load checkpoints/gpt2_345m --out-seq-length 1024 --temperature 1.0 --genfile samples.json --num-samples 2 --top_p 0.9 --recompute
[2022-01-17 17:22:05,729] [INFO] [launch.py:80:main] WORLD INFO DICT: {'127.0.0.1': [1, 2]}
[2022-01-17 17:22:05,729] [INFO] [launch.py:86:main] nnodes=1, num_local_procs=2, node_rank=0
[2022-01-17 17:22:05,729] [INFO] [launch.py:99:main] global_rank_mapping=defaultdict(<class 'list'>, {'127.0.0.1': [0, 1]})
[2022-01-17 17:22:05,729] [INFO] [launch.py💯main] dist_world_size=2
[2022-01-17 17:22:05,729] [INFO] [launch.py:102:main] Setting CUDA_VISIBLE_DEVICES=1,2
using world size: 2, data-parallel-size: 2, tensor-model-parallel size: 1, pipeline-model-parallel size: 1
using torch.float16 for parameters ...
------------------------ arguments ------------------------
accumulate_allreduce_grads_in_fp32 .............. False
adam_beta1 ...................................... 0.9
adam_beta2 ...................................... 0.999
adam_eps ........................................ 1e-08
adlr_autoresume ................................. False
adlr_autoresume_interval ........................ 1000
apply_query_key_layer_scaling ................... True
apply_residual_connection_post_layernorm ........ False
attention_dropout ............................... 0.1
attention_softmax_in_fp32 ....................... False
bert_binary_head ................................ True
bert_load ....................................... None
bf16 ............................................ False
bias_dropout_fusion ............................. True
bias_gelu_fusion ................................ True
biencoder_projection_dim ........................ 0
biencoder_shared_query_context_model ............ False
block_data_path ................................. None
checkpoint_activations .......................... False
checkpoint_in_cpu ............................... False
checkpoint_num_layers ........................... 1
clip_grad ....................................... 1.0
consumed_train_samples .......................... 0
consumed_train_tokens ........................... 0
consumed_valid_samples .......................... 0
contigious_checkpointing ........................ False
cpu_optimizer ................................... False
cpu_torch_adam .................................. False
curriculum_learning ............................. False
data_impl ....................................... infer
data_parallel_size .............................. 2
data_path ....................................... None
dataloader_type ................................. single
DDP_impl ........................................ local
decoder_seq_length .............................. None
deepscale ....................................... False
deepscale_config ................................ None
deepspeed ....................................... False
deepspeed_activation_checkpointing .............. False
deepspeed_config ................................ None
deepspeed_mpi ................................... False
distribute_checkpointed_activations ............. False
distributed_backend ............................. nccl
embedding_path .................................. None
encoder_seq_length .............................. 1024
eod_mask_loss ................................... False
eval_interval ................................... 1000
eval_iters ...................................... 100
evidence_data_path .............................. None
exit_duration_in_mins ........................... None
exit_interval ................................... None
ffn_hidden_size ................................. 4096
finetune ........................................ False
fp16 ............................................ True
fp16_lm_cross_entropy ........................... False
fp32_residual_connection ........................ False
genfile ......................................... samples.json
global_batch_size ............................... 8
greedy .......................................... False
hidden_dropout .................................. 0.1
hidden_size ..................................... 1024
hysteresis ...................................... 2
ict_head_size ................................... None
ict_load ........................................ None
img_dim ......................................... 224
indexer_batch_size .............................. 128
indexer_log_interval ............................ 1000
init_method_std ................................. 0.02
init_method_xavier_uniform ...................... False
initial_loss_scale .............................. 4294967296
kv_channels ..................................... 64
layernorm_epsilon ............................... 1e-05
lazy_mpu_init ................................... None
load ............................................ checkpoints/gpt2_345m
local_rank ...................................... 0
log_batch_size_to_tensorboard ................... False
log_interval .................................... 100
log_learning_rate_to_tensorboard ................ True
log_loss_scale_to_tensorboard ................... True
log_num_zeros_in_grad ........................... False
log_params_norm ................................. False
log_timers_to_tensorboard ....................... False
log_validation_ppl_to_tensorboard ............... False
loss_scale ...................................... None
loss_scale_window ............................... 1000
lr .............................................. 0.00015
lr_decay_iters .................................. 320
lr_decay_samples ................................ None
lr_decay_style .................................. cosine
lr_decay_tokens ................................. None
lr_warmup_fraction .............................. 0.01
lr_warmup_iters ................................. 0
lr_warmup_samples ............................... 0
make_vocab_size_divisible_by .................... 128
mask_prob ....................................... 0.15
masked_softmax_fusion ........................... True
max_position_embeddings ......................... 1024
memory_centric_tiled_linear ..................... False
merge_file ...................................... ../deepspeed_megatron/gpt_files/gpt2-merges.txt
micro_batch_size ................................ 4
min_loss_scale .................................. 1.0
min_lr .......................................... 0.0
mmap_warmup ..................................... False
no_load_optim ................................... True
no_load_rng ..................................... True
no_save_optim ................................... None
no_save_rng ..................................... None
num_attention_heads ............................. 16
num_channels .................................... 3
num_classes ..................................... 1000
num_layers ...................................... 24
num_layers_per_virtual_pipeline_stage ........... None
num_samples ..................................... 2
num_workers ..................................... 2
onnx_safe ....................................... None
openai_gelu ..................................... False
optimizer ....................................... adam
out_seq_length .................................. 1024
override_lr_scheduler ........................... False
params_dtype .................................... torch.float16
partition_activations ........................... False
patch_dim ....................................... 16
pipeline_model_parallel_size .................... 1
profile_backward ................................ False
query_in_block_prob ............................. 0.1
rampup_batch_size ............................... None
rank ............................................ 0
recompute ....................................... True
remote_device ................................... none
reset_attention_mask ............................ False
reset_position_ids .............................. False
retriever_report_topk_accuracies ................ []
retriever_score_scaling ......................... False
retriever_seq_length ............................ 256
sample_input_file ............................... None
sample_output_file .............................. None
sample_rate ..................................... 1.0
save ............................................ None
save_interval ................................... None
scatter_gather_tensors_in_pipeline .............. True
scattered_embeddings ............................ False
seed ............................................ 1234
seq_length ...................................... 1024
sgd_momentum .................................... 0.9
short_seq_prob .................................. 0.1
split ........................................... 969, 30, 1
split_transformers .............................. False
synchronize_each_layer .......................... False
temperature ..................................... 1.0
tensor_model_parallel_size ...................... 1
tensorboard_dir ................................. None
tensorboard_log_interval ........................ 1
tensorboard_queue_size .......................... 1000
tile_factor ..................................... 1
titles_data_path ................................ None
tokenizer_type .................................. GPT2BPETokenizer
top_k ........................................... 0
top_p ........................................... 0.9
train_iters ..................................... 5000
train_samples ................................... None
train_tokens .................................... None
use_checkpoint_lr_scheduler ..................... False
use_contiguous_buffers_in_ddp ................... False
use_cpu_initialization .......................... None
use_one_sent_docs ............................... False
use_pin_memory .................................. False
virtual_pipeline_model_parallel_size ............ None
vocab_extra_ids ................................. 0
vocab_file ...................................... ../deepspeed_megatron/gpt_files/gpt2-vocab.json
weight_decay .................................... 0.01
world_size ...................................... 2
zero_allgather_bucket_size ...................... 0.0
zero_contigious_gradients ....................... False
zero_reduce_bucket_size ......................... 0.0
zero_reduce_scatter ............................. False
zero_stage ...................................... 1.0
-------------------- end of arguments ---------------------
setting number of micro-batches to constant 1

building GPT2BPETokenizer tokenizer ...
padded vocab (size: 50257) with 47 dummy tokens (new size: 50304)
initializing torch distributed ...
initializing tensor model parallel with size 1
initializing pipeline model parallel with size 1
setting random seeds to 1234 ...
initializing model parallel cuda seeds on global rank 0, model parallel rank 0, and data parallel rank 0 with model parallel seed: 3952 and data parallel seed: 1234
compiling dataset index builder ...
make: Entering directory /qfs/people/shar703/scripts/mega_ai/Megatron-DeepSpeed/megatron/data' make: Nothing to be done for default'.
make: Leaving directory `/qfs/people/shar703/scripts/mega_ai/Megatron-DeepSpeed/megatron/data'

done with dataset index builder. Compilation time: 0.071 seconds
compiling and loading fused kernels ...
Detected CUDA files, patching ldflags
Emitting ninja build file /qfs/people/shar703/scripts/mega_ai/Megatron-DeepSpeed/megatron/fused_kernels/build/build.ninja...
Building extension module scaled_upper_triang_masked_softmax_cuda...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module scaled_upper_triang_masked_softmax_cuda...
Detected CUDA files, patching ldflags
Emitting ninja build file /qfs/people/shar703/scripts/mega_ai/Megatron-DeepSpeed/megatron/fused_kernels/build/build.ninja...
Building extension module scaled_masked_softmax_cuda...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module scaled_masked_softmax_cuda...
Detected CUDA files, patching ldflags
Emitting ninja build file /qfs/people/shar703/scripts/mega_ai/Megatron-DeepSpeed/megatron/fused_kernels/build/build.ninja...
Building extension module fused_mix_prec_layer_norm_cuda...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module fused_mix_prec_layer_norm_cuda...
done with compiling and loading fused kernels. Compilation time: 7.641 seconds
building GPT model ...
number of parameters on (tensor, pipeline) model parallel rank (0, 0): 354871296
loading checkpoint from checkpoints/gpt2_345m at iteration 5000
could not load the checkpoint
[Errno 2] No such file or directory: 'checkpoints/gpt2_345m/iter_0005000/mp_rank_00/model_optim_rng.pt'

Start with tools/generate_samples_gpt.py and the checkpoint path checkpoints/gpt2_345m, then reproduce the command shown in the issue with the same distributed settings. The payload does not include the checkpoint error itself, so identify the failing load or generation step before defining the fix; done means the command completes successfully.