AI21 Labs Melatih Model Bahasa dengan 178 Miliar Parameter Menggunakan Instans P4d Amazon EC2, PyTorch

AI21 Labs menggunakan machine learning untuk mengembangkan model bahasa yang berfokus pada pemahaman makna, dan pada tahun 2021, layanan ini menetapkan tujuan untuk melatih Jurassic-1 Jumbo yang baru dirilis, yaitu model bahasa autoregresif dengan 178 miliar parameter. Developer yang mendaftar untuk pengujian beta akan mendapatkan akses ke Jurassic-1 Jumbo dan dapat segera mulai menyesuaikan model untuk kasus penggunaan mereka. Perusahaan rintisan perangkat lunak ingin melatih model secara efisien, sehingga mereka melirik Amazon Web Services (AWS), dan membangun solusi menggunakan Amazon Elastic Compute Cloud (Amazon EC2), suatu layanan web yang menyediakan kapasitas komputasi yang aman dan dapat diubah ukurannya di cloud. Memilih Amazon EC2 memberi perusahaan kendali atas proses pelatihan, termasuk alokasi simpul.

Untuk fungsionalitas komputasi dan jaringan yang kuat, perusahaan tersebut memilih Instans P4d Amazon EC2, yang memberikan throughput tinggi dan jaringan latensi rendah untuk pelatihan machine learning dan aplikasi komputasi performa tinggi di cloud. Dengan menggunakan Instans P4d Amazon EC2, AI21 Labs memperoleh performa dan memori yang diperlukan dengan mendistribusikan pelatihan model ke ratusan GPU untuk menghadirkan pemrosesan bahasa alami sebagai layanan melalui model Jumbo Jurassic-1. Karena perusahaan tersebut sekarang melatih dan mengendalikan model skala besarnya sendiri, perusahaan dapat bekerja untuk mengembangkan model baru pada skala yang sama dan berinovasi dengan lebih mudah.

Didirikan pada tahun 2017, AI21 Labs mengejar misi hibrida: melakukan penelitian pemrosesan bahasa alami dan mengembangkan produk yang didukung kecerdasan buatan (AI) untuk membaca dan menulis. Produk andalannya, Wordtune, adalah asisten penulisan dan pengeditan cerdas yang diluncurkan pada Oktober 2020, dan telah berkembang untuk mendukung hampir satu juta pengguna. Produk utama lainnya, AI21 Studio, menawarkan akses API ke model bahasa Jurassic-1 milik perusahaan serta pengembangan model kustom. “Kami adalah bagian dari kelompok kecil perusahaan yang menawarkan model bahasa sebagai layanan, memberdayakan siapa pun mulai dari developer independen hingga perusahaan multinasional untuk membangun aplikasi dan layanan di atas teknologi pemrosesan bahasa alami yang canggih,” kata Yoav Shoham, cofounder dan co-CEO di AI21 Labs. “Selain itu, kami mengejar inovasi ilmiah dan mengatasi tantangan rekayasa perangkat lunak yang ditimbulkan oleh model dengan ukuran dan kompleksitas ini.”

Untuk melatih megamodel deep learning pertamanya secara efisien dan mendukung kebutuhan penskalaan dan performa model yang tinggi, AI21 Labs membutuhkan komputasi yang kuat, kecepatan jaringan yang efisien, serta akses ke dukungan dan panduan teknis. Untuk alasan ini, pada awal tahun 2021 perusahaan mulai menerapkan solusi di AWS, memilih untuk melatih model menggunakan Instans P4d Amazon EC2. Instans ini diterapkan dalam klaster hyperscale yang disebut Amazon EC2 UltraClusters, menyediakan lebih dari 4.000 GPU NVIDIA A100, infrastruktur jaringan tanpa pemblokiran skala Petabita, dan throughput tinggi, penyimpanan latensi rendah. 

