DeepSpeed: Melatih Model Raksasa di GPU Terbatas dengan ZeRO dan Offload

Halo temen-temen, kali ini kita masuk ke topik yang agak berat tapi sangat berguna kalau kalian mulai main model besar, namanya DeepSpeed. Ini library buatan Microsoft yang tugasnya bikin training model raksasa jadi mungkin di hardware yang jumlahnya terbatas. Kalau kalian pernah dapat error CUDA out of memory padahal batch size udah dikecilin sampai satu, atau pernah lihat model 7 miliar parameter dan mikir "ini nggak mungkin muat di GPU saya", DeepSpeed adalah jawaban yang perlu kalian pelajari.

Inti keajaibannya ada di teknologi bernama ZeRO, singkatan dari Zero Redundancy Optimizer. Ide dasarnya cerdas dan sebenarnya sederhana: di training multi-GPU biasa, setiap GPU menyimpan salinan lengkap dari model, gradien, dan state optimizer. Itu boros luar biasa. ZeRO memecah semua itu dan membagi potongannya ke seluruh GPU, sehingga tidak ada data yang disimpan berulang kali.

Di tutorial ini kita bahas kenapa memori training itu jauh lebih besar dari yang kalian kira, tiga tahap ZeRO dan kapan memakai masing-masing, instalasi, file konfigurasi JSON, integrasi dengan HuggingFace Trainer dan dengan Accelerate, CPU offload, sampai tips debugging yang bakal menghemat berjam-jam waktu kalian.

Introduction

Sebelum ke kode, kalian harus paham dulu memori GPU kalian sebenarnya dipakai buat apa.

Ke Mana Perginya Memori GPU

Banyak orang mengira memori training itu didominasi bobot model. Padahal bukan. Mari kita hitung untuk model 1 miliar parameter yang dilatih dengan Adam dalam presisi campuran.

Bobot model dalam fp16 memakan sekitar 2 byte per parameter, jadi 2 GB.

Gradien dalam fp16 juga 2 byte per parameter, jadi 2 GB lagi.

State optimizer Adam menyimpan salinan bobot fp32, momentum fp32, dan variance fp32. Itu 4 plus 4 plus 4, jadi 12 byte per parameter, alias 12 GB.

Totalnya sudah 16 GB hanya untuk state, sebelum menghitung aktivasi. Aktivasi (hasil antara tiap layer yang disimpan untuk backward pass) bisa menambah beberapa GB lagi tergantung batch size dan panjang sekuens.

Sekarang bayangkan kalian punya 4 GPU. Di training data-parallel biasa, semua angka di atas dikalikan empat karena setiap GPU menyimpan salinan penuh. 64 GB terpakai untuk menyimpan informasi yang sebenarnya cuma bernilai 16 GB. Inilah redundansi yang dihapus ZeRO.

Tiga Tahap ZeRO

ZeRO membagi pekerjaannya dalam tiga tahap yang bisa kalian pilih.

Tahap 1 membagi state optimizer ke seluruh GPU. Ini menghapus porsi memori terbesar (12 GB dari contoh di atas) dengan overhead komunikasi yang sangat kecil. Ini pilihan aman untuk hampir semua orang.

Tahap 2 membagi state optimizer dan gradien. Penghematan lebih besar, overhead komunikasi masih ringan. Ini titik manis yang paling sering dipakai orang.

Tahap 3 membagi state optimizer, gradien, dan juga parameter model itu sendiri. Artinya tidak ada satu GPU pun yang memegang model utuh. Parameter dikumpulkan sesaat ketika sebuah layer dieksekusi, lalu dilepas lagi. Penghematan memorinya dramatis, tapi komunikasi antar-GPU jadi jauh lebih sibuk, sehingga butuh interkoneksi cepat seperti NVLink atau InfiniBand supaya tidak lambat.

Di atas itu ada offload, yaitu memindahkan state optimizer atau bahkan parameter ke RAM CPU atau NVMe. Ini bikin kalian bisa melatih model yang benar-benar tidak muat di GPU, dengan konsekuensi kecepatan turun signifikan.

Aturan praktis yang aku pakai: mulai dari tahap 2. Kalau masih kehabisan memori, naik ke tahap 3. Kalau masih kurang juga, aktifkan CPU offload. Jangan langsung lompat ke tahap 3 plus offload karena kalian akan membayar kecepatan tanpa perlu.

