В последнее время мы наблюдаем большой рост вариаций модели Transformer (Efficient Transformers: A Survey). Однако использование улучшенных вариантов требует пользовательских сложных реализаций, скрытых во множестве плотных библиотек. В результате мы все еще видим неоптимизированных трансформеров в дикой природе, с простыми блоками внимания или редкими блоками внимания, без операций с плавленным трансформатором, что приводит к увеличению времени обучения и увеличению затрат на обучение.
xFormers предоставляет простое компонуемое решение для создания Transformers, абстрагируя оптимизированные ядра графического процессора (которые включаются автоматически для вас), объединенные слои Transformer, которые выходят за рамки того, что может предложить PyTorch, оставаясь бесшовным для пользователя (просто torch.nn.Module). В сочетании с Lightning это создает чрезвычайно мощный инструмент для обучения моделей с высоким уровнем гибкости.
Ниже мы рассмотрим пример реализации Vision Transformer с использованием xFormers и Lightning, а также несколько советов по масштабированию вашей модели на нескольких графических процессорах с минимальными изменениями кода. Сосредоточьтесь на исследованиях, а не на шаблонах!
Что такое xFormers?
xFormers сводится к модульной библиотеке компонентов преобразователя, дополнительно предоставляя фабрики конфигураций для гибкого создания экземпляра модели преобразователя. Абстракции в xFormers также позволяют вам легко реализовывать свои собственные модули.
Быстрое прототипирование с помощью xFormers и Lightning
Ниже приведен фрагмент модуля Vision Transformer Lightning, адаптированный из этого удивительного репозитория для включения xFormers. Использование конфигураций xFormer позволяет беспрепятственно менять блоки внимания и другие элементы модели Transformer.
Мы определяем конфиг как простой словарь (здесь мы определяем только кодировщик) и передаем аргументы, управляющие всем, от типа внимания до стиля отсева или нормы слоя. Это упрощает настройку преобразователя из конструктора LightningModules.
Например, чтобы переключиться на другой механизм внимания, все, что нам нужно сделать, это передать attention="linformer" конструктору VisionTransformer.
В сочетании с модулем Lightning CIFAR мы можем очень легко обучить эту модель! Мы определяем DataModule + преобразования данных, которые мы хотели бы использовать. Затем мы создаем экземпляр нашего VisionTransformer и Lightning Trainer. Наконец, мы вызываем fit, а Lightning берет на себя все остальное!
Если вы используете установку Triton и работаете на Ampere GPU, xFormers автоматически включит очень эффективные ядра CUDA для ускорения обучения бесплатно, а также автоматически включит объединенные слои для типичных операций Transformer.
Сначала мы определяем наш импорт:
Затем мы определяем наши преобразования данных и модуль данных:
Наконец, мы определяем наш VisionTransformer и тренируемся с помощью Lightning Trainer!
Чтобы дать представление о том, насколько важна возможность замены блоков во время экспериментов, ниже приведены данные о пропускной способности в xFormers с различными блоками внимания при различной длине последовательности. Это чрезвычайно важно для таких задач, как суммирование текста в больших диапазонах токенов, что в конечном итоге снижает затраты.
Используйте DeepSpeed для масштабирования трансформаторов
Включение DeepSpeed обеспечит распределенную оптимизацию памяти, что позволит масштабировать модель до больших размеров. Для PyTorch Lightning все, что требуется, — это аргумент, передаваемый в Lightning Trainer!
Мы выбираем DeepSpeed ZeRO Stage 2 при обучении прежде всего потому, что это хороший компромисс между скоростью и оптимизацией памяти. С помощью ZeRO Stage 3 мы достигаем больших размеров модели, однако наблюдаем снижение времени обучения и даже больше, если мы включаем разгрузку. При обучении с нуля в масштабе это может стать узким местом, и масштабирование количества графических процессоров обычно является лучшим фактором для контроля. Дополнительную информацию см. в обширной документации по моделированию параллелизма!
Совет при использовании DeepSpeed Stage 3 — использовать хук configure_sharded_model в LightningModule для определения ваших модулей. Под капотом это значительно экономит время инициализации за счет разделения вашей модели между графическими процессорами на модуль.
Будущая работа в PyTorch устранит необходимость в таком хуке в будущем (см. Метаустройство для получения дополнительной информации).
Следующие шаги
Мы надеемся, что xFormers и Lightning сделают эффективные модели трансформаторов стандартными, поскольку размеры моделей продолжают увеличиваться до триллионов, предоставляя исследователям инструменты для творчества, экспериментируя с их собственными компонентами трансформаторов.
Посмотреть все доступные модули и подробные примеры можно в документации xFormers. Взгляните на пример microGPT, чтобы увидеть, как xFormers и Lightning можно использовать в НЛП. Если вас интересуют распределенные оптимизации при масштабировании на несколько графических процессоров, ознакомьтесь с Документацией по модели Lightning Parallel!
Спасибо Бенджамину Лефодо и всей команде xFormers! Огромное спасибо Phoeby Naren за иллюстрацию :)
