Introdução ao Cloud TPU
As unidades de processamento de tensor (TPUs) são circuitos integrados de aplicativo específico (ASICs) 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 modelos com mais eficiência ao usar 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 (HBM) no chip para permitir o uso de modelos e tamanhos de lote maiores. Elas podem ser conectadas em grupos chamados frações, que escalonam as cargas de trabalho com pouca ou nenhuma alteração no código.
O código executado em TPUs precisa ser compilado pelo compilador de álgebra linear acelerada (XLA). O XLA é um compilador JIT (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 dele em um código de máquina de TPU. O restante do programa é executado na máquina host da TPU. O compilador XLA faz parte da imagem da VM de TPU que é executada em uma máquina host de TPU.
Para mais informações sobre unidades de processamento de tensor, consulte O que são TPUs.
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 machine learning. Em geral, é possível decidir qual hardware é melhor para a carga de trabalho com base nas diretrizes a serem seguidas.
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 com muitas operações personalizadas do TensorFlow escritas em C++.
- Modelos limitados pelas 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 cálculos 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 de lote efetivo grandes.
- Modelos com embeddings ultragrandes, comuns em cargas de trabalho avançadas de classificação e recomendação.
Os Cloud TPUs não são adequados para as seguintes cargas de trabalho:
- Programas de álgebra linear que exigem 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 redes neurais com operações personalizadas no loop de treinamento principal.
TPUs no Google Cloud
É possível usar TPUs com as VMs do Cloud TPU, o Google Kubernetes Engine e a Vertex AI. A tabela a seguir lista os recursos de cada serviço do Google Cloud.
| Serviço doGoogle Cloud | Recursos |
|---|---|
| Cloud TPU | Introdução às VMs do Cloud TPU |
| Google Kubernetes Engine | |
| Vertex AI |
Práticas recomendadas para o desenvolvimento de modelos
Um programa com cálculos dominados por operações não matriciais, como adição, reformatação ou concatenação, provavelmente não atingirá uma alta utilização de MXU. A seguir, conheça algumas diretrizes que ajudam você 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 cálculos na MXU com mais eficiência. A estrutura do hardware da MXU (uma matriz sistólica de 128x128) e o design do subsistema de memória da TPU (que preferencialmente usa dimensões 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 adequados para a divisão em blocos, enquanto outros requerem a realização de reformatações antes da divisão. Em geral, as operações de reformatação são limitadas pela memória no Cloud TPU.
Formatos
O compilador XLA compila um gráfico de ML a tempo para o primeiro lote. Se os lotes seguintes tiverem formatos diferentes, o modelo não vai funcionar. Compilar novamente o gráfico sempre que o formato muda é um processo muito lento. Portanto, qualquer modelo que tenha tensores com formatos dinâmicos não é adequado para as TPUs.
Preenchimento
Nos programas de alto desempenho do Cloud TPU, cálculos densos podem ser facilmente divididos em blocos de 128x128. Quando um cálculo 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 de TPU.
- O preenchimento aumenta a quantidade de armazenamento em memória no chip necessária para um tensor, o que pode resultar em um 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 usando a ferramenta op_profile. Para evitar o preenchimento, escolha dimensões de tensor que sejam adequadas para TPUs.
Dimensões
É muito importante escolher as dimensões de tensores adequadas para conseguir o máximo de desempenho do hardware de 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 deles deve ser múltiplo de 128. Caso contrário, o compilador preenche um desses valores até chegar em 128. De preferência, o tamanho do lote e as dimensões do recurso devem ser múltiplos de oito para permitir um alto desempenho do subsistema de memória.
Introdução ao Cloud TPU
- Configurar uma conta do Google Cloud
- Ativar a API Cloud TPU
- Conceder ao Cloud TPU acesso a 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
Pedir ajuda
Para receber ajuda, entre em contato com o suporte do Cloud TPU. Se você tiver um projeto ativo do Google Cloud , prepare-se para fornecer estas informações:
- O ID do projeto do Google Cloud .
- O nome da TPU, se aplicável.
- Outras informações relevantes.
A seguir
Quer saber mais sobre o Cloud TPU? Confira estes recursos: