Colossal-AI позволяет тренировать огромные модели ИИ на обычном ноутбуке



🤷🏼 Обучение с нуля или даже тонкая настройка больших моделей ИИ обычно недоступна рядовым пользователям. Она требует десятков или сотен графических процессоров и глубоких познаний архитектуры искусственного интеллекта.



🙋🏻‍♂️ Компания HPC-AI Tech попыталась обойти этот барьер. В 2021 году она представила Colossal-AI — единую систему глубокого обучения, которую в теории может освоить даже инженер-любитель с обычным лэптопом.



📉 Инструмент на основе фреймворка PyTorch использует гетерогенный подход в обучении нейросетей. В его основе — оптимизация потребления памяти во время тренировки. Здесь поддерживаются тензорный параллелизм и контрольные точки активации.



💻 Для работы Colossal-AI не обязательно нужны суперкомпьютеры. Инструмент может обучать большие модели с 1,5 млрд параметров на ноутбуке с RTX 206 (6 ГБ) на борту. А компьютер с графическим чипом RTX 3090 (24 ГБ) способен натренировать нейросеть с 18 млрд параметров.



💽 Недавно у Colossal-AI появился модуль Gemini — аналог Zero-offload от майкрософтской библиотеки DeepSpeed. По словам разработчиков, их менеджер памяти превосходит конкурента, потому что эффективнее распределяет данные.



🆓 Например, Zero-offload статически делит данные модели между памятью CPU и GPU, и их расположение остается постоянным для разных конфигураций обучения. Поэтому когда память GPU закончится, система выдаст ошибку, даже если память процессора будет не заполнена. При этом Colossal-AI сможет завершить тренировку, перекинув часть информации на ЦП.



🔁 Для обучения моделей Colossal-AI делит итерации на два этапа: прогревочный и непрогревочный. На первом система собирает информацию об использовании памяти, а на втором — использует ее для перераспределения тензоров. Это позволяет свести к минимуму объем перемещаемых данных между CPU и GPU.



Исходный код проекта открыт, поэтому воспользоваться возможностями Colossal-AI может любой желающий.