Belajar ML Engineer - Distributed Training

#Pendahuluan

Setelah 20 episode membangun sistem ML yang lengkap, ada batas yang cepat atau lambat kalian temukan: satu GPU tidak cukup. Entah karena dataset terlalu besar (training terlalu lama), atau model terlalu besar (tidak muat di VRAM), atau keduanya. Jawabannya adalah distributed training — menyebar pekerjaan ke banyak GPU, baik dalam satu mesin maupun lintas mesin.

Distributed training adalah area yang menakutkan bagi banyak engineer: istilah-istilahnya membingungkan dan kesalahannya mahal (GPU idle = uang terbakar). Episode ini membedah strategi paralelisme satu per satu, kapan memakai yang mana, dan bagaimana memulainya dengan tool standar PyTorch, Ray, dan TorchX.

#Dua Masalah, Dua Strategi

Distributed training menjawab dua masalah berbeda dengan dua strategi berbeda:

Masalah	Strategi	Analogi
Data terlalu besar / training lambat	Data parallelism	4 koki mengerjakan 4 bagian adonan berbeda
Model tidak muat di satu GPU	Model parallelism	Satu koki memegang resep, tugas berat dipecah

#Data Parallelism

Setiap GPU memegang salinan model lengkap, menerima subset data yang berbeda, menghitung gradien, lalu menyinkronkan gradien antar GPU. Lebih banyak GPU = throughput lebih tinggi, selama komunikasi tidak menjadi bottleneck.

Implementasi modern: DDP (DistributedDataParallel). Ia lebih baik dari DataParallel lama karena memakai komunikasi point-to-point (NCCL) per layer, bukan broadcast seluruh model di tiap iterasi:

DDP: data parallelism

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel
 
def train_worker(rank: int, world_size: int):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    model = DistributedDataParallel(build_model().to(rank), device_ids=[rank])
    loader = make_loader(rank, world_size)   # tiap rank dapat subset data
    # training loop seperti biasa; gradien disinkronkan otomatis

#Model Parallelism

Saat model tidak muat di satu GPU, model dipecah. Dua varian:

Pipeline parallelism: layer-layer dibagi antar GPU; data mengalir seperti pipa (GPU 1 → 2 → 3). Latensi naik karena lapisan berurutan, tapi throughput tetap dapat diskalakan.
Tensor parallelism: satu layer dipecah melintasi beberapa GPU — setiap GPU menghitung sebagian dari layer tersebut, lalu hasilnya digabung. Butuh komunikasi sangat cepat per operasi (NCCL di node yang sama); ini teknik yang dipakai untuk melatih LLM besar.

Untuk LLM, kombinasi ketiganya adalah cara kerja industri: tensor parallelism di dalam node (NVLink cepat), pipeline parallelism antar node, data parallelism di atas semuanya. Untungnya, kalian jarang menulis ini manual — framework menanganinya.

#Kapan Memakai Strategi Apa: Decision Guide

Situasi	Strategi Awal
Model muat di 1 GPU, training lambat	Data parallelism (DDP)
Model muat di 1 GPU, batch besar tidak muat	Gradient accumulation
Model muat di 1 node, butuh throughput	DDP, lalu tuning batch
Model tidak muat 1 GPU	Tensor/pipeline parallelism
8 GPU, model besar	Tensor (intra-node) + data
Lintas node	Pipeline + data parallelism

Aturan penting: jangan menambah kompleksitas sebelum mengukur. Kurva scaling jarang linier — cek GPU utilization; jika sudah 95%+ di 1 GPU dan komunikasi tidak dominan, DDP baru memberi manfaat nyata.

#Gradient Accumulation: Tip Tanpa Distributed

Sebelum belanja GPU, ada teknik yang sering cukup: gradient accumulation — mensimulasikan batch besar dengan menjumlahkan gradien beberapa batch kecil. Berguna saat model muat tapi batch yang dibutuhkan lebih besar dari memori:

Gradient accumulation

accum_steps = 4
optimizer.zero_grad()
for i, (xb, yb) in enumerate(loader):
    loss = loss_fn(model(xb), yb) / accum_steps   # normalisasi
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

#Ray: Komputasi Terdistribusi yang Serba Bisa

