Repository: oumi-ai/oumi
File: src/oumi/quantize/bnb_quantizer.py (Python)

Code before:
            quantized_size_bytes=quantized_size,
            output_path=output_path,
            format_type=config.output_format,
        )

    def _quantize_model(self, config: QuantizationConfig):
        """Quantize model using BitsAndBytes."""
        from transformers import BitsAndBytesConfig  # type: ignore

        logger.info(
            f"Loading model for BitsAndBytes quantization: {config.model.model_name}"
        )
        logger.info("📥 Loading base model...")

        # Configure quantization based on method
        quantization_config = self._get_quantization_config(config.method)

        logger.info(f"🔧 Using {config.method} quantization")

        # Load and quantize model
        model = AutoModelForCausalLM.from_pretrained(
            config.model.model_name,
            quantization_config=quantization_config,
            device_map="auto",
            torch_dtype=torch.float16,
            trust_remote_code=True,
            **(config.model.model_kwargs or {}),
        )

        tokenizer = AutoTokenizer.from_pretrained(
            config.model.tokenizer_name or config.model.model_name,
            trust_remote_code=True,
        )

        return model, tokenizer

    def _get_quantization_config(self, method: str):
        """Get BitsAndBytes quantization config based on method."""
        from transformers import BitsAndBytesConfig

        if method == "bnb_4bit":
            return BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
            )
        elif method == "bnb_8bit":
            return BitsAndBytesConfig(
                load_in_8bit=True,
                llm_int8_threshold=6.0,

Review comment (suggestion):
nit: you can use the parameters from the config (trust_remote_code, device_map, torch_dtype) if you think the user might need to change them

Code after:
            quantized_size_bytes=quantized_size,
            output_path=output_path,
            format_type=config.output_format,
        )

    def _quantize_model(self, config: QuantizationConfig):
        """Quantize model using BitsAndBytes."""
        logger.info(
            f"Loading model for BitsAndBytes quantization: {config.model.model_name}"
        )
        logger.info("📥 Loading base model...")

        # Configure quantization based on method
        quantization_config = self._get_quantization_config(config.method)

        logger.info(f"🔧 Using {config.method} quantization")

        # Load and quantize model
        torch_dtype = config.model.torch_dtype
        if torch_dtype == torch.float32:
            torch_dtype = torch.float16

        model = AutoModelForCausalLM.from_pretrained(
            config.model.model_name,
            quantization_config=quantization_config,
            device_map=config.model.device_map,
            torch_dtype=torch_dtype,
            trust_remote_code=config.model.trust_remote_code,
            **(config.model.model_kwargs or {}),
        )

        tokenizer = AutoTokenizer.from_pretrained(
            config.model.tokenizer_name or config.model.model_name,
            trust_remote_code=config.model.trust_remote_code,
            **(config.model.tokenizer_kwargs or {}),
        )

        return model, tokenizer

    def _get_quantization_config(self, method: str):
        """Get BitsAndBytes quantization config based on method."""
        from transformers import BitsAndBytesConfig

        if method == "bnb_4bit":
            return BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
            )
        elif method == "bnb_8bit":