Uma análise detalhada sobre SparseCore para modelos de embedding grandes (LEM)

O SparseCore é um processador em blocos especializado projetado para aceleração de alto desempenho de cargas de trabalho que envolvem acesso e computação irregulares e esparsos à memória, principalmente em grandes conjuntos de dados armazenados em memória de alta largura de banda (HBM). Ele é excelente em tarefas como pesquisas de incorporação, mas também acelera várias outras cargas de trabalho dinâmicas e esparsas.

1. Introdução ao SparseCore

Principais recursos da arquitetura:

• Arquitetura em blocos: consiste em vários blocos de computação (cada bloco é uma unidade completa do Dataflow com memória local e unidade de processamento próprias), permitindo o processamento paralelo.
• Execução dinâmica: oferece suporte nativo a fluxo de controle e acessos à memória dependentes de dados, o que é crucial para dados esparsos.
• Processamento de vetores: usa tarefas de vetores pequenos (8 ou 16 elementos, dependendo da versão do hardware) para computação eficiente.
• Controle centralizado: um único sequenciador SparseCore orquestra tarefas em todos os blocos, garantindo operações sincronizadas.
• Suporte para resumo de dados: inclui operações entre faixas especializadas que são úteis para tarefas como classificação, filtragem e somas de prefixo.
• Hierarquia de memória: usa estrategicamente a HBM para armazenar grandes conjuntos de dados e a memória de bloco de notas local (SPMEM) para organizar dados acessados com frequência, reduzindo significativamente a latência da HBM.

Resumo das especificações:

2. Pré-processamento de host do SparseCore

A preparação eficaz de dados é fundamental para o desempenho do SparseCore, e é aí que o pré-processamento do host desempenha um papel vital. Ele abrange várias funcionalidades importantes:

• Transformação de dados:
  • Aplique as transformações necessárias aos dados de entrada brutos.
  • Gerenciar transformações de ID, o que é particularmente importante ao lidar com o empilhamento de recursos ou tabelas.
  • Converta os dados de entrada para o formato esparso de coordenada (COO), detalhado na seção a seguir.
  • Particione os dados para uma distribuição eficiente entre os diferentes SparseCores disponíveis no chip.
• Validação de limite:
  • Verifique se as características dos dados de entrada (por exemplo, número de IDs) estão de acordo com os limites operacionais predefinidos do SparseCore, como max_ids_per_partition e max_unique_ids_per_partition.
  • Se os dados de entrada excederem esses limites, a camada de pré-processamento do host poderá tentar segmentar os dados em minilotes menores que se encaixem nas restrições.
• Transferência de dados:
  • Copie com eficiência os dados processados e validados para a memória de alta largura de banda (HBM) da TPU, preparando-os para a execução do SparseCore.

Como entender o empilhamento de tabelas:

O empilhamento de tabelas é uma técnica de otimização significativa em que várias tabelas de incorporação são combinadas logicamente para aumentar a eficiência da pesquisa de incorporação. Esse processo geralmente é gerenciado automaticamente pela estrutura de ML subjacente.

• Agrupamento de recursos: isso ocorre quando vários recursos distintos compartilham a mesma tabela de incorporação. Um exemplo comum é usar um único dicionário de incorporação para vários recursos categóricos, como CEPs de diferentes contextos.
• Empilhamento de tabelas: nesse cenário, várias tabelas de incorporação distintas são empilhadas. As tabelas que compartilham a mesma dimensão de incorporação e configuração de otimizador geralmente são agrupadas.

A principal vantagem do empilhamento de tabelas é a criação de um tamanho de lote efetivo maior para as operações nessas tabelas empilhadas. Isso reduz a sobrecarga computacional e pode ser eficaz para ocultar latências de comunicação entre chips (ICI, na sigla em inglês). Para um desempenho ideal, é recomendável usar um número moderado de tabelas empilhadas (geralmente de 5 a 100).

3. Conversão para tensores COO

Antes de serem processados pelo SparseCore, os dados geralmente são convertidos em um formato de tensor esparso de coordenadas (COO). O formato COO é uma maneira de representar matrizes esparsas de forma eficiente, geralmente usando três matrizes:

• row_ids: uma matriz que contém os índices de linha para cada elemento diferente de zero. No contexto do processamento em lote, isso geralmente corresponde à dimensão de lote.
• col_ids: uma matriz que contém os índices de coluna para cada elemento diferente de zero. Para embeddings, geralmente são os valores de recurso ou ID.
• values (opcional): uma matriz que contém os valores reais dos elementos diferentes de zero nas coordenadas (row, col) correspondentes. Para cálculos de limite (discutidos mais adiante) relacionados a contagens de ID, esses valores (ganhos) geralmente não são considerados.

Exemplo:

Considere uma matriz esparsa de entrada que representa lotes de IDs:

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

Após a conversão para o formato COO (e possivelmente após a remoção de IDs duplicados na mesma amostra):

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

Essa conversão é fundamental para a forma como o SparseCore processa e distribui o trabalho. Os col_ids, em particular, são essenciais para determinar a qual partição específica do SparseCore um ID pertence, permitindo fragmentação e pesquisa eficientes.

4. SparsecoreConfig: a API de alto nível

APIs de embedding específicas do framework:

