В рамках курса Stanford CS231n прошла лекция, посвященная одной из самых актуальных тем в современном машинном обучении — крупномасштабному распределенному обучению. Эксперты обсуждают, как за последние 10 лет индустрия перешла от обучения моделей на одном графическом процессоре (GPU) к кластерам, объединяющим десятки тысяч устройств. В качестве сквозного примера используется архитектура Llama 3 405B от Meta, на примере которой разбираются алгоритмы, позволяющие эффективно использовать вычислительные мощности уровня экзофлопс.
🖥️ Анатомия GPU: От графики к тензорным ядрам 3:06
Современные нейросети обучаются на специализированных сопроцессорах — GPU (Graphics Processing Units). Изначально созданные для рендеринга пикселей, они превратились в универсальные параллельные процессоры. Особую роль в этом сыграла компания NVIDIA и её основатель Дженсен Хуанг (Jensen Huang), чьё имя носит аудитория в Стэнфорде, где проходит лекция .
Архитектура современного флагмана NVIDIA H100 демонстрирует сложную иерархию памяти и вычислительных блоков:
- HBM (High Bandwidth Memory): 80 ГБ памяти с пропускной способностью около 3 ТБ/с .
- L2 Cache: 50 МБ памяти, расположенной ближе к ядрам для быстрого доступа .
- Streaming Multiprocessors (SM): 132 активных мультипроцессора (физически их 144, но из-за дефектов производства — «биннинга» — гарантируется работа 132) .
- Tensor Cores: Специализированные схемы для матричного умножения. Одно тензорное ядро H100 способно выполнять умножение матриц (размера 16x4 на 4x8) за один такт .
Лектор подчеркивает, что именно тензорные ядра являются «магическим ингредиентом» производительности. Они работают в смешанной точности (mixed precision): умножение происходит в 16-битных форматах, а сложение — в 32-битных . Ошибка в коде PyTorch, связанная с приведением типов, может привести к тому, что вычисления уйдут на обычные FP32-ядра, что замедлит процесс в 20 раз .
🏗️ От одного чипа к экзомасштабным кластерам 15:20
За последние 12 лет пиковая производительность одного устройства выросла примерно в 1000 раз: от 5 терафлопс у NVIDIA K40 (2013 год) до почти 5000 терафлопс у новой архитектуры Blackwell (B200) . Однако настоящий масштаб достигается за счет объединения устройств в огромные дата-центры, которые следует воспринимать как один гигантский компьютер .
Структура кластера Llama 3 (Meta) служит эталоном современной инфраструктуры:
- Server (Узел): 8 GPU в одном корпусе. Скорость общения между ними — 900 ГБ/с .
- Rack (Стойка): 2 сервера (16 GPU), высота около 2 метров .
- Pod (Под): 192 стойки (3072 GPU) с высокоскоростными интерконнектами .
- Cluster (Кластер): 8 подов, в сумме 24 576 GPU .
Общая мощность такого кластера составляет 24 экзофлопса (10 в 18-й степени операций в секунду) . Помимо NVIDIA, на рынке присутствуют Google с их собственными чипами TPU (Tensor Processing Units), AMD с ускорителем MI325X и AWS с чипом Trainium . Однако, по мнению спикера, NVIDIA и Google остаются безусловными лидерами отрасли.
🔄 Пять степеней параллелизма 27:42
Обучение гигантских моделей требует разделения вычислений по четырем осям тензора (Batch, Sequence, Dimension) и оси слоев (Layer) . Это порождает пять основных стратегий параллелизма:
1. Data Parallelism (DP)
Самый простой метод: каждая GPU получает свою часть мини-батча, но хранит полную копию весов модели . После вычисления градиентов устройства выполняют операцию all-reduce, чтобы усреднить результаты и синхронно обновить веса. Проблема DP — ограниченность памятью одной GPU. 1 миллиард параметров требует около 8 ГБ памяти (с учетом градиентов и состояния оптимизатора), поэтому модель на 10-15 миллиардов параметров просто не поместится в одну H100 .
2. Fully Sharded Data Parallelism (FSDP)
Решение проблемы памяти: веса модели, градиенты и состояния оптимизатора распределяются («шардируются») между всеми GPU . В процессе прямого прохода слой за слоем нужные веса временно копируются на все устройства, используются и сразу удаляются. Это позволяет тренировать модели любого размера, ограничиваясь лишь суммарной памятью всего кластера .
3. Pipeline Parallelism (PP)
Слои нейросети делятся между группами GPU. Например, первые 10 слоев на первой видеокарте, следующие 10 — на второй. Чтобы GPU не простаивали в ожидании данных от предыдущих слоев (проблема «пузыря»), батч данных разбивается на «микро-батчи», которые запускаются каскадом .
4. Tensor Parallelism (TP)
Разделение конкретных матричных операций внутри одного слоя между несколькими GPU . Каждое устройство вычисляет лишь часть результата одного умножения матриц. Это требует очень быстрого соединения между чипами, поэтому обычно ограничено пределами одного сервера (8 GPU).
5. Context Parallelism (CP)
Относительно новый метод, необходимый для работы с длинными текстами (100k+ токенов). Последовательность данных разбивается на части, и разные GPU отвечают за разные фрагменты текста в одной и той же модели .
📈 MFU: Главная метрика эффективности 58:17
При настройке десятков параметров (размер батча, степень шардирования, количество микро-батчей) исследователи ориентируются на MFU (Model FLOPs Utilization) .
- HFU (Hardware FLOPs Utilization): Показывает, какой процент теоретической мощности железа вы используете в чистых вычислениях.
- MFU: Показывает, какой процент мощности идет на полезную работу по продвижению обучения модели (без учета накладных расходов на коммуникацию и повторные вычисления).
По словам лектора, MFU выше 30% считается хорошим результатом, а выше 40% — выдающимся . В отчете Llama 3 указано, что Meta удалось достичь 38-43% MFU на кластере из 16 000 GPU . Парадоксально, но на новых поколениях железа MFU часто падает: вычислительная мощность чипов растет быстрее (в 3 раза), чем пропускная способность памяти и интерконнекта (в 2 раза), что создает «узкое горлышко» в коммуникациях .
🛠️ Практические советы и компенсация памяти 53:36
Для тех, кто планирует обучение, лектор предлагает следующий «рецепт масштабирования»:
- До 1 млрд параметров: Используйте обычный Data Parallelism .
- Свыше 1 млрд параметров: Переходите на FSDP .
- Проблемы с памятью активаций: Включите Activation Checkpointing. Этот метод удаляет промежуточные результаты (активации) после прямого прохода и заново вычисляет их во время обратного . Это экономит память в обмен на 33% дополнительного времени вычислений.
- Масштаб уровня Meta: Используйте 4D-параллелизм. Llama 3 405B обучалась с использованием 8-way TP, 16-way CP, 16-way PP и 8-way DP одновременно .
В завершение подчеркивается: современное глубокое обучение — это не только математика нейросетей, но и сложнейшая дисциплина на стыке компьютерной архитектуры и системного программирования.