Introdução ao conversor de inferência do Cloud TPU v5e

Introdução

O conversor de inferência do Cloud TPU prepara e otimiza um modelo do TensorFlow 2 (TF2) para inferência de TPU. O conversor é executado em um shell de VM local ou de TPU. O shell da VM da TPU é recomendado porque vem pré-instalado com as ferramentas de linha de comando necessárias para o conversor. Ele usa um SavedModel exportado e executa as seguintes etapas:

1. Conversão de TPU: adiciona TPUPartitionedCall e outras operações de TPU ao modelo para torná-lo veiculável na TPU. Por padrão, um modelo exportado para inferência não tem essas operações e não pode ser exibido na TPU, mesmo que tenha sido treinado nela.
2. Lotes: adiciona operações em lote ao modelo para ativar o agrupamento em gráfico e melhorar a capacidade.
3. Conversão BFloat16: ela converte o formato de dados do modelo de float32 para bfloat16, melhorando o desempenho computacional e reduzindo o uso de memória de alta largura de banda (HBM, na sigla em inglês) na TPU.
4. Otimização de formas de E/S: otimiza as formas do tensor para dados transferidos entre a CPU e a TPU para melhorar a utilização da largura de banda.

Ao exportar um modelo, os usuários criam aliases de função para todas as funções que querem executar na TPU. Elas passam essas funções ao conversor, e o conversor as coloca na TPU e as otimiza.

O conversor de inferência do Cloud TPU está disponível como uma imagem do Docker que pode ser executada em qualquer ambiente com o Docker instalado.

Tempo estimado para concluir as etapas mostradas acima: cerca de 20 minutos a 30 minutos

Pré-requisitos

1. O modelo precisa ser TF2 e exportado no formato SavedModel.
2. O modelo precisa ter um alias de função para a função da TPU. Consulte o exemplo de código para saber como fazer isso. Os exemplos a seguir usam tpu_func como o alias da função da TPU.
3. Verifique se a CPU da sua máquina é compatível com as instruções Advanced Vector eXtensions (AVX), já que a biblioteca do TensorFlow (a dependência do Cloud TPU Inference Converter) é compilada para usar as instruções do AVX. A maioria das CPUs é compatível com o AVX.
   1. Você pode executar lscpu | grep avx para verificar se o conjunto de instruções do AVX tem suporte.

Antes de começar

Antes de iniciar a configuração, faça o seguinte:

• Criar um novo projeto: no console do Google Cloud, na página do seletor de projetos, escolha ou crie um projeto do Cloud.

• Configurar uma VM de TPU: crie uma nova VM de TPU usando o console do Google Cloud ou o gcloud ou use uma VM de TPU atual para executar a inferência com o modelo convertido na VM da TPU.

  • Verifique se a imagem da VM da TPU é baseada no TensorFlow. Por exemplo, --version=tpu-vm-tf-2.11.0.
  • O modelo convertido será carregado e veiculado nessa VM da TPU.

• Verifique se você tem as ferramentas de linha de comando necessárias para usar o conversor de inferência do Cloud TPU. É possível instalar o SDK Google Cloud e o Docker localmente ou usar uma VM de TPU que tenha esse software instalado por padrão. Você usa essas ferramentas para interagir com a imagem do conversor.

  Conecte-se à instância com SSH usando o seguinte comando:

  gcloud compute tpus tpu-vm ssh ${tpu-name} --zone ${zone} --project ${project-id}
  

Configuração do ambiente

Configure o ambiente usando o shell da VM da TPU ou o shell local.

Shell da VM da TPU

• No shell da VM da TPU, execute os seguintes comandos para permitir o uso do Docker não raiz:

  sudo usermod -a -G docker ${USER}
  newgrp docker
  

• Inicialize os auxiliares de credenciais do Docker:

  gcloud auth configure-docker \
    us-docker.pkg.dev
  

Shell local

No shell local, configure o ambiente seguindo estas etapas:

• Instale o SDK Cloud, que inclui a ferramenta de linha de comando gcloud.

• Instale o Docker:

• Permitir o uso não raiz do Docker:

  sudo usermod -a -G docker ${USER}
  newgrp docker
  

• Faça login no ambiente:

  gcloud auth login
  

• Inicialize os auxiliares de credenciais do Docker:

  gcloud auth configure-docker \
      us-docker.pkg.dev
  

• Extraia a imagem do Docker do conversor de inferência:

    CONVERTER_IMAGE=us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
    docker pull ${CONVERTER_IMAGE}
    

Imagem do conversor

A imagem é usada para fazer conversões de modelo únicas. Defina os caminhos do modelo e ajuste as opções do conversor de acordo com suas necessidades. A seção Exemplos de uso fornece vários casos de uso comuns.

Observação :no código a seguir, é possível usar um bucket do Cloud Storage para input_model_dir e output_model_dir sem especificar mount no comando docker.

docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      function_alias: "tpu_func"
    }
    batch_options {
      num_batch_threads: 2
      max_batch_size: 8
      batch_timeout_micros: 5000
      allowed_batch_sizes: 2
      allowed_batch_sizes: 4
      allowed_batch_sizes: 8
      max_enqueued_batches: 10
    }
'

Inferência com o modelo convertido na VM da TPU

# Initialize the TPU
resolver = tf.distribute.cluster_resolver.TPUClusterResolver("local")
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

# Load the model
model = tf.saved_model.load(${CONVERTED_MODEL_PATH})

# Find the signature function for serving
serving_signature = 'serving_default' # Change the serving signature if needed
serving_fn = model.signatures[serving_signature]
# Run the inference using requests.
results = serving_fn(**inputs)
logging.info("Serving results: %s", str(results))