• JAX: https://github.com/jax-ml/jax-tpu-embedding (link em inglês)
• TensorFlow: https://www.tensorflow.org/recommenders/api_docs/python/tfrs/layers/embedding/TPUEmbedding
• Keras: https://keras.io/keras_rs/api/embedding_layers/distributed_embedding

O SparsecoreConfig ou mecanismos equivalentes, como flags do XLA, serve como uma interface de alto nível para controlar uma ampla variedade de comportamentos do SparseCore. É fundamental entender esses parâmetros para ajustar o desempenho de forma eficaz e garantir a operação correta dos modelos.

• disable_table_stacking: bool = False
  • Explicação: essa flag controla se o empilhamento automático de tabelas impede que o framework empilhe tabelas, o que pode levar a redução da performance devido ao aumento dos overheads e à diminuição da capacidade de ocultar a latência da interconexão entre chips (ICI, na sigla em inglês).
  • Padrão: False (implicando que o empilhamento de tabelas geralmente é ativado por padrão quando a estrutura o oferece suporte).
• max_ids_per_chip_per_sample: int = 64
  • Explicação: esse parâmetro estabelece um limite máximo global para o número total de IDs de incorporação que um único chip pode processar de uma amostra no lote de entrada, agregados em todas as tabelas. É um mecanismo para gerenciar recursos no nível do chip, antes que limites mais granulares por tabela ou por partição sejam considerados. O ajuste desse valor geralmente depende de características específicas do modelo e da capacidade geral do sistema.
  • Padrão: 64.
• max_ids_per_table: Optional[Dict[str, int]] = None
  • Explicação: esse parâmetro especifica o número máximo de IDs de incorporação (que podem incluir duplicados) que podem ser processados para cada tabela lógica, considerando todas as partições em todos os SparseCores. Esse é um limite mais amplo do que max_ids_per_partition. Se uma tabela T for dividida em P partições, esse limite será aplicado à soma dos IDs direcionados a todas as partições P. Ela geralmente está relacionada a max_ids_per_partition_per_sample e ao tamanho geral do lote.
  • Configuração: geralmente configurada usando um arquivo de limites (por exemplo, usando a flag xla_sparse_core_max_ids_file), em que max_ids_per_partition é definido. Esse conceito de tabela é um método para definir os limites de partição (max_ids e max_uniques).
  • Padrão: None. O valor pode ser inferido dos limites por partição ou de outras configurações se não for fornecido explicitamente.
• max_unique_ids_per_table: Optional[Dict[str, int]] = None
  • Explicação: análogo a max_ids_per_table, mas esse parâmetro especifica o número máximo de IDs exclusivos para cada tabela lógica. Essa é uma configuração crítica para dimensionar adequadamente os buffers no dispositivo usados no processamento de ID exclusivo e nas operações de vetor subsequentes.
  • Configuração: também é definida em um arquivo de limites ou derivada de max_unique_ids_per_partition_per_sample.
  • Padrão: None.
• allow_id_dropping: bool = False
  • Explicação: essa flag booleana controla a remoção de IDs quando o número de IDs encontrados nos dados de entrada (limites observados) ultrapassa os limites definidos durante a compilação (por exemplo, max_ids_per_partition).
    • Se True: os IDs que excederiam os limites serão descartados silenciosamente. Normalmente, os IDs em uma partição são processados em ordem classificada, e qualquer ID que exceda o limite da mini-lote designada é descartado. Isso permite que o programa continue a execução, mas pode ter um impacto negativo na acurácia do modelo.
    • Se False: um erro é acionado, e o processo provavelmente será encerrado se os limites observados excederem os limites compilados. Essa abordagem garante que todos os dados sejam processados, mas exige que os limites sejam configurados de forma mais conservadora.
  • Padrão: False (causando um erro em caso de estouro em vez de remoção silenciosa de dados).

• initialize_tables_on_host: bool = True

  • Explicação: essa flag determina se as tabelas de incorporação são inicializadas na CPU do host antes de serem transferidas para a memória de alta largura de banda (HBM) da TPU. A prática padrão é que as tabelas sejam inicializadas no host. Definir como True segue essa convenção. Se ele fosse definido como False, isso implicaria um mecanismo de inicialização no dispositivo, que poderia ter diferentes implicações de desempenho ou pré-requisitos de inicialização específicos.

• enable_fast_table_initialization: bool = False

  • Explicação: inicializa as tabelas diretamente na TPU. Isso pode ajudar a reduzir os tempos de inicialização do modelo.

5. Pipeline para desempenho

O pipeline é uma técnica de otimização de desempenho que permite a execução simultânea de operações no TensorCore (TC) e no SparseCore (SC). Ao sobrepor esses cálculos, a taxa de transferência geral pode ser significativamente melhorada.

