[Bugfix] Include device index in compile cache paths (#38962)
Signed-off-by: Neil Schemenauer <nas@arctrix.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
This commit is contained in:
co-authored by
Tyler Michael Smith
parent
463aa5e30f
commit
6dcc5d7cae
@@ -1074,7 +1074,10 @@ class VllmBackend:
|
||||
self.compilation_config.cache_dir = cache_dir
|
||||
rank = vllm_config.parallel_config.rank
|
||||
dp_rank = vllm_config.parallel_config.data_parallel_index
|
||||
local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}", self.prefix)
|
||||
dev = torch.accelerator.current_device_index()
|
||||
local_cache_dir = os.path.join(
|
||||
cache_dir, f"rank_{rank}_{dp_rank}_dev{dev}", self.prefix
|
||||
)
|
||||
os.makedirs(local_cache_dir, exist_ok=True)
|
||||
self.compilation_config.local_cache_dir = local_cache_dir
|
||||
|
||||
|
||||
@@ -559,7 +559,8 @@ def _support_torch_compile(
|
||||
|
||||
rank = self.vllm_config.parallel_config.rank
|
||||
dp_rank = self.vllm_config.parallel_config.data_parallel_index
|
||||
cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
|
||||
dev = torch.accelerator.current_device_index()
|
||||
cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}_dev{dev}")
|
||||
aot_compilation_path = os.path.join(cache_dir, "model")
|
||||
if not envs.VLLM_DISABLE_COMPILE_CACHE:
|
||||
loaded_fn = _try_load_aot_compiled_fn(self, aot_compilation_path)
|
||||
|
||||
Reference in New Issue
Block a user