Parameter Sharding

Data parallelism menggandakan model ke setiap perangkat — dan menggandakan pula komponen yang paling rakus memori: state optimizer Adam sebesar 8 byte per parameter harus ada lengkap di setiap salinan. Pada model 1 miliar parameter dengan mixed precision, tiap perangkat membawa sekitar 14,9 GB state padahal perangkat lain memegang salinan yang identik. Parameter sharding — dipopulerkan keluarga ZeRO (Zero Redundancy Optimizer) — menghapus duplikasi itu: komponen yang tadinya digandakan penuh di setiap perangkat dipecah merata, lalu dirakit kembali lewat komunikasi hanya pada saat dibutuhkan. Chapter ini menghitung penghematan setiap stage ZeRO dengan aritmetika yang ditampilkan penuh, menyimulasikan sharding tensor sungguhan di CPU, lalu menempatkan FSDP PyTorch sebagai realisasi API-nya.

Tiga Stage ZeRO: Apa yang Dipecah #

ZeRO bekerja berlapis. Stage 1 hanya memecah state optimizer (8P untuk Adam), stage 2 menambahkan gradien (2P pada mixed precision), dan stage 3 melengkapinya dengan parameter serta master weight. Untuk model 1 miliar parameter float16 + Adam float32 pada 4 perangkat, hitungannya sebagai berikut:

Angka 16P di baseline berasal dari lima komponen: parameter fp16 (2P), master weight fp32 (4P, tempat pembaruan presisi penuh — konsepnya dibahas di chapter Mixed Precision Training ), gradien fp16 (2P), serta m dan v milik Adam (masing-masing 4P). Berkas zero_sharding.py menghitung seluruh tabel itu secara nyata, lalu menyimulasikan shard dan all-gather pada tensor sungguhan:

"""Parameter sharding ala ZeRO: aritmetika memori + simulasi shard 4 device."""

import torch

torch.manual_seed(64)

# ---------- 1. Aritmetika memori ZeRO untuk model 1 miliar parameter ----------
P = 1_000_000_000
ND = 4  # jumlah device untuk sharding

print("=== Model 1 miliar parameter, mixed precision fp16 + Adam fp32 ===")
komponen = [
    ("parameter fp16", 2),
    ("master weight fp32", 4),
    ("gradien fp16", 2),
    ("Adam m (fp32)", 4),
    ("Adam v (fp32)", 4),
]
print(f"{'Komponen':22s} {'Rumus':>8s} {'Total (GB)':>12s}")
for nama, faktor in komponen:
    print(f"{nama:22s} {faktor}P {faktor * P / 1024**3:>10.1f} GB")

print()
print(f"{'Konfigurasi':38s} {'Total':>9s} {'Per device (ND={ND})':>12s}")
tahapan = [
    ("Baseline (tanpa sharding)", 16),
    ("ZeRO stage 1: optimizer", 2 + 4 + 2 + 8 / ND),
    ("ZeRO stage 2: + gradien", 2 + 4 + (2 + 8) / ND),
    ("ZeRO stage 3: + parameter", (2 + 4 + 2 + 8) / ND),
]
for nama, faktor in tahapan:
    total = faktor * P
    print(f"{nama:38s} {faktor:>6.1f}P {total / 1024**3:>8.1f} GB "
          f"{total / ND / 1024**3:>8.1f} GB")
print("hemat stage 3 vs baseline : "
      f"{(1 - (2 + 4 + 2 + 8) / ND / 16):.0%} memori state")

# ---------- 2. Simulasi shard tensor nyata ke 4 "device" ----------
print()
print("=== Simulasi sharding tensor nyata (4 device) ===")
W = torch.randn(1024, 1024, dtype=torch.float16)  # 1.048.576 elemen
byte_asli = W.numel() * W.element_size()
print(f"tensor W        : shape {tuple(W.shape)} fp16 = {byte_asli:,} byte")

shard = {}
for perangkat in range(ND):
    potongan = W.chunk(ND, dim=0)[perangkat].clone()  # salinan terpisah
    shard[perangkat] = potongan
    print(f"device {perangkat}   : shape {tuple(potongan.shape)} "
          f"= {potongan.numel() * potongan.element_size():,} byte")

# ---------- 3. All-gather: rakit kembali parameter utuh ----------
gather = torch.cat([shard[i] for i in range(ND)], dim=0)
print(f"all-gather      : shape {tuple(gather.shape)} "
      f"= {gather.numel() * gather.element_size():,} byte")
print(f"identik dengan tensor asli: {torch.equal(gather, W)}")

# ---------- 4. Sharding optimizer state: pembaruan per device ----------
print()
print("=== Pembaruan Adam ter-shard (stage 1) ===")
m = {i: torch.zeros_like(shard[i]) for i in range(ND)}
v = {i: torch.zeros_like(shard[i]) for i in range(ND)}
g = {i: torch.randn_like(shard[i]) for i in range(ND)}  # gradien ter-reduce-scatter
for i in range(ND):
    m[i].mul_(0.9).add_(g[i])
    v[i].mul_(0.999).add_(g[i] ** 2)