• Mecanismo: em uma etapa de treinamento padrão que envolve pesquisas de incorporação esparsas (processadas pelo SC) e cálculos de camada densa (processados pelo TC), o pipeline permite que o SC trabalhe na parte da etapa i (por exemplo, transmissão direta ou indireta) enquanto o TC processa simultaneamente uma parte diferente da mesma etapa i ou até mesmo partes de etapas adjacentes, como i-1 ou i+1.
• Impacto nos gradientes: o SparseCore pode operar em gradientes "desatualizados". Por exemplo, os gradientes calculados durante a fase de backpropagation da etapa i podem não ser totalmente atualizados e visíveis para o SC até a etapa i+2.
• Compensação entre desempenho e valores numéricos: essa execução sobreposta pode levar a acelerações substanciais, potencialmente até uma melhoria de 2x no tempo de etapa do dispositivo. No entanto, as mudanças sutis nos valores numéricos (embedding_weights) resultantes do uso de gradientes desatualizados podem influenciar o comportamento de convergência do modelo ou a acurácia final alcançada. A aceitabilidade dessa troca depende muito do modelo e geralmente requer validação empírica.
• Flag de controle: o encadeamento pode ser controlado por tf_xla_disable_full_embedding_pipelining. Definir essa flag como true desativa o pipeline completo (sobreposição de computação TensorCore e SparseCore), enquanto defini-la como false (ou se a semântica da flag implicar ativação quando for falsa) a ativa.

Fluxo conceitual de pipeline:

• Sem pipelining (fluxo sequencial simplificado):

  Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i

• Com o encadeamento (fluxo simplificado sobreposto):

  Time ->
  Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
  Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1
  

  Observação: as etapas reais de pipeline implementadas no hardware e no compilador podem ser mais complexas, geralmente envolvendo pré-loops, loops de execução principais e pós-loops para gerenciar dependências de dados e garantir a correção.

6. A função do XLA

A álgebra linear acelerada (XLA, na sigla em inglês) é o compilador específico ao domínio que traduz gráficos computacionais de alto nível, geralmente de frameworks como o TensorFlow, em código de máquina altamente otimizado e adaptado para TPUs. Isso inclui gerar as instruções para operações destinadas ao SparseCore.

Funções principais no contexto SparseCore:

• Compilação de operações esparsas: o XLA é responsável por compilar operações de pesquisa de incorporação (como SparseDenseMatmulOp) e outros cálculos esparsos em programas SparseCore executáveis de baixo nível.
• Integração de limites: usa os limites operacionais configurados (por exemplo, max_ids_per_partition, max_unique_ids_per_partition, geralmente fornecidos por um arquivo de limites especificado por flags como xla_sparse_core_max_ids_file) para determinar estaticamente os tamanhos e alocar buffers de memória no dispositivo, principalmente no SPMEM.
• Otimizações segmentadas: o XLA realiza um conjunto de otimizações projetadas especificamente para a arquitetura SparseCore. Isso pode incluir programação de instruções, transformações de layout de memória e fusão de operações para maximizar a eficiência.
• Controle usando flags: muitos aspectos do comportamento do SparseCore, parâmetros de ajuste e estratégias de otimização são expostos e controlados por flags do XLA (por exemplo, xla_sparse_core_estimate_max_ids para estimativa de limite ou xla_sc_detect_nan para depuração).

Status do código aberto:

No momento, a implementação do Sparsecore é interna e veiculada usando libtpu.so.

Relatórios e diagnósticos de erros:

Falhas de compilação relacionadas a configurações do SparseCore ou restrições de recursos geralmente se manifestam como erros de compilação XLA:TPU. Essas mensagens de erro podem fornecer insights valiosos sobre problemas como limites definidos muito altos para o SPMEM disponível ou o uso de configurações não compatíveis.

7. Como os limites se traduzem em tabelas no SparseCore

No SparseCore, os "limites" são parâmetros de configuração fundamentais que se referem principalmente a duas configurações por partição para cada tabela fragmentada (distribuída) nos SparseCores disponíveis:

• max_ids_per_partition: define o número máximo de IDs totais (incluindo duplicados) que um único SparseCore pode enviar ou processar para uma partição específica de uma determinada tabela em uma única etapa computacional.
• max_unique_ids_per_partition: define o número máximo de IDs exclusivos que um único SparseCore pode enviar ou processar para um

Tradução para layout e processamento de tabelas físicas:

• Estratégia de fragmentação de tabela: as tabelas de incorporação geralmente são "mod-sharded" em todos os SparseCores do sistema. Isso significa que cada SparseCore se torna responsável por um subconjunto distinto do vocabulário (linhas) de cada tabela. Um ID j geralmente é atribuído a SparseCore_k com base em uma fórmula como k = j % num_total_sparse_cores.
• Definição de "partição": neste contexto, uma "partição" se refere ao segmento específico de uma tabela de incorporação para o qual um único SparseCore processa pesquisas.
• Alocação de buffer SPMEM: esses limites são usados pelo compilador XLA para dimensionar e alocar buffers de forma estática na memória de rascunho do dispositivo (SPMEM). Os buffers são dimensionados para que todos os dados necessários relacionados aos IDs de uma determinada partição (até os limites especificados de max_ids e max_unique_ids) possam ser carregados no SPMEM para processamento. Isso é especialmente importante para cálculos não elementares, como a redução de IDs duplicados em uma partição (por exemplo, ao criar uma representação de linha esparsa compactada [CSR, na sigla em inglês]), em que todo o conjunto de dados relevante para os IDs dessa partição precisa estar disponível na memória rápida.

• Limites compilados x limites observados:

  • Limites observados: é o número real de IDs encontrados para cada partição durante a execução, com base nos dados de entrada que estão sendo processados.
  • Se os limites observados excederem os limites compilados, isso poderá levar à remoção de IDs (se allow_id_dropping estiver ativado) ou a erros.

