Интересное что-то

Forwarded from .ml

Скучали? А мы-таки собрали пост про неэффективное использование вычислительных ресурсов ⬇️

Если вы хотите выжать максимум из своей GPU, нужно знать, как устроена видеокарта. Если совсем на пальцах, у неё есть:

📌 DRAM — большая, но медленная память. Расположена на отдельном чипе или чипах. В современных GPU объём может достигать 80 ГБ и больше.
📌 Streaming Multiprocessors (SM) — непосредственно вычислительные модули с CUDA и Tensor-ядрами. Позволяют запускать операции параллельно, распределяя пайплайны вычислений между собой.
📌 SRAM — быстрая, но маленькая память (обычно сотни Кб). Находится внутри вычислительных блоков.

Чтобы выжать максимум производительности, нужно учитывать особенности архитектуры. Чтобы видеокарта что-то посчитала, ей нужно, чтобы данные для вычисления оказались в SRAM.

Но SRAM маленький, и хранить там все данные невозможно. Поэтому обычно данные сначала копируются из DRAM в SRAM, затем производится вычисление, а после этого результат снова копируется из SRAM в DRAM.

Например, стандартная цепочка PyTorch-вызовов (без TorchDynamo) всегда будет работать так, что будет происходить перегонка байтиков туда-сюда: DRAM-SRAM-DRAM-SRAM-... .

Но ведь есть такие операции, которые можно спокойно выполнить без копирования промежуточных результатов в DRAM и из DRAM. Например, перемножить один кусочек матрицы на другой, и затем применить к результату функцию активации. Такой операции будет достаточно одного только копирования входных данных из DRAM и сохранения итогового результата в DRAM, а матричное умножение и, например, ReLU можно применить друг за другом, используя лишь SRAM.

📝 Для более тонкого контроля над памятью можно писать кастомные GPU-ядра: с нуля, используя библиотеки CUDA или с помощью Triton.

Что такое Triton? 🛠

Это программный интерфейс, который позволяет писать кастомные GPU-ядра без прямого использования CUDA.

Чем он хорош:

📍Код пишется на некотором сабсете Python, следовательно порог входа не такой высокий, а работает это дело через JIT-компиляцию.
📍Обеспечивает гибкий контроль над памятью и параллелизмом — мы сами решаем, когда ходить в DRAM, а когда не ходить, и имеем больше контроля над тем, как мы будем параллелить вычисления.

Возвращаясь к примеру выше, Triton позволяет реализовать выполнение некоторой последовательности операций без копирования промежуточных результатов в/из DRAM.

Как написать кернел?

Кернел — это функция с декоратором triton.jit. Их главная особенность — маленькие программки, которые могут быть запущены параллельно. Каждая запущенная копия будет иметь свой pid-идентификатор. Можно сделать так, чтобы каждый идентификатор обрабатывал не все данные, а отдельный фрагмент (вспоминаем, что SRAM-то маленький, а ещё нужно как-то уметь в параллелизм между блоками).

Пример кернела из официальной документации.

В кернеле мы используем:

👾 Указатели на входные и выходные данные.
👾 Размер вектора.
👾 Размер блока (Block Size) — количество элементов, которые обрабатываются одним PID-ом.
👾 PID — идентификатор запущенной программки.

Также в примере используют бинарную маску при чтении и записи в DRAM, чтобы не выйти за пределы нужной нам памяти — например, когда размер вектора не кратен параметру block size.

А также, нужно реализовать небольшой враппер для того, чтобы запускать наш кернел — ниже оставлю пример кода 👇

32 views18:16