Transformers documentation
LiteRT
LiteRT
LiteRT (formerly TensorFlow Lite) is Google’s runtime for on-device inference. The model format is .tflite and language models ship as one .litertlm file for the LiteRT-LM runtime.
Export a Transformers model with litert-torch. It lowers the torch.export graph to LiteRT directly, and not through ONNX or a TensorFlow SavedModel.
pip install litert-torch
export_hf loads a language model from the Hub, quantizes the weights to int8 by default, and writes model.litertlm.
litert-torch export_hf \
--model="HuggingFaceTB/SmolLM2-135M-Instruct" \
--output_dir="./smollm2_litertlm"litert_torch.convert traces a model with sample inputs and exports a .tflite file. The returned object also runs it, so the export can be checked in place.
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())) # capitalTransformers integration
- from_pretrained() loads the model weights in safetensors format.
- litert-torch runs torch.export and lowers the graph to LiteRT operators.
export_hfadds the KV cache, prefill and decode signatures, and int8 quantization. - AutoTokenizer loads the tokenizer.
export_hfpacks it and the chat template into the.litertlmfile. - At runtime,
.tfliteruns on LiteRT and.litertlmon LiteRT-LM, from Kotlin, Swift, C++, or Python (ai-edge-litertandlitert-lm-api). The oldertflite-runtimewheels stop at Python 3.11.
Transformers v4 documented
optimum-cli export tflite, which converted through TensorFlow. It was removed with TensorFlow support in v5 (#40760) and is not part of Optimum 2.x.
Resources
- LiteRT and LiteRT-LM docs
- Convert PyTorch models and GenAI models guides