Секреты эффективных LLM: учет ресурсов и оптимизация в PyTorch

Stanford Online 140 тыс. 1 ч 19 мин 10 мин 10.04.2025
Главное

В Стэнфордском университете в рамках продвикутого курса CS336 «Моделирование языков с нуля» состоялась лекция, посвященная практическим основам разработки больших языковых моделей и точному учету вычислительных ресурсов. Преподаватель курса подробно разобрал низкоуровневую механику популярного фреймворка PyTorch, принципы работы с тензорами и методику быстрых оценок, известную как «расчеты на салфетке». Ключевой идеей материала стало понимание того, как неэффективное использование памяти и вычислительных мощностей напрямую конвертируется в финансовые потери при масштабировании нейросетей.

🧮 Экономика ИИ и «расчеты на салфетке» 0:43

Когда речь заходит об обучении крупных языковых моделей, инженеры часто сталкиваются с необходимостью быстрой оценки требуемых мощностей. Для этого лектор предлагает использовать так называемые «расчеты на салфетке» (napkin math). В качестве примера приводится задача: сколько времени займет обучение плотной модели-трансформера на 70 миллиардов параметров с использованием 15 триллионов токенов на кластере из 1024 ускорителей NVIDIA H100?

Базовая формула для оценки общего объема необходимых вычислений (в операциях с плавающей запятой, FLOP) выглядит следующим образом: количество операций равно шести произведениям числа параметров на число токенов. Зная паспортную производительность процессора H100 и заложив реалистичную эффективность использования оборудования (Model FLOPs Utilization, MFU) на уровне 0,5, можно рассчитать доступный пул операций в день для всего кластера. Поделив общий объем необходимых FLOP на суточную производительность 1024 карт, мы получаем результат — примерно 144 дня.

Другой важный практический вопрос: модель какого максимального размера можно обучить на стандартном узле из 8 ускорителей H100 с использованием популярного оптимизатора AdamW, если действовать «в лоб»? Каждая карта H100 обладает 80 гигабайтами быстрой памяти HBM. Для хранения параметров, градиентов и состояний оптимизатора в стандартном режиме требуется около 16 байт на один параметр. Простой расчет показывает предел — около 40 миллиардов параметров. Преподаватель подчеркивает, что эта оценка является грубой, поскольку она не учитывает память под активации, размер батча и длину контекста, которые критически важны для выполнения практических заданий курса.

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

💾 Анатомия тензоров: биты, типы данных и скрытые ловушки памяти 5:14

Фундаментальным строительнымлом для хранения любых данных в глубоком обучении — будь то параметры, градиенты, состояния оптимизатора или активации — являются тензоры. С точки зрения потребления памяти размер тензора определяется исключительно количеством содержащихся в нем элементов и выбранным типом данных.

📊 Эволюция точности: от золотого стандарта до трюков NVIDIA

Историческим «золотым стандартом» вычислений является тип float32 (FP32), или одинарная точность. На каждый элемент выделяется 32 бита: 1 бит на знак, 8 бит на экспоненту (определяет динамический диапазон) и 23 бита на фракцию (мантиссу, отвечающую за точность значений). Лектор с иронией отмечает путаницу в терминологии: специалисты по машинному обучению называют float32 «полной точностью», в то время как ученые из сферы суперкомпьютерных вычислений лишь посмеются над этим, поскольку сами используют как минимум float64. Глубокое обучение, по словам преподавателя, «весьма небрежно» к точности, поэтому FP32 — это максимум, который может потребоваться в ML-задачах. Чтобы проиллюстрировать масштабы, лектор приводит факт: всего одна матрица в полносвязном слое (FFN) нейросети GPT-3 занимает порядка 2,3 гигабайта памяти.

Для экономии памяти и ускорения вычислений были разработаны альтернативные форматы:

🧠 Смешанная точность (Mixed Precision) как компромисс

Полноценное обучение исключительно в низком разрешении чревато сбоями. Как объясняет лектор, для хранения весов и состояний оптимизатора по-прежнему необходим формат float32, иначе процесс обучения «пойдет наперекосяк». Решением становится смешанная точность (mixed precision).

В рамках этого подхода параметры временно переводятся в bfloat16 для быстрого выполнения проходов matmul (матричного умножения) в составе прямого шага. При этом критически важные компоненты, такие как механизм внимания (attention), а также накопление градиентов со временем, требуют строгого сохранения во float32.

⚙️ Вычислительная архитектура PyTorch: CPU против GPU 14:20

