Llama dari Nol, Cara Mengimplementasikan Paper Tanpa Menangis

 (blog.briankitano.com)
2 poin oleh GN⁺ 2023-08-10 | 1 komentar | Bagikan ke WhatsApp
  • Brian Kitano membuat sendiri versi mini Llama dengan TinyShakespeare, dan menyimpulkan bahwa implementasi paper paling aman dilakukan dengan mulai dari model kecil, lalu mengganti komponen satu per satu sambil melatih dan mengevaluasi setiap kali ada perubahan
  • Ia lebih dulu menyiapkan fungsi bantu validasi seperti pembagian data, pembuatan batch, evaluasi loss, dan fungsi generasi, lalu memastikan model sederhana bisa dikompilasi dan dilatih sebelum menambahkan komponen Llama
  • RMSNorm, RoPE, dan SwiGLU ditambahkan secara berurutan, sambil memeriksa apakah tiap layer bekerja sesuai harapan lewat shape tensor, sifat rumus, dan attention map
  • Saat causal mask dihapus dari attention berbasis RoPE, validation loss turun hingga 0.16, tetapi kualitas generasi memburuk karena terjadi kebocoran informasi akibat model melihat token masa depan
  • Versi mini Llama final terdiri dari 4 blok dengan sekitar 2,37 juta parameter, menurunkan validation loss hingga sekitar 1.0, dan menunjukkan bahwa aliran gradient serta learning rate schedule juga perlu diperiksa

Mulai dari kecil dan bangun keyakinan secara iteratif

  • Inti implementasi paper adalah memulai dari model kecil, mengganti komponen satu per satu, lalu mengulang pelatihan dan evaluasi setiap kali ada perubahan
  • Pertama, ia menyiapkan fungsi-fungsi bantu untuk memeriksa model secara kuantitatif
    • pembagian data
    • loop pelatihan
    • visualisasi loss
    • evaluasi validation loss
  • Alih-alih langsung memindahkan semua komponen paper sekaligus, ia juga menyiapkan fungsi evaluasi kualitatif untuk melihat hasil generasi dari model sederhana dan cepat yang sudah pernah ia implementasikan
  • Layer tensor diperiksa dengan .shape, assert, dan plt.imshow; alih-alih langsung mengoptimalkan perkalian matriks sejak awal, ia lebih dulu mengecek hasil yang diharapkan secara manual lalu mengefisienkannya dengan fungsi torch
  • Pengujian perlu dilakukan dengan mengubah ukuran batch, panjang sekuens, dan dimensi embedding; kode yang hanya benar untuk satu ukuran bisa rusak saat inferensi

Dataset dan pengaturan dasar

  • Target implementasi adalah versi sangat diperkecil dari Llama milik Meta AI, dan data latih yang dipakai adalah TinyShakespeare
  • Llama dilatih dengan 1.4T token, tetapi di sini digunakan TinyShakespeare dengan skala sekitar 1,11 juta karakter
  • Llama asli menggunakan tokenizer byte-pair encoding SentencePiece, tetapi implementasi ini memakai tokenizer tingkat karakter yang sederhana
    • ukuran vocabulary adalah 65
    • karena dataset kecil, cara penyimpanan di memori tidak dioptimalkan secara khusus
  • Dictionary MASTER_CONFIG dipakai untuk mengelola pengaturan model seperti vocab_size, batch_size, context_window, dan d_model
    • tujuannya untuk mengurangi konstanta dan magic number serta membuat kode lebih mudah dibaca
  • Fungsi get_batches membagi data menjadi train 80%, val 10%, test 10%, lalu membuat input x dan label y yang bergeser satu karakter dari titik awal acak

