[BUG] DeepSpeed Zero Inference (stage 3) Stuck When One Process Doesn't Execute `model.generate()`

Describe the bug
I attempted to perform inference on the LLaMA2 70B model using DeepSpeed with Zero optimization (stage 3) across multiple GPUs (NVIDIA V100). While the initial setup appeared successful, I encountered a critical issue: if any process within the distributed environment does not execute the model.generate() function, it leads to a deadlock situation, causing processes on other GPUs to remain stuck indefinitely.

To Reproduce
Steps to reproduce the behavior:

1. Modify the script in DeepSpeedExamples to support running DeepSpeed with Zero optimization on multiple GPUs.
2. Add an infinite loop to a specific rank before the model.generate() function
3. (This modification also generates the bug): Have one process finish generating prompts while others are still generating, and the finished one does not continue generating the next prompt. In this case, even with a barrier to prevent the finished process from deleting its model, the bug still occurs.
4. Use DeepSpeed v0.12.6 and Torch v2.0.1+cu117
5. Execute the script with deepspeed --num_gpus 8 ...

Expected behavior
The inference should not be stuck.

I understand that Zero optimization stage 3 allows multiple GPUs to share the model weights by utilizing faster P2P communication such as NVLink to avoid transferring repeated data from CPU to GPU via PCIe.

Is there any solution or workaround that allows all GPUs to continue working, even if not all processes are executing the function model.generate()? This would help maintain parallelism and prevent the deadlock situation currently observed.

ds_report output

{
    "fp16": {
        "enabled": true
    }, 
    "bf16": {
        "enabled": false
    }, 
    "zero_optimization": {
        "stage": 3, 
        "stage3_prefetch_bucket_size": 2.684355e+08, 
        "stage3_param_persistence_threshold": 8.192000e+03, 
        "stage3_max_live_parameters": 2.684355e+08, 
        "offload_param": {
            "device": "cpu", 
            "pin_memory": true
        }
    }, 
    "steps_per_print": 2.000000e+03, 
    "train_batch_size": 16, 
    "wall_clock_breakdown": false
}

Screenshots
In the provided screenshot, an infinite loop for rank 0 (utilizing GPU0) avoids executing the model.generate() function, demonstrating that other processes are also stuck, with no fluctuation in power usage.

System info:

• OS: Linux 3.10.0-1127.el7.x86_64
• GPU count and types: NVIDIA V100 (32GB) * 8
• Hugging Face Transformers: version 4.33.0.dev0
• Python 3.10.8

Start with DeepSpeedExamples/inference/huggingface/zero_inference/run_model.py and reproduce the multi-GPU case using DeepSpeed v0.12.6, Torch 2.0.1+cu117, and deepspeed --num_gpus 8. Compare execution when one rank does not call model.generate() or finishes prompts early. Done means other ranks do not remain stuck and inference can continue.