Código de error: E2001

Categoría: Tiempo de compilación: RHSDataType no admitido en hardware

Este error se produce cuando el tipo de datos que se usa para el operando del lado derecho (RHS) en una multiplicación de matrices (p.ej., jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul o el operador @) no es compatible de forma nativa con la generación específica de TPU que se usa.

Ejemplos de mensajes de error:

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 de XLA: TPU

Descripción general

La unidad de multiplicación de matrices (MXU) de la TPU admite de forma nativa las operaciones de Float32 en todas las generaciones de hardware.

Sin embargo, la compatibilidad nativa con BFloat16 y otros tipos de datos cuantificados (p.ej., Int4, Int8 o Float8) varía según la generación de hardware. Este error se activa cuando tu kernel intenta asignar una multiplicación de matrices a la MXU con un tipo de datos que tu generación específica de TPU no tiene los circuitos físicos para ejecutar.

Por lo general, este error indica que el paso de Canonicalización del compilador, que intenta convertir automáticamente los tipos no admitidos en tipos admitidos (p.ej., a través de la emulación de software), no pudo encontrar una regla de conversión válida o no pudo hacerlo porque se inhabilitó el Modo de compatibilidad.

Depuración

Para resolver este error, debes alinear tus tipos de datos con las capacidades de tu hardware. Tienes las siguientes opciones:

1. Cómo convertir a tipos nativos

La solución más confiable es convertir manualmente tus operandos a un tipo de datos compatible con hardware (como Float32 o BFloat16 en TPU v4+) dentro de tu kernel antes de la operación matmul.

  • Motivo: Float32 es el tipo de datos universal que admite de forma nativa la MXU en todas las generaciones de TPU.
  • Compensación: Esto tiene un costo de VPU (unidad de procesamiento vectorial): los ciclos necesarios para realizar la conversión, pero garantiza que tu kernel se ejecutará en el hardware actual.

2. Cómo verificar el modo de compatibilidad

Por lo general, el compilador puede controlar automáticamente estos problemas de discrepancia de tipos en el modo de compatibilidad, que está habilitado de forma predeterminada. Verifica dos veces la configuración de XLA para asegurarte de que --xla_mosaic_compat_mode no esté configurado como falso.

Esto actúa como un "polyfill" que inyecta secuencias de emulación de software para las operaciones que tu hardware no admite de forma nativa.

Qué habilita el modo de compatibilidad:

  • MatMuls de precisión mixta: Permite combinar operandos de números enteros con acumuladores de números de punto flotante insertando automáticamente operaciones de conversión (p.ej., extender números enteros a Float32 antes de la multiplicación de matrices).
  • Emulación de baja precisión: En ciertas generaciones de hardware, emula tipos no admitidos, como el punto flotante 4-bit (4E2M1FN) o el punto flotante 8-bit (8E4M3FN), extendiéndolos a tipos admitidos, como BFloat16 o Float32, antes de la ejecución.

Ten en cuenta que este modo prioriza la compatibilidad por sobre el rendimiento máximo, ya que la emulación requiere instrucciones adicionales para convertir los formatos de datos antes de que la MXU pueda operar con ellos.

3. Actualiza el hardware o solicita asistencia

Si tu algoritmo requiere estrictamente un rendimiento nativo para tipos como Int4 o Float8 sin la sobrecarga de la conversión o la emulación, deberás ejecutarlo en una generación de TPU más reciente con compatibilidad nativa.

Solicitud de función: Si crees que tu hardware admite esta operación o si el compilador no encuentra una ruta de emulación válida incluso en el modo de compatibilidad, envía una solicitud de función. Por lo general, garantizamos que las operaciones sean compatibles con versiones posteriores. Por lo tanto, si tu kernel se ejecuta en una generación de TPU, debería ejecutarse en todas las generaciones futuras, pero no se garantiza que tenga emulación para generaciones anteriores (para algunas de las cuales las conversiones serían muy costosas).