Memastikan kompilasi dan pelatihan dengan model dasar

  • Model pertama adalah SimpleBrokenModel, yang terdiri dari embedding dan jaringan feed-forward sederhana
    • nn.Embedding
    • Linear
    • ReLU
    • Linear
  • Dalam implementasi paper, mengatakan model “berfungsi” berarti harus memenuhi dua syarat sekaligus
    • kompilasi: shape tensor cocok antar-layer
    • pelatihan: loss benar-benar turun
  • Fungsi evaluate_loss mengambil sampel 10 batch dari split train dan val untuk menghitung rata-rata loss
  • Setelah dilatih selama 1000 epoch, SimpleBrokenModel memiliki validation loss sekitar 3.94, hampir tidak turun dari cross-entropy awal 4.17
  • Penyebabnya adalah nilai yang sudah melalui softmax dimasukkan ke F.cross_entropy
    • F.cross_entropy milik PyTorch menerima logits yang belum dinormalisasi secara langsung
    • SimpleModel yang menghapus softmax berhasil menurunkan validation loss hingga sekitar 2.51
  • Setelah itu ditambahkan fungsi generate untuk memeriksa langsung karakter yang dibuat model, dan model dasar pun setidaknya berada dalam kondisi loss yang menurun meski belum sempurna

Komponen Llama 1: RMSNorm

  • Dibanding Transformer asli, Llama memakai tiga perubahan arsitektur utama
    • pre-normalization RMSNorm
    • Rotary embeddings
    • fungsi aktivasi SwiGLU
  • Transformer asli menggunakan BatchNormalization, sedangkan Llama memakai RMSNorm, yang menskalakan dengan variance tanpa melakukan centering pada vektor
  • Jika Transformer asli menerapkan normalisasi pada keluaran layer attention sebagai post-normalization, Llama justru menerapkannya lebih dulu pada input sebagai pre-normalization
  • RMSNorm yang diimplementasikan mengasumsikan shape input (batch, seq_len, d_model)
  • Hasil RMSNorm diuji lewat sifat bahwa norma layer menjadi akar kuadrat dari jumlah elemen layer
    • assert
    • row-wise comparison
    • torch.allclose
  • SimpleModel_RMS, yaitu model dasar yang ditambah RMSNorm, sedikit menurunkan validation loss ke sekitar 2.5015

Komponen Llama 2: RoPE dan causal mask

  • RoPE adalah metode positional encoding untuk Transformer yang merepresentasikan posisi token sebagai rotasi embedding
  • get_rotary_matrix membuat matriks rotasi per posisi untuk context window dan dimensi embedding
  • Implementasi RoPE diuji dengan sifat berikut
    • inner product dari dua vektor yang diputar pada posisi m dan n harus cocok dengan rotasi posisi relatif n-m
  • RoPEAttentionHead membuat w_q, w_k, dan w_v, menerapkan rotasi RoPE pada query dan key, lalu memakai F.scaled_dot_product_attention
  • Perbedaan shape tensor antara waktu pelatihan dan waktu inferensi perlu diperhatikan
    • saat pelatihan, sering kali bentuknya sesuai konfigurasi seperti (config['batch_size'], config['context_window'], config['d_model'])
    • saat inferensi, model bisa memproses satu contoh seperti (1, 1, config['d_model'])
    • di dalam forward, pengindeksan harus berdasarkan shape dari input, bukan dari nilai konfigurasi model
  • Model yang menambahkan multi-head attention RoPE tanpa causal mask menurunkan validation loss secara drastis hingga 0.1623, tetapi hasil generasinya buruk seperti OOOO... dan IIII...
  • Setelah attention map diperiksa, terlihat semua posisi merujuk ke semua posisi lain, sehingga pada prediksi token berikutnya terjadi kebocoran informasi karena model melihat token masa depan
  • Setelah diganti dengan RoPEMaskedAttentionHead yang menerapkan is_causal=True pada F.scaled_dot_product_attention, perhatian ke area segitiga atas yang mewakili masa depan menjadi nyaris 0
  • Setelah causal mask diterapkan, validation loss menjadi 2.0815, lalu turun lagi menjadi 1.8985 dengan pelatihan lebih lama

