Tambahkan Accelerate ke kode Anda

Setiap framework pelatihan terdistribusi memiliki cara kerjanya sendiri yang dapat mengharuskan penulisan banyak kode kustom untuk menyesuaikannya dengan kode pelatihan PyTorch dan lingkungan pelatihan Anda. Accelerate menawarkan cara yang ramah untuk berinteraksi dengan framework pelatihan terdistribusi ini tanpa perlu mempelajari detail spesifik dari masing-masing framework. Accelerate menangani detail-detail tersebut untuk Anda, sehingga Anda dapat fokus pada kode pelatihan dan skalakannya ke lingkungan pelatihan terdistribusi apa pun.

Dalam tutorial ini, Anda akan belajar cara mengadaptasi kode PyTorch yang sudah ada dengan Accelerate dan memulai perjalanan Anda menuju pelatihan pada sistem terdistribusi dengan mudah! Anda akan memulai dengan loop pelatihan PyTorch dasar (dianggap bahwa semua objek pelatihan seperti model dan optimizer sudah disiapkan) dan secara bertahap mengintegrasikan Accelerate ke dalamnya.

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Akselerator

Accelerator adalah kelas utama untuk menyesuaikan kode Anda agar dapat bekerja dengan Accelerate. Kelas ini mengetahui tentang konfigurasi distribusi yang Anda gunakan, seperti jumlah proses yang berbeda dan jenis perangkat keras Anda. Kelas ini juga menyediakan akses ke banyak metode yang diperlukan untuk memungkinkan kode PyTorch Anda bekerja di lingkungan pelatihan distribusi apa pun serta untuk mengelola dan mengeksekusi proses di berbagai perangkat.

Itulah mengapa Anda harus selalu memulai dengan mengimpor dan membuat instance Accelerator di skrip Anda.

from accelerate import Accelerator

accelerator = Accelerator()

Accelerator juga mengetahui perangkat mana yang harus dituju untuk memindahkan objek PyTorch Anda, sehingga disarankan untuk membiarkan Accelerate menangani hal ini untuk Anda.

- device = "cuda"
+ device = accelerator.device
  model.to(device)

Menyiapkan objek PyTorch

Selanjutnya, Anda perlu menyiapkan objek PyTorch Anda (model, optimizer, scheduler, dll.) untuk pelatihan terdistribusi. Metode prepare() menangani penempatan model Anda ke dalam kontainer yang sesuai (seperti GPU tunggal atau multi-GPU) untuk konfigurasi pelatihan Anda, menyesuaikan optimizer dan scheduler untuk menggunakan AcceleratedOptimizer dan AcceleratedScheduler dari Accelerate, serta membuat dataloader baru yang dapat dipecah (sharded) di antara proses.

Accelerate hanya menyiapkan objek yang mewarisi dari kelas PyTorch masing-masing, seperti torch.optim.Optimizer.

Objek PyTorch dikembalikan dalam urutan yang sama seperti saat dikirim.

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

Loop pelatihan

Terakhir, hapus panggilan to(device) untuk input dan target di dalam loop pelatihan karena kelas DataLoader dari Accelerate secara otomatis menempatkan mereka pada perangkat yang tepat. Anda juga harus mengganti jalur backward() yang biasa dengan metode backward() dari Accelerate yang menskalakan gradien untuk Anda dan menggunakan metode backward() yang sesuai tergantung pada konfigurasi terdistribusi Anda (misalnya, DeepSpeed atau Megatron).

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

Rangkum semuanya dan loop pelatihan Accelerate baru Anda seharusnya sekarang terlihat seperti ini!

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

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

Fitur pelatihan

Accelerate menawarkan fitur tambahan - seperti akumulasi gradien, pemotongan gradien, pelatihan presisi campuran, dan lainnya - yang dapat Anda tambahkan ke skrip Anda untuk meningkatkan eksekusi pelatihan. Mari kita jelajahi ketiga fitur ini.

Akselerasi gradien

Gradient accumulation memungkinkan Anda melatih pada ukuran batch yang lebih besar dengan mengakumulasi gradien selama beberapa batch sebelum memperbarui bobot. Hal ini dapat berguna untuk mengatasi keterbatasan memori. Untuk mengaktifkan fitur ini di Accelerate, tentukan parameter gradient_accumulation_steps di kelas Accelerator dan tambahkan context manager accumulate() ke skrip Anda.

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

Pemangkasan gradien

Gradient clipping adalah teknik untuk mencegah “ledakan gradien”, dan Accelerate menawarkan:

• clipgrad_value() untuk memotong gradien ke nilai minimum dan maksimum
• clipgrad_norm() untuk menormalisasi gradien ke nilai tertentu

Presisi campuran

Mixed precision mempercepat pelatihan dengan menggunakan tipe data presisi rendah seperti fp16 (half-precision) untuk menghitung gradien. Untuk kinerja terbaik dengan Accelerate, loss harus dihitung di dalam model Anda (seperti pada model Transformers) karena komputasi di luar model dihitung dalam presisi penuh.

Atur tipe presisi campuran yang akan digunakan di Accelerator, lalu gunakan context manager autocast() untuk secara otomatis mengonversi nilai ke tipe data yang ditentukan.

Accelerate memungkinkan presisi campuran otomatis, sehingga autocast() hanya diperlukan jika ada operasi presisi campuran lain di luar yang dilakukan pada loss oleh backward() yang sudah menangani penskalaan.

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)

Simpan dan muat

Accelerate juga dapat menyimpan dan memuat model setelah pelatihan selesai, atau Anda juga dapat menyimpan state model dan optimizer yang dapat berguna untuk melanjutkan pelatihan.

Model

Setelah semua proses selesai, buka bungkus model dengan metode unwrap_model() sebelum menyimpannya karena metode prepare() membungkus model Anda ke dalam antarmuka yang tepat untuk pelatihan terdistribusi. Jika Anda tidak membuka bungkus model, penyimpanan kamus status model juga akan menyimpan lapisan tambahan potensial dari model yang lebih besar dan Anda tidak akan dapat memuat kembali bobot ke model dasar Anda.

Anda harus menggunakan metode save_model() untuk membungkus dan menyimpan kamus status model. Metode ini juga dapat menyimpan model ke dalam checkpoint yang dipartisi atau ke dalam format safetensors.

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)

load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})

Status

Selama pelatihan, Anda mungkin ingin menyimpan keadaan saat ini dari model, optimizer, generator acak, dan berpotensi juga scheduler laju belajar agar dapat dipulihkan dalam skrip yang sama. Anda harus menambahkan metode save_state() dan load_state() ke skrip Anda untuk menyimpan dan memuat keadaan.

Untuk lebih menyesuaikan di mana dan bagaimana state disimpan melalui save_state(), gunakan kelas ProjectConfiguration. Misalnya, jika automatic_checkpoint_naming diaktifkan, setiap checkpoint yang disimpan akan disimpan di Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}.

Item stateful lainnya yang akan disimpan harus didaftarkan dengan metode register_for_checkpointing() agar dapat disimpan dan dimuat. Setiap objek yang diteruskan ke metode ini untuk disimpan harus memiliki fungsi load_state_dict dan state_dict.

Jika Anda telah menginstal torchdata>=0.8.0, Anda juga dapat meneruskan use_stateful_dataloader=True ke DataLoaderConfiguration Anda. Hal ini memperluas kelas DataLoader dari Accelerate dengan fungsi load_state_dict dan state_dict, serta membuat Accelerator.save_state dan Accelerator.load_state juga melacak sejauh mana dataset pelatihan telah dibaca saat mempersistenkan model.