Compilare i modelli con l'SDK Google Tensor

L'SDK Google Tensor ti consente di ottimizzare ed eseguire modelli di machine learning (ML) direttamente sulla TPU dedicata di Google Tensor. Oltre ai modelli ML standard, puoi compilare e accelerare i modelli linguistici all'interno del workflow LiteRT.

Per alcuni modelli preottimizzati, forniamo anche un file binario di opzioni del compilatore aggiuntive che puoi includere nel tuo flusso di lavoro. Per ottenere prestazioni ottimali su Google Tensor, ti consigliamo di utilizzare i seguenti flag di ottimizzazione durante la compilazione:

Flag di compilazione per Google Tensor

Ottimizza il processo di compilazione del modello applicando flag di configurazione specifici per personalizzare le prestazioni e l'utilizzo delle risorse. Puoi utilizzare questi parametri nel flusso di lavoro Python LiteRT per modificare il comportamento di compilazione per i modelli PyTorch e TFLite.

Flag Requisito Descrizione Valore predefinito
google_tensor_truncation_type Facoltativo Imposta il tipo di dati di destinazione per le operazioni in virgola mobile.
Valori supportati: auto (impostazione predefinita), bfloat16, half, no_truncation		 Automatica
google_tensor_sharding_intensity Facoltativo Controlla la suddivisione del modello per l'elaborazione parallela.
Opzioni:minima, moderata, estesa, massima.		 minimal
google_tensor_int64_to_int32 Facoltativo Imposta su true per consentire al compilatore di convertire gli interi a 64 bit in interi a 32 bit, il che potrebbe essere necessario per alcuni modelli. Falso
google_tensor_enable_large_model_support Facoltativo Imposta su true se il modello è più grande di 2 GB. Falso
google_tensor_enable_4bit_compilation Facoltativo Consente la compilazione di modelli con operazioni di convoluzione contenenti parametri a 4 bit. Falso
google_tensor_extra_options_path Facoltativo Passa opzioni del compilatore aggiuntive come file binario. "" (vuoto)

Puoi utilizzare questi flag con il flusso Python LiteRT come mostrato nei seguenti snippet di codice:

  • Quando compili un modello PyTorch utilizzando 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))

  • Quando compili un modello TFLite

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

Esempio di utilizzo

Nell'esempio seguente viene utilizzato il flag 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"
  )

Per ulteriori informazioni, consulta LiteRT AOT Colab.

Compilare modelli linguistici per Google Tensor

Per compilare i modelli linguistici per Google Tensor, segui le istruzioni riportate in Compilazione AOT della NPU.

Per esportare LLM per le TPU Google Tensor, segui l'esempio per i flag aggiuntivi richiesti per la compilazione della NPU.

Esempio:

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