Добро пожаловать в третью часть нашего углубленного руководства по профилированию в PyTorch. Если вы следили за нашей серией, то уже знаете, как читать трассировки профилировщика и использовать их для оптимизации кода. В первой части мы разобрали базовые математические операции, такие как сложение и умножение, увидев, как таблица профилировщика выявляет «узкие места» (hotspots), а трассировка показывает порядок выполнения алгоритмов во времени. Во второй части мы обернули эти операции в слой `nn.Linear`, собрали многослойный перцептрон (MLP) и профилировали как стандартные, так и вручную настроенные ядра (kernels).
С точки зрения архитектуры Transformer, логичным следующим шагом для профилирования является еще один фундаментальный алгоритм — механизм внимания (attention). Несмотря на свою печально известную квадратичную сложность, существует множество хитрых приемов, которые делают его быстрым. Наша цель здесь — не охватить каждый трюк детально, а показать, как каждый из них выглядит под микроскопом профилировщика. Мы будем использовать GPU для запуска скриптов, но вы можете легко повторить эти эксперименты на инфраструктуре Hugging Face, используя Dev Mode с Spaces.
01Наивное внимание: разбор по косточкам
Механизм внимания оперирует тремя основными компонентами: Запросами (Queries, Q), Ключами (Keys, K) и Значениями (Values, V). Взаимодействие между ними можно записать как короткую последовательность шагов:- Построение оценок внимания (scores): умножение матриц $Q \cdot K^T$.
- Масштабирование оценок: умножение на скаляр $1/\sqrt{d_k}$.
- Применение причинной маски (causal mask): замена определенных элементов на $-\infty$, чтобы модель не «видела» будущие токены.
- Нормализация с помощью softmax: получение весов внимания (attn).
- Перевзвешивание значений: умножение весов внимания на значения $V$.
class NaiveCausalAttention(nn.Module):
def __init__(self, head_dim):
super().__init__()
self.scale = 1.0 / math.sqrt(head_dim)
def forward(self, q, k, v, mask):
scores = torch.matmul(q, k.transpose(-2, -1))
scores = scores * self.scale
scores = scores.masked_fill(mask, float("-inf"))
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, v)
return out- Ядро matmul (для $Q \cdot K^T$).
- Ядро mul (для масштабирования).
- Операцию для маскирования.
- Ядро softmax.
- Ядро matmul (для $Attn \cdot V$).
mul: масштабирование.masked_fill: причинная маскировка.softmax: ядро softmax.
- Matmul (запрос и ключ).
- Mul (масштабирование).
- Копирование памяти (Memcpy) 🤔.
- Причинная маскировка.
- Softmax (генерация весов внимания).
- Matmul (веса внимания и значения).
masked_fill, который создает новый тензор.
masked_fill_). Это стандартная конвенция для операций, изменяющих тензор на месте, без создания новой копии в памяти.02Наивное внимание с операциями «на месте» (In-place)
Все, что мы меняем, — это заменяемmasked_fill на masked_fill_ (обратите внимание на завершающее подчеркивание, конвенцию PyTorch для операций «на месте»), и запускаем тот же скрипт.
def forward(self, q, k, v, mask):
# q, k, v: [batch, heads, seq, head_dim]
scores = torch.matmul(q, k.transpose(-2, -1)) # [batch, heads, seq, seq]
scores = torch.mul(scores, self.scale)
scores.masked_fill_(mask, float("-inf")) # In-place!
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, v) # [batch, heads, seq, head_dim]
return outИсточник: Hugging Face ↗