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

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

ZeRO делится на три стадии:

• Stage 1 — шардирование оптимизаторских состояний (например, моментов в Adam) между устройствами. Параметры и градиенты дублируются на всех устройствах.

• Stage 2 — дополнительно к шардированию оптимизаторских состояний, шардируются также градиенты. Параметры остаются дублированными.

• Stage 3 — шардируются параметры, градиенты и оптимизаторские состояния. Это максимальное разделение, позволяющее обучать очень большие модели, которые не помещаются в память одного устройства.

Пример: при fine-tuning больших языковых моделей (LLM) ZeRO Stage 3 позволяет эффективно использовать память GPU, распределяя все данные оптимизации, что снижает требования к памяти и ускоряет обучение.