Cómo convertir modelos de IA generativa de PyTorch para la inferencia en el dispositivo

La API de LiteRT Torch Generative es una biblioteca de alto rendimiento diseñada para crear y convertir modelos de PyTorch basados en transformadores al formato LiteRT/LiteRT-LM. Esto permite a los desarrolladores implementar sin problemas modelos de IA generativa, en particular modelos de lenguaje grandes (LLM), para la generación de texto e imágenes integrado en el dispositivo con facilidad.

La API de Torch Generative admite la conversión de modelos para CPU, GPU y NPU. Si combinas la API de Torch Generative con LiteRT-LM, puedes crear aplicaciones responsivas y enfocadas en la privacidad que ejecuten modelos generativos por completo integrado en el dispositivo.

Convierte desde la biblioteca de Hugging Face Transformer

La extensión de exportación de Hugging Face de LiteRT Torch proporciona una ruta optimizada para convertir modelos de IA generativa directamente desde la biblioteca de Hugging Face Transformers al formato LiteRT-LM. En comparación con las APIs de LiteRT Torch Generative que te proporcionan bloques de construcción de PyTorch para compilar y optimizar modelos personalizados, esta herramienta controla las complejidades de descargar pesos, traducir arquitecturas de modelos de PyTorch y aplicar técnicas de optimización, como optimizaciones de gráficos y cuantización en un solo flujo de trabajo. Genera un archivo .litertlm, que está optimizado para la inferencia integrado en el dispositivo en CPU, GPU y NPU con el entorno de ejecución de LiteRT-LM.

Requisitos previos

Antes de usar la extensión de exportación, asegúrate de tener la siguiente configuración:

  • Instala el paquete de Python de LiteRT Torch. La extensión de exportación de Hugging Face está integrada directamente en el paquete litert-torch.
  • (Opcional) Para la compilación de NPU, instala las extensiones del SDK de LiteRT NPU con pip install ai-edge-litert[npu-sdk]. Para obtener más detalles, puedes seguir la Colab de compilación AOT de LiteRT NPU.
  • El entorno de Hugging Face está configurado si deseas cargar desde el hub de Hugging Face directamente. La herramienta export_hf usa los mecanismos de autenticación de transformadores estándar, como HF_TOKEN o la CLI. Mira el ejemplo:

Para descargar modelos protegidos (como Gemma o Llama), debes autenticarte con Hugging Face con la CLI o una variable de entorno:

# Set your Hugging Face token as an environment variable
export HF_TOKEN="your_hugging_face_token"

# Or use the Hugging Face CLI login
hf auth login

Uso básico

Puedes usar export_hf con la línea de comandos o la API de Python. La herramienta descargará automáticamente el modelo de Hugging Face o lo cargará desde la ruta de acceso local proporcionada, lo rastreará, aplicará optimizaciones predeterminadas y lo convertirá en un archivo .litertlm compatible con la inferencia de CPU y GPU.

Interfaz de línea de comandos (CLI)

Usa el comando litert-torch export_hf. Debes proporcionar el ID del modelo de Hugging Face y el directorio de salida elegido.

litert-torch export_hf \
  --model=google/gemma-3-270m-it \
  --output_dir=/tmp/gemma3-270m-it-litertlm

Para exportar un modelo local o personalizado, también puedes pasar la ruta de acceso al punto de control de safetensor:

litert-torch export_hf \
  --model=/path/to/safetensor/dir \
  --output_dir=/my_custom_litertlm

API de Python

Para la integración en secuencias de comandos o notebooks de Python, importa el export módulo de litert_torch.generative.export_hf.

from litert_torch.generative.export_hf import export

export.export(
    model='google/gemma-3-270m-it',
    output_dir='/tmp/gemma3-270m-it-litertlm',
)

Implementación integrado en el dispositivo con LiteRT-LM

Una vez que hayas exportado correctamente tu modelo a un archivo .litertlm, puedes implementarlo directamente integrado en el dispositivo con LiteRT-LM para obtener una ejecución de alto rendimiento en CPU y GPU. Consulta los detalles para usar la API de LiteRT-LM. Para la aceleración de NPU, consulta la guía de compilación AOT de NPU.

Arquitecturas admitidas

La herramienta export_hf verifica las siguientes arquitecturas de modelos de Transformers. Para verificarlo, consulta el campo model_type en config.json.

  • Gemma 3 (Gemma3ForCausalLM)
  • Gemma 3n (Gemma3nForCausalLM)
  • Gemma 4 (Gemma4ForCausalLM)
  • Llama (LlamaForCausalLM)
  • Mistral (MistralForCausalLM)
  • Qwen 2/2.5 (Qwen2ForCausalLM)
  • Qwen 3 (Qwen3ForCausalLM)
  • SmolLM 3 (SmolLM3ForCausalLM)

