Conversion script!

Should work the same, but much speedier and fixes the tensor issue.

convert.py

@@ -0,0 +1,327 @@
+# converter.py
+
+import sys
+import torch
+import safetensors.torch as st
+import logging
+import math
+import tflite.Model
+import tflite.SubGraph
+from tflite.TensorType import TensorType
+
+# Set up logging
+logger = logging.getLogger(__name__)
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
+    level=logging.INFO
+)
+
+# Define scale and size mappings
+name_of_tensor_type = {
+    0: "FLOAT32",
+    9: "INT8   ",
+    17: "INT4   ",
+}
+
+dtype_for_tensor_type = {
+    0: torch.float32,
+    9: torch.int8,
+    17: torch.uint8,  # Because torch.int4 doesn't exist
+}
+
+size_for_tensor_type = {
+    0: 4,
+    9: 1,
+    17: 0.5,
+}
+
+# Function to update target tensor names
+def update_target_name(target_name: str) -> str:
+    """Updates the target name to match the tensor name convention."""
+    def reverse_replace(theStr: str, a, b):
+        return theStr.replace(b, a)
+    
+    target_name = reverse_replace(target_name, ".weight", ".w")
+    target_name = reverse_replace(target_name, 
+        "model.layers.", "params.lm.transformer.x_layers_"
+    )
+
+    target_name = reverse_replace(target_name, 
+        "mlp.gate_proj", "ff_layer.ffn_layer1_gate"
+    )
+    target_name = reverse_replace(target_name, "mlp.up_proj", "ff_layer.ffn_layer1")
+    target_name = reverse_replace(target_name, "mlp.down_proj", "ff_layer.ffn_layer2")
+
+    target_name = reverse_replace(target_name,
+        "post_layer_norm.weight", "post_layer_norm.scale"
+    )
+    target_name = reverse_replace(target_name,
+        "post_attention_layernorm", "post_layer_norm"
+    )
+    
+    target_name = reverse_replace(target_name, 
+        "pre_layer_norm.weight", "pre_layer_norm.scale"
+    )
+    target_name = reverse_replace(target_name, "input_layernorm", "pre_layer_norm")
+    
+    target_name = reverse_replace(target_name, "self_attn.q_proj", "self_attention.q")
+    target_name = reverse_replace(target_name, "self_attn.k_proj", "self_attention.k")
+    target_name = reverse_replace(target_name, "self_attn.v_proj", "self_attention.v")
+    target_name = reverse_replace(target_name, "self_attn.o_proj", "self_attention.post")
+    target_name = reverse_replace(target_name, 
+        "model.embed_tokens", "params.lm.softmax.logits_ffn"
+    )
+    target_name = reverse_replace(target_name, "final_ln.weight", "final_ln.scale")
+    target_name = reverse_replace(target_name, "model.norm", "params.lm.final_ln")
+    
+    return target_name
+
+# Optimized dequantization for INT4
+def convert_quantized_int4_to_fp(quantized_data, scale_data, dims, dim_scale, dtype):
+    zero_point = 8
+
+    # Reshape quantized data to 1D tensor
+    quantized_data = quantized_data.view(-1)
+
+    # Extract low and high 4 bits
+    low_bits = (quantized_data & 0x0F).type(torch.int8)
+    high_bits = (quantized_data >> 4).type(torch.int8)
+
+    # Concatenate low and high bits
+    int4_values = torch.stack((low_bits, high_bits), dim=1).view(-1)
+    int4_values = int4_values - zero_point  # Adjust zero point
+
+    # Apply scaling
+    scaled_data = int4_values.type(dtype) * scale_data
+
+    # Reshape to original dimensions
+    scaled_data = scaled_data.view(dims[0], dims[1])
+
+    return scaled_data
+
+# Function to dequantize INT8
+def convert_quantized_int8_to_fp(quantized_data, scale_data, dims, dim_scale, dtype):
+    zero_point = 0  # Assuming zero_point=0 for int8
+
+    # Reshape quantized data to 1D tensor
+    quantized_data = quantized_data.view(-1).type(torch.int8)
+    
+    # Handle scale_data based on dim_scale
+    if dim_scale:
+        # Per-column scaling
+        scale_data = scale_data.repeat_interleave(2)
+    else:
+        # Per-row scaling
+        scale_data = scale_data.repeat_interleave(2)
+    
+    # Convert scale_data to the same dtype
+    scale_data = scale_data.to(dtype=dtype)
+
+    # Apply scaling
+    scaled_data = (quantized_data - zero_point).type(dtype) * scale_data
+
+    # Reshape to original dimensions
+    scaled_data = scaled_data.view(dims[0], dims[1])
+
+    return scaled_data
+
+def main():
+    # Check command-line arguments
+    if len(sys.argv) < 3:
+        print("Usage: python converter.py <path_to_tflite_model> <output_safetensors_file> [fp32|fp16|bf16]")
+        sys.exit(1)
+
+    tflite_model_path = sys.argv[1]
+    output_safetensors_path = sys.argv[2]
+    dtype_arg = sys.argv[3] if len(sys.argv) >= 4 else "fp32"
+
+    if dtype_arg == "fp32":
+        TARGET_DTYPE = torch.float32
+    elif dtype_arg == "fp16":
+        TARGET_DTYPE = torch.float16
+    elif dtype_arg == "bf16":
+        TARGET_DTYPE = torch.bfloat16
+    else:
+        print("Unsupported dtype. Choose from fp32, fp16, bf16.")
+        sys.exit(1)
+
+    logger.info(f"Starting conversion with TARGET_DTYPE={TARGET_DTYPE}")
+
+    # Read the TFLite model
+    with open(tflite_model_path, "rb") as input_file:
+        buf = bytearray(input_file.read())
+
+    model: tflite.Model.Model = tflite.Model.Model.GetRootAs(buf)
+    graph: tflite.SubGraph.SubGraph = model.Subgraphs(0)
+
+    # Initialize dictionaries to hold tensors
+    i4_tensors = {}
+    i8_tensors = {}
+    fp32_tensors = {}
+    scale_tensors = {}
+    tensor_dims = {}
+
+    # Read and sort tensors
+    for i in range(graph.TensorsLength()):
+        tensor = graph.Tensors(i)
+        tensor_name = tensor.Name().decode("utf-8")
+        tensor_type: TensorType = tensor.Type()
+
+        if tensor_name.endswith(".w_quantized_scale"):
+            scale_tensors[tensor_name] = tensor
+        elif tensor_type == TensorType.INT4:
+            i4_tensors[tensor_name] = tensor
+        elif tensor_type == TensorType.INT8:
+            i8_tensors[tensor_name] = tensor
+        elif tensor_type == TensorType.FLOAT32:
+            fp32_tensors[tensor_name] = tensor
+
+        tensor_buf_size = tensor.Shape(0)
+        tensor_size = tensor_buf_size // size_for_tensor_type[tensor_type]
+        
+        shape = None
+        if (".self_attention.q." in tensor_name
+            or ".self_attention.post." in tensor_name) and tensor_size == 4_194_304:
+            shape = (2048, 2048)
+        elif (".self_attention.k." in tensor_name
+              or ".self_attention.v." in tensor_name) and tensor_size == 524_288:
+            shape = (256, 2048)
+        elif (".ff_layer.ffn_layer1_gate." in tensor_name
+              or ".ff_layer.ffn_layer1." in tensor_name) and tensor_size == 25_165_824:
+            shape = (12_288, 2048)
+        elif ".ff_layer.ffn_layer2." in tensor_name and tensor_size == 25_165_824:
+            shape = (2048, 12_288)
+        elif "params.lm.softmax.logits_ffn.w" == tensor_name and tensor_size == 524_550_144:
+            shape = (256_128, 2048)
+        # LayerNorm weights are of shape {1, 1, 2048}
+        elif "layer_norm" in tensor_name and tensor_size == 2048:
+            shape = (1, 1, 2048)
+        else:
+            # Default to 1D if shape is unknown
+            pass
+
+        tensor_dims[tensor_name] = shape
+
+    # Dictionary to hold dequantized tensors
+    tensor_dict = {}
+
+    # Dequantize FP32 tensors
+    for tensor_name, tensor in fp32_tensors.items():
+        logger.info(f"Saving fp32 {tensor_name}...")
+        buffer_meta = model.Buffers(tensor.Buffer())
+        dims = tensor_dims.get(tensor_name)
+
+        target_name = update_target_name(tensor_name)
+
+        tensor_data = torch.frombuffer(buffer=buf, 
+                                       dtype=torch.float32, 
+                                       offset=buffer_meta.Offset(),
+                                       count=buffer_meta.Size() // 4)
+        
+        # Assign reshaped tensor back
+        if dims is not None:
+            tensor_data = tensor_data.reshape(dims)
+
+        if TARGET_DTYPE != torch.float32:
+            tensor_data = tensor_data.to(dtype=TARGET_DTYPE)
+
+        tensor_dict[target_name] = tensor_data
+
+    del fp32_tensors
+
+    # Dequantize INT8 tensors
+    for tensor_name, quantized_tensor in i8_tensors.items():
+        buffer_meta = model.Buffers(quantized_tensor.Buffer())
+        scale_tensor_name = tensor_name + "_quantized_scale"
+        scale_buf_meta = model.Buffers(scale_tensors[scale_tensor_name].Buffer())
+        dims = tensor_dims.get(tensor_name)
+
+        logger.info(f"Dequantizing int8 {dims} {tensor_name}...")
+
+        target_name = update_target_name(tensor_name)
+
+        quantized_buf = torch.frombuffer(buffer=buf, 
+                                         dtype=torch.int8, 
+                                         offset=buffer_meta.Offset(),
+                                         count=buffer_meta.Size())
+        
+        scale_buf = torch.frombuffer(buffer=buf,
+                                     dtype=torch.float32,
+                                     offset=scale_buf_meta.Offset(),
+                                     count=scale_buf_meta.Size() // 4)
+        
+        # MediaPipe TfLiteWeightAccessor::BuildWeightsMapFromTfliteModel sets
+        # dim_scale=0, so we do the same.
+        tensor_data = convert_quantized_int8_to_fp(
+            quantized_data=quantized_buf,
+            scale_data=scale_buf,
+            dims=dims,
+            dim_scale=0,
+            dtype=TARGET_DTYPE
+        )
+        
+        tensor_dict[target_name] = tensor_data
+
+        del quantized_buf, scale_buf
+
+    del i8_tensors
+
+    # Dequantize INT4 tensors
+    for tensor_name, quantized_tensor in i4_tensors.items():
+        buffer_meta = model.Buffers(quantized_tensor.Buffer())
+        scale_tensor_name = tensor_name + "_quantized_scale"
+        scale_buf_meta = model.Buffers(scale_tensors[scale_tensor_name].Buffer())
+        dims = tensor_dims.get(tensor_name)
+
+        logger.info(f"Dequantizing int4 {dims} {tensor_name}...")
+
+        target_name = update_target_name(tensor_name)
+
+        quantized_buf = torch.frombuffer(buffer=buf, 
+                                         dtype=torch.uint8, 
+                                         offset=buffer_meta.Offset(),
+                                         count=buffer_meta.Size())
+        
+        scale_buf = torch.frombuffer(buffer=buf,
+                                     dtype=torch.float32,
+                                     offset=scale_buf_meta.Offset(),
+                                     count=scale_buf_meta.Size() // 4)
+        
+        # Special handling for 'logits_ffn.w_quantized_scale'
+        if 'logits_ffn.w_quantized_scale' in tensor_name:
+            # Assuming two scale factors per row, average them
+            if scale_buf.numel() % 2 != 0:
+                logger.error(f"Scale data size for {tensor_name} is not even. Cannot average.")
+                sys.exit(1)
+            scale_data = scale_buf.view(-1, 2).mean(dim=1)  # Average every two scale factors
+            # Repeat each scale factor twice to match the two int4 values
+            scale_data = scale_data.repeat_interleave(2)
+        else:
+            # General handling: per-row scaling, repeat each scale factor twice
+            scale_data = scale_buf.repeat_interleave(2)
+        
+        # Convert and reshape quantized_data
+        tensor_data = convert_quantized_int4_to_fp(
+            quantized_data=quantized_buf,
+            scale_data=scale_data,
+            dims=dims,
+            dim_scale=0,
+            dtype=TARGET_DTYPE
+        )
+        
+        tensor_dict[target_name] = tensor_data
+
+        del quantized_buf, scale_buf
+
+    del i4_tensors
+    del scale_tensors
+
+    del buf, model, graph
+
+    # Save all tensors to the safetensors file
+    logger.info(f"Saving to {output_safetensors_path}...")
+    st.save_file(tensor_dict, output_safetensors_path)
+    logger.info(f"Success! Saved to {output_safetensors_path}")
+
+if __name__ == "__main__":
+    main()