load custom model fix

scripts/eval.py

@@ -205,15 +205,12 @@ def load_custom_model(args) -> HFLM:
         from modeling_myolmoe import MyOlmoeForCausalLM, MyOlmoeConfig
         logger.info("Successfully imported MyOlmoeForCausalLM and MyOlmoeConfig")
         
-        # 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
+        # Check if config is a dataclass
         if not hasattr(MyOlmoeConfig, '__dataclass_fields__'):
             logger.warning("MyOlmoeConfig is not a dataclass, this may cause issues")
         
-        # Register with correct model type
+        # Register the custom model class with the correct config
+        from transformers import AutoConfig, AutoModelForCausalLM
         AutoConfig.register("myolmoe", MyOlmoeConfig)
         AutoModelForCausalLM.register(MyOlmoeConfig, MyOlmoeForCausalLM)
         logger.info("Registered MyOlmoeForCausalLM with MyOlmoeConfig")
@@ -223,36 +220,35 @@ def load_custom_model(args) -> HFLM:
         logger.error("Make sure the custom model code is available in the specified path")
         raise
     
-    # Load model manually to avoid HFLM wrapper issues
+    # Load model manually to avoid wrapper issues
+    logger.info("Loading model manually to avoid wrapper issues...")
+    
     try:
-        logger.info("Loading model manually to avoid wrapper issues...")
-        
-        # Load tokenizer first
+        # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(
             args.model_path,
             trust_remote_code=args.trust_remote_code
         )
         
-        # Load config with explicit class
-        config = MyOlmoeConfig.from_pretrained(
+        # Load config using the custom config class
+        model_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')}")
+        # Debug information
+        logger.info(f"Loaded config type: {type(model_config)}")
+        logger.info(f"Config model_type: {model_config.model_type}")
         
         # Load model instance
         model_instance = MyOlmoeForCausalLM.from_pretrained(
             args.model_path,
-            config=config,
+            config=model_config,
             trust_remote_code=args.trust_remote_code,
-            torch_dtype=torch.bfloat16 if args.dtype == "bfloat16" else "auto",
-            low_cpu_mem_usage=True
+            torch_dtype=torch.bfloat16 if args.dtype == "bfloat16" else "auto"
         )
         
-        # Create HFLM wrapper with pre-loaded model
+        # Create HFLM with pre-loaded model
         model = HFLM(
             pretrained=model_instance,
             tokenizer=tokenizer,
@@ -261,29 +257,12 @@ def load_custom_model(args) -> HFLM:
             max_batch_size=args.max_batch_size
         )
         
-        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
-
+        raise
+    
+    logger.info("Custom model loaded successfully")
+    return model
 
 def validate_model_config(model_path: str, trust_remote_code: bool = False) -> Dict[str, Any]:
     """