Sinkronisasi gradien

Modul terdistribusi PyTorch beroperasi dengan berkomunikasi bolak-balik di antara semua GPU dalam sistem Anda. Komunikasi ini membutuhkan waktu, dan memastikan semua proses mengetahui status satu sama lain terjadi pada titik pemicu tertentu saat menggunakan modul ddp.

Titik pemicu ini ditambahkan ke model PyTorch, khususnya pada metode forward() dan backward() mereka. Hal ini terjadi ketika model dibungkus dengan DistributedDataParallel:

import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel

model = nn.Linear(10, 10)
ddp_model = DistributedDataParallel(model)

Di Accelerate, konversi ini terjadi secara otomatis saat memanggil prepare() dan memasukkan model Anda.

+ from accelerate import Accelerator
+ accelerator = Accelerator()
  import torch.nn as nn
- from torch.nn.parallel import DistributedDataParallel

  model = nn.Linear(10,10)
+ model = accelerator.prepare(model)

Perlambatan dalam akumulasi gradien

Anda kini memahami bahwa PyTorch menambahkan hook ke metode forward dan backward dari model PyTorch Anda saat pelatihan dalam setup terdistribusi. Namun, bagaimana hal ini berisiko memperlambat kode Anda?

Dalam DDP (distributed data parallel), urutan spesifik di mana proses dilakukan dan dijalankan diharapkan terjadi pada titik-titik tertentu, dan hal ini juga harus terjadi pada waktu yang kurang lebih bersamaan sebelum melanjutkan.

Contoh paling langsung adalah saat Anda memperbarui parameter model melalui optimizer.step(). Tanpa akumulasi gradien, semua instans model perlu menghitung, mengolasi, dan memperbarui gradien mereka sebelum beralih ke batch data berikutnya. Saat melakukan akumulasi gradien, Anda mengakumulasi n gradien loss dan melewati optimizer.step() hingga n batch tercapai. Karena semua proses pelatihan hanya perlu disinkronkan pada saat optimizer.step() dipanggil, tanpa modifikasi apa pun pada langkah pelatihan Anda, komunikasi antar-proses yang tidak perlu ini dapat menyebabkan perlambatan yang signifikan.

Bagaimana Anda dapat menghindari overhead ini?

Menyelesaikan masalah perlambatan

Karena Anda melewati pembaruan parameter model saat melatih pada batch-batch ini, gradiennya tidak perlu disinkronkan hingga titik di mana optimizer.step() benar-benar dipanggil. PyTorch tidak dapat secara otomatis mengetahui kapan Anda perlu melakukan ini, tetapi mereka menyediakan alat untuk membantu melalui context manager no_sync yang ditambahkan ke model Anda setelah dikonversi ke DDP.

Di dalam context manager ini, PyTorch akan melewati sinkronisasi gradien ketika .backward() dipanggil, dan pemanggilan pertama .backward() di luar context manager ini akan memicu sinkronisasi. Lihat contoh di bawah:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for index, batch in enumerate(dataloader):
    inputs, targets = batch
    # Trigger gradient synchronization on the last batch
    if index != (len(dataloader) - 1):
        with ddp_model.no_sync():
            # Gradients only accumulate
            outputs = ddp_model(inputs)
            loss = loss_func(outputs)
            accelerator.backward(loss)
    else:
        # Gradients finally sync
        outputs = ddp_model(inputs)
        loss = loss_func(outputs)
        accelerator.backward(loss)
        optimizer.step()

Di Accelerate untuk menjadikan ini API yang dapat dipanggil terlepas dari perangkat pelatihan (meskipun mungkin tidak melakukan apa-apa jika Anda tidak berada di sistem terdistribusi!), ddp_model.no_sync diganti dengan no_sync() dan beroperasi dengan cara yang sama:

  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

  for index, batch in enumerate(dataloader):
      inputs, targets = batch
      # Trigger gradient synchronization on the last batch
      if index != (len(dataloader)-1):
-         with ddp_model.no_sync():
+         with accelerator.no_sync(model):
              # Gradients only accumulate
              outputs = ddp_model(inputs)
              loss = loss_func(outputs, targets)
              accelerator.backward(loss)
      else:
          # Gradients finally sync
          outputs = ddp_model(inputs)
          loss = loss_func(outputs)
          accelerator.backward(loss)
          optimizer.step()
          optimizer.zero_grad()

Seperti yang mungkin Anda perkirakan, fungsi accumulate() membungkus pemeriksaan kondisional ini dengan melacak nomor batch saat ini, sehingga Anda mendapatkan API akumulasi gradien akhir:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for batch in dataloader:
    with accelerator.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

Sebagai hasilnya, Anda sebaiknya menggunakan accelerator.accumulate atau accelerator.no_sync dalam hal pilihan API.

Seberapa besar perlambatan yang terjadi, dan kesalahan-kesalahan mudah yang bisa Anda lakukan

Untuk menyiapkan contoh yang realistis, pertimbangkan pengaturan berikut:

• Dua node TPU tunggal T4 dan satu node dengan dua GPU
• Setiap GPU adalah T4, dan dihosting di GCP
• Skrip yang digunakan adalah modifikasi dari skrip Contoh NLP
• Ukuran batch per GPU adalah 16, dan gradien diakumulasi setiap 4 langkah

Semua skrip tersedia di repositori ini.

Jika tidak berhati-hati terhadap sinkronisasi gradien dan komunikasi GPU, sejumlah besar waktu dapat terbuang saat GPU tersebut saling berkomunikasi selama periode yang tidak diperlukan.

Seberapa besar?

Referensi:

• Baseline: tidak menggunakan praktik sinkronisasi yang dibahas di sini
• no_sync secara tidak tepat: no_sync hanya di sekitar panggilan backward, bukan forward
• no_sync: menggunakan pola no_sync dengan benar
• accumulate: menggunakan accumulate() dengan benar

Berikut adalah rata-rata detik per batch yang dihitung dari 29 batch data untuk setiap konfigurasi, baik pada satu node maupun pada konfigurasi dua node:

Seperti yang dapat Anda lihat, jika Anda tidak berhati-hati dalam mengatur sinkronisasi gradien, Anda dapat mengalami perlambatan lebih dari 2x selama pelatihan!

Jika Anda khawatir untuk memastikan semuanya dilakukan dengan benar, kami sangat menyarankan untuk memanfaatkan fungsi accumulate() dan meneruskan gradient_accumulation_steps atau gradient_accumulation_plugin ke objek Accelerator agar Accelerate dapat menangani hal ini untuk Anda.

no_sync memerlukan memori GPU tambahan saat menggunakan FSDP

Perlu diperhatikan bahwa tidak menyinkronkan gradien dapat berdampak buruk saat melakukan pelatihan FSDP. Seperti yang telah diperingatkan di torch, context manager no_sync untuk FSDP akan membutuhkan memori tambahan.

Oleh karena itu, dalam situasi yang intensif memori saat menggunakan FSDP, kami merekomendasikan untuk mengatur sync_each_batch menjadi True di GradientAccumulationPlugin untuk menonaktifkan no_sync.

Lihat contoh di bawah di mana kami melakukan fine-tuning Mixtral (47B parameter) pada 8 GPU A100-80GB. Kami melihat bahwa bahkan untuk gradient_accumulation_steps=2 yang moderat, kami dengan cepat mengalami out-of-memory (OOM) jika no_sync diaktifkan. Sekali lagi, ini disebabkan oleh overhead memori tambahan akibat no_sync dari FSDP. Namun, jika no_sync dinonaktifkan melalui sync_each_batch=True, maka konsumsi memori untuk gradient_accumulation_steps=16 kembali ke tingkat gradient_accumulation_steps=1.

Menonaktifkan no_sync berarti akan ada perlambatan akibat sinkronisasi data tambahan, seperti yang dijelaskan pada bagian-bagian sebelumnya dari panduan ini.