Ai2 заменила приоритетный планировщик GPU на бюджеты времени и fair-share

Команда AI Infrastructure в Ai2 отвечает за вычислительные мощности института на GPU, в первую очередь для крупных распределённых обучений. Она описывает работу как пирамиду из четырёх метрик: доступность (насколько исправно оборудование), занятость (доля доступного времени, отданная конкретной задаче), эффективность выбора, то есть влияние (как часто ресурсы получают самые ценные задачи), и утилизация (доля мощности GPU, реально использованная за время задачи). Пост посвящён именно влиянию.

Масштаб такой: Ai2 управляет тысячами GPU NVIDIA H100, B200 и B300, собранных в кластеры от 88 до 1024 GPU. Ими пользуются около 150 внутренних исследователей: обучение больших языковых и визуально-языковых моделей, симуляции для обучения с подкреплением в робототехнике, дообучение для научных агентных сценариев. Спрос на GPU-часы превышает предложение: по поданным задачам в любой момент заявок на 2-3 раза больше GPU, чем есть в наличии.

Раньше работал приоритетный планировщик: у каждой команды был лимит одновременных GPU для задач, защищённых от вытеснения, а от вытеснения можно было вообще отказаться. Это породило предсказуемые патологии. Пользователи «сидели» на GPU, запуская пустые задачи, к которым можно подключиться при необходимости, потому что отладочные задачи не удавалось запускать с достаточно малой задержкой. Случилась инфляция приоритетов: в итоге 100% задач шли с приоритетом HIGH, и низкие приоритеты вообще не получали GPU-времени. А дежурные инженеры тратили большую часть времени на тикеты, договариваясь об организованной остановке невытесняемых задач на хостах с известными проблемами обслуживания.

Сначала команда пыталась ужесточить контроль над выставлением приоритетов и выдавала важным проектам GPU в монопольное пользование. Позже авторы признали, что получили классическую «трагедию общих ресурсов». Монополии оказались слишком грубым инструментом: из-за сезонности исследований GPU простаивали, пока другая команда ждала мощности. Авторы сравнивают это с ручным решением задачи о рюкзаке.

Новая схема строится на бюджетах времени. Вместо выдачи командам конкретных GPU руководство заранее распределяет доли GPU-времени между исследовательскими усилиями исходя из оценки их вероятного влияния, как инвестор. Бюджеты иерархичны: решения внутри проекта принимает ведущий исследователь, внутри программы, главный исследователь (principal investigator), между программами, ведущий менеджер программ или сам CEO. Каждый запрос на GPU-время должен быть профинансирован из бюджета, иначе он не защищён от вытеснения. Любой трюк с получением GPU теперь списывается из бюджета пользователя: задача-«сидельщик» тратит бюджет команды впустую. Цель, сделать обман планировщика дороже, чем честный спор за больший бюджет. На иллюстративной схеме проект A1 имеет гарантированную долю 35% от общей мощности.

Над бюджетами работает иерархический планировщик fair-share. Сам алгоритм не нов: он восходит к Hadoop Fair Scheduler 2009 года и применяется в SLURM Fair Tree и YARN Fair Scheduler. Новое, входные данные: дерево повторяет структуру исследовательских программ, а веса, это бюджеты, задаваемые менеджерами, а не статические квоты. Планировщик отслеживает занятость в скользящем окне (по умолчанию 7 дней) и ставит задачи из недоиспользованных долей выше задач из переиспользованных. За неделю каждая группа должна получить своё время, если постоянно подаёт задачи с достаточным спросом. Различаются два вида занятости: «выделенная» списывается с бюджета и защищена от вытеснения в течение минимального времени работы, а «невыделенная» не списывается ни с какого бюджета, не защищена с самого начала и может быть вытеснена любым выделенным запросом. Так GPU остаются полностью занятыми, даже если распределение долей не совпало со спросом.

Третий элемент, «контракт планирования». Обучающие задачи идут часами, днями и неделями, и без ограничений получивший GPU мог держать их неделю и дольше, именно это делало возможным «сидение» и мешало обслуживанию. Теперь задача при подаче обязана объявить минимальное время работы, кратчайшее время, нужное для заметного прогресса. В течение этого времени она защищена от вытеснения; затем планировщик вправе перебалансировать нагрузку и автоматически вернуть возобновляемые задачи в очередь. Можно указать нулевое минимальное время: тогда время невыделенное, всегда вытесняемое, но бесплатное для бюджета. Жизненный цикл: подача с минимальным временем и признаком возобновляемости; планирование по fair-share с весом по отношению фактической занятости к выделенному времени в окне; работа в течение минимального времени за счёт бюджета; продолжение, пока доли продолжают отдавать задаче приоритет (тоже за счёт бюджета); возможное вытеснение и возврат в очередь; завершение с освобождением ресурсов.

