Код ошибки: E2001

Категория: Время компиляции: Неподдерживаемый тип данных RHS на оборудовании

Эта ошибка возникает, когда тип данных, используемый для операнда правой части матрицы при умножении (например, 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 Backends: TPU

Обзор

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

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

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

Отладка

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

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

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

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

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

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

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

Что включает в себя режим совместимости:

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

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

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

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

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