По умолчанию при создании тензоров в PyTorch они размещаются в оперативной памяти центрального процессора (CPU RAM). Лектор предупреждает, что работа исключительно на CPU делает процесс обучения на порядки медленнее, лишая инженера преимуществ графических ускорителей. Перенос данных на графический чип (GPU) должен осуществляться явно с помощью метода .to(). Пересылка данных между RAM и памятью GPU требует осязаемого времени, что необходимо учитывать при проектировании систем. При тестировании на кластере H100 с 80 ГБ памяти разница в выделенном объеме до и после явного создания тензора на GPU должна в точности соответствовать размеру матрицы.

🧭 Устройство тензора «под капотом» и магия представлений (Views)

Внутри PyTorch тензор представляет собой математический объект, который на системном уровне является лишь указателем на выделенный непрерывный блок памяти. То, как именно этот одномерный массив интерпретируется как многомерная матрица, определяется метаданными — шагами по осям (strides). Страйд указывает, сколько элементов в физической памяти нужно пропустить, чтобы перейти к следующему элементу в данном измерении. Для двумерной матрицы $4 \times 4$ страйд по строкам (stride 0) равен 4, а по столбцам (stride 1) равен 1.

Такая архитектура позволяет множеству тензоров ссылаться на один и тот же физический участок памяти. Операции срезов, изменения формы (.view()) или транспонирования не копируют данные, а лишь создают новое «представление» (view) с измененными страйдами. Однако лектор призывает к осторожности:

📈 Матричное умножение: истинная цена вычислений 23:53

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

✍️ Инструменты Einops и einsum: спасение от путаницы в индексах

Классический код на PyTorch часто изобилует конструкциями вроде x.transpose(-2, -1), где программист оперирует отрицательными индексами осей. По мнению преподавателя, это частый источник трудноловимых багов, поскольку комментарии к коду быстро устаревают. Элегантным решением является использование библиотек Einops и Jax typing. Они позволяют разработчику декларативно описывать имена измерений в виде понятных строк (например, batch, seq, hidden).

Функция einsum (нотация Эйнштейна) превращает матричное умножение в гибко контролируемый процесс с автоматическим суммированием по неназванным осям. Лектор подтверждает, что встроенный компилятор torch.compile успешно оптимизирует такие выражения, определяя наилучший порядок сокращения размерностей, что зачастую превосходит по эффективности код, написанный вручную.

🏎️ Модель эффективности: метрика MFU и маркетинг чипмейкеров

Для оценки стоимости вычислений лектор приводит базовое правило: количество FLOP при перемножении матриц всегда равно двум произведениям трех участвующих размерностей. Все остальные операции в глубоком обучении (сложения, активации) растут линейно и меркнут на фоне matmuls. В реальном мире производительность оценивается через коэффициент использования аппаратных FLOP (Model FLOPs Utilization, MFU). Он рассчитывается как отношение теоретически полезных операций модели к паспортной пиковой скорости чипа.

Показатель MFU выше 0,5 считается в индустрии хорошим результатом, тогда как значение в 5% говорит о крайне неэффективном коде. Добиться 100% утилизации невозможно из-за накладных расходов на передачу данных. Комментируя заявления производителей оборудования, лектор раскрывает маркетинговый трюк NVIDIA: заявляемые пиковые 1979 терафлоп в секунду для H100 предполагают использование так называемой «структурированной разреженности» (sparsity) по схеме «2 из 4 элементов равны нулю». В реальной практике плотных языковых моделей эта технология почти не применяется, поэтому реальную паспортную цифру нужно сразу делить ровно надвое.

📐 Математика обратного прохода: откуда берется «множитель 6» 49:43

Чтобы понять, почему на этапе napkin math затраты на обучение оценивались через коэффициент 6, необходимо детально разобрать фазу вычисления градиентов. На примере простой двухслойной линейной сети лектор демонстрирует распределение вычислительной нагрузки.

Для прямого прохода (forward pass) через слои с весами требуется выполнить объем вычислений, эквивалентный удвоенному произведению числа входных данных на общее количество параметров модели. Однако на этапе обратного распространения ошибки (backward pass) ситуация усложняется из-за применения правила дифференцирования сложной функции (chain rule). Чтобы обновить веса, системе необходимо вычислить два типа градиентов:

  1. Градиент функции потерь относительно самих весовых коэффициентов (для выполнения шага оптимизатора).
  2. Градиент относительно промежуточных активаций (чтобы передать ошибку на предыдущие слои сети).

