Mixed precision

Topik ini membahas mixed precision, yaitu teknik menjalankan sebagian operasi pelatihan pada tipe data presisi rendah sementara bagian lain tetap presisi penuh. Dua komponen PyTorch yang terlibat dibahas satu per satu, yaitu torch.autocast yang memilih presisi untuk setiap operasi dan torch.amp.GradScaler yang mencegah gradien kecil pada float16 jatuh ke nol. Konteks kapan teknik ini memberi keuntungan nyata juga dirinci, karena keuntungannya tidak muncul pada semua kondisi. Seluruh contoh dijalankan pada build CPU dengan bfloat16, sehingga dapat direproduksi tanpa GPU.

Presisi tunggal dan presisi campuran#

PyTorch menyimpan tensor dan bobot model dalam float32 secara default, yaitu format presisi tunggal (single precision) berukuran 32 bit. Mixed precision mengganti sebagian operasi ke format 16 bit, baik float16 (fp16) maupun bfloat16 (bf16), sehingga setiap angka hanya menempati dua byte, bukan empat byte. Penghematan itu membuka dua keuntungan sekaligus, yaitu memori yang lebih longgar untuk batch yang lebih besar dan laju komputasi yang lebih tinggi pada GPU yang menyediakan unit khusus presisi rendah. Kedua format 16 bit tersebut tidak identik, dan perbedaannya menentukan cara penggunaannya.

Perbedaan terpenting ada pada rentang nilainya. float16 menyimpan eksponen yang lebih pendek daripada float32, sehingga nilai sangat kecil mudah jatuh ke nol (underflow) dan nilai sangat besar mudah meluap. bfloat16 mempertahankan rentang eksponen float32, sehingga tidak rentan terhadap kedua gejala tersebut, dengan konsekuensi presisi mantissa yang lebih rendah. Perbedaan inilah yang menjelaskan mengapa loss scaling hanya diperlukan untuk float16.

Memilih presisi per operasi dengan autocast#

torch.autocast adalah context manager yang memilih dtype untuk setiap operasi di dalam bloknya sesuai daftar bawaan PyTorch. Operasi yang cocok dengan presisi rendah seperti nn.Linear dan konvolusi dijalankan pada bfloat16 atau fp16, sedangkan operasi yang sensitif seperti penghitungan loss tetap pada float32. Bobot model tidak pernah diubah, sehingga pembaruan bobot tetap dihitung dalam presisi penuh. Skrip berikut menguji perilaku tersebut pada CPU:

import torch
import torch.nn as nn

linear = nn.Linear(4, 3)
x = torch.randn(2, 4)  # shape (2, 4)

# tanpa autocast: keluaran mengikuti dtype bobot, yaitu float32
print("Tanpa autocast :", linear(x).dtype)

with torch.autocast("cpu", dtype=torch.bfloat16):
    y = linear(x)  # shape (2, 3)
    print("Dengan autocast:", y.dtype)

# konsekuensi presisi rendah pada memori
a = torch.randn(2, 4)
print("float32 :", a.element_size(), "byte per elemen")
print("bfloat16:", a.to(torch.bfloat16).element_size(), "byte per elemen")

Tanpa autocast : torch.float32
Dengan autocast: torch.bfloat16
float32 : 4 byte per elemen
bfloat16: 2 byte per elemen

Keluaran menunjukkan dua hal. Keluaran nn.Linear berpindah dari torch.float32 ke torch.bfloat16 begitu operasinya berada di dalam konteks autocast, tanpa perubahan apa pun pada definisi model. Pemeriksaan element_size() kemudian menegaskan konsekuensinya pada memori, yaitu dua byte per elemen untuk bfloat16 dibandingkan empat byte untuk float32. Perilaku ini berlaku sama pada CUDA, dengan float16 menggantikan bfloat16 sebagai dtype pilihan.

GradScaler dan loss scaling#

Pada float16, gradien bernilai kecil mudah jatuh ke nol karena rentang formatnya sempit. Solusi PyTorch adalah loss scaling, yaitu mengalikan loss dengan faktor skala $s$ sebelum backward sehingga gradien ikut membesar, lalu mengurungkan faktor itu sebelum optimizer memperbarui bobot:

GradScaler mengurus ketiga langkah tersebut, yaitu mengalikan loss, mengurungkan skala pada gradien, dan menyesuaikan nilai $s$. Skala dinaikkan bila tidak ada gradien inf atau nan dan diturunkan bila ada. Alur lengkap satu langkah pelatihan dengan GradScaler ditunjukkan pada diagram berikut.

Pada GPU dengan float16, pola lengkapnya ditulis sebagai berikut.

# tidak-dijalankan: sandbox tidak memiliki GPU CUDA
scaler = torch.amp.GradScaler("cuda")

for x_batch, y_batch in loader:
    optimizer.zero_grad()
    with torch.autocast("cuda", dtype=torch.float16):
        y_pred = model(x_batch)           # operasi linear berjalan di fp16
        loss = criterion(y_pred, y_batch)
    scaler.scale(loss).backward()         # gradien ikut terdiskala
    scaler.step(optimizer)                # skala diurungkan sebelum step
    scaler.update()                       # skala s diperbarui

Contoh berikut menjalankan pola yang sama pada CPU dengan bfloat16, sehingga dapat dieksekusi apa adanya. Strukturnya identik dengan pola CUDA di atas, dan yang berbeda hanya pemilihan device serta dtype.

import torch
import torch.nn as nn

torch.manual_seed(0)

# data sintetis: 64 sampel, 10 fitur, target regresi satu nilai
x = torch.randn(64, 10)                       # shape (64, 10)
w_asli = torch.randn(10, 1)
y = x @ w_asli + 0.1 * torch.randn(64, 1)     # shape (64, 1)

model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scaler = torch.amp.GradScaler("cpu")

# periksa dtype hasil forward di dalam konteks autocast
with torch.autocast("cpu", dtype=torch.bfloat16):
    y_pred = model(x)                         # shape (64, 1)
    loss = criterion(y_pred, y)
print("dtype y_pred:", y_pred.dtype, "| dtype loss:", loss.dtype)

for epoch in range(5):
    optimizer.zero_grad()
    with torch.autocast("cpu", dtype=torch.bfloat16):
        y_pred = model(x)                     # shape (64, 1)
        loss = criterion(y_pred, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    print(f"epoch {epoch + 1}: loss {loss.item():.4f}")

print("Skala GradScaler (CPU):", scaler.get_scale())

dtype y_pred: torch.bfloat16 | dtype loss: torch.float32
epoch 1: loss 5.1801
epoch 2: loss 4.8980
epoch 3: loss 4.6345
epoch 4: loss 4.3843
epoch 5: loss 4.1315
Skala GradScaler (CPU): 65536.0

Baris pertama keluaran memperlihatkan sifat khas mixed precision, yaitu dtype campuran dalam satu langkah, dengan prediksi pada bfloat16 dan loss tetap float32. Loss turun mulus selama lima epoch, sehingga presisi rendah tidak merusak konvergensi pada model kecil ini. Nilai skala yang tetap 65536.0 juga sesuai teori, karena bfloat16 tidak memerlukan loss scaling, sehingga GradScaler pada CPU terutama menyiapkan API yang sama dengan CUDA.

Kapan mixed precision penting#

Mixed precision memberi keuntungan nyata pada tiga kondisi, yaitu pelatihan pada GPU modern, model dengan banyak parameter, dan batch yang besar. Pada kondisi tersebut, penghematan memori memungkinkan batch lebih besar atau model lebih besar pada perangkat keras yang sama, sedangkan percepatan datang dari unit komputasi presisi rendah pada GPU. Sebaliknya, pada model kecil yang dilatih pada CPU sebagaimana contoh jalur ini, mixed precision terutama menjadi latihan API, karena prosesor belum tentu menjalankan bfloat16 lebih cepat daripada float32. Kualitas konvergensi juga bukan tujuannya, karena hasil akhir keduanya praktis setara.

Keputusan praktisnya mengikuti perangkat keras dan ukuran pelatihan. Pelatihan pada GPU dengan model yang cukup besar untuk memenuhi memori menjadikan autocast pilihan default, dan pilihannya makin aman bila dtype yang tersedia adalah bfloat16 yang tidak menuntut loss scaling. Pelatihan pada CPU untuk pembelajaran tetap dapat menggunakan jalur yang sama, tetapi penilaian efisiensinya sebaiknya menunggu perangkat keras yang memang dimaksudkan.

Kesalahan umum#

Tiga kesalahan berikut paling sering muncul saat mixed precision dipraktikkan pertama kali.

Yang pertama adalah memanggil optimizer.step() langsung setelah scaler.scale(loss).backward(). Gradien pada .grad masih membawa faktor skala, sehingga bobot melompat jauh dari posisinya. Perbaikannya, pasangkan scaler.step(optimizer) dengan scaler.update() dan jangan memanggil optimizer.step() secara terpisah.

Yang kedua adalah tetap menggunakan API lama torch.cuda.amp.autocast() dan torch.cuda.amp.GradScaler(). API tersebut sudah usang (deprecated) pada PyTorch 2.x, sedangkan penggantinya torch.amp.autocast dan torch.amp.GradScaler menerima device secara eksplisit. Perbaikannya, pindahkan seluruh pemanggilan ke awalan torch.amp sebagaimana contoh topik ini.

Yang ketiga adalah mengharapkan percepatan pada CPU. Jalur autocast CPU digunakan di sini agar contohnya dapat dieksekusi, sedangkan keuntungan kecepatannya baru muncul pada GPU. Perbaikannya, sesuaikan ekspektasi dengan mengukur waktu pelatihan sebelum dan sesudah autocast pada perangkat yang benar-benar digunakan.

Ringkasan#

• torch.autocast memilih presisi per operasi di dalam bloknya, sedangkan bobot model dan pembaruannya tetap berada pada float32.
• GradScaler mengalikan loss dengan faktor $s$ sebelum backward, mengurungkannya sebelum step, dan menyesuaikan $s$ secara otomatis. Loss scaling diperlukan untuk float16, bukan untuk bfloat16.
• Keuntungan mixed precision baru terasa pada GPU, model besar, dan batch besar. Pada CPU, contoh mixed precision dengan bfloat16 berfungsi sebagai latihan API.

Referensi#

• Automatic Mixed Precision package — PyTorch — Dokumentasi resmi torch.autocast dan torch.amp.GradScaler beserta daftar perilaku per operasi.
• Automatic Mixed Precision recipe — PyTorch — Resep resmi penggunaan autocast dan GradScaler di dalam perulangan pelatihan.