Ray adalah framework komputasi terdistribusi yang lebih luas dari sekadar training — ia mengelola cluster, menjadwalkan task, dan menyediakan Ray Train untuk distributed training dengan API tinggi:

Ray Train dengan TorchTrainer

from ray.train.torch import TorchTrainer, TorchConfig
from ray.train import ScalingConfig
 
 
def train_func(config):
    import torch.distributed as dist
    # kode training biasa, Ray menangani device & DDP setup
    ...
 
trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
    torch_config=TorchConfig(backend="nccl"),
)
result = trainer.fit()

Keunggulan Ray: satu ekosistem untuk training (Ray Train), hyperparameter search (Ray Tune), serving (Ray Serve), dan data pipeline — mempermudah operasional dibanding menumpuk banyak framework.

#TorchX dan Teknologi Pengganti

TorchX (dari Meta) menyediakan aplikator yang menjalankan job training PyTorch di berbagai backend (local, kubernetes, slurm) dengan API yang konsisten — model JSON app yang bisa dijadwalkan:

TorchX: definisikan app & jadwalkan

from torchx.specs import AppDef, Replicas, Driver, Resource
import torchx
 
app = AppDef(
    name="train-job",
    roles=[
        Replicas(
            role="trainer", num_replicas=4,
            image="registry.example/ml-train:1.4.2",
            resource=Resource(cpu=8, gpu=4, memMB=65536),
            entrypoint="python -m src.train --ddp",
        )
    ],
)
torchx.run(app, scheduler="kubernetes")

Dunia distributed training 2026 punya banyak lapisan: training libraries (PyTorch DDP, FSDP, DeepSpeed), clusters (Ray, TorchX, Horovod), dan platform (Kubeflow, SageMaker, Vertex AI). Sebagai MLE, kuasai satu jalur sampai nyaman — DDP + Ray adalah titik masuk yang paling rasional.

#Framework Stateful: FSDP

Untuk model besar yang masih muat di satu node, FSDP (Fully Sharded Data Parallelism) adalah pendekatan modern: parameter, gradien, dan optimizer state di-shard antar GPU, sehingga ukuran model yang bisa dilatih jauh lebih besar dari satu GPU — sambil tetap memakai paradigma data parallelism yang sederhana:

FSDP (ringkas)

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
model = FSDP(build_model())

FSDP adalah teknik di balik fine-tune model 7B+ di satu node dengan beberapa GPU. Ia evolusi dari DDP: bukan hanya data yang dibagi, tapi juga state model.

Tip

Aturan emas distributed training: selalu uji scaling di 2 GPU dulu sebelum membeli 8. Kebanyakan masalah (komunikasi tidak efisien, batch terlalu kecil per GPU, bottleneck data loading) muncul sudah jelas di 2 GPU — dan jauh lebih murah untuk diperbaiki.

#Common Pitfalls

Pitfall	Dampak	Pencegahan
GPU idle karena data loading lambat	Utilisasi rendah, uang terbuang	`num_workers` & prefetch
Batch per GPU terlalu kecil	Komunikasi mendominasi, scaling buruk	Perbesar batch per GPU / grad accum
DDP tanpa seeded sharding	Data duplikat antar rank	`DistributedSampler`
Komunikasi lintas node lambat	Tensor parallelism melambat	Tensor paralel dalam node saja
Langsung menambah GPU tanpa ukur	Scaling sub-linear	Cek util GPU & scaling di 2 GPU

#Penutup

Pada episode 21 ini, kalian telah menguasai training terdistribusi:

Data parallelism (DDP) membagi data, sinkronkan gradien — untuk training yang lambat.
Model parallelism (pipeline & tensor) membagi model — untuk model yang tidak muat.
Gradient accumulation sebagai solusi murah sebelum belanja GPU.
Ray (Train + Tune + Serve) dan TorchX menyatukan orkestrasi distributed.
FSDP membuka training model besar di satu node.

Di episode 22 selanjutnya kita akan membahas Serving at Scale — Kubernetes, autoscaling inference, dan multi-tenant serving: cara menyajikan model ke banyak pengguna dengan elastisitas dan isolasi yang benar. Sampai jumpa di episode 22!