В MLX везде float32, а ответы расходятся: проверьте два вычислительных пути

Одинаковый float32 на входе и выходе ещё не означает, что весь расчёт внутри MLX выполнен с полной точностью этого формата. Если после смены размера пакета ответы нейросети слегка разошлись, сначала проверьте выбранный вычислительный путь: причина может находиться ниже уровня модели.

Тип массива — часть картины

В документации MLX 0.32.4 описан режим пониженной точности для семейства матричных операций на подходящем оборудовании. Он может затрагивать умножение матриц, свёртки и внимание. Массивы при этом остаются float32. Конкретный набор операций и величина расхождений зависят от устройства и вычислительного backend — программного слоя исполнения.

Представьте проверку распознавания изображений: вы подали один и тот же кадр отдельно, а затем вместе с двумя другими. Веса и подготовка изображения одинаковы, но выходные значения немного отличаются. Это повод локализовать расхождение, прежде чем переобучать модель или менять данные.

На Metal отдельный кандидат — алгоритм Winograd для свёрток 3×3 с шагом 1. MLX выбирает его с учётом общего размера входа, включая размер пакета. Поэтому одинаковый кадр в пакетах из одного и трёх изображений может пройти разными вычислительными путями. Это объяснение возможного механизма, а не диагноз любой такой ошибки.

Меняйте одну настройку за раз

Сначала сохраните входы, веса, версию MLX, устройство и результаты исходного запуска. Для своей задачи заранее задайте допустимую абсолютную и относительную погрешность. Рядом с разницей выходных массивов проверьте прикладной результат: например, изменился ли выбранный класс. Малое численное расхождение и ошибка решения — разные показатели.

Затем повторите свой проверочный скрипт в новом процессе. Согласно справочнику переменных окружения, MLX_ENABLE_TF32=0 отключает разрешение пониженной точности соответствующих операций. Отдельный запуск с MLX_CONV_WINOGRAD=0 переключает свёртки с Winograd на GEMM. Ниже check.py обозначает ваш скрипт с зафиксированными данными:

MLX_ENABLE_TF32=0 python check.py
MLX_CONV_WINOGRAD=0 python check.py

Не включайте оба переключателя сразу в первом сравнении: так труднее понять причину изменения. Для примера с изображением сравните одиночный и пакетный режимы при каждой настройке, сохраняя сам кадр неизменным. Если расхождение осталось, продолжайте искать источник; отсутствие эффекта тоже полезно сужает проверку.

Выбор по результату, а не по названию режима

Оставлять более строгий режим постоянно имеет смысл, если он помогает выполнить ваш критерий точности при приемлемом времени работы. Измерьте обе стороны компромисса на своей нагрузке. Эти флаги не обещают побитовой воспроизводимости всей нейросети: полезный итог проверки — установленная причина расхождения и понятный допуск для конкретной задачи.