load custom model fix

scripts/eval.py

@@ -205,8 +205,15 @@ def load_custom_model(args) -> HFLM:
         from modeling_myolmoe import MyOlmoeForCausalLM, MyOlmoeConfig
         logger.info("Successfully imported MyOlmoeForCausalLM and MyOlmoeConfig")
         
-        # CRITICAL FIX: Register the custom model class with the correct config
+        # CRITICAL FIX: Ensure the config is properly registered as a dataclass
         from transformers import AutoConfig, AutoModelForCausalLM
+        from dataclasses import dataclass
+        
+        # Make sure the config is a proper dataclass
+        if not hasattr(MyOlmoeConfig, '__dataclass_fields__'):
+            logger.warning("MyOlmoeConfig is not a dataclass, this may cause issues")
+        
+        # Register with correct model type
         AutoConfig.register("myolmoe", MyOlmoeConfig)
         AutoModelForCausalLM.register(MyOlmoeConfig, MyOlmoeForCausalLM)
         logger.info("Registered MyOlmoeForCausalLM with MyOlmoeConfig")
@@ -216,52 +223,36 @@ def load_custom_model(args) -> HFLM:
         logger.error("Make sure the custom model code is available in the specified path")
         raise
     
-    # Load model configuration using the correct config class
-    config = MyOlmoeConfig.from_pretrained(
-        args.model_path,
-        trust_remote_code=args.trust_remote_code
-    )
-    
-    logger.info("Model will use default top-k routing configuration")
-    
-    # Create HFLM with explicit model class specification
+    # Load model manually to avoid HFLM wrapper issues
     try:
-        model = HFLM(
-            pretrained=args.model_path,
-            device=args.device,
-            batch_size=args.batch_size,
-            max_batch_size=args.max_batch_size,
-            dtype=args.dtype,
-            trust_remote_code=args.trust_remote_code,
-            # Pass the custom model class explicitly
-            backend="causal",
-            model_kwargs={"torch_dtype": torch.bfloat16 if args.dtype == "bfloat16" else "auto"}
-        )
-    except Exception as e:
-        logger.error(f"Failed to create HFLM wrapper: {e}")
-        # Alternative approach: load model manually then wrap
-        logger.info("Trying alternative loading approach...")
+        logger.info("Loading model manually to avoid wrapper issues...")
         
-        # Load tokenizer and model manually
+        # Load tokenizer first
         tokenizer = AutoTokenizer.from_pretrained(
             args.model_path,
             trust_remote_code=args.trust_remote_code
         )
         
-        # Use the correct config class
+        # Load config with explicit class
         config = MyOlmoeConfig.from_pretrained(
             args.model_path,
             trust_remote_code=args.trust_remote_code
         )
         
+        # Verify config is valid
+        logger.info(f"Loaded config type: {type(config)}")
+        logger.info(f"Config model_type: {getattr(config, 'model_type', 'unknown')}")
+        
+        # Load model instance
         model_instance = MyOlmoeForCausalLM.from_pretrained(
             args.model_path,
             config=config,
             trust_remote_code=args.trust_remote_code,
-            torch_dtype=torch.bfloat16 if args.dtype == "bfloat16" else "auto"
+            torch_dtype=torch.bfloat16 if args.dtype == "bfloat16" else "auto",
+            low_cpu_mem_usage=True
         )
         
-        # Create HFLM with pre-loaded model
+        # Create HFLM wrapper with pre-loaded model
         model = HFLM(
             pretrained=model_instance,
             tokenizer=tokenizer,
@@ -269,9 +260,29 @@ def load_custom_model(args) -> HFLM:
             batch_size=args.batch_size,
             max_batch_size=args.max_batch_size
         )
-    
-    logger.info("Custom model loaded successfully")
-    return model
+        
+        logger.info("Custom model loaded successfully")
+        return model
+        
+    except Exception as e:
+        logger.error(f"Failed to load custom model: {e}")
+        logger.error("Trying fallback approach...")
+        
+        # Fallback: Try loading as standard transformers model
+        try:
+            model = HFLM(
+                pretrained=args.model_path,
+                device=args.device,
+                batch_size=args.batch_size,
+                max_batch_size=args.max_batch_size,
+                dtype=args.dtype,
+                trust_remote_code=True  # Force trust remote code
+            )
+            logger.info("Fallback loading successful")
+            return model
+        except Exception as fallback_e:
+            logger.error(f"Fallback also failed: {fallback_e}")
+            raise e
 
 
 def validate_model_config(model_path: str, trust_remote_code: bool = False) -> Dict[str, Any]: