В команде AI Infrastructure в Ai2 мы занимаемся предоставлением вычислительных мощностей на базе GPU, ориентируясь на крупные распределенные задачи обучения. Мы рассматриваем эту задачу как пирамиду из четырех взаимосвязанных метрик. Основой является доступность: как часто оборудование готово к работе. Следующий уровень — это занятость: доля доступного времени, отведенная под конкретную задачу. Далее идет влияние: как часто наиболее ценные задачи получают ресурсы. Вершина пирамиды — это использование: доля мощности GPU, задействованная на протяжении всего времени выполнения задачи. Эта статья посвящена улучшению влияния наших решений по планированию. Недавно мы заменили планировщик на основе приоритетов на систему, включающую бюджеты времени GPU, иерархическое распределение ресурсов и контракт на временное использование. В результате мы перенесли обсуждение того, сколько времени GPU должно быть выделено каждому исследовательскому проекту, с операционной задачи на прозрачный административный процесс бюджетирования.

Проблема переполнения

В Ai2 мы управляем тысячами GPU NVIDIA H100, B200 и B300, сгруппированными в кластеры размером от 88 до 1024 GPU. Эти кластеры предназначены для крупномасштабного распределенного обучения AI-моделей и обслуживают около 150 внутренних исследователей, работающих в различных областях AI, включая полный цикл обучения LLM и VLM, симуляцию обучения с подкреплением в робототехнике и постобучение для научных задач. Как и во многих лабораториях, спрос на время GPU значительно превышает предложение. На основании поданных рабочих нагрузок в любой момент времени у нас есть запросы на 2-3 раза больше GPU, чем доступно. Исторически мы использовали планировщик на основе приоритетов, позволяя рабочим нагрузкам отказываться от возможности прерывания. Каждой команде был установлен лимит на количество одновременно используемых GPU для задач, защищенных от прерывания. Прерываемые рабочие нагрузки могли превышать этот лимит на неиспользуемых GPU. Эта стратегия привела к предсказуемым проблемам. Например, мы наблюдали случаи "заселения" GPU, когда пользователи оставляли неактивные нагрузки, к которым могли подключиться, когда возникала необходимость. Это происходило, потому что исследователи не могли запустить отладочные нагрузки с достаточной низкой задержкой, чтобы решить проблемы в реальном времени. Мы также наблюдали инфляцию приоритетов, когда в конечном итоге 100% запланированных рабочих нагрузок использовали высокий приоритет. Это означало, что нагрузки с более низким приоритетом вообще не получали времени GPU.

Трагедия общин

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

Бюджеты вместо расписаний

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

Иерархический справедливый планировщик

В дополнение к инструменту бюджетирования времени GPU мы построили иерархический справедливый планировщик для управления фактической занятостью распределений в рамках программы. Алгоритм здесь не новый — иерархическая справедливость на временном интервале является частью наследия, которое восходит к Hadoop Fair Scheduler 2009 года. Однако для нас новыми являются входные данные: дерево отражает структуру исследовательской программы, а веса — это бюджеты, установленные менеджерами, а не статические квоты. Планировщик отслеживает занятость за скользящую временную рамку (по умолчанию 7 дней) и сортирует рабочие нагрузки от недоиспользуемых распределений к перегруженным.

Контракт на планирование

Дополнительной особенностью распределенного обучения, которая усложняет справедливое распределение ресурсов, является то, что рабочие нагрузки могут выполняться очень долго. Обучающие задания часто работают часами, днями и иногда неделями. После назначения рабочая нагрузка может оставаться на своих GPU на неделю или более, не предоставляя возможности другим получить свое бюджетное время. Чтобы решить эти проблемы, мы ввели "контракт на планирование". В обмен на доступ к кластеру рабочая нагрузка должна объявить свое минимальное время работы, или самое короткое время, необходимое для достижения значимого прогресса. В течение этого времени рабочая нагрузка защищена от прерывания.

Результаты

С начала развертывания мы наблюдаем, что пользователи и команды последовательно получают свое выделенное время GPU. Мы подсчитываем время, причитающееся команде, как ее распределение, ограниченное часами в зависимости от фактического спроса. В течение 30-дневного тестового периода команды получили 98% времени GPU, которое им причиталось, и 13 из 15 команд получили 95% или более, причем худший случай составил 90%. Занятость кластера оставалась стабильной на уровне 98% до и после изменений, с превышением спроса над предложением на 2-3 раза в обоих периодах. 18% доставленного времени GPU было незанято, что позволило поддерживать высокую занятость в периоды, когда профинансированные случаи использования не были готовы к запуску.

Проблемы

Кривая обучения была более крутой, чем мы предполагали. Мы развернули изменения поэтапно, поэтому в первые дни исследователи сталкивались с различным поведением в зависимости от того, какой кластер они использовали. Кроме того, наши интерфейсы сохранили некоторые старые термины, значение которых изменилось. Документация не решила путаницы. Что сработало, так это проведение живых объяснительных сессий, предоставляющих форум для исследователей задать вопросы и для инженерной команды предоставить более глубокие описания того, как и почему планировщик принимал свои решения о приоритете.

Будущее

Мы стремимся к вершине пирамиды: использованию. Нам нужно убедиться, что начальная загрузка, контрольные точки и сами приложения для обучения выполняются как можно более эффективно, максимизируя ценность запланированного времени, которое получает каждая рабочая нагрузка. Если вы хотите решать подобные задачи в тесном сотрудничестве с исследователями, мы приглашаем вас рассмотреть открытые инженерные вакансии в Ai2.