A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 03:30 UTC
[V1] Only print cudagraph tqdm on rank 0 with `is_global_first_rank` (#19516)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 1, 2025 at 15:02 UTC
be250bbc67973766e546e0e3d8abb21e5caa2b1f
2 files changed
+38
-4
vllm/distributed/parallel_state.py
+31
@@ -1315,6 +1315,37 @@ def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
1315
return [x == 1 for x in aggregated_data.tolist()]
1316
1317
1318
+def is_global_first_rank() -> bool:
1319
+ """
1320
+ Check if the current process is the first rank globally across all
1321
+ parallelism strategies (PP, TP, DP, EP, etc.).
1322
+
1323
+ Unlike group-specific checks like `get_tensor_model_parallel_rank() == 0`
1324
+ or `get_pp_group().is_first_rank`, this function checks the global rank
1325
+ across all parallelism dimensions.
1326
+
1327
+ Returns:
1328
+ bool: True if this is the global first rank (rank 0), False otherwise.
1329
+ Returns True if distributed is not initialized (single process).
1330
+ """
1331
+ try:
1332
+ # If world group is available, use it for the most accurate check
1333
+ global _WORLD
1334
+ if _WORLD is not None:
1335
+ return _WORLD.is_first_rank
1336
+
1337
+ # If torch distributed is not initialized, assume single process
1338
+ if not torch.distributed.is_initialized():
1339
+ return True
1340
+
1341
+ # Fallback to torch's global rank
1342
+ return torch.distributed.get_rank() == 0
1343
+
1344
+ except Exception:
1345
+ # If anything goes wrong, assume this is the first rank
1346
+ return True
1347
+
1348
+
1349
def _node_count(pg: Union[ProcessGroup, StatelessProcessGroup]) -> int:
1350
"""
1351
Returns the total number of nodes in the process group.
vllm/v1/worker/gpu_model_runner.py
+7
-4
@@ -26,7 +26,7 @@ from vllm.distributed.kv_transfer import (get_kv_transfer_group,
26
has_kv_transfer_group)
27
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
28
from vllm.distributed.parallel_state import (
29
- get_pp_group, get_tp_group, graph_capture,
29
+ get_pp_group, get_tp_group, graph_capture, is_global_first_rank,
30
prepare_communication_buffer_for_model)
31
from vllm.forward_context import (DPMetadata, get_forward_context,
32
set_forward_context)
@@ -2285,9 +2285,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2285
# can reuse the memory pool allocated for the large shapes.
2286
with graph_capture(device=self.device):
2287
full_cg = self.full_cuda_graph
2288
- for num_tokens in tqdm(reversed(self.cudagraph_batch_sizes),
2289
- desc="Capturing CUDA graphs",
2290
- total=len(self.cudagraph_batch_sizes)):
2288
+ # Only rank 0 should print progress bar during capture
2289
+ compilation_cases = reversed(self.cudagraph_batch_sizes)
2290
+ if is_global_first_rank():
2291
+ compilation_cases = tqdm(list(compilation_cases),
2292
+ desc="Capturing CUDA graph shapes")
2293
+ for num_tokens in compilation_cases:
2294
# We skip EPLB here since we don't want to record dummy metrics
2295
for _ in range(
2296
self.compilation_config.cudagraph_num_of_warmups):