FlashAttention, pelatihan multi-GPU, dan kuantisasi bobot

Infrastruktur AI modern bergantung pada kernel atensi yang sadar perangkat keras, paralelisme multi-GPU terdistribusi, dan kuantisasi bit rendah untuk melatih dan melayani model berukuran masif secara efisien. Ketiga teknik ini menjembatani kesenjangan antara skala model secara teoretis dan batas fisik GPU.

1. FlashAttention-2 dan FlashAttention-3 pada GPU Hopper (H100)

Meskipun FlashAttention asli mengatasi hambatan bandwidth memori, ia masih menyisakan banyak performa komputasi mentah yang tidak terpakai pada GPU modern. FlashAttention-2 (2023) dan FlashAttention-3 (2024) merombak cara thread GPU dan mesin perangkat keras berkoordinasi.

• FlashAttention-2 (Pembagian Kerja yang Lebih Baik).

• Mengurangi Operasi Non-Matmul. Tensor Core GPU menjalankan perkalian matriks (matmul) sekitar 16× lebih cepat daripada core CUDA standar menjalankan operasi skalar seperti penskalaan ulang Softmax. FlashAttention-2 menulis ulang matematika Softmax daring agar penskalaan ulang ditunda hingga akhir loop dalam, sehingga mengurangi instruksi non-matmul yang lambat.
• Paralelisasi pada Panjang Sekuens. Versi asli memparalelkan pekerjaan terutama pada ukuran batch dan head atensi, sehingga core GPU menganggur saat pelatihan dengan sekuens panjang dan ukuran batch kecil. FlashAttention-2 menambahkan paralelisme pada dimensi panjang sekuens dan membalik urutan loop (loop luar atas blok Q, loop dalam atas K dan V), sehingga meminimalkan penulisan shared memory yang lambat.

• FlashAttention-3 (Pemanfaatan Perangkat Keras Hopper H100).

• Perpindahan Data Asinkron (TMA). Arsitektur Hopper dari NVIDIA memperkenalkan Tensor Memory Accelerator (TMA), unit perangkat keras khusus yang menyalin data antara HBM global dan SRAM on-chip tanpa menyibukkan thread komputasi.
• Spesialisasi Warp (Model Producer-Consumer). FlashAttention-3 membagi warp GPU (kelompok 32 thread) ke dalam peran khusus. Warp Producer menggunakan TMA untuk mengambil lebih awal (prefetch) tile K dan V berikutnya ke SRAM, sementara warp Consumer memasukkan tile saat ini ke Tensor Core (instruksi WGMMA). Ini membuat transfer memori tumpang tindih dengan perhitungan matriks dan Softmax secara bersamaan.
• Dukungan Presisi Rendah FP8. FlashAttention-3 membuka Tensor Core FP8 milik Hopper, menggandakan throughput mentah dibandingkan FP16. Untuk mencegah hilangnya presisi akibat aktivasi outlier di FP8, ia memakai kuantisasi blok (menskalakan tile kecil secara independen) dan pemrosesan inkoheren (mengalikan Q dan K dengan matriks ortogonal acak untuk melunakkan nilai ekstrem sebelum dikuantisasi).

2. Strategi Penskalaan Multi-GPU (Paralelisme 3D)

Ketika sebuah model memiliki puluhan atau ratusan miliar parameter, bobot, gradien, dan status optimizer-nya melebihi VRAM satu GPU. Pelatihan skala besar menggabungkan tiga strategi paralelisasi yang berbeda, yang sering disebut Paralelisme 3D.

• Data Parallelism & ZeRO / FSDP. Pada Data Parallelism standar, setiap GPU menyimpan salinan penuh model sehingga memboroskan memori. Kerangka kerja modern memakai ZeRO (Zero Redundancy Optimizer) atau PyTorch FSDP (Fully Sharded Data Parallel), yang me-shard parameter, gradien, dan status optimizer di seluruh GPU, lalu mengumpulkan bobot setiap layer tepat pada waktunya untuk komputasi.

