Belajar Deep Learning - Training di GPU & Distributed Training

#Pendahuluan

Di episode 12 kita membangun transformer kecil yang berjalan di satu GPU (atau CPU). Sekarang pertanyaannya: bagaimana melatih model yang lebih besar dan lebih cepat? Jawabannya bukan "beli GPU lebih besar", melainkan memanfaatkan yang sudah ada secara optimal — dan jika masih kurang, membagi beban ke banyak GPU.

Ini keterampilan yang menentukan karier di 2026: tim riset melatih model 7B-70B parameter dengan multi-GPU, dan ML Engineer yang paham mixed precision serta memori GPU selalu lebih bernilai. Episode ini membekali kalian dari optimasi satu GPU sampai distributed training — konsep yang sama dipakai Hugging Face dan semua lab besar.

#Mixed Precision: Lebih Cepat dan Hemat Memori

#FP32 vs FP16 vs BF16

Secara default PyTorch memakai FP32 (32-bit float). GPU modern (Volta+) jauh lebih cepat menghitung FP16 (16-bit) — hampir 2x throughput di banyak kasus — dan memakai setengah memori. Masalahnya FP16 punya jangkauan kecil: angka sangat kecil menjadi nol, angka besar menjadi infinity.

BF16 (bfloat16) memecah masalah itu: rentang eksponen sama dengan FP32, hanya presisi mantissa yang dikurangi. Untuk training, BF16 hampir selalu lebih stabil daripada FP16. GPU NVIDIA Ampere (A100/A10/L4) ke atas mendukung BF16; GPU lebih lama hanya FP16.

#Automatic Mixed Precision (AMP)

AMP menyimpan bobot master di FP32 sementara, melakukan forward/backward dalam FP16/BF16, dan menaikkan kembali gradien ke FP32 saat update:

Mixed precision dengan torch.amp

from torch.amp import autocast, GradScaler
 
scaler = GradScaler("cuda", enabled=True)
 
for x, y in loader:
    optimizer.zero_grad()
    with autocast("cuda", dtype=torch.bfloat16):
        loss = loss_fn(model(x), y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Aturan pakai: BF16 cukup tanpa GradScaler (tidak ada underflow); FP16 butuh GradScaler untuk menangani gradien kecil. Hasilnya: training ~1.5-2x lebih cepat dengan kualitas hampir identik.

#Optimasi Memori

#Gradient Accumulation

Jika batch ideal 64 tapi GPU hanya muat batch 32, jangan langsung mengecilkan batch (kualitas training berubah). Gradient accumulation menjumlahkan gradien dari beberapa batch kecil sebelum update:

Gradient accumulation

accum_steps = 2  # simulasi batch 64 dari batch 32
 
for i, (x, y) in enumerate(loader):
    loss = loss_fn(model(x), y) / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Bagi loss dengan accum_steps agar skala gradien konsisten. Efek sampingnya: gradien lebih bising karena update lebih jarang; akurasi biasanya tetap baik.

#Strategi Memori Lain

• Gradient checkpointing: membuang aktivasi antara dan menghitung ulang saat backward — hemat memori besar dengan biaya komputasi ekstra.
• Gradient clipping (episode 6) juga mengurangi kebutuhan memori buffer.
• Batch size kecil + lebih banyak iterasi: opsi terakhir karena mengubah dinamika training.

#Distributed Training

#DataParallel (DP) — Mulai di Sini, Tapi Hati-hati

DP membagi batch ke semua GPU, lalu replikasi model di tiap GPU dengan sinkronisasi gradien lewat satu thread. Mudah dipakai, tetapi scaling-nya buruk (prosesor utama jadi bottleneck) dan tidak efisien — cukup untuk belajar konsep.

DataParallel (tidak disarankan produksi)

model = nn.DataParallel(model)  # satu baris, tapi lambat di banyak GPU

#DistributedDataParallel (DDP) — Standar

DDP mereplikasi model dan data di tiap proses (satu proses per GPU). Tiap proses training dengan batch-nya sendiri, lalu gradien disinkronkan antar proses (all-reduce) di setiap step. Skalanya jauh lebih baik daripada DP. Alur standar dengan torchrun:

Init DDP

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
 
dist.init_process_group("nccl")   # nccl untuk NVIDIA GPU
model = DDP(model, device_ids=[local_rank])
 
for x, y in loader:
    loss = loss_fn(model(x), y)   # DataLoader sudah samar
    loss.backward()               # gradien auto all-reduce antar GPU
    optimizer.step()

Jalankan dengan:

torchrun --nproc_per_node=4 train.py

Detail krusial: DistributedSampler memastikan tiap proses melihat subset data yang berbeda; sampler.set_epoch(epoch) wajib dipanggil tiap epoch agar shuffling antar epoch tetap acak dan tidak bertabrakan.

#FSDP — Sharding Parameter

DDP masih menuntut seluruh model muat di tiap GPU. Untuk model besar, FSDP (Fully Sharded Data Parallel) membagi parameter, gradien, dan optimizer state ke semua GPU — tiap GPU hanya menyimpan bagiannya. Mirip ZeRO dari DeepSpeed. Inilah cara melatih model 7B+ di beberapa GPU:

FSDP

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
 
model = FSDP(model, auto_wrap_policy=size_based_auto_wrap_policy(min_num_params=1e6))

FSDP menambah komunikasi antar GPU (parameter di-gather saat forward), sehingga untuk model kecil justru bisa lebih lambat dari DDP. Aturan praktis: DDP untuk model yang muat di satu GPU, FSDP/ZeRO untuk yang tidak muat.

#Common Pitfalls

• Seed tidak sama antar proses: inisialisasi acak berbeda di tiap GPU merusak konsistensi. Set seed + dist.barrier() setelah setup.
• Lupa set_epoch: DistributedSampler tanpa set epoch menghasilkan urutan data sama tiap epoch.
• BatchNorm di DDP: statistik running-nya tidak sinkron antar proses secara default. Pertimbangkan SyncBatchNorm jika batch per GPU kecil.
• Mixed precision di CPU: autocast("cuda") di mesin tanpa GPU error — guard dengan torch.cuda.is_available().

#Penutup

Pada episode 13 ini, kalian telah menguasai training skala.

Inti yang harus dibawa pulang:

• Mixed precision (BF16/FP16 + AMP) memberi 1.5-2x percepatan dengan kualitas hampir sama.
• Gradient accumulation mensimulasikan batch besar saat VRAM terbatas.
• DDP = replikasi data & sinkronisasi gradien; standar untuk model muat di satu GPU.
• FSDP/ZeRO = sharding parameter untuk model yang melebihi satu GPU.
• Optimalkan satu GPU dulu sebelum pindah ke distributed.

Di episode 14 selanjutnya kita akan membuat training bisa dilacak dan dituning: experiment tracking & hyperparameter tuning — TensorBoard, Weights & Biases, MLflow, Optuna, dan learning rate finder. Sampai jumpa di episode 14!