Каждая из этих операций по структуре является полноценным матричным умножением и требует точно такого же объема FLOP, как и прямой шаг. В итоге обратный проход обходится ровно в четыре произведения объема данных на число параметров. Суммируя прямой шаг (2x) и обратный шаг (4x), мы получаем фундаментальную константу — $6 \times \text{данные} \times \text{параметры}$.

🏗️ Сборка модели: инициализация, оптимизаторы и управление памятью 59:21

Переходя к сборке полноценного тренировочного цикла, лектор подчеркивает важность правильной инициализации весов (nn.Parameter). Попытка инициализировать параметры стандартным распределением Гаусса (Unit Gaussian) «в лоб» приводит к катастрофическому разгону дисперсии значений на выходе, пропорционально квадратному корню из скрытой размерности слоев. Для стабилизации обучения применяется инициализация Ксавье (Xavier initialization), масштабирующая веса обратно пропорционально квадратному корню из числа входов, а также усечение хвостов распределения в жестких границах от -3 до 3.

В качестве важной инженерной практики лектор рекомендует всегда жестко фиксировать случайные зерна (random seeds) для каждого источника стохастичности (инициализация, дропаут, порядок данных). По его мнению, детерминизм — лучший друг инженера при отладке распределенных систем. Кроме того, при работе с огромными датасетами (например, тексты LLaMA объемом 2,8 терабайта) запрещено загружать файлы в память целиком; вместо этого используется механизм отображения файлов в память (mmap), считывающий данные с диска строго по требованию.

⚡ Оптимизаторы и их аппетит к памяти

Выбор алгоритма оптимизации напрямую влияет на доступный бюджет памяти устройства. Обычный градиентный спуск (SGD) не требует хранения истории, но современные LLM обучаются на сложных адаптивных алгоритмах вроде Adam, который сохраняет как скользящее среднее градиентов (момент), так и среднее квадратов градиентов.

На примере более простого алгоритма AdaGrad лектор показывает, что для его работы в памяти необходимо постоянно хранить накопленную сумму квадратов градиентов, что фактически удваивает накладные расходы на хранение параметров. Итоговая формула физической памяти складывается из четырех компонентов, умноженных на размерность типа данных (4 байта для FP32):

$$\text{Память} = 4 \times (\text{Параметры} + \text{Активации} + \text{Градиенты} + \text{Состояния оптимизатора})$$

Чтобы спастись от нехватки памяти под активации, инженеры вынуждены применять технику контрольных точек (activation checkpointing), при которой активации промежуточных слоев удаляются из памяти сразу после прямого шага и заново вычисляются во время обратного прохода.

🔮 Будущее вычислений: квантование и синергия железа с архитектурой

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

Лектор отмечает важную закономерность: современный дизайн нейросетевых архитектур полностью завязан на специфику железа NVIDIA. Если инженерам удастся заставить обучение стабильно работать в формате int4 (целые числа), это принесет колоссальный прирост скорости. При этом лектор указывает на фундаментальную асимметрию: обучать модели в сверхнизкой точности невероятно тяжело, однако уже готовую, обученную модель можно подвергнуть агрессивному квантованию для инференса без существенной потери качества, получив кратный выигрыш в эффективности.

💬 Цитаты

«По мнению лектора, в академической среде студенты часто привыкают просто запускать код и смотреть, что произойдет, не задумываясь об эффективности.»

Лектор курса CS336 03:25

«Обучать модели в сверхнизкой точности невероятно тяжело, однако уже готовую, обученную модель можно подвергнуть агрессивному квантованию для инференса.»

Лектор курса CS336 01:18:38
👥 Спикер
🔗 Упомянутые сайты и проекты
📖 Термины
FLOP
Количество элементарных операций с плавающей запятой, мера объема вычислений.
MFU
Коэффициент полезного использования теоретической вычислительной мощности графического ускорителя моделью.
Страйд (Stride)
Шаг по физической памяти, определяющий расстояние между соседними элементами тензора в конкретном измерении.
Квантование
Перевод весов нейросети в форматы с более низкой точностью (например, из float32 в int4) для ускорения работы.
Исчезновение порядка (Underflow)
Ситуация, когда число слишком мало для представления в данном формате точности и округляется до нуля.
📊 Цифры
⚖️ Другая сторона
Искусственный интеллект PyTorch Стэнфордский университет учет ресурсов тензоры квантование