• Cálculo de limites: o processo de determinar limites adequados envolve uma análise cuidadosa da distribuição dos dados de entrada. Para qualquer tabela (vamos chamar de T1, que pode fazer parte de uma tabela empilhada maior T):

  1. O lote de entrada (por exemplo, um SparseTensor 2D de formato [BatchSize, MaxSequenceLength]) é dividido inicialmente entre os SparseCores disponíveis. Por exemplo, se um TensorCore for pareado com dois SparseCores, cada SparseCore poderá receber um sub-lote de formato [BatchSize/2, MaxSequenceLength].
  2. Em seguida, esse sublote é convertido no formato COO, gerando row_ids e col_ids.
  3. IDs duplicados na mesma amostra (ou seja, entradas com o mesmo row_id e col_id) são removidos.
  4. Para cada col_id restante (em uma amostra), o SparseCore de destino responsável por esse ID é determinado usando a regra de mod-sharding: target_sc_id = col_id % num_total_sparse_cores.
  5. Uma contagem é mantida do número total de IDs (ids_per_sparse_core[target_sc_id]++) e do número de IDs exclusivos (unique_ids_per_sparse_core[target_sc_id]++, depois de garantir a exclusividade para esse target_sc_id específico) destinados a cada target_sc_id.
  6. O max_ids_per_partition da tabela T1 é definido como max(ids_per_sparse_core_array).
  7. Da mesma forma, o max_unique_ids_per_partition da tabela T1 é definido como max(unique_ids_per_sparse_core_array).
  8. Se a tabela T1 for um componente de uma tabela agrupada, outras transformações, como rotações ou mudanças, poderão ser aplicadas às distribuições de ID antes de somar as estatísticas de todas as tabelas constituintes. Isso ajuda a equilibrar a carga entre os chips.

Definir esses limites corretamente é um ato de equilíbrio: limites mais baixos podem levar a um desempenho maior (já que menos dados precisam ser processados por etapa e a pressão do SPMEM é reduzida), mas, se definidos muito baixos, podem resultar em mini-lotes excessivos ou remoção indesejável de IDs.

8. Como cada SparseCore se comunica

A comunicação do SparseCore, principalmente no contexto do processamento de uma lista de IDs para pesquisas de incorporação, depende de vários mecanismos coordenados:

• Fragmentação de mod e roteamento implícito:
  • As tabelas de incorporação são fragmentadas por módulo em todos os SparseCores do sistema.
  • Quando o host fornece um lote de dados de entrada (que é pré-processado no formato COO, incluindo col_ids), o valor col_id é usado para determinar qual SparseCore é responsável por esse ID específico: target_sc_id = col_id % num_total_sparse_cores.
  • Cada SparseCore recebe e processa apenas o subconjunto de IDs mapeados para as partições de vocabulário atribuídas. A etapa de pré-processamento do host é crucial para preparar os dados de forma que cada SparseCore possa identificar e operar facilmente nos IDs relevantes.
• Distribuição de dados por host:
  • A lógica de pré-processamento do host particiona o lote de entrada geral e distribui as partes relevantes de row_ids e col_ids (junto com os recursos ou pesos associados, se aplicável) para a memória (HBM) diretamente acessível por cada SparseCore ou para uma HBM compartilhada de onde os SparseCores vão buscar os dados necessários.
• Processamento intra-SparseCore:
  • Depois que um SparseCore recebe o conjunto designado de IDs para uma determinada partição de tabela, ele realiza operações como a remoção de duplicação desses IDs e a coleta dos vetores de embedding correspondentes. Esses cálculos são principalmente locais, executados nos próprios blocos do SparseCore e usando o SPMEM local.
• Comunicação entre SparseCore (All-to-All):
  • Após a fase inicial de processamento (como pesquisas de incorporação), um padrão de comunicação "todos para todos" pode ser usado para combinar ou redistribuir resultados em SparseCores (por exemplo, antes de inserir ativações em uma camada TensorCore que espera entrada correspondente a todas as posições de amostra originais). Isso é vital para reconstruir o conjunto completo de ativações se o lote de entrada original foi distribuído para processamento paralelo.
• Comunicação com TensorCores:
  • Os SparseCores se comunicam com os TensorCores para enviar ativações de incorporação (durante a transmissão para frente) e receber gradientes (durante a transmissão para trás). Essa interação é organizada pelo programa compilado em XLA e geralmente envolve HBM como um buffer intermediário. A estratégia de pipeline (discutida anteriormente) influencia muito o tempo e a sincronização dessa comunicação SC-TC.

Em essência, a "distribuição" inicial de IDs para os SparseCores adequados é feita principalmente pelo esquema de fragmentação e pelas etapas de pré-processamento do host. A comunicação subsequente envolve o SparseCores operando nos dados locais, possivelmente seguido por operações de comunicação coletiva, como all-to-all, se os dados precisarem ser trocados ou reordenados globalmente entre os SparseCores antes de serem processados pelos TensorCores.

9. Gerenciamento de memória do SparseCore

Cada SparseCore gerencia de maneira eficiente vários tipos distintos de memória para realizar os cálculos:

