LiteRT

LiteRT (anteriormente TensorFlow Lite) é o runtime do Google para inferência em dispositivo. O formato do modelo é .tflite e os modelos de linguagem são distribuídos como um único arquivo .litertlm para o runtime LiteRT-LM.

Exporte um modelo do Transformers com litert-torch. Ele converte o grafo do torch.export diretamente para o LiteRT, sem passar por ONNX ou um SavedModel do TensorFlow.

pip install litert-torch

litert-torch export_hf \
    --model="HuggingFaceTB/SmolLM2-135M-Instruct" \
    --output_dir="./smollm2_litertlm"

import litert_torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

model_id = "google-bert/bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForMaskedLM.from_pretrained(model_id).eval()
inputs = tokenizer("Paris is the [MASK] of France.", return_tensors="pt", padding="max_length", max_length=128)

litert_model = litert_torch.convert(model, sample_kwargs=dict(inputs))
litert_model.export("bert.tflite")

outputs = litert_model(**{name: tensor.numpy() for name, tensor in inputs.items()})
mask_index = inputs["input_ids"][0].tolist().index(tokenizer.mask_token_id)
print(tokenizer.decode(outputs["logits"][0, mask_index].argmax()))  # capital

Integração com Transformers

1. from_pretrained() carrega os pesos do modelo no formato safetensors.
2. litert-torch executa torch.export e reduz o grafo para operadores do LiteRT. export_hf adiciona o cache de KV, as assinaturas de prefill e decode, e a quantização int8.
3. AutoTokenizer carrega o tokenizer. export_hf empacota-o junto com o chat template no arquivo .litertlm.
4. Em tempo de execução, .tflite roda no LiteRT e .litertlm no LiteRT-LM, a partir de Kotlin, Swift, C++ ou Python (ai-edge-litert e litert-lm-api). As rodas mais antigas do tflite-runtime param na Python 3.11.

A documentação do Transformers v4 descrevia o optimum-cli export tflite, que realizava a conversão por meio do TensorFlow. Ele foi removido junto com o suporte ao TensorFlow na v5 (#40760) e não faz parte do Optimum 2.x.

Recursos

• Documentação do LiteRT e do LiteRT-LM
• Guias Converter modelos do PyTorch e Modelos de GenAI