Compiler des modèles avec le SDK Google Tensor

Le SDK Google Tensor vous permet d'optimiser et d'exécuter des modèles de machine learning (ML) directement sur le TPU dédié de Google Tensor. En plus des modèles de ML standards, vous pouvez compiler et accélérer des modèles de langage dans votre workflow LiteRT.

Pour certains modèles pré-optimisés, nous fournissons également un fichier binaire d'options de compilation supplémentaires que vous pouvez inclure dans votre workflow. Pour obtenir des performances optimales sur Google Tensor, nous vous recommandons d'utiliser les options d'optimisation suivantes lors de la compilation :

Options de compilation pour Google Tensor

Optimisez le processus de compilation de votre modèle en appliquant des options de configuration spécifiques pour adapter les performances et l'utilisation des ressources. Vous pouvez utiliser ces paramètres dans votre workflow Python LiteRT pour ajuster le comportement de compilation des modèles PyTorch et TFLite.

Option Exigence Description Valeur par défaut
google_tensor_truncation_type Facultatif Définit le type de données cible pour les opérations à virgule flottante.
Valeurs acceptées : auto (par défaut), bfloat16, half, no_truncation		 auto
google_tensor_sharding_intensity Facultatif Contrôle l'agressivité avec laquelle le modèle est divisé pour le traitement parallèle.
Options : minimal, moderate, extensive, maximum.		 minimal
google_tensor_int64_to_int32 Facultatif Définissez la valeur sur "true" pour autoriser le compilateur à convertir les entiers 64 bits en entiers 32 bits, ce qui peut être nécessaire pour certains modèles. Faux
google_tensor_enable_large_model_support Facultatif Définissez la valeur sur "true" si votre modèle est supérieur à 2 Go. Faux
google_tensor_enable_4bit_compilation Facultatif Permet la compilation de modèles avec des opérations de convolution contenant des paramètres de 4 bits. Faux
google_tensor_extra_options_path Facultatif Transmet des options de compilation supplémentaires sous forme de fichier binaire. "" (vide)

Vous pouvez utiliser ces options avec le flux Python LiteRT, comme illustré dans les extraits de code suivants :

  • Lors de la compilation d'un modèle PyTorch à l'aide de 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))

  • Lors de la compilation d'un modèle TFLite

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

Exemple d'utilisation

Dans l'exemple suivant, l'option google_tensor_truncation_type="half" est utilisée :

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

Pour en savoir plus, consultez LiteRT AOT Colab.

Compiler des modèles de langage pour Google Tensor

Pour compiler des modèles de langage pour Google Tensor, suivez les instructions de Compilation AOT NPU.

Pour exporter des LLM pour les TPU Google Tensor, suivez l'exemple des options supplémentaires requises pour la compilation NPU.

Exemple :

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