Compilar modelos com o SDK do Google Tensor

Com o SDK do Google Tensor, é possível otimizar e executar modelos de machine learning (ML) diretamente na TPU dedicada do Google Tensor. Além dos modelos padrão de ML, é possível compilar e acelerar modelos de linguagem no fluxo de trabalho do LiteRT.

Para alguns modelos pré-otimizados, também fornecemos um arquivo binário extra de opções de compilador que pode ser incluído no seu fluxo de trabalho. Para alcançar a performance ideal no Google Tensor, recomendamos usar as seguintes flags de otimização durante a compilação:

Flags de compilação para o Google Tensor

Otimize o processo de compilação do modelo aplicando flags de configuração específicas para adaptar o desempenho e o uso de recursos. Você pode usar esses parâmetros no fluxo de trabalho do Python LiteRT para ajustar o comportamento de compilação dos modelos do PyTorch e do TFLite.

Sinalização Requisito Descrição Valor padrão
google_tensor_truncation_type Opcional Define o tipo de dados de destino para operações de usar pontos flutuantes.
Valores aceitos:auto (padrão), bfloat16, half, no_truncation		 automático
google_tensor_sharding_intensity Opcional Controla a intensidade com que o modelo é dividido para processamento paralelo.
Opções:mínima, moderada, extensa, máxima.		 minimalista
google_tensor_int64_to_int32 Opcional Defina como "true" para permitir que o compilador converta números inteiros de 64 bits em números inteiros de 32 bits, o que pode ser necessário para alguns modelos. Falso
google_tensor_enable_large_model_support Opcional Defina como "true" se o modelo tiver mais de 2 GB. Falso
google_tensor_enable_4bit_compilation Opcional Permite a compilação de modelos com operações de convolução que contêm parâmetros de 4 bits. Falso
google_tensor_extra_options_path Opcional Transmite opções extras do compilador como um arquivo binário. "" (vazio)

É possível usar essas flags com o fluxo do Python LiteRT, conforme mostrado nos snippets de código a seguir:

  • Ao compilar um modelo do PyTorch usando 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))

  • Ao compilar um modelo tflite

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

Exemplo de uso

No exemplo a seguir, a flag google_tensor_truncation_type="half" é usada:

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

Consulte LiteRT AOT Colab para mais informações.

Compilar modelos de linguagem para o Google Tensor

Para compilar modelos de linguagem para o Google Tensor, siga as instruções em Compilação AOT da NPU.

Para exportar LLMs para TPUs do Google Tensor, siga o exemplo das flags adicionais necessárias para a compilação da NPU.

Exemplo:

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