Komponen Llama 3: SwiGLU dan menumpuk blok

  • Llama mengganti nonlinieritas ReLU dengan fungsi aktivasi SwiGLU
  • SwiGLU yang diimplementasikan adalah swish-gated linear unit yang menggunakan dua transformasi linear dan parameter beta yang dapat dipelajari
  • RopeModel yang memasukkan SwiGLU ke bagian feed-forward memiliki 592.706 parameter dan validation loss sekitar 1.8963
  • Setelah itu dibuat LlamaBlock untuk menggabungkan susunan berikut ke dalam satu blok
    • pre-normalization RMSNorm
    • masked RoPE multi-head attention
    • residual connection
    • pre-normalization RMSNorm
    • feed-forward SwiGLU
    • residual connection
  • Model Llama final menetapkan n_layers=4 dan menumpuk 4 LlamaBlock dengan nn.Sequential berbasis OrderedDict
  • Jumlah parameter model final adalah 2.370.246, dan hasil pelatihannya sebagai berikut
    • setelah pelatihan awal 4 layer, validation loss 1.5532
    • setelah dilatih lebih lanjut hingga 10.000 epoch, validation loss 1.1479
    • setelah pelatihan tambahan, validation loss 0.9997
    • loss untuk satu batch pada split test adalah 1.2358

Hasil generasi dan pemeriksaan debugging

  • Model final dapat menghasilkan nama, pemenggalan baris, dan potongan kata yang mirip gaya Shakespeare, tetapi kualitas kalimat sebenarnya masih terbatas
  • Loss cross-entropy bisa diintuisikan dari sudut pandang pemilihan token
    • loss awal 4.17 hampir setara dengan pemilihan acak pada vocabulary berukuran 65
    • loss 1.08 dapat ditafsirkan seperti memilih secara acak dari sekitar 2,9 token
  • Aliran gradient diperiksa dengan fungsi show_grads
    • fungsi ini menghitung rasio gradient bernilai absolut kecil pada setiap parameter
    • jika gradient sebagian besar parameter tidak mendekati 0, maka alirannya tergolong baik
  • Llama asli menggunakan learning schedule Cosine Annealing, tetapi pada implementasi ini hasil eksperimennya justru lebih buruk
  • Dalam eksperimen Cosine Annealing, bahkan pada tolerance yang sangat rendah, attention bias hampir tidak menerima sinyal, dan karena penyebabnya belum jelas, lebih aman memulai dari pendekatan yang sederhana dalam implementasi nyata

Bacaan terkait

1 komentar

 
GN⁺ 2023-08-10
Pendapat di Hacker News

  • Tampaknya ada bug dalam implementasi SwiGLU: dalam makalah rujukan, beta pada feed-forward network adalah konstanta, bukan nilai yang dapat dilatih, dan ditulis sebagai FFnSwiGLU = Swish1...
    Ini mengacu pada persamaan 6 di https://arxiv.org/pdf/2002.05202.pdf
    Dalam implementasi resmi llama pun beta konstanta dihilangkan: https://github.com/facebookresearch/llama/blob/main/llama/mo...
    Jika melihat baris-baris "feedforward.1.beta', 0.0" pada log blog, beta merosot menjadi 0 selama pelatihan, padahal seharusnya konstanta 1

    • Ini menunjukkan betapa sulitnya mengimplementasikan jaringan saraf Transformer secara tepat. Kesalahan bisa terjadi di banyak tahap, dan biasanya hanya tampak sebagai “performa sedikit lebih buruk dari aslinya”, sehingga sulit dipastikan
      Jaringan juga sering beradaptasi terhadap perubahan, disengaja atau tidak, dan setelah pelatihan berbagai variasi arsitektur bisa berperilaku mirip, sehingga kadang ambigu apakah harus benar-benar sama dengan aslinya
      Salah satu cara menemukan kesalahan seperti ini adalah mencocokkan output secara persis dengan implementasi referensi. Seperti model tiny-random dari HuggingFace, meskipun bobotnya acak, output harus sama persis; jika berbeda, itu sinyal adanya bug
      Namun cara ini terutama efektif untuk bug yang muncul saat inferensi, sedangkan masalah yang hanya terjadi saat pemrosesan data, optimizer, atau pelatihan lebih sulit ditangkap
    • Menurut saya nilai bias pada Transformer pada umumnya cenderung kurang cocok
      Secara pribadi saya menduga ini karena sifatnya yang autoregresif dan mirip ODE, tetapi saya tidak cukup yakin

  • Pekerjaannya bagus, tetapi SimpleBrokenModel dan SimpleModel awal memiliki cukup banyak komputasi yang terbuang. Urutannya adalah embedding 65 -> 128, linear 128 -> 128, ReLU, linear 128 -> 65; karena tidak ada nonlinieritas di antara dua layer pertama dan keduanya linier, layer linier kedua pada dasarnya tidak berguna
    Model ini pada akhirnya sama dengan MLP klasik dengan satu hidden layer, dan dari sisi FLOPS, operasi 128*128=16k terbuang dari total 128*128+65*128=24k

    • Sepertinya bukan hanya saya yang masih mempelajari nonlinieritas. Saya penasaran apakah perbaikan terbaik di sini adalah memasukkan ReLU atau SwiGLU di antara embedding dan layer linier pertama, atau sekadar menghapus layer liniernya
      Layer embedding adalah struktur khusus yang mengubah indeks token menjadi vektor embedding, jadi sepertinya tidak bisa dihapus

  • Secara keseluruhan ini menunjukkan prinsip-prinsip dasar dengan baik. Saya terutama suka kalimat “pakailah .shape secara religius. assert dan plt.imshow adalah temanmu”, dan pra-kondisi serta pasca-kondisi shape harus selalu di-assert
    Saya juga penasaran apakah bear atau typeguard mendukung pemeriksaan semacam ini sebagai decorator
    Namun bagian “pilih model yang kecil, sederhana, dan cepat, lalu buat helper untuk mengevaluasi secara kualitatif” sepertinya maksudnya evaluasi kuantitatif. Dengan begitu ada baseline numerik untuk dibandingkan dengan teknik yang lebih canggih
    Saran untuk mengimplementasikan komponen makalah satu per satu juga perlu lebih presisi. Makalah biasanya mencoba beberapa perubahan sekaligus lalu menunjukkan kontribusi tiap elemen melalui ablation experiment, jadi menurut saya lebih baik mulai dari perubahan arsitektur inti, lalu mengevaluasi setiap perubahan atomik sesuai urutan dampak terbesar dalam ablation experiment sambil menjaga dependensi

    • Alih-alih bear atau typeguard, berkat https://peps.python.org/pep-0646/, sebagian hal bisa langsung dimasukkan lewat anotasi tipe Python
      Misalnya, shape per sumbu bisa diekspresikan dalam tipe seperti ndarray[float, Dim1, *Shape], dan shape keluaran bisa dioverload sesuai nilai axis
    • Saya tidak terlalu tahu PyTorch, tetapi terakhir kali saya mengecek belum demikian, sedangkan Jax mendukung pemeriksaan runtime dasar untuk shape matriks melalui bear / typeguard
      Meski begitu, Python tampaknya sulit sebagus Julia. Sistem tipe Julia bisa jauh lebih mudah menjamin apakah ukuran matriks cocok

  • Saya penasaran apa prinsip penggunaan SwiGLU alih-alih ReLU. Entah para penulis sekadar mencoba semua fungsi nonlinier yang memungkinkan, atau ada alasan yang lebih mendalam

    • Seperti banyak riset lain, jika tidak ada penjelasan jelas yang didukung studi ketat, kemungkinan besar mereka melakukan pencarian hill-climbing secara acak terhadap perubahan satu baris yang tampak keren, lalu berhenti ketika sudah waktunya menulis makalah dan melakukan ablation experiment

  • Karena bearblog sedang terkena DDoS, saya tinggalkan repositorinya: https://github.com/bkitano/llama-from-scratch

  • Dari sudut pandang orang yang sedang belajar AI, saya mencoba merangkum secara singkat istilah-istilah yang muncul di tulisan itu. Token adalah pengenal bilangan bulat yang merepresentasikan potongan teks, dan dalam LLM, potongan karakter yang sering digunakan dikelompokkan dalam ukuran kosakata yang terbatas
    Fungsi loss adalah nilai yang mengukur selisih antara prediksi dan jawaban benar; makin rendah makin baik. PyTorch adalah library untuk menangani tensor dan jaringan neural, sedangkan tensor adalah array angka multidimensi yang mencakup skalar, vektor, dan matriks
    Jaringan neural adalah struktur koneksi neuron yang memiliki bobot dan bias, dan layer linear adalah struktur sederhana di mana semua input dan output saling terhubung. ReLU adalah fungsi aktivasi seperti Math.max(0, x); jika hanya menumpuk layer linear, pada akhirnya hasilnya sama saja dengan satu fungsi linear, jadi non-linearitas ditambahkan untuk meningkatkan kemampuan belajar
    Gradien adalah besaran perubahan numerik yang dihitung selama pelatihan untuk membuat model lebih akurat, dan batch normalization adalah metode yang membantu pelatihan dengan menyesuaikan angka-angka yang mengalir. Positional encoding memberi tahu posisi relatif token-token dalam bentuk vektor
    Operator @ di Python adalah alias untuk __matmul__ dan digunakan untuk perkalian matriks. Epoch berarti melatih seluruh dataset satu kali, sedangkan batch adalah jumlah data yang dimasukkan sekaligus sebelum pembaruan parameter
    Attention adalah inti yang membuat LLM bekerja; ia memproses token input secara paralel untuk membuat tensor perantara, lalu menggunakannya untuk menghasilkan token output

    • Di luar bidang ini, orang mungkin tidak tahu apa yang dimaksud dengan “Karpathy”. Jika Andrej Karpathy diperkenalkan bersama konteks seperti “komunikator sains sekaligus peneliti”, maksud untuk merujuk pada tulisan atau videonya akan menjadi lebih jelas
    • Bagi pemula, token lebih akurat dipahami bukan sekadar pengenal bilangan bulat untuk potongan teks, melainkan lebih dekat ke potongan kata yang cukup umum sehingga berguna dengan sendirinya
      Misalnya, writ yang sama-sama muncul dalam writing, written, dan writer bisa menjadi satu token, dan writer bisa ditokenisasi menjadi writ dan er
      Embedding adalah tahap yang mengubah token-token seperti ini menjadi representasi angka yang unik
    • Jika fungsi-fungsi linear dikomposisikan, hasilnya tetap fungsi linear. Jadi jika semuanya linear, meski menumpuk banyak layer, semua layer selain satu sebenarnya menjadi pemborosan; untuk menghindarinya dibutuhkan non-linearitas
    • Selain seri video Karpathy dan accompanying repo, saya penasaran apakah ada materi atau buku lain yang sangat membantu dalam perjalanan belajar
    • Saya penasaran apa sebenarnya yang dilakukan batch normalization, dan bagaimana itu membantu

  • Jika ada implementasi dan checkpoint yang sudah ada untuk model tersebut, cara paling efektif untuk memastikan implementasi sendiri benar adalah memuat checkpoint itu lalu membandingkan outputnya
    Jika outputnya tidak cocok, biasanya ada detail implementasi yang salah, dan kita bisa menelusuri setiap layer secara sistematis untuk menemukan perbedaan sebenarnya. Dalam proses itu, bisa saja kita menemukan keanehan pada implementasi yang sudah ada
    Ini adalah soal model itu sendiri, sedangkan pelatihan adalah sumbu yang terpisah. Meski begitu, jika hyperparameter dibuat cukup mirip, saat implementasi modelnya benar biasanya hasilnya akan cukup baik

  • Baik cara membaca paper maupun isi paper tersebut sama-sama bagus, dan saya juga merekomendasikan seri Makemore dari Karpathy

  • Saran-saran ringkasnya sangat bagus, dan menurut saya saran untuk melakukan assert pada shape tensor berlaku juga untuk library aljabar linear umum mana pun. Saat menulis kode aljabar linear yang kompleks, sangat penting untuk bergerak dalam langkah kecil dan menulis kode secara defensif
    Memprogram aljabar linear di bahasa-bahasa arus utama itu mengerikan karena tidak ada pemeriksaan shape saat compile time. Shape tensor seharusnya menjadi bagian dari tipe, dan jika mencoba mengalikan 3x4 dengan 3x4 tanpa transpose, kompilasinya sendiri seharusnya gagal
    Gagal pada operasi dengan dimensi yang tidak cocok setelah menjalankan komputasi panjang itu benar-benar skenario terburuk
    Menurut saya perangkat pada tensor PyTorch juga seharusnya bertipe statis. Saat ini, jika mencoba mengalikan tensor di memori CPU dengan tensor di memori GPU, yang muncul adalah error runtime