Побочный эффект авторы назвали более важным, чем ожидали: нездоровые хосты могут разгружаться, когда у задач истекает минимальное время, поэтому ремонт можно полностью автоматизировать. Это сократило число ремонтов, требовавших участия человека, на 74%. Пользователь Chris Clark говорит, что новый планировщик создаёт ощущение дополнительных 30% вычислений: неиспользованную часть лимита можно позже «выбрать» всплеском без вытеснения. Это его субъективная оценка, а не измеренный показатель.

Ключевые факты

• Ai2 управляет тысячами GPU H100, B200 и B300 в кластерах от 88 до 1024 GPU для около 150 внутренних исследователей; заявок на GPU в 2, 3 раза больше, чем есть в наличии.
• Старый приоритетный планировщик привёл к «сидению» на GPU пустыми задачами, инфляции приоритетов (в итоге 100% задач шли с HIGH) и перегрузке дежурных инженеров.
• Новая схема: бюджеты GPU-времени, которые менеджеры распределяют по дереву программ, плюс иерархический fair-share со скользящим окном (по умолчанию 7 дней).
• «Контракт планирования»: каждая задача объявляет минимальное время работы, в течение которого защищена от вытеснения; после этого её можно вытеснить и вернуть в очередь.
• По словам авторов, автоматизация ремонта за счёт разгрузки хостов сократила число ремонтов с участием человека на 74% (исходное число в тексте не указано).

Почему это важно

Нехватка GPU, общая проблема лабораторий, и Ai2 откровенно описывает, как привычная схема приоритетов выродилась: когда приоритет HIGH ничего не стоит, им пользуются все, а низкие приоритеты голодают. Главная идея поста, перенести спор о том, сколько GPU-времени заслуживает проект, из ручных ситуативных решений дежурных в прозрачный процесс бюджетирования. Авторы подчёркивают, что сам алгоритм fair-share не нов; ценность, в том, как к нему подведены вводные: дерево повторяет структуру исследовательских программ, а веса задают менеджеры.

Кому это важно

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

Как это применить

Пост описывает подход одной организации, а не готовый инструмент. Из него можно взять принципы: каждый запрос на GPU должен быть профинансирован бюджетом, иначе он не защищён от вытеснения; решения о долях принимают менеджеры, лучше всех знающие компромиссы (на каждом уровне иерархии, свой руководитель); задачи объявляют минимальное время работы и после него могут быть возвращены в очередь; невыделенное время остаётся бесплатным, но вытесняемым. Авторы напоминают, что близкие по духу алгоритмы уже есть в SLURM (Fair Tree) и YARN (Fair Scheduler). Важное условие: возобновляемые задачи, которые можно автоматически ставить в очередь после вытеснения.

Можно ли доверять

Это первоисточник: пост команды AI Infrastructure самой Ai2 о своей системе. Цифры (74% меньше ремонтов с участием человека, 2-3-кратный избыток спроса, 100% задач с HIGH), собственные данные авторов без независимой проверки. У 74% не приведены базовое число и период измерения. «Дополнительные 30%», субъективное ощущение пользователя, а не измеренный прирост, а 35% у проекта A1, иллюстративный пример на схеме. Раздел о симуляциях в доступной части материала обрывается, поэтому результаты симуляций и итоговые выводы здесь не пересказаны.

Риски и подводные камни

Авторы сами предупреждают, что изменения политики планирования могут иметь непредвиденные последствия: задача с нулевой суммой, где время одному исследователю означает отнятие у другого. Схема требует постоянного процесса пересмотра бюджетов, по словам авторов, они продолжают его дорабатывать, и частых возможностей для исследователей отстаивать нужное им время. Сама логика зависит от решений менеджеров о долях и от честного объявления минимального времени работы. Задачи с нулевым минимальным временем всегда могут быть вытеснены, а после минимального времени планировщик вправе перебалансировать нагрузку, автоматически возвращая возобновляемые задачи в очередь.

«Новый планировщик создаёт ощущение, что у нас на 30% больше вычислительных мощностей.»

— Chris Clark, пользователь нового планировщика Ai2