Código do erro: E2001

Categoria:tempo de compilação: tipo de dados RHS não compatível com hardware

Esse erro ocorre quando o tipo de dados usado para o operando do lado direito (RHS) em uma multiplicação de matrizes (por exemplo, jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul ou o operador @) não tem suporte nativo da geração de TPU específica que está sendo usada.

Exemplo de mensagens de erro:

INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...

The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>

Back-ends do XLA:TPU

Visão geral

A unidade de multiplicação de matrizes (MXU) da TPU oferece suporte nativo a operações Float32 em todas as gerações de hardware.

No entanto, o suporte nativo para BFloat16 e outros tipos de dados quantizados (por exemplo, Int4, Int8 ou Float8) varia de acordo com a geração de hardware. Esse erro é acionado quando o kernel tenta mapear uma multiplicação de matrizes para a MXU usando um tipo de dados que a geração de TPU específica não tem o circuito físico para executar.

Esse erro geralmente indica que a passagem de canonicalização do compilador, que tenta converter automaticamente tipos não compatíveis em compatíveis (por exemplo, por emulação de software), não conseguiu encontrar uma regra de conversão válida ou foi impedida de fazer isso porque o modo de compatibilidade estava desativado.

Depuração

Para resolver esse erro, alinhe os tipos de dados com os recursos do hardware. Você tem as seguintes opções:

1. Converter para tipos nativos

A correção mais confiável é converter manualmente os operandos para um tipo de dados com suporte de hardware (como Float32 ou BFloat16 na TPU v4 ou mais recente) no kernel antes da operação matmul.

  • Por quê:Float32 é o tipo de dados universal com suporte nativo da MXU em todas as gerações de TPU.
  • Compensação:isso tem um custo de VPU (unidade de processamento de vetor) nos ciclos necessários para realizar a conversão, mas garante que o kernel será executado no hardware atual.

2. Verificar o modo de compatibilidade

Normalmente, o compilador pode processar automaticamente esses problemas de incompatibilidade de tipos no modo de compatibilidade , que é ativado por padrão. Verifique novamente as configurações do XLA para garantir que --xla_mosaic_compat_mode não esteja definido como falso.

Isso funciona como um "polyfill", injetando sequências de emulação de software para operações que o hardware não oferece suporte nativo.

O que o modo de compatibilidade permite:

  • MatMuls de precisão mista:permite misturar operandos inteiros com acumuladores flutuantes inserindo automaticamente operações de conversão (por exemplo, estendendo números inteiros para Float32 antes do matmul).
  • Emulação de baixa precisão: em determinadas gerações de hardware, emula tipos não compatíveis, como ponto flutuante 4-bit (4E2M1FN) ou ponto flutuante 8-bit (8E4M3FN), estendendo-os para tipos compatíveis, como BFloat16 ou Float32, antes da execução.

Esse modo prioriza a compatibilidade em relação ao desempenho máximo, já que a emulação exige instruções adicionais para converter formatos de dados antes que a MXU possa operar neles.

3. Fazer upgrade do hardware ou solicitar suporte

Se o algoritmo exigir desempenho nativo para tipos como Int4 ou Float8 sem a sobrecarga de conversão ou emulação, será necessário executar em uma geração de TPU mais recente com suporte nativo.

Solicitação de recurso:se você acredita que o hardware oferece suporte a essa operação ou se o compilador está perdendo um caminho de emulação válido, mesmo no modo de compatibilidade, registre uma solicitação de recurso. Normalmente, garantimos que as operações sejam compatíveis com versões futuras. Portanto, se o kernel for executado em uma geração de TPU, ele será executado em todas as gerações futuras, mas não é garantido que ele tenha emulação para gerações mais antigas (para algumas das quais as conversões seriam muito caras).