Instalasi

DeepSpeed butuh Linux, PyTorch dengan CUDA, dan compiler C++. Di Windows sebaiknya pakai WSL2.

python -m venv venv
source venv/bin/activate

pip install torch
pip install deepspeed

Verifikasi instalasi dan lihat operator mana yang tersedia:

dsreport

Perintah itu akan mencetak tabel berisi status tiap ekstensi. Tidak semua harus hijau. Yang penting torch dan deepspeed terdeteksi, dan kalau kalian mau memakai CPU offload, pastikan cpuadam bisa dikompilasi.

Kalau kompilasi ekstensi gagal, biasanya karena CUDAHOME tidak diset atau versi compiler tidak cocok. Set variabel lingkungan berikut sering menyelesaikan masalah:

export CUDAHOME=/usr/local/cuda
export DSBUILDCPUADAM=1

{
  "trainmicrobatchsizepergpu": 4,
  "gradientaccumulationsteps": 8,
  "gradientclipping": 1.0,


  "bf16": {
    "enabled": true
  },

  "zerooptimization": {
    "stage": 2,
    "overlapcomm": true,

    "contiguousgradients": true,
    "reducebucketsize": 5e7,
    "allgatherbucketsize": 5e7
  },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weightdecay": 0.01

    }
  },

  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmupminlr": 0,

      "warmupmaxlr": 2e-5,

      "warmupnumsteps": 100,

      "totalnumsteps": 5000

    }
  },

  "stepsperprint": 50,

  "wallclockbreakdown": false

}

Simpan sebagai dsconfig.json. Beberapa penjelasan penting.

Untuk tahap 3 dengan offload, konfigurasi zero-nya jadi seperti ini:

"zerooptimization": {
  "stage": 3,
  "offloadoptimizer": { "device": "cpu", "pinmemory": true },

  "offloadparam": { "device": "cpu", "pinmemory": true },

  "overlapcomm": true,
  "contiguousgradients": true,

  "stage3gather16bitweightsonmodelsave": true

}

Perhatikan stage3gather16bitweightsonmodelsave. Tanpa opsi ini, checkpoint yang kalian simpan hanya berisi shard, bukan model utuh, dan kalian akan bingung waktu memuatnya kembali.

from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    TrainingArguments, Trainer, DataCollatorForLanguageModeling,
)
from datasets import loaddataset

modelname = "meta-llama/Llama-3.2-1B"

tokenizer = AutoTokenizer.frompretrained(modelname)

tokenizer.padtoken = tokenizer.eostoken

model = AutoModelForCausalLM.frompretrained(modelname)


ds = loaddataset("wikitext", "wikitext-2-raw-v1", split="train[:2000]")

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True, maxlength=512)


ds = ds.map(tokenize, batched=True, removecolumns=ds.columnnames)


args = TrainingArguments(
    outputdir="out",
    perdevicetrainbatchsize=4,
    gradientaccumulationsteps=8,
    numtrainepochs=1,
    loggingsteps=50,

    savestrategy="epoch",
    bf16=True,
    deepspeed="dsconfig.json",     # ini satu-satunya baris DeepSpeed

)

trainer = Trainer(
    model=model,
    args=args,
    traindataset=ds,
    datacollator=DataCollatorForLanguageModeling(tokenizer, mlm=False),

)
trainer.train()

deepspeed --numgpus=4 train.py

Satu hal yang perlu kalian tahu: kalau kalian mengisi optimizer dan scheduler di JSON, itu akan menimpa pengaturan di TrainingArguments. Supaya tidak bingung, kalian bisa memakai nilai khusus "auto" di JSON untuk membiarkan Trainer yang menentukan.

"optimizer": {
  "type": "AdamW",
  "params": { "lr": "auto", "weightdecay": "auto" }

}

accelerate config

Jawab ya saat ditanya apakah ingin memakai DeepSpeed, lalu tunjuk file dsconfig.json kalian atau isi pertanyaan tahap ZeRO yang muncul. Setelah itu training script kalian yang sudah memakai Accelerator() jalan tanpa perubahan sama sekali:

accelerate launch train.py

Ini kombinasi favoritku karena kode training tetap PyTorch murni dan bisa dijalankan tanpa DeepSpeed waktu debugging di laptop.

Cara 3: API DeepSpeed Langsung