Exemplos de uso

Adicionar um alias de função para a função da TPU

1. Encontre ou crie uma função no modelo que encapsule tudo o que você quer executar na TPU. Se @tf.function não existir, adicione-o.
2. Ao salvar o modelo, forneça SaveOptions como abaixo para dar a model.tpu_func um alias func_on_tpu.
3. É possível transmitir esse alias de função para o conversor.

class ToyModel(tf.keras.Model):
  @tf.function(
      input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
  def tpu_func(self, x):
    return x * 1.0

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'func_on_tpu': model.tpu_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

Converter um modelo com várias funções de TPU

É possível colocar várias funções na TPU. Basta criar vários aliases de função e transmiti-los em converter_options_string para o conversor.

tpu_functions {
  function_alias: "tpu_func_1"
}
tpu_functions {
  function_alias: "tpu_func_2"
}

Quantização.

Importante: esse recurso é experimental e pode estar sujeito a mudanças sem aviso prévio. Os usuários precisam avaliar cuidadosamente o desempenho do modelo quantizado antes de usá-lo na produção. O suporte da TPU v5e para a quantização é fornecido no TensorFlow.

A quantização é uma técnica que reduz a precisão dos números usados para representar os parâmetros de um modelo. Isso resulta em um tamanho de modelo menor e computação mais rápida. Um modelo quantizado fornece ganhos em capacidade de inferência, bem como menor uso de memória e tamanho de armazenamento, ao custo de pequenas quedas na precisão.

O novo recurso de quantização pós-treinamento do TensorFlow, voltado à TPU, foi desenvolvido a partir de um recurso semelhante existente no TensorFlow Lite, usado para segmentar dispositivos móveis e de borda. Para saber mais sobre a quantização em geral, consulte o documento do TensorFlow Lite.

Conceitos de quantização

Esta seção define conceitos especificamente relacionados à quantização com o conversor de inferência.

Os conceitos relacionados a outras configurações de TPU (por exemplo, fatias, hosts, chips e TensorCores) são descritos na página Arquitetura do sistema de TPU.

• Quantização pós-treinamento (PTQ, na sigla em inglês): a PTQ é uma técnica que reduz o tamanho e a complexidade computacional de um modelo de rede neural sem afetar significativamente a acurácia dele. A PTQ converte as ativações e os pesos de ponto flutuante de um modelo treinado em números inteiros de precisão mais baixa, como números inteiros de 8 ou 16 bits. Isso pode causar uma redução significativa no tamanho do modelo e na latência de inferência, além de gerar uma pequena perda na acurácia.

• Calibração: a etapa de calibragem para a quantização é o processo de coleta de estatísticas no intervalo de valores que os pesos e as ativações de um modelo de rede neural levam. Essas informações são usadas para determinar os parâmetros de quantização do modelo, que são os valores usados para converter as ativações e os pesos de ponto flutuante em números inteiros.

• Conjunto de dados representativo: um conjunto de dados representativo para quantização é um conjunto de dados pequeno que representa os dados de entrada reais do modelo. Ela é usada durante a etapa de calibração da quantização para coletar estatísticas sobre o intervalo de valores que os pesos e as ativações do modelo vão realizar. O conjunto de dados representativo precisa satisfazer as seguintes propriedades:

  • Ele precisa representar corretamente as entradas reais do modelo durante a inferência. Isso significa que ele precisa abranger o intervalo de valores que o modelo provavelmente verá no mundo real.
  • Eles devem fluir coletivamente por cada ramificação de condicionais (como tf.cond), se houver. Isso é importante porque o processo de quantização precisa processar todas as entradas possíveis no modelo, mesmo que elas não estejam explicitamente representadas no conjunto de dados representativo.
  • Ele precisa ser grande o suficiente para coletar estatísticas suficientes e reduzir erros. Como regra geral, é recomendável usar mais de 200 amostras representativas.

  O conjunto de dados representativo pode ser um subconjunto do conjunto de treinamento ou um conjunto separado, projetado especificamente para representar as entradas reais do modelo. A escolha de qual conjunto de dados usar depende do aplicativo específico.

• Quantização de intervalo estático (SRQ, na sigla em inglês): o SRQ determina o intervalo de valores para os pesos e as ativações de um modelo de rede neural uma vez, durante a etapa de calibragem. Isso significa que o mesmo intervalo de valores é usado para todas as entradas do modelo. Isso pode ser menos preciso do que a quantização de intervalo dinâmico, especialmente para modelos com uma ampla gama de valores de entrada. No entanto, a quantização de intervalo estático requer menos computação em tempo de execução do que a quantização de intervalo dinâmico.

• Quantização de intervalo dinâmico (DRQ, na sigla em inglês): a DRQ determina o intervalo de valores para os pesos e as ativações de um modelo de rede neural para cada entrada. Isso permite que o modelo se adapte ao intervalo de valores dos dados de entrada, o que pode melhorar a acurácia. No entanto, a quantização de intervalo dinâmico requer mais computação no ambiente de execução do que a quantização de intervalo estático.

  Recurso	Quantização de intervalo estático	Quantização de intervalo dinâmico
  Intervalo de valores	Determinado uma vez, durante a calibragem	Determinado para cada entrada
  Acurácia	Pode ser menos preciso, especialmente para modelos com uma ampla gama de valores de entrada.	Podem ser mais precisos, especialmente para modelos com uma ampla gama de valores de entrada.
  Complexidade	Mais simples	Mais complexo
  Computação no ambiente de execução	Menos computação	Mais computação

• Quantização somente de peso: esse tipo é um tipo de quantização que só quantifica os pesos de um modelo de rede neural, deixando as ativações em ponto flutuante. Essa pode ser uma boa opção para modelos sensíveis à acurácia, porque ajuda a preservar a acurácia.

Como usar a quantização

A quantização pode ser aplicada com a configuração e a definição de QuantizationOptions às opções do conversor. As principais opções são:

• Tags: coleção de tags que identificam o MetaGraphDef no SavedModel para quantização. Não é necessário especificar se você tem apenas um MetaGraphDef.
• signature_keys: sequência de chaves que identifica SignatureDef que contêm entradas e saídas. Se não for especificado, ["serving_default"] será usado.
• quantization_method: método de quantização a ser aplicado. Se não for especificada, a quantização STATIC_RANGE vai ser aplicada.
• op_set: precisa ser mantido como XLA. Essa é a opção padrão no momento, não é necessário especificar.
• Representação_de_conjuntos de dados: especifica o conjunto de dados usado para calibrar os parâmetros de quantização.

Como criar o conjunto de dados representativo

Um conjunto de dados representativo é essencialmente um iterável de amostras. Onde um exemplo é um mapa de: {input_key: input_value}. Exemplo:

representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]

