torchgpipe

Torchgpipe

Implementasi GPIPE di Pytorch. Ini dioptimalkan untuk CUDA daripada 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 )

APA ITU GPIPE?

GPIPE adalah perpustakaan paralelisme pipa yang dapat diskalakan yang diterbitkan oleh Google Brain, yang memungkinkan pelatihan yang efisien dari model besar yang memakan memori. Menurut kertas, GPIPE dapat melatih model 25x yang lebih besar dengan menggunakan perangkat 8x (TPU), dan melatih model 3.5x lebih cepat dengan menggunakan perangkat 4X.

GPIPE: Pelatihan yang efisien dari jaringan saraf raksasa menggunakan paralelisme pipa

Google Amoebanet-B terlatih dengan parameter 557m di atas GPIPE. Model ini telah mencapai 84,3% top-1 dan 97,0% akurasi top-5 pada tolok ukur klasifikasi ImageNet (kinerja canggih pada Mei 2019).

Penggunaan GPIP (a) Paralelisme pipa dan (b) komputasi otomatis propagasi ke depan selama backpropagation, karenanya memanfaatkan pelatihan model besar. Kami merujuk pada (b) sebagai pos pemeriksaan, mengikuti terminologi terkenal di komunitas Pytorch.

Paralelisme pipa

GPIPE membagi model menjadi beberapa partisi dan menempatkan setiap partisi pada perangkat yang berbeda untuk menempati lebih banyak kapasitas memori. Dan itu membagi mini-batch menjadi beberapa mikro-batch untuk membuat partisi bekerja sejajar mungkin.

Pos pemeriksaan

Pos Pemeriksaan diterapkan pada setiap partisi untuk meminimalkan konsumsi memori secara keseluruhan dengan model. Selama perambatan ke depan, hanya tensor di batas -batas antara partisi yang diingat. Semua tensor perantara lainnya diuapkan, dan dihitung ulang selama propagasi back bila perlu.

Penggunaan

Saat ini, Torchgpipe membutuhkan lingkungan berikut:

• Python 3.6+
• Pytorch 1.1+

Untuk menggunakan torchgpipe, instal melalui PYPI:

$ pip install torchgpipe

Untuk melatih modul dengan gpipe, cukup bungkus dengan torchgpipe.GPipe . Modul Anda harus nn.Sequential karena gpipe akan secara otomatis membagi modul menjadi partisi dengan lapisan berturut -turut. Argumen balance menentukan jumlah lapisan di setiap partisi. Argumen chunks menentukan jumlah mikro-batch. Input, output, dan tensor menengah harus Tensor atau Tuple[Tensor, ...] .

Kode contoh di bawah ini menunjukkan cara membagi modul dengan empat lapisan menjadi empat partisi masing -masing memiliki satu lapisan. Kode ini juga membagi mini-batch menjadi 8 mikro-batch:

 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 )

Dokumentasi

Kunjungi torchgpipe.readthedocs.io untuk informasi lebih lanjut termasuk referensi API.

Benchmarking

Detail lengkap dan lebih banyak tolok ukur tersedia di Torchgpipe.readthedocs.io.

Benchmark Akurasi ResNet-101

Ukuran batch	Torchgpipe	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 harus transparan untuk tidak memperkenalkan penyetelan hiperparameter tambahan. Untuk memverifikasi transparansi, kami mereproduksi laju kesalahan Top-1 dari ResNet-101 di ImageNet, seperti yang dilaporkan pada Tabel 2 (c) dari SGD minibatch besar yang akurat oleh Goyal et al.

Tolok ukur u-net (b, c)

Percobaan	U-net (b, c)	Parameter	Penggunaan memori
Baseline	(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

Tabel ini menunjukkan bagaimana GPIPE memfasilitasi penskalaan model U-Net. Baseline menunjukkan garis dasar tanpa paralelisme pipa atau pos pemeriksaan, dan pipa -1 , -2 , -4 , -8 menunjukkan bahwa model dilatih dengan gpipe dengan jumlah partisi yang sesuai.

Di sini kami menggunakan arsitektur U-Net yang disederhanakan. Ukuran model ditentukan oleh hiperparameter B dan C yang sebanding dengan jumlah lapisan dan filter, masing -masing.

Tolok ukur kecepatan U-Net (5, 64)

Percobaan	Throughput	Mempercepat
Baseline	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 ×

Untuk memverifikasi efisiensi dengan koneksi lewati, kami mengukur throughput U-NET dengan berbagai jumlah perangkat. Kami memilih untuk menggunakan u-net karena memiliki beberapa koneksi lompatan panjang.

Benchmark Amoebanet-D (18, 256)

Percobaan	Throughput	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 : Jumlah partisi, m : jumlah mikro-batch)

Tabel ini menunjukkan tolok ukur kecepatan yang direproduksi pada Amoebanet-D (18, 256), seperti yang dilaporkan pada Tabel 2 GPIPE oleh Huang et al. Perhatikan bahwa kami mengganti k di koran dengan n .

Catatan

Proyek ini fungsional, tetapi antarmuka belum dikonfirmasi. Semua API publik dapat berubah tanpa peringatan sampai v0.1.0.

Penulis dan Lisensi

Proyek Torchgpipe dikembangkan oleh Heungsub Lee, Myungryong Jeong, dan Chiheon Kim di Kakao Brain, dengan Sungbin Lim, Ildoo Kim, Woonhyuk Baek, dan bantuan Boogeon Yoon. Ini didistribusikan di bawah lisensi BSD 3-Clause.

Kutipan

Jika Anda menerapkan perpustakaan ini ke proyek dan riset apa pun, silakan kutip kode kami:

 @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}
}