Fix: Remove device_map to prevent meta tensor initialization - ACE-Step custom model needs direct device placement

acestep/handler.py

@@ -495,7 +495,6 @@ class AceStepHandler:
                             torch_dtype=self.dtype,
                             low_cpu_mem_usage=False,  # Disable memory-efficient weight loading
                             _fast_init=False,  # Disable fast initialization (prevents meta device)
-                            device_map={"": device},  # Explicitly map all components to target device
                         )
                         attn_implementation = candidate
                         break
@@ -569,9 +568,10 @@ class AceStepHandler:
                 # Determine target device for VAE
                 vae_device = device if not self.offload_to_cpu else "cpu"
                 self.vae = AutoencoderOobleck.from_pretrained(
-                    vae_checkpoint_path,
-                    device_map={"":  vae_device}  # Explicitly map to target device
+                    vae_checkpoint_path
                 )
+                # Move VAE to target device
+                self.vae = self.vae.to(vae_device)
                 if not self.offload_to_cpu:
                     # Keep VAE in GPU precision when resident on accelerator.
                     vae_dtype = self._get_vae_dtype(device)
@@ -602,9 +602,10 @@ class AceStepHandler:
                 # Determine target device for text encoder
                 text_encoder_device = device if not self.offload_to_cpu else "cpu"
                 self.text_encoder = AutoModel.from_pretrained(
-                    text_encoder_path,
-                    device_map={"":  text_encoder_device}  # Explicitly map to target device
+                    text_encoder_path
                 )
+                # Move text encoder to target device
+                self.text_encoder = self.text_encoder.to(text_encoder_device)
                 if not self.offload_to_cpu:
                     self.text_encoder = self.text_encoder.to(self.dtype)
                 else:

acestep/llm_inference.py

@@ -278,9 +278,10 @@ class LLMHandler:
             target_device = device if not self.offload_to_cpu else "cpu"
             self.llm = AutoModelForCausalLM.from_pretrained(
                 model_path,
-                trust_remote_code=True,
-                device_map={"":  target_device}  # Explicitly map to target device
+                trust_remote_code=True
             )
+            # Move model to target device
+            self.llm = self.llm.to(target_device)
             if not self.offload_to_cpu:
                 self.llm = self.llm.to(self.dtype)
             else:
@@ -3024,13 +3025,13 @@ class LLMHandler:
                 self._hf_model_for_scoring = AutoModelForCausalLM.from_pretrained(
                     model_path,
                     trust_remote_code=True,
-                    torch_dtype=self.dtype,
-                    device_map={"":  str(device)}  # Explicitly map to vLLM device
+                    torch_dtype=self.dtype
                 )
+                # Move model to vLLM device
+                self._hf_model_for_scoring = self._hf_model_for_scoring.to(device)
                 load_time = time.time() - start_time
                 logger.info(f"HuggingFace model loaded in {load_time:.2f}s")
                 
-                # Already on device from device_map
                 self._hf_model_for_scoring.eval()
                 
                 logger.info(f"HuggingFace model for scoring ready on {device}")
@@ -3054,13 +3055,13 @@ class LLMHandler:
                 self._hf_model_for_scoring = AutoModelForCausalLM.from_pretrained(
                     model_path,
                     trust_remote_code=True,
-                    torch_dtype=self.dtype,
-                    device_map={"":  device}  # Explicitly map to target device
+                    torch_dtype=self.dtype
                 )
+                # Move model to target device
+                self._hf_model_for_scoring = self._hf_model_for_scoring.to(device)
                 load_time = time.time() - start_time
                 logger.info(f"HuggingFace model loaded in {load_time:.2f}s")
                 
-                # Already on device from device_map
                 self._hf_model_for_scoring.eval()
                 
                 logger.info(f"HuggingFace model for scoring ready on {device}")