Os conjuntos de dados representativos precisam ser salvos como arquivos TFRecord usando a classe TfRecordRepresentativeDatasetSaver atualmente disponível no pacote tf-nightly pip. Exemplo:

# Assumed tf-nightly installed.
import tensorflow as tf
representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]
tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
       path_map={'serving_default': '/tmp/representative_dataset_path'}
    ).save({'serving_default': representative_dataset})

Examples

O exemplo a seguir quantiza o modelo com a chave de assinatura de serving_default e o alias de função de tpu_func:

Observação: no código a seguir, defina CONVERTER_IMAGE como us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:nightly para usar o build noturno.

docker run \
  --mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
  --mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
  ${CONVERTER_IMAGE} \
  --input_model_dir=/tmp/input \
  --output_model_dir=/tmp/output \
  --converter_options_string=' \
    tpu_functions { \
      function_alias: "tpu_func" \
    } \
    external_feature_configs { \
      quantization_options { \
        signature_keys: "serving_default" \
        representative_datasets: { \
          key: "serving_default" \
          value: { \
            tfrecord_file_path: "${TF_RECORD_FILE}" \
          } \
        } \
      } \
    } '

Adicionar lotes

O conversor pode ser usado para adicionar lotes a um modelo. Para ver uma descrição das opções de lote que podem ser ajustadas, consulte Definição das opções de lote.

Por padrão, o conversor agrupa todas as funções de TPU no modelo. Ela também pode agrupar assinaturas e funções fornecidas pelo usuário para melhorar ainda mais o desempenho. Qualquer função da TPU, função ou assinatura fornecida pelo usuário que é agrupada em lote precisa atender aos requisitos rígidos de forma da operação de lote.

O conversor também pode atualizar as opções de lote atuais. Confira abaixo um exemplo de como adicionar lotes a um modelo. Para mais informações sobre lotes, consulte Análise detalhada de lotes.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

Desativar otimizações de forma de bfloat16 e E/S

BFloat16 e Otimizações de forma de E/S são ativadas por padrão. Se não funcionarem bem com o modelo, eles poderão ser desativados.

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

Relatório de conversões

É possível encontrar esse relatório de conversão no registro depois de executar o conversor de inferência. Veja um exemplo abaixo.

-------- Conversion Report --------
TPU cost of the model: 96.67% (2034/2104)
CPU cost of the model:  3.33% (70/2104)

Cost breakdown
================================
%         Cost    Name
--------------------------------
3.33      70      [CPU cost]
48.34     1017    tpu_func_1
48.34     1017    tpu_func_2
--------------------------------

Neste relatório, estimamos o custo computacional do modelo de saída na CPU e na TPU e detalhamos o custo da TPU em cada função, o que deve refletir sua seleção das funções de TPU nas opções de conversor.

Se você quiser usar melhor a TPU, teste a estrutura do modelo e ajuste as opções de conversor.

Observação :no relatório de conversão atual, todas as operações nas funções da TPU são contabilizadas como custo de TPU.

Perguntas frequentes

Que funções devo colocar na TPU?

É melhor colocar o máximo possível do modelo na TPU. A grande maioria das operações é executada mais rapidamente na TPU.

Caso o modelo não tenha operações, strings ou tensores esparsos incompatíveis com TPU, a melhor estratégia costuma ser colocar o modelo inteiro na TPU. E você pode fazer isso encontrando ou criando uma função que encapsule todo o modelo, criando um alias de função para ela e transmitindo isso para o conversor.

Caso o modelo contenha peças que não funcionam na TPU (por exemplo, operações, strings ou tensores esparsos incompatíveis com TPU), a escolha de funções de TPU vai depender de onde está a parte incompatível.

• Se estiver no início ou no fim do modelo, é possível refatorá-lo para mantê-lo na CPU. Exemplos são os estágios de pré e pós-processamento de strings. Para mais informações sobre como mover o código para a CPU, consulte "Como mover uma parte do modelo para a CPU?". Ele mostra uma maneira típica de refatorar o modelo.
• Se ele estiver no meio do modelo, é melhor dividi-lo em três partes, conter todas as operações incompatíveis com TPU na parte do meio e executar na CPU.
• Se ele for um tensor esparso, considere chamar tf.sparse.to_dense na CPU e transmitir o tensor denso resultante para a parte da TPU do modelo.

