Converter modelos de IA generativa do PyTorch para inferência no dispositivo

A API generativa do Torch LiteRT é uma biblioteca de alta performance projetada para criar e converter modelos do PyTorch baseados em transformadores para o formato LiteRT/LiteRT-LM. Isso permite que os desenvolvedores implantem modelos de IA generativa, especificamente modelos de linguagem grandes (LLMs), para geração de texto e imagem no dispositivo com facilidade.

A API generativa do Torch oferece suporte à conversão de modelos para CPU, GPU e NPU. Ao parear a API generativa do Torch com LiteRT-LM, é possível criar aplicativos responsivos e focados na privacidade que executam modelos generativos totalmente no dispositivo.

Converter da biblioteca de transformadores do Hugging Face

A extensão de exportação do Hugging Face do Torch LiteRT oferece um caminho simplificado para converter modelos de IA generativa diretamente da biblioteca de transformadores do Hugging Face para o formato LiteRT-LM. Em comparação com APIs generativas do Torch LiteRT que fornecem blocos de construção do pytorch para criar e otimizar modelos personalizados, essa ferramenta lida com as complexidades de download de pesos, tradução de arquiteturas de modelos do PyTorch e aplicação de técnicas de otimização, como otimizações de gráficos e quantização em um único fluxo de trabalho. Ela gera um arquivo .litertlm, que é otimizado para inferência no dispositivo em CPU, GPU e NPU usando o ambiente de execução LiteRT-LM.

Pré-requisitos

Antes de usar a extensão de exportação, verifique se você tem a seguinte configuração:

  • Instale o pacote do Python do Torch LiteRT. A extensão de exportação do Hugging Face é integrada diretamente ao pacote litert-torch.
  • (Opcional) Para a compilação de NPU, instale as extensões do SDK do LiteRT NPU usando pip install ai-edge-litert[npu-sdk]. Para mais detalhes, siga o Colab de compilação AOT do LiteRT NPU.
  • O ambiente do Hugging Face está configurado se você pretende carregar diretamente do hub do Hugging Face. A ferramenta export_hf usa os mecanismos de autenticação de transformadores padrão, como HF_TOKEN ou CLI. Veja um exemplo:

Para fazer o download de modelos restritos (como Gemma ou Llama), é necessário fazer a autenticação com o Hugging Face usando a CLI ou uma variável de ambiente:

# 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

É possível usar export_hf na linha de comando ou na API Python. A ferramenta vai fazer o download automático do modelo do Hugging Face ou carregar o modelo do caminho local fornecido, rastreá-lo, aplicar otimizações padrão e convertê-lo em um .litertlm arquivo compatível com a inferência de CPU e GPU.

Interface de linha de comando (CLI)

Use o comando litert-torch export_hf. É necessário fornecer o ID do modelo do Hugging Face e o diretório de saída escolhido.

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

Para exportar um modelo local ou personalizado, também é possível transmitir o caminho para o ponto de verificação do safetensor:

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

API Python

Para integração em scripts ou notebooks Python, importe o 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',
)

Implantação no dispositivo com o LiteRT-LM

Depois de exportar o modelo para um arquivo .litertlm, é possível implantá-lo diretamente no dispositivo usando o LiteRT-LM para execução de alta performance na CPU e na GPU. Confira detalhes sobre como usar a API LiteRT-LM. Para aceleração de NPU, consulte o guia de compilação AOT de NPU.

Arquiteturas compatíveis

A ferramenta export_hf verifica as seguintes arquiteturas de modelo de transformadores. Isso pode ser verificado conferindo o campo model_type em 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)

Configurações avançadas

Embora seja possível explorar as opções avançadas disponíveis nas flags de extensão, a seguir estão alguns botões comuns que você pode tentar.

Modelos de linguagem de visão

Para modelos compatíveis, é possível definir --task=image_text_to_text e --export_vision_encoder para carregar e exportar o modelo de codificador de visão.

Arquiteturas compatíveis:

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

Configuração de quantização

Os modelos de IA generativa costumam ser muito grandes para serem executados com eficiência em dispositivos de borda sem otimização. Por padrão, export_hf aplica a receita de quantização dynamic_wi8_afp32 usando o AI Edge Quantizer, que quantiza pesos para INT8 por canal, mantendo as ativações em FP32.

É possível substituir esse comportamento padrão usando a flag --quantization_recipe (ou o parâmetro quantization_recipe no Python). Você pode fornecer o nome de uma receita integrada do AI Edge Quantizer ou especificar o caminho para uma receita JSON personalizada.

Exemplo:

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

Substituição de modelo do Jinja

O modelo do Jinja que acompanha o modelo de transformadores pode não ser compatível com o LiteRT-LM (por exemplo, modelos Gemma4). É possível definir a flag use_jinja_template como False ou usar a opção jinja_chat_template_override para substituir o modelo.

Exemplo:

 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

Compilação AOT de NPU

Além da CPU e da GPU, também é possível segmentar aceleradores de NPU compatíveis ao exportar seus modelos, fornecendo as opções específicas da NPU.

Google Tensor

Pré-requisitos: siga a página do SDK do Google Tensor para a configuração do ambiente de desenvolvimento.

Para exportar LLMs segmentando TPUs do Google Tensor, siga o exemplo das flags adicionais necessárias para a compilação de TPU.

Exemplo:

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 mais informações, consulte Compilar modelos com o SDK do Google Tensor.

Ambiente de execução de IA da Qualcomm:

Pré-requisitos: siga a integração do LiteRT Qualcomm para instruções de configuração do SDK e dispositivos compatíveis.

Exemplo:

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 da MediaTek:

Pré-requisitos: siga a integração do LiteRT MediaTek para instruções de configuração do SDK e dispositivos compatíveis.

Exemplo:

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

Intel OpenVINO

Pré-requisitos: siga a integração do LiteRT Intel OpenVINO para instruções de configuração do SDK e dispositivos compatíveis.

Exemplo:

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

Recriar e converter usando a API generativa do Torch LiteRT

A API generativa do Torch LiteRT também fornece blocos de construção para criar e otimizar modelos personalizados do PyTorch, incluindo, entre outros, camadas normalizadoras, atenções e outros módulos básicos. Se o modelo não estiver coberto por a extensão de exportação do Hugging Face do Torch LiteRT, é possível criar seus próprios modelos para serem compatíveis com o LiteRT e o LiteRT-LM.

Há exemplos de modelos, incluindo LLMs, modelos de difusão e modelos de ASR. Confira e implante seu próprio modelo.

Para mais informações, consulte o repositório do GitHub da API generativa do Torch.