• Memória de bloco de notas (SPMEM):
  • Nature: uma SRAM local relativamente pequena, mas muito rápida, disponível exclusivamente para cada SparseCore. É importante observar que o SPMEM não é um cache. O uso dele é gerenciado e orquestrado explicitamente pelo compilador XLA.
  • Finalidade: o SPMEM é usado para "organizar dados de forma oportunista". Isso inclui entradas, saídas e resultados intermediários necessários para cálculos contínuos de SC. A transferência de dados para a SPMEM reduz significativamente a alta latência normalmente associada ao acesso à HBM.
  • Dimensionamento: conforme discutido na seção "Limites", os buffers do SPMEM são dimensionados de forma estática no momento da compilação. Esse dimensionamento é baseado em parâmetros como max_ids_per_partition e max_unique_ids_per_partition. Essa alocação estática garante que, para qualquer operação em uma partição de tabela (como redução de CSR), todos os dados necessários para os IDs dessa partição (até os limites definidos) possam caber no SPMEM.
  • Otimizações do compilador: o compilador XLA incorpora otimizações sofisticadas para determinar precisamente a quantidade de dados e quais elementos de dados específicos precisam ser armazenados em SPMEM para ocultar efetivamente a latência de HBM e maximizar o desempenho.
  • Restrição de alocação dinâmica: no momento, o compilador SparseCore não oferece suporte à alocação dinâmica de bloco de notas. Isso destaca a importância crítica do dimensionamento estático com a configuração cuidadosa dos limites.
• Memória de alta largura de banda (HBM):
  • Nature: um recurso de memória grande e compartilhada acessível a todos os SparseCores, TensorCores e o sistema host. As tabelas de incorporação principais são armazenadas em HBM.
  • Uso da pilha: as operações do SparseCore geralmente exigem armazenamento temporário na HBM para resultados intermediários que não cabem na SPMEM limitada ou precisam ser transmitidos entre estágios maiores do pipeline de processamento. O uso da pilha de HBM durante as transmissões para frente e para trás pode ser estimado da seguinte forma:
    • Pilha HBM de transmissão direta (tabela única) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 bytes
    • Pilha HBM de transmissão para trás (tabela única) ≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 bytes
  • Uso de heap: a HBM também acomoda o heap, que é gerenciado pelo host. O heap armazena dados como pesos de camadas densas, constantes usadas pelo modelo e dados de entrada pré-buscados. O uso do heap tende a aumentar com o número de etapas para as quais o host faz pré-busca de dados (controlado pela flag maximum_parallel_iterations). Embora mais pré-busca possa melhorar a performance ao sobrepor transferências de host para dispositivo com computação de dispositivo, ela também consome mais HBM.
  • Serialização para otimização de HBM: a flag xla_sc_num_serialized_tables_to_optimize_hbm oferece um mecanismo para controlar quantos dados de tabelas são mantidos "ativos" na memória de pilha HBM a qualquer momento. Aumentar esse número serializa o processamento de mais tabelas, o que pode reduzir o uso máximo da pilha HBM, mas pode prejudicar o desempenho devido à redução do paralelismo.
• Memória de vetor (VMEM):
  • A VMEM é uma memória de bloco de notas local usada exclusivamente pelo TC (TensorCore). Embora a VMEM não seja gerenciada diretamente pelo SparseCore, ela é uma parte integrante do ecossistema de memória com que o SC interage, principalmente pelo TensorCore.

Estratégia geral de gerenciamento de memória:

A principal estratégia de gerenciamento de memória do SparseCore gira em torno do uso do SPMEM pequeno e rápido para os dados "ativos" que estão sendo processados por um bloco do SparseCore, minimizando assim os acessos à HBM mais lenta. Os limites configurados são o principal mecanismo para garantir que o SPMEM não transborde. A HBM é usada para armazenar grandes tabelas de incorporação e dados temporários que excedem a capacidade da SPMEM ou precisam ser compartilhados entre diferentes unidades de processamento ou etapas do pipeline. O compilador XLA é responsável por orquestrar toda a movimentação de dados e alocação de buffer com base nesses princípios arquitetônicos e nos limites configurados pelo usuário.

10. Gargalos de desempenho e memória

Para alcançar a performance ideal com o SparseCore, é necessário entender claramente os possíveis gargalos e como resolvê-los. Eles podem surgir no host, no próprio SparseCore ou na interação dele com os TensorCores.

Gargalos comuns de desempenho:

• Gargalo do host:
  • Problema: a CPU host pode não conseguir pré-processar os dados e enviá-los para a TPU com rapidez suficiente, o que leva ao subaproveitamento dos SparseCores e TensorCores. Esse é um limitador de desempenho frequente.
  • Mitigação: monitore a utilização da CPU do host e as métricas do pipeline de entrada. Otimize as rotinas de pré-processamento e carregamento de dados do lado do host (consulte as dicas de conversão de COO). Ajuste a flag maximum_parallel_iterations para refinar a pré-busca de dados.
• Sincronização TC/SC abaixo do ideal (falta de pipeline):
  • Problema: se o pipeline entre o TensorCore e o SparseCore estiver desativado ou não estiver funcionando de maneira eficiente, uma unidade poderá passar muito tempo esperando a outra, reduzindo a taxa de transferência geral do sistema.
  • Mitigação: verifique se o encadeamento está ativado (por exemplo, tf_xla_disable_full_embedding_pipelining = false ou equivalente).
