Обучение на нескольких GPU

Обучение на нескольких GPU в LibreYOLO — это PyTorch DistributedDataParallel: один процесс на GPU, у каждого полная копия модели и своя часть каждого батча, а градиенты усредняются по рангам на каждом шаге.

Запуск на двух GPU

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

Python

from libreyolo import LibreYOLO # Защита __main__ обязательна: каждый порождённый воркер заново# импортирует этот модуль, и без неё обучение запустилось бы рекурсивно.if __name__ == "__main__":    model = LibreYOLO("LibreYOLO9s.pt")    model.train(        data="my-dataset.yaml",        epochs=100,        batch=32,     # глобальный батч: 16 изображений на GPU при двух GPU        device="0,1",    )

Если устройств больше одного, а окружения torchrun нет, train() модели сохраняет веса во временный файл, вычисляет autobatch, если он запрошен, и порождает по одному процессу-воркеру на GPU через torch.multiprocessing.spawn. Каждый воркер заново импортирует класс модели, пересобирает её из сохранённых весов и идёт обычным путём для одного устройства, потому что внутри порождённого воркера переменные окружения torchrun уже выставлены. Когда запуск завершается, лучший чекпойнт ранга 0 загружается обратно в экземпляр модели у вызывающей стороны.

device принимает "0,1", [0, 1], 0, "cuda:0", "cpu", "mps" и "auto". Порождение процессов запускается только для списка из нескольких индексов CUDA.

Защита `main` обязательна

Порождённые воркеры заново импортируют модуль, из которого их породили. Без защиты if __name__ == "__main__": этот импорт заново выполняет вызов обучения, и каждый воркер порождает собственных воркеров. Библиотека распознаёт этот случай и выбрасывает ошибку, вместо того чтобы дать рекурсии продолжиться:

spawn_ddp_train() was called from inside a spawned subprocess. This usually
means your script calls model.train(device=...) at the top level without a
'if __name__ == "__main__":' guard.

Всё, что попадает в воркер, проходит через pickle, поэтому callbacks= должен быть picklable. Класс уровня модуля подойдёт, замыкание или лямбда — нет, и ошибка прямо об этом говорит и указывает на встроенные логгеры как на альтернативу.

batch — это глобальный батч

batch — это число изображений на один шаг оптимизатора по всем GPU. Загрузчик данных каждого ранга строится с размером batch // world_size и DistributedSampler, поэтому batch=32 на двух GPU означает 16 изображений на GPU, а не 32.

Батч, который не делится нацело на world size, приводит к ошибке, а не к тихому обучению с другим размером:

batch=6 is the global batch and must be divisible by world_size=4: each rank
trains at batch // world_size, so this value would silently train at a
different global batch than requested. Use batch=4 or batch=8.

Градиенты усредняет сам DDP, поэтому значение функции потерь передаётся дальше без масштабирования. Домножение его на world size сверх этого раздуло бы эффективную скорость обучения примерно во столько раз, сколько GPU.

Autobatch под DDP

batch=-1 работает и возвращает глобальный батч, делящийся на world size.

Python

from libreyolo import LibreYOLO if __name__ == "__main__":    model = LibreYOLO("LibreYOLO9s.pt")    # Замеряется один раз на GPU 0, результат кратен world size.    model.train(data="my-dataset.yaml", batch=-1, device="0,1")

На пути с порождением процессов замер идёт в родительском процессе на первом устройстве, ещё до появления воркеров, поэтому каждый воркер получает конкретное целое число и межпроцессная координация не нужна. Под torchrun замер делает ранг 0 и рассылает результат одним тензором типа long.

Замер определяет, сколько помещается на одной GPU, и умножает результат на world size. Если задан nbs, глобальный батч ограничивается сверху значением nbs и округляется вниз до кратного world size, поэтому добавление GPU уменьшает число шагов накопления, а не батч на одну GPU. Как устроен сам замер, описано на странице Гиперпараметры.

SyncBatchNorm

При DDP слои BatchNorm каждого ранга видят только свою часть батча. При размере batch // world_size эта часть может оказаться настолько маленькой, что накопленная статистика ухудшит сошедшуюся модель по сравнению с запуском на одной GPU.

sync_bn=True превращает каждый BatchNorm в SyncBatchNorm, так что статистика считается по глобальному батчу. Преобразование происходит только при активном распределённом режиме, поэтому на запуск с одной GPU этот флаг не влияет ни в одном из положений.

Для свёрточных семейств с большим числом слоёв BatchNorm он и так включён по умолчанию: YOLOX, YOLOv7, YOLOv9 и его варианты, YOLO-NAS, PicoDet, RTMDet и FOMO. У всех остальных семейств по умолчанию выключен. Если модель содержит BatchNorm, sync_bn выключен, а батч на ранг меньше 16, тренер выдаёт предупреждение.

Python

from libreyolo import LibreYOLO if __name__ == "__main__":    model = LibreYOLO("LibreRTDETRr18.pt")    model.train(        data="my-dataset.yaml",        batch=32,        device="0,1",        sync_bn=True,    )

У sync_bn нет флага в CLI. Это аргумент Python.

Запуск через torchrun

torchrun тоже работает и подходит лучше, когда запуском процессов уже управляет планировщик кластера. Напишите скрипт под одно устройство и дайте torchrun выставить переменные окружения с рангами.

train.py

from libreyolo import LibreYOLO if __name__ == "__main__":    model = LibreYOLO("LibreYOLO9s.pt")    model.train(data="my-dataset.yaml", epochs=100, batch=32)

Запуск

torchrun --nproc_per_node=2 train.py

Не совмещайте эти два способа. Если окружение torchrun присутствует, device="0,1" процессы не порождает: тренер берёт cuda:LOCAL_RANK, а количеством процессов управляет torchrun.

Поведение рангов

Все побочные эффекты — за рангом 0. Он определяет каталог запуска и рассылает выбранное имя, чтобы все ранги сошлись на одном, пишет чекпойнты и артефакты, вызывает пользовательские колбэки и логгеры. Остальные ранги обучаются и отдают градиенты.

Каждый ранг инициализирует генератор случайных чисел загрузчика данных и аугментаций своим значением, выведенным из заданного seed, поэтому ранги не выбирают одинаковые аугментации.

Платформа и бэкенд

Бэкенд выбирается автоматически: NCCL, когда доступны и CUDA, и NCCL, иначе Gloo. Под Windows NCCL не собирается, поэтому запуски на Windows получают Gloo без какой-либо настройки. Группа процессов инициализируется с таймаутом в три часа.

Что не работает под DDP

Захват CUDA-графа. cuda_graph=True пишет одну строку в лог, и обучение идёт в eager-режиме. См. Производительность обучения.
Профилировщик обучения. profile=True игнорируется с предупреждением.

Автоматическое порождение процессов поддерживают не все семейства. Таких двадцать четыре — все обучаемые семейства детекции, классификации, семантической сегментации и восстановления. Семейство без такой поддержки, получив устройство из нескольких GPU, выбрасывает ошибку, называя API модели и команду torchrun, а не обучается молча на одной GPU.

Смотрите также

Гиперпараметры — про batch, nbs и возобновление.
Логгеры экспериментов — про требование picklable к колбэкам.
Обучение на арендованной GPU — про аренду машины с несколькими GPU.