Bagaimana Amazon Search M5 menghemat 30% biaya pelatihan LLM dengan menggunakan AWS Trainium | Layanan Web Amazon

Selama beberapa dekade, Amazon telah memelopori dan berinovasi dalam pembelajaran mesin (ML), menghadirkan pengalaman menyenangkan bagi pelanggannya. Sejak awal, Amazon telah menggunakan ML untuk berbagai kasus penggunaan seperti rekomendasi buku, pencarian, dan deteksi penipuan. Mirip dengan industri lainnya, kemajuan perangkat keras yang dipercepat telah memungkinkan tim Amazon untuk mengejar arsitektur model menggunakan jaringan saraf dan pembelajaran mendalam (DL).

Program M5 dalam Amazon Search memiliki strategi pembelajaran penemuan untuk Amazon dan membangun model berskala besar di berbagai bahasa, multi-lokal, multi-entitas, multitugas, dan multi-modal seperti teks, gambar, dan video. Program M5 telah menyajikan penyematan universal dan model fondasi skala besar ke ratusan tim ML di Amazon sambil mempertahankan kontrol ketat terhadap optimalisasi biaya. Untuk mencapai hal ini, tim M5 secara berkala mengevaluasi teknik baru untuk mengurangi biaya.

Seperti banyak organisasi ML, akselerator sebagian besar digunakan untuk mempercepat pelatihan dan inferensi DL. Saat AWS meluncurkan akselerator yang dibuat khusus dengan rilis pertama Inferensi AWS pada tahun 2020, tim M5 dengan cepat memulainya memanfaatkannya untuk menerapkan beban kerja produksi dengan lebih efisien, menghemat biaya dan mengurangi latensi. Tahun lalu, AWS meluncurkannya Pelatihan AWS akselerator, yang mengoptimalkan kinerja per biaya untuk mengembangkan dan membangun model DL generasi berikutnya. Dalam postingan ini, kami membahas bagaimana M5 mampu mengurangi biaya untuk melatih model mereka sebesar 30%, dan berbagi beberapa praktik terbaik yang kami pelajari selama ini.

Contoh pelatihan

Dengan kemajuan dalam akselerator yang dibuat khusus, Amazon juga menyediakan akselerator menarik dalam bentuk AWS Inferentia dan Trainium. Sesuai dengan namanya, chip ini masing-masing dioptimalkan untuk melebihi kebutuhan beban kerja inferensi dan pelatihan. Untuk pelatihan model fondasi skala besar yang ukurannya mencapai miliaran parameter, Trainium Instans Trn1 dan Trn1n adalah pilihan ideal karena karakteristiknya. Instans Trn1 didukung oleh teknologi canggih NeuronCore-v2, dan memiliki banyak sekali komputasi akselerator dan memori. Instans Trn1n juga dapat dipilih untuk jumlah bandwidth jaringan yang lebih besar (1,600 Gbs), sehingga cocok untuk pelatihan berkinerja tinggi dengan mempertimbangkan optimalisasi biaya.

Untuk menggunakan akselerator, Anda memerlukan lapisan perangkat lunak untuk mendukungnya. Dengan chip Trn dan Inf, SDK Neuron AWS membuka akselerator yang dibuat khusus oleh Amazon dengan bantuan PyTorch XLA. PyTorch XLA mengubah mode bersemangat PyTorch menjadi implementasi berbasis grafik mode malas. Grafik ini kemudian digunakan dan dikompilasi lebih lanjut untuk digunakan dengan akselerator. PyTorch Neuron (bagian dari Neuron SDK) memungkinkan pengguna PyTorch untuk melatih model mereka di Trainium NeuronCores dengan beberapa baris kode.

Model dan beban kerja

Tim M5 melatih dan menerapkan model dasar dan representasi universal untuk membantu berbagai tim di Amazon dalam menghadirkan kesenangan Amazon.com pelanggan. Salah satu model tersebut adalah model encoder teks yang diikuti oleh multi-layer perceptron (MLP) dengan interaksi fitur eksplisit atau implisit yang ditentukan oleh arsitektur jaringan saraf dengan ratusan juta parameter yang dapat dilatih. Model ini dilatih pada miliaran token, dan digunakan untuk menghasilkan jutaan penyematan dalam pengaturan inferensi batch offline. Penyematan ini merupakan masukan ke layanan Amazon tingkat 1 yang berhubungan dengan pelanggan.