• Gargalos induzidos por limites:
  • Problema:
    • Limites muito baixos: podem acionar o mini-batching excessivo (divisão de lotes de entrada em vários sublotes menores para atender aos limites restritos). Embora isso mantenha a correção, cada minilote introduz alguma sobrecarga de processamento, o que pode diminuir a velocidade da execução geral. Se allow_id_dropping for verdadeiro, limites muito baixos também poderão levar à remoção de IDs, o que afeta a acurácia do modelo.
    • Limites muito altos (mas ainda adequados): embora limites muito altos possam impedir o mini-loteamento, eles podem aumentar a pressão da SPMEM desnecessariamente se as características reais dos dados raramente se aproximarem desses valores máximos. Elas também podem levar a um uso maior da pilha HBM do que o estritamente necessário.
    • Falhas de compilação: se os limites configurados exigirem mais SPMEM ou pilha HBM do que a memória física disponível, a compilação vai falhar.
  • Mitigação: verifique se os limites estão definidos corretamente.
• Distorção na distribuição de dados:
  • Problema: se determinadas partições do SparseCore receberem consistentemente um número desproporcionalmente maior de IDs em comparação com outras (indicando uma distribuição ruim de IDs), esses SparseCores sobrecarregados vão se tornar gargalos de desempenho.
  • Mitigação: a troca de IDs durante o processo de minilotes pode ajudar a aliviar esse problema em tabelas empilhadas, especialmente aquelas com tabelas de usuários "frequentes". Analise as distribuições de ID com cuidado para definir limites adequados e equilibrados por tabela.
• Problemas de empilhamento de tabelas:
  • Problema:
    • Poucas tabelas empilhadas: talvez não seja suficiente para ocultar a latência da ICI ou reduzir adequadamente as sobrecargas de processamento.
    • Muitas tabelas empilhadas: pode resultar na criação de tabelas lógicas muito grandes que se tornam difíceis de gerenciar ou podem exceder os limites de recursos disponíveis.
  • Mitigação:
    • Garanta o número ideal de tabelas para empilhamento. Uma diretriz geral sugere um "ponto ideal" de 5 a 100 tabelas para empilhamento.
• Numéricos/quantização ineficientes:
  • Problema: usar precisão FP32 completa quando formatos de precisão mais baixa, como BF16 ou números inteiros quantizados, seriam suficientes (e oferecem computação mais rápida) pode ser um gargalo de desempenho.
  • Mitigação: use opções de menor precisão. No entanto, a quantização tem uma sobrecarga e pode exigir um ajuste cuidadoso dos parâmetros para manter a acurácia do modelo.
• Saturação da largura de banda de HBM:
  • Problema: o movimento excessivo de dados de e para a HBM, possivelmente causado por larguras de recursos muito pequenas (resultando em alta sobrecarga de padding), padrões de acesso à memória ineficientes ou um número extremamente grande de pesquisas, pode saturar a largura de banda da HBM disponível.
  • Mitigação: aumentar o número de TPUs pode ajudar na saturação da largura de banda da HBM.

Gargalos de memória comuns:

• Estouro de SPMEM (falha na compilação):
  • Problema: se max_ids_per_partition e max_unique_ids_per_partition estiverem definidos como muito altos, o compilador XLA poderá não alocar SPMEM suficiente, resultando em erros de compilação como: "Fixed size allocations (...) do not fit in TileSpmem (...)". Além disso, se o termo (sample_count * feature_width) / kNumTiles (em que kNumTiles é o número de blocos por SC) for muito grande para reunir operandos de preparo no SPMEM de bloco, poderão ocorrer erros como "Gather operand too large...".
  • Mitigação: reduza o tamanho do lote ou aumente o número de chips usados para processamento.
• Estouro de pilha do HBM (tempo de execução ou compilação):
  • Problema: se a combinação de feature_width, max_unique_nz_per_row e logical_replica_count levar a requisitos de memória de pilha HBM que excedam a HBM disponível, isso poderá causar erros de falta de memória (OOM) durante a execução ou a compilação.
  • Mitigação: ajuste a flag xla_sc_num_serialized_tables_to_optimize_hbm para reduzir o uso da pilha HBM ao serializar o processamento de tabelas. Isso geralmente tem um custo de desempenho.
• Exaustão do heap HBM:
  • Problema: causado principalmente por pesos de camada densa muito grandes, inúmeras constantes armazenadas na memória ou pré-busca de entrada excessivamente agressiva (maximum_parallel_iterations alto).
  • Mitigação: monitore o uso do heap usando ferramentas como o XProf Memory Viewer.
• Sobrecarga de padding:
  • Problema: as tabelas de incorporação são preenchidas para serem alinhadas a 32B (equivalente a 8 pontos flutuantes) na dimensão do atributo. Consequentemente, larguras de recursos pequenas (por exemplo, 1 ponto flutuante) geram uma sobrecarga de padding significativa (por exemplo, 7/8 do espaço de buffer alocado é padding), o que leva ao desperdício de HBM. A dimensão do vocabulário das tabelas também é preenchida para ser um múltiplo do número de SparseCores no sistema. No entanto, esse impacto geralmente é insignificante para tabelas com um tamanho de vocabulário suficientemente alto.

Fatores gerais que afetam o desempenho e a memória:

• Topologia: o número de chips disponíveis e a arquitetura de interconexão deles.
• Tamanho do lote: afeta diretamente o sample_count por SparseCore, o que influencia o consumo de memória e a carga de computação.
• Formatação de dados: garantir um layout de dados eficiente no dispositivo é crucial para um desempenho ideal.

11. Analisar um perfil do SparseCore

Analisar um perfil de desempenho é uma etapa fundamental para identificar gargalos e descobrir oportunidades de otimização nas suas cargas de trabalho do SparseCore.

1. Extrair um rastreamento:
   • Use ferramentas de criação de perfil, como o XProf, para capturar um trace de execução detalhado enquanto o modelo está sendo treinado ou executando inferência. Esse rastreamento vai fornecer uma linha do tempo das operações que ocorrem no host, nos TensorCores e nos SparseCores.
2. Examine o visualizador de traces (por exemplo, no XProf ou no TensorBoard):
   • Atividade do host: analise a atividade do host. Há lacunas significativas na atividade da TPU? Essas lacunas podem indicar que o host é um gargalo e não consegue enviar dados com rapidez suficiente. Analise a performance do seu pipeline de entrada.
   • Atividade do TensorCore (TC) e do SparseCore (SC):
     • Analise as linhas do tempo de execução do TC e do SC. Elas estão operando em paralelo, indicando um pipeline eficaz? Ou há períodos prolongados em que uma unidade fica ociosa, esperando a outra?
     • Identifique as operações que consomem mais tempo (operações de execução mais longa) no SC e no TC.
     • Os resultados de rastreamento visual (geralmente mostrando blocos coloridos que representam diferentes operações ao longo do tempo, como TPU:0 SparseCore 1 (pid 1005)) são muito úteis para identificar visualmente operações dominantes e períodos ociosos.
   • Análise do tempo de etapa: observe o tempo geral de etapa e entenda como ele é distribuído ou dividido entre o processamento do host, o cálculo de SC e o cálculo de TC.
3. Análise de memória (XProf Memory Viewer):
   • Uso de heap: use ferramentas como a guia "Memory Viewer" do XProf para inspecionar o uso de heap da HBM. Isso pode ajudar a determinar se pesos de modelos grandes, constantes ou pré-busca de entrada excessivamente agressiva estão consumindo uma quantidade excessiva de HBM. Ativar flags como --vmodule=best_fit_allocator=1 pode fornecer registros de uso máximo do heap.
   • Uso da pilha (indireto): embora o perfilamento direto da pilha de HBM possa ser complexo, se você encontrar erros de falta de memória e o uso do heap parecer razoável, o esgotamento da pilha de HBM (geralmente devido a limites ou larguras de recursos muito grandes) é um forte suspeito. As fórmulas fornecidas para o uso da pilha de HBM podem ajudar a estimar isso.
4. Procure padrões específicos:
   • Mini-batching: se os limites forem excedidos com frequência, você poderá observar evidências de mini-batching no rastreamento (por exemplo, um número maior de operações de SC menores do que o esperado para o tamanho do lote global). Isso geralmente pode ser inferido dos registros ou observando as contagens de invocações de determinadas operações.
   • Descarte de ID: se o descarte de ID estiver ativado e ocorrendo, os registros do sistema poderão fornecer indicações disso. Isso também seria um sinal claro de que os limites configurados são muito restritivos para os dados de entrada.
   • Tempos de compilação: tempos de recompilação estendidos, principalmente se a otimização direcionada por feedback (FDO, na sigla em inglês) estiver ativada e ajustando os limites com frequência, podem adicionar uma sobrecarga significativa ao tempo total de treinamento.
5. Correlacionar com flags e configuração:
   • Relacione o comportamento observado no perfil às suas configurações do SparseCore (configurações em arquivos de limites, flags do XLA). Por exemplo, se xla_sc_num_serialized_tables_to_optimize_hbm for definido como um valor alto, você poderá esperar uma performance mais lenta do SC, mas um consumo menor da pilha HBM.
6. Processo iterativo:
   • A criação de perfis costuma ser um processo iterativo de refinamento. Faça uma mudança específica (ajuste um limite, ative ou desative um recurso), capture um novo perfil e compare com o anterior para ver o impacto da sua modificação.

12. Flags gerais de depuração

Várias flags podem ser ativadas para ajudar na depuração de problemas relacionados à execução do SparseCore. É importante observar que ativar essas verificações geralmente causa uma penalidade de performance e, portanto, elas normalmente devem ser desativadas para execuções de produção.

• Verificações de ID (fora do intervalo):
  • Flag: xla_sparse_core_enable_id_bound_check = true
  • Finalidade: permite verificações no sistema host para detectar se algum ID de incorporação nos dados de entrada está fora do intervalo de vocabulário válido definido para uma determinada tabela de incorporação. Isso ajuda a detectar problemas relacionados a dados de entrada incorretos ou corrompidos.
• Verificador de NaN:
  • Flag: xla_sc_detect_nan = true
  • Finalidade: permite a detecção de valores NaN (não um número) em dados de ponto flutuante processados no SparseCore. Se um NaN for detectado nas entradas ou saídas de várias transmissões do compilador, essa flag vai gerar um erro. Esses erros geralmente fornecem informações sobre onde o NaN foi encontrado.
