Что такое DeepSpeed ZeRO Stage 1, 2, 3?

DeepSpeed ZeRO (Zero Redundancy Optimizer) — это технология оптимизации распределённого обучения больших моделей, которая снижает требования к памяти и увеличивает масштабируемость. Она делится на три стадии:

• Stage 1: Разделение оптимизатора — параметры оптимизатора (градиенты, моменты) распределяются между GPU, уменьшая дублирование.
• Stage 2: Разделение оптимизатора и градиентов — дополнительно к оптимизатору градиенты распределяются, что ещё больше снижает потребление памяти.
• Stage 3: Полное разделение параметров, оптимизатора и градиентов — каждый GPU хранит только часть параметров модели, оптимизатора и градиентов, что позволяет обучать модели, которые не помещаются в память одного устройства.

Таким образом, с каждой следующей стадией уменьшается дублирование данных в памяти, что позволяет эффективно масштабировать обучение на большое количество GPU.