Outro fator a ser considerado é o uso de HBM. As tabelas de embedding podem usar muito HBM. Se elas ultrapassarem a limitação de hardware da TPU, precisarão ser colocadas na CPU, junto com as operações de pesquisa.

Sempre que possível, só deve haver uma função de TPU em uma assinatura. Se a estrutura do modelo exigir a chamada de várias funções de TPU por solicitação de inferência recebida, considere a latência adicional do envio de tensores entre a CPU e a TPU.

Uma boa maneira de avaliar a seleção das funções da TPU é consultar o Relatório de conversões. Ela mostra a porcentagem de computação colocada na TPU e o detalhamento do custo de cada função da TPU.

Como mover uma parte do modelo para a CPU?

Se o modelo tiver partes que não podem ser exibidas na TPU, será necessário refatorar o modelo para movê-las para a CPU. Aqui está um exemplo de brinquedo. O modelo é um modelo de linguagem com uma etapa de pré-processamento. O código para definições e funções de camada é omitido para simplificar.

class LanguageModel(tf.keras.Model):
  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.bert_layer(word_ids)

Esse modelo não pode ser exibido diretamente na TPU por dois motivos. Primeiro, o parâmetro é uma string. Segundo, a função preprocess pode conter muitas operações de string. Ambas não são compatíveis com TPU.

Para refatorar esse modelo, crie outra função chamada tpu_func para hospedar o bert_layer de uso intensivo de computação. Em seguida, crie um alias de função para tpu_func e transmita-o para o conversor. Dessa forma, tudo em tpu_func será executado na TPU e tudo o que restar em model_func será executado na CPU.

class LanguageModel(tf.keras.Model):
  @tf.function
  def tpu_func(self, word_ids):
    return self.bert_layer(word_ids)

  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.tpu_func(word_ids)

O que devo fazer se o modelo tiver operações, strings ou tensores esparsos incompatíveis com TPU?

A maioria das operações padrão do TensorFlow é compatível com a TPU, mas algumas incluindo strings e tensores esparsos não são compatíveis. O conversor não verifica se há operações incompatíveis com TPU. Portanto, um modelo com essas operações pode passar a conversão. No entanto, ao executá-lo para inferência, ocorrerão erros como abaixo.

'tf.StringToNumber' op isn't compilable for TPU device.

Se o modelo tiver operações incompatíveis com TPU, elas precisarão ser colocadas fora da função da TPU. Além disso, a string é um formato de dados incompatível com a TPU. Portanto, as variáveis do tipo string não podem ser colocadas na função da TPU. Além disso, os parâmetros e os valores de retorno da função da TPU não podem ser digitados como string. Da mesma forma, evite colocar tensores esparsos na função da TPU, incluindo nos parâmetros e valores de retorno dela.

Geralmente, não é difícil refatorar a parte incompatível do modelo e movê-la para a CPU. Veja um exemplo.

Como dar suporte a operações personalizadas no modelo?

Se operações personalizadas forem usadas no modelo, o conversor poderá não reconhecê-las e não converter o modelo. Isso ocorre porque a biblioteca de operações da operação personalizada, que contém a definição completa, não está vinculada ao conversor.

Como o código do conversor ainda não é aberto, não é possível criar o conversor com uma operação personalizada.

O que devo fazer se eu tiver um modelo do TensorFlow 1?

O conversor não oferece suporte a modelos do TensorFlow 1. Os modelos do TensorFlow 1 precisam ser migrados para o TensorFlow 2.

Preciso ativar a ponte MLIR ao executar meu modelo?

A maioria dos modelos convertidos pode ser executada com a ponte MLIR TF2XLA mais recente ou a ponte TF2XLA original.

Como converter um modelo que já foi exportado sem um alias de função?

Se um modelo foi exportado sem um alias de função, a maneira mais fácil é exportá-lo novamente e criar um alias de função. Mesmo que a reexportação não seja uma opção, ainda será possível converter o modelo fornecendo um concrete_function_name. No entanto, identificar o concrete_function_name correto exige algum trabalho de detetive.

Os aliases de função são um mapeamento de uma string definida pelo usuário para um nome de função concreto. Eles facilitam a referência a uma função específica no modelo. O conversor aceita aliases de função e nomes de funções concretos brutos.

Os nomes concretos de funções podem ser encontrados examinando a saved_model.pb.

O exemplo a seguir mostra como colocar na TPU uma função concreta chamada __inference_serve_24.

sudo docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      concrete_function_name: "__inference_serve_24"
    }'

Como resolver um erro de restrição da constante de tempo de compilação?

Para treinamento e inferência, o XLA exige que as entradas de determinadas operações tenham uma forma conhecida no tempo de compilação da TPU. Isso significa que, quando o XLA compila a parte da TPU do programa, as entradas para essas operações precisam ter um formato conhecido estaticamente.

Há duas maneiras de resolver esse problema.

• A melhor opção é atualizar as entradas da operação para que tenham um formato estaticamente conhecido no momento em que o XLA compilar o programa de TPU. Essa compilação acontece logo antes da parte da TPU do modelo ser executada. Isso significa que a forma precisará ser estaticamente conhecida no momento em que o TpuFunction estiver prestes a ser executado.
• Outra opção é modificar a TpuFunction para não incluir mais a operação problemática.

Por que estou recebendo um erro de forma de lote?

O lote tem requisitos de forma rígidos que permitem que as solicitações recebidas sejam agrupadas na 0a dimensão (ou seja, a dimensão de lote). Esses requisitos de forma vêm da operação de lote do TensorFlow e não podem ser atenuados.

O não cumprimento desses requisitos resultará em erros como:

