Modelle mit dem Google Tensor SDK kompilieren

Mit dem Google Tensor SDK können Sie ML-Modelle (Machine Learning) optimieren und direkt auf der dedizierten TPU von Google Tensor ausführen. Neben Standard-ML-Modellen können Sie Sprachmodelle in Ihrem LiteRT-Workflow kompilieren und beschleunigen.

Für bestimmte voroptimierte Modelle stellen wir auch eine zusätzliche Binärdatei mit Compileroptionen zur Verfügung, die Sie in Ihren Workflow einbeziehen können. Um eine optimale Leistung auf Google Tensor zu erzielen, empfehlen wir, bei der Kompilierung die folgenden Optimierungs-Flags zu verwenden:

Kompilierungs-Flags für Google Tensor

Sie können den Prozess der Modellkompilierung optimieren, indem Sie bestimmte Konfigurations-Flags anwenden, um die Leistung und die Ressourcennutzung anzupassen. Sie können diese Parameter in Ihrem LiteRT-Python-Workflow verwenden, um das Kompilierungsverhalten für PyTorch- und TFLite-Modelle anzupassen.

Flag Anforderung Beschreibung Standardwert
google_tensor_truncation_type Optional Legt den Zieldatentyp für Gleitkommaoperationen fest.
Unterstützte Werte: auto (Standard), bfloat16, half, no_truncation		 Automatisch
google_tensor_sharding_intensity Optional Steuert, wie aggressiv das Modell für die parallele Verarbeitung aufgeteilt wird.
Optionen: minimal, moderate, extensive, maximum.		 minimal
google_tensor_int64_to_int32 Optional Auf „true“ setzen, damit der Compiler 64-Bit-Ganzzahlen in 32-Bit-Ganzzahlen konvertieren kann. Dies kann für einige Modelle erforderlich sein. Falsch
google_tensor_enable_large_model_support Optional Auf „true“ setzen, wenn Ihr Modell größer als 2 GB ist. Falsch
google_tensor_enable_4bit_compilation Optional Aktiviert die Kompilierung von Modellen mit Faltungsoperationen, die 4-Bit-Parameter enthalten. Falsch
google_tensor_extra_options_path Optional Übergibt zusätzliche Compileroptionen als Binärdatei. "" (leer)

Sie können diese Flags mit dem LiteRT-Python-Flow verwenden, wie in den folgenden Code-Snippets gezeigt:

  • Beim Kompilieren eines PyTorch-Modells mit ai_edge_torch

    compiled_models = (
  ai_edge_torch.experimental_add_compilation_backend(tensor_g5_target,
  flag_name1=value1,
  flag_name2=value2, ...).convert(
    channel_last_selfie_segmentation,
    sample_input))

  • Beim Kompilieren eines TFLite-Modells

    compiled_models = aot_lib.aot_compile(
    tflite_model_path,
    target=[tensor_g5_target],
    flag_name1=value1,
    flag_name2=value2, ...)

Verwendungsbeispiel

Im folgenden Beispiel wird das Flag google_tensor_truncation_type="half" verwendet:

  compiled_models = aot_lib.aot_compile(
      tflite_model_path,
      target=[tensor_g5_target],
      keep_going=False,
      google_tensor_truncation_type="half"
  )

Weitere Informationen finden Sie unter LiteRT AOT Colab.

Sprachmodelle für Google Tensor kompilieren

Folgen Sie der Anleitung unter NPU AOT-Kompilierung, um Sprachmodelle für Google Tensor zu kompilieren.

Wenn Sie LLMs für Google Tensor-TPUs exportieren möchten, folgen Sie dem Beispiel für die zusätzlichen Flags, die für die NPU-Kompilierung erforderlich sind.

Beispiel :

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