Kalau kalian butuh kontrol penuh, ini bentuk aslinya.

import deepspeed
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

model = nn.Sequential(
    nn.Linear(512, 2048), nn.ReLU(),
    nn.Linear(2048, 2048), nn.ReLU(),
    nn.Linear(2048, 10),
)

x = torch.randn(4096, 512)
y = torch.randint(0, 10, (4096,))
loader = DataLoader(TensorDataset(x, y), batchsize=4)


modelengine, optimizer, , scheduler = deepspeed.initialize(

    model=model,
    modelparameters=model.parameters(),
    config="dsconfig.json",

)

criterion = nn.CrossEntropyLoss()

for epoch in range(3):
    for xb, yb in loader:
        xb = xb.to(modelengine.device)
        yb = yb.to(modelengine.device)


        loss = criterion(modelengine(xb), yb)

        modelengine.backward(loss)   # bukan loss.backward()

        modelengine.step()           # sudah termasuk zerograd dan scheduler


    if modelengine.globalrank == 0:

        print(f"epoch {epoch} loss {loss.item():.4f}")

modelengine.savecheckpoint("checkpoints", tag="epochfinal")

Perhatikan tiga perbedaan penting dari PyTorch biasa. Kalian memanggil modelengine.backward(loss) bukan loss.backward(). Kalian memanggil modelengine.step() yang sudah mengurus optimizer step, scheduler step, dan zero grad sekaligus. Dan kalian tidak boleh memanggil optimizer.zerograd() manual, karena akan mengacaukan akumulasi gradien internal.

Checkpoint DeepSpeed itu tersharding. Isinya bukan satu file pytorchmodel.bin melainkan banyak file per rank. Untuk mengubahnya jadi bobot standar yang bisa dimuat siapa saja, DeepSpeed menyediakan skrip bawaan.

python -m deepspeed.utils.zerotofp32 checkpoints/epochfinal modelfp32.bin

Atau dari dalam Python:

from deepspeed.utils.zerotofp32 import loadstatedictfromzerocheckpoint

model = loadstatedictfromzerocheckpoint(model, "checkpoints/epochfinal")

Kalau kalian memakai HuggingFace Trainer dengan stage3gather16bitweightsonmodelsave aktif, trainer.savemodel() sudah menghasilkan model utuh, jadi langkah konversi ini tidak diperlukan.

Menghitung Kebutuhan Memori Sebelum Training

Ini fitur yang jarang orang tahu tapi sangat menghemat waktu. DeepSpeed bisa memperkirakan kebutuhan memori sebelum kalian menjalankan apa pun.

from transformers import AutoModel
from deepspeed.runtime.zero.stage3 import estimatezero3modelstatesmemneedsalllive


model = AutoModel.frompretrained("meta-llama/Llama-3.2-1B")
estimatezero3modelstatesmemneedsalllive(model, numgpuspernode=4, numnodes=1)

Aktifkan gradient checkpointing kalau aktivasi yang membebani, bukan state optimizer. Di HuggingFace cukup gradientcheckpointing=True di TrainingArguments. Ini menukar sekitar 30 persen kecepatan dengan penghematan memori aktivasi yang besar.

Perhatikan bahwa CPU offload butuh RAM besar dan CPU yang tidak lemah. Kalau RAM host kalian pas-pasan, offload justru bisa membuat proses kalian mati kehabisan memori sistem.

Jangan gabungkan deepspeed launcher dengan torchrun. Pilih satu. Kalau memakai Accelerate, cukup accelerate launch.

Selalu cek dsreport setelah upgrade PyTorch atau CUDA. Ekstensi DeepSpeed dikompilasi terhadap versi tertentu dan bisa diam-diam rusak setelah upgrade.

Checkpoint tahap 3 tersharding dan perlu dikonversi dengan zeroto_fp32 sebelum dipakai di tempat lain.

Gunakan estimator memori bawaan sebelum menyewa GPU, dan gabungkan dengan gradient checkpointing kalau aktivasi yang jadi masalah.

Coba mulai dari fine-tuning model 1 miliar parameter dengan ZeRO tahap 2 di dua GPU. Begitu kalian lihat model yang tadinya out of memory sekarang jalan mulus, kalian bakal paham kenapa library ini jadi tulang punggung banyak proyek LLM open source. Selamat mencoba.