# Apple MLX: стандартный GroupNorm обнуляет конечные FP16-входы

Локальный технический отчёт, 10 сентября 2026 года. Завершена проверка
наблюдения из предыдущего BatchNorm-аудита. В MLX 0.32.2 подтверждена
ошибка низкоточной статистики в стандартной ветви GroupNorm, подготовлен
и проверен изолированный патч. Ничего не опубликовано и не отправлено.

## Контрпример на четырёх числах

```python
import mlx.core as mx
import mlx.nn as nn
mx.set_default_device(mx.cpu)
x = mx.array([[[-128., 128., -128., 128.]]], dtype=mx.float16)
print(nn.GroupNorm(1, 4, affine=False)(x))
# Исходный MLX: [0,0,0,0]. Правильно после округления FP16: [-1,1,-1,1].
```

Среднее равно 0, дисперсия равна 16384. Входы, дисперсия и нормализованный
выход представимы в FP16. Но промежуточная сумма четырёх квадратов равна
65536 и переполняет этот тип до деления на 4. Выход становится нулевым.

| Вариант | Полученный результат |
|---|---|
| Исходный GroupNorm, стандартная группировка | `[0,0,0,0]` |
| Исходный GroupNorm, `pytorch_compatible=True`, одна группа | `[-1,1,-1,1]` |
| Исправленный GroupNorm, стандартная группировка | `[-1,1,-1,1]` |

[Измерения установленного wheel](wheel-reproduction.json),
[минимальное сравнение до/после](paired-reproduction.json),
[скрипт сравнения](reproduce.py). Добавка eps=10⁻⁵ учтена; точные вещественные
выходы чуть отличаются от ±1, но округляются к этим числам в FP16.

Подтверждены также неверные градиенты gamma и входа. Это не только потеря
точности при выводе: стандартный слой в ряде случаев превращает ненулевой
результат и производные в нули. В первом [пробнике](probe.py) пример
`[-256,0,512]` даёт нулевой выход и нулевой gamma-градиент.

## Почему нельзя просто переключить режим группировки

