Você conhece a TPU do Google?

1. Introdução

(Caso seja um profissional experiente em Machine Learning ou Deep Learning, pode pular esse capítulo)

TPU significa Tensor Processing Unit ou unidades de processamento de tensores! Tensores? Ok, isso é outra coisa que precisamos esclarecer para poder explicar o resto.

Vamos começar com um exemplo de Visão Computacional clássica onde, diferente de um modelo generativo multimodal, precisamos treinar um modelo RNA (Rede Neural Artificial) para uma tarefa específica.

Normalmente trabalhamos com matrizes como formato de entrada e saída dos modelos. Mas por que? Porque é muito mais fácil executar operações em lote dentro de uma rede neural. Vamos tomar como exemplo a seguinte imagem de um gato em escala de cinza que queremos classificar.

Essa imagem é representada internamente por uma matriz de 47x47 elementos onde cada elemento é um número que corresponde à tonalidade do pixel e você pretende treinar um modelo super especializado em reconhecer essa categoria (ou classe) de objeto. Ex: Pessoa, gato, cachorro, etc.

Esse tipo de treinamento mais específico é chamado de supervisionado e para treinar o modelo você utiliza um conjunto de exemplos (ou dataset) conhecidos. Enquanto vai treinando e ajustando o modelo para ir aprendendo os padrões de cada exemplo de cada classe, você precisa em linhas gerais executar duas tarefas:

1. Medir a diferença da saída desejada versus a saída gerada pelo modelo.
2. Ajustar os pesos (parâmetros) da rede dependendo de quão perto ou quão distante você está nessa medida.

A saída de um modelo de classificação é um vetor onde cada elemento representa uma similaridade com objeto apresentado na entrada. Ou seja, qual classe é a mais “provável”.

Portanto, se um modelo treinado tem a capacidade de classificar uma imagem em 1000 classes, por exemplo, essa saída vai ser um vetor de 1000 elementos.

Daí eu lhes pergunto: Faz sentido fazer essa comparação elemento a elemento com dois vetores de 1000 elementos cada?

E durante o treinamento para ajustar os parâmetros da rede de acordo com o erro atual adivinhem só qual o formato desses parâmetros? Vetores!

Porém, nesse caso temos um problema maior pois os parâmetros têm uma relação direta com cada entrada e no nosso exemplo já é uma matriz de 47x47.

Daí eu lhes pergunto de novo: Faz sentido fazer esse ajuste dos parâmetros elemento a elemento?

Só que para piorar esse cenário, a imagem de entrada de um modelo robusto de Visão Computacional quase nunca é na verdade em escalada de cinza. Geralmente é em RGB. Ou seja, 3 canais de cores. Isto é, pode multiplicar por três tudo que eu falei!

E isso não para por aí: precisamos lembrar que um treinamento desse tipo utiliza geralmente milhares de imagens. Ou seja, multiplique tudo isso por milhares.

Daí eu lhes pergunto pela última vez: Faz sentido fazer tudo que falei elemento a elemento?

É claro que não! Por isso é mais interessante fazer em lote: matrizes de matrizes. Isto é, tensores!

2. Visão de Arquitetura

Agora que já sabemos de onde viemos, vamos falar de onde estamos e para onde estamos indo.

2.1. CPUs

Antes de entrar nos detalhes de arquiteturas mais modernas para processamento de tensores, precisamos entender quais são os limitadores que sempre existiram na computação em processadores mais tradicionais como a CPU (Central Processing Unit).

Durante a execução de um modelo em uma CPU, ela carrega dados e pesos (ambos geralmente são tensores) da memória, executa uma operação de cada vez (por núcleo) usando esses valores e armazena o resultado de volta na memória.

Ou seja, o acesso à memória é lento quando comparado à velocidade da operação e geralmente limita o rendimento da aplicação.

2.2. GPUs

Para obter maior rendimento, a GPU (Graphics Processing Unit) contém milhares de Unidades Lógicas Aritméticas (ALUs) em um único processador. Uma GPU moderna geralmente contém entre 2.500 e 5.000 ALUs. O grande número de processadores significa que você pode executar milhares de operações simultaneamente.

Essa arquitetura de GPU funciona bem em aplicações com paralelismo massivo, como operações matriciais em uma rede neural.

Em uma carga de trabalho de treinamento típica para aprendizado profundo, uma GPU pode fornecer uma taxa de transferência (throughput) numa ordem de grandeza maior do que uma CPU.

Mas a GPU ainda é um processador de propósito geral que precisa suportar muitos aplicativos e softwares diferentes. Portanto, as GPUs têm o mesmo problema que as CPUs. Para cada operação nas milhares de ALUs, uma GPU deve acessar a memória compartilhada para ler dados, pesos e armazenar os resultados.

2.3. TPUs

O Google projetou a TPU (Tensor Processing Unit) para ser um processador especializado em multiplicação de matrizes para cargas de trabalho de redes neurais. Ou seja, as TPUs não são de uso geral para qualquer tipo de aplicação.

As TPUs contêm milhares de multiplicadores-acumuladores que são conectados diretamente uns aos outros para formar uma grande matriz física na máquina host.

