В Databricks мы постоянно разрабатываем и оцениваем новые методы исследований в области ИИ, чтобы сделать наши продукты более интеллектуальными, эффективными и полезными для клиентов.
Наши исследования охватывают такие продукты, как Genie, Smart Routing и AI Functions. Недавно мы опубликовали информацию о том, как мы ускорили поиск в 3 раза, разработали новые методы автономного обучения с подкреплением и расширили границы извлечения данных из сложных документов.
Чтобы создавать отличные модели ИИ и управлять сотнями параллельных исследовательских задач, нам нужна инфраструктура, способная не отставать от темпа. Именно поэтому мы используем Databricks AI Runtime (AIR).
В этой статье мы расскажем, как AIR помогает нам быстро экспериментировать, масштабировать ресурсоемкие задачи обучения и адаптироваться к обучению под управлением агентов, а также поделимся опытом управления ресурсами GPU в исследовательских проектах.
Проблема: множество экспериментов и совершенно разные потребности в вычислениях
Наша исследовательская команда разрабатывает модели, лежащие в основе Databricks AI Functions, включая ai_parse_document и ai_extract. Например, функции Databricks AI обеспечивают высочайшее качество синтаксического анализа документов благодаря моделям, обученным понимать корпоративные документы. Достижение такого качества требует масштабных исследований различных архитектур, наборов данных, методов обучения и гиперпараметров, охватывающих десятки тысяч запусков обучения и оценки.
Мы используем контролируемое дообучение (SFT), автономное обучение с подкреплением (RL) и обучение с подкреплением в реальном времени (online RL) для создания моделей размером от нескольких миллиардов до почти триллиона параметров. В зависимости от рабочей нагрузки мы применяем графические процессоры NVIDIA H100 или B300, доступные через AIR. Наши крупнейшие запуски обучения могут требовать совместной работы сотен GPU в течение дней или недель, а также высокоскоростной сети для поддержания эффективности обучения.
Поддержка таких рабочих нагрузок создает три проблемы:
- Отказоустойчивость для длительных задач обучения. Обучение может потребовать тысяч часов работы GPU. Сбой одного узла или сети RDMA может оставить сотни графических процессоров простаивать и заставить нас перезапускаться с более ранней контрольной точки, теряя часы прогресса. Для проектов обучения с подкреплением, таких как KARL, надежная инфраструктура имеет решающее значение.
- Быстрое экспериментирование и прототипирование. Улучшение функций ИИ означает тестирование тысяч подходов: от различных моделей и методов обучения до рецептов данных. Нам необходимо быстро запускать эксперименты без необходимости многократно подготавливать инфраструктуру или перестраивать окружение.
- Управление парком GPU. Мы балансируем между несколькими исследовательскими проектами с различными требованиями к вычислениям, приоритетами и сроками, сохраняя при этом гибкость для реагирования на новые методы исследований или модели с открытым исходным кодом.
Внедрение AI Runtime для решения этих задач
Наша исследовательская команда была «клиентом номер ноль» для AI Runtime с момента создания продукта в прошлом году. Мы используем AIR для разработки, обучения и тестирования моделей, лежащих в основе ИИ-продуктов Databricks, полагаясь на ее инфраструктуру GPU и интеграцию с нашими существующими инструментами.
Встроенная отказоустойчивость GPU
AIR выполняет предварительные проверки назначенных узлов GPU перед началом обучения, чтобы убедиться в их исправности и возможности эффективного взаимодействия. Она также мониторит состояние узлов во время выполнения, обнаруживая такие проблемы, как сбои связи NCCL, тепловое ограничение или троттлинг питания, а также проблемы с драйверами GPU. Если узел выходит из строя и приводит к сбою задачи, AIR обнаруживает проблему и изолирует затронутый узел, чтобы последующие задачи не планировались на той же проблемной инфраструктуре.
Это особенно важно для таких проектов, как KARL, где сложная генерация развертывания и постобучение моделей с более чем 100 миллиардами параметров могут требовать тысяч часов работы GPU. Вместо многократного устранения неполадок инфраструктуры мы можем сосредоточиться на улучшении наших моделей.
Быстрая итерация с помощью CLI AIR и пользовательских контейнеров
Исследования требуют от нас частой смены конфигураций моделей, подходов к обучению и кода. С помощью CLI AIR мы можем программно запускать рабочие нагрузки, указывать необходимые GPU и позволять AIR управлять подготовкой и планированием ресурсов, включая высокоскоростную сеть для распределенного обучения в облаках.
Мы также полагаемся на пользовательские контейнеры Docker для локальной разработки. AIR позволяет использовать те же образы для обучения, обеспечивая совместимость без необходимости воссоздавать наше окружение. Мы можем загружать нашу кодовую базу, выполнять пользовательский код в этих контейнерах и использовать Databricks Artifact Registry для регистрации и совместного использования образов в различных экспериментах. Инвестиции AIR в задержку запуска и оптимизацию контейнеров помогают нам запускать задачи обучения за считанные секунды, сокращая циклы обратной связи при тестировании новых идей.
MLflow и логирование для отслеживания экспериментов
Мы используем MLflow для отслеживания наших исследовательских экспериментов. AIR автоматически создает и управляет экспериментами MLflow для наших задач обучения, избавляя от необходимости отдельной настройки.
Мы можем отслеживать статус задач, транслировать и проверять логи обучения, а также получать метрики для понимания хода экспериментов. Это также дает нашим ИИ-агентам доступ к информации, необходимой для оценки результатов и определения следующих шагов.
Бесшовный и управляемый доступ к исследовательским данным
Все наши исследовательские данные хранятся в Unity Catalog. Благодаря подключению Unity Catalog через FUSE, AIR позволяет нам получать доступ к этим данным так, будто они являются частью локальной файловой системы, даже при работе в разных рабочих пространствах Databricks.
Это делает наши данные доступными, отслеживаемыми, разделяемыми и регулируемыми без изменения того, как наш код обучения взаимодействует с файлами. В сочетании с MLflow и нашими средами на базе Docker это обеспечивает единый рабочий процесс от локальной разработки до распределенного обучения.
Управление емкостью GPU в исследовательских проектах
Наши рабочие нагрузки имеют совершенно разные потребности в ресурсах. Для эксперимента по отладке может потребоваться всего несколько GPU, в то время как для крупного запуска обучения — сотни в течение дней или недель. Различные проекты также имеют конкурирующие сроки и приоритеты.
Возможности AIR по расстановке приоритетов и преемции рабочих нагрузок помогают нам распределять ресурсы между проектами, уделять первоочередное внимание чувствительным ко времени экспериментам и использовать избыточную емкость, когда она доступна. Например, свободные ресурсы могут поддерживать дополнительные исследовательские эксперименты в выходные дни без ущерба для более приоритетных задач. наша цель — поддерживать доступность ресурсов GPU для важных экспериментов без необходимости исследователям вручную координировать инфраструктуру.
Масштабирование на базе AIR с обучением под управлением агентов
Популяризированный Андреем Карпатым, autoresearch меняет подход к разработке моделей ИИ. Вместо того чтобы исследователи вручную запускали и оценивали каждый эксперимент, мы все чаще используем агентов для предложения изменений, запуска экспериментов, оценки результатов и определения того, что попробовать дальше.
Это изменение драматически увеличило объем наших экспериментов. В рамках исследовательских задач по функциям ИИ количество задач обучения выросло в 8 раз.
Поскольку агенты берут на себя все большую часть процесса исследований, наша инфраструктура должна не отставать. Медленный запуск или запутанная ошибка конфигурации могут быть терпимы, когда исследователь запускает несколько экспериментов. Когда агенты запускают сотни таких задач, те же проблемы быстро становятся узкими местами.
Исследования под управлением агентов требуют быстрого запуска задач, программных интерфейсов, понятных и полезных сообщений об ошибках, надежного выполнения и доступных результатов обучения.
Наш рабочий процесс автоисследований следует непрерывному циклу:
- Предложение: агенты определяют потенциальные улучшения для моделей или подходов к обучению.
- Запуск испытаний: они отправляют множество экспериментов в AIR параллельно.
- Обучение: они оценивают метрики обучения, чтобы определить, что повышает производительность.
- Обновление плана: они используют результаты для уточнения своего подхода и запуска новых экспериментов.
AIR особенно хорошо подходит для этого рабочего процесса, так как агенты могут взаимодействовать с ней программно и использовать преимущества быстрого запуска задач, встроенной надежности и возможности масштабирования. С помощью интерфейса командной строки Databricks и databricks-air-skills наши агенты могут запускать параллельные рабочие нагрузки, контролировать статус задач, получать метрики MLflow и выполнять итерации на основе результатов.
Это позволяет нам проводить больше экспериментов, в то время как наши исследователи сосредоточены на разработке новых идей и улучшении качества моделей. По мере того как мы продвигаемся к исследованиям под управлением агентов, AIR предоставляет нам инфраструктуру для поддержки как более крупных моделей, так и более высокой скорости экспериментов.
Начните работу сегодня
Наш опыт использования AI Runtime показал нам, насколько важна надежная и гибкая GPU-инфраструктура для исследований. От крупномасштабного обучения и контролируемого доступа к данным до экспериментов под управлением агентов — AIR позволяет нам тратить меньше времени на управление инфраструктурой и больше времени на разработку лучших моделей.
Если вы обучаете или дорабатываете свои собственные модели, ознакомьтесь с документацией AI Runtime, чтобы начать работу.
AI Runtime также обеспечивает основу для наших исследований в таких областях, как ускоренный поиск, обучение с подкреплением и понимание документов. Ознакомьтесь с другими нашими материалами на странице Databricks AI Research.