Configuración avanzada

Si bien puedes explorar las opciones avanzadas disponibles en las marcas de extensión, las siguientes son algunas opciones comunes que puedes probar.

Modelos de lenguaje de visión

Para los modelos compatibles, puedes configurar --task=image_text_to_text y --export_vision_encoder para cargar y exportar el modelo de codificador de visión.

Arquitecturas admitidas:

  • Gemma 3 (Gemma3ForConditionalGeneration)
  • Gemma 4 (Gemma4ForConditionalGeneration)

Configuración de cuantización

Los modelos de IA generativa suelen ser demasiado grandes para ejecutarse de manera eficiente en dispositivos perimetrales sin optimización. De forma predeterminada, export_hf aplica la receta de cuantización dynamic_wi8_afp32 con AI Edge Quantizer, que cuantiza los pesos a INT8 por canal mientras mantiene las activaciones en FP32.

Puedes anular este comportamiento predeterminado con la marca --quantization_recipe (o el quantization_recipe parámetro en Python). Puedes proporcionar el nombre de una receta integrada de AI Edge Quantizer o especificar la ruta de acceso a una receta JSON personalizada.

Ejemplo:

litert-torch export_hf \
  --model=google/gemma-3-270m-it \
  --output_dir=/tmp/gemma3-270m-it-litertlm \
  --quantization_recipe=/path/to/my/quantization_recipe.json

Anulación de plantilla de Jinja

Es posible que la plantilla de Jinja que viene con el modelo de transformadores no sea compatible con LiteRT-LM (p.ej., modelos de Gemma4). Puedes configurar la marca use_jinja_template en False o usar la opción jinja_chat_template_override para anular la plantilla.

Ejemplo:

 litert-torch export_hf \
   --model=google/gemma-4-E2B-it \
  --output_dir=/tmp/gemma4_2b_litertlm \
  --externalize_embedder \
  --jinja_chat_template_override=litert-community/gemma-4-E2B-it-litert-lm

Compilación AOT de NPU

Además de la CPU y la GPU, también puedes segmentar los aceleradores de NPU compatibles cuando exportas tus modelos proporcionando las opciones específicas de NPU.

Google Tensor

Requisitos previos: Sigue la página de Google Tensor SDK para configurar el entorno de desarrollo.

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

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}'

Para obtener más información, consulta Cómo compilar modelos con el SDK de Google Tensor.

Entorno de ejecución de IA de Qualcomm:

Requisitos previos: Sigue la integración de LiteRT Qualcomm para obtener instrucciones de configuración del SDK y dispositivos compatibles.

Ejemplo:

litert-torch export-hf \
  --model=google/gemma-3-270m-it \
  --output_dir=/tmp/gemma3-270m-google-tensor-g5 \
  --split_cache \
  --externalize_embedder \
  --quantization_recipe='' \
  --aot_backend=qualcomm \
  --aot_soc_model=SM8750

NeuroPilot de MediaTek:

Requisitos previos: Sigue la integración de LiteRT MediaTek para obtener instrucciones de configuración del SDK y dispositivos compatibles.

Ejemplo:

litert-torch export-hf \
  --model=google/gemma-3-270m-it \
  --output_dir=/tmp/gemma3-270m-google-tensor-g5 \
  --split_cache \
  --externalize_embedder \
  --aot_backend=mediatek \
  --aot_soc_model=MT8189

OpenVINO de Intel

Requisitos previos: Sigue la integración de LiteRT Intel OpenVINO para obtener instrucciones de configuración del SDK y dispositivos compatibles.

Ejemplo:

litert-torch export-hf \
  --model=google/gemma-3-270m-it \
  --output_dir=/tmp/gemma3-270m-google-tensor-g5 \
  --split_cache \
  --externalize_embedder \
  --aot_backend=intel_openvino \
  --aot_soc_model=PTL

Vuelve a crear y convertir con la API de LiteRT Torch Generative

La API de LiteRT Torch Generative también proporciona bloques de construcción para compilar y optimizar modelos personalizados de PyTorch, incluidas, entre otras, capas de normalización, atenciones y otros módulos básicos. Si tu modelo no está cubierto por la extensión de exportación de Hugging Face de LiteRT Torch, puedes compilar tus propios modelos para que sean compatibles con LiteRT y LiteRT-LM.

Hay ejemplos de modelos, incluidos LLM, modelos de difusión y modelos de ASR. No dudes en consultarlos y, luego, implementar tu propio modelo.

Para obtener más información, consulta el repositorio de GitHub de la API de Generative Torch.