Title: Помилка полягала в одній осі

Прихована помилка в реалізації chunked-scan у PyTorch дозволяє певним гібридним моделям підглядати майбутні токени щоразу, коли відсутні швидкі об'єднані ядра (fused kernels), що ставить під загрозу такі чекпоінти, як Zamba2-1.2B та Nemotron-H-8B. Дефект спричинений однією редукцією, виконаною не за тією віссю, і він виникає на тому шляху виконання, на який покладаються більшість CI-пайплайнів та запусків на CPU.

Чому ця помилка важлива

Гібридні моделі та моделі простору станів (state-space models) більше не залежать виключно від механізму attention; вони також використовують лінійні рекуренції, скани (scans) та згортки (convolutions). Операція маскування, яка блокує майбутні токени в матриці attention, не гарантує автоматично причинно-наслідковий зв'язок (causality) для цих додаткових операцій. Коли швидкі об'єднані ядра, які зазвичай правильно обробляють scan, відсутні, PyTorch перемикається на шлях chunked-scan на чистому Python. Цей резервний шлях містить помилку переплутування осей (axis-mixup), що дозволяє інформації з пізніших позицій протікати назад під час інференсу.

Цей витік є непомітним. Він не призводить до збою моделі та не викликає очевидної помилки. Натомість він робить показники loss та perplexity штучно низькими, оскільки модель фактично «шахраює» — бачить ті самі токени, які вона має передбачити. Будь-яка подальша оцінка, яка покладається на ці метрики, будується на хибному фундаменті.

Як було виявлено проблему

Дослідники порівняли два прямі проходи (forward passes) через одну й ту саму модель:

  1. Випадкова послідовність токенів.
  2. Ідентична послідовність із одним зміненим токеном.

Вони виміряли різницю в прихованих станах (hidden states) шар за шаром. Перевірка масок не виявила нічого, але пошаровий аудит із внесенням помилок (fault injection) виявив 192 із 192 внесених помилок, що підтвердило наявність витоку.

Швидка перевірка бібліотеки transformers показала:

  • Zamba2-1.2B має витік, коли розмір чанка (chunk size) встановлено на 256.
  • Nemotron-H-8B має витік при розмірі чанка 128.
  • Більшість інших чекпоінтів не показали витоку за тих самих умов.

Дефект криється у коді chunked-scan, який виконується щоразу, коли відсутні необов'язкові об'єднані ядра. Це включає:

  • Усі виконання на CPU.
  • Середовища GPU без специфічних пакетів fused-kernel.
  • Стандартні інсталяції PyTorch, які не містять додаткових залежностей.

Оскільки помилка з'являється лише за відсутності цих ядер, вона може виникати в середовищах CI та на процесорах CPU.

Хто під загрозою та яка ціна

Будь-яка команда, яка навчає, донавчає (fine-tunes) або оцінює гібридні моделі без використання об'єднаних ядер, ризикує опублікувати завищені показники продуктивності. Уявне покращення loss або perplexity є ілюзорним; модель фактично «зазирнула вперед». Для дослідницьких груп це може призвести до введення в оману тверджень про досягнення результатів рівня state-of-the-art. Для комерційного впровадження це може спричинити помилки під час генерації, які модель насправді ніколи не вивчала.

Контраргумент

Деякі розробники стверджують, що правильно застосованої причинно-наслідкової маски (causal mask) достатньо для запобігання будь-якому витоку майбутніх токенів. Ця помилка спростовує таке твердження: скани, згортки та певні шари нормалізації можуть повністю обійти маску. Переплутування осей у chunked-scan показує, що причинно-наслідковий зв'язок має забезпечуватися всюди, де протікають дані, а не лише в матриці attention.

На що звернути увагу далі

  • Гігієна залежностей: встановлюйте пакети швидких об'єднаних ядер на всіх вузлах навчання та інференсу, особливо в CI-пайплайнах.
  • Скрипти аудиту: дотримуйтесь двоетапного аудиту, рекомендованого дослідниками:
    • Внесіть відому помилку в чекпоінт, який ви тестуєте, як позитивний контроль.
    • Запускайте послідовності, довжина яких перевищує розмір чанка або вікна моделі; короткі послідовності ніколи не виявлять витік.

Проведення аудиту на будь-якому гібридному або state-space чекпоінті з довжиною послідовності, що перевищує розмір чанка, покаже, чи залишається модель вразливою.

Висновок: Навіть модель, яка проходить усі стандартні тести, може непомітно «шахраювати», якщо шлях виконання перемикається на помилкову реалізацію. Перевірка наявності швидких об'єднаних ядер та явне тестування на витік за межами масок attention тепер є необхідними кроками перед тим, як довіряти метрикам будь-якої гібридної моделі.