torchgpipe

Torchgpipe

Uma implementação GPIPE em Pytorch. É otimizado para CUDA em vez 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 )

O que é Gpipe?

O GPIPE é uma biblioteca de paralelismo de pipeline escalável publicada pelo Google Brain, que permite treinamento eficiente de grandes modelos de consumo de memória. De acordo com o artigo, o GPIPE pode treinar um modelo maior de 25x usando dispositivos 8x (TPU) e treinar um modelo 3.5X mais rápido usando dispositivos 4x.

GPIPE: Treinamento eficiente de redes neurais gigantes usando o paralelismo do pipeline

Google treinado AMOEBANET-B com parâmetros de 557m sobre o GPIPE. Esse modelo alcançou 84,3% de top-1 e 97,0% de precisão Top-5 na referência de classificação do ImageNet (o desempenho de última geração em maio de 2019).

O GPIPE usa (a) paralelismo do pipeline e (b) recomputação automática da propagação direta durante a retropropagação, aproveita o treinamento de um grande modelo. Nós nos referimos a (b) como conquistamento de verificação, seguindo a conhecida terminologia na comunidade de Pytorch.

Paralelismo do pipeline

O GPIPE divide um modelo em várias partições e coloca cada partição em um dispositivo diferente para ocupar mais capacidade de memória. E divide um mini-lote em vários micro-lotes para fazer com que as partições funcionem o mais paralelo possível.

Conquistamento de checking

A oposição de verificação é aplicada a cada partição para minimizar o consumo geral de memória por um modelo. Durante a propagação direta, apenas os tensores nos limites entre as partições são lembrados. Todos os outros tensores intermediários são volatilizados e recomputados durante a retropropagação quando necessário.

Uso

Atualmente, o Torchgpipe requer os seguintes ambientes:

• Python 3.6+
• Pytorch 1.1+

Para usar o Torchgpipe, instale -o via Pypi:

$ pip install torchgpipe

Para treinar um módulo com gpipe, basta embrulhá -lo com torchgpipe.GPipe . Seu módulo deve ser nn.Sequential pois o GPIPE dividirá automaticamente o módulo em partições com camadas consecutivas. O argumento balance determina o número de camadas em cada partição. O argumento chunks especifica o número de micro-lotes. Os tensores de entrada, saída e intermediários devem ser Tensor ou Tuple[Tensor, ...] .

O código de exemplo abaixo mostra como dividir um módulo com quatro camadas em quatro partições, cada uma com uma única camada. Este código também divide um mini-lote em 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 )

Documentação

Visite Torchgpipe.readthedocs.io para obter mais informações, incluindo as referências da API.

Benchmarking

Os detalhes completos e mais referências estão disponíveis no Torchgpipe.readThedocs.io.

Referência de precisão Resnet-101

Tamanho do lote	Torchgpipe	nn.dataparallelal	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 / D
4K	22.13 ± 0,09	N / D	N / D

O GPIPE deve ser transparente para não introduzir ajuste adicional de hiperparâmetro. Para verificar a transparência, reproduzimos a taxa de erro TOP-1 de resnet-101 no ImageNet, conforme relatado na Tabela 2 (c) de SGD preciso e de grande minibatch por Goyal et al.

Referência de memória U-Net (B, C)

Experimentar	U-net (b, c)	Parâmetros	Uso de memória
linha de base	(6, 72)	362.2m	20.3 Gib
Pipeline-1	(11, 128)	2.21b	20.5 Gib
Pipeline-2	(24, 128)	4.99b	43.4 Gib
Pipeline-4	(24, 160)	7.80b	79.1 Gib
Pipeline-8	(48, 160)	15.82b	154.1 Gib

A tabela mostra como o GPIPE facilita os modelos de escala de U-Net. A linha de base indica a linha de base sem o paralelismo da tubulação nem o check -se, e o pipeline -1 , -2 , -4 , -8 indica que o modelo é treinado com GPIPE com o número correspondente de partições.

Aqui usamos uma arquitetura de rede U simplificada. O tamanho de um modelo é determinado pelos hiperparâmetros B e C, que são proporcionais ao número de camadas e filtros, respectivamente.

U-NET U (5, 64) referência de velocidade

Experimentar	Taxa de transferência	Acelerar
linha de base	28.500/s	1 ×
Pipeline-1	24.456/s	0,858 ×
Pipeline-2	35.502/s	1,246 ×
Pipeline-4	67.042/s	2.352 ×
Pipeline-8	88.497/s	3.105 ×

Para verificar a eficiência nas conexões SKIP, medimos a taxa de transferência da U-Net com vários números de dispositivos. Optamos por usar a U-Net, pois possui várias conexões longas de pular.

AMOEBANET-D (18, 256) Benchmark de velocidade

Experimentar	Taxa de transferência	Torchgpipe	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 partições, m : número de micro-lotes)

A tabela mostra a referência de velocidade reproduzida em Amoebanet-D (18, 256), conforme relatado na Tabela 2 de Gpipe por Huang et al. Observe que substituímos K no papel por n .

Notas

Este projeto é funcional, mas a interface ainda não está confirmada. Todas as APIs públicas estão sujeitas a alterações sem aviso prévio até v0.1.0.

Autores e licenciamento

O Projeto Torchgpipe é desenvolvido por Heungsub Lee, Myungryong Jeong e Chiheon Kim em Kakao Brain, com Sungbin Lim, Ildoo Kim, Woonhyuk Baek e a ajuda de Boogeon Yoon. É distribuído sob a licença BSD de 3 cláusulas.

Citação

Se você aplicar esta biblioteca a qualquer projeto e pesquisa, cite nosso 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}
}