Introdução ao Cloud TPU
As unidades de processamento de tensor (TPUs) são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Elas são usadas para acelerar as cargas de trabalho de machine learning. Para mais informações sobre o hardware de TPU, consulte Arquitetura de TPU. O Cloud TPU é um serviço da Web que disponibiliza TPUs como recursos de computação escalonáveis no Google Cloud.
As TPUs treinam seus modelos com mais eficiência se usarem hardwares projetados para executar operações de matriz de grande escala, geralmente encontradas em algoritmos de machine learning. As TPUs têm memória de alta largura de banda no chip (HBM), o que permite usar modelos e tamanhos de lote maiores. As TPUs podem ser conectadas em grupos chamados de fatias, que escalonam as cargas de trabalho com poucas ou nenhuma alteração no código.
O código executado em TPUs precisa ser compilado pelo compilador de álgebra linear acelerada (XLA, na sigla em inglês). A XLA é um compilador just-in-time que usa o gráfico emitido por um aplicativo de framework de ML e compila os componentes de álgebra linear, perda e gradiente do gráfico em código de máquina TPU. O restante do programa é executado na máquina host da TPU. O compilador XLA faz parte da imagem da VM da TPU que é executada em uma máquina host da TPU.
Quando usar TPUs
Os Cloud TPUs são otimizados para cargas de trabalho específicas. Em algumas situações, é recomendável usar GPUs ou CPUs em instâncias do Compute Engine para executar cargas de trabalho de aprendizado de máquina. Em geral, é possível decidir qual hardware é melhor para sua carga de trabalho com base nas diretrizes a seguir:
CPUs
- prototipagem rápida que requer flexibilidade máxima
- modelos simples que não levam muito tempo para serem treinados
- modelos pequenos com tamanhos de lotes menores
- Modelos que contêm muitas operações personalizadas do TensorFlow escritas em C++
- modelos limitados por E/S disponíveis ou pela largura de banda de rede do sistema host;
GPUs
- Modelos com um número significativo de operações personalizadas do PyTorch/JAX que precisam ser executadas pelo menos parcialmente nas CPUs
- Modelos com operações do TensorFlow que não estão disponíveis no Cloud TPU (consulte a lista de operações do TensorFlow disponíveis)
- modelos de médio a grande porte com tamanhos de lote maiores
TPUs
- modelos com preponderância de computações matriciais
- Modelos sem operações personalizadas do PyTorch/JAX no loop de treinamento principal
- modelos com treinamento de semanas ou meses
- modelos grandes com tamanhos grandes de lote efetivo
- Modelos com embeddings ultragrandes comuns em cargas de trabalho de classificação avançada e recomendação
Os Cloud TPUs não são adequados para as seguintes cargas de trabalho:
- Programas algébricos lineares que requerem ramificação frequente ou contêm muitas operações algébricas de elemento
- Cargas de trabalho que exigem aritmética de alta precisão
- Cargas de trabalho de rede neural com operações personalizadas no loop de treinamento principal
TPUs em Google Cloud
É possível usar TPUs nas VMs do Cloud TPU, no Google Kubernetes Engine e no Vertex AI. A tabela a seguir lista os recursos de cada serviço Google Cloud.
| serviçoGoogle Cloud | Recursos |
|---|---|
| Cloud TPU | Começar a usar as VMs do Cloud TPU |
| Google Kubernetes Engine | |
| Vertex AI |
Práticas recomendadas para o desenvolvimento de modelos
Um programa com computação dominada por operações não matriciais, como "add", "reshape" ou "concatenate", provavelmente não atingirá uma alta utilização de MXU. Confira a seguir algumas diretrizes que ajudam a escolher e criar modelos adequados para o Cloud TPU.
Layout
O compilador XLA executa transformações no código, incluindo a divisão do produto de matrizes em blocos menores para executar computações na MXU com mais eficiência. A estrutura do hardware da MXU (um arranjo sistólico de 128x128) e o design do subsistema de memória da TPU (que preferencialmente usa dimensões que são múltiplas de oito) são usados pelo compilador XLA para aumentar a eficiência da divisão em blocos. Por consequência, determinados layouts são mais propícios para a divisão em blocos, enquanto outros requerem a realização de remodelações antes da divisão. As operações de remodelação geralmente são limitadas pela memória no Cloud TPU.
Formas
O compilador XLA compila um gráfico de ML a tempo para o primeiro lote. Se os lotes subsequentes tiverem formas diferentes, o modelo não vai funcionar. Compilar novamente o gráfico sempre que a forma muda é um processo muito lento. Portanto, qualquer modelo que tenha tensores com formas dinâmicas não é adequado para TPUs.
Preenchimento
Programas de alto desempenho do Cloud TPU são aqueles com computação densa que pode ser facilmente dividida em blocos de 128 x 128. Quando uma computação de matriz não ocupa uma MXU inteira, o compilador preenche os tensores com zeros. O preenchimento apresenta duas desvantagens:
- Os tensores preenchidos com zeros utilizam pouco o núcleo da TPU.
- O preenchimento aumenta a quantidade de armazenamento em memória on-chip necessária para um tensor, o que pode resultar em erro de falta de memória em casos extremos.
O preenchimento é realizado automaticamente pelo compilador XLA quando necessário. No entanto, é possível determinar a quantidade de preenchimento por meio da ferramenta op_profile. Para evitar o preenchimento, escolha dimensões de tensor que sejam adequadas para as TPUs.
Dimensões
Escolher as dimensões de tensores adequadas ajuda muito para conseguir o máximo de desempenho do hardware da TPU, especialmente da MXU. O compilador XLA tenta usar o tamanho do lote ou uma dimensão de recurso para maximizar o uso da MXU. Portanto, um desses deve ser um múltiplo de 128. Caso contrário, o compilador preencherá um desses valores até chegar em 128. De preferência, o tamanho do lote e as dimensões da característica devem ser múltiplos de oito, permitindo, assim, conseguir um alto desempenho do subsistema de memória.
Primeiros passos com o Cloud TPU
- Configurar uma Google Cloud conta
- Ativar a API Cloud TPU
- Conceda ao Cloud TPU acesso aos seus buckets do Cloud Storage
- Executar um cálculo básico em uma TPU
- Treinar um modelo de referência em uma TPU
- Analisar seu modelo
Como pedir ajuda
Entre em contato com o suporte do Cloud TPU. Se você tiver um projeto Google Cloud ativo, esteja preparado para fornecer as seguintes informações:
- O ID do projeto Google Cloud
- O nome da TPU, se houver
- Outras informações que você quer fornecer
A seguir
Quer saber mais sobre o Cloud TPU? Os recursos a seguir podem ajudar: