Mixed precision dan kecepatan

autocast, numerik bfloat16, dan kebiasaan pemanfaatan GPU.

Daftar isi

Mixed precision menghitung sebagian operasi pelatihan pada presisi rendah seperti bfloat16 atau float16, sehingga pelatihan pada GPU menjadi lebih cepat dan hemat memori tanpa mengorbankan konvergensi. Topik ini menjelaskan pembagian presisi yang menyelamatkan numeriknya, memeragakan autocast pada CPU dengan bfloat16, dan menetapkan kebiasaan pemanfaatan GPU yang menyertainya. Setelah selesai, keputusan mengaktifkan mixed precision dan batas-batasnya memiliki dasar yang jelas.

Dua presisi dan pembagian perannya#

Presisi penuh float32 menjadi standar karena rentangnya aman untuk akumulasi gradien, tetapi operasinya paling mahal. Presisi rendah bfloat16 atau float16 mempercepat perkalian matriks secara besar pada GPU modern dan memangkas ukuran aktivasi di memori. Pembagian peran mixed precision menyelesaikan pertentangan itu, yaitu perkalian matriks besar berjalan pada presisi rendah yang cepat, sedangkan operasi yang sensitif seperti penjumlahan loss, softmax, dan normalisasi tetap pada float32.

Perbedaan kedua format presisi rendah menentukan pilihan. Bfloat16 memiliki rentang eksponen yang sama dengan float32 sehingga tidak memerlukan penanganan kelebihan rentang khusus, dan format ini menjadi bawaan pilihan pada perangkat yang mendukungnya, sedangkan float16 memiliki rentang yang lebih sempit dan menuntut scaling gradien agar tidak tenggelam ke nol. Pada PyTorch, torch.autocast mengurus pemilihan presisi per operasi, dan GradScaler menemani float16 dengan penskalaan yang otomatis.

autocast pada CPU dengan bfloat16#

autocast tidak eksklusif milik GPU, dan peragaannya dapat dijalankan penuh pada CPU dengan bfloat16. Kode berikut menjalankan layer linier besar di dalam dan di luar autocast, membandingkan presisi keluarannya, dan mengukur kesamaan arah hasilnya.

python Contoh: autocast bfloat16 pada CPU
import torch
import torch.nn as nn

torch.manual_seed(0)
lin = nn.Linear(256, 256)
x = torch.randn(64, 256)
with torch.no_grad():
    y32 = lin(x)
    with torch.autocast("cpu", dtype=torch.bfloat16):
        ybf = lin(x)
selisih = (y32.float() - ybf.float()).abs()
print("presisi keluaran autocast:", ybf.dtype)
print("selisih rata-rata        :", round(float(selisih.mean()), 6))
print("korelasi arah            :", round(float(torch.corrcoef(torch.stack([y32.flatten(), ybf.flatten()]))[0, 1]), 4))
text Keluaran
presisi keluaran autocast: torch.bfloat16
selisih rata-rata        : 0.001253
korelasi arah            : 1.0

Keluaran autocast bernilai bfloat16 dengan selisih rata-rata yang sangat kecil terhadap float32, dan korelasi arahnya nyaris satu, sehingga sinyalnya tetap utuh pada tingkat yang dipakai pelatihan. Pada GPU, pola pemanggilannya sama persis dengan mengganti perangkatnya, dan itulah alasan kebiasaan yang dibangun di sini berpindah tanpa perubahan kode.

Siklus pelatihan dengan autocast#

Penempatan autocast mengikuti tiga aturan yang tidak berubah antar perangkat. Pertama, forward pass dan perhitungan loss berada di dalam blok with torch.autocast(...), sedangkan backward di luar bloknya karena autograd membawa informasi presisi dari grafik maju. Kedua, dengan float16 pada CUDA, GradScaler membungkus loss.backward() dan optimizer.step() untuk mencegah gradien yang tenggelam, sedangkan bfloat16 tidak membutuhkannya. Ketiga, evaluasi tetap dibungkus no_grad, dan metrik dihitung dalam float32 agar akurasinya tidak ikut menurun.

Pola lengkapnya pada GPU dituliskan sebagai berikut dan sengaja tidak dijalankan pada penulisan jalur ini karena sandbox tidak memiliki CUDA:

python
# tidak-dijalankan
import torch

scaler = torch.amp.GradScaler("cuda")
for xb, tb in loader:
    optimizer.zero_grad()
    with torch.autocast("cuda", dtype=torch.float16):
        loss = loss_fn(model(xb.to("cuda")), tb.to("cuda"))
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Kebiasaan yang menyertainya menyangkut dua hal di luar presisi. Ukuran batch dinaikkan selagi memori yang dihemat oleh aktivasi presisi rendah memungkinkannya, karena langkah yang membawa lebih banyak contoh menaikkan pemanfaatan GPU lebih daripada presisinya sendiri. Selain itu, profil jarak sempit yang sebenarnya, yaitu apakah bottleneck pelatihan berada di perhitungan, transfer data, atau pemuatan data, diperiksa sebelum optimasi apa pun, karena mixed precision hanya menolong bottleneck perhitungan.

Batas dan pemeriksaan konvergensi#

Mixed precision bukan tanpa risiko, dan risikonya muncul pada operasi yang akumulatif. Penjumlahan loss pada batch besar, kumulasi gradien yang berlangsung lama, dan statistik batch normalization yang variansinya kecil adalah contoh yang perlu dijaga tetap pada float32, dan autocast sudah menangani sebagian besar dari daftar ini secara bawaan. Gejala yang perlu diwaspadai adalah loss yang menjadi nan pada langkah-langkah awal, dan langkah pertama penanganannya adalah menurunkan learning rate atau beralih ke bfloat16 yang rentangnya lebih aman.

Pemeriksaan konvergensi yang jujur membandingkan run presisi penuh dengan run mixed precision pada benih dan data yang sama, dan perbandingannya memakai kaidah pelacakan eksperimen pada topik Pelacakan eksperimen. Selisih metrik akhir yang dalam kisaran variansi antar benih dinyatakan aman, sedangkan pergeseran yang lebih besar menuntut audit operasi yang sensitif sebelum kecepatannya diambil.

Ringkasan#

  • Mixed precision menjalankan perkalian matriks pada presisi rendah dan menyimpan operasi sensitif pada float32, sehingga GPU lebih cepat dan hemat memori.
  • Bfloat16 berbagi rentang eksponen float32 dan tidak butuh GradScaler, sedangkan float16 menuntut penskalaan gradien otomatis.
  • Forward dan loss berada di dalam blok autocast, backward di luar, dan pola ini berpindah antar perangkat tanpa perubahan kode.
  • Konvergensinya diverifikasi dengan perbandingan run berbenih sama, dan bottleneck diperiksa profil sebelum optimasi mana pun.

Referensi#