使用 Google Tensor SDK 編譯模型

Google Tensor SDK 可讓您直接在 Google Tensor 的專用 TPU 上，最佳化及執行機器學習 (ML) 模型。除了標準機器學習模型，您還可以在 LiteRT 工作流程中編譯及加速語言模型。

對於某些預先最佳化模型，我們也會提供額外的編譯器選項二進位檔案，供您納入工作流程。如要讓 Google Tensor 發揮最佳效能，建議在編譯期間使用下列最佳化標記：

Google Tensor 的編譯旗標

套用特定設定旗標，調整效能和資源用量，進而最佳化模型編譯程序。您可以在 LiteRT Python 工作流程中使用這些參數，調整 PyTorch 和 TFLite 模型的編譯行為。

檢舉 規定 說明 預設值
google_tensor_truncation_type 選用 設定浮點運算的目標資料型別。
支援的值：auto (預設)、bfloat16、half、no_truncation		 自動
google_tensor_sharding_intensity 選用 控制模型分割程度，以利平行處理。
選項：最少、中等、大量、最多。		 極簡
google_tensor_int64_to_int32 選用 設為 true 可讓編譯器將 64 位元整數轉換為 32 位元整數，這對某些模型來說可能很有必要。 否
google_tensor_enable_large_model_support 選用 如果模型大於 2 GB，請設為 true。 否
google_tensor_enable_4bit_compilation 選用 可編譯含有 4 位元參數的捲積運算模型。 否
google_tensor_extra_options_path 選用 以二進位檔案形式傳遞額外的編譯器選項。 "" (空白)

您可以搭配 LiteRT Python 流程使用這些旗標，如下列程式碼片段所示：

  • 使用 ai_edge_torch 編譯 pytorch 模型時

    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))

  • 編譯 tflite 模型時

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

使用範例

在以下範例中，我們使用 google_tensor_truncation_type="half" 標記：

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

詳情請參閱 LiteRT AOT Colab。

為 Google Tensor 編譯語言模型

如要編譯 Google Tensor 的語言模型，請按照「NPU AOT 編譯」中的操作說明進行。

如要匯出 Google Tensor TPU 的 LLM，請參閱範例，瞭解 NPU 編譯所需的額外標記。

範例：

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