Compila modelos con el SDK de Google Tensor

El SDK de Google Tensor te permite optimizar y ejecutar modelos de aprendizaje automático (AA) directamente en la TPU dedicada de Google Tensor. Además de los modelos de AA estándar, puedes compilar y acelerar modelos de lenguaje dentro de tu flujo de trabajo de LiteRT.

Para ciertos modelos preoptimizados, también proporcionamos un archivo binario de opciones de compilador adicionales que puedes incluir en tu flujo de trabajo. Para lograr un rendimiento óptimo en Google Tensor, te recomendamos que uses las siguientes marcas de optimización durante la compilación:

Marcas de compilación para Google Tensor

Optimiza el proceso de compilación de tu modelo aplicando marcas de configuración específicas para adaptar el rendimiento y el uso de recursos. Puedes usar estos parámetros en tu flujo de trabajo de Python de LiteRT para ajustar el comportamiento de compilación de los modelos de PyTorch y TFLite.

Marcar Requisito Descripción Valor predeterminado
google_tensor_truncation_type Opcional Establece el tipo de datos de destino para las operaciones de punto flotante.
Valores admitidos: auto (predeterminado), bfloat16, half, no_truncation		 automático
google_tensor_sharding_intensity Opcional Controla la agresividad con la que se divide el modelo para el procesamiento paralelo.
Opciones: minimal, moderate, extensive, maximum.		 minimalista
google_tensor_int64_to_int32 Opcional Se establece como verdadero para permitir que el compilador convierta números enteros de 64 bits en números enteros de 32 bits, lo que puede ser necesario para algunos modelos. Falso
google_tensor_enable_large_model_support Opcional Se establece como verdadero si tu modelo tiene más de 2 GB. Falso
google_tensor_enable_4bit_compilation Opcional Permite la compilación de modelos con operaciones de convolución que contienen parámetros de 4 bits. Falso
google_tensor_extra_options_path Opcional Pasa opciones de compilador adicionales como un archivo binario. "" (vacío)

Puedes usar estas marcas con el flujo de Python de LiteRT como se muestra en los siguientes fragmentos de código:

  • Cuando se compila un modelo de PyTorch con ai_edge_torch

    compiled_models = (
  ai_edge_torch.experimental_add_compilation_backend(tensor_g5_target,
  flag_name1=value1,
  flag_name2=value2, ...).convert(
    channel_last_selfie_segmentation,
    sample_input))

  • Cuando se compila un modelo de TFLite

    compiled_models = aot_lib.aot_compile(
    tflite_model_path,
    target=[tensor_g5_target],
    flag_name1=value1,
    flag_name2=value2, ...)

Ejemplo de uso

En el siguiente ejemplo, se usa la marca google_tensor_truncation_type="half":

  compiled_models = aot_lib.aot_compile(
      tflite_model_path,
      target=[tensor_g5_target],
      keep_going=False,
      google_tensor_truncation_type="half"
  )

Consulta LiteRT AOT Colab para obtener más información.

Compila modelos de lenguaje para Google Tensor

Para compilar modelos de lenguaje para Google Tensor, sigue las instrucciones de la compilación de AOT de NPU.

Para exportar LLMs para las TPU de Google Tensor, sigue el ejemplo de las marcas adicionales necesarias para la compilación de NPU.

Ejemplo:

litert-torch export-hf \
  --model=google/gemma-3-270m-it \
  --output_dir=/tmp/gemma3-270m-google-tensor-g5 \
  --split_cache \
  --externalize_embedder \
  --prefill_lengths=128, \
  --cache_length=1280 \
  --quantization_recipe="weight_only_wi8_afp32" \
  --aot_backend=GOOGLE \
  --aot_soc_model=Tensor_G5 \
  --aot_compilation_config_dict='{"google_tensor_enable_large_model_support": True}'