УСКОРЕНИЕ ДООБУЧЕНИЕ БОЛЬШИХ ЯЗЫКОВЫХ МОДЕЛЕЙ (LLM) ПРИ ПОМОЩИ ФРЕЙМВОРКА AXOLOTL

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

Ключевые слова: большие языковые модели, дообучение, Axolotl, оптимизация, параллелизм последовательностей, LoRA.

Дообучение больших языковых моделей (LLM) под конкретные задачи стало неотъемлемой частью современного рабочего процесса в области искусственного интеллекта. Однако этот процесс остаётся одним из самых ресурсоёмких этапов разработки: даже использование параметрически эффективных методов вроде LoRA требует значительных вычислительных мощностей, тщательной настройки гиперпараметров и глубокого понимания внутреннего устройства трансформеров. Сложность многократно возрастает, когда речь заходит о работе с моделями, содержащими десятки или сотни миллиардов параметров, либо о необходимости обрабатывать сверхдлинные контексты – например, при анализе многотомных документов или длительных диалогов. Именно здесь на помощь приходит фреймворк Axolotl, который был специально разработан для того, чтобы демократизировать доступ к эффективному дообучению и ускорить его в разы, предоставляя при этом удобный интерфейс и богатый набор оптимизаций как для обучения, так и для последующего развёртывания моделей.

Axolotl представляет собой открытый инструмент, построенный вокруг идеи унификации всех этапов пост-тренировочного пайплайна. Под пост-тренировкой здесь понимается широкий спектр техник: от классического полного fine-tuning и адаптации с помощью LoRA/QLoRA до более сложных методов обучения с предпочтениями, таких как DPO, IPO или KTO, а также дообучения с подкреплением. Разработчики Axolotl поставили перед собой цель абстрагировать пользователя от необходимости писать сотни строк кода на PyTorch и вручную комбинировать библиотеки вроде Transformers, PEFT, DeepSpeed и Flash Attention. Вместо этого всё управление осуществляется через единый конфигурационный файл в формате YAML, где можно описать модель, набор данных, параметры оптимизатора, техники ускорения и даже логирование экспериментов. Такой подход не только снижает порог входа для новичков, но и кардинально сокращает время, затрачиваемое на настройку окружения и отладку, позволяя исследователям и инженерам сосредоточиться на сути задачи.

Одной из главных «киллер-фич» Axolotl, напрямую влияющих на скорость дообучения, является реализация кастомных оптимизированных ядер для LoRA. Стандартная реализация LoRA в библиотеке PEFT, хоть и экономит память, часто не оптимизирована по скорости выполнения из-за того, что операции с низкоранговыми адаптерами и замороженными весами выполняются последовательно. Axolotl, вдохновляясь проектом Unsloth, внедрил собственные autograd-функции, которые сливают вычисления для линейных слоёв и механизма внимания. Например, при прямом проходе через слой с адаптером вычисления объединяются таким образом, чтобы минимизировать количество обращений к памяти и эффективно использовать тензорные ядра современных GPU. Для обратного прохода также используются специально написанные ядра на языке Triton, которые оптимизируют вычисление градиентов для функций активации вроде SwiGLU, широко применяемых в архитектурах LLaMA и её производных. Включение этих оптимизаций в конфигурационном файле достигается добавлением нескольких флагов: lora_mlp_kernel, lora_qkv_kernel и lora_o_kernel. При их активации можно наблюдать прирост скорости обучения до 30–50% без какого-либо ущерба для финального качества модели. Более того, эти улучшения работают не только на одиночных GPU, но и в распределённых средах, включая Data Parallelism с DeepSpeed и FSDP, что делает их применимыми для широкого круга задач.

Чтобы лучше понять, как выглядит типичная конфигурация Axolotl, рассмотрим пример настройки для дообучения небольшой инструктивной модели, например, LLaMA-3.2-1B-Instruct, с использованием QLoRA и всех ускоряющих техник. В файле конфигурации указывается base_model, путь к датасету в формате sharegpt, параметры загрузки в 4-битной точности для экономии памяти, настройки LoRA с рангом 32 и включёнными оптимизированными ядрами, размер микро-батча 2, накопление градиентов для имитации большего батча, планировщик скорости обучения косинус, включённый gradient checkpointing для экономии памяти, флаг flash_attention и параметры логирования. Всё это занимает не более 50–60 строк человекочитаемого текста. При этом в фоновом режиме Axolotl автоматически загрузит токенизатор и модель, применит все необходимые патчи для поддержки длинных контекстов, настроит оптимизатор с 8-битным Adam, и можно запускать обучение одной командой accelerate launch -m axolotl.cli.train config.yml. Такой уровень автоматизации и скорости настройки недостижим при использовании сырых библиотек.

В заключение можно сказать, что Axolotl представляет собой один из наиболее продуманных и эффективных инструментов для пост-тренировки больших языковых моделей на сегодняшний день. Его подход к ускорению дообучения базируется на трёх китах: во-первых, это внедрение низкоуровневых оптимизаций (кастомные ядра LoRA, Flash Attention, Liger Kernel), которые максимально эффективно используют аппаратное обеспечение; во-вторых, это реализация передовых методов распределённых вычислений, таких как Sequence Parallelism, позволяющих работать с контекстами невообразимой ранее длины; в-третьих, это тесная интеграция с инструментами компрессии моделей (LLM Compressor, квантизация), что даёт возможность получать не только быстро обученные, но и быстрые в инференсе модели.