Nvidia veröffentlicht Werkzeug zum Mixed-Precision-Training in PyTorch
Das Tool mit dem Namen Apex hilft beim Mischen von 16-bit- und 32-bit-Gleitkommazahlen in PyTorch.
(Bild: Shutterstock)
- Rainald Menge-Sonnentag
Der Grafikkartenhersteller Nvidia hat mit Apex ein Tool für das Mixed-Precision-Training in PyTorch herausgegeben. Der Name steht für "A PyTorch Extension". Mit dem Tool lassen sich Berechnungen mit Gleitkommazahlen in 32-bit (FP32) und 16-bit (FP16) mischen. Apex enthält eine Bibliothek zum Automatisieren des Mixed-Precision-Trainings und eine Wrapper-Klasse zum Optimieren von FP16-Berechnungen.
Die richtige Mischung
Üblicherweise setzen Machine-Learning-Frameworks wie TensorFlow, Keras und auch PyTorch standardmäßig auf FP32. Im Vergleich dazu bietet FP16 eine geringere Präzision, benötigt gleichzeitig aber nur halb soviel Speicher. Entsprechend verdoppelt sich der Durchsatz bei der Speicherbandbreite. Noch stärker wirkt sich die Berechnung auf Volta Tensor Cores bei generalisierten Matrix-Matrix-Multiplikationen (GEMM) aus, bei denen FP16 laut Nvidia einen Vorteil vom Faktor Acht gegenüber FP32 bietet.
Der Kniff liegt somit in der richtigen Mischung aus FP32- und FP16-Arithmetik in Deep Neural Networks (DNNs). Nvidia hat dazu 2017 einen Blogbeitrag veröffentlicht. Vereinfacht gesagt arbeitet der Großteil des optimierten neuronalen Netzes mit FP16-Werten, während wenige Operationen wie größere Reduktionsvorgänge in FP32 ablaufen.
Automatisiert und optimiert
PyTorch bringt von Haus aus eine Funktion für das Mixed-Precision-Training mit: Der Aufruf von half() wandelt die Parameter von Modulen beziehungsweise die Daten von Tensoren von FP32 in FP16 um. Anschließend erfolgen alle Berechnungen mit 16-bit-Gleitkomma-Arithmetik. Da PyTorch eine Anbindung an die Nvidia-Bibliotheken cuBLAS und cuDNN mitbringt, kommt die Umwandlung unter anderem der GEMM zugute.
Ein Bestandteil von Apex ist die Bibliothek Amp (Automatic Mixed Precision), die das Mixed-Precision-Training weitgehend automatisieren soll. Dafür setzt sie auf eine Kombination aus einem Whitelist- und einer Blacklist-Ansatz. Funktionen auf der Whitelist arbeiten mit FP16-Argumenten und umfassen beispielsweise Matrix-Multiplikationen oder Convolution-Funktionen, für die 16-bit ausreichend präzise ist und die am meisten davon profitieren. Blacklist-Elemente benötigen FP32-Inputs, da die Präzision von 16-bit-Werten potenziell nicht ausreicht. Beispiele sind Verlustfunktionen zum Optimieren des Lernalgorithmus (Loss Functions).
Nicht gekennzeichnete Funktionen haben keine besondere Präferenz: FP16 reicht aus, aber FP32 bringt keine relevante Performanceeinbußen. Die Bibliothek stellt für diese Funktionen lediglich sicher, dass alle Argumente denselben Typ aufweisen. Eine genaue Erklärung der Arbeitsweise von Amp lässt sich dem GitHub-Repository entnehmen.
Die Python-Klasse FP16_Optimizer ist der zweite Bestandteil von Apex. Sie dient als Wrapper für vorhandene PyTorch-Optimierer, die das statische oder dynamische Loss Scaling übernimmt. Weitere Automatisierungen wie bei Amp entfallen beim Einsatz des Wrapper, der damit Entwicklern mehr Kontrolle lässt. FP16_Optimizer lässt sich mit wenigen Anpassungen in vorhandenen PyTorch-Code integrieren.
Weitere Details finden sich im Nvidia-Blog. Der Sourcecode für die Bibliothek Amp und die Klasse FP16_Optimizer liegen auf GitHub. (rme)
