Mixed Precision Training

"""Mixed precision di CPU: autocast bfloat16 dan mekanik GradScaler.

Bagian 1: anatomi tiga dtype lewat torch.finfo (presisi vs rentang).
Bagian 2: dua mode gagal float16 — luapan (inf) dan nilai kecil membulat
ke nol — dibandingkan dengan bfloat16 yang rentangnya selaras float32.
Bagian 3: forward pass di dalam torch.autocast("cpu", bfloat16) —
bobot tetap float32, keluaran berubah dtype, hasil sedikit bergeser.
Bagian 4: training mini float32 vs autocast bfloat16 pada masalah yang
sama dengan bobot awal identik — loss berbeda tipis, keduanya konvergen.
Bagian 5: siklus GradScaler scale -> unscale_ -> step -> update.
"""

import torch
import torch.nn as nn
from sklearn.datasets import make_moons

torch.manual_seed(42)

# ---------- 1. Anatomi tiga dtype ----------
print(f"{'dtype':16s} {'bit':>4s} {'resolusi':>10s} {'min':>13s} {'maks':>13s}")
for dt in (torch.float32, torch.float16, torch.bfloat16):
    info = torch.finfo(dt)
    print(f"{str(dt):16s} {info.bits:4d} {info.resolution:>10.0e} "
          f"{info.min:>13.3e} {info.max:>13.3e}")

# ---------- 2. Dua mode gagal float16 ----------
a = torch.tensor([300.0], dtype=torch.float16)
b = torch.tensor([300.0], dtype=torch.bfloat16)
print()
print("300.0 x 300.0 : float16 =", (a * 300.0).item(),
      "| bfloat16 =", (b * 300.0).item())

kecil = torch.tensor([1e-8])
print("gradien 1e-8  : float16 =", f"{kecil.to(torch.float16).item():.2e}",
      "| bfloat16 =", f"{kecil.to(torch.bfloat16).item():.2e}",
      "| float32 =", f"{kecil.item():.2e}")
print("1e-8 x 65536  : float16 =",
      f"{(kecil * 65536.0).to(torch.float16).item():.2e}")

# ---------- 3. autocast CPU: tipe berubah, hasil bergeser tipis ----------
model = nn.Linear(64, 32)
x = torch.randn(8, 64)
with torch.autocast("cpu", dtype=torch.bfloat16):
    y_auto = model(x)
y_biasa = model(x)
print()
print("bobot model      :", model.weight.dtype)
print("keluaran autocast:", y_auto.dtype)
print("keluaran biasa   :", y_biasa.dtype)
print("selisih maksimum :", (y_auto.float() - y_biasa).abs().max().item())

# ---------- 4. Training mini: float32 vs autocast bfloat16 ----------
X, y = make_moons(n_samples=300, noise=0.25, random_state=42)
X_t = torch.tensor(X, dtype=torch.float32)
y_t = torch.tensor(y, dtype=torch.long)


def buat_model():
    torch.manual_seed(123)  # bobot awal identik untuk kedua percobaan
    return nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2))


def latih(pakai_autocast, judul):
    model = buat_model()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
    criterion = nn.CrossEntropyLoss()
    riwayat = {}
    for epoch in range(1, 31):
        optimizer.zero_grad()
        if pakai_autocast:
            with torch.autocast("cpu", dtype=torch.bfloat16):
                logits = model(X_t)
                loss = criterion(logits, y_t)
        else:
            logits = model(X_t)
            loss = criterion(logits, y_t)
        loss.backward()
        optimizer.step()
        if epoch in (1, 10, 30):
            riwayat[epoch] = loss.item()
    with torch.no_grad():
        akurasi = (model(X_t).argmax(dim=1) == y_t).float().mean().item()
    print(f"{judul:26s} ep1 {riwayat[1]:7.4f} | ep10 {riwayat[10]:7.4f} | "
          f"ep30 {riwayat[30]:7.4f} | akurasi {akurasi:.3f}")
    return riwayat


print()
print("training: make_moons 300 | full-batch | SGD lr=0.5 | 30 epoch")
print("logits autocast keluar sebagai", end=" ")
model_coba = buat_model()
with torch.autocast("cpu", dtype=torch.bfloat16):
    print(model_coba(X_t).dtype)
r32 = latih(False, "float32 penuh")
rbf = latih(True, "autocast bfloat16")
print("loss ep30 berbeda?", abs(r32[30] - rbf[30]) > 0,
      f"(selisih {abs(r32[30] - rbf[30]):.4f})")

# ---------- 5. Siklus GradScaler (pola float16 di GPU) ----------
print()
scaler = torch.amp.GradScaler("cpu")
w = torch.nn.Parameter(torch.tensor([0.001]))
optimizer = torch.optim.SGD([w], lr=0.1)
loss = (w * 1e-8).sum()          # gradien analitis = 1e-8
scaler.scale(loss).backward()    # loss diskalakan lalu backward
print("gradien terskala  :", f"{w.grad.item():.2e}")
scaler.unscale_(optimizer)       # kembalikan skala gradien
print("gradien unscale_  :", f"{w.grad.item():.2e}")
scaler.step(optimizer)           # update hanya bila aman
scaler.update()                  # sesuaikan skala utk langkah berikutnya
print("w setelah step    :", f"{w.item():.12f}",
      "(0.001 - 0.1 x 1e-8)")
print("skala berikutnya  :", scaler.get_scale())

dtype             bit   resolusi           min          maks
torch.float32      32      1e-06    -3.403e+38     3.403e+38
torch.float16      16      1e-03    -6.550e+04     6.550e+04
torch.bfloat16     16      1e-02    -3.390e+38     3.390e+38

300.0 x 300.0 : float16 = inf | bfloat16 = 90112.0
gradien 1e-8  : float16 = 0.00e+00 | bfloat16 = 1.00e-08 | float32 = 1.00e-08
1e-8 x 65536  : float16 = 6.55e-04

bobot model      : torch.float32
keluaran autocast: torch.bfloat16
keluaran biasa   : torch.float32
selisih maksimum : 0.004613995552062988

training: make_moons 300 | full-batch | SGD lr=0.5 | 30 epoch
logits autocast keluar sebagai torch.bfloat16
float32 penuh              ep1  0.5467 | ep10  0.3238 | ep30  0.3089 | akurasi 0.863
autocast bfloat16          ep1  0.5465 | ep10  0.3237 | ep30  0.3087 | akurasi 0.863
loss ep30 berbeda? True (selisih 0.0001)

gradien terskala  : 6.55e-04
gradien unscale_  : 1.00e-08
w setelah step    : 0.000999999000 (0.001 - 0.1 x 1e-8)
skala berikutnya  : 65536.0