[BUG] autotuning hangs out when an error occurs

Describe the bug

Autotuning hangs out at exp 13 (z3_gas1_tmbspg130):

[2023-08-07 03:46:00,874] [INFO] [scheduler.py:378:run_experiment] Launching exp_id = 12, exp_name = z3_gas1_tmbspg72, with resource = neox-autotune-worker-0:0,1,2,3@neox-autotune-worker-1:0,1,2,3, and ds_config = /mnt/4_2/autotune_results/z3_gas1_tmbspg72/ds_config.json
  0%|          | 0/1 [00:00<?, ?it/s]neox-autotune-worker-0: Warning: Permanently added 'neox-autotune-worker-0,10.233.72.138' (ECDSA) to the list of known hosts.
neox-autotune-worker-1: Warning: Permanently added 'neox-autotune-worker-1,10.233.71.222' (ECDSA) to the list of known hosts.
[2023-08-07 03:49:02,320] [INFO] [scheduler.py:430:clean_up] Done cleaning up exp_id = 12 on the following workers: neox-autotune-worker-0,neox-autotune-worker-1
[2023-08-07 03:49:02,321] [INFO] [scheduler.py:393:run_experiment] Done running exp_id = 12, exp_name = z3_gas1_tmbspg72, with resource = neox-autotune-worker-0:0,1,2,3@neox-autotune-worker-1:0,1,2,3
100%|██████████| 1/1 [03:06<00:00, 186.63s/it]
[2023-08-07 03:49:07,362] [INFO] [scheduler.py:344:run_experiment] Scheduler wrote ds_config to /mnt/4_2/autotune_results/z3_gas1_tmbspg130/ds_config.json, /mnt/4_2/autotune_results/z3_gas1_tmbspg130/ds_config.json
[2023-08-07 03:49:07,413] [INFO] [scheduler.py:351:run_experiment] Scheduler wrote exp to /mnt/4_2/autotune_results/z3_gas1_tmbspg130/exp.json,/mnt/4_2/autotune_results/z3_gas1_tmbspg130/exp.json
[2023-08-07 03:49:07,439] [INFO] [scheduler.py:378:run_experiment] Launching exp_id = 13, exp_name = z3_gas1_tmbspg130, with resource = neox-autotune-worker-0:0,1,2,3@neox-autotune-worker-1:0,1,2,3, and ds_config = /mnt/4_2/autotune_results/z3_gas1_tmbspg130/ds_config.json

stderr.log:

