BMTrain: большой модели — эффективное обучение

BMTrain — это набор инструментов для эффективного обучения больших моделей.

Основные разделы сайта:

1. Обзор.
2. Документация.
3. Установка.
4. Использование.
5. Производительность.
6. Английский вариант.

Последние новости:

— 14 июня 2022 года выпущена версия BMTrain 0.1.7. Поддерживается оптимизация ZeRO-2!

— 30 марта 2022 года вышла версия BMTrain 0.1.2. Адаптирована работа с инструментами OpenPrompt и OpenDelta.

— 16 марта 2022 года опубликована версия BMTrain 0.1.1. Это первый стабильный выпуск, в котором исправлены проблемы бета-версии.

— 11 февраля 2022 года представлена первая бета-версия BMTrain 0.0.15.

Обзор

BMTrain — это эффективный инструмент для обучения больших моделей с сотнями миллиардов параметров. BMTrain позволяет обучать модели распределённо, сохраняя при этом лаконичность кода.

Документация

На сайте readthedocs.io можно найти дополнительную информацию об инструменте.

Установка

Есть два способа установки:

1. С помощью pip: pip install bmtrain.
2. Из исходного кода: скачайте пакет и запустите pip install (в будущем метод установки через setup.py будет заменён на более современный).

Установка BMTrain может занять несколько минут, так как требует компиляции C/CUDA исходного кода. Рекомендуется компилировать BMTrain непосредственно в тренировочной среде, чтобы избежать потенциальных проблем с разными средами.

Использование

Шаги по использованию BMTrain:

1. Включите BMTrain.

   • В начале кода инициализируйте BMTrain с помощью функции init_distributed.
   • Не используйте модули PyTorch для распределённого обучения при использовании BMTrain, включая torch.distributed.init_process_group и связанные функции связи.

2. Используйте оптимизацию ZeRO.

   • Замените torch.nn.Module на bmtrain.DistributedModule.
   • Замените torch.nn.Parameter на bmtrain.DistributedParameter.
   • Используйте bmtrain.CheckpointBlock для блоков трансформатора.

3. Оптимизируйте связь.

   • Замените torch.nn.ModuleList на bmtrain.TransformerBlockList.
   • Вместо цикла for module in self.module_list: x = module(x, ...) используйте x = self.module_list(x, ...). Нам удалось реализовать эти детали. Вам нужно только добавить оптимизатор и стратегию настройки скорости обучения в менеджер оптимизатора (optim_manager) с помощью функции add_optimizer, а затем позволить optim_manager выполнить операции zero_grad(), backward(), clip_grad_norm() и step().

Если вы не используете смешанную точность обучения, вам не нужно масштабировать потери, просто установите loss_scale в None в конструкторе OptimManger(loss_scale=None). Это также является параметром по умолчанию для конструктора OptimManager.

Если вы используете смешанную точность обучения, масштабирование потерь — это распространённая техника в смешанной точности обучения. Мы масштабировали потери в optim_manager.backward(loss), чтобы избежать переполнения градиента. Просто установите loss_scale на число с плавающей запятой в конструкторе OptimManager. Loss_scale будет автоматически настраиваться в процессе обучения на основе градиента.

Производительность

Мы обучили модель GPT-2 с 13 миллиардами параметров, используя 4 сервера, каждый сервер имеет 8 видеокарт V100. Мы протестировали пропускную способность каждого GPU во время процесса обучения, результаты приведены в таблице ниже.

Структура модели:

• 40 слоёв;
• 128 голов внимания;
• размер скрытого слоя 5120;
• длина последовательности 512.

ZeROa(mp=b) означает DeepSpeed + Megatron ZeRO stage a и параллелизм моделей = b.

Знак «-» в таблице означает превышение объёма памяти.

Поддержка моделей

Мы уже перенесли большинство распространённых моделей NLP в BMTrain. Вы можете найти список поддерживаемых моделей в проекте ModelCenter.

Сообщество открытого исходного кода

Приглашаем участников вносить свой вклад, следуя нашему руководству по внесению вклада.

Вы также можете связаться с нами через другие платформы:

• QQ группа: 735930538;
• официальный сайт: https://www.openbmb.org;
• Weibo: http://weibo.cn/OpenBMB;
• Twitter: https://twitter.com/OpenBMB.

Лицензия с открытым исходным кодом

Этот инструмент использует лицензию Apache 2.0 с открытым исходным кодом.

Другие примечания

Пакет инструментов BMTrain вносит изменения в PyTorch на нижнем уровне. Если ваша программа выдаёт неожиданные результаты, вы можете сообщить соответствующую информацию в issue.