Repository: huggingface/optimum-benchmark
Issue category: Performance Issue

Timeout with multiple AMD GPUs tensor parallelism in vLLM
## Problem Description

When attempting to run Optimum Benchmark in vLLM using tensor parallelism across multiple AMD GPUs (MI210), I encounter a timeout error from NCCL watchdog. However, the benchmark works fine with a single AMD GPU in vLLM, and the vLLM API server functions correctly when using tensor parallelism.
During the benchmark the VRAM% and GPU usage % from rocm-smi were 90-100%.

## Environment

- GPUs: AMD MI210 (x4)
- ROCm version: 6.1.2
- Framework: vLLM (tried 0.4.3 and 0.5.3)
- Optimum Benchmark version: 0.4.0

## Tested Configurations

1. Multi-GPU setup (failing)
2. Single-GPU setup (working)
3. API Server Inference with Tensor Parallelism [(from ROCm Documentation)](https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/llm-inference-frameworks.html#vllm-inference) (working)

## Docker Images Tested

1. Official vLLM docker for ROCm: [Dockerfile.rocm](https://github.com/vllm-project/vllm/blob/main/Dockerfile.rocm)
2. ROCm's fork of vLLM: [Dockerfile.rocm](https://github.com/ROCm/vllm/blob/main/Dockerfile.rocm)

Both Docker images resulted in the same issue with Optimum Benchmark on multiple GPUs.

## Optimum Benchmark configuration

```
defaults:
  - benchmark
  - scenario: inference
  - launcher: process
  - backend: vllm
  - _self_

name: vllm_llama

launcher:
  device_isolation: true
  device_isolation_action: warn

backend:
  device: cuda
  device_ids: 0,1,2,3
  no_weights: false
  serving_mode: offline
  model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
  engine_args:
    dtype: float16
    tensor_parallel_size: 4
    enforce_eager: true
    distributed_executor_backend: mp
    disable_custom_all_reduce: true

scenario:
  latency: true
  memory: false
  energy: false
  input_shapes:
    batch_size: 1
    sequence_length: 128
  generate_kwargs:
    max_new_tokens: 100
    min_new_tokens: 100
```

## Logs

- Multi-GPU issue log
[tp4_failure.log](https://github.com/user-attachments/files/16797139/tp4_failure.log)

Preview of the error:
```
[ISOLATED-PROCESS][2024-08-29 06:56:29,604][inference][INFO] -  + Running Text Generation latency tracking
[ISOLATED-PROCESS][2024-08-29 06:56:29,605][latency][INFO] -    + Tracking latency using CPU performance counter
[rank2]:[E ProcessGroupNCCL.cpp:574] [Rank 2] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1561, OpType=BROADCAST, NumelIn=1, NumelOut=1, Timeout(ms)=600000) ran for 600043 milliseconds before timing out.
[rank2]:[E ProcessGroupNCCL.cpp:588] [Rank 2] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank2]:[E ProcessGroupNCCL.cpp:594] [Rank 2] To avoid data inconsistency, we are taking the entire process down.
[rank2]:[E ProcessGroupNCCL.cpp:1358] [PG 0 Rank 2] NCCL watchdog thread terminated with exception: [Rank 2] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1561, OpType=BROADCAST, NumelIn=1, NumelOut=1, Timeout(ms)=600000) ran for 600043 milliseconds before timing out.
```

- Single-GPU successful run log
[single_gpu_success.log](https://github.com/user-attachments/files/16797118/single_gpu_success.log)

- API Server successful run log
[tp4_http_success.log](https://github.com/user-attachments/files/16797121/tp4_http_success.log)



Functions to edit:
optimum_benchmark/trackers/latency.py:LatencyTracker.__init__

Patch:
diff --git a/optimum_benchmark/trackers/latency.py b/optimum_benchmark/trackers/latency.py
index cb236413..1e0f1e95 100644
--- a/optimum_benchmark/trackers/latency.py
+++ b/optimum_benchmark/trackers/latency.py
@@ -121,7 +121,8 @@ def __init__(self, device: str, backend: str):
         self.device = device
         self.backend = backend
         self.is_asynchronous = self.backend == "pytorch" and self.device == "cuda"
-        self.is_distributed = is_torch_distributed_available() and torch.distributed.is_initialized()
+        self.is_distributed = (self.backend != "vllm" and
+                               is_torch_distributed_available() and torch.distributed.is_initialized())
 
         if self.is_asynchronous:
             LOGGER.info("\t+ Tracking latency using Pytorch CUDA events")