Масштабируемое обучение: параллелизм и оптимизация данных в CS336

В восьмой лекции курса Stanford CS336 (весна 2025) преподаватель подробно рассматривает архитектурные подходы к распределенному обучению языковых моделей. В центре внимания — минимизация узких мест при передаче данных между GPU и узлами, а также реализация различных стратегий параллелизма с использованием библиотек PyTorch и NCCL.

⚙️ Уровни иерархии памяти и коммуникации

Эффективность обучения напрямую зависит от того, насколько грамотно организовано перемещение данных внутри и между устройствами. В NVIDIA-экосистеме существует четкая иерархия, которую необходимо учитывать при проектировании вычислений:

• L1-кэш: Самый быстрый, но крайне ограниченный объем памяти внутри потокового мультипроцессора (SM).
• HBM (High-Bandwidth Memory): Основная память на GPU, имеющая гораздо больший объем, но работающая медленнее, чем L1.
• NVLink: Прямое соединение между GPU внутри одного узла, позволяющее передавать данные, минуя CPU.
• NVSwitch: Технология для соединения GPU между разными узлами, исключающая использование медленной сети Ethernet.

🔄 Коллективные операции: примитивы распределенного обучения

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

• Broadcast: Передача данных с одного ранга (устройства) на все остальные.
• Scatter: Разделение данных и распределение разных частей по разным устройствам.
• Gather: Сбор данных с разных устройств на одно конкретное.
• Reduce: Выполнение ассоциативной операции (сумма, минимум, максимум) над данными со всех устройств с последующим сохранением результата на одном ранге.
• All-Gather: Сбор всех данных со всех устройств на все устройства.
• Reduce-Scatter: Комбинация редукции и распределения, где результат операции «редукции» сохраняется в распределенном виде.
• All-Reduce: Операция, эквивалентная reduce + all-gather.

Все эти операции в PyTorch эффективно реализованы через библиотеку NCCL (NVIDIA Collective Communication Library), которая оптимизирует пути передачи данных, исходя из топологии аппаратного обеспечения.

📊 Стратегии распределенного обучения

Для масштабирования обучения моделей на множество GPU применяются три основных метода параллелизма:

1. Параллелизм данных (Data Parallelism)

Модель копируется на все GPU, но каждый из них получает лишь часть общего пакета данных (batch).

• На каждом шаге происходит синхронизация градиентов с помощью операции all-reduce.
• В результате все GPU обновляют веса одинаково, сохраняя идентичность параметров модели.

2. Тензорный параллелизм (Tensor Parallelism)

Модель разбивается по скрытым размерностям матриц.

• Каждый GPU владеет лишь частью весовых матриц каждого слоя.
• Во время прямого прохода происходит активный обмен активациями между GPU через all-gather, что требует высокой пропускной способности сети.

3. Конвейерный параллелизм (Pipeline Parallelism)

Модель разрезается на части по глубине (слоям).

• Для снижения простоя («пузырей» в конвейере) данные разбиваются на микро-батчи.
• Используются коммуникации «точка-точка» (send/receive) для передачи результатов работы одного набора слоев на следующий.

🔮 Взгляд в будущее и альтернативы

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