1. Os tensores de entrada em lote precisam ter pelo menos uma dimensão.
2. As dimensões das entradas precisam ser iguais.
3. Os tensores de entrada em lote fornecidos em uma determinada invocação de operação precisam ter um tamanho de dimensão 0 igual.
4. A 0a dimensão do tensor de saída em lote não é igual à soma dos tamanhos de 0a dimensão dos tensores de entrada.

Para atender a esses requisitos, considere fornecer uma função ou assinatura diferente para o lote. Também pode ser necessário modificar as funções existentes para atender a esses requisitos.

Se uma função estiver sendo loteada, verifique se as formas de input_signature de @tf.function todas têm "None" na dimensão 0. Se uma assinatura estiver sendo agrupada em lote, verifique se todas as entradas têm -1 na dimensão 0.

Para uma explicação completa sobre por que esses erros estão acontecendo e como resolvê-los, consulte Análise detalhada em lote.

Problemas conhecidos

A função da TPU não pode chamar indiretamente outra função da TPU

Embora o conversor possa lidar com a maioria dos cenários de chamada de função em todo o limite CPU-TPU, há um caso extremo raro em que ele falharia. É quando uma função da TPU chama indiretamente outra função da TPU.

Isso ocorre porque o conversor modifica o autor da chamada direta de uma função da TPU, deixando de chamar a própria função da TPU para chamar um stub de chamada da TPU. O stub de chamada contém operações que só podem funcionar na CPU. Quando uma função da TPU chama qualquer função que chame o autor da chamada direto, essas operações da CPU podem ser executadas na TPU para execução, o que gera erros de kernel ausentes. Esse caso é diferente de uma função de TPU que chama diretamente outra função de TPU. Nesse caso, o conversor não modifica nenhuma função para chamar o stub de chamada para que ele possa funcionar.

No Converter, implementamos a detecção desse cenário. Se você vir o erro a seguir, isso significa que seu modelo atingiu esse caso extremo:

Unable to place both "__inference_tpu_func_2_46" and "__inference_tpu_func_4_68"
on the TPU because "__inference_tpu_func_2_46" indirectly calls
"__inference_tpu_func_4_68". This behavior is unsupported because it can cause
invalid graphs to be generated.

A solução geral é refatorar o modelo para evitar esse cenário de chamada de função. Se você achar isso difícil, entre em contato com a equipe de suporte do Google para mais informações.

Referência

Opções do conversor no formato Protobuf

message ConverterOptions {
  // TPU conversion options.
  repeated TpuFunction tpu_functions = 1;

  // The state of an optimization.
  enum State {
    // When state is set to default, the optimization will perform its
    // default behavior. For some optimizations this is disabled and for others
    // it is enabled. To check a specific optimization, read the optimization's
    // description.
    DEFAULT = 0;
    // Enabled.
    ENABLED = 1;
    // Disabled.
    DISABLED = 2;
  }

  // Batch options to apply to the TPU Subgraph.
  //
  // At the moment, only one batch option is supported. This field will be
  // expanded to support batching on a per function and/or per signature basis.
  //
  //
  // If not specified, no batching will be done.
  repeated BatchOptions batch_options = 100;

  // Global flag to disable all optimizations that are enabled by default.
  // When enabled, all optimizations that run by default are disabled. If a
  // default optimization is explicitly enabled, this flag will have no affect
  // on that optimization.
  //
  // This flag defaults to false.
  bool disable_default_optimizations = 202;

  // If enabled, apply an optimization that reshapes the tensors going into
  // and out of the TPU. This reshape operation improves performance by reducing
  // the transfer time to and from the TPU.
  //
  // This optimization is incompatible with input_shape_opt which is disabled.
  // by default. If input_shape_opt is enabled, this option should be
  // disabled.
  //
  // This optimization defaults to enabled.
  State io_shape_optimization = 200;

  // If enabled, apply an optimization that updates float variables and float
  // ops on the TPU to bfloat16. This optimization improves performance and
  // throughtput by reducing HBM usage and taking advantage of TPU support for
  // bfloat16.
  //
  // This optimization may cause a loss of accuracy for some models. If an
  // unacceptable loss of accuracy is detected, disable this optimization.
  //
  // This optimization defaults to enabled.
  State bfloat16_optimization = 201;

  BFloat16OptimizationOptions bfloat16_optimization_options = 203;

  // The settings for XLA sharding. If set, XLA sharding is enabled.
  XlaShardingOptions xla_sharding_options = 204;
}

message TpuFunction {
  // The function(s) that should be placed on the TPU. Only provide a given
  // function once. Duplicates will result in errors. For example, if
  // you provide a specific function using function_alias don't also provide the
  // same function via concrete_function_name or jit_compile_functions.
  oneof name {
    // The name of the function alias associated with the function that
    // should be placed on the TPU. Function aliases are created during model
    // export using the tf.saved_model.SaveOptions.
    //
    // This is a recommended way to specify which function should be placed
    // on the TPU.
    string function_alias = 1;

    // The name of the concrete function that should be placed on the TPU. This
    // is the name of the function as it found in the GraphDef and the
    // FunctionDefLibrary.
    //
    // This is NOT the recommended way to specify which function should be
    // placed on the TPU because concrete function names change every time a
    // model is exported.
    string concrete_function_name = 3;

    // The name of the signature to be placed on the TPU. The user must make
    // sure there is no TPU-incompatible op under the entire signature.
    string signature_name = 5;

    // When jit_compile_functions is set to True, all jit compiled functions
    // are placed on the TPU.
    //
    // To use this option, decorate the relevant function(s) with
    // @tf.function(jit_compile=True), before exporting. Then set this flag to
    // True. The converter will find all functions that were tagged with
    // jit_compile=True and place them on the TPU.
    //
    // When using this option, all other settings for the TpuFunction
    // will apply to all functions tagged with
    // jit_compile=True.
    //
    // This option will place all jit_compile=True functions on the TPU.
    // If only some jit_compile=True functions should be placed on the TPU,
    // use function_alias or concrete_function_name.
    bool jit_compile_functions = 4;
  }

}

