В области больших языковых моделей развернулась настоящая гонка масштабов, где технологические гиганты соревнуются в количестве параметров своих нейросетей. В своём видео исследователь Янник Килчер (Yannic Kilcher) подробно разбирает инженерную статью компании Google, посвящённую системе GShard, которая позволила обучить гигантскую модель на 600 миллиардов параметров всего за четыре дня. Этот фреймворк кардинально меняет подход к распределению вычислений, предлагая автоматическое разделение нейросети на независимых «экспертов».
🔄 Смена парадигмы: почему Google не пошёл по пути OpenAI 0:00
Компания OpenAI удивила мир своей моделью GPT-3 со 175 миллиардами параметров, но Google значительно превзошёл этот показатель, создав языковую модель на 600 миллиардов параметров. Статья Дмитрия Лепихина и его коллег «GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding» описывает не просто архитектуру, а масштабную инженерную систему. В отличие от OpenAI, которая в GPT-3 увеличивала глубину сети (достигая примерно 100 слоёв) и размер контекста до 2048 токенов, разработчики из Google выбрали принципиально иной путь масштабирования.
Полносвязные глубокие архитектуры сталкиваются с серьёзными ограничениями эффективности вычислений и времени прямого и обратного проходов. Максимальная плотная (dense) модель Google имела всего 2,3 миллиарда параметров, но её обучение с помощью фреймворка GPipe на 2048 платах TPU требовало шести недель. Вместо бесконечного наращивания слоёв инженеры Google в архитектуре GShard ограничили глубину сети всего 36 слоями, а контекстное окно — 1024 токенами, перенеся весь упор на расширение внутренних блоков модели.
🤖 Архитектура смеси экспертов (MoE) и жёсткая маршрутизация 8:30
Стандартный трансформер состоит из блоков внимания и полносвязных слоёв (feedforward), через которые последовательно проходят все токены текста. Инженеры Google внедрили концепцию «смеси экспертов» (Mixture of Experts, MoE), заменив единый полносвязный слой множеством параллельных независимых нейросетевых блоков — экспертов. В максимальной конфигурации GShard задействует до 2048 таких экспертов, работающих параллельно.
Главным нововведением стал механизм маршрутизации, распределяющий токены между экспертами:
- Вместо «мягкого» распределения (soft routing), когда каждый токен обрабатывается каждым экспертом, применяется жёсткая маршрутизация (hard routing).
- Каждый токен отправляется максимум к одному или двум наиболее подходящим экспертам, что позволяет колоссально экономить вычислительные ресурсы.
- В систему намеренно подмешивается случайный шум, благодаря чему с определённой вероятностью токен посылается только одному эксперту. По мнению Янника Килчера, именно этот шум обеспечивает стабильность градиентного спуска при обратном распространении ошибки через дискретные пути.
- Дополнительная функция потерь (auxiliary loss) штрафует систему, если токены распределяются неравномерно, предотвращая перегрузку отдельных популярных экспертов.
🎛️ Под капотом GShard: автоматическое разделение и магия XLA 23:24
Поскольку 2048 экспертов физически невозможно уместить в память одного вычислительного ускорителя, модель приходится разделять (шардировать) между множеством распределённых компьютеров. Традиционный механизм внимания дублируется и синхронизируется на всех машинах с использованием параллелизма по данным, тогда как сами эксперты изолированно живут на разных устройствах. Программирование такого распределения вручную — сложнейшая задача, но фреймворк GShard автоматизирует этот процесс.
Для описания тензорных операций разработчики активно используют нотацию Эйнштейна (einsum). В коде GShard разработчику достаточно использовать две ключевые функции:
replicate: указывает системе, что тензор весов необходимо скопировать и синхронизировать на всех вычислительных узлах.split: разделяет вычисления и данные по определённой оси (например, оси экспертов) между независимыми машинами.
Компилятор XLA (Accelerated Linear Algebra) берёт на себя всю работу по авторазбиению кода. По словам Янника Килчера, фреймворк работает поразительно умно: если матрица не помещается в память, GShard автоматически организует циклы для поблочного вычисления и пересылки слайсов данных. Более того, система самостоятельно реализует алгоритмы «обмена ореолом» (halo exchanges) для распределённых свёрточных операций, избавляя программистов от написания низкоуровневых примитивов передачи данных.
🌍 Мультиязычный перевод и феномен языкового трансфера 47:43
Google протестировал разработанную систему на задаче мультиязычного машинного перевода со 100 языков на английский, используя гигантский внутренний датасет из 25 миллиардов пар документов, из которых для финального обучения отобрали 13 миллиардов примеров. При совместном обучении столь разных языков возникает феномен положительного языкового трансфера (positive language transfer). Например, при обучении схожих языков (таких как немецкий и швейцарский немецкий) модель может переносить общие лингвистические конструкции, что резко повышает качество перевода для редких языков с малым объёмом данных.
Однако у мультиязычных моделей есть и обратная сторона:
- Для популярных языков (с большим объёмом ресурсов) общая вместимость модели становится узким горлышком, из-за чего качество перевода может падать по сравнению с отдельной специализированной сетью.
- Решить эту проблему можно только радикальным увеличением числа параметров.
- Существует опасность избыточной ёмкости: если экспертов становится слишком много, жёсткая маршрутизация начнёт отправлять редкие языки к одним экспертам, а популярные — к другим. В таком случае, как подчёркивает Килчер, положительный трансфер знаний исчезнет. Оптимальный баланс для данной задачи лежит где-то между 500 и 2048 экспертами.
⚡ Фантастическая эффективность и недосягаемый триллион 56:02
Результаты тестов GShard показывают колоссальный отрыв в эффективности обучения. Модель на 600 миллиардов параметров была обучена на 2048 ядрах TPU v3 всего за 4 дня, что потребовало 22 «ядро-года» (TPU core years) вычислений. Для сравнения, плотная модель на 2,3 миллиарда параметров требовала 235 ядро-лет, а обучение 100 отдельных двуязычных базовых моделей заняло бы 29 ядро-лет. С ростом количества слоёв эффективность использования токенов растёт, однако расширение пула экспертов демонстрирует логарифмическую кривую полезности.
Инженеры Google также предприняли попытку создать модель с 1 триллионом параметров, добавив 60 слоёв и сохранив 2000 экспертов. К сожалению, довести её обучение до конца не удалось из-за возникшей численной нестабильности (numerical instability). По мнению Янника Килчера, сообщество трансформеров сейчас находится в ожидании своего «ResNet-момента» — архитектурного прорыва, который позволит стабильно обучать бесконечно глубокие и масштабные сети аналогично тому, как это произошло со свёрточными сетями в компьютерном зрении.