Публичный commit на момент чтения:
`81ba1c6a0e50a9268b931579c2d4f1158b9aab5a`.
В [GroupNorm](https://github.com/ml-explore/mlx/blob/81ba1c6a0e50a9268b931579c2d4f1158b9aab5a/python/mlx/nn/layers/normalization.py#L236)
стандартный и совместимый режимы намеренно задают разные группы каналов.
При C=6 и G=3, с нумерацией каналов от нуля:

| Группа | Стандартный режим: c mod G | Совместимый: floor(c/(C/G)) |
|---|---|---|
| 0 | 0,3 | 0,1 |
| 1 | 1,4 | 2,3 |
| 2 | 2,5 | 4,5 |

При G=1 или G=C состав групп совпадает. При промежуточных G простое
включение `pytorch_compatible=True` изменило бы математическую функцию.

В стандартном режиме mean/var вычисляются в dtype входа. Совместимый режим
уже вызывает `mx.fast.layer_norm`; его
[CPU fallback](https://github.com/ml-explore/mlx/blob/81ba1c6a0e50a9268b931579c2d4f1158b9aab5a/mlx/fast.cpp#L347)
рассчитывает статистику в float32. Отсюда наблюдаемое различие при G=1.

## Патч

[Изменён только GroupNorm._group_norm](groupnorm-low-precision.patch).
Для непустых FP16/BF16 входов после прежнего разделения на группы ось
групп переставляется, чтобы нормализация каждой группы шла по последней
оси. Затем вызывается `mx.fast.layer_norm` без собственных affine-параметров,
и оси возвращаются обратно.

Состав групп, spatial-оси, внешние gamma/beta, итоговый dtype и форма
сохраняются в проверенных случаях. Float32/float64, integer/complex и
уже существующий совместимый путь оставлены прежними. Пустые входы тоже
идут по прежней ветви — важная поправка после проверки первой версии.

Тесты исполняют точные сохранённые Python-модули до/после на установленном
бинарном MLX 0.32.2, CPU. Это не полная сборка текущего main.
[Метаданные и SHA256](source-metadata.json),
[исходный класс установленного wheel](installed-groupnorm-source.py),
[построение патча](build_patch.py), [исправленный модуль](patched-normalization.py).
Часть публичных файлов повторно использована из предыдущего пакета только
после совпадения точного commit и контрольных сумм.

## Независимые проверки

[Эталон](audit_common.py) задаёт группы непосредственно через принадлежность
каналов, без повторения reshape/transpose реализации. Средние и дисперсии
вычисляются скалярной Decimal-арифметикой с точностью 70 цифр по уже
квантованным входам. JVP/VJP и gamma/beta проверяются аналитическими
производными нормализации по соответствующей группе.

| Набор | Объём | До | После |
|---|---|---:|---:|
| [Основной](regression.py): 332 сценария forward, 24 сценария градиентов | 1424 assertions | 45 несовпадений | 0 |
| [Совместимость](compatibility.py) | 25 общих проверок | 0 | 0 |
| Дополнительные сравнения неизменённых путей и крайних входов | 16 проверок после | — | 0 |

1424 — число assertions, не отдельных векторов. Из них 428 сравнивают
численные результаты, остальные проверяют dtype, форму и неизменность
входов. 45 исходных несовпадений: 35 выходов, 4 входных VJP, 3 gamma-градиента,
3 входных JVP; 36 случаев float16 и 9 bfloat16. Часть относится к точности
статистики, поэтому эти числа не равны количеству переполнений или отдельных
дефектов.

Проверены два режима группировки, G=1/2/3 и singleton-группы, NC/NLC/NHWC,
четыре dtype, affine on/off и неодинаковые gamma/beta, малые и большие
амплитуды, константные входы и малые положительные eps. Включены
переставленные spatial-оси, обращённые каналы и broadcast batch.

Совместимость включает точный публичный `test_group_norm` с настоящими
unittest-assertions, явную проверку состава групп, пустые входы, NaN/Inf,
неизменённые dtype-пути и большой общий сдвиг в стандартном float64-режиме.
Для исходного совместимого float64-пути допуски учитывают его существующую
float32-арифметику; это не проверка полноценной точности float64.

[Основные результаты](regression-results.json),
[совместимость и оставшиеся ограничения](compatibility-results.json),
[журнал последнего прогона](compatibility-final.log).
CPU-время последних успешных наборов — примерно 0.106 и 0.018 секунды.
Это время тестов, не benchmark скорости слоя на модели.

## Выявленная регрессия первого прототипа

Первая версия прошла основной набор непустых входов, но направляла пустой
массив `(2,0,6)` в быстрый путь. Вычисление этого пути не завершилось
в пределах CPU-лимита. Исходный GroupNorm возвращал `[[],[]]`.

Первый широкий запуск совместимости остановлен лимитом 30 CPU-секунд;
диагностический — лимитом 5. Отдельная проверка первого прототипа на пустом
входе остановлена через установленный лимит 2 CPU-секунды. После добавления
проверки `x.size` пустой массив снова возвращается корректно, полный набор
совместимости завершается за доли секунды.

Сохранены [первая версия](patched-normalization-v1.py),
[её патч](groupnorm-low-precision-v1.patch),
[успешные непустые тесты v1](regression-v1-results.json),
[диагностический журнал](compatibility-diagnostic.log),
[исходный пустой вход](empty-before.log),
[пустой вход v1](empty-after-v1.log),
[пустой вход после исправления](empty-after-final.log),
[статусы ограниченных запусков](bounded-failures.json).
Эти остановленные процессы не засчитаны как успешные тесты. Бесконечность
внутреннего цикла отдельно не доказывалась; установлен факт незавершения
в пределах заданных лимитов.

Отдельная ошибка первого численного стенда: практически нулевой остаток
производной −3.97·10⁻¹⁴ сравнивался с нулём в неизменённой float32-ветви
без достаточного допуска на сокращение. Добавлена абсолютная оценка
8·epsilon_float32·масштаб производной. [Первый результат](regression-first-results.json)
и [заметки стенда](harness-notes.txt) сохранены. Этот случай не считается
ошибкой MLX; основной патч для его устранения не менялся.

## Проверка дубликатов

Четыре [GitHub-поиска](duplicate-search.json) вернули шесть уникальных
issues/PR. Прочитаны их тексты и 12 доступных обычных комментариев.
Комментарий к #3613 не удалось получить после DNS/connection timeout;
это ограничение поиска сохранено в [журнале запросов](duplicate-comments.json)
и [первом ответе](duplicate-comments-first.json).

- [#3653](https://github.com/ml-explore/mlx/pull/3653) исправляет проверку
  числа групп и размерностей. Наши входы удовлетворяют этим ограничениям.
- [#4312](https://github.com/ml-explore/mlx/pull/4312) касается eps.
  В контрпримере используется обычное положительное eps=10⁻⁵.
- #3819 и #3613 затрагивают документацию нормализации и другой loss.
- [#2490](https://github.com/ml-explore/mlx/issues/2490) обсуждает снижение
  точности/памяти модели, без этого минимального контрпримера.
- [#3702](https://github.com/ml-explore/mlx/issues/3702) в итоге локализован
  авторами в стороннем custom Metal GroupNorm с отсутствующим барьером.
  Нынешний случай воспроизведён на CPU в штатном Python-слое MLX.

Соседний [mlx-examples #1434](https://github.com/ml-explore/mlx-examples/issues/1434)
касается разных eps при переносе Stable Diffusion. Здесь eps у сравниваемых
функций одинаково. Численный механизм родственен уже известному
[FP16 InstanceNorm #4228](https://github.com/ml-explore/mlx/issues/4228)
и предыдущему локальному BatchNorm-отчёту.

Точного дубликата этого стандартного GroupNorm-контрпримера в просмотренных
материалах не найдено. Это подтверждение ещё одного затронутого слоя;
абсолютная новизна класса ошибок или право на выплату не установлены.

## Границы результата и точка продолжения

Патч не исправляет дисперсии за пределами float32. На BF16/float32 входах
порядка 2⁸⁰ нулевые ответы остаются; они записаны в `remaining_limits`.

Также зарегистрировано соседнее наблюдение: существующий совместимый
float64-путь при общем сдвиге 2⁴⁰ теряет различия координат из-за приведения
в fast.layer_norm к float32. Для смещений `[-2,-1,1,2]` он даёт четыре нуля
вместо примерно `[-1.264909,-0.632454,0.632454,1.264909]`. Стандартный
float64-путь остаётся корректным в проверке. Это отдельная точка продолжения;
её собственный патч и проверка дубликатов здесь не завершены.

GPU, compiled-путь, производительность на больших массивах и полное обучение
модели не проверены. Новый transpose/fast-путь требует соответствующего
ревью и измерения перед использованием на больших моделях.
Патч остаётся исследовательским. Рабочая копия MLX не изменена.

## Воспроизведение

Из этой папки, последовательно:

```sh
python3 build_patch.py
/Users/khamit/Documents/math/software-audits/current-stack-2026-09-07/.venv/bin/python reproduce.py
/Users/khamit/Documents/math/software-audits/current-stack-2026-09-07/.venv/bin/python regression.py
/Users/khamit/Documents/math/software-audits/current-stack-2026-09-07/.venv/bin/python compatibility.py
python3 validate_artifacts.py
```

Скрипты задают CPU и один численный поток. Лимит основного набора —
30 CPU-секунд, совместимости — 5, отдельной диагностики пустого входа — 2.
GPU и параллельные вычислительные процессы не запускались.
Для повторения тестов сеть не требуется. Доступ за пределами песочницы
использован для инициализации установленного MLX и отдельно для чтения
публичных источников. [Проверка артефактов](artifact-validation.json),
[контрольные суммы](SHA256SUMS.json).