Работаем со свертками в PyTorch с помощью библиотеки CUTLASS и алгоритма Implicit GEMM

Библиотека CUTLASS — это набор C++ шаблонов для реализации высокопроизводительного GEMM в коде. Она предоставляет структурные блоки, из которых можно собрать или просто вызвать операцию GEMM. Поддерживает вычисление смешанной точности, использование TensorCores и других примитивов, доступных для быстрого вычисления. В отличие от cuBLAS, это open source-библиотека. Ее относительно просто интегрироватьь и модифицировать под свои задачи. 

Подход основан на тайлинге — разбиении задачи на мелкие подзадачи и распараллеливании. Как устроен процесс работы:

• Процедура начинается с работы с глобальной памятью: выбираем блоки данных (тайлы) из глобальной памяти для умножения матриц.

• Затем используются примитивы для переноса этих данных в shared-память, где происходит тайлинг на этом уровне.

• После выполняется работа на уровне варпов и регистров с использованием TensorCores или CUDA Cores. 

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

Библиотека предоставляет различные уровни специализации и конфигурирования: Device-level, Kernel-level, Block-level, warp, Instruction. Весь API представлен в виде шаблонов, из которых можно набирать те типы, которые потом инстанцируются для реализации нужного тайлинга. 

Какие еще инструменты могут расширить функциональность PyTorch для работы с большими свертками? Как выбрать алгоритм, подходящий для обучения моделей? Узнаете из статьи →