Infrastruktur untuk penggunaan pipa produksi Batch AWS dengan strategi antrian pembagian yang adil, menggunakan cluster multi-node trn1.32xlarge berkemampuan EFA sebagai komputasi untuk pelatihan model. Secara fungsional, alur produksi melakukan pelatihan model tambahan, evaluasi model terlatih, dan inferensi batch offline pada model terlatih, semuanya menggunakan PyTorch sebagai pustaka DL yang mendasarinya.

Anda

Memuaskan pelanggan kami adalah prinsip utama. Mengingat sifat saluran yang berhubungan dengan pelanggan, semua perjanjian tingkat layanan (SLA) harus dipenuhi tanpa regresi. Kami mengidentifikasi dua kriteria penerimaan penting untuk mengadaptasi jalur produksi GPU kami yang ada dan mentransisikannya ke Trainium:

• Kualitas model – Kualitas model kami berdampak langsung pada pengalaman pelanggan. Kami mensyaratkan perbedaan kualitas model antara GPU dan Trainium kurang dari 0.1%.
• Throughput pelatihan – Kami berulang kali melatih model kami secara berkala untuk memberikan pengalaman terbaru kepada pelanggan kami. Kami mensyaratkan bahwa konvergensi model harus dicapai dalam jangka waktu yang telah ditentukan (misalnya 1 minggu) untuk memenuhi SLA produksi kami.

Di bagian berikut, kami berbagi perjalanan kami dalam bekerja mundur dari kriteria ini, dan pembelajaran kami untuk mendukung beban kerja produksi skala Amazon.

Naskah pelatihan

Sebelum memulai pelatihan model, kita perlu melakukan perubahan pada skrip pelatihan agar sesuai dengan XLA. Mengingat ukuran modelnya, kami menggunakan data terdistribusi paralel (DDP) untuk melatih model. DDP memungkinkan kami meningkatkan throughput pelatihan model dengan meningkatkan jumlah mesin yang digunakan untuk menjalankan pelatihan model, tanpa perubahan kode apa pun. Kami mengikuti instruksi yang diberikan di Tutorial pelatihan MLP Neuron PyTorch untuk menambahkan konstruksi khusus XLA dalam skrip pelatihan kami. Perubahan kode ini mudah diterapkan. Berikut adalah beberapa pembelajaran teknis yang signifikan dari latihan ini yang sangat meningkatkan throughput model kami:

• Penempatan xm.mark_step() - xm.mark_step() mengkompilasi dan menjalankan grafik komputasi yang dikumpulkan dengan malas. Memanggil mark_step terlalu sering akan menghasilkan lebih banyak grafik kecil, sedangkan penerapannya terlalu sedikit akan menghasilkan lebih sedikit grafik tetapi besar. Bergantung pada aplikasi Anda, throughput dan implementasi pelatihan model Anda akan bervariasi berdasarkan penempatan Anda xm.mark_step(). Implementasi kami menempatkan satu xm.mark_step() setelah langkah maju dan mundur, dan satu lagi setelah langkah pengoptimal.
• Pembungkus pemuat data dengan pemuat perangkat multiprosesor XLA – Ini adalah langkah penting yang dapat dengan mudah dilewatkan. Pemuat perangkat multiprosesor torch_xla.distributed.parallel_loader.MpDeviceLoader memuat data pelatihan pada setiap perangkat XLA dengan opsi untuk melakukan pramuat dan tumpang tindih pemuatan data dengan pengoperasian perangkat untuk meningkatkan throughput. Pemuat perangkat juga memanggil xm.mark_step() dan karena itu mampu membuat grafik untuk memuat data ke perangkat dari host.

Kompilasi untuk Pelatihan

Secara tradisional, siklus pengembangan model dengan GPU melibatkan perubahan pada model atau skrip pelatihan dan menjalankannya secara langsung di perangkat GPU. Akselerator seperti Trainium yang menggunakan XLA memerlukan langkah tambahan sebelum pelatihan model dapat dijalankan di akselerator. Grafik komputasi XLA hanya dapat dijalankan setelah dikompilasi. Secara umum, ada dua cara untuk melakukan kompilasi ini: Ahead of Time (AOT), yang mana Anda melacak dan mengkompilasi semua grafik terlebih dahulu lalu menjalankannya, atau Just In Time (JIT), yang mana grafik ditelusuri, dikompilasi, dan dijalankan sebagaimana adanya. ditemui. Neuron SDK menyediakan kedua hal ini secara langsung. Biasanya, kompilasi AOT dilakukan terlebih dahulu. Grafik kemudian dijalankan setelah kompilasi ini. Jika grafik baru ditemukan, runtime Neuron akan memanggil kompilasi JIT sebelum menjalankannya. Untuk melakukan kompilasi AOT, Neuron SDK menyediakan neuron_parallel_compile, utilitas kompilasi yang mengekstrak grafik dari uji coba skrip pelatihan dan melakukan kompilasi AOT paralel.

Aspek penting dari kompilasi AOT adalah memastikan bahwa tidak ada grafik komputasi baru yang dibuat selama pelatihan. Salah satu sumber grafik komputasi baru (dan karenanya kompilasi ulang) adalah bentuk dinamis dari kumpulan pelatihan selama pelatihan model. Kami menemukan bahwa penggunaan bentuk statis dan kumpulan ukuran tetap menghilangkan kompilasi waktu pelatihan dan sangat meningkatkan hasil pelatihan tanpa memengaruhi akurasi model. Dengan menerapkan batasan tersebut pada pelatihan, kami mengamati bahwa hanya diperlukan 4–5 langkah pelatihan model, satu langkah validasi model, dan pemeriksaan model satu kali untuk menelusuri semua grafik selama kompilasi AOT. Penting untuk dicatat bahwa Neuron SDK terus berkembang, dan di masa depan akan mendukung bentuk dinamis juga.

Selanjutnya grafik yang telah disusun disimpan dalam Cache Persisten Neuron pada disk atau dalam an Layanan Penyimpanan Sederhana Amazon (Amazon S3) ember. Hal ini sangat berguna untuk beban kerja produksi di mana arsitektur model dan konfigurasi pelatihan tidak berubah. Oleh karena itu, biaya kompilasi hanya dikeluarkan satu kali. Menggunakan cache semudah mengatur tanda lingkungan:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Kompiler Neuron juga menyediakan tiga opsi pengoptimalan tingkat kompiler (O1, O2, O3) untuk menyeimbangkan waktu kompilasi dan throughput model yang dijalankan. O1 memungkinkan pengoptimalan inti pada grafik komputasi dan meminimalkan waktu kompilasi, O3 memberikan throughput eksekusi model yang lebih baik dengan mengorbankan waktu kompilasi yang lebih tinggi, dan O2 (opsi default) adalah keseimbangan di antara keduanya. Untuk kasus penggunaan kami, kami menggunakan pengoptimalan O1 dan mengamati pengurangan waktu kompilasi sebesar 86% tanpa perubahan pada metrik akurasi model, sambil mengamati pengurangan throughput sekitar 5–7% dibandingkan dengan pengoptimalan default (O2). Bergantung pada kasus penggunaannya, Anda dapat memilih tingkat pengoptimalan yang berbeda.

Ringkasnya, kami menggunakan flag berikut untuk kompilasi:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Kompatibilitas pos pemeriksaan

Ketika kompilasi berhasil diselesaikan, kita dapat melanjutkan untuk melatih model kita di Trainium. Seperti disebutkan sebelumnya, kami melatih model kami secara bertahap, artinya kami memuat titik pemeriksaan model yang telah dilatih sebelumnya dan melanjutkan pelatihan dengan data baru. PyTorch dan PyTorch XLA memungkinkan transisi mulus antar akselerator melalui interoperabilitas pos pemeriksaan. Memiliki fleksibilitas perpindahan antara GPU dan Trainium memungkinkan kami memuat model GPU sebelumnya dengan lancar dan berlatih di mesin Trainium. Hal ini penting untuk memastikan bahwa kami dapat menginisialisasi model kami dengan model terbaik yang telah dilatih sebelumnya tanpa adanya waktu henti produksi atau hilangnya akurasi model.

Karena model GPU disimpan menggunakan utilitas penyimpanan model PyTorch standar, kami dapat menggunakan utilitas pemuatan pos pemeriksaan PyTorch untuk memuat model GPU pada perangkat Trainium.

Misalnya pada GPU/CPU, Anda dapat menyimpan model dengan kode berikut:

