[debug] zeroGPU

rosetta/baseline/multi_stage.py

@@ -59,24 +59,20 @@ class TwoStageInference:
         if context_path == "google/gemma-3-1b-it":
             torch._dynamo.config.cache_size_limit = 64
             self.context_model = AutoModelForCausalLM.from_pretrained(
-                context_path, torch_dtype=torch.bfloat16, sliding_window=4096,
-                # device_map={"": self.device}, 
-            ).to(self.device)
+                context_path, torch_dtype=torch.bfloat16, device_map={"": self.device}, sliding_window=4096
+            )
         else:
             self.context_model = AutoModelForCausalLM.from_pretrained(
-                context_path, 
-                torch_dtype=torch.bfloat16, 
-                # device_map={"": self.device}
-            ).to(self.device)
+                context_path, torch_dtype=torch.bfloat16, device_map={"": self.device}
+            )
         # Apply generation config to context model
         apply_generation_config(self.context_model, self.generation_config)
         
         # Load answer LLM
         self.answer_tokenizer = AutoTokenizer.from_pretrained(answer_path)
         self.answer_model = AutoModelForCausalLM.from_pretrained(
-            answer_path, torch_dtype=torch.bfloat16, 
-            # device_map={"": self.device}
-        ).to(self.device)
+            answer_path, torch_dtype=torch.bfloat16, device_map={"": self.device}
+        )
         # Apply generation config to answer model
         apply_generation_config(self.answer_model, self.generation_config)
     

rosetta/utils/evaluate.py

@@ -313,8 +313,8 @@ def load_hf_model(model_name: str, device: torch.device, generation_config: Opti
         model = AutoModelForCausalLM.from_pretrained(
             str(model_name),
             torch_dtype=torch.bfloat16,
-            # device_map={"": device}
-    ).eval().to(device)
+            device_map={"": device}
+    ).eval()
     
     # Apply generation config
     apply_generation_config(model, generation_config)
@@ -352,8 +352,8 @@ def load_rosetta_model(model_config: Dict[str, Any], eval_config: Dict[str, Any]
     slm_model = AutoModelForCausalLM.from_pretrained(
         str(slm_model_path),
         torch_dtype=torch.bfloat16,
-        # device_map={"": device}
-    ).eval().to(device)
+        device_map={"": device}
+    ).eval()
     
     # Apply generation config to SLM
     apply_generation_config(slm_model, generation_config)
@@ -362,15 +362,15 @@ def load_rosetta_model(model_config: Dict[str, Any], eval_config: Dict[str, Any]
         llm_model = AutoModelForCausalLM.from_pretrained(
             str(llm_model_path),
                 torch_dtype=torch.bfloat16,
-                # device_map={"": device},
+                device_map={"": device},
                 sliding_window=4096
-            ).eval().to(device)
+            ).eval()
     else:
         llm_model = AutoModelForCausalLM.from_pretrained(
             str(llm_model_path),
             torch_dtype=torch.bfloat16,
-            # device_map={"": device}
-        ).eval().to(device)
+            device_map={"": device}
+        ).eval()
     
     # Apply generation config to LLM
     apply_generation_config(llm_model, generation_config)