print(f"state m + v device 0 : {m[0].numel() * m[0].element_size() * 2:,} byte "
      f"(bukan {2 * W.numel() * 2:,} byte seperti tanpa sharding)")
print(f"rasio state optimizer per device: 1/{ND} dari total")

Output yang Dihasilkan:

=== Model 1 miliar parameter, mixed precision fp16 + Adam fp32 ===
Komponen                  Rumus   Total (GB)
parameter fp16         2P        1.9 GB
master weight fp32     4P        3.7 GB
gradien fp16           2P        1.9 GB
Adam m (fp32)          4P        3.7 GB
Adam v (fp32)          4P        3.7 GB

Konfigurasi                                Total Per device (ND={ND})
Baseline (tanpa sharding)                16.0P     14.9 GB      3.7 GB
ZeRO stage 1: optimizer                  10.0P      9.3 GB      2.3 GB
ZeRO stage 2: + gradien                   8.5P      7.9 GB      2.0 GB
ZeRO stage 3: + parameter                 4.0P      3.7 GB      0.9 GB
hemat stage 3 vs baseline : 75% memori state

=== Simulasi sharding tensor nyata (4 device) ===
tensor W        : shape (1024, 1024) fp16 = 2,097,152 byte
device 0   : shape (256, 1024) = 524,288 byte
device 1   : shape (256, 1024) = 524,288 byte
device 2   : shape (256, 1024) = 524,288 byte
device 3   : shape (256, 1024) = 524,288 byte
all-gather      : shape (1024, 1024) = 2,097,152 byte
identik dengan tensor asli: True

=== Pembaruan Adam ter-shard (stage 1) ===
state m + v device 0 : 1,048,576 byte (bukan 4,194,304 byte seperti tanpa sharding)
rasio state optimizer per device: 1/4 dari total

Mari kita telaah kode di atas baris per baris:

1. Baris 8-21 (tabel komponen): lima komponen mixed precision dihitung dari rumus ber-koefisien 2P hingga 4P; angka GB ditampilkan langsung dari faktor * P / 1024**3 — parameter fp16 1,9 GB, Adam 3,7 GB masing-masing — tanpa satu pun angka yang diketik manual.
2. Baris 25-36 (tabel stage): stage 1 memecah hanya optimizer (8/ND), stage 2 menambahkan gradien ((2+8)/ND), stage 3 memecah semuanya (16/ND). Hasilnya 14,9 GB → 9,3 GB → 7,9 GB → 3,7 GB per seluruh cluster; kolom “per device” membagi total dengan ND=4. Stage 3 memangkas 75% memori state dibanding baseline.
3. Baris 41-50 (shard tensor nyata): tensor W berukuran 1024×1024 fp16 (2.097.152 byte) dipecah lewat W.chunk(4, dim=0) — setiap “device” menyimpan salinan terpisah berukuran (256, 1024) atau 524.288 byte, tepat seperempat aslinya.
4. Baris 53-56 (all-gather): torch.cat merakit kembali keempat shard menjadi tensor (1024, 1024); torch.equal membuktikan hasilnya identik bit-per-bit — inilah operasi yang dijalankan sebelum forward pada stage 3, ketika parameter utuh dibutuhkan sesaat untuk hitungan lalu dilepas lagi.
5. Baris 61-69 (pembaruan ter-shard): masing-masing device memperbarui m dan v hanya pada bagiannya sendiri lewat operasi in-place mul_/add_ — state optimizer device 0 hanya 1.048.576 byte, seperempat dari 4.194.304 byte tanpa sharding. Gradien g diasumsikan sudah terbagi lewat reduce-scatter (jumlah gradien antar device lalu disebar balik dalam bentuk terpecah).

Harga yang Dibayar dan FSDP #

Sharding bukan diskon tanpa syarat. Stage 3 menambah komunikasi all-gather parameter di tiap forward dan reduce-scatter gradien di tiap backward — volume komunikasinya sebanding dengan ukuran parameter, jauh melebihi DDP yang “hanya” menggabungkan gradien. Praktik umumnya: pakai stage 1/2 bila model masih muat (komunikasi ringan), naikkan ke stage 3 hanya ketika parameter sendiri tidak muat. Di PyTorch, realisasi pola ini adalah torch.distributed.fsdp.FullyShardedDataParallel (FSDP) yang memecah parameter, gradien, dan state optimizer antar perangkat secara otomatis — padanan ZeRO stage 3 — dan menjadi standar pelatihan model besar bersama DDP.

Dengan sharding, masalah memori state terpecahkan. Masalah berikutnya sudah menunggu di arah lain: model sequence panjang, yang memori atensinya tumbuh kuadratik mengikuti panjang teks. Chapter Efficient Attention menutup daftar teknik dengan cara menulis ulang perhitungan attention itu sendiri.

Rujukan resmi FSDP tersedia di docs.pytorch.org/docs/stable/fsdp.html ; makalah asli ZeRO dapat dibaca di arxiv.org/abs/1910.02054 .