Трансформеры жрут память как не в себя. Это не новость. Но пока все паковали квантизацию и тащили flash-attention, инженеры из лаборатории системы машинного обучения MIT решили зайти с другого конца. Они предлагают переписать весь блок трансформера - нормализацию, активации, остаточные связи - как единый GEMM-эпилог. Звучит как ересь? Возможно. Но цифры говорят сами за себя.
CODA (Compressed Operator via Data-Aware epilogue) - метод, который объединяет все bottleneck-операции внутри блока трансформера в один fused kernel, выполняемый как постобработка после основного матричного умножения.
Если вы когда-нибудь писали свой трансформер вручную, вы знаете эту боль: слой за слоем вы читаете и пишете одни и те же тензоры. LayerNorm - обращение к памяти. GELU - обращение. Dropout - еще одно. Residual - давай еще. Добавьте сюда автоматическую дифференциацию, которая плодит промежуточные градиенты, и получаете классический memory-bound сценарий, где чистый FLOPs простаивает в ожидании данных.
Что предлагает CODA и почему это не очередной велосипед
Команда исследователей заметила, что все эти операции - нормализация, элементные функции активации, остаточные связи - математически могут быть выражены как операции над выходом одного большого GEMM. Суть проста: вместо того чтобы выполнять пять разных kernel'ов, каждый из которых загружает и выгружает данные через глобальную память, они сливают их в один этап.
Это достигается за счет расширения так называемого "epilogue" - финальной стадии типичного GEMM-ядра, обычно отвечающей за простые бродкасты или аккумуляцию bias. CODA делает epilogue программируемым и дает ему доступ ко всем внутренним состояниям матричного умножения. На практике это означает, что вы можете внутри одного kernel'а выполнить LayerNorm + GELU + residual, не сохраняя промежуточные буферы.
Важный нюанс: CODA не трогает self-attention. Он работает только с feed-forward и transformer block в целом. Но именно эти части занимают до 60% времени обучения в моделях с большим hidden size.
Результаты, которые они приводят в препринте (а полноценного репозитория, к сожалению, пока нет), впечатляют: на модели GPT-2 XL (1.5B) обучение ускорилось на 27% при сохранении точности. На маленьких моделях (например, 350M) выигрыш скромнее - около 12%, что объясняется меньшим размером матриц и, как следствие, меньшим числом memory-bound операций.
| Модель | Ускорение обучения | Снижение используемой памяти | Потеря точности |
|---|---|---|---|
| GPT-2 XL (1.5B) | 27% | ~18% | <0.1% |
| GPT-2 Medium (350M) | 12% | ~10% | <0.05% |
| LLaMA-7B (адаптированная) | 22% | ~15% | <0.1% |
Как это связано с тем, что мы уже умеем
Тренд на fusion операций в deep learning не нов. FlashAttention, cuBLAS epilogue, Triton - все это движется в сторону уменьшения количества обращений к HBM. Написать transformer с нуля на CUDA - это прекрасный способ прочувствовать, как операции гуляют по памяти. CODA идет дальше: он не просто фузит операции внутри одного ядра, а полностью переосмысливает логику вычислений.
Особенно интересно это выглядит в контексте альтернативных архитектур. Например, сверточные декодеры против трансформеров - там авторы вообще избавляются от self-attention, но проблемы memory-bound остаются в feed-forward. CODA может дать второе дыхание старым добрым трансформерам, прежде чем их окончательно похоронят.
Почему я пока не хлопаю в ладоши
Да, идея красивая. Да, цифры красивые. Но есть "но": на сегодняшний день нет публичного репозитория с кодом. Препринт - это хорошо, но без воспроизводимых результатов и бенчмарков на современном железе (H100/B200) доверия мало. Исследователи использовали A100, а на H100 с его трансформер-движком картина может быть иной.
Кроме того, CODA требует изменений в самом ядре GEMM. Это не drop-in решение. Для интеграции в PyTorch или JAX придется писать пользовательские CUDA/Triton ядра. Если у вас нет опыта - можно попробовать научить маленькие модели генерировать ядра, но это пока экзотика.
Ирония в том, что многие компании уже давно используют проприетарные версии похожих подходов. NVIDIA cuDNN, например, имеет встроенные fused kernels для transformer blocks. Но открытого, универсального метода с детальным анализом не было. CODA - первый шаг к стандартизации.
Что дальше: прогноз на июнь 2026
Скорее всего, метод либо выльется в open-source реализацию на Triton, либо останется в виде академической статьи, которую будут цитировать пару лет. Если авторы выложат код - обязательно протестируйте на своей задаче. Особенно если вы запускаете 80B модель на ограниченном бюджете памяти - экономия в 15% может быть решающей.
Мой совет: не ждите готовых решений. Берите идею CODA, открывайте Triton и пишите свой fused epilogue. Заодно вспомните, как работает transformer изнутри - создание своего трансформера на C++ с нуля даст вам понимание, где теряется производительность. А потом вы сможете сделать то, что обещают авторы CODA - на практике.
