Function get_kbit_device_map

src/alignment/model_utils.py:38–40

Useful for running inference with quantized models by setting `device_map=get_peft_device_map()`

()

def get_kbit_device_map() -> Dict[str, int] | None:
    """Useful for running inference with quantized models by setting `device_map=get_peft_device_map()`"""
    return {"": get_current_device()} if torch.cuda.is_available() else None


def get_quantization_config(model_args: ModelArguments) -> BitsAndBytesConfig | None: