Model Parallelism

Data parallelism menggandakan model ke setiap perangkat — strategi yang langsung kandas bila modelnya sendiri tidak muat di memori satu perangkat. Model bahasa berskor puluhan miliar parameter, ditambah state optimizer 16 byte per parameter (lihat chapter Keterbatasan Memori GPU ), memang melebihi kapasitas GPU tunggal mana pun. Model parallelism mengambil jalan sebaliknya dari DP: datanya tetap utuh, tetapi lapisan model dipecah ke beberapa perangkat — tahap awal di satu device, tahap akhir di device lain, dan aktivasi mengalir di antara keduanya. Chapter ini membedah penempatan manual per tahap, membuktikan aliran forward-backward melintasi perbatasan device lewat simulasi dua “perangkat” di CPU, lalu membandingkan posisinya terhadap data parallelism.

Pecah Lapisan, Alirkan Aktivasi #

Ide dasarnya sederhana: device cpu:0 membawa lapisan 1–4, device cpu:1 membawa lapisan 5–7. Forward pass berjalan seperti relay — tensor masuk diproses tahap awal di perangkat pertama, hasilnya dikirim melintasi perbatasan (di GPU nyata: h.to("cuda:1")), lalu diproses tahap akhir di perangkat kedua. Backward menjalankan relay yang sama secara mundur: gradien terhitung di perangkat kedua, dikirim balik melintasi perbatasan, lalu meneruskan perhitungan gradien di perangkat pertama. Berkas model_parallel.py menyimulasikan keseluruhan aliran itu dengan dua tahap sungguhan di CPU:

"""Model parallelism simulasi dua perangkat: layer dipecah, aktivasi mengalir."""

import torch

torch.manual_seed(21)

D_IN, D_H, D_OUT, B = 16, 96, 4, 32

# ---------- Model utuh sebagai pembanding ----------
model_utuh = torch.nn.Sequential(
    torch.nn.Linear(D_IN, D_H),
    torch.nn.GELU(),
    torch.nn.Linear(D_H, D_H),
    torch.nn.GELU(),
    torch.nn.Linear(D_H, D_H),
    torch.nn.GELU(),
    torch.nn.Linear(D_H, D_OUT),
)

# ---------- Pecah lapisan ke dua "perangkat" ----------
tahap0 = torch.nn.Sequential(*(model_utuh[i] for i in range(4)))  # cpu:0
tahap1 = torch.nn.Sequential(*(model_utuh[i] for i in range(4, 7)))  # cpu:1


def byte_param(modul):
    return sum(p.numel() * p.element_size() for p in modul.parameters())


p_utuh = sum(p.numel() for p in model_utuh.parameters())
print("=== Pembagian parameter ke dua perangkat ===")
print(f"model utuh      : {p_utuh:,} parameter = {byte_param(model_utuh):,} byte")
print(f"device cpu:0    : tahap0 = {byte_param(tahap0):,} byte")
print(f"device cpu:1    : tahap1 = {byte_param(tahap1):,} byte")
print(f"puncak per device: {max(byte_param(tahap0), byte_param(tahap1)):,} byte "
      f"({byte_param(model_utuh) / max(byte_param(tahap0), byte_param(tahap1)):.1f}x lebih kecil)")


# ---------- Forward chain manual antar device ----------
def pindah(h, asal, tujuan):
    byte = h.numel() * h.element_size()
    print(f"  transfer aktivasi {tuple(h.shape)} {asal} -> {tujuan}: {byte:,} byte")
    return h  # di GPU nyata: h.to(torch.device(tujuan))


x = torch.randn(B, D_IN)
print()
print("=== Alur forward antar device ===")
print(f"x masuk di cpu:0, shape {tuple(x.shape)}")
h = tahap0(x)  # seluruh lapisan tahap0 berjalan di cpu:0
h1 = pindah(h, "cpu:0", "cpu:1")
out = tahap1(h1)  # lapisan tahap1 berjalan di cpu:1
print(f"keluaran akhir di cpu:1, shape {tuple(out.shape)}")
print(f"identik dengan model utuh: {torch.equal(out, model_utuh(x))}")

