O Google usa tecnologia de IA na tradução de conteúdos para seu idioma de preferência. As traduções com IA podem ter erros.

Otimizações de desempenho da TPU7x (Ironwood)

Este guia descreve vários métodos para otimizar o desempenho com a TPU7x (Ironwood) gerenciando de forma eficiente o movimento de dados entre o sistema de memória multicamadas. Isso inclui técnicas como treinamento de baixa precisão, fragmentação, otimização de comunicação, rematerialização de ativação, ajuste de memória vetorial com escopo (VMEM) e kernels de acelerador personalizados.

Para otimizar o desempenho com a TPU7x, primeiro você precisa conhecer a arquitetura Ironwood, especificamente a hierarquia de memória e a topologia de interconexão. Para mais informações, consulte TPU7x (Ironwood).

Treinamento de baixa precisão com FP8

O FP8 (ponto flutuante de 8 bits) é um formato de dados numéricos eficiente usado principalmente para acelerar o treinamento e a inferência de modelos. Ao representar números usando 8 bits, em vez dos formatos padrão de 16 bits (FP16 ou BF16) e 32 bits (FP32), as TPUs podem processar dados de forma significativamente mais rápida e usar menos memória.

A TPU7x é compatível com aceleração de hardware integrada para tipos de dados FP8, oferecendo um desempenho teórico máximo de 4.614 TFLOPS por chip. Esse recurso pode acelerar significativamente os tempos de treinamento de ponta a ponta. Para operações compatíveis, especialmente multiplicações de matrizes densas que são comuns para cargas de trabalho de IA, o uso de FP8 pode gerar melhorias de desempenho de 1,3x em relação ao treinamento padrão de BF16. Em comparação com o BF16, o FP8 dobra o pico de FLOPs e reduz pela metade o consumo de memória para pesos e ativações. O FP8 deve ser uma das principais ferramentas de ajuste para cargas de trabalho vinculadas à computação e cenários limitados pela capacidade de memória ou pela largura de banda.

O uso do FP8 oferece os seguintes benefícios de desempenho:

Redução da pressão na memória de alta largura de banda (HBM): um consumo de memória menor permite que modelos maiores ou com caches de chave-valor maiores durante a inferência caibam inteiramente nos 192 GiB de HBM. Isso evita o descarregamento de dados para a memória do host, que é mais lenta e cara.
Aumento do tamanho efetivo do lote: ao reduzir a memória necessária para ativações, o FP8 permite o uso de tamanhos de lote maiores. Isso melhora o paralelismo de dados e pode levar a uma maior capacidade de processamento e melhor utilização das unidades de computação.
Requisitos de largura de banda de memória mais baixos: mover metade da quantidade de dados para cada operação reduz a demanda no caminho de dados de HBM para MXU. Em sistemas em que a movimentação de dados é um gargalo comum, isso ajuda a manter as MXUs saturadas de trabalho.

O uso de FP8 com degradação zero ou limitada no desempenho requer a seleção cuidadosa de técnicas de quantização. Confira algumas práticas recomendadas para o treinamento de FP8:

Granularidade de escalonamento: comece com o escalonamento por tensor como referência. Se houver problemas de qualidade ou desempenho, mude para o dimensionamento por eixo. O escalonamento de subcanais pode ser desnecessário.
Modo de escalonamento: o escalonamento dinâmico, que calcula os fatores de escalonamento no momento da execução, é um bom padrão para manter a qualidade. Embora o escalonamento estático possa oferecer um aumento significativo de desempenho ao eliminar cálculos, ele requer uma criação de perfil cuidadosa para determinar os fatores de escalonamento corretos e pode não ser adequado para todos os casos de uso, especialmente quando as configurações do modelo mudam. Por outro lado, alguns modelos e configurações robustos podem fixar a escala no limite de FP8 para pesos ou ativações, permitindo reduzir a sobrecarga de quantização, mantendo a precisão e melhorando o desempenho.
Formatos FP8 (E4M3 e E5M2): uma abordagem comum e eficaz é usar uma combinação de formatos FP8. Por exemplo, use E4M3 para pesos e ativações na passagem direta para aproveitar a maior precisão do E4M3 e use E5M2 para gradientes na passagem reversa para acomodar o intervalo dinâmico mais amplo dos gradientes.
Arredondamento: usar o arredondamento para o par mais próximo (RNE, na sigla em inglês) em vez do arredondamento estocástico para gradientes pode manter a qualidade e oferecer melhor desempenho e reprodutibilidade.
Ativação do FP8 no MaxText: o MaxText oferece suporte ao treinamento de FP8 pela biblioteca de quantização QWIX. Para ativar a quantização, defina a seguinte flag na configuração: use_qwix_quantization=true.

