Tensor Parallelism dan DeepSpeed ZeRO Stage 3: Biar Model Besar Gak Pake Pusing!

Bayangin kamu lagi bantuin teman angkat kotak besar yang berat banget. Kalau cuma kamu sendiri yang angkat, pasti capek dan lama banget kan? Nah, itu kira-kira gambaran dari Tensor Parallelism di dunia AI. Jadi, kalau model transformer raksasa punya bobot (matriks) yang super gede, kita gak mungkin suruh satu GPU aja yang ngolah semua. Makanya, bobot itu dibagi-bagi ke beberapa GPU biar kerjaannya ringan dan cepet selesai.

Dalam Tensor Parallelism, matriks bobot itu dipecah jadi beberapa potongan—bisa berdasarkan kolom atau baris—dan setiap GPU cuma handle satu potongan. Misalnya ada 4 GPU, bobotnya dipotong jadi 4 bagian. Tiap GPU cuma ngolah potongan itu aja. Abis itu hasilnya digabung lagi jadi output penuh. Kayak kamu dan temanmu angkat kotak yang sudah dipotong-potong, terus di mobil kalian satukan lagi jadi kotak utuh.

• Pembagian Potongan Bobot (Sharding)
  Bobot $W$ besar, misal ukuran $d_{out} \times d_{in}$, dibagi jadi beberapa shard. Tiap GPU dapat satu shard. Proses perkalian matriks $X \times W$ juga dilakukan di potongan masing-masing GPU.
• Penggabungan Hasil
  Kalau bobot dipotong per kolom, hasil perkalian tiap GPU berupa kolom parsial yang dijajarkan (concatenate) jadi output penuh. Kalau dipotong per baris, hasilnya perlu dijumlahkan (all-reduce) supaya output lengkap.
• Keuntungan
  Hemat memori per GPU, komputasi lebih cepat karena paralel, dan gampang skala kalau butuh lebih banyak GPU.
• Tipe Pembagian
  - Column-Parallel (CP): Bobot dipotong per kolom, output tinggal dijajarkan.
  - Row-Parallel (RP): Bobot dipotong per baris, output perlu dijumlahkan all-reduce.
• Contoh di Hugging Face Transformers
  Kamu tinggal set konfigurasi tensor parallel, misal `--tensor-parallel-size=4`, terus library otomatis atur pembagian bobot dan aktivasi ke GPU. Kamu tetap panggil `.generate()` atau `.forward()` biasa, dan di belakang layar pembagian kerja sudah di-handle.

DeepSpeed ZeRO Stage 3 itu ibarat gudang super gede yang penuh dengan barang (parameter model), buku catatan (optimizer state), dan nota kerja (gradien). Kalau cuma satu orang yang pegang semuanya, memori GPU bisa penuh banget. ZeRO Stage 3 hadir buat bagi-bagi semua itu ke beberapa GPU sehingga gak ada yang bawa semuanya sekaligus.

• Barang, Catatan, dan Nota
  - Barang = parameter model (bobot jaringan)
  - Catatan = optimizer state (momentum, variance)
  - Nota = gradien (hasil backprop)
• Stage 1 sampai 3
  - Stage 1: Bagikan optimizer state saja.
  - Stage 2: Bagikan optimizer state dan gradien.
  - Stage 3: Bagikan parameter model, optimizer state, dan gradien secara penuh. Ini paling ekstrem dan efisien.
• Cara Kerja ZeRO Stage 3
  - Forward pass: GPU cuma panggil potongan parameter yang dia pegang, kalau perlu potongan dari GPU lain, pinjam dulu (all-gather), lalu setelah selesai kembalikan.
  - Backward pass: Gradien dan optimizer state juga di-shard, update parameter cuma di potongan lokal, tanpa salinan lengkap.
• Manfaat
  Memori per GPU turun drastis, bisa training model yang jauh lebih besar, batch size bisa naik, throughput juga makin tinggi.
• ZeRO Inference
  DeepSpeed juga punya versi ZeRO buat inferensi. Model besar tetap bisa dipakai di cluster GPU yang lebih kecil, dengan bobot dipanggil per lapisan saat dibutuhkan, terus dikembalikan setelah selesai.
• Kolaborasi dengan Tensor Parallelism
  Tensor Parallelism bantu pecah komputasi lapisan, ZeRO Stage 3 pecah parameter dan optimizer state. Kalau digabung, tiap GPU tanggung jawab potongan bobot dan komputasi lapisan, komunikasi dioptimasi biar efisien.
• Implementasi di Hugging Face
  Konfigurasinya bisa via file JSON seperti:
  

  {
    "zero_optimization": {
      "stage": 3,
      "offload_param": {
        "device": "cpu",
        "pin_memory": true
      },
      "allgather_bucket_size": 500000000,
      "reduce_bucket_size": 500000000
    },
    "train_batch_size": 8,
    "train_micro_batch_size_per_gpu": 1
  }

  
  Di Python tinggal inisialisasi model dan deepspeed:
  

  from transformers import AutoModelForCausalLM
  import deepspeed
  
  model = AutoModelForCausalLM.from_pretrained("nama-model-sahabatAI")
  ds_config = "deepspeed_config.json"
  model, optimizer, _, _ = deepspeed.initialize(
      config=ds_config,
      model=model,
      model_parameters=model.parameters()
  )

  
  Untuk inferensi:
  

  import torch
  from transformers import AutoModelForCausalLM, AutoTokenizer
  
  tokenizer = AutoTokenizer.from_pretrained("nama-model-sahabatAI")
  model = AutoModelForCausalLM.from_pretrained(
      "nama-model-sahabatAI",
      torch_dtype=torch.float16,
      device_map="auto",
      use_deepspeed=True,
      deepspeed_config="deepspeed_infer.json"
  )
  
  inputs = tokenizer("Halo, SahabatAI!", return_tensors="pt").to(model.device)
  outputs = model.generate(**inputs)
  print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Ringkasnya:

• Tensor Parallelism = Memecah matriks bobot jadi potongan-potongan, tiap GPU angkat satu potongan, terus hasilnya digabung lagi.
• ZeRO Stage 3 = Memecah parameter, optimizer state, dan gradien jadi bagian kecil-kecil, tiap GPU cuma bawa sebagian kecil, pinjam bagian lain saat perlu, terus balikin setelah selesai.

Dengan kombinasi ini, SahabatAI dan model transformer raksasa lainnya bisa jalan lancar di cluster GPU yang gak perlu super mahal atau gede banget memorinya. Jadi, gak cuma hemat biaya, tapi juga ngebut banget buat training dan inferensi. Bayangin aja, kayak punya tim angkut kotak dan penjaga gudang yang pinter banget ngatur tugas masing-masing, biar semua kerjaan beres tanpa drama!