Skip to content

8 bit quantization performance issues, 4b-it #5

Description

@dscarmo

Hello, thanks for the amazing work on medgemma.

I have been working with other multimodal models such as CheXagent. It performs pretty well under 8 bit and 4 bit quantization using bitsandbytes.

When attempting the same with your reference code (which works well for bfloat16), the model outputs gibberish. Maybe I am missing something on initialization or on the forward? Here is the code:

Here are the relevant methods, loading:


def load_model(self):
        """Load the MedGemma model and processor."""
        print("Loading MedGemma model... This may take a few minutes.")
        
        # Get HuggingFace token from environment
        hf_token = os.getenv('HF_TOKEN')
        if hf_token:
            print("✅ Using HuggingFace token for authentication")
        else:
            print("⚠️  No HF_TOKEN found in environment variables")
            print("   If the model requires authentication, please set HF_TOKEN in your .env file")
        
        try:
            # Configure model loading parameters based on quantization
            if self.use_8bit:
                print("🔄 Loading model in 8-bit mode for reduced memory usage...")
                try:
                    import bitsandbytes as bnb
                    print("✅ Using bitsandbytes for 8-bit quantization")
                except ImportError:
                    raise ImportError(
                        "bitsandbytes is required for 8-bit quantization. "
                        "Please install it using: pip install bitsandbytes>=0.41.0"
                    )
                
                load_kwargs = {
                    "load_in_8bit": True,
                    "device_map": "auto",
                    "token": hf_token,
                }
            else:
                print("🔄 Loading model in full precision mode...")
                load_kwargs = {
                    "torch_dtype": torch.bfloat16,
                    "device_map": "auto",
                    "token": hf_token,
                }
            
            self.model = AutoModelForImageTextToText.from_pretrained(
                self.model_id,
                **load_kwargs
            )
            self.processor = AutoProcessor.from_pretrained(
                self.model_id,
                token=hf_token,
            )
            
            print("✅ Model loaded successfully!")
            if self.use_8bit:
                print("ℹ️  Model is running in 8-bit mode for reduced memory usage")
            
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            if "authentication" in str(e).lower() or "token" in str(e).lower():
                print("💡 This might be an authentication issue. Please:")
                print("   1. Create a .env file in the medgemma folder")
                print("   2. Add your HuggingFace token: HF_TOKEN=your_token_here")
                print("   3. Get a token from: https://huggingface.co/settings/tokens")
            raise

Forward:


def analyze_image(self, 
                     image: Image.Image | List[Image.Image], 
                     user_prompt: str = "Describe this X-ray",
                     system_prompt: str = "You are an expert radiologist.",
                     max_new_tokens: int = 1024) -> str:
        """
        Analyze one or more medical images and generate a report.
        
        Args:
            image (PIL.Image.Image | List[PIL.Image.Image]): The medical image(s) to analyze
            user_prompt (str): The user's question or instruction
            system_prompt (str): The system prompt defining the AI's role
            max_new_tokens (int): Maximum number of tokens to generate
            
        Returns:
            str: The generated medical analysis
        """
        if self.model is None or self.processor is None:
            raise RuntimeError("Model not loaded. Call load_model() first.")
        
        # Convert single image to list for consistent handling
        images = [image] if isinstance(image, Image.Image) else image
        
        # Add information about number of images to the prompt
        if len(images) > 1:
            user_prompt = f"{user_prompt} (Analyzing {len(images)} images together)"
        
        # Prepare the message content with all images
        content = [{"type": "text", "text": user_prompt}]
        for img in images:
            content.append({"type": "image", "image": img})
        
        messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": system_prompt}]
            },
            {
                "role": "user",
                "content": content
            }
        ]
        
        inputs = self.processor.apply_chat_template(
            messages, 
            add_generation_prompt=True, 
            tokenize=True,
            return_dict=True, 
            return_tensors="pt"
        ).to(self.model.device, dtype=torch.bfloat16)
        
        input_len = inputs["input_ids"].shape[-1]
        
        with torch.inference_mode():
            generation = self.model.generate(
                **inputs, 
                max_new_tokens=max_new_tokens, 
                do_sample=False,
                disable_compile=True  # Compilation disabled due to compatibility issues with current Transformers version
            )
            generation = generation[0][input_len:]
        
        decoded = self.processor.decode(generation, skip_special_tokens=True)
        return decoded

As I mentioned, this code works fine when not using quantization. My use case is deploying these models on 8 GB of GPU budget, for low resource hospitals (low income country).

Metadata

Metadata

Assignees

No one assigned

    Labels

    type:performancePerformance, efficiency, latency, or memory usage

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions