fix device_map (#505)

Author: coding-famer

lmms_eval/models/internvl2.py

@@ -204,8 +204,8 @@ def __init__(
             self._device = torch.device(f"cuda:{accelerator.local_process_index}")
             self.device_map = f"cuda:{accelerator.local_process_index}"
 
-        self._model = AutoModel.from_pretrained(self.path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=True, device_map=device_map).eval()
-        self._tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=True, device_map=device_map)
+        self._model = AutoModel.from_pretrained(self.path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=True, device_map=self.device_map).eval()
+        self._tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=True, device_map=self.device_map)
 
         if accelerator.num_processes > 1:
             assert accelerator.distributed_type in [DistributedType.FSDP, DistributedType.MULTI_GPU, DistributedType.DEEPSPEED], "Unsupported distributed type provided. Only DDP and FSDP are supported."

lmms_eval/models/qwen2_audio.py

@@ -58,7 +58,7 @@ def __init__(
         self._model = Qwen2AudioForConditionalGeneration.from_pretrained(
             pretrained,
             torch_dtype="auto",
-            device_map=device_map,
+            device_map=self.device_map,
         ).eval()
 
         self.processor = AutoProcessor.from_pretrained(pretrained)