[Bugfix] Fix torch.compile() error when using MultiprocessingGPUExecutor (#5229)

This commit is contained in:
zifeitong 2024-06-03 20:55:50 -07:00 committed by GitHub
parent f42a006b15
commit a58f24e590
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -34,6 +34,9 @@ class MultiprocessingGPUExecutor(DistributedGPUExecutor):
# Ensure that VLLM_INSTANCE_ID is set, to be inherited by workers # Ensure that VLLM_INSTANCE_ID is set, to be inherited by workers
os.environ["VLLM_INSTANCE_ID"] = get_vllm_instance_id() os.environ["VLLM_INSTANCE_ID"] = get_vllm_instance_id()
# Disable torch async compiling which won't work with daemonic processes
os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
from torch.cuda import device_count from torch.cuda import device_count
assert world_size <= device_count(), ( assert world_size <= device_count(), (
"please set tensor_parallel_size to less than max local gpu count") "please set tensor_parallel_size to less than max local gpu count")