[https://nvbugs/5549081][fix] Fix device id assignment for some vision models (NVIDIA#8070)

Signed-off-by: Chang Liu (Enterprise Products) <[email protected]>
Signed-off-by: Chang Liu <[email protected]>

Author: chang-l

tensorrt_llm/_torch/models/modeling_hyperclovax.py

@@ -726,7 +726,8 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]):
         self.vision_config = self.pretrained_config.vision_config
 
         model_path = self.pretrained_config._name_or_path
-        self.device = f"cuda:{model_config.mapping.rank}"
+        # TODO: use config.mapping.get_local_rank() instead
+        self.device = f"cuda:{torch.cuda.current_device()}"
 
         hf_model_config = AutoConfig.from_pretrained(model_path,
                                                      trust_remote_code=True)

tensorrt_llm/_torch/models/modeling_llama.py

@@ -999,7 +999,8 @@ def __init__(self, model_config: ModelConfig[Llama4Config], *args,
                  **kwargs):
         super().__init__()
         self.pretrained_config = model_config.pretrained_config
-        self.device = f"cuda:{model_config.mapping.rank}"
+        # TODO: use config.mapping.get_local_rank() instead
+        self.device = f"cuda:{torch.cuda.current_device()}"
 
         self.dtype = self.pretrained_config.text_config.torch_dtype
 

tensorrt_llm/_torch/models/modeling_llava_next.py

@@ -295,7 +295,8 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
         super().__init__()
         self.model_config = model_config
         self.pretrained_config = model_config.pretrained_config
-        self.device = f"cuda:{model_config.mapping.rank}"
+        # TODO: use config.mapping.get_local_rank() instead
+        self.device = f"cuda:{torch.cuda.current_device()}"
         model_path = self.pretrained_config._name_or_path
 
         # Determine the actual local path for model files