LiteRT

LiteRT

LiteRT는 구글이 만든 온디바이스 추론용 런타임이에요. Transformers 모델을 .tflite 같은 경량 포맷으로 내보내 모바일 기기에서 바로 돌리는 방법을 살펴볼게요.

출처: 문서

본문

LiteRT(구 TensorFlow Lite)는 구글의 온디바이스 추론용 런타임이에요. 모델 포맷은 .tflite이고, 언어 모델은 LiteRT-LM 런타임용으로 하나의 .litertlm 파일로 배포돼요.

Transformers 모델은 litert-torch로 내보내요. 이 도구는 torch.export 그래프를 ONNX나 TensorFlow SavedModel을 거치지 않고 바로 LiteRT로 내려줘요.

pip install litert-torch

export_hf는 Hub에서 언어 모델을 불러와 기본적으로 가중치를 int8로 양자화한 뒤 model.litertlm을 작성해요.

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

litert_torch.convert는 샘플 입력으로 모델을 추적(trace)하고 .tflite 파일을 내보내요. 반환된 객체가 그대로 실행도 해 주기 때문에 내보낸 결과를 바로 확인할 수 있어요.

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

Transformers 통합 과정

  1. from_pretrained()가 모델 가중치를 safetensors 포맷으로 불러와요.
  2. litert-torch가 torch.export를 실행하고 그래프를 LiteRT 연산자로 내려줘요. export_hf는 여기에 KV cache, prefill·decode 시그니처, int8 양자화를 추가해요.
  3. AutoTokenizer가 토크나이저를 불러와요. export_hf는 토크나이저와 채팅 템플릿을 .litertlm 파일에 함께 넣어요.
  4. 실행 시점에는 .tflite가 LiteRT에서, .litertlm이 LiteRT-LM에서 돌아가요. Kotlin, Swift, C++, Python(ai-edge-litert와 litert-lm-api) 모두에서 사용할 수 있어요. 예전 tflite-runtime wheel은 Python 3.11까지만 지원해요.

[!NOTE] Transformers v4에서는 TensorFlow를 거치는 optimum-cli export tflite를 문서로 소개했었어요. 이 경로는 v5에서 TensorFlow 지원과 함께 제거됐고(#40760), Optimum 2.x에는 포함되지 않아요.

리소스 (Resources)

더 알아보기 (Learn more)