message BatchOptions {
  // Number of scheduling threads for processing batches of work. Determines
  // the number of batches processed in parallel. This should be roughly in line
  // with the number of TPU cores available.
  int32 num_batch_threads = 1;

  // The maximum allowed batch size.
  int32 max_batch_size = 2;

  // Maximum number of microseconds to wait before outputting an incomplete
  // batch.
  int32 batch_timeout_micros = 3;

  // Optional list of allowed batch sizes. If left empty,
  // does nothing. Otherwise, supplies a list of batch sizes, causing the op
  // to pad batches up to one of those sizes. The entries must increase
  // monotonically, and the final entry must equal max_batch_size.
  repeated int32 allowed_batch_sizes = 4;

  // Maximum number of batches enqueued for processing before requests are
  // failed fast.
  int32 max_enqueued_batches = 5;

  // If set, disables large batch splitting which is an efficiency improvement
  // on batching to reduce padding inefficiency.
  bool disable_large_batch_splitting = 6;

  // Experimental features of batching. Everything inside is subject to change.
  message Experimental {
    // The component to be batched.
    // 1. Unset if it's for all TPU subgraphs.
    // 2. Set function_alias or concrete_function_name if it's for a function.
    // 3. Set signature_name if it's for a signature.
    oneof batch_component {
      // The function alias associated with the function. Function alias is
      // created during model export using the tf.saved_model.SaveOptions, and is
      // the recommended way to specify functions.
      string function_alias = 1;

      // The concreate name of the function. This is the name of the function as
      // it found in the GraphDef and the FunctionDefLibrary. This is NOT the
      // recommended way to specify functions, because concrete function names
      // change every time a model is exported.
      string concrete_function_name = 2;

      // The name of the signature.
      string signature_name = 3;
    }
  }

  Experimental experimental = 7;
}

message BFloat16OptimizationOptions {
  // Indicates where the BFloat16 optimization should be applied.
  enum Scope {
    // The scope currently defaults to TPU.
    DEFAULT = 0;
    // Apply the bfloat16 optimization to TPU computation.
    TPU = 1;
    // Apply the bfloat16 optimization to the entire model including CPU
    // computations.
    ALL = 2;
  }

  // This field indicates where the bfloat16 optimization should be applied.
  //
  // The scope defaults to TPU.
  Scope scope = 1;

  // If set, the normal safety checks are skipped. For example, if the model
  // already contains bfloat16 ops, the bfloat16 optimization will error because
  // pre-existing bfloat16 ops can cause issues with the optimization. By
  // setting this flag, the bfloat16 optimization will skip the check.
  //
  // This is an advanced feature and not recommended for almost all models.
  //
  // This flag is off by default.
  bool skip_safety_checks = 2;

  // Ops that should not be converted to bfloat16.
  // Inputs into these ops will be cast to float32, and outputs from these ops
  // will be cast back to bfloat16.
  repeated string filterlist = 3;
}

message XlaShardingOptions {
  // num_cores_per_replica for TPUReplicateMetadata.
  //
  // This is the number of cores you wish to split your model into using XLA
  // SPMD.
  int32 num_cores_per_replica = 1;

  // (optional) device_assignment for TPUReplicateMetadata.
  //
  // This is in a flattened [x, y, z, core] format (for
  // example, core 1 of the chip
  // located in 2,3,0 will be stored as [2,3,0,1]).
  //
  // If this is not specified, then the device assignments will utilize the same
  // topology as specified in the topology attribute.
  repeated int32 device_assignment = 2;

  // A serialized string of tensorflow.tpu.TopologyProto objects, used for
  // the topology attribute in TPUReplicateMetadata.
  //
  // You must specify the mesh_shape and device_coordinates attributes in
  // the topology object.
  //
  // This option is required for num_cores_per_replica > 1 cases due to
  // ambiguity of num_cores_per_replica, for example,
  // pf_1x2x1 with megacore and df_1x1
  // both have num_cores_per_replica = 2, but topology is (1,2,1,1) for pf and
  // (1,1,1,2) for df.
  // - For pf_1x2x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,2,1,1]
  //   device_coordinates=flatten([0,0,0,0], [0,1,0,0])
  // - For df_1x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,1,1,2]
  //   device_coordinates=flatten([0,0,0,0], [0,0,0,1])
  // - For df_2x2, mesh shape and device_coordinates looks like:
  //   mesh_shape = [2,2,1,2]
  //   device_coordinates=flatten(
  //    [0,0,0,0],[0,0,0,1],[0,1,0,0],[0,1,0,1]
  //    [1,0,0,0],[1,0,0,1],[1,1,0,0],[1,1,0,1])
  bytes topology = 3;
}

Análise detalhada de lotes

Os lotes são usados para melhorar a capacidade e o uso da TPU. Isso permite que várias solicitações sejam processadas ao mesmo tempo. Durante o treinamento, os lotes podem ser feitos usando tf.data. Durante a inferência, isso geralmente é feito com a adição de uma operação no gráfico que agrupa as solicitações recebidas. A operação espera até que tenha solicitações suficientes ou que o tempo limite seja atingido antes de gerar um lote grande com base nas solicitações individuais. Consulte Definição das opções de lote para mais informações sobre as diferentes opções de lote que podem ser ajustadas, incluindo tamanhos e tempos limite.

