Обучение крупных моделей глубокого обучения остаётся уделом немногих: чем больше параметров, тем дороже инфраструктура. Например, чтобы обучить модель на 10 миллиардов параметров, обычно требуется узел уровня DGX-2 с 19 видеокартами NVIDIA V100 — такая конфигурация стоит больше 100 тысяч долларов. Для большинства исследователей и даже многих университетских лабораторий это недостижимо.
Команда из Калифорнийского университета в Мерседе совместно с Microsoft предложила способ обойти это ограничение — технологию ZeRO-Offload. Она позволяет обучать модели с миллиардами параметров на одной видеокарте без переработки самой модели. По сути, это гибридная GPU-CPU технология обучения, сочетающая высокую вычислительную эффективность с почти линейной масштабируемостью.
В чём проблема существующих подходов
Основная сложность крупномасштабного обучения — так называемые «состояния модели»: параметры, градиенты и состояния оптимизатора, которые занимают огромный объём памяти GPU. Существующие методы гетерогенного обучения, снижающие нагрузку на память видеокарты, в основном рассчитаны на активационную память небольших сверточных сетей и плохо подходят для по-настоящему больших моделей.
Стандартный способ масштабирования — обычный параллелизм данных — требует дублирования и данных, и вычислений на каждом устройстве, что делает его непригодным для гетерогенного обучения. ZeRO-Offload решает эту проблему иначе: технология задействует и память, и вычислительные мощности CPU для разгрузки GPU, эффективно работая совместно с параллелизмом данных на базе ZeRO. При этом в памяти CPU хранится всего одна копия состояний оптимизатора независимо от степени параллелизма по данным — это и даёт масштабируемость вплоть до 128 GPU.
Как устроен ZeRO-Offload
Разработчики опирались на три принципа: эффективность, масштабируемость и удобство использования. Исследователи нашли оптимальную схему разделения данных и вычислений между CPU и GPU: градиенты, состояния оптимизатора и сами вычисления оптимизатора переносятся на процессор, а параметры модели вместе с прямым и обратным проходом остаются на видеокарте.
По результатам экспериментов такой подход даёт десятикратный прирост допустимого размера модели при минимальных издержках на коммуникацию и ограниченной нагрузке на CPU. В итоге на одной видеокарте NVIDIA V100 удалось обучить модель на 13 миллиардов параметров, достигнув производительности 40 терафлопс.
Доступность
ZeRO-Offload уже включена в открытую библиотеку DeepSpeed для PyTorch и добавляется в существующие пайплайны обучения буквально несколькими строками кода. Авторы описали технологию в статье «ZeRO-Offload: Democratizing Billion Scale Model Training», опубликованной на arXiv.
Источник: MarkTechPost.