Para executar as operações de matrizes, as TPUs carregam os parâmetros (linhas tracejadas vermelhas na figura a seguir) da memória HBM (High Bandwidth Memory ou memória de altíssima velocidade) para as Unidades de Multiplicação de Matrizes (MXU).

Em seguida, a TPU carrega os dados da memória HBM e, à medida que cada multiplicação é executada, o resultado é passado para o próximo multiplicador-acumulador.

A saída é a soma de todos os resultados de multiplicação entre os dados e os parâmetros. Não é necessário acesso à memória principal durante o processo de multiplicação de matrizes.

Dessa forma, as TPUs podem atingir uma alta taxa de transferência computacional em operações de redes neurais.

2.4. NPUs

E ainda existe a NPU (Neural Processing Unit) que tem um foco maior na inferência eficiente do que em treinamento. Geralmente utilizada em IA embarcada na borda (dispositivos com processadores mais leves). Mas isso é assunto para um outro artigo em breve.

3. Desenvolvimento

Se você já é uma pessoa que desenvolve em PyTorch com experiência em CUDA da NVIDIA, a transição para o ambiente de TPU pode ser simplificada. O foco principal não está no código em si, mas sim na utilização de versões compatíveis de frameworks e bibliotecas que gerenciam a compilação e execução no hardware especializado da TPU.

E se você está iniciando agora sua jornada em Machine Learning ou Deep Learning, é o momento certo para aprender a escolher os frameworks e bibliotecas corretos para trabalhar com TPUs em Google Cloud.

A visão Google Cloud para a demanda de processamento massivo de tensores que existe hoje vai desde a camada mais baixa de aceleradores até formatos de provisionamento estático e dinâmico desses recursos.

Porém, neste artigo vamos focar em alguns componentes principais da camada de software para desenvolvimento e execução de modelos.

3.1. XLA

Ao contrário do ambiente CUDA, onde cada operação matemática costuma disparar um comando específico (kernel) para o hardware, as TPUs utilizam o compilador XLA (Accelerated Linear Algebra). Ele agrupa essas operações antes da execução, criando um código altamente eficiente e específico para a arquitetura da TPU, eliminando a necessidade de gerenciamento manual de kernels. Ele pega o grafo computacional do seu modelo (que normalmente seria executado passo a passo) e o compila em operações de matriz altamente eficientes e específicas para o hardware da TPU.

A principal ferramenta para habilitar o uso da TPU é a biblioteca PyTorch/XLA que é o equivalente ao uso de CUDA, mas direcionado ao hardware da TPU.

Para instalar a versão XLA do PyTorch em uma VM TPU em GCP execute os seguintes passos.

sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip install numpy
pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

OBS: Nos exemplos que vamos executar mais adiante utilizamos uma VM TPU da série v5e e ela precisa ser provisionada com a seguinte versão de software: v2-alpha-tpuv5-lite

Em PyTorch/CUDA, você move seus tensores e modelos para o dispositivo com .cuda() ou .to(‘cuda’). Em PyTorch/XLA, você usará os dispositivos XLA.

• Importação: Importe o módulo PyTorch/XLA: import torch_xla.core.xla_model as xm
• Obtendo o Dispositivo TPU: Em vez de torch.device(“cuda”), obtenha o dispositivo TPU com: device = xm.xla_device()
• Movendo Tensores/Modelos: Mova seus dados e modelos para o dispositivo XLA: model.to(device) e tensor.to(device).

Veja um código-exemplo universal de alocação de device TPU com XLA e GPU com CUDA.

def get_pytorch_device():
    """Detects and configures the best available PyTorch device."""
    device = torch.device('cpu')
    accelerator_type = 'cpu' #fallback
    xm_module = None # To return the xm module if it's a TPU
    num_devices_detected = 1 # Default for CPU

    # Try TPU
    if use_xla:
        try:
            device = xm.xla_device() # Attempts to allocate the XLA device
            accelerator_type = 'tpu'
            xm_module = xm
            # If xm.xla_device() succeeded, we can now count the supported devices
            num_devices_detected = len(xm.get_xla_supported_devices()) 
            print(f"Detected and configured for TPU: {device} ({num_devices_detected} TPUs available)")
            return device, accelerator_type, num_devices_detected, xm_module
        except RuntimeError as e:
            # Catches the error if xm.xla_device() cannot find a TPU
            print(f"INFO: Could not configure TPU via xm.xla_device(): {e}")
            pass # Continues to the next attempt (GPU)

    # Try GPU
    if torch.cuda.is_available():
        device = torch.device('cuda')
        accelerator_type = 'gpu'
        num_devices_detected = torch.cuda.device_count() # Actual number of GPUs in the system
        print(f"Detected and configured for GPU: {device} ({num_devices_detected} GPUs available)")
        return device, accelerator_type, num_devices_detected, None

    raise RuntimeError("ERROR: No TPU/GPU detected. Ensure the environment is configured correctly.")

Depois da alocação, você utiliza o objeto device no restante do seu código PyTorch.

Caso você queira fazer uma comparação mais direta, o XLA também permite compilar seu código para GPU. Basta instalar a versão PyTorch/XLA para abstrair o CUDA na máquina com GPU da NVIDIA e alocar o device com XLA.

# CUDA 12.1
pip install torch-xla==2.1.* --index-url https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch-xla.html

3.2. JAX

Devido à arquitetura da TPU, que é otimizada para multiplicação de matrizes em larga escala, o desempenho é maximizado ao usar grandes tamanhos de lote (batch sizes).

O JAX (Just-in-Time Autograd XLA) é uma biblioteca de computação numérica de alto desempenho em Python, construída sobre a sintaxe familiar de Python e NumPy. Ele é projetado para alto desempenho em aceleradores como GPUs e TPUs.

O poder central do JAX reside nas suas transformações de função que compilam e vetorizam automaticamente o código Python padrão.

Embora seja poderoso para Machine Learning (ML), o JAX é também uma biblioteca fundamental para qualquer domínio que exija operações numéricas aceleradas. O JAX usa o XLA por padrão, garantindo performance excepcional e portabilidade de hardware. Ele compila funções em código otimizado para a TPU (ou GPU).

Para configurar uma VM TPU para uso do JAX execute o seguinte comando:

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

O JAX foi projetado para velocidade em dispositivos únicos e escalabilidade em múltiplos aceleradores. Aproveitando o XLA, a distribuição de computações requer pouca configuração adicional, pois o XLA lida com a comunicação e sincronização.

O JAX simplifica a orquestração e o gerenciamento de recursos em hardware especializado em cenários de processamento massivo e pode automaticamente disparar e coordenar múltiplas threads em paralelo, uma para cada núcleo de TPU (ou chip XLA).

O código principal do JAX se concentra em transformar o código Python em funções compiladas (que rodam em TPU ou GPU via XLA), enquanto o XLA resolve a camada de orquestração de dispositivos.

Segue um exemplo de como disparar várias operações em paralelo para o mesmo conjunto de dados distribuindo o trabalho entre os vários aceleradores e sincronizando os resultados com poucas linhas de código:

pmapped_op = jax.pmap(operation_func, in_axes=(0,) * len(args), out_axes=0)

execution_times = []

# Repeat the tests in order to get average execution time  
for i in range(num_tests):
    start_time = time.time()

    # _ = pmapped_op(*args)
    # In non-benchmark applications, '.block_until_ready()' (or similar explicit synchronizations)
    # is typically removed to allow asynchronous execution. This maximizes overall throughput
    # by letting the CPU prepare the next tasks while the accelerator is busy.
    # Synchronization will happen implicitly only when the result is actually needed (e.g., fetching to CPU).
    _ = pmapped_op(*args).block_until_ready()

    end_time = time.time()
    execution_times.append(end_time - start_time)
    print(f"\rTest {i+1}/{num_tests}: {execution_times[-1]*1000:.6f} ms", end="")

O uso do JAX em conjunto com o XLA e as TPUs permite atingir altas taxas de transferência e eficiência computacional, especialmente em tarefas que envolvem multiplicação de matrizes em larga escala e orquestração paralela.

Se quiser utilizar o JAX em GPU para uma comparação direta, execute o seguinte comando:

# CUDA 12.1
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

4. Microbenchmarks

A seguir é apresentada uma simulação de operações com tensores onde foram executados as seguintes operações 50, 100, 150 e 200 vezes utilizando duas matrizes de A, B e C de 4096x4096 cada:

• Adição: A + B
• Multiplicação: A x B
• Simulação de Rede Neural: Relu(A x B + C)

Foram utilizadas uma VM TPU v5e com 8 chips e uma VM G2 com 8 GPUs L4 da NVidia. Ambas com 384 GB de memória.

Podemos observar que devido a arquitetura de carga e acesso à memória a TPU obteve uma vantagem no tempo total a cada carga e também sofreu menos com o aumento de carga conforme a quantidade de testes foi aumentando.

5. Conclusão

A Unidade de Processamento de Tensores (TPU) do Google representa um avanço arquitetônico crucial para o desenvolvimento de Machine Learning em larga escala. Ao ser projetada especificamente para operações massivas de multiplicação de matrizes, a TPU supera as limitações de acesso à memória inerentes às arquiteturas de CPU e GPU de propósito geral, entregando um alto throughput computacional.

A adoção de frameworks como PyTorch/XLA e JAX maximiza essa eficiência, simplificando a orquestração paralela e garantindo que times de Engenharia de Machine Learning possam focar na lógica do modelo, aproveitando ao máximo o hardware especializado. Além de facilitar a migração de modelos e aplicações que utilizam PyTorch/CUDA.

Os microbenchmarks demonstraram que, ao utilizar a combinação de TPU, XLA e JAX, é possível obter ganhos significativos em tempo total de execução, consolidando as TPUs como a escolha ideal para cargas de trabalho de treinamento e fine-tuning que exigem a maior eficiência possível.

Como sugestão para iniciar a experimentação utilizando PyTorch/XLA sem ter que fazer toda a configuração de ambiente do zero, é possível utilizar o Google Colab onde basta configurar o runtime para TPU.