torch.save(model.state_dict(), PATH)

Kemudian Anda memuat kembali model tersebut di Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

Demikian pula, Anda dapat menyimpan model di Trainium dengan kode berikut:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

Dan memuat kembali model pada GPU/CPU:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

Faktanya, karena kami menggunakan DDP untuk pelatihan model, pemuatan model tidak bergantung pada jumlah mesin yang digunakan untuk melatih pos pemeriksaan sebelumnya. Hal ini memungkinkan kami untuk menskalakan armada Trn1 secara horizontal tanpa perubahan kode atau efek buruk pada pelatihan model. Pos pemeriksaan berbasis PyTorch ini dapat langsung digunakan atau bahkan dibuat skrip obor untuk kasus penggunaan inferensi di AWS Inferentia2 atau akselerator lainnya.

Stabilitas operasional

Tidak dapat cukup ditekankan bahwa menjalankan beban kerja dalam produksi memerlukan beberapa SLA yang harus dipenuhi. Untuk kasus penggunaan kami, selain kualitas model dan SLA throughput pelatihan, pipeline produksi harus stabil secara operasional, yang berarti waktu henti dan gangguan minimal selama pelatihan model, evaluasi, dan inferensi.

Seperti halnya pipeline berbasis GPU yang sudah ada, kami menambahkan banyak mekanisme untuk membuat operasional pipeline stabil. Sebelum memulai pelatihan model, kami menjalankan beberapa uji kewarasan untuk menilai kesehatan mesin. Pengujian ini umumnya mencakup operasi tensor sederhana untuk memverifikasi kesehatan perangkat akselerator. Kami telah mengamati bahwa untuk pelatihan terdistribusi, penting untuk menjalankan pengujian guna memverifikasi komunikasi kolektif antar instance. Kami menggunakan Rangkaian tes NCCOM dari Neuron SDK untuk mencapai hal ini, menjalankan berbagai operasi seperti all-gather, all-reduce, dan less-scatter.

Bahkan setelah mengikuti saran yang kami sebutkan, kami mengamati bahwa masalah sementara tidak dapat dihindari dalam saluran apa pun, apa pun akselerator yang mendasarinya. Untuk membangun ketahanan dalam alur pelatihan apa pun, kami merekomendasikan untuk membangun mekanisme percobaan ulang untuk menyelesaikan potensi masalah ini. Kita gunakan Percobaan ulang otomatis AWS Batch untuk mencoba kembali pekerjaan yang mengalami kegagalan sementara selama pelatihan model. Pengulangan ulang ini dapat memakan biaya besar jika terjadi kegagalan menjelang akhir pelatihan. Untuk mengatasi masalah ini, kami telah mengadaptasi skrip pelatihan kami untuk memuat titik pemeriksaan model yang telah dilatih sebelumnya dan melanjutkan pelatihan dari titik tersebut. Dengan fungsi ini, kami dapat memulai kembali tugas pelatihan yang gagal secara agresif dengan overhead minimal.

Dengan adanya mekanisme ketahanan ini, kami mampu mencapai tingkat keberhasilan 98.5% untuk beban kerja kami di Trn1, sebanding dengan tingkat keberhasilan pipeline GPU kami saat ini.

Hasil

Untuk memvalidasi keakuratan model kami, kami menginisialisasi dua model dari titik pemeriksaan GPU yang sama, dan melatih satu model di Trainium dan model lainnya pada GPU yang sebanding. Kedua model dilatih dengan hyperparameter pelatihan yang sama. Kumpulan data yang digunakan untuk penghitungan metrik adalah kumpulan data ketidaksepakatan, dan kami mengevaluasi keakuratan model pada kumpulan data ini setiap N langkah global. Sumbu X adalah langkah global, dan sumbu Y adalah akurasi model. Kami mengamati perbedaan kurang dari 0.1% dalam akurasi model di setiap titik pada grafik berikut.

Selain itu, untuk mengevaluasi efektivitas biaya pelatihan model, kami lebih memilih untuk membandingkan waktu jam dinding yang dibutuhkan untuk mencapai konvergensi model. Kami percaya ini memberikan pandangan yang lebih praktis tentang penghematan biaya dibandingkan dengan ukuran seperti biaya per token, pencapaian FLOPS/dolar, dan faktor lainnya. Mengingat waktu pelatihan trn1.32xl dan sebanding Cloud komputasi elastis Amazon (Amazon EC2), kami mengamati bahwa Trainium menawarkan biaya hingga 30% lebih murah untuk memodelkan konvergensi.

