Код ошибки: E2001

Категория: Compile Time: Unsupported RHSDataType on Hardware

Эта ошибка возникает, когда тип данных, используемый для правого операнда в матричном умножении (например, jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul или оператор @), не поддерживается нативно используемым поколением TPU.

Примеры сообщений об ошибках

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<...>

Бэкэнды XLA: TPU

Обзор

Матричный умножитель (MXU) TPU изначально поддерживает операции Float32 на всех поколениях оборудования.

Однако встроенная поддержка BFloat16 и других квантованных типов данных (например, Int4, Int8 или Float8) зависит от поколения оборудования. Эта ошибка возникает, когда ядро пытается сопоставить умножение матриц с MXU, используя тип данных, для которого в вашем поколении TPU нет физических схем.

Эта ошибка обычно означает, что этап канонизации компилятора, на котором выполняется попытка автоматически преобразовать неподдерживаемые типы в поддерживаемые (например, с помощью программной эмуляции), не смог найти подходящее правило преобразования или не смог выполнить преобразование, поскольку режим совместимости был отключен.

Отладка

Чтобы устранить эту ошибку, приведите типы данных в соответствие с возможностями оборудования. У вас есть следующие варианты:

1. Преобразование в нативные типы

Самый надежный способ исправить ошибку – вручную преобразовать операнды в поддерживаемый оборудованием тип данных (например, Float32 или BFloat16 на TPU v4+) внутри ядра до операции matmul.

  • Причина изменений. Float32 – это универсальный тип данных, который изначально поддерживается MXU во всех поколениях TPU.
  • Компромисс. Это связано с затратами на VPU (Vector Processing Unit) – циклами, необходимыми для выполнения приведения типов, но гарантирует, что ваше ядро будет работать на текущем оборудовании.

2. Проверьте режим совместимости

Обычно компилятор может автоматически устранять такие несоответствия типов в режиме совместимости, который включен по умолчанию. Проверьте конфигурацию XLA и убедитесь, что параметр --xla_mosaic_compat_mode не имеет значение false.

Это своего рода "полифил", который добавляет последовательности программной эмуляции для операций, которые не поддерживаются вашим оборудованием.

Что позволяет делать режим совместимости

  • Матричные умножения со смешанной точностью. Позволяет смешивать операнды типа Integer с аккумуляторами типа Float, автоматически вставляя операции приведения типов (например, расширяя целые числа до Float32 перед матричным умножением).
  • Эмуляция с низкой точностью. На некоторых поколениях оборудования эмулирует неподдерживаемые типы, например числа с плавающей запятой 4-bit (4E2M1FN) или 8-bit (8E4M3FN), путем их преобразования в поддерживаемые типы, например BFloat16 или Float32, перед выполнением.

Обратите внимание, что в этом режиме приоритет отдается совместимости, а не максимальной производительности, поскольку эмуляция требует дополнительных инструкций для преобразования форматов данных, прежде чем MXU сможет их обработать.

3. Как обновить оборудование или запросить поддержку

Если вашему алгоритму требуется именно производительность нативных типов, например Int4 или Float8, без дополнительных затрат на преобразование или эмуляцию, вам нужно использовать TPU более нового поколения с нативной поддержкой.

Запрос на добавление функции. Если вы считаете, что ваше оборудование поддерживает эту операцию или что компилятор не может найти подходящий путь эмуляции даже в режиме совместимости, отправьте запрос на добавление функции. Обычно мы гарантируем прямую совместимость операций. Поэтому, если ваше ядро работает на TPU определенного поколения, оно должно работать и на всех последующих поколениях, но нет гарантии, что оно будет эмулироваться на более старых поколениях (для некоторых из них приведение типов будет очень дорогостоящим).