Meta представляет LlamaRL: масштабируемый фреймворк для обучения больших языковых моделей с помощью RL на базе PyTorch

Роль reinforcement learning в тонкой настройке больших языковых моделей (LLM)

Reinforcement learning (RL) стал мощным инструментом для тонкой настройки больших языковых моделей (LLM), чтобы сделать их поведение более интеллектуальным. Эти модели уже способны выполнять широкий спектр задач: от обобщения текста до генерации кода. RL помогает адаптировать их выходные данные на основе структурированной обратной связи.

Поскольку растёт спрос на модели, которые должны быть не только точными, но и соответствовать сложным предпочтениям или правилам, RL предоставляет важнейший механизм для повышения их производительности. Следовательно, RL стал центральным компонентом в процессе пост-тренинга многих продвинутых систем LLM.

Проблемы масштабирования RL для LLM

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

Заметные компоненты включают модели политик, системы оценки вознаграждений и критики. Размеры моделей исчисляются сотнями миллиардов параметров, и такие проблемы, как использование памяти, задержка передачи данных и время простоя GPU, представляют собой сложные инженерные задачи. Без эффективного проектирования эти ограничения препятствуют применению RL к новым, более крупным моделям.

Ограничения предыдущих фреймворков RL для LLM

Предыдущие решения сталкивались с трудностями либо из-за жёсткости, либо из-за неэффективности при масштабировании. Традиционные синхронные фреймворки выполняют генерацию и обучение последовательно, что часто приводит к простою GPU из-за несоответствия продолжительности задач.

Инструменты, такие как DeepSpeed-Chat, используют гибридные стратегии памяти, но требуют, чтобы модели использовали общее пространство памяти, что приводит к узким местам в производительности во время генерации. Некоторые распределённые методы пытаются разделить компоненты, но всё ещё полагаются на тяжёлые инструменты оркестрации, что ограничивает гибкость.

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

LlamaRL от Meta: распределённый асинхронный фреймворк RL на базе PyTorch

Исследователи Meta представили LlamaRL, полностью асинхронный и распределённый фреймворк для обучения с подкреплением. Он предназначен для обучения массивных LLM на кластерах от нескольких до тысяч GPU. Они создали LlamaRL полностью на базе PyTorch и реализовали дизайн с одним контроллером для упрощения координации.

Этот дизайн обеспечивает модульную настройку. Отдельные исполнители управляют каждым компонентом RL — таким как генератор, тренер и модель вознаграждения — и работают параллельно. Такая асинхронная настройка сокращает время ожидания на протяжении всего конвейера RL. Она также позволяет независимо оптимизировать параллелизм моделей и использование памяти.

Ключевые особенности: выгрузка, эффективность использования памяти и асинхронное выполнение

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

Распределённый прямой доступ к памяти (DDMA) поддерживает эту выгрузку. Он использует NVIDIA NVLink для синхронизации весов менее чем за две секунды — даже для моделей с 405 миллиардами параметров.

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

Реальные тесты производительности: ускорение в 10,7 раз для моделей на 405B

LlamaRL демонстрирует значительное улучшение скорости обучения без ущерба для качества. Для модели с 8 миллиардами параметров и 256 GPU время выполнения этапа обучения сокращается с 22,45 секунды до 8,90 секунды. Для модели на 70 миллиардов параметров время сокращается с 82,32 до 20,67 секунд.

Наиболее впечатляющим является то, что для модели с 405 миллиардами параметров на 1024 GPU LlamaRL сокращает время выполнения шага RL с 635,8 до 59,5 секунд и достигает ускорения в 10,7 раз по сравнению с синхронной базовой линией.

Эти достижения являются результатом не только асинхронного выполнения, но и его стратегий по разделению памяти и вычислений. Тестовые оценки на MATH и GSM8K подтверждают, что LlamaRL поддерживает стабильную производительность. Некоторые метрики даже показывают небольшое улучшение.

Заключительные мысли: LlamaRL как масштабируемый путь вперёд в обучении LLM

Это исследование представляет собой практическое и масштабируемое решение одной из наиболее значимых проблем — обучения больших языковых моделей (LLM) с помощью reinforcement learning.

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

Источник

Оставьте комментарий