Código do erro: E2001

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

Esse erro ocorre quando o tipo de dados usado para o operando lado direito (RHS, na sigla em inglês) em uma multiplicação de matrizes (por exemplo, jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul ou o operador @) não é compatível nativamente com a geração específica de TPU em uso.

Exemplos 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 do hardware. Esse erro é acionado quando o kernel tenta mapear uma multiplicação de matrizes para a MXU usando um tipo de dados que sua geração específica de TPU não tem o circuito físico para executar.

Esse erro geralmente indica que a transmissão de canonizaçã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 as capacidades do hardware. Você tem as seguintes opções:

1. Fazer cast para tipos nativos

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

  • Por quê:Float32 é o tipo de dados universal compatível nativamente com a MXU em todas as gerações de TPU.
  • Compensação:isso vem com um custo de VPU (unidade de processamento vetorial), os ciclos necessários para realizar a transmissão, mas garante que seu 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 de XLA para garantir que --xla_mosaic_compat_mode não esteja definido como "false".

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

O que o modo de compatibilidade permite:

  • MatMuls de precisão mista:permite misturar operandos inteiros com acumuladores de ponto flutuante 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 sem suporte, como ponto flutuante 4-bit (4E2M1FN) ou 8-bit (8E4M3FN), estendendo-os para tipos compatíveis, como BFloat16 ou Float32 antes da execução.

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

3. Fazer upgrade de hardware ou pedir suporte

Se o algoritmo exigir desempenho nativo para tipos como Int4 ou Float8 sem o custo 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 seu hardware é compatível com essa operação ou se o compilador não tem um caminho de emulação válido mesmo no modo de compatibilidade, envie 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 vai funcionar em todas as gerações futuras. Mas não há garantia de que ele tenha emulação para gerações mais antigas, para algumas das quais as transmissões seriam muito caras.