Kompilowanie modeli za pomocą pakietu Google Tensor SDK

Pakiet Google Tensor SDK umożliwia optymalizowanie i uruchamianie modeli uczenia maszynowego bezpośrednio na dedykowanym procesorze TPU Google Tensor. Oprócz standardowych modeli uczenia maszynowego możesz kompilować i przyspieszać modele językowe w ramach przepływu pracy LiteRT.

W przypadku niektórych wstępnie zoptymalizowanych modeli udostępniamy też dodatkowy plik binarny opcji kompilatora, który możesz uwzględnić w swoim przepływie pracy. Aby uzyskać optymalną wydajność na procesorze Google Tensor, podczas kompilacji zalecamy używanie tych flag optymalizacji:

Flagi kompilacji dla procesora Google Tensor

Zoptymalizuj proces kompilacji modelu, stosując określone flagi konfiguracji, aby dostosować wydajność i wykorzystanie zasobów. Możesz używać tych parametrów w przepływie pracy LiteRT w Pythonie, aby dostosować zachowanie kompilacji zarówno w przypadku modeli PyTorch, jak i TFLite.

Flaga Wymaganie Opis Wartość domyślna
google_tensor_truncation_type Opcjonalny Ustawia docelowy typ danych dla operacji zmiennoprzecinkowych.
Obsługiwane wartości: auto (domyślna), bfloat16, half, no_truncation		 auto
google_tensor_sharding_intensity Opcjonalny Określa, jak agresywnie model jest dzielony na potrzeby przetwarzania równoległego.
Opcje: minimal, moderate, extensive, maximum.		 minimal
google_tensor_int64_to_int32 Opcjonalny Ustaw wartość true, aby umożliwić kompilatorowi konwertowanie liczb całkowitych 64-bitowych na 32-bitowe, co może być konieczne w przypadku niektórych modeli. Fałsz
google_tensor_enable_large_model_support Opcjonalny Ustaw wartość true, jeśli model jest większy niż 2 GB. Fałsz
google_tensor_enable_4bit_compilation Opcjonalny Umożliwia kompilowanie modeli z operacjami konwolucji zawierającymi parametry 4-bitowe. Fałsz
google_tensor_extra_options_path Opcjonalny Przekazuje dodatkowe opcje kompilatora jako plik binarny. "" (pusty)

Możesz używać tych flag w przepływie LiteRT w Pythonie, jak pokazano w tych fragmentach kodu:

  • Podczas kompilowania modelu PyTorch za pomocą 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))

  • Podczas kompilowania modelu TFLite

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

Przykład zastosowania

W tym przykładzie użyto flagi 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"
  )

Więcej informacji znajdziesz w Colab LiteRT AOT.

Kompilowanie modeli językowych dla procesora Google Tensor

Aby skompilować modele językowe dla procesora Google Tensor, postępuj zgodnie z instrukcjami w sekcji Kompilacja NPU AOT.

Aby wyeksportować modele LLM na procesory TPU Google Tensor, postępuj zgodnie z przykładem dotyczącym dodatkowych flag wymaganych do kompilacji NPU.

Przykład:

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