Sakana AI у співпраці з NVIDIA: дозволяє GPU пропускати 80% неефективних обчислень великих моделей, прискорення обробки H100 на 30%

За даними моніторингу Beating, Sakana AI у співпраці з NVIDIA відкрили формат розріджених даних під назвою TwELL та відповідний прискорювальний ядро, що успішно дозволяє GPU при роботі з великими моделями пропускати ті «недійсні обчислення», результати яких близькі до нуля. Ця схема без втрати точності моделі дозволяє підвищити швидкість обробки H100 до 30%, максимізувати швидкість тренування до 24% та значно зекономити пікову пам’ять відеокарти.

Передавальні шари великих моделей (FFN) споживають більшу частину параметрів і обчислювальної потужності. Але насправді, під час кожного генерування тексту, понад 80% нейронів перебувають у «сплячому режимі» (з активністю, близькою до нуля), не роблячи внеску у кінцевий результат. Якщо можна пропустити ці нейрони, можна зекономити величезну кількість обчислювальних ресурсів. Однак сучасні GPU природно оптимізовані для обчислень щільних матриць, і при традиційних методах відбору розріджених даних витрати на пошук і зчитування даних з пам’яті з’їдають усю економію.

Формат TwELL саме створений для подолання цього «жорсткого» обмеження апаратного забезпечення. Він повністю відповідає паралельній логіці GPU: замість того, щоб збирати розріджені дані по регіонах, він ділить дані на маленькі блоки (tile), які GPU найкраще обробляє. Таким чином, обчислювальні ядра GPU можуть безпосередньо пакувати корисні дані локально, повністю усуваючи часозатратне глобальне зчитування з пам’яті та ідеально інтегруючись у сучасну конвеєрну систему прискорення чипа.

У тестах на модель з 1,5 мільярда параметрів достатньо було додати невелике регуляризаційне обмеження під час тренування, щоб знизити пропорцію нейронів, що дійсно потребують обчислень, до менш ніж 2%, а результати за сімома завданнями знижувалися незначно. Дані також виявили закономірність: чим більша модель, тим більше нейронів перебуває у сплячому режимі (у моделі з 2 мільярдами параметрів частка розріджених нейронів на 38% менша, ніж у моделі з 500 мільйонами). Це означає, що у майбутньому, при створенні ще більших моделей, ця оптимізація для апаратного забезпечення відкриє ще більш значні можливості для підвищення продуктивності.