# ---------- Backward: gradien mengalir balik melintasi perbatasan ----------
target = torch.randn(B, D_OUT)
loss = ((out - target) ** 2).mean()
loss.backward()
n0 = sum(float(p.grad.norm()) for p in tahap0.parameters() if p.grad is not None)
n1 = sum(float(p.grad.norm()) for p in tahap1.parameters() if p.grad is not None)
print()
print("=== Backward melintasi perbatasan device ===")
print(f"norma gradien tahap0 (cpu:0): {n0:.4f} (bukan 0: gradien sampai)")
print(f"norma gradien tahap1 (cpu:1): {n1:.4f}")

Output yang Dihasilkan:

=== Pembagian parameter ke dua perangkat ===
model utuh      : 20,644 parameter = 82,576 byte
device cpu:0    : tahap0 = 43,776 byte
device cpu:1    : tahap1 = 38,800 byte
puncak per device: 43,776 byte (1.9x lebih kecil)

=== Alur forward antar device ===
x masuk di cpu:0, shape (32, 16)
  transfer aktivasi (32, 96) cpu:0 -> cpu:1: 12,288 byte
keluaran akhir di cpu:1, shape (32, 4)
identik dengan model utuh: True

=== Backward melintasi perbatasan device ===
norma gradien tahap0 (cpu:0): 0.1142 (bukan 0: gradien sampai)
norma gradien tahap1 (cpu:1): 0.5349

Mari kita telaah kode di atas baris per baris:

1. Baris 21-22 (pemecahan tahap): model_utuh[i] for i in range(4) merujuk module yang sama dengan model utuh — bukan salinan — sehingga kedua tahap benar-benar memecah satu model, bukan membangun model baru. Tahap0 membawa dua blok Linear+GELU pertama, tahap1 membawa sisanya.
2. Baris 29-35 (pembagian memori): dari 82.576 byte parameter, cpu:0 membawa 43.776 byte dan cpu:1 membawa 38.800 byte — puncak per perangkat hanya 1,9x lebih kecil dari model utuh, karena pemecahan menimbang kapasitas, bukan sekadar membelah dua. Di GPU nyata angka ini dikalikan state optimizer: memori per perangkat turun karena 16 × P hanya ditanggung oleh pecahannya masing-masing.
3. Baris 39-42 (pindah): fungsi mencatat transfer aktivasi antar device — 12.288 byte untuk tensor (32, 96) — dan di GPU nyata baris kuncinya adalah h.to(torch.device("cuda:1")); transfer ini adalah biaya komunikasi tetap yang tidak ada pada data parallelism.
4. Baris 49-53 (forward relay): h = tahap0(x) berjalan penuh di cpu:0, hasilnya dipindah, lalu out = tahap1(h1) berjalan di cpu:1. torch.equal membuktikan keluaran identik bit-per-bit dengan model utuh — pemecahan tidak mengubah hasil, hanya lokasi hitung.
5. Baris 56-64 (backward relay): setelah loss.backward(), kedua tahap memiliki gradien bukan nol (0,1142 dan 0,5349). Autograd melacak graph yang sama melintasi perbatasan — di GPU nyata, gradien dari cuda:1 dikirim balik ke cuda:0 lewat komunikasi antar device sebelum tahap0 dapat menghitung gradiennya.

DP vs MP: Dua Sisi Berlawanan #

Tabel terakhir menyingkap kelemahan struktural model parallelism murni: pada satu saat hanya satu perangkat yang bekerja — saat tahap0 memproses batch, cpu:1 menunggu, dan sebaliknya. Perangkat kedua dalam contoh kita menganggur separuh waktu. DP dan MP juga sering dipakai bersama pada sistem nyata (data dipecah beberapa kelompok, model dipecah dalam tiap kelompok).

Kelemahan “perangkat menganggur bergantian” inilah yang dijawab teknik lanjutan: bila alih-alih menunggu satu batch penuh selesai, kita mengalirkan banyak micro-batch secara bertingkat lewat pipa, kedua perangkat bisa sibuk hampir sepanjang waktu. Chapter berikutnya membahas Pipeline Parallelism — lengkap dengan diagram gelembung (bubble) dan hitungan efisiensinya.

Rujukan resmi pola penempatan model multi-perangkat tersedia di docs.pytorch.org/docs/stable/notes/cuda.html .