$cat 4_2/autotune_results/z3_gas1_tmbspg130/stderr.log | tail -n 10
neox-autotune-worker-0:   File "/workspace/gpt-neox/megatron/model/transformer.py", line 856, in forward
neox-autotune-worker-0:     attention_output, attention_bias = self.attention(
neox-autotune-worker-0:   File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1517, in _call_impl
neox-autotune-worker-0:     result = forward_call(*args, **kwargs)
neox-autotune-worker-0:   File "/workspace/gpt-neox/megatron/model/transformer.py", line 685, in forward
neox-autotune-worker-0:     context_layer = self.attention(
neox-autotune-worker-0:   File "/workspace/gpt-neox/megatron/model/transformer.py", line 415, in attention
neox-autotune-worker-0:     matmul_result = torch.baddbmm(
neox-autotune-worker-0: torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 16.25 GiB (GPU 2; 79.15 GiB total capacity; 45.24 GiB already allocated; 16.21 GiB free; 61.53 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
pdsh@neox-autotune-worker-0: neox-autotune-worker-0: ssh exited with exit code 1

stdout.log:

$cat 4_2/autotune_results/z3_gas1_tmbspg130/stdout.log | tail -n 10
neox-autotune-worker-0: [2023-08-07 03:50:09,175] [INFO] [checkpointing.py:529:forward] Activation Checkpointing Information
neox-autotune-worker-0: [2023-08-07 03:50:09,175] [INFO] [checkpointing.py:530:forward] ----Partition Activations True, CPU CHECKPOINTING False
neox-autotune-worker-0: [2023-08-07 03:50:09,175] [INFO] [checkpointing.py:531:forward] ----contiguous Memory Checkpointing False with 24 total layers
neox-autotune-worker-0: [2023-08-07 03:50:09,175] [INFO] [checkpointing.py:533:forward] ----Synchronization True
neox-autotune-worker-0: [2023-08-07 03:50:09,175] [INFO] [checkpointing.py:534:forward] ----Profiling time in checkpointing False
neox-autotune-worker-0: [2023-08-07 03:50:19,797] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 25123
neox-autotune-worker-0: [2023-08-07 03:50:19,805] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 25124
neox-autotune-worker-0: [2023-08-07 03:50:19,940] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 25125
neox-autotune-worker-0: [2023-08-07 03:50:20,315] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 25127
neox-autotune-worker-0: [2023-08-07 03:50:20,316] [ERROR] [launch.py:321:sigkill_handler] ['/usr/bin/python', '-u', 'train.py', '--local_rank=3', '/mnt/6-7B_autotune.yml', '--vocab_file', '/mnt/data/gpt2-vocab.json', '--merge_file', '/mnt/data/gpt2-merges.txt', '--data_path', '/mnt/data/enwik8/enwik8_text_document', '--deepspeed_config', 'eyJtb2RlbF9wYXJhbGxlbF9zaXplIjogMSwgIm9wdGltaXplciI6IHsidHlwZSI6ICJBZGFtIiwgInBhcmFtcyI6IHsibHIiOiAwLjAwMDEyLCAiYmV0YXMiOiBbMC45LCAwLjk1XSwgImVwcyI6IDFlLTA4fX0sICJhdXRvdHVuaW5nIjogeyJlbmFibGVkIjogdHJ1ZSwgIm1ldHJpY19wYXRoIjogIi90OWsvbW50LzRfMi9hdXRvdHVuZV9yZXN1bHRzL3ozX2dhczFfdG1ic3BnMTMwL21ldHJpY3MuanNvbiJ9LCAiemVyb19vcHRpbWl6YXRpb24iOiB7InN0YWdlIjogM30sICJmcDE2IjogeyJmcDE2IjogdHJ1ZSwgImVuYWJsZWQiOiB0cnVlLCAibG9zc19zY2FsZSI6IDAsICJsb3NzX3NjYWxlX3dpbmRvdyI6IDEwMDAsICJoeXN0ZXJlc2lzIjogMiwgIm1pbl9sb3NzX3NjYWxlIjogMX0sICJncmFkaWVudF9hY2N1bXVsYXRpb25fc3RlcHMiOiAxLCAidHJhaW5fbWljcm9fYmF0Y2hfc2l6ZV9wZXJfZ3B1IjogMTMwLCAidHJhaW5fYmF0Y2hfc2l6ZSI6IDEwNDB9'] exits with return code = 1

In fact, although ‘CUDA out of memory’ can cause errors, it does not always cause the experiment to hang.

At exp2 (z1_gas1_tmbspg121), deepspeed reported the CUDA out of memory error and proceed to the next experiment.

[2023-08-07 03:21:10,703] [INFO] [scheduler.py:344:run_experiment] Scheduler wrote ds_config to /mnt/4_2/autotune_results/z1_gas1_tmbspg121/ds_config.json, /mnt/4_2/autotune_results/z1_gas1_tmbspg121/ds_config.json
[2023-08-07 03:21:10,786] [INFO] [scheduler.py:351:run_experiment] Scheduler wrote exp to /mnt/4_2/autotune_results/z1_gas1_tmbspg121/exp.json, /mnt/4_2/autotune_results/z1_gas1_tmbspg121/exp.json
[2023-08-07 03:21:10,869] [INFO] [scheduler.py:378:run_experiment] Launching exp_id = 2, exp_name = z1_gas1_tmbspg121, with resource = neox-autotune-worker-0:0,1,2,3@neox-autotune-worker-1:0,1,2,3, and ds_config = /mnt/4_2/autotune_results/z1_gas1_tmbspg121/ds_config.json
  0%|          | 0/1 [00:00<?, ?it/s]neox-autotune-worker-1: Warning: Permanently added 'neox-autotune-worker-1,10.233.71.222' (ECDSA) to the list of known hosts.
neox-autotune-worker-0: Warning: Permanently added 'neox-autotune-worker-0,10.233.72.138' (ECDSA) to the list of known hosts.
pdsh@neox-autotune-worker-0: neox-autotune-worker-1: ssh exited with exit code 1
[2023-08-07 03:22:58,028] [INFO] [scheduler.py:430:clean_up] Done cleaning up exp_id = 2 on the following workers: neox-autotune-worker-0,neox-autotune-worker-1
[2023-08-07 03:22:58,029] [INFO] [scheduler.py:393:run_experiment] Done running exp_id = 2, exp_name = z1_gas1_tmbspg121, with resource = neox-autotune-worker-0:0,1,2,3@neox-autotune-worker-1:0,1,2,3
[2023-08-07 03:23:05,716] [INFO] [scheduler.py:225:parse_results] The experiment exp_id = 2, exp_name = z1_gas1_tmbspg121, did not run successfully with error = CUDA out of memory. Tried to allocate 46.44 GiB (GPU 3; 79.15 GiB total capacity; 52.45 GiB already allocated; 10.46 GiB free; 67.28 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
, thus a metrics.txt does not exist for it. Check the stderr.log in /mnt/4_2/autotune_results/z1_gas1_tmbspg121

z1_gas1_tmbspg121 stderr.log:

neox-autotune-worker-1:   File "/workspace/gpt-neox/megatron/training.py", line 811, in train
neox-autotune-worker-1:     loss_dict, skipped_iter = train_step(
neox-autotune-worker-1:   File "/workspace/gpt-neox/megatron/training.py", line 708, in train_step
neox-autotune-worker-1:     loss = forward_step(
neox-autotune-worker-1:   File "/workspace/gpt-neox/megatron/training.py", line 375, in forward_step
neox-autotune-worker-1:     loss = cross_entropy(outputs, (labels, loss_mask),
neox-autotune-worker-1:   File "/workspace/gpt-neox/megatron/model/gpt2_model.py", line 71, in cross_entropy
neox-autotune-worker-1:     losses = mpu.vocab_parallel_cross_entropy(output.float().contiguous(), labels)
neox-autotune-worker-1: torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 46.44 GiB (GPU 1; 79.32 GiB total capacity; 52.45 GiB already allocated; 10.63 GiB free; 67.28 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
pdsh@neox-autotune-worker-0: neox-autotune-worker-1: ssh exited with exit code 1

Is the key of the problem 'who exited with exit code 1'?

# z1_gas1_tmbspg121
pdsh@neox-autotune-worker-0: neox-autotune-worker-1: ssh exited with exit code 1

# z3_gas1_tmbspg130
pdsh@neox-autotune-worker-0: neox-autotune-worker-0: ssh exited with exit code 1

Start in scheduler.py around run_experiment, clean_up, and parse_results, comparing the exp_id 13 hang with the exp_id 2 failure path shown in the issue. Trace how the worker exit code and CUDA out-of-memory error are handled. Done means a failed experiment is cleaned up, reported, and does not prevent the scheduler from proceeding.