• Tensor Parallelism (Intra-Layer). Alih-alih menduplikasi perkalian matriks feed-forward yang besar Y = XW, Tensor Parallelism membagi W per kolom ke beberapa GPU (W = [W_1, W_2]). Setiap GPU menghitung separuh keluarannya secara bersamaan dan bertukar hasil parsial melalui NVLink berkecepatan tinggi (900 GB/s pada H100).

• Pipeline Parallelism (Inter-Layer). Layer 1–8 berada di GPU 1, Layer 9–16 di GPU 2, dan seterusnya. Untuk mencegah GPU hilir menganggur sambil menunggu layer sebelumnya selesai ("pipeline bubble"), kerangka kerja memotong setiap batch pelatihan menjadi micro-batch yang lebih kecil agar semua GPU tetap sibuk seperti lini perakitan.

3. Kuantisasi (GPTQ, AWQ) dan Kernel Dekuantisasi Langsung (On-the-Fly)

Saat inferensi (menghasilkan token satu per satu), Large Language Model dibatasi oleh bandwidth memori, bukan oleh komputasi. Memuat 70 miliar parameter 16-bit (140 GB) dari VRAM hanya untuk menghasilkan satu token membuat Tensor Core menunggu transfer memori. Kuantisasi memampatkan bobot dari float 16-bit menjadi integer 4-bit (INT4), sehingga mengurangi ukuran memori dan waktu transfer hingga hampir 4×.

• GPTQ (Kuantisasi Pasca-Pelatihan melalui Kalibrasi Hessian).

• Pembulatan sederhana bobot 16-bit menjadi integer 4-bit merusak akurasi model. GPTQ memasukkan dataset kalibrasi kecil melalui model dan memakai matematika orde dua (matriks Hessian invers) untuk mengukur seberapa sensitif setiap bobot.
• Ia mengkuantisasi bobot kolom demi kolom; setiap kali pembulatan satu bobot menimbulkan galat, GPTQ secara matematis menyesuaikan bobot yang belum dikuantisasi di baris tersebut untuk meniadakan galat itu.

• AWQ (Activation-Aware Weight Quantization).

• Para peneliti menemukan bahwa sekitar 0,1% hingga 1% kanal bobot jauh lebih penting daripada sisanya karena mengalikan nilai aktivasi yang besar.
• Alih-alih mempertahankan bobot kritis tersebut dalam 16-bit (yang menciptakan tata letak memori presisi campuran yang berantakan dan lambat), AWQ menaikkan skala kanal bobot yang menonjol itu sebelum kuantisasi dan menurunkan skala aktivasi input yang sesuai saat runtime, sehingga akurasi terjaga sementara 100% bobot tetap dalam format 4-bit yang seragam.

• Cara Kernel CUDA Kustom Melakukan Dekuantisasi Langsung (W4A16).

• GPU tidak dapat langsung mengalikan aktivasi 16-bit (FP16) dengan bobot integer 4-bit (INT4). Sebagai gantinya, kernel inferensi khusus (seperti Marlin atau ExLlama) menjalankan eksekusi W4A16 (Bobot 4-bit, Aktivasi 16-bit):
• 1. Pembacaan HBM Terkompresi: Kernel memuat bobot 4-bit yang dipak (delapan nilai INT4 dipak dalam satu integer 32-bit) dari VRAM global yang lambat, sehingga memangkas penggunaan bandwidth memori sebesar 4×.
• 2. Pembongkaran di Tingkat Register: Di dalam register GPU yang sangat cepat, kernel CUDA memakai pergeseran bit (>>) dan bitmask (&) yang cepat untuk membongkar integer INT4 dan menerapkan faktor skala/zero-point yang tersimpan untuk mengubahnya kembali menjadi nilai FP16 secara langsung.
• 3. Matematika Tensor Core Seketika: Bobot FP16 yang baru dibongkar langsung dimasukkan ke Tensor Core bersama aktivasi FP16 tanpa pernah menulis kembali bobot 16-bit yang sudah diperluas ke VRAM.