🖥 Как бы вы реализовали функцию потерь в PyTorch?

В PyTorch функции потерь могут быть реализованы путем создания подкласса класса nn.Module и переопределения метода forward. Метод forward принимает на вход прогнозируемый выход и фактический выход и возвращает значение потерь.

Приведем пример кода:

import torch
import torch.nn as nn

class CustomLoss(nn.Module):
def __init__(self):
super(MyLoss, self).__init__()

def forward(self, output, target):

loss = ... # compute the loss

return loss

Теперь, чтобы использовать функцию потерь, необходимо инициализировать ее и передать в качестве аргумента параметру criterion оптимизатора в цикле обучения.

model = ...
optimizer = ...
criterion = CustomLoss()

# цикл обучения
for epoch in range(num_epochs):

optimizer.zero_grad()
outputs = model(data)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

...

#pytorch #junior

@machinelearning_interview

🌟 MetaShuffling от PyTorch: ускоряем вывод Llama 4 MoE без лишних вычислений и задержек.

PyTorch представил MetaShuffling — решение для ускорения вывода в Llama 4 MoE, которое решает проблемы динамизма и разреженности моделей, связанных с маршрутизацией токенов. Вместо традиционных методов padding или slicing, MetaShuffling использует переупорядочивание токенов по экспертам, избавляясь от ненужных операций и синхронизации между CPU и GPU. Это снижает использование памяти и устраняет задержки, связанные с обработкой «пустышек» или множественными запусками ядер.

В основе решения - идея группировки токенов, назначенных одному эксперту, в непрерывные блоки. Такой подход позволяет использовать dense tensors вместо разреженных структур, сохраняя статичные формы данных.

Благодаря этому MetaShuffling совместим с механизмами графов (CUDAGraph, torch.compile), которые ускоряют исполнение, избегая повторных синхронизаций. Решение особенно эффективно для Llama 4, где каждый MoE-слой активирует лишь часть экспертов, что на практике создает динамические нагрузки.

▶️Ключевыми инновациями стали оптимизация ядер GroupedGEMM и IndexShuffling:

🟢GroupedGEMM, написанный на Triton, обрабатывает несколько матриц в одном вызове, используя статические и динамические разбиения размеров, позволяя пропускать неактивных экспертов и «лишние» токены без дополнительных затрат.

🟢IndexShuffling, в свою очередь, выполняет сортировку токенов и подсчет их количества на каждом эксперте за один проход, что по тестам оказалось в 5–13 раз быстрее, чем стандартные реализации PyTorch.

Результаты экспериментов на H100 80GB выглядят многообещающими.

Prefill Llama 4 Maverick с FP8 GroupedGEMM достигает 1,197 TFlops при 286 мкс, что близко к теоретическому пределу GPU.

В задачах декодирования метрики также демонстрируют высокую эффективность: 44,88 TFlops за 59 мкс. Даже при малом количестве токенов (128) MetaShuffling показывает 80% использования пропускной способности памяти.

Для multi-host сценариев MetaShuffling предлагает гибкость между «динамическими» и «статичными» формами данных. В режиме eager (без графов) используется минимальное заполнение без синхронизации CPU-GPU.

В graph mode — статичные формы с минимальным паддингом, что сокращает сетевой трафик и память. Также реализована дедупликация коммуникаций, которая распределяет нагрузку между узлами, снижая задержки в распределенных вычислениях.

MetaShuffling поддерживает интеграцию с FBGEMM Generative AI Kernel Library, позволяя применять оптимизации для vLLM и SGLang.


📌 Полная статья в блоге Pytorch


@ai_machinelearning_big_data

#AI #ML #MetaShuffling #Pytorch