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 caiba na memória de rascunho local do SparseCore (Tilespmem).
Exemplo 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).
A contagem de réplicas lógicas é um parâmetro de compilador interno que determina como os lotes de entrada são particionados para gerenciar a pressão de alocação de rascunho. O compilador tenta dividir a carga de trabalho em partes menores (réplicas) para que os buffers intermediários necessários para cada parte caibam na memória de rascunho limitada do SparseCore. Geralmente, uma contagem de réplicas lógicas mais alta 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 caibam 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 embedding não exclusivos em todos os SparseCores.max_unique_nz_per_row: o número máximo de IDs de pesquisa de embedding exclusivos.
Depuração
Para resolver esse erro, é necessário reduzir a pressão de memória no rascunho 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 conjunto de dados.
Você pode aplicar a otimização direcionada por feedback (FDO) para determinar os valores ideais para esses parâmetros.
2. Reduzir o tamanho do lote
O sample_count é derivado diretamente do tamanho do lote global. A redução do 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.