Fragmentação e paralelismo

A fragmentação é o processo de dividir um modelo grande ou os dados de treinamento dele em partes menores e distribuí-las em vários chips ou núcleos de TPU. Escolher a estratégia de fragmentação certa é importante para alcançar um alto desempenho na TPU7x.

Uma abordagem simples que apenas maximiza o grau de paralelismo geralmente resulta em um desempenho ruim, porque fica limitada pela comunicação. A melhor abordagem geralmente é selecionar a estratégia de fragmentação mais simples que atenda às restrições de memória, pois isso minimiza a sobrecarga de comunicação e permite que as unidades de computação sejam usadas de maneira eficiente.

Antes de selecionar uma estratégia de fragmentação, a primeira etapa em qualquer esforço de ajuste de desempenho precisa ser uma análise de intensidade aritmética. Essa análise determina se um determinado cálculo é limitado pela computação, pela largura de banda da memória ou pela largura de banda de interconexão. Ela é calculada como a proporção de operações de usar pontos flutuantes para os bytes de dados que precisam ser movidos.

Uma alta intensidade aritmética indica uma carga de trabalho vinculada à computação. Uma baixa intensidade aritmética sugere uma carga de trabalho limitada por memória ou comunicação, em que o desempenho é limitado pela velocidade com que os dados podem ser movidos da HBM ou pela rede ICI. Essa análise informa o tamanho ideal do lote e a estratégia de fragmentação. Por exemplo, uma carga de trabalho vinculada à comunicação não se beneficiará de uma estratégia de fragmentação que introduza ainda mais comunicação, como o paralelismo de tensor de alto grau.

Framework de decisão da estratégia de fragmentação

O MaxText oferece várias estratégias de fragmentação. A escolha ideal depende da arquitetura do modelo, do comprimento da sequência e da necessidade de equilibrar a carga computacional com a sobrecarga de comunicação.

Paralelismo de dados totalmente fragmentados (FSDP): é a estratégia padrão preferida para paralelismo de dados. O FSDP fragmenta os pesos do modelo, os gradientes e os estados do otimizador nos dispositivos de paralelismo de dados. Durante a computação, cada dispositivo executa uma operação All-Gather para recuperar os pesos completos necessários para o microlote local. O FSDP é altamente eficaz, desde que o tamanho do lote por dispositivo seja grande o suficiente para ocultar a latência dessa comunicação All-Gather. Para modelos de Mix de Especialistas (MoE), o cálculo da intensidade aritmética precisa considerar a esparsidade.
Paralelismo de tensor (TP): o TP fragmenta tensores individuais em vários dispositivos. Normalmente, os tensores são matrizes de peso em perceptron multicamadas (MLP) e blocos de atenção. A alta intensidade aritmética do hardware (11,5 mil) impõe um requisito muito alto às dimensões do modelo para tornar o TP viável em relação à ICI, e tentar usar o TP pode resultar em um sistema limitado pela comunicação.
Paralelismo de especialistas (EP, na sigla em inglês): essa é a estratégia padrão e necessária para treinar modelos de MoE. O EP fragmenta as camadas "especialistas" em um conjunto de dispositivos, e um coletivo de comunicação All-to-All é usado para rotear tokens para o dispositivo especialista designado. O EP pode ser eficiente se a dimensão do MLP do modelo for grande o suficiente para se aproximar do limite.
Paralelismo de contexto (CP): o CP é uma estratégia especializada essencial para treinar modelos com sequências muito longas. A função principal dele é gerenciar o consumo de memória das ativações, que cresce de forma quadrática com o comprimento da sequência e pode exceder a capacidade da HBM. O CP fragmenta a dimensão de sequência dos tensores de ativação, o que permite o uso de um tamanho de lote fracionário por dispositivo. Como o CP introduz mais comunicação do que o FSDP, a regra geral é usar o grau mínimo de CP necessário para atender às restrições de memória e garantir que o fragmento do eixo do lote permaneça um número inteiro.