Kesimpulan

Ada banyak faktor yang perlu dipertimbangkan saat mengevaluasi berbagai akselerator untuk beban kerja DL Anda. Beberapa yang paling penting adalah kualitas model, keluaran, biaya, dan ketersediaan. Sangat penting untuk memastikan bahwa kualitas dan throughput model Anda tidak dikorbankan berdasarkan akselerator yang Anda pilih.

Berkat kemitraan dan kolaborasi kami dengan tim Annapurna Neuron, tim Amazon Search M5 telah mampu menghemat biaya hingga 30% dengan berpindah ke Trainium. Tim ini mampu menggunakan Trainium dan mencapai kualitas model serta keseimbangan keluaran dengan akselerator serupa di pasar. Interoperabilitas pos pemeriksaan dan perubahan kode minimal dengan dukungan XLA memungkinkan M5 memilih di antara beberapa akselerator untuk beban kerjanya. Hal ini memungkinkan tim M5 memanfaatkan kekuatan komputasi Trainium yang besar, dan membangun solusi akselerator agnostik untuk memuaskan pelanggan Amazon.com. Dari segi operasional, Trainium telah terbukti mampu mendukung layanan tier-1 pada skala Amazon. Tim M5 terus memindahkan lebih banyak beban kerja ke Trainium untuk menyediakan model terbaik bagi Amazon dengan biaya terendah.

Singkatnya, tim M5 telah mampu melakukan pelatihan ML tingkat produksi yang hemat biaya dengan menambahkan Trainium ke armada akselerator. Kami mendorong Anda untuk melihat Trainium dan perangkat Neuron lainnya seperti AWS Inferentia untuk mendapatkan manfaat silikon Amazon yang dibuat khusus untuk beban kerja ML. Mulailah dengan mudah menggunakan salah satu dari banyak tutorial yang menampilkan berbagai model, seperti Llama 2, tersedia di Trainium.

Tentang Penulis

Abhinandan Patni adalah Insinyur Perangkat Lunak Senior di Amazon Search. Dia berfokus pada membangun sistem dan alat untuk pelatihan pembelajaran mendalam terdistribusi yang skalabel dan inferensi waktu nyata.

James Taman adalah Arsitek Solusi di Amazon Web Services. Dia bekerja dengan Amazon.com untuk merancang, membangun, dan menerapkan solusi teknologi di AWS, dan memiliki minat khusus pada AI dan pembelajaran mesin. Di waktu luangnya, dia senang mencari budaya baru, pengalaman baru, dan mengikuti perkembangan tren teknologi terkini. Anda dapat menemukannya di LinkedIn.

Jerry Mannil adalah seorang insinyur perangkat lunak di Amazon Search. Dia berupaya meningkatkan efisiensi, ketahanan, dan skalabilitas infrastruktur pelatihan terdistribusi.

Ken Su adalah seorang insinyur perangkat lunak di Amazon Search. Dia berupaya meningkatkan efisiensi pelatihan dan alur kerja pelatihan terdistribusi yang terukur. Di luar pekerjaan, dia suka hiking dan tenis.

RJ adalah seorang Insinyur di Amazon. Dia membangun dan mengoptimalkan sistem untuk sistem terdistribusi untuk pelatihan dan berupaya mengoptimalkan sistem adopsi guna mengurangi latensi untuk Inferensi ML. Di luar pekerjaannya, dia menjajaki penggunaan AI Generatif untuk membuat resep makanan.

• Konten Bertenaga SEO & Distribusi PR. Dapatkan Amplifikasi Hari Ini.
• PlatoData.Jaringan Vertikal Generatif Ai. Berdayakan Diri Anda. Akses Di Sini.
• PlatoAiStream. Intelijen Web3. Pengetahuan Diperkuat. Akses Di Sini.
• PlatoESG. Karbon, teknologi bersih, energi, Lingkungan Hidup, Tenaga surya, Penanganan limbah. Akses Di Sini.
• PlatoHealth. Kecerdasan Uji Coba Biotek dan Klinis. Akses Di Sini.
• Sumber: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/