import torch
import bitsandbytes as bnb
from transformers.cache_utils import DynamicCache

# Ensure that DynamicCache and set are safe during serialization
torch.serialization.add_safe_globals([DynamicCache, set])

def quantize_and_save(past_key_values, save_path="cache_nf4.pth"):
    quantized_past_key_values = []
    for k, v in past_key_values:
        # Convert tensors to float16 for quantization
        k, v = k.to(torch.float16), v.to(torch.float16)
        
        # Perform 4-bit quantization using NF4
        k_quant, k_absmax = bnb.functional.quantize_4bit(k, quant_type="nf4")
        v_quant, v_absmax = bnb.functional.quantize_4bit(v, quant_type="nf4")
        
        # Extract float values from the scaling factors
        k_absmax_value = k_absmax.item()  # Directly extract the float value
        v_absmax_value = v_absmax.item()  # Directly extract the float value
        
        # Append quantized tensors and their corresponding absmax values
        quantized_past_key_values.append(((k_quant, k_absmax_value), (v_quant, v_absmax_value)))

    # Save the quantized data to disk
    torch.save(quantized_past_key_values, save_path)
    print("Cache saved successfully with NF4 quantization and absmax.")

def load_and_dequantize(save_path="cache_nf4.pth"):
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        loaded_past_key_values = torch.load(save_path, map_location=device)
        print("NF4-quantized cache loaded successfully.")
    except FileNotFoundError:
        print("No NF4 quantized cache file found.")
        loaded_past_key_values = []

    past_key_values = DynamicCache()

    if loaded_past_key_values:
        for layer_idx, ((k_quant, k_absmax), (v_quant, v_absmax)) in enumerate(loaded_past_key_values):
            # Convert absmax values back to tensors
            k_absmax_tensor = torch.tensor(k_absmax, dtype=torch.float32, device=device)
            v_absmax_tensor = torch.tensor(v_absmax, dtype=torch.float32, device=device)

            # Define output shapes based on quantized tensor dimensions
            key_out_shape = (k_quant.shape[0] * 8, k_quant.shape[1])
            value_out_shape = (v_quant.shape[0] * 8, v_quant.shape[1])

            # Dequantize the key and value tensors
            key_states = bnb.functional.dequantize_4bit(
                k_quant,
                absmax=k_absmax_tensor,
                out=torch.empty(key_out_shape, dtype=torch.float16, device=device),
                quant_type="nf4"
            )
            value_states = bnb.functional.dequantize_4bit(
                v_quant,
                absmax=v_absmax_tensor,
                out=torch.empty(value_out_shape, dtype=torch.float16, device=device),
                quant_type="nf4"
            )

            # Update the DynamicCache with the dequantized tensors
            past_key_values.update(key_states, value_states, layer_idx)

    print("DynamicCache populated with dequantized float16 past_key_values.")
    return past_key_values

# Assuming `past_key_values` is already defined and populated somewhere in your code
quantize_and_save(past_key_values)

# Load the quantized cache
past_key_values = load_and_dequantize()






---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[13], line 71
     68     return past_key_values
     70 # Assuming `past_key_values` is already defined and populated somewhere in your code
---> 71 quantize_and_save(past_key_values)
     73 # Load the quantized cache
     74 past_key_values = load_and_dequantize()

Cell In[13], line 15, in quantize_and_save(past_key_values, save_path)
     12 k, v = k.to(torch.float16), v.to(torch.float16)
     14 # Perform 4-bit quantization using NF4
---> 15 k_quant, k_absmax = bnb.functional.quantize_4bit(k, quant_type="nf4")
     16 v_quant, v_absmax = bnb.functional.quantize_4bit(v, quant_type="nf4")
     18 # Extract float values from the scaling factors

File /opt/conda/lib/python3.10/site-packages/bitsandbytes/functional.py:1182, in quantize_4bit(A, absmax, out, blocksize, compress_statistics, quant_type, quant_storage)
   1156 """
   1157 Quantize tensor A in blocks of 4-bit values.
   1158 
   (...)
   1179     The quantization state to undo the quantization.
   1180 """
   1181 if A.device.type != "cuda":
-> 1182     raise NotImplementedError(f"Device type not supported for FP4 quantization: {A.device.type}")
   1183 if quant_type not in ["fp4", "nf4"]:
   1184     raise NotImplementedError(f"4-bit quantization data type {quant_type} is not implemented.")

NotImplementedError: Device type not supported for FP4 quantization: cpu









why it's using cpu i need to direct passs in gpu