[Перевод] Оптимизация ядра WebGPU для перемножения матриц и достижения производительности свыше 1ТФЛОПС

info@vsetut.pro

Стать автором

Вернуться

05.12.2024 21:57:21 | Хабр

Хабы: Программирование, Высокая производительность, Python, Процессоры

Пишем Surfgrad, высокопроизводительную библиотеку для автоматического дифференцирования выражений при помощи WebGPU.

Я работаю в компании Nomic, и многие из моих коллег заняты созданием больших TSNE-подобных визуализаций, работающих в браузере. При визуализации таких двумерных карт возникает две проблемы: проецировать эти конструкции (напр. TSNE и UMAP) в 2D-координатную систему протекает медленно и требует больших затрат оперативной памяти, особенно по мере того, как вы увеличиваете датасет и пытаетесь визуализировать в браузере миллионы точек данных, не расплавив при этом ноутбук невзначай.

Отобразить в браузере миллионы точек данных, не расплавив компьютер — та ещё задача. Мне доводилось слышать, что многие проблемы с масштабированием удаётся решать при помощи инструмента Deepscatter, разработанного Беном Шмидтом.

Но многие из таких разговоров, которые мне известны, вертятся вокруг Typescript и великолепия WebGPU как такового. Готовя эту статью, я не смог найти ни одной библиотеки для автоматического дифференцирования выражений, которая была бы написана с применением WebGPU. Но было бы упущением не назвать здесь два репозитория с функционально схожим наполнением: webGPT (библиотека на основе трансформеров, приспособлена только для логического вывода) и webgpu-blas (ядра для быстрого перемножения матриц под webGPU). Поэтому, в качестве самообразования и желая получше изучить WebGPU и Typescript, я решил написать Surfgrad, высокопроизводительную библиотеку для автоматического дифференцирования выражений под управлением WebGPU. Она обеспечивает тензорные операции в браузере. Как понятно по названию и по принципу работы, она во многом сделана по примеру tinygrad и micrograd.

Подробнее