Comments 13
Что то на сложном. Начал читать, подумал что тут triton - это triton server nvidia.
Классная статья, для понимания в образовательных целях.
Автор красавчик)
Спасибо) На самом деле реализовать такой проект ещё гораздо сложнее, чем его читать) Просто сейчас тема ускорения трансформеров очень актуальная - самый сок, как говорится, однако по ней практически нет нормальных туториалов, а текущие реализации имеют много ограничений. Поэтому хочется надеяться, что данный проект принесёт пользу многим ml- разработчикам.
но и для студентов благодаря возможности работать в Google Colab и Kaggle,
У меня вопрос, как студента в специализации Data Science. А можно ли провернуть такой же проект с созданием своей LLM (или даже VLM), но на Jupyter Notebook или Marimo? Просто мне нравится делать проект а-ля визуализация данных на юпайтере, но можно ли провернуть такое на реализацию своего аналога Gemini или ChatGPT?
Спасибо, дорогой Друг!!! Ты спас мою 2080 ti 22gb от забвения!
занятия по 10 часов в сутки результат дал все-таки.
Планируется ли добавить поддержку видеокарт из "красного" лагеря (AMD) и NPU, или это просто принципиальные ограничения выбранного движка?
Хороший вопрос: я тоже про это изначально думал, но пока что решил отказаться по двум причинам:
1) У меня нет AMD разных поколений для тестов :)
2) По правде говоря, пока что я не знаю точно сколько нужно внести изменений в исходных код, чтобы сделать поддержку AMD: возможно придётся переписывать flash attention с полного нуля под эту задачу, что в таком случае займёт много времени, которого как всегда нет.
Если вы хотите запустить flash attention на AMD, то это можно сделать через проект tri dao, но там будут те же ограничения, что и в других реализациях.
Огромное спасибо за поддержку Turing и Windows. Это буквально демократизация современных LLM-оптимизаций, так как официальный FA2 отрезал огромный пласт разработчиков на T4 (Colab/Kaggle) и потребительских картах 20-й серии. Отличный пример того, как research engineering решает реальные боли сообщества.
Спасибо за интересный пост!
А какое ускорение достигается в сравнении с наивной реализацией attention?
Было бы полезно иметь цифры для операции внимания отдельно, и end-2-end ускорение для какого-нибудь трансформера (типа Квена-3).
Очень хороший и важный вопрос, ответ на который будет очень большим и не однозначным, но перед этим хочу упомянуть особенности, которые влияют на проведение качественных тестов производительности:
Начиная с версии 3.3.0, в triton появился баг, связанный с увеличенным потреблением shared memory, что особенно заметно на turing и ampere: ранее помещавшиеся блоки в память теперь туда не помещаются. Для решения данной проблемы пришлось прибегнуть к созданию двух разных конфигураций с их автоматическим определением для версий triton до 3.2.0 включительно и после - отсюда и появились 2 режима установки: modern и legacy. Прикол в том, что страрые gpu типа turing и ampere в таком в modern режиме работают субоптимально на уровне vanilla attention для небольших моделей (до сотни миллионов параметров) и немного быстрее для больших (около миллиарда) моделей (GPT, BERT). Поэтому стоит учитывать, что для таких режимов производительность будет разной. Логичным будет вопрос: зачем тогда вообще нужно было писать конфиги для старых gpu в modern режиме? Это необходимо для случаев, если у вас в связке используются разнородные gpu, например ampere и blackwel. О том как тренировать модели в таком случае - будет немного ниже.
Результаты будут сильно зависеть как от количества параметров модели, так и от размера данных, не говоря уже о параметрах для обучения.
Для проведения таких тестов необходимо большое количество различных gpu и желательно по несколько штук, чтобы сравнить результаты не только на больших данных, но и в разных режимах и на разных моделях с разным количеством параметров. Пока что это физически и финансово невозможно (для меня) обеспечить такую материальную базу. Именно поэтому на данный момент раздел с бенчмарками отсутствует (об этом я писал на github в конце readme), а если проект пойдёт в гору, то позже этот раздел появится.
Я специально сделал такое большое вступление, чтобы у вас или читателей не возникло впечатление, что я халатно отношусь к проекту - на самом деле наоборот :) Теперь по поводу результатов: из того что я тестил на своих данных (их было меньше чем в тестах tri dao), на ampere (6xRTX 3090) получилось быстрее примерно в 5.8 раз для GPT-2 small и в 5.2 для GPT-2 medium в сравнении с vanilla attention, на hopper и blackwell цифры были практически одинаковые: в 7 раз на small и 6.6 на medium. Сравнивая его с другими реализациями flash attention на небольших моделях и данных, результаты были плюс-минус похожими.
Данный алгоритм я сделал задолго до того как его выложил в open-source и он уже использовался в проектах, когда приходилось создавать и учить с нуля бота по аналогии с ChatGPT на разнородных GPU. В связи с этим пришлось организовать обучение таким образом, чтобы нагрузка на gpu распределялась не равномерно (иначе самая быстрая будет ожидать самые медленные), а в таких пропорциях, чтобы они завершали вычисления практически в одно и тоже время (минимизировать простой). Конкретно в этом случае у меня не было возможности провести сравнение скорости обучения в сравнении с vanilla attention, но даже так понятно, что прирост был ощутимый. К сожалению, большего пока не могу сказать по этой теме. Как-то так :)
Ускоряем LLM по максимуму. Как я создал кросс-платформенный Flash Attention с поддержкой Turing+ архитектур и не только