Pendekatan perusahaan dioptimalkan lebih lanjut melalui GPUDirect RDMA dengan latensi rendah dan bandwidth tinggi, bersama dengan Elastic Fabric Adapter (EFA), suatu antarmuka jaringan untuk instans Amazon EC2 yang memungkinkan pelanggan menjalankan aplikasi yang membutuhkan komunikasi antarsimpul tingkat tinggi dalam skala besar di AWS. Karena ukuran modelnya, tim perlu menggunakan pemrosesan paralel untuk mencapai waktu pelatihan yang efisien, sehingga mereka melirik kemampuan jaringan di AWS untuk mendukung pelatihan terdistribusi dan paralelisme model. “Instans P4d Amazon EC2 menawarkan jaringan berperforma tinggi 400 Gbps di EFA,” kata Opher Lieber, technical lead Jurassic di AI21 Labs. “Kecepatan jaringan GPU-ke-GPU secara langsung memengaruhi kemampuan untuk menskalakan secara efisien dan tetap hemat biaya saat melakukan penskalaan ke ratusan GPU.”

AI21 Labs dimulai dengan memunculkan basis kode pada Instans P4d Amazon EC2 yang diaktifkan untuk EFA. Kemudian menguji dan memverifikasi performa dan penskalaan yang efisien dari pendekatan pelatihan multisimpul. Selanjutnya, tim meluncurkan pelatihan cepat untuk model ukuran penuh—yang menggunakan ratusan GPU—untuk memverifikasi fungsi dan performa. Dari sana, pelatihan model Jurassic-1 Jumbo di AWS dapat dimulai. Untuk orkestrasi, perusahaan memilih solusi internal yang mengalokasikan instans menggunakan kit pengembangan perangkat lunak AWS—AWS SDK for Python (Boto3), yang memudahkan integrasi aplikasi, pustaka, atau skrip Python milik pelanggan dengan berbagai layanan AWS.

Untuk penyimpanan, AI21 Labs memilih Amazon Simple Storage Service (Amazon S3), yang menawarkan skalabilitas, ketersediaan data, keamanan, dan performa terdepan di industri. “Kami dapat mencapai performa yang sangat baik di Amazon S3 menggunakan bantuan dari tim AWS—jadi ini adalah pilihan yang mudah dari segi performa dan harga,” kata Lieber. Tim menggunakan bucket Amazon S3 untuk menyimpan dan memuat titik pemeriksaan secara efisien dan terdistribusi. Untuk mencatat kemajuan dan peristiwa pelatihan, tim menggunakan Amazon CloudWatch, layanan pemantauan dan observabilitas. 

Saat menerapkan solusi, AI21 Labs memanfaatkan dukungan dari AWS. Timnya berkonsultasi dengan spesialis AWS yang memberikan panduan tentang pertanyaan dan masalah terkait tingkat layanan, arsitektur, dan perangkat keras. Selain itu, perusahaan juga meningkatkan performa Jurassic-1 Jumbo menggunakan PyTorch on AWS, kerangka kerja deep learning sumber terbuka yang memudahkan pengembangan model machine learning dan menerapkannya ke produksi. 

AI21 Labs menyelesaikan pelatihan selama beberapa bulan, dan berakhir pada Juni 2021. Megamodel baru tersebut, sebuah model bahasa autoregresif dengan 178 miliar parameter, yang sebanding dengan penawaran perusahaan pesaing. Ini juga menawarkan kosakata 256.000 item yang berbeda yang menyediakan kemampuan representasi teks yang diperluas serta dukungan untuk entitas bernama. Sekarang perusahaan tersebut menawarkan Jurassic-1 Jumbo (bersama dengan pendampingnya, Jurassic-1 Large, yang memiliki 7 miliar parameter) dalam versi beta terbuka melalui penawaran AI21 Studio perusahaan. Dengan menggunakan layanan ini, berbagai developer bisa membangun produk pada model Jurassic-1 Jumbo, dan AI21 Labs telah melihat penggunaannya di banyak industri, termasuk pemasaran, pembuatan konten, gim, penelitian medis, otomotif, telekomunikasi, dan keuangan.

Karena AI21 Labs memiliki akses langsung ke modelnya, AI21 Labs dapat beradaptasi dan berinovasi tanpa bergantung pada pihak ketiga dan dapat mengeksplorasi tujuan inovasi berkelanjutan, yang merupakan bagian penting dari misinya. AI21 Labs saat ini membuat prototipe model tambahan, yang rencananya juga akan dilatih dalam skala besar. “Melatih dan memiliki megamodel kami sendiri akan terus menjadi faktor pembeda yang penting dalam penawaran Wordtune dan AI21 Studio kami,” kata Shoham.