DeepSpeed-ZeRO-Stufen: Welcher Trainingszustand wird aufgeteilt?
Automatische Übersetzung
Dieser Artikel wurde automatisch aus der englischen Originalversion übersetzt.
ZeRO-1 verteilt den Optimizer-Zustand auf datenparallele GPUs, ZeRO-2 zusätzlich die Gradienten und ZeRO-3 zusätzlich die Modellparameter. Teste die niedrigste Stufe, bei der die Workload in den Speicher passt und einen akzeptablen Throughput erreicht. Eine höhere Stufe spart mehr Speicher für den Modellzustand, kann aber mehr Kommunikation erfordern.
Diese Stufen ändern, wo der Trainingszustand gespeichert wird. Sie reduzieren den Speicherbedarf der Activations nicht automatisch.
Ein Speicherbeispiel mit ausdrücklich genannten Annahmen
Angenommen, das Training verwendet FP16-Weights und -Gradienten sowie FP32-Master-Weights für Adam und zwei Optimizer-Momente. Pro Parameter werden dann 2 Bytes für Weights, 2 für Gradienten und 12 für den Optimizer-Zustand benötigt: insgesamt 16 Bytes.
Bei 7.5 Milliarden Parametern auf acht GPUs ergibt die Berechnung für den Modellzustand:
| Stufe | Auf GPUs verteilter Zustand | GB für den Modellzustand pro GPU |
|---|---|---|
| Gewöhnliche Datenparallelität | Keiner | 120 |
| ZeRO-1 | Optimizer-Zustand | 41.25 |
| ZeRO-2 | Optimizer-Zustand und Gradienten | 28.125 |
| ZeRO-3 | Alle drei Komponenten | 15 |
Die Werte sind dezimale GB. Activations, vorübergehend zusammengeführte Parameter, zusätzlicher Speicherbedarf des Allocators und Runtime-Puffer sind nicht enthalten. Verwendet der Trainings-Stack andere Speicherpräzisionen, muss die Berechnung angepasst werden.
Die ZeRO-Veröffentlichung erklärt die Aufteilung der Zustände und analysiert die Kommunikation. ZeRO-3 rekonstruiert benötigte Parameter durch All-Gather-Operationen; die Analyse der vollständigen Aufteilung in der Veröffentlichung ergibt etwa das 1.5-Fache des Kommunikationsvolumens gewöhnlicher Datenparallelität. Dieses Verhältnis gilt unter den Annahmen der Analyse und ist keine Garantie für jedes Netzwerk und jede Konfiguration.
Anhand des Speicherprofils auswählen
Wenn die Parameter in den Speicher passen und die Optimizer-Zustände den Speicherfehler verursachen, teste Stufe 1 oder 2. Wenn vollständig replizierte Parameter nicht hineinpassen, teste Stufe 3. Wenn Activations den Speicherbedarf bestimmen, kombiniere die gewählte Stufe mit Activation Checkpointing oder einem kleineren Trainings-Batch.
DeepSpeed unterstützt auch das Auslagern von Zuständen auf CPU und NVMe. Prüfe die aktuelle ZeRO-Konfiguration und miss die Übertragungskosten auf dem tatsächlichen Host und Speichersystem. Die zusätzliche Kapazität durch Auslagerung kann helfen, aber langsame Übertragungen können den Throughput senken.
Der ZeRO-Abschnitt im LLM Engineering Guide enthält eine ausführbare Berechnung für andere Parameterzahlen und GPU-Gruppen.