Por padrão, o conversor insere a operação de lote diretamente antes da computação da TPU. Ela une as funções de TPU fornecidas pelo usuário e qualquer computação de TPU preexistente no modelo com operações de lote. É possível substituir esse comportamento padrão informando ao conversor quais funções e/ou assinaturas precisam ser agrupadas.

O exemplo a seguir mostra como adicionar os lotes padrão.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

Lotes de assinaturas

Observação: esse recurso ainda está em fase experimental. Após o teste, o campo será removido do wrapper experimental.

O lote de assinaturas agrupa todo o modelo, começando pelas entradas da assinatura e indo até as saídas da assinatura. Ao contrário do comportamento padrão de lotes do conversor, o lote de assinaturas agrupa a computação da TPU e da CPU. Isso proporciona um ganho de desempenho de 10% a 20% durante a inferência em alguns modelos.

Como todos os lotes, o lote de assinaturas tem requisitos de forma rígidos. Para ajudar a garantir que esses requisitos de forma sejam atendidos, as entradas de assinatura precisam ter formas que tenham pelo menos duas dimensões. A primeira dimensão é o tamanho do lote e deve ter um tamanho de -1. Por exemplo, (-1, 4), (-1) ou (-1, 128, 4, 10) são formas de entrada válidas. Se isso não for possível, use o comportamento de lote padrão ou o agrupamento de funções.

Para usar o agrupamento de assinaturas, forneça os nomes das assinaturas como signature_name usando o BatchOptions.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    signature_name: "serving_default"
  }
}

Lotes de funções

Observação: esse recurso ainda está em fase experimental. Após o teste, o campo será removido do wrapper experimental.

O lote de funções pode ser usado para informar ao conversor quais funções devem ser agrupadas em lotes. Por padrão, o conversor agrupa todas as funções da TPU em lote. O lote de funções modifica esse comportamento padrão.

O agrupamento de funções pode ser usado para a computação em lote da CPU. Muitos modelos têm uma melhoria de desempenho quando a computação da CPU é feita em lote. A melhor maneira de processar a computação em lote é usando lotes de assinatura. No entanto, isso pode não funcionar em alguns modelos. Nesses casos, é possível usar o lote de funções para agrupar parte da computação da CPU, além do cálculo da TPU. Observe que essa operação não pode ser executada na TPU. Portanto, qualquer função de lote fornecida precisa ser chamada na CPU.

O agrupamento de funções também pode ser usado para atender aos requisitos rígidos de forma impostos pela operação de lote. Nos casos em que as funções da TPU não atendem aos requisitos de forma da operação de lote, o agrupamento de funções pode ser usado para instruir o conversor a agrupar diferentes funções.

Para usar isso, gere um function_alias para a função que será agrupada em lote. Para fazer isso, encontre ou crie uma função no modelo que envolva tudo o que você quer agrupar. Verifique se essa função atende aos requisitos rígidos de forma impostos pela operação de lote. Adicione @tf.function se ela ainda não tiver um. É importante fornecer o input_signature ao @tf.function. A 0a dimensão precisa ser None porque é a dimensão do lote. Portanto, ela não pode ter um tamanho fixo. Por exemplo, [None, 4], [None] ou [None, 128, 4, 10] são formas de entrada válidas. Ao salvar o modelo, forneça SaveOptions como os mostrados abaixo para dar a model.batch_func um alias "batch_func". Em seguida, transmita esse alias de função para o conversor.

class ToyModel(tf.keras.Model):
  @tf.function(input_signature=[tf.TensorSpec(shape=[None, 10],
                                              dtype=tf.float32)])
  def batch_func(self, x):
    return x * 1.0

  ...

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'batch_func': model.batch_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

Em seguida, transmita os function_alias(s) usando as BatchOptions.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    function_alias: "batch_func"
  }
}

Definição das opções de lotes

• num_batch_threads (número inteiro): número de linhas de execução de programação para processamento de lotes de trabalho. Determina o número de lotes processados em paralelo. Isso deve estar aproximadamente alinhado ao número de núcleos de TPU disponíveis.
• max_batch_size: (número inteiro) tamanho máximo de lote permitido. Pode ser maior que allowed_batch_sizes para usar a divisão grande de lote.
• batch_timeout_micros: (número inteiro) número máximo de microssegundos a aguardar antes de gerar um lote incompleto.
• allowed_batch_sizes: (lista de números inteiros) se a lista não estiver vazia, ela preencherá os lotes até o tamanho mais próximo na lista. A lista precisa estar aumentada monotonicamente e o elemento final precisa ser menor ou igual a max_batch_size.
• max_enqueued_batches: (número inteiro) número máximo de lotes enfileirados para processamento antes que as solicitações falhem rapidamente.

Como atualizar as opções de lotes atuais

Para adicionar ou atualizar as opções de lotes, execute a imagem do Docker especificando batch_options e definindo disable_default_optimizations como verdadeiro usando a sinalização --converter_options_string. As opções de lote serão aplicadas a todas as funções de TPU ou operações de lote preexistentes.

Observação: não é possível atualizar opções de lote definindo tpu_functions juntos usando --converter_options_string.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}
disable_default_optimizations=True

Requisitos de formas em lotes

Os lotes são criados ao concatenar tensores de entrada nas solicitações com a dimensão de lote (0a). Os tensores de saída são divididos ao longo da enésima dimensão. Para executar essas operações, a operação de lote tem requisitos rigorosos de forma para entradas e saídas.

Tutorial