A tabela a seguir mapeia tipos de carga de trabalho comuns para a estratégia de fragmentação ideal:

Tipo de carga de trabalho	Fragmentação primária recomendada	Fragmentação secundária	Principais gargalos	Justificativa
Modelo denso: sequência curta	FSDP	N/A	Rematerialização, FF Matmuls	O FSDP oferece o melhor equilíbrio. Com sequências curtas, a memória de ativação pode não ser uma grande preocupação. A chave é um lote global grande o suficiente para ocultar o peso All-Gather do FSDP. À medida que o tamanho do lote aumenta, o tamanho da ativação aumenta, e uma política de rematerialização adequada é necessária para garantir que essa configuração não fique sem memória.
Modelo denso: sequência longa	FSDP	CP	Flash attention, memória de ativação	A memória de ativação se torna a principal restrição. O CP é necessário para ativar tamanhos de lote fracionários por dispositivo e evitar problemas de falta de memória (OOMs, na sigla em inglês). A atenção flash é a principal fonte de computação e tempo desperdiçado.
Modelo MoE: sequência curta	FSDP + EP	N/A	All-to-All (roteamento de especialistas), rematerialização	Os modelos de MoE exigem EP para fragmentar os especialistas. A comunicação All-to-All para roteamento de tokens é um grande gargalo que precisa ser sobreposto. A rematerialização também é uma fonte significativa de desperdício.
Modelo MoE: escala muito grande	FSDP + EP + PP	Paralelismo de modelos (MP)	Todos os gargalos mencionados anteriormente, além de bolhas de pipeline	Para modelos que excedem a memória de um único pod, o PP é necessário para fragmentar camadas entre pods. Isso introduz overheads de comunicação da DCN e de bolhas de pipeline. Essa é uma configuração altamente complexa que exige um ajuste cuidadoso.

Otimização da comunicação

O principal mecanismo para sobrepor comunicação e computação na TPU7x é chamado de SparseCore Collective Offloading. A arquitetura Ironwood inclui unidades SparseCore dedicadas, que atuam como linhas de execução de controle independentes capazes de gerenciar a movimentação de dados na malha ICI. Isso permite que as operações de comunicação coletiva (como All-Gather ou Reduce-Scatter) sejam executadas em paralelo com as principais computações que ocorrem nos TensorCores. Esse é o método recomendado para coletivos assíncronos na TPU7x. Use as flags recomendadas para ativar o descarregamento para os coletivos mais comuns.

Rematerialização de ativação

A rematerialização de ativação, também conhecida como checkpoint de gradiente, é uma técnica fundamental para reduzir o consumo de HBM de um modelo. Em vez de armazenar todas as ativações intermediárias da passagem direta na HBM para serem usadas durante a passagem reversa, ela salva apenas algumas ativações principais (pontos de verificação) e recalcula as outras sob demanda durante a passagem reversa. Isso economiza uma quantidade significativa de memória ao custo de um aumento na computação (aproximadamente 25% a 30% mais FLOPs para um bloco de transformador padrão).

A decisão de como aplicar a rematerialização de forma agressiva é um parâmetro de ajuste crítico que depende inteiramente do gargalo principal, que geralmente varia com o comprimento da sequência.

Para cargas de trabalho de sequência longa (como 128 mil): nesses casos, o tamanho dos tensores de ativação é o principal consumidor de HBM. A carga de trabalho normalmente é vinculada à memória. Portanto, aplicar uma política de rematerialização agressiva é altamente benéfico. A economia de memória permite que o treinamento continue sem erros de falta de memória e também permite tamanhos de lote maiores, e a sobrecarga computacional da recomputação é uma compensação que vale a pena.

Para cargas de trabalho de sequência curta (como 8k): nesses casos, a memória de ativação é uma preocupação muito menor, e é mais provável que a carga de trabalho seja limitada pela computação. A sobrecarga computacional da rematerialização pode ser a maior fonte de ineficiência.

Como ajustar as políticas de rematerialização no MaxText

O MaxText oferece controle granular sobre a rematerialização por meio de um conjunto de políticas predefinidas e personalizadas, configuradas usando a flag remat_policy.

