Mengompilasi model dengan Google Tensor SDK

Google Tensor SDK memungkinkan Anda mengoptimalkan dan menjalankan model machine learning (ML) langsung di TPU khusus Google Tensor. Selain model ML standar, Anda dapat mengompilasi dan mempercepat model bahasa dalam alur kerja LiteRT.

Untuk model pra-dioptimalkan tertentu, kami juga menyediakan file biner opsi compiler tambahan yang dapat Anda sertakan dalam alur kerja Anda. Untuk mencapai performa optimal di Google Tensor, sebaiknya gunakan tanda pengoptimalan berikut selama kompilasi:

Flag kompilasi untuk Google Tensor

Optimalkan proses kompilasi model Anda dengan menerapkan tanda konfigurasi tertentu untuk menyesuaikan performa dan penggunaan resource. Anda dapat menggunakan parameter ini dalam alur kerja Python LiteRT untuk menyesuaikan perilaku kompilasi untuk model PyTorch dan TFLite.

Bendera Persyaratan Deskripsi Nilai Default
google_tensor_truncation_type Opsional Menetapkan jenis data target untuk operasi floating point.
Nilai yang didukung: otomatis (default), bfloat16, setengah, tanpa_pemotongan
otomatis
google_tensor_sharding_intensity Opsional Mengontrol seberapa agresif model dibagi untuk pemrosesan paralel.
Opsi: minimal, sedang, ekstensif, maksimum.
minimal
google_tensor_int64_to_int32 Opsional Setel ke benar (true) agar compiler dapat mengonversi bilangan bulat 64-bit menjadi bilangan bulat 32-bit, yang mungkin diperlukan untuk beberapa model. Salah
google_tensor_enable_large_model_support Opsional Tetapkan ke benar (true) jika model Anda lebih besar dari 2 GB. Salah
google_tensor_enable_4bit_compilation Opsional Memungkinkan kompilasi model dengan operasi konvolusi yang berisi parameter 4 bit. Salah
google_tensor_extra_options_path Opsional Meneruskan opsi compiler tambahan sebagai file biner. "" (Kosong)

Anda dapat menggunakan tanda ini dengan alur Python LiteRT seperti yang ditunjukkan dalam cuplikan kode berikut:

  • Saat mengompilasi model pytorch menggunakan 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))
    
  • Saat mengompilasi model tflite

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

Contoh penggunaan

Dalam contoh berikut, flag google_tensor_truncation_type="half" digunakan:

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

Lihat LiteRT AOT Colab untuk mengetahui informasi selengkapnya.

Mengompilasi model bahasa untuk Google Tensor

Untuk mengompilasi model bahasa untuk Google Tensor, ikuti petunjuk di bagian Kompilasi AOT NPU.

Untuk mengekspor LLM bagi TPU Tensor Google, ikuti contoh untuk flag tambahan yang diperlukan untuk kompilasi NPU.

Contoh:

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