Para entender esses requisitos, é útil primeiro entender como os lotes são executados. No exemplo abaixo, estamos agrupando em lote uma op tf.matmul simples.

def my_func(A, B)
    return tf.matmul(A, B)

A primeira solicitação de inferência produz as entradas A e B com as formas (1, 3, 2) e (1, 2, 4), respectivamente. A segunda solicitação de inferência produz as entradas A e B com as formas (2, 3, 2) e (2, 2, 4).

O tempo limite de lotes foi atingido. O modelo aceita um tamanho de lote de três. Portanto, as solicitações de inferência 1 e 2 são agrupadas sem qualquer preenchimento. Os tensores em lote são formados pela concatenação das solicitações no 1 e 2 ao longo da dimensão do lote (0a). Como o A do 1 tem uma forma de (1, 3, 2) e o A do 2 tem uma forma de (2, 3, 2), quando eles são concatenados com a dimensão do lote (0a), a forma resultante é (3, 3, 2).

O tf.matmul é executado e produz uma saída com a forma (3, 3, 4).

A saída de tf.matmul é agrupada em lote. Portanto, ela precisa ser dividida novamente em solicitações separadas. A operação de lote faz isso dividindo ao longo da dimensão do lote (0a) de cada tensor de saída. Ela decide como dividir a 0a dimensão com base na forma das entradas originais. Como as formas da solicitação no 1 têm uma 0a dimensão de 1, a saída tem uma 0a dimensão de 1 para uma forma de (1, 3, 4). Como as formas da solicitação no 2 têm uma 0a dimensão de 2, a saída tem uma 0a dimensão de 2 para uma forma de (2, 3, 4).

Requisitos de forma

Para executar a concatenação de entrada e a divisão de saída descritas acima, a op de lote tem os seguintes requisitos de forma:

1. As entradas para lotes não podem ser escalares. Para concatenar na dimensão 0, os tensores precisam ter pelo menos duas dimensões.

   No tutorial acima. Nem A nem B são escalares.

   Se esse requisito não for atendido, ocorrerá um erro como: Batching input tensors must have at least one dimension. Uma correção simples para esse erro é transformar o vetor em um escalar.

2. Em diferentes solicitações de inferência (por exemplo, invocações de execução de sessão diferentes), os tensores de entrada com o mesmo nome têm o mesmo tamanho para cada dimensão, exceto a 0a dimensão. Isso permite que as entradas sejam concatenadas corretamente ao longo da 0a dimensão

   No tutorial acima, a solicitação A da solicitação 1 tem a forma (1, 3, 2). Isso significa que qualquer solicitação futura precisa produzir uma forma com o padrão (X, 3, 2). A solicitação 2 atende a esse requisito com (2, 3, 2). Da mesma forma, a solicitação B da solicitação 1 tem a forma (1, 2, 4). Portanto, todas as solicitações futuras precisam produzir uma forma com o padrão (X, 2, 4).

   Se esse requisito não for atendido, ocorrerá um erro como: Dimensions of inputs should match.

3. Para uma determinada solicitação de inferência, todas as entradas precisam ter o mesmo tamanho de dimensão 0. Se diferentes tensores de entrada para a operação de lote tiverem dimensões 0 diferentes, a operação de lote não saberá como dividir os tensores de saída.

   No tutorial acima, todos os tensores da solicitação 1 têm um tamanho de 1 como a dimensão 0. Isso informa à operação de lote que a saída precisa ter um tamanho de dimensão 0 de 1. Da mesma forma, os tensores da solicitação n.o 2 têm um tamanho de 0 a 2, portanto, a saída terá um tamanho de 2 de 0 a dimensão. Quando a operação de lote divide a forma final de (3, 3, 4), ela produz (1, 3, 4) para a solicitação no 1 e (2, 3, 4) para a solicitação no 2.

   O não cumprimento desse requisito resultará em erros como: Batching input tensors supplied in a given op invocation must have equal 0th-dimension size.

4. O tamanho 0 da forma de cada tensor de saída precisa ser a soma de todos os tamanhos 0 dos tensores de entrada, mais qualquer preenchimento introduzido pela operação de lote para atender à próxima maior allowed_batch_size. Isso permite que a operação em lote divida os tensores de saída ao longo da 0a dimensão com base na 0a dimensão dos tensores de entrada.

   No tutorial acima, os tensores de entrada têm uma 0a dimensão de 1 da solicitação no 1 e 2 da solicitação no 2. Portanto, cada tensor de saída precisa ter uma 0a dimensão de 3, porque 1+2=3. O tensor de saída (3, 3, 4) atende a esse requisito. Se 3 não fosse um tamanho de lote válido, mas 4 fosse, a operação de lote teria que preencher a 0a dimensão das entradas de 3 a 4. Nesse caso, cada tensor de saída teria que ter 4 como tamanho de 0.

   O não cumprimento desse requisito resultará em um erro como: Batched output tensor's 0th dimension does not equal the sum of the 0th dimension sizes of the input tensors.

Como resolver erros de requisitos de forma

Para atender a esses requisitos, considere fornecer uma função ou assinatura diferente para o lote. Também pode ser necessário modificar as funções existentes para atender a esses requisitos.

Se uma função estiver sendo agrupada em lote, verifique se as formas de input_signature de @tf.function todas têm None na dimensão 0 (também conhecida como dimensão do lote). Se uma assinatura estiver sendo agrupada em lote, verifique se todas as entradas têm -1 na 0a dimensão.

A op BatchFunction não é compatível com SparseTensors como entradas ou saídas. Internamente, cada tensor esparso é representado como três tensores separados, que podem ter diferentes tamanhos de 0a dimensão.