Categoria:CompileTime: SparseCore No Viable Logical Replica Count
Esse erro ocorre quando o compilador XLA:SparseCore não consegue determinar uma
configuração válida de contagem de réplicas lógicas que permita que a carga de trabalho se ajuste à
memória de bloco de notas local do SparseCore (Tilespmem).
Exemplos de mensagens de erro:
XLA:TPU compile permanent error. Compilation failure: No viable logical replica count for the embedding table with metadata: max_nz_per_row = 141352, max_unique_nz_per_row = 8, feature_width = 8, sample_count = 204800 (last tried split factor for vector splitting = 1, last tried split factor for sample dimension splitting = 1, fixed_size_allocation_bytes = 410880, row_dependent_size_allocation_bytes = 1696224, total_spmem_size_bytes = 524288) ...
Back-ends do XLA:TPU
Visão geral
Esse erro é específico para casos de uso do SparseCore, principalmente para modelos de embedding grandes (LEMs, na sigla em inglês).
A contagem de réplicas lógicas é um parâmetro interno do compilador que determina como os lotes de entrada são particionados para gerenciar a pressão de alocação do bloco de notas. O compilador tenta dividir a carga de trabalho em partes menores (réplicas) para que os buffers intermediários necessários para cada parte se encaixem na memória de bloco de notas limitada do SparseCore. Em geral, uma contagem maior de réplicas lógicas reduz a pressão de alocação processando lotes menores de dados por vez.
Esse erro indica que, mesmo depois de tentar várias configurações de divisão, o compilador não conseguiu encontrar uma configuração em que os buffers necessários se encaixassem na memória Tilespmem. O tamanho da alocação é determinado por uma combinação de:
sample_count: o número de IDs de pesquisa de embedding atribuídos a cada SparseCore (derivado do tamanho do lote).feature_width: o tamanho da dimensão do embedding.max_nz_per_row: o número máximo de IDs de pesquisa de incorporação não exclusivos em todos os SparseCores.max_unique_nz_per_row: o número máximo de IDs de pesquisa de incorporação exclusivos.
Depuração
Para resolver esse erro, reduza a pressão sobre a memória no bloco de notas do SparseCore.
1. Melhorar as estimativas de metadados
O compilador aloca memória com base em max_nz_per_row e max_unique_nz_per_row. Se esses valores forem estimados de forma conservadora (ou seja, definidos muito mais altos do que os dados reais exigem), o compilador vai reservar espaço desnecessário, causando esse erro. Verifique se esses parâmetros refletem com precisão a distribuição real de IDs do seu conjunto de dados.
Considere aplicar a otimização direcionada por feedback (FDO) para determinar os valores ideais desses parâmetros.
2. Reduzir o tamanho do lote
O sample_count é derivado diretamente do tamanho global do lote. Reduzir o tamanho do lote diminui a quantidade de dados que cada SparseCore precisa processar por etapa, reduzindo assim o tamanho dos buffers de rascunho necessários.