Allen Institute for AI (AI2) выпустил Olmo-core 3, значительное обновление инфраструктуры для обучения больших языковых моделей (LLM) с архитектурой Mixture-of-Experts (MoE). Главная цель релиза — масштабирование обучения в диапазон триллионов параметров без потери вычислительной эффективности. Это основа для следующего поколения моделей Olmo, которые перейдут на MoE-архитектуру.
Ключевая проблема MoE-моделей при масштабировании — накладные расходы на коммуникацию и маршрутизацию данных. Olmo-core 3 решает её за счёт перехода от FSDP к DDP и внедрения специализированных оптимизаций маршрутизации.
01Архитектурный сдвиг: от FSDP к DDP
Предыдущие версии Olmo-core использовали Fully Sharded Data Parallelism (FSDP), который требовал постоянного сбора и перераспределения весов модели для каждого мини-батча. Olmo-core 3 переключился на Distributed Data Parallelism (DDP). Эксперты остаются резидентными (постоянно загруженными) на GPU, а данные маршрутизируются к ним. Это устраняет дорогостоящие операции сбора весов.
В бенчмарке на 8 GPU NVIDIA B300 новая архитектура показала пропускную способность 52 000 токенов/с на GPU против 19 400 токенов/с в старой реализации. Прирост составил 2.7x при сохранении числа активных параметров на токен (~3.2B).
02Техники масштабирования и оптимизации
Olmo-core 3 комбинирует три основных метода распределения модели и состояния оптимизатора:
- Expert Parallelism: Распределение пула экспертов по GPU. Каждая карта хранит только часть экспертов.
- Pipeline Parallelism: Разделение слоёв модели между группами GPU для снижения требований к памяти.
- Distributed Optimizer: Распределение состояния оптимизатора по узлам кластера вместо дублирования на каждой карте.
Для оптимизации маршрутизации и вычислений внедрены:
- Rowwise expert parallelism: Данные направляются напрямую в буферы ввода экспертов, минимизируя перестановку данных.
- GPU-resident routing: Метаданные маршрутизации хранятся на GPU, позволяя CPU планировать задачи без ожидания копирования данных.
- Grouped GEMM: Объединение множества мелких вычислений экспертов в более крупные операции для повышения эффективности GPU.
03Поддержка MXFP8 и результаты бенчмарков
Фреймворк поддерживает формат MXFP8 (Mixed Precision FP8). Это снижает объём перемещаемых данных и вычислительную нагрузку, если затраты на конвертацию форматов не превышают экономию.
На 4 GPU NVIDIA B300 включение MXFP8 дало следующие результаты по сравнению с BF16:
- Рост пропускной способности обучения на заметную величину.
- Снижение пикового потребления VRAM.
Основной прирост обеспечили операции feed-forward и перемещение данных между экспертами, а не механизм внимания (attention).
04Масштабирование до триллионов параметров
Olmo-core 3 продемонстрировал способность работать с моделями огромного размера. Основные тесты проводились на GPU NVIDIA B300 (архитектура Blackwell).
| Конфигурация | Параметры (всего) | Активные параметры | GPU | Пропускная способность |
|---|---|---|---|---|
| Базовый MoE | 47B | 3.2B | 8x B300 | 52k токенов/с/GPU |
| Масштабный тест | Over 1 Trillion | — | — | — |
| DeepEP v2 (эксперимент) | 2.38 Trillion | — | — | — |
Для конфигурации в 1.2 трлн параметров использовалась случайная маршрутизация для оценки производительности системы. Также проведён эксперимент с альтернативной коммуникационной библиотекой DeepEP v2, позволивший запустить модель на 2.38 трлн параметров (тест ёмкости, а не полноценного обучения).
05Кому подойдёт / что запустится
Olmo-core 3 ориентирован на исследовательские лаборатории и энтузиастов с доступом к высокопроизводительным кластерам NVIDIA (B200/B300/H100/H200). Для локального запуска на потребительских видеокартах (RTX 4090/3090) фреймворк не предназначен из-за требований к межкарточной коммуникации и объёму VRAM.
Код фреймворка открыт и доступен на GitHub. Для успешного внедрения требуется понимание принципов распределённого обучения и настройка параметров маршрутизации под конкретное железо.
Источник: Hugging Face ↗