[Minor] Fix false warning when TP=1 (#2674)

This commit is contained in:
Woosuk Kwon
2024-01-30 14:39:40 -08:00
committed by GitHub
parent bbe9bd9684
commit 105a40f53a

View File

@ -29,6 +29,10 @@ def init_custom_ar() -> None:
return
rank = get_tensor_model_parallel_rank()
world_size = get_tensor_model_parallel_world_size()
if world_size == 1:
# No need to initialize custom allreduce for single GPU case.
return
if world_size not in _SUPPORTED_WORLD_SIZES:
logger.warn(
"Custom allreduce is disabled due to an unsupported world size: "