Repository: oumi-ai/oumi
File: src/oumi/inference/vllm_inference_engine.py (Python)

Code before:
        if not (
            math.isfinite(gpu_memory_utilization)
            and gpu_memory_utilization > 0
            and gpu_memory_utilization <= 1.0
        ):
            raise ValueError(
                "GPU memory utilization must be within (0, 1]. Got "
                f"{gpu_memory_utilization}."
            )

        # Check if BitsAndBytes quantization was requested.
        bitsandbytes_quantization = False
        if model_params.model_kwargs:
            bitsandbytes_quantization_kwargs = ["load_in_4bit", "load_in_8bit"]
            for key in bitsandbytes_quantization_kwargs:
                if model_params.model_kwargs.get(key):
                    bitsandbytes_quantization = True

        vllm_kwargs = {}

        # Ensure BitsAndBytes keys for vllm.LLM are properly set.
        if bitsandbytes_quantization:
            if quantization and quantization != "bitsandbytes":
                raise RuntimeError(
                    "The model kwargs include `bitsandbytes` quantization keys "
                    f"(such as {', '.join(bitsandbytes_quantization_kwargs)}), but a "
                    "different quantization method was requested: `{quantization}`."
                )
            else:
                quantization = "bitsandbytes"
                vllm_kwargs["load_format"] = "bitsandbytes"
                logger.info("VLLM engine loading a `bitsandbytes` quantized model.")

        if tensor_parallel_size <= 0:
            if torch.cuda.device_count() > 1:
                tensor_parallel_size = torch.cuda.device_count()
            else:
                tensor_parallel_size = 1

        self._lora_request = None
        if model_params.adapter_model:
            # ID should be unique for this adapter, but isn't enforced by vLLM.
            self._lora_request = LoRARequest(
                lora_name="oumi_lora_adapter",
                lora_int_id=1,
                lora_path=model_params.adapter_model,
            )
            logger.info(f"Loaded LoRA adapter: {model_params.adapter_model}")
            lora_rank = get_lora_rank(model_params.adapter_model)
            vllm_kwargs["max_lora_rank"] = lora_rank
            logger.info(f"Setting vLLM max LoRA rank to {lora_rank}")

Review comment (suggestion):
Does load_in_4bit and load_in_8bit always mean `bitsandbytes` ? Surely other quantization methods use that ? Maybe we can update the message to say that only bitsandbytes is supported to load in 4/8 bit at this toime

Code after:
        if not (
            math.isfinite(gpu_memory_utilization)
            and gpu_memory_utilization > 0
            and gpu_memory_utilization <= 1.0
        ):
            raise ValueError(
                "GPU memory utilization must be within (0, 1]. Got "
                f"{gpu_memory_utilization}."
            )

        # Check if any quantization keys are set.
        quantization_keys_set = False
        if model_params.model_kwargs:
            quantization_kwargs = ["load_in_4bit", "load_in_8bit"]
            for key in quantization_kwargs:
                if model_params.model_kwargs.get(key):
                    quantization_keys_set = True

        vllm_kwargs = {}

        # If quantization requested but undefined, default to BitsAndBytes.
        if quantization_keys_set:
            if not quantization or quantization == "bitsandbytes":
                quantization = "bitsandbytes"
                vllm_kwargs["load_format"] = "bitsandbytes"
                logger.info("VLLM engine loading a `bitsandbytes` quantized model.")

        if tensor_parallel_size <= 0:
            if torch.cuda.device_count() > 1:
                tensor_parallel_size = torch.cuda.device_count()
            else:
                tensor_parallel_size = 1

        self._lora_request = None
        if model_params.adapter_model:
            # ID should be unique for this adapter, but isn't enforced by vLLM.
            self._lora_request = LoRARequest(
                lora_name="oumi_lora_adapter",
                lora_int_id=1,
                lora_path=model_params.adapter_model,
            )
            logger.info(f"Loaded LoRA adapter: {model_params.adapter_model}")
            lora_rank = get_lora_rank(model_params.adapter_model)
            vllm_kwargs["max_lora_rank"] = lora_rank
            logger.info(f"Setting vLLM max LoRA rank to {lora_rank}")

        if max_num_seqs is not None:
            vllm_kwargs["max_num_seqs"] = max_num_seqs

        self._tokenizer = build_tokenizer(model_params)
        self._llm = vllm.LLM(