# Copyright 2024 The MediaPipe Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Functions to perform the checkpoint conversion.""" import contextlib import ctypes import os from typing import Any, List, Optional from absl import logging import numpy as np from mediapipe.tasks.python.core import mediapipe_c_bindings from mediapipe.tasks.python.core import mediapipe_c_utils from mediapipe.tasks.python.core import serial_dispatcher from mediapipe.tasks.python.genai.converter import converter_base from mediapipe.tasks.python.genai.converter import converter_factory from mediapipe.tasks.python.genai.converter import external_dependencies from mediapipe.tasks.python.genai.converter import quantization_util jnp = external_dependencies.jnp _CTYPES_SIGNATURES = ( mediapipe_c_utils.CStatusFunction( func_name='MpLlmConverterGenerateCpuTfLite', core_argtypes=( ctypes.c_char_p, # model_type ctypes.c_char_p, # weight_path ctypes.c_char_p, # vocab_model_file ctypes.c_bool, # is_quantized ctypes.c_char_p, # output_tflite_file ), ), mediapipe_c_utils.CStatusFunction( func_name='MpLlmConverterGenerateGpuTfLite', core_argtypes=( ctypes.c_char_p, # model_type ctypes.c_char_p, # weight_path ctypes.c_char_p, # vocab_model_file ctypes.c_bool, # is_quantized ctypes.c_bool, # obfuscate ctypes.c_char_p, # output_tflite_file ctypes.c_int, # lora_rank ctypes.c_char_p, # lora_weight_path ctypes.c_char_p, # lora_output_tflite_file ctypes.c_char_p, # lora_main_model_type ctypes.c_char_p, # image_encoder_file ctypes.c_char_p, # image_adapter_file ctypes.c_char_p, # submodel_type ctypes.c_bool, # use_dynamic_ple ctypes.c_bool, # apply_srq ctypes.c_int, # block_size ), ), mediapipe_c_utils.CStatusFunction( func_name='MpLlmConverterConvertHfTokenizer', core_argtypes=( ctypes.c_char_p, # vocab_model_file ctypes.c_char_p, # output_vocab_file ), ), ) def _safe_encode_str(value: str | None) -> bytes: """Encodes a string to bytes, returning an empty byte string if None.""" return value.encode('utf-8') if value is not None else b'' class ConversionConfig(object): """Config for checkpoint conversion. Attributes: input_ckpt: Directory or path for the input checkpoint. ckpt_format: Checkpoint format, e.g. 'safetensors', 'pytorch'. model_type: Name of the model, e.g. GEMMA_2B. backend: Target backend to run the model. Can be either "cpu" or "gpu". output_dir: Where the output file(s) to be stored. is_quantized: Whether the checkpoint is already quantized. If the checkpoint is already quantized, the converter will not quantize it again and it will ignore the quantization parameters. is_symmetric: Whether to quantize symmetrically. attention_quant_bits: Target quantization bits for the attention layers. feedforward_quant_bits: Target quantization bits for the feedforward layers. embedding_quant_bits: Target quantization bits for the embedding layers. combine_file_only: Whether to combine the weight files only (assuming the weight files are already existed). vocab_model_file: The file path to the 1) SentencePiece vocab model; 2) Hugging Face BPE tokenizer files; 1) is applicable for the Gemma model and 2) is applicable for other models. When 2) is used, the provided path is expected to point to a directory that contains both tokenizer.json and tokenizer_config.json files. obfuscate: Whether to obfuscate the model. output_tflite_file: (optional) the output tflite filename. If not provided, the output will be `model.tflite` stored in the output_dir. fp16_scale: A scalar value between [0, 1]. Some models can run into activation overflow issue when running in 16-bit floating point mode. To solve this, we need to scale down the weights of certain layers. See go/llm-on-device-fp16 for more detailed explanation. lora_ckpt: The directory or path for the lora checkpoint. Required in order to convert the lora weights. lora_rank: An integer representing the rank of LoRA. Required in order to convert the lora weights.If not provided, then the converter assumes there is no LoRA weights. Note that only the GPU backend supports LoRA. lora_alpha: A float representing the scale of LoRA weights. lora_output_tflite_file: A string indicating the name of the generated tflite file for the LoRA weight. Only applicable when the lora_rank is not zero. lora_main_model_type: If the LoRA will be applied to a submodel packaged with a main model, what the main model type is. image_encoder_file: A string with the name of the image encoder tflite file. image_adapter_file: A string with the name of the image adapter tflite file. submodel_type: Name of submodel, e.g. GEMMA_2B. use_fake_weights: Whether to use fake weights. If set to True, the weights will be filled with zeros. use_dynamic_ple: Whether any PLE embeddings should be loaded dynamically. Default is true, which will cause embeddings to only be loaded into VRAM on demand. use_mse_quant: Whether to use MSE quantization for recomputing scales. block_size: Default of 0 is off. Can also be 32 or 128, for using blockwise Q4_0 compression on the models and layers which support this. apply_srq: Default of None will apply SRQ for pre-quantized models only. Will override this behavior if set explicitly. """ def __init__( self, input_ckpt: str, ckpt_format: str, model_type: str, backend: str, output_dir: str, is_quantized: bool = False, is_symmetric: bool = True, attention_quant_bits: int = 8, feedforward_quant_bits: int = 8, embedding_quant_bits: int = 8, combine_file_only: bool = False, vocab_model_file: str = '', obfuscate: bool = False, output_tflite_file: Optional[str] = None, fp16_scale: Optional[float] = None, block_size: int = 0, lora_ckpt: Optional[str] = None, lora_rank: Optional[int] = None, lora_alpha: Optional[float] = None, lora_output_tflite_file: Optional[str] = None, lora_main_model_type: Optional[str] = None, image_encoder_file: Optional[str] = None, image_adapter_file: Optional[str] = None, submodel_type: Optional[str] = None, use_fake_weights: bool = False, use_dynamic_ple: bool = True, use_mse_quant: bool = False, apply_srq: Optional[bool] = None, ): self.input_ckpt = input_ckpt self.ckpt_format = ckpt_format self.model_type = model_type self.backend = backend if os.path.isfile(output_dir): raise ValueError('Output directory mush not point to an existing file.') if not os.path.isdir(output_dir): logging.info('Creating output directory: %s', output_dir) os.makedirs(output_dir, exist_ok=True) self.output_dir = output_dir self.is_quantized = is_quantized self.is_symmetric = is_symmetric self.attention_quant_bits = attention_quant_bits self.feedforward_quant_bits = feedforward_quant_bits self.embedding_quant_bits = embedding_quant_bits self.combine_file_only = combine_file_only self.vocab_model_file = vocab_model_file self.obfuscate = obfuscate self.image_encoder_file = image_encoder_file self.image_adapter_file = image_adapter_file self.submodel_type = submodel_type self.use_fake_weights = use_fake_weights self.use_dynamic_ple = use_dynamic_ple self.use_mse_quant = use_mse_quant self.apply_srq = apply_srq if output_tflite_file: parent_dir = os.path.dirname(output_tflite_file) if not os.path.isdir(parent_dir): logging.info('Creating tflite parent directory: %s', parent_dir) os.makedirs(parent_dir, exist_ok=True) self.output_tflite_file = output_tflite_file else: self.output_tflite_file = os.path.join(output_dir, 'model.tflite') self.fp16_scale = None if block_size not in (0, 32, 128): raise ValueError('Block size must be one of (0 (off), 32, 128)') else: self.block_size = block_size self.lora_ckpt = lora_ckpt self.lora_rank = lora_rank self.lora_alpha = lora_alpha self.lora_output_tflite_file = lora_output_tflite_file self.lora_main_model_type = lora_main_model_type if (self.lora_ckpt is None) ^ (self.lora_rank is None): raise ValueError( 'lora_ckpt and lora_rank must be either both provided or both not' ' provided.' ) if self.lora_rank is not None: if backend == 'cpu': raise ValueError('LoRA is not supported for CPU backend.') lora_applicable_models = [ 'GEMMA_2B', 'GEMMA2_2B', 'PHI_2', 'GEMMA3_1B', 'GEMMA3_4B', 'GEMMA3_12B', 'GEMMA3_27B', 'GEMMA3_300M', ] if model_type not in lora_applicable_models: raise ValueError( 'LoRA is only applicable for the model_type:' f' {", ".join(lora_applicable_models)}, but get model_type:' f' {model_type}.' ) @contextlib.contextmanager def filemanager(filename: str, mode: str): try: with open(filename, mode) as f: yield f finally: pass class _LlmConverter: """Bundles all conversion logic for LLM models.""" def __init__(self, lib: Any): self._lib = serial_dispatcher.SerialDispatcher(lib, _CTYPES_SIGNATURES) def quantize_by_actions( self, actions: List[converter_base.QuantizationAction], backend: str, is_symmetric: bool, use_mse_quant: bool = False, block_size: int = 0, ): """Quantizes the weights by actions. Args: actions: A list of QuantizationAction that contains the information and tensor values to be quantized. backend: Target backend to run the model. Can be either "cpu" or "gpu". is_symmetric: Whether to quantize symmetrically. use_mse_quant: Whether to use MSE quantization for recomputing scales. block_size: 0 to disable, and 32 or 128 to use blockwise Q4_0 for recomputing scales. Returns: A dictionary that maps from the updated tensor names to the quantized tensor values + a boolean that indicates whether the tensor values need to be packed (only applicable for the 4-bit quantized weights). """ output_tensors = {} qvalue_suffix = '_quantized_value' scale_suffix = '_quantized_scale' zp_suffix = '_quantized_zp' for action in actions: if action.tensor_value is None: continue # The dtype needs to be compared in string as it is a custom numpy dtype. # Explicitly cast the bfloat16 and float16 dtype to float32 to make sure # its value is converted and serialized correctly. if ( str(action.tensor_value.dtype) == 'bfloat16' or action.tensor_value.dtype == np.float16 ): action.tensor_value = action.tensor_value.astype(np.float32) if ( (not action.is_quantized) and action.tensor_value.dtype != np.float32 and action.tensor_value.dtype != np.int8 ): raise ValueError( 'All tensors should be casted to either float32 or int8, but got: ' '%s' % action.tensor_value.dtype ) if action.is_quantized: pack = action.tensor_value.dtype == jnp.int4 if qvalue_suffix in action.target_name: target_name = action.target_name[: -len(qvalue_suffix)] # Stores the quantized value in int8 for 4-bit quantization. if pack: action.tensor_value = action.tensor_value.astype(jnp.int8) output_tensors[target_name] = (action.tensor_value, pack) elif ( scale_suffix in action.target_name or zp_suffix in action.target_name ): output_tensors[action.target_name] = ( action.tensor_value, False, ) else: output_tensors[action.target_name] = (action.tensor_value, False) if action.quantize_axis: quant_bits = action.quantize_bits quant_axis = action.quantize_axis.copy() if ( block_size > 0 and quant_bits == 4 and quant_axis != [0] and quant_axis != [1] and quant_axis != [0, 1] ): # Turn off 4bit blockwise quant for unsupported quantization axis # configurations print( 'Disabling blockwise 4bit quant for unsupported quantization' ' axis: ', quant_axis, ' on tensor ', action.target_name, ) quant_bits = 8 pack = quant_bits == 4 if action.tensor_value.dtype == np.int8: if backend == 'cpu' and pack: raise ValueError( 'Converting pre-quantized checkpoint into 4-bit is not' ' supported for CPU backend.' ) output_tensors[action.target_name] = (action.tensor_value, pack) else: if is_symmetric: bs = block_size if pack else 0 target_var, scale = quantization_util.quantize_tensor( var=action.tensor_value, axis=action.quantize_axis, sym=is_symmetric, number_bits=quant_bits, use_mse_quant=use_mse_quant, # blockwise quant applies to symmetric 4bit only block_size=bs, ) if bs > 0: if len(scale.shape) > 3: raise ValueError( 'Blockwise int4 quantization not yet supported for' ' 4-dimensional or higher scales.' ) # Reshape and transpose output from blockwise quant. # We want per-block scale values as the last dimension. if quant_axis == [0] or quant_axis == [0, 1]: if len(scale.shape) <= 2: scale = scale.transpose() else: scale = scale.transpose([1, 2, 0]) elif quant_axis == [1]: if len(scale.shape) == 3: # Only tested in one case, so not sure if this generalizes. # Seems to behave identically to our experimental # quant_axis = [0, 1] support so far. scale = scale.transpose([2, 0, 1]) # Reshape to fold in the per-block dimension target_var = target_var.reshape(action.tensor_value.shape) output_tensors[action.target_name] = (target_var, pack) output_tensors[action.target_name + scale_suffix] = ( scale, False, ) zp = None else: target_var, scale, zp = quantization_util.quantize_tensor( var=action.tensor_value, axis=action.quantize_axis, sym=is_symmetric, number_bits=action.quantize_bits, use_mse_quant=use_mse_quant, ) if backend == 'cpu' and pack: target_var, scale, zp = quantization_util.update_to_uint4( target_var, scale, zp ) output_tensors[action.target_name] = (target_var, pack) output_tensors[action.target_name + scale_suffix] = (scale, False) if zp is not None: output_tensors[action.target_name + zp_suffix] = (zp, False) else: output_tensors[action.target_name] = (action.tensor_value, False) return output_tensors def combined_weight_bins_to_tflite( self, model_type: str, backend: str, weight_path: str, output_tflite_file: str, obfuscate: bool, vocab_model_file: str, lora_rank: Optional[int] = None, lora_weight_path: Optional[str] = None, lora_output_tflite_file: Optional[str] = None, lora_main_model_type: Optional[str] = None, image_encoder_file: Optional[str] = None, image_adapter_file: Optional[str] = None, submodel_type: Optional[str] = None, use_dynamic_ple: Optional[bool] = None, apply_srq: Optional[bool] = None, block_size: int = 0, ): """Combines weight files to tflite file.""" if backend == 'cpu': if lora_rank is not None: logging.fatal('LoRA is not supported for CPU backend.') self._lib.MpLlmConverterGenerateCpuTfLite( model_type.encode('utf-8'), weight_path.encode('utf-8'), vocab_model_file.encode('utf-8'), True, output_tflite_file.encode('utf-8'), ) elif backend == 'gpu': self._lib.MpLlmConverterGenerateGpuTfLite( model_type.encode('utf-8'), weight_path.encode('utf-8'), vocab_model_file.encode('utf-8'), True, obfuscate, output_tflite_file.encode('utf-8'), 0 if lora_rank is None else lora_rank, _safe_encode_str(lora_weight_path), _safe_encode_str(lora_output_tflite_file), _safe_encode_str(lora_main_model_type), _safe_encode_str(image_encoder_file), _safe_encode_str(image_adapter_file), _safe_encode_str(submodel_type), True if use_dynamic_ple is None else use_dynamic_ple, False if apply_srq is None else apply_srq, block_size, ) else: raise ValueError(f'Unsupported backend: {backend}') def convert_bpe_vocab(self, vocab_model_file: str, output_dir: str) -> str: """Converts the BPE vocab model file to SPM format. Args: vocab_model_file: The input BPE vocab model file path. output_dir: The output directory to store the SPM vocab model file. Returns: The path to the output SPM vocab model file. """ if not os.path.isdir(vocab_model_file): raise ValueError( 'The input BPE vocab model file path is expected to be a directory' ' that contains both tokenizer.json and tokenizer_config.json files.' ) output_vocab_file = os.path.join(output_dir, 'spm.model') self._lib.MpLlmConverterConvertHfTokenizer( vocab_model_file.encode('utf-8'), output_vocab_file.encode('utf-8'), ) return output_vocab_file def sort_layer_info(self, layer_info_file: str) -> None: """Loads and sorts the layer info file.""" layer_info = [] with filemanager(layer_info_file, 'r') as finfo: for line in finfo: line = line.strip() if line: layer_info.append(line) layer_info = list(set(layer_info)) layer_info.sort() with filemanager(layer_info_file, 'w') as finfo: for line in layer_info: finfo.write(line + '\n') finfo.write('\n') def maybe_quantize_and_write_tensors_to_bins( self, ckpt_loader: converter_base.CkptLoaderBase, config: ConversionConfig, ) -> None: """Quantizes the weight tensors according to the loader and writes them to bins.""" actions = ckpt_loader.load_to_actions() for action in actions: # Quantize the weight quantized_tensors = self.quantize_by_actions( action, config.backend, config.is_symmetric, config.use_mse_quant, block_size=config.block_size, ) del action # Write the tensors into file(s). writer = converter_factory.create_writer( writer_type='weight_bins', output_dir=config.output_dir, backend=config.backend, ) writer.write_variables(quantized_tensors, config.use_fake_weights) del quantized_tensors del writer def convert_checkpoint(self, config: ConversionConfig) -> None: """Converts the checkpoint to tflite file.""" logging.info('input folder: %s', config.input_ckpt) if os.path.isdir(config.vocab_model_file): vocab_model_path = self.convert_bpe_vocab( config.vocab_model_file, config.output_dir ) else: vocab_model_path = config.vocab_model_file if not config.combine_file_only: # Load the layer weights and prepare the quantization configurations. loader = converter_factory.create_ckpt_loader( config.ckpt_format, ckpt_path=config.input_ckpt, is_quantized=config.is_quantized, is_symmetric=config.is_symmetric, backend=config.backend, attention_quant_bits=config.attention_quant_bits, feedforward_quant_bits=config.feedforward_quant_bits, embedding_quant_bits=config.embedding_quant_bits, special_model=config.model_type, fp16_scale=config.fp16_scale, ) self.maybe_quantize_and_write_tensors_to_bins(loader, config) if config.lora_ckpt is not None and config.lora_ckpt != config.input_ckpt: # If lora ckpt and the input ckpt is the same. The lora conversion is # handled in the previous loader. lora_loader = converter_factory.create_ckpt_loader( config.ckpt_format, ckpt_path=config.lora_ckpt, is_quantized=config.is_quantized, is_symmetric=config.is_symmetric, backend=config.backend, attention_quant_bits=config.attention_quant_bits, feedforward_quant_bits=config.feedforward_quant_bits, embedding_quant_bits=config.embedding_quant_bits, special_model=config.model_type, ) self.maybe_quantize_and_write_tensors_to_bins(lora_loader, config) self.sort_layer_info(os.path.join(config.output_dir, 'layer_info.txt')) self.combined_weight_bins_to_tflite( config.model_type, config.backend, weight_path=config.output_dir, output_tflite_file=config.output_tflite_file, obfuscate=config.obfuscate, vocab_model_file=vocab_model_path, lora_rank=config.lora_rank, lora_weight_path=config.output_dir, lora_output_tflite_file=config.lora_output_tflite_file, lora_main_model_type=config.lora_main_model_type, image_encoder_file=config.image_encoder_file, image_adapter_file=config.image_adapter_file, submodel_type=config.submodel_type, use_dynamic_ple=config.use_dynamic_ple, # Default SRQ behavior is to match is_quantized apply_srq=config.is_quantized if config.apply_srq is None else config.apply_srq, block_size=config.block_size, ) def __enter__(self) -> '_LlmConverter': """Returns `self` upon entering the runtime context.""" return self def __exit__(self, *args) -> None: """Shuts down the LlmConverter on exit of the context manager. Args: *args: Unused. Raises: RuntimeError: If the LLM converter failed to close. """ del args # Unused. self._lib.close() def convert_checkpoint(config: ConversionConfig) -> None: """Converts the checkpoint to tflite file.""" lib = mediapipe_c_bindings.load_raw_library(_CTYPES_SIGNATURES) with _LlmConverter(lib) as converter: converter.convert_checkpoint(config) def quantize_by_actions( actions: List[converter_base.QuantizationAction], backend: str, is_symmetric: bool, use_mse_quant: bool = False, block_size: int = 0, ): """Quantizes the weights by actions. Args: actions: A list of QuantizationAction that contains the information and tensor values to be quantized. backend: Target backend to run the model. Can be either "cpu" or "gpu". is_symmetric: Whether to quantize symmetrically. use_mse_quant: Whether to use MSE quantization for recomputing scales. block_size: 0 to disable, and 32 or 128 to use blockwise quantization for recomputing scales. Returns: A dictionary that maps from the updated tensor names to the quantized tensor values + a boolean that indicates whether the tensor values need to be packed (only applicable for the 4-bit quantized weights). """ lib = mediapipe_c_bindings.load_raw_library(_CTYPES_SIGNATURES) with _LlmConverter(lib) as converter: return converter.quantize_by_actions( actions, backend, is_symmetric, use_mse_quant, block_size )