• Verificador de limites (acesso à memória):
  • Flag: xla_sc_assert_level=bounds
  • Finalidade: essa flag ativa uma ferramenta no estilo ASAN (AddressSanitizer) que reescreve instruções de acesso à memória (como carregamentos/armazenamentos de VMEM e operações de DMA) para incluir verificações dinâmicas. Essas verificações verificam se o acesso à memória está dentro dos limites alocados da região de memória de destino.
  • Comportamento: se um acesso à memória fora dos limites for detectado, a execução vai falhar.
  • Atenção: é possível que esse verificador produza falsos positivos, por exemplo, devido a padrões de acesso complexos que não são totalmente compreendidos pelo verificador. Essa transformação é aplicada em uma etapa avançada do processo de compilação do back-end.
• Verificador de buffer (corrupção de memória):
  • Flags:
    • xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
    • xla_tpu_verify_launch_id_across_cores=true
  • Finalidade: essas flags ajudam a garantir que os buffers de memória não sejam corrompidos ou substituídos inadvertidamente por operações não relacionadas. O Buffer Sanitizer verifica o conteúdo dos buffers para garantir que eles não estão mudando de forma inesperada.

13. Suporte à quantização

O SparseDenseMatmulOp do SparseCore foi projetado para oferecer suporte a operações em tabelas de incorporação usando tipos de dados de ponto flutuante de 32 bits (FP32) e inteiros. Embora o treinamento de modelos seja normalmente realizado usando a precisão FP32 para tabelas de incorporação, a quantização pós-treinamento (PTQ, na sigla em inglês) pode ser aplicada. A PTQ permite o uso de tipos de dados de menor precisão (como números inteiros de 8 bits) para inferência, o que pode melhorar o desempenho e reduzir o uso de memória.

Quantização simulada:

O SparseDenseMatmulOp pode ser configurado para realizar "quantização simulada". Nesse modo operacional, os vetores de embedding são primeiro quantizados para uma precisão menor e depois dequantizados para uma precisão maior (por exemplo, FP32) antes de serem usados em cálculos subsequentes. Essa técnica permite treinar modelos considerando os efeitos do ruído de quantização. O treinamento com quantização simulada pode melhorar a acurácia do modelo final quando ele é totalmente quantizado para inferência.

Atributos de configuração para SparseDenseMatmulOp (para quantização):

• quantization_config_num_buckets = 256
  • Esse atributo especifica o número de buckets ou níveis discretos em que um número de ponto flutuante de 32 bits será quantizado. Por exemplo, ao quantizar para números inteiros de 8 bits, normalmente se especifica 2^8 =256 buckets.
• quantization_config_low = -X.X
  • Esse atributo define o valor de ponto flutuante mínimo no intervalo de quantização. Todos os valores de entrada abaixo desse mínimo especificado serão cortados para esse valor mínimo durante a quantização.
• quantization_config_high = Y.Y
  • Esse atributo define o valor de ponto flutuante máximo no intervalo de quantização. Todos os valores de entrada acima desse máximo especificado serão cortados para esse valor máximo durante a quantização.

Interação numérica e de pipeline:

O comportamento numérico do modelo pode mudar dependendo se o encadeamento entre o TensorCore e o SparseCore está ativado. Se o encadeamento estiver ativo, os gradientes processados pelo SparseCore poderão estar "desatualizados" (de uma iteração anterior). Isso pode interagir com o processo de quantização e afetar potencialmente a dinâmica de treinamento do modelo ou a acurácia final.

14. Recursos futuros e melhorias recentes

O ecossistema do SparseCore está sujeito a desenvolvimento e melhorias contínuos.

Roteiro:

• Minilotes de amostra-dimensão:
  • Isso está planejado como um recurso complementar às funcionalidades atuais de mini-batches de dimensão de vocabulário.
  • Isso permitiria mais particionamento das entradas de incorporação ao longo da dimensão da amostra. Isso seria feito com a introdução de loops no dispositivo que podem filtrar e processar pesquisas de um subconjunto de amostras por vez. Esse recurso pode ser útil para gerenciar contagens de ID por amostra muito grandes ou para melhorar o balanceamento de carga em unidades de processamento.
• Melhor suporte para incorporação com menos de oito números inteiros por linha (larguras de recursos pequenas):
  • O design atual geralmente usa um padding significativo para larguras de recursos de incorporação que são menores que 8 floats (o que corresponde a 32 bytes). Esse padding pode levar ao desperdício de HBM e à possível subutilização de recursos de computação. As melhorias futuras visam reduzir essa ineficiência em tabelas com dimensões de recursos pequenas.

Melhorias recentes:

• Preparação de operandos de coleta em HBM:
  • Essa otimização ajuda a reduzir a pressão na memória de rascunho compartilhada (SPMEM) ao permitir que algumas entradas ou saídas de operações de coleta sejam armazenadas na HBM maior.
• Redução no uso da memória da pilha:
  • Implementamos melhorias para reduzir o consumo de memória da pilha HBM durante as operações do SparseCore, idealmente sem afetar negativamente o desempenho ou a capacidade de transferência geral.

Essas melhorias se concentram em melhorar o desempenho, a eficiência da memória e a flexibilidade operacional do SparseCore para uma variedade ainda maior de cargas de trabalho esparsas.