2BP: 2-stufige Rückwärtspropagation

2BP: 2-Stage Backpropagation

May 28, 2024

Autoren: Christopher Rae, Joseph K. L. Lee, James Richings

cs.AI

Zusammenfassung

Mit zunehmender Größe und Komplexität von Deep Neural Networks (DNNs) überschreiten sie oft die Speicherkapazität eines einzelnen Beschleunigers, was die Aufteilung der Modellparameter auf mehrere Beschleuniger erforderlich macht. Die Pipeline-Parallelität ist eine häufig verwendete Aufteilungsstrategie zur Schulung großer DNNs. Allerdings werden aktuelle Implementierungen der Pipeline-Parallelität unbeabsichtigt durch die automatischen Differentiationswerkzeuge, die von ML-Frameworks bereitgestellt werden, ausgebremst. Dieses Papier stellt das 2-stufige Backpropagation (2BP) vor. Durch die Aufteilung des Rückpropagierungsschritts in zwei separate Phasen können wir die Leerlauf-Rechenzeit reduzieren. Wir haben 2BP an verschiedenen Modellarchitekturen und Pipelining-Zeitplänen getestet und in allen Fällen eine Steigerung der Durchsatzleistung erzielt. Unter Verwendung von 2BP konnten wir im Vergleich zu traditionellen Methoden bei der Schulung eines LLaMa-ähnlichen Transformers mit 7 Milliarden Parametern auf 4 GPUs eine Steigerung der Durchsatzleistung um das 1,70-fache erreichen.

English

As Deep Neural Networks (DNNs) grow in size and complexity, they often exceed the memory capacity of a single accelerator, necessitating the sharding of model parameters across multiple accelerators. Pipeline parallelism is a commonly used sharding strategy for training large DNNs. However, current implementations of pipeline parallelism are being unintentionally bottlenecked by the automatic differentiation tools provided by ML frameworks. This paper introduces 2-stage backpropagation (2BP). By splitting the backward propagation step into two separate stages, we can reduce idle compute time. We tested 2BP on various model architectures and pipelining schedules, achieving increases in throughput in all cases. Using 2BP, we were able to achieve a 1.70x increase in throughput compared to traditional methods when training a LLaMa-like transformer with 7 billion parameters across 4 GPUs.