«Сделать так, чтобы большие языковые модели на суперкомпьютерах обучались быстрее»

Феликс Смирнов проходит обучение в аспирантской школе по техническим наукам НИУ ВШЭ и является лауреатом стипендии Президента Российской Федерации для аспирантов и адъюнктов. В настоящее время он — стажер-исследователь международной лаборатории суперкомпьютерного атомистического моделирования и многомасштабного анализа МИЭМ. Молодой ученый рассказал о своем исследовании.
Мои проекты в лаборатории напрямую связаны с темой диссертационного исследования, которое посвящено эффективности коллективных GPU-коммуникаций при распределенном обучении LLM на суперкомпьютерных системах. Глобально изучением суперкомпьютерных коммуникаций я начал заниматься на первом курсе магистратуры, когда передо мной встал вопрос об их производительности в различных высокоскоростных сетях и прикладных программных пакетах.
Цель работы — повысить скорость распределенного обучения LLM за счет увеличения производительности коллективных GPU-коммуникаций на нескольких его сетевых этапах. Достижение этой цели позволит оптимизировать рабочие нагрузки LLM, что обеспечит экономию временных и вычислительных ресурсов суперкомпьютера. Такая экономия имеет практическую значимость, позволяя в перспективе более эффективно применять ИИ в различных сферах промышленности и бизнеса.
Для достижения цели была поставлена глобальная задача — разработать комплексный метод анализа производительности рабочих нагрузок LLM: от макроскопической перспективы (общее поведение приложения и межузловые обмены данными) до микроскопического анализа (производительность и использование ресурсов на уровне отдельных вычислительных ядер GPU). Его разработка позволит выявить проблемные места на сетевых этапах обмена данными в процессе обучения, устранение которых повысит эффективность коммуникаций. Новизна метода состоит в том, что он будет основан на модели аппаратно-программного коммуникационного стека, где каждый из ее уровней вносит вклад в конечную производительность коллективных операций при обучении LLM.
Применяемые инструменты
Используются инструменты для анализа производительности параллельных кодов, освоенные в ходе выполнения наших лабораторных проектов. Некоторые из них нужны для сбора метрик и данных о работе коммуникаций внутри программы (например, данных о времени их выполнения или размере переданных через них сообщений в байтах). Тут можно выделить HPCToolkit и Score-P, на основе работы которых генерируются файлы профилей и временных диаграмм процесса обучения LLM.
Также используются инструменты для анализа этих профилей и временных диаграмм: HPCViewer, Vampir и Cube. Они упорядочивают и визуализируют собранные данные, что позволяет делать первичные выводы об узких местах обучения. Не стоит забывать и о базовых сетевых тестах, которые могут оценить задержки коммуникаций. Например, это набор микробенчмарков OSU или ReproMPI.
Неотъемлемым инструментом и одновременно объектом исследования выступает сам суперкомпьютер. Он предоставляет необходимую среду из десятков вычислительных узлов, оснащенных множеством GPU. На этой инфраструктуре запускаются тяжелые распределенные вычисления и тестируется эффективность межузлового взаимодействия при передаче больших объемов данных в процессе обучения LLM.
Где уже были представлены результаты исследования
Основные результаты работы были апробированы на нескольких международных и всероссийских профильных суперкомпьютерных конференциях. Среди них: International Conference on Parallel Processing & Applied Mathematics (PPAM), Международная конференция «Суперкомпьютерные дни в России» (RSD) и Всероссийская научная конференция с международным участием «Параллельные вычислительные технологии» (ПаВТ). Результаты также были опубликованы издательством Springer в сборниках трудов конференций и в рецензируемом журнале The International Journal of High Performance Computing Applications (IJHPCA), который входит в первый уровень «Белого списка».