Method load_base_model

optillm/inference.py:1024–1111

(self, model_id: str, quantize: bool = True)

        return model

    def load_base_model(self, model_id: str, quantize: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
        def _load_model():
            logger.info(f"Loading base model: {model_id}")
            
            device = self.device_manager.get_optimal_device()
            logger.info(f"Using device: {device}")
            
            # Load tokenizer
            tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, token=os.getenv("HF_TOKEN"))
            
            # Base kwargs for model loading
            model_kwargs = {
                "trust_remote_code": True,
                "device_map": "auto" if 'cuda' in device else device
            }
            
            # Configure device-specific optimizations
            if 'cuda' in device:
                compute_capability = torch.cuda.get_device_capability(0)
                if compute_capability[0] >= 8:
                    model_kwargs["torch_dtype"] = torch.bfloat16
                elif compute_capability[0] >= 7:
                    model_kwargs["torch_dtype"] = torch.float16
                    
                # Check for flash attention availability
                try:
                    import flash_attn
                    has_flash_attn = True
                    logger.info("Flash Attention 2 is available")
                    model_kwargs["attn_implementation"] = "flash_attention_2"
                except ImportError:
                    has_flash_attn = False
                    logger.info("Flash Attention 2 is not installed - falling back to default attention")
                    
            elif 'mps' in device:
                # Special handling for Gemma models which have NaN issues with float16 on MPS
                if 'gemma' in model_id.lower():
                    model_kwargs["torch_dtype"] = torch.float32
                    logger.info("Using MPS device with float32 for Gemma model (float16 causes NaN)")
                else:
                    model_kwargs["torch_dtype"] = torch.float16
                    logger.info("Using MPS device with float16 precision")
            else:
                # CPU can use FP16 if available
                if hasattr(torch.cpu, 'has_fp16') and torch.cpu.has_fp16:
                    model_kwargs["torch_dtype"] = torch.float16
                    logger.info("Using CPU device with float16 precision")
                else:
                    model_kwargs["torch_dtype"] = torch.float32
                    logger.info("Using CPU device with float32 precision - FP16 not supported")
            
            # Load model with configured optimizations
            try:
                model = AutoModelForCausalLM.from_pretrained(
                    model_id,
                    token=os.getenv("HF_TOKEN"),
                    **model_kwargs
                )