Эффективное обучение на нескольких GPU
Multi-gpu сетап может быть полезен для дополнительного масштабирования и оптимизации обучения, если даже после техник увеличения эффективности на одной gpu, производительность обучения вас не устраивает.
Оптимизация достигается за счет разного рода параллелизма:
дока huggngface: https://huggingface.co/docs/transformers/perf_train_gpu_many
гайд accelerate: https://huggingface.co/docs/accelerate/index
Прежде чем пробовать конкретный метод, нужно понять несколько моментов:
- помещается ли модель полностью на одну gpu
- помещается ли самый жирный слой модели на одну gpu
- gpu внутри одной ноды или нет
- насколько быстрая связь между нодами
В конца статьи будут даны конкретные рекомендации в зависимости от вашего сетапа.
Data parallelism
Есть два вида параллелизации данных: Data Parallel (DP) и Distributed Data Parallel (DDP) - полезная статья.
Почти всегда имеет смысл использовать DDP, потому что он быстрее и лучше утилизирует gpu. DP эффективнее DDP, если у gpu нет nvlink - это технология для обменивания данными без задержек и с высокой пропускной способностью между gpu и другими gpu или cpu, но у современных карточек, особенно для работы с ИИ, эта технология есть. DDP также поддерживает обучение на разных нодах.
Алгоритм Distributed Data Parallel (DDP):
- В начале обучения основной процесс копирует модель с 0-ой gpu на остальные
то есть мы запускаем столько одинаковых процессов, сколько у нас свободных gpu (world size) - Во время обучения берем очередной батч от даталоадера и делим его между всеми gpu
- forward на каждом процессе проходит независимо
- после backward мы усредняем градиенты со всех gpu и делаем одинаковое обновление весов для каждой модели
- повторяем 2, 3, 4
Таким образом в DDP взаимодействие между карточками происходит только во время отправки градиентов.
ZeRO Data Parallelism
Начну с хорошей аналогии, которую украл отсюда.
Механизм похож на стратегию распределение вещей в групповом походе: человек А несет палатку, человек B несет горелку и котелок, а человек С несет топор. Каждую ночь они делятся тем, что у них есть, и получают от других то, чего у них нет, а утром снова упаковывают свое распределенное снаряжение и продолжают путь. Это и есть ZeRO DP/Sharded DDP. Сравните эту стратегию с более простой, при которой каждый человек несет свою собственную палатку, горелку, котелок и топор (аналогично DataParallel (DP и DDP) в PyTorch), что было бы гораздо менее эффективно.
- равномерно распределяет параметры слоя модели между gpu, то есть хранится на gpu суммарно только одна копия модели, в отличие от DDP, где хранится world size копий модели
- во время forward слои подгружают те параметры, которые им нужно для расчета, а после подчищают все не свои параметры
Здесь мы разделяем между gpu веса конкретного слоя, что делает метод похожим на tensor parallelism.
Tensor Parallelism
Смысл в том, чтобы разбить слои модели между gpu и делать вычисления на разных gpu параллельно, а собирать тензор вместе только там, где это необходимо.
Но из-за того, что gpu нужно часто общаться между собой, то такой метод параллелизации лучше использовать внутри одной ноды, чтобы связь между gpu было быстрой.
Model and pipeline parallelism
В model parallelism, также известной как vertical parallelism, мы распределяем последовательности слоев модели между разными gpu, в отличие от ZeRO и от tensor parallelism.
- gpu 0 делает обычный forward слоями 1, 2, 3
- передаем данные на gpu1
- gpu1 делает обычный forward слоями 4, 5, 6
- backward делаем аналогично, но в другом порядке
Основная проблема в том, что работает только одна gpu в любой момент времени, остальные ждут.
Поэтому придумали pipeline parallelism, в котором данные батча нарезаются на куски - на картинке схема pipeline и отличия от model parallelism.
Таким образом мы минимизируем простои GPU.
Такой подход похож на аккумуляцию градиентов, поэтому часто вводят понятие micro-batch вместе с обычным batch size и mini-batch. Например, если у нас batch size 1024, при этому кол-во DP процессов 4, то наш mini-batch 256, при этом если кол-во чанков 32, то micro-batch равен 8. То есть каждая gpu в момент времени работает с размером батча 8, при этом шаги оптимизации происходят раз в 1024 семпла.
Количество чанков - гиперпараметр, поэтому его стоит настраивать, ведь при количестве чанков 1, имеем обычный model paralelism, что неэффективно, также как неэффективно брать очень большое количество чанков.
Data Parallelism + Pipeline Parallelism
Можно объединить DP и PP: gpu0 передает часть вычислений по принципу PP gpu2, при этом такая связка копируется на gpu1 и gpu3 по принципу DP
Data Parallelism + Pipeline Parallelism + Tensor Parallelism
Можно объединять и все 3 типа параллелизма, никто не запрещает, если хватит gpu…
Стратегии выбора техник параллелизации
- Модель помещается на одну GPU:
Используем Data Parallel (DP) или ZeRO параллелизм. - Модель не помещается на одну GPU:
Выбираем между Pipeline Parallel (PP), ZeRO или Tensor Parallel (TP).
Если nvlink внутри ноды быстрый - то будет примерно одинаковый перфоманс, в противном случае pipeline parallelism может выигрывать. - Если самый жирный слой модели не помещается на одну GPU:
Тогда даже pipeline parallelism не сможет обработать модель - в таком случае используем tensor parallelism или ZeRO + техники оптимизации обучения на одной GPU. - стратегии для multi-Node / multi-GPU сетапа