Lensa ML

Distributed Training

Langkah 1 dari 6: Kenapa Perlu Didistribusikan?

Memori dan komputasi terlalu besar untuk satu GPU

Token training2.00T
Training state, 16 byte per parameter × 6.74B parameter2 + 2 + 12 B = 107.8 GB80 GB GPUweight fp16gradient fp16state Adam fp32Waktu training pada 150 TFLOP/s per GPU, scaling sempurna (ilustratif)1 GPU17.1 tahun64 GPU98 hari1,024 GPU6 hari
Training state
107.8 GB
Melebihi satu GPU 80 GB
1.35×
Komputasi, 6ND
8.1e22 FLOP
Di satu GPU
17.1 tahun

Training dengan 2,00T token butuh 8,1e22 FLOP, sekitar 17,1 tahun di satu GPU pada laju ini. 107,8 GB state juga sudah melebihi kapasitas GPU sebelum langkah pertama berjalan.

Geser Token training dari 50B sampai 2T dan perhatikan batang waktunya. Angka 16 byte per parameter mengikuti hitungan Adam mixed precision di paper ZeRO: 2 untuk weight fp16, 2 untuk gradient fp16, dan 12 untuk master weight, momentum, dan variance fp32. Aktivasi menambah di atas itu dan membesar seiring ukuran batch dan panjang sequence. Aturan 6ND menghitung sekitar 2 FLOP per parameter per token untuk forward pass dan 4 untuk backward pass. Llama 2 7B dilatih dengan 2T token, jadi komputasinya sekitar 8,1e22 FLOP.

Langkah 2 dari 6: Data Parallelism

Salinan penuh, batch dibagi, gradient di-all-reduce

GPU (replika data-parallel)8
GPU 0128 seqGPU 1128 seqGPU 2128 seqGPU 3128 seqGPU 4128 seqGPU 5128 seqGPU 6128 seqGPU 7128 seqring all-reduce13.5 GB gradientTiap GPU menyimpan 107,8 GB penuh (merah: di atas 80 GB)Global batch 1.024 sequence berisi 4.096 token, dibagi rata
State per GPU
107.8 GB
Sequence per GPU
128
Terkirim per GPU, per langkah
23.6 GB
Komputasi per langkah
2.4 menit

Pada 8 GPU, masing-masing hanya menghitung 128 sequence (2,4 menit), sedangkan all-reduce tetap mengirim 23,6 GB per GPU, mendekati batas 2×. Lewat tautan 50 GB/s, itu sekitar 0,47 s per langkah (ilustratif), dan setiap salinan tetap butuh 107,8 GB.

Geser GPU dari 1 ke 64 dan bandingkan byte yang terkirim dengan waktu komputasi. Pada ring all-reduce, setiap GPU mengirim gradient-nya dalam N chunk: pertama menjumlahkan keliling ring (reduce-scatter), lalu meneruskan chunk yang sudah dijumlahkan (all-gather), sehingga totalnya 2(N − 1)/N kali gradient. Angka itu nyaris tidak bertambah seiring N, jadi metode ini mudah di-scale. Framework seperti PyTorch DDP menumpangkan all-reduce pada backward pass, dimulai dari gradient layer terakhir selagi layer sebelumnya masih dihitung.

Langkah 3 dari 6: ZeRO Sharding

Bagi optimizer state, gradient, dan weight

GPU data-parallel N8
Sharding
Yang disimpan GPU 0 dari tiap kelompok tensor (bagian yang disorot)weight fp1613.5 GB, salinan penuhgradient fp1613.5 GB, salinan penuhstate Adam80.9 GB / 8 = 10.1 GBTotal GPU 037.1 GB80 GBAktivasi dan buffer tidak dihitung
State per GPU
37.1 GB
Lebih kecil
2.9×
Terkirim per GPU, per langkah
23.6 GB
Muat di 80 GB
Ya

ZeRO-1 membagi optimizer state menjadi 8 shard, sehingga GPU 0 turun ke 37,1 GB. Komunikasinya sama dengan data parallelism biasa, 23,6 GB per GPU per langkah, dan 27,0 GB yang tidak di-shard menjadi batas bawahnya.

Geser GPU data-parallel N dan ganti antara Tanpa, ZeRO-1, ZeRO-2, dan ZeRO-3. Rumus per GPU dari Rajbhandari et al. adalah (4 + K/N)Ψ byte untuk stage 1, (2 + (2 + K)/N)Ψ untuk stage 2, dan (4 + K)Ψ/N untuk stage 3, dengan K = 12 untuk Adam. Contoh dari paper itu sendiri, model 7,5B di 64 GPU, turun dari 120 GB ke 1,9 GB. Stage 1 dan 2 tidak menambah lalu lintas karena all-reduce pada dasarnya reduce-scatter yang diikuti all-gather, dan ZeRO memperbarui tiap shard di antara keduanya.

Langkah 4 dari 6: Tensor Parallelism

Bagi tiap perkalian matriks ke beberapa GPU

GPU tensor-parallel t4
Tautan antar-GPU
input Xnaik: blok kolomakt. lokalturun: blok barisparsialall-reduceXGPU 0GPU 0GPU 1GPU 1GPU 2GPU 2GPU 3GPU 3ΣZSatu blok MLP dibagi 4; attention dibagi dengan cara yang sama, per headIlustratif: 150 TFLOP/s per GPU, NVLink 450 GB/s, jaringan 50 GB/s
Weight per GPU, satu layer
50.6M
State per GPU
27.0 GB
Terkirim per GPU, per layer
201 MB
Komunikasi ÷ komputasi
0.05×

Setiap GPU mengalikan dengan 1/4 dari tiap matriks, dan hasil parsialnya dijumlahkan oleh 4 all-reduce per layer, 201 MB per GPU untuk satu sequence. Lewat NVLink, itu butuh 0,45 ms, dibanding 8,3 ms komputasi.

Pindahkan GPU tensor-parallel t ke 1, 2, 4, dan 8, lalu ganti tautannya. Megatron-LM membagi matriks MLP pertama per kolom, sehingga setiap GPU menerapkan fungsi non-linear pada irisannya sendiri tanpa pertukaran, lalu membagi matriks kedua per baris sehingga setiap GPU menghasilkan jumlah parsial. Attention dibagi per head dengan pola yang sama. Hasilnya, ada satu all-reduce setelah attention dan satu setelah MLP di forward pass, ditambah dua lagi di backward pass. Karena pertukaran ini berada di jalur kritis, t jarang melebihi 8 GPU dalam satu server.

Langkah 5 dari 6: Pipeline Parallelism

Bagi layer, alirkan micro-batch

Pipeline stage p4
Micro-batch m4
Jadwal GPipe: satu baris per stage, waktu berjalan ke kananGPU 0L0–7GPU 1L8–15GPU 2L16–23GPU 3L24–3111223344112233441122334411223344waktu, 21 unit (backward dihitung 2× forward); abu-abu = bubble menganggurforward, bernomor per micro-batchbackward
Bubble (waktu menganggur)
42.9%
Layer per GPU
8
State per GPU
27.0 GB
Unit sibuk per GPU
12 dari 21

4 micro-batch di 4 stage menyisakan bubble 42,9%, (p − 1)/(m + p − 1). Antar-stage hanya aktivasi satu micro-batch yang berpindah, 34 MB untuk sequence 4.096 token.

Geser Pipeline stage p dan Micro-batch m, lalu perhatikan segitiga abu-abu di awal dan akhir. GPipe (Huang et al., 2018) menyimpan aktivasi semua m micro-batch sampai backward pass-nya berjalan, jadi memori membesar seiring m. Jadwal 1F1B dari PipeDream dan Megatron-LM menyelang-nyeling satu forward dan satu backward per stage, sehingga paling banyak p micro-batch berjalan bersamaan dengan bubble yang sama. Jadwal interleaved memberi setiap GPU beberapa potongan layer yang lebih kecil dan memangkas bubble lebih jauh.

Langkah 6 dari 6: 3D Parallelism

Tensor di dalam node, pipeline antar-node, data di atasnya

Tensor parallel
Pipeline stage
ZeRO pada DP
Replika data-parallel8
Replika data-parallel →Tiap kotak satu kelompok tensor-parallel berisi 8 GPU, dalam satu nodestage 0, L0–7stage 1, L8–15stage 2, L16–23stage 3, L24–318 TP × 4 PP × 8 DP = 256 GPU di 32 nodeKonfigurasi ilustratif; hari dihitung pada 150 TFLOP/s per GPU, tanpa komunikasi
GPU (node)
256 (32)
State per GPU
1.16 GB
Pipeline bubble
2.3%
2T token butuh
25 hari

256 GPU masing-masing menyimpan 1,16 GB state, menyisakan ruang untuk aktivasi. Tiap pipeline mendapat 128 micro-batch, jadi bubble-nya 2,3%. Pada skala ini, 2T token butuh sekitar 25 hari dengan scaling ideal.

Klik satu preset, lalu ubah sakelar dan Replika data-parallel. Narayanan et al. (2021) memberi resep yang umum dipakai: jaga tensor parallelism tetap di dalam satu node, tempat NVLink cepat; pakai pipeline stage antar-node, karena hanya mengirim aktivasi ke tetangganya; lalu perbesar skala dengan data parallelism. Mereka melatih model 1 triliun parameter dengan cara ini di 3.072 GPU. Derajat data-parallel yang lebih besar berarti micro-batch per pipeline lebih sedikit pada global batch yang tetap, jadi bubble naik seiring membesarnya cluster.

Apa itu distributed training dalam deep learning?

Distributed training adalah melatih satu neural network di banyak GPU sekaligus dengan membagi pekerjaannya. Data parallelism membagi batch, tensor parallelism membagi tiap perkalian matriks, dan pipeline parallelism membagi tumpukan layer. Metode sharding seperti ZeRO membagi training state itu sendiri, jadi tidak ada GPU yang harus menyimpan semuanya.

Semua model bahasa besar dilatih dengan cara ini, karena memori dan komputasinya melebihi satu akselerator. Llama 2, GPT-3, dan model sejenis berjalan di ribuan GPU sekaligus, biasanya dengan tensor parallelism di dalam tiap server 8-GPU, pipeline parallelism antarserver, dan data parallelism di atasnya. Alat yang sama, seperti PyTorch FSDP, DeepSpeed, dan Megatron-LM, juga membuat fine-tuning model 7B bisa dilakukan di beberapa GPU saja.

Ide utama

  • Training model 7B dengan Adam dalam mixed precision butuh sekitar 108 GB untuk weight, gradient, dan optimizer state, lebih dari satu GPU 80 GB.
  • Data parallelism menyalin model ke setiap GPU dan melakukan all-reduce pada gradient, yang mempercepat komputasi tetapi tidak menghemat memori.
  • ZeRO membagi optimizer state, lalu gradient, lalu weight ke GPU data-parallel, sehingga memori per GPU turun hingga N kali.
  • Tensor parallelism membagi matriks tiap layer dan butuh all-reduce di dalam setiap layer, jadi tetap berada di dalam satu node bertautan cepat.
  • Pipeline parallelism menempatkan layer ke dalam stage dan menyisakan bubble menganggur sebesar (p − 1)/(m + p − 1) dari setiap langkah.
  • Run besar menggabungkan ketiganya, dengan jumlah GPU total sama dengan hasil kali derajat tensor, pipeline, dan data-parallel.

Pertanyaan umum

Kapan butuh distributed training?

Saat training state tidak muat di satu GPU, atau satu GPU akan terlalu lama. Training penuh model 7B dengan Adam dalam mixed precision butuh sekitar 108 GB untuk weight, gradient, dan optimizer state, lebih dari kartu 80 GB, bahkan sebelum menghitung aktivasi. Fine-tuning dengan LoRA atau QLoRA melatih parameter yang jauh lebih sedikit dan sering muat di satu GPU, jadi mungkin tidak perlu distributed training sama sekali.

Apa bedanya data parallelism dan model parallelism?

Data parallelism memberi setiap GPU seluruh model dan irisan batch yang berbeda. Model parallelism membagi modelnya sendiri, baik di dalam layer (tensor parallelism) maupun antar-layer (pipeline parallelism), jadi setiap GPU hanya menyimpan sebagian weight. Data parallelism lebih sederhana dan dicoba lebih dulu, lalu model parallelism ditambahkan ketika model atau state-nya tidak lagi muat.

Apakah FSDP sama dengan ZeRO?

PyTorch Fully Sharded Data Parallel menerapkan ide yang sama dengan ZeRO stage 3. Setiap GPU menyimpan satu shard weight, gradient, dan optimizer state, lalu mengumpulkan layer utuh tepat sebelum dipakai. Strategi sharding-nya juga punya mode yang setara dengan ZeRO-2, yaitu menyimpan weight utuh tetapi membagi gradient dan optimizer state.

Apakah global batch yang lebih besar mengubah hasil training?

Ya, bisa. Lebih banyak GPU data-parallel dengan batch per GPU yang sama akan menaikkan global batch. Akibatnya noise gradient turun, dan biasanya learning rate serta warmup perlu disetel ulang. Setelah melewati critical batch size yang bergantung pada model, tambahan data per langkah hanya memberi sedikit kemajuan, jadi menambah GPU tidak lagi mempersingkat training.

Berapa GPU yang dibutuhkan untuk fine-tuning model 7B?

Fine-tuning penuh dengan Adam butuh sekitar 108 GB state ditambah aktivasi, jadi biasanya berjalan di 2 sampai 8 GPU 80 GB dengan sharding ZeRO atau FSDP. Metode hemat parameter seperti LoRA membekukan weight dasar dan melatih adapter kecil, sehingga memorinya mendekati 13,5 GB weight fp16. Dengan weight dasar 4-bit (QLoRA), satu GPU 24 GB sering sudah cukup.

Penjelasan terkait

Referensi

  1. Rajbhandari et al. (2019), "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models"
  2. Shoeybi et al. (2019), "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism"
  3. Huang et al. (2018), "GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism"
  4. Narayanan et al. (2021), "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM"

Terakhir diperbarui