un — Kembangkan Model Bahasa: Backprop di CUDA Kustom [DESCRIPTION /]

un

tamu

1 / ?

kembali ke pelajaran

Gradien Lokal Berlipat Ganda

Kernel Forward & Backward

Forward Pass

Forward pass ANDREA-120M membawa input melalui urutan operasi:

x = embed(token_ids)         # embedding token
for layer in 12_layers:
x = x + attn(LN(x))      # sublayer perhatian
x = x + mlp(LN(x))       # sublayer MLP
logits = LN(x) @ embed.T     # proyeksi output terikat
loss   = cross_entropy(logits, targets)

Setiap operasi membaca tensor input & menghasilkan tensor output. Forward pass berakhir pada satu skalar: cross-entropy loss untuk batch ini.

Backward Pass

Pelatihan memperbarui bobot ke arah yang menurunkan loss. Untuk mendapatkan arah pembaruan, engine membutuhkan:

dL/dW untuk setiap W yang dapat dipelajari di dalam model

Aturan rantai memberikan ini. Untuk sebuah rantai loss = f(g(h(x))):

dL/dx = (dL/df) * (df/dg) * (dg/dh) * (dh/dx)

Setiap faktor adalah gradien lokal: seberapa besar output dari satu operasi berubah ketika inputnya berubah sedikit. Mengalikan gradien lokal secara mundur melalui graf menyebarkan sinyal loss ke setiap bobot.

Diferensiasi Mode Balik

Backprop menghitung gradien dalam urutan balik: dimulai dari dL/dlogits = 1, kemudian berjalan mundur melalui cross-entropy, kemudian proyeksi output, kemudian layer norm, kemudian dua belas blok transformer, kemudian embeddings. Pada setiap langkah, kalikan gradien masuk dengan Jakobian lokal.

Mode balik efisien ketika output adalah skalar tunggal (loss) & ada banyak input (bobot). Satu pass mundur menghasilkan gradien untuk setiap bobot dalam model. Mode maju membutuhkan satu pass per bobot; untuk ANDREA-120M dengan ~120M bobot, mode maju tidak memungkinkan.

Mengapa Mode Balik

ANDREA-120M memiliki ~120M bobot & menghasilkan satu skalar loss per langkah pelatihan. Bandingkan *reverse-mode automatic differentiation* dengan *forward-mode*. Nyatakan (1) mode mana yang menghasilkan semua gradien bobot dalam satu *backward pass*; (2) berapa banyak *forward-mode pass* yang diperlukan untuk menghitung semua 120M gradien bobot; (3) mode mana yang digunakan ANDREA & mengapa.

Setiap Operasi Forward Mendapat Kembaran Backward

Disiplin Pasangan

microgpt_cuda.cu menyediakan dua kernel CUDA untuk setiap operasi: satu yang menghitung output forward, satu yang menghitung gradien input berdasarkan gradien output. Pasangannya satu-satu:

Kernel maju	Kernel mundur	Operasi
`k_embed_fwd`	`k_embed_bwd`	Pencarian embedding token
`k_layernorm_fwd`	`k_layernorm_bwd`	Normalisasi lapisan
`k_attn_qkv_fwd`	`k_attn_qkv_bwd`	Proyeksi Q, K, V
`k_attn_fwd`	`k_attn_bwd`	Perhatian dot-product terukur
`k_attn_out_fwd`	`k_attn_out_bwd`	Proyeksi keluaran W_O
`k_mlp_fwd`	`k_mlp_bwd`	MLP (dengan GELU)
`k_residual_add`	`k_residual_add_bwd`	Koneksi residual
`k_loss_fwd`	`k_loss_bwd`	Fungsi kerugian cross-entropy

Delapan pasangan operasi mencakup transformer lengkap. Ditambah beberapa kernel utilitas: k_grad_norm_partial, k_grad_norm_final, k_grad_scale untuk clipping gradien (lihat aktivitas 75).

Tugas Sebuah Kernel Backward

Diberikan gradien yang mengalir dari lapisan berikutnya (grad_output), sebuah kernel backward menghitung:

1. grad_input: gradien dengan respecto tensor input dari operasi. Ini diteruskan lebih lanjut ke belakang.

2. grad_weight: gradien dengan respecto parameter yang dapat dipelajari dalam operasi. Ini masuk ke dalam state optimizer.

Keduanya dihitung dalam satu peluncuran kernel. Thread CUDA bekerja sama pada tile-tile tensor gradien secara paralel.

Tensor yang Disimpan

Komputasi mundur sering membutuhkan nilai dari forward pass. Misalnya, k_layernorm_bwd membutuhkan mean & variance yang dihitung selama forward; k_mlp_bwd membutuhkan pre-aktivasi GELU. Engine pelatihan menyimpan ini di buffer khusus selama forward, kemudian membacanya selama mundur.

Biaya memori: kira-kira bentuk yang sama dengan output forward untuk setiap tensor yang disimpan. Untuk ANDREA-120M dengan batch=8, seq=1024, d_model=768, satu tensor yang disimpan adalah 8 × 1024 × 768 × 4 bytes = 25 MB. Di seluruh 12 lapisan & beberapa tensor yang disimpan per lapisan, aktivasi mendominasi VRAM selama pelatihan (~5-10 GB pada kartu 24 GB).

Melacak Satu Langkah Backward

ANDREA-120M menyelesaikan forward pass melalui satu blok transformer. Lacak apa yang terjadi selama backward pass melalui blok yang sama (dalam struktur pre-norm: `x = x + Attention(LN(x))` kemudian `x = x + MLP(LN(x))`). Sebutkan kernel backward dalam urutan mereka dijalankan, & nyatakan kernel forward mana yang dipasangkan dengan masing-masing. Tutupi setidaknya 4 kernel.

Di Mana Gradien Disimpan di Memori

Satu Tensor Gradien Per Tensor Bobot

Setiap tensor bobot yang dapat dipelajari di ANDREA-120M memiliki tensor gradien yang cocok dengan bentuk yang identik. Untuk setiap blok:

W_Q       [768, 768]     ↔   grad_W_Q       [768, 768]
W_K       [768, 768]     ↔   grad_W_K       [768, 768]
W_V       [768, 768]     ↔   grad_W_V       [768, 768]
W_O       [768, 768]     ↔   grad_W_O       [768, 768]
W_1       [768, 3072]    ↔   grad_W_1       [768, 3072]
W_2       [3072, 768]    ↔   grad_W_2       [3072, 768]
LN1.gamma [768]          ↔   grad_LN1.gamma [768]
LN1.beta  [768]          ↔   grad_LN1.beta  [768]
LN2.gamma [768]          ↔   grad_LN2.gamma [768]
LN2.beta  [768]          ↔   grad_LN2.beta  [768]

Plus token embeddings, position embeddings, & a final layer norm. Total memori buffer gradien cocok dengan memori bobot: ~120M floats, ~480 MB pada FP32, ~240 MB pada FP16.

Akumulasi Melalui Microbatch

batch_size ANDREA = 8 muat di VRAM pada FP16. Batch efektif yang lebih besar memerlukan akumulasi gradien: jalankan beberapa forward+backward pass pada batch kecil, menjumlahkan gradien ke buffer yang sama, lalu lakukan satu langkah optimizer.

for microbatch in range(n_microbatches):
forward(microbatch)
backward()           # MENAMBAHKAN ke buffer grad, tidak menimpa
scale_grads(1.0 / n_microbatches)  # rata-rata di seluruh microbatches
optimizer_step()
zero_grads()             # reset untuk langkah pelatihan berikutnya

Kernel backward menggunakan semantik +=, bukan =. Setiap panggilan menambahkan kontribusi gradien ke buffer yang ada; buffer menyimpan jumlah berjalan hingga zero_grads() membersihkannya.

Status Optimizer

AdamW (aktivitas 73) menyimpan dua buffer lagi per bobot: momen pertama m & momen kedua v. Total memori waktu pelatihan:

bobot:    1× jumlah bobot
gradien:  1× jumlah bobot
Adam m:     1× jumlah bobot
Adam v:     1× jumlah bobot
acts tersimpan: ~2-4× tergantung lapisan & batch
──────────────────────────────────────────
total:      ~6-8× jumlah bobot

ANDREA-120M pada FP16: ~240 MB × 4 buffer (bobot, grad, m, v) + ~5-10 GB aktivasi = ~10-12 GB total. Nyaman di bawah batas 24 GB RTX 4090. ANDREA-12M dilatih dalam 1.4 GB; penskalaan parameter 10× membawa ~10× memori.

Menentukan Ukuran Buffer Gradien

ANDREA-120M memiliki ~120.000.000 bobot & menggunakan akumulasi gradien di seluruh 4 mikro-batch per langkah pelatihan. Hitung: (a) ukuran buffer gradien dalam MB pada FP16; (b) total memori untuk bobot + gradien + Adam m + Adam v pada FP16; (c) berapa banyak panggilan `forward()` + `backward()` terpisah yang dijalankan per langkah pelatihan. Tunjukkan perhitungan Anda.

Kontrol Penuh Atas Memori & Presisi

Biaya Kerangka Generik

PyTorch & JAX membuat autograd nyaman: tulis kode Python, dapatkan gradien secara otomatis. Biayanya: lapisan dispatch generik antara kode Anda & CUDA. Setiap operasi melewati overhead interpreter Python, pembukuan kerangka kerja, & pemilihan kernel dinamis. Untuk melatih model bahasa kecil pada satu GPU, overhead itu penting.

Biaya konkret yang dihindari ANDREA:

1. Latensi interpreter Python. Setiap operasi PyTorch melintasi batas Python/C++. Untuk ~100 peluncuran kernel per langkah pelatihan pada ~9 langkah/menit, itu ~900 penyeberangan batas per menit. Dispatch tingkat C menghilangkan ini.

2. Ketidakpastian alokator framework. Alokator caching PyTorch memberikan throughput yang baik secara rata-rata tetapi memori puncak yang tidak dapat diprediksi. Mesin pelatihan ANDREA mengalokasikan pra setiap buffer saat startup; tidak ada realokasi selama pelatihan, tidak ada fragmentasi, tidak ada OOM mendadak pada langkah 100K.

3. Pemilihan kernel generik. PyTorch memilih kernel saat runtime melalui heuristik. ANDREA memilih kernel saat waktu kompilasi, disesuaikan dengan ukuran tile tensor core RTX 4090.

4. Pipa presisi campuran. Jalur cuBLAS FP16 ANDREA-120M & eksperimen tensor core FP8 E4M3 ANDREA memerlukan kontrol presisi yang tepat atas tensor mana yang berada pada presisi mana. Kerangka kerja generik menyediakan kontrol ini melalui API berlapis; penulisan CUDA kustom menuliskannya secara langsung.

Kompromi

Biaya CUDA kustom: lebih banyak kode untuk ditulis, lebih banyak bug untuk ditemukan, tidak ada ekosistem komunitas. microgpt_cuda.cu milik ANDREA adalah ~6000 baris CUDA tulisan tangan yang memakan waktu berbulan-bulan untuk di-debug. Setiap operasi baru memerlukan penulisan kernel forward, kernel backward, & tes.

Yang didapat ANDREA:

- Reproduktibilitas penuh. Pipeline pelatihan adalah satu binary C plus satu proxy Python. Tidak ada pergeseran versi di berbagai rilis PyTorch, tidak ada ketidakcocokan versi CUDA dengan wheel framework.

- Resume bit-exact. SIGTERM memicu penulisan checkpoint yang menangkap setiap tensor tepat seperti yang dilihat GPU. Resume melanjutkan trajektori loss yang sama seperti run sebelumnya.

- Memori yang dapat diprediksi. ANDREA-120M dilatih selama 200K langkah tanpa OOM. Memori dihitung sejak startup engine.

- Akses perangkat keras langsung. Ukuran tile tensor core, pengaturan FP8 E4M3, salinan memori asinkron: semuanya dapat diakses langsung di CUDA, tidak transparan di framework generik.

Reproduktibilitas Sebagai Misi

Bagian 9 whitepaper ANDREA mencantumkan tumpukan reproduktibilitas lengkap:

Mesin pelatihan: microgpt/microgpt_cuda.cu
Proxy pelatihan: microgpt/training_proxy.py
Konfigurasi eksperimen: experiments/ANDREA-*-TRAIN.json
Pipeline data: scripts/pull-hermes3.py, scripts/prep-megachat.py
Dashboard: scripts/live-loss-dashboard.html
Spesifikasi Bandit: docs/FIREHOSE-BANDIT.md
Dokumentasi model: docs/ANDREA.md

Persyaratan perangkat keras: satu GPU NVIDIA dengan ≥8 GB VRAM (RTX 3060 atau lebih baik). Siapa pun dapat mereproduksi ANDREA-12M dari artefak ini. Jalur CUDA kustom adalah bagian dari alasan mengapa: tidak ada pembekuan versi framework, tidak ada kejutan dependensi lima tahun dari sekarang.

Sinyal & Checkpoint

Loop pelatihan CUDA merespons dua sinyal POSIX:

- SIGTERM: tulis checkpoint segera, lalu keluar. Digunakan saat menghentikan pelatihan dengan bersih.

- SIGUSR1: tulis checkpoint segera, lanjutkan pelatihan. Digunakan selama polish pivot di v3 untuk menangkap state tanpa mengganggu jalannya.

Format checkpoint: [int32 step][int32 n_params][n_params × float32 weights][n_params × float32 m][n_params × float32 v]. Penghitung langkah, jumlah bobot, kemudian bobot diikuti oleh momen Adam. Melanjutkan secara bit-exact. Proxy mengarsipkan .samples.json & .state.json secara terpisah pada polish; .loss.json tidak pernah diarsipkan (itu mengakumulasi riwayat pelatihan penuh).

Mengapa Tidak PyTorch

ANDREA bisa menggunakan autograd PyTorch daripada menulis `microgpt_cuda.cu` secara manual. Berikan dua alasan teknik yang berbeda mengapa ANDREA memilih CUDA kustom. Satu alasan harus merujuk pada kontrol memori atau presisi; yang lain harus merujuk pada reproduktibilitas, ketergantungan framework, atau pemeliharaan jangka panjang.