torchgpipe
v0.0.7
Una implementación de GPIPE en Pytorch. Está optimizado para CUDA en lugar de TPU.
from torchgpipe import GPipe
model = nn . Sequential ( a , b , c , d )
model = GPipe ( model , balance = [ 1 , 1 , 1 , 1 ], chunks = 8 )
output = model ( input )GPIPE es una biblioteca de paralelismo de tubería escalable publicada por Google Brain, que permite un entrenamiento eficiente de modelos grandes que consumen memoria. Según el documento, GPIPE puede entrenar un modelo de 25x más grande utilizando dispositivos 8X (TPU) y entrenar un modelo 3.5x más rápido utilizando dispositivos 4x.
GPIPE: entrenamiento eficiente de redes neuronales gigantes que utilizan el paralelismo de la tubería
Google entrenó a Amoebanet-B con parámetros de 557m sobre GPIPE. Este modelo ha alcanzado el 84.3% Top-1 y 97.0% de precisión top-5 en el punto de referencia de clasificación de Imagenet (el rendimiento de vanguardia a mayo de 2019).
GPIPE usa (a) paralelismo de la tubería y (b) recomputación automática de la propagación hacia adelante durante la propagación de retroceso, por lo tanto, aprovecha el entrenamiento de un modelo grande. Nos referimos a (b) como punto de control, siguiendo la conocida terminología en la comunidad de Pytorch.
Actualmente, TorchgPipe requiere los siguientes entornos:
Para usar TORCHGPIPE, instálelo a través de PYPI:
$ pip install torchgpipe Para entrenar un módulo con GPIPE, simplemente envuélvelo con torchgpipe.GPipe . Su módulo debe ser nn.Sequential , ya que GPIPE dividirá automáticamente el módulo en particiones con capas consecutivas. El argumento balance determina el número de capas en cada partición. El argumento chunks especifica el número de micro-lotes. La entrada, la salida y los tensores intermedios deben ser Tensor o Tuple[Tensor, ...] .
El siguiente código de ejemplo muestra cómo dividir un módulo con cuatro capas en cuatro particiones, cada una con una sola capa. Este código también divide un mini lote en 8 micro-lotes:
from torchgpipe import GPipe
model = nn . Sequential ( a , b , c , d )
model = GPipe ( model , balance = [ 1 , 1 , 1 , 1 ], chunks = 8 )
for input in data_loader :
output = model ( input )Visite TORCHGPIPE.ReadThedocs.io para obtener más información, incluidas las referencias de API.
Los detalles completos y más puntos de referencia están disponibles en Torchgpipe.readthedocs.io.
| Tamaño por lotes | cigratoria | nn.dataparallel | Goyal et al. |
|---|---|---|---|
| 256 | 21.99 ± 0.13 | 22.02 ± 0.11 | 22.08 ± 0.06 |
| 1k | 22.24 ± 0.19 | 22.04 ± 0.24 | N / A |
| 4K | 22.13 ± 0.09 | N / A | N / A |
GPIPE debe ser transparente para no introducir un ajuste adicional de hiperparámetro. Para verificar la transparencia, reproducimos la tasa de error Top-1 de ResNet-101 en ImageNet, como se informó en la Tabla 2 (c) de SGD de minibatch de gran y grande por Goyal et al.
| Experimento | U-Net (B, C) | Parámetros | Uso de memoria |
|---|---|---|---|
| base | (6, 72) | 362.2m | 20.3 GIB |
| tubería-1 | (11, 128) | 2.21b | 20.5 GIB |
| tubería-2 | (24, 128) | 4.99b | 43.4 GIB |
| tubería-4 | (24, 160) | 7.80b | 79.1 GIB |
| tubería-8 | (48, 160) | 15.82b | 154.1 GIB |
La tabla muestra cómo GPIPE facilita la escala de modelos de red en U. La línea de base denota la línea de base sin paralelismo de la tubería ni el punto de control, y la tubería -1 , -2 , -4 , -8 denota que el modelo esté entrenado con GPIPE con el número correspondiente de particiones.
Aquí utilizamos una arquitectura simplificada de Net U. El tamaño de un modelo está determinado por los hiperparámetros B y C que son proporcionales al número de capas y filtros, respectivamente.
| Experimento | Rendimiento | Acelerar |
|---|---|---|
| base | 28.500/s | 1 × |
| tubería-1 | 24.456/s | 0.858 × |
| tubería-2 | 35.502/s | 1.246 × |
| tubería-4 | 67.042/s | 2.352 × |
| tubería-8 | 88.497/s | 3.105 × |
Para verificar la eficiencia con las conexiones de omisión, medimos el rendimiento de U-Net con varios números de dispositivos. Elegimos usar U-Net ya que tiene varias conexiones de omisión largas.
| Experimento | Rendimiento | cigratoria | Huang et al. |
|---|---|---|---|
| n = 2, M = 1 | 26.733/s | 1 × | 1 × |
| n = 2, M = 4 | 41.133/s | 1.546 × | 1.07 × |
| n = 2, M = 32 | 47.386/s | 1.780 × | 1.21 × |
| n = 4, M = 1 | 26.827/s | 1.006 × | 1.13 × |
| n = 4, M = 4 | 44.543/s | 1.680 × | 1.26 × |
| n = 4, M = 32 | 72.412/s | 2.711 × | 1.84 × |
| n = 8, M = 1 | 24.918/s | 0.932 × | 1.38 × |
| n = 8, M = 4 | 70.065/s | 2.625 × | 1.72 × |
| n = 8, M = 32 | 132.413/s | 4.966 × | 3.48 × |
( N : Número de particiones, M : número de micro-lotes)
La tabla muestra el punto de referencia de velocidad reproducido en Amoebanet-D (18, 256), como se informa en la Tabla 2 de GPIPE por Huang et al. Tenga en cuenta que reemplazamos K en el papel con n .
Este proyecto es funcional, pero la interfaz aún no está confirmada. Todas las API públicas están sujetas a cambios sin previo aviso hasta V0.1.0.
El proyecto Torchgpipe es desarrollado por Heungsub Lee, Mygryong Jeong y Chiheon Kim en Kakao Brain, con Sungbin Lim, Ildoo Kim, Woonhyuk Baek y la ayuda de Boogeon Yoon. Se distribuye bajo la licencia BSD de 3 cláusulas.
Si aplica esta biblioteca a cualquier proyecto e investigación, cite nuestro código:
@article{kim2020torchgpipe,
title={torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models},
author={Chiheon Kim and Heungsub Lee and Myungryong Jeong and Woonhyuk Baek and Boogeon Yoon and Ildoo Kim and Sungbin Lim and Sungwoong Kim},
year={2020},
eprint={2004.09910},
archivePrefix={arXiv}
}