Políticas predefinidas

O MaxText oferece as seguintes políticas integradas:

full: a política mais agressiva, que rematerializa quase tudo. Isso minimiza o uso de HBM, mas maximiza a sobrecarga de recomputação. Ideal para cenários de sequências longas com extrema restrição de memória.
minimal: a política menos agressiva, que armazena a maioria das ativações. Isso maximiza o uso da HBM, mas minimiza a recomputação. Ideal para cargas de trabalho de sequência curta e vinculadas à computação em que a memória não é uma preocupação.
Políticas intermediárias: opções como save_dot_with_context_except_mlp, save_qkv_proj e save_out_proj oferecem várias compensações ao criar pontos de verificação seletivos das saídas de operações de produto escalar caras, enquanto rematerializam operações mais baratas por elemento.

Políticas personalizadas

Para ter mais controle, defina remat_policy como custom. Isso permite especificar o comportamento de camadas individuais no módulo de decodificação do modelo. Cada camada pode receber um dos três comportamentos:

device: a ativação é armazenada na HBM no dispositivo de TPU.
remat: a ativação é descartada e será rematerializada durante o caminho de volta.
offload: a ativação é movida da HBM para a memória do host da CPU, liberando a HBM ao custo da latência de transferência do PCIe.

Ajuste de VMEM com escopo

O desempenho do kernel, como a atenção de flash, depende dos tamanhos de bloco selecionados no kernel, que são limitados pela memória vetorial (VMEM) disponível. Cada um dos dois TensorCores em um chip TPU7x tem 64 MiB de memória vetorial (VMEM). Essa capacidade de VMEM pode ser dividida entre o escopo atual (VMEM com escopo) e a pré-busca de peso futura. Aumentar a VMEM com escopo permite aumentar os tamanhos dos blocos no kernel, o que pode reduzir as interrupções de memória e aumentar o desempenho dos kernels. É possível alterar o tamanho da VMEM com escopo definindo xla_tpu_scoped_vmem_limit_kib (em LIBTPU_INIT_ARGS), que pode ser usado para analisar o desempenho do kernel e os limites de desempenho de ponta a ponta. A otimização do tamanho do VMEM com escopo pode afetar indiretamente o desempenho do kernel personalizado do Pallas, já que o aumento do VMEM com escopo desbloqueia um espaço de pesquisa de hiperparâmetros maior para tamanhos de blocos no kernel.

Kernels Tokamax

O Tokamax, uma biblioteca de kernels JAX de alto desempenho com muitos kernels de TPU altamente otimizados, resolve vários gargalos comuns específicos de hardware:

Atenção rápida: é usada como a principal implementação de atenção para eliminar o gargalo de HBM da atenção padrão e usa a implementação de atenção mais eficiente em TPUs.
Multiplicação de matriz agrupada (GMM, na sigla em inglês) do Megablox: para cargas de trabalho de MoE, o Megablox lida com multiplicações de matriz agrupadas de maneira eficiente, computando sobre a representação de ativações irregulares. Ele mapeia de forma eficiente a dimensão irregular, calculando multiplicações de matrizes entre grupos irregulares de linhas em LHS e a matriz de especialista correspondente, evitando a necessidade de preencher lotes para um tamanho fixo.
Ajuste empírico com tune-jax: a biblioteca tune-jax tem utilitários para realizar pesquisas empíricas de tamanhos de bloco ideais. Os tamanhos de kernel padrão geralmente não são ideais. O ajuste permite escolher tamanhos de blocos de VMEM compatíveis com o hardware para maximizar a utilização dele.
Estimativa máxima de logits: o kernel de atenção do Tokamax Splash pode ser ainda mais otimizado definindo um valor para max_logit_const. Se definido, ele substitui o cálculo de redução do logit máximo durante a operação softmax de atenção (softmax(Q * KT)), reduzindo parte da sobrecarga computacional e de sincronização. No MaxText, ele é implementado pela configuração use_max_logits_estimate, que pode ser definida como None (desativado) ou um pontuação flutuante. Verifique se o intervalo de logit do seu modelo específico permanece compatível com a estimativa para evitar o estouro numérico. O teste de convergência é recomendado se esse valor estiver definido.