Inferensi terdistribusi

Inferensi terdistribusi dapat dibagi menjadi dua kategori:

1. Memuat seluruh model ke setiap GPU dan mengirimkan potongan batch melalui salinan model di setiap GPU secara bergantian
2. Memuat bagian-bagian model ke setiap GPU dan memproses satu input pada satu waktu

Panduan ini membahas kurung pertama. Untuk yang kedua, lihat Big Model Inference.

Mengirim potongan batch secara otomatis ke setiap model yang dimuat

Ini adalah solusi yang paling intensif dalam penggunaan memori, karena setiap GPU harus menyimpan salinan lengkap model di memori pada waktu tertentu.

Biasanya saat melakukan ini, pengguna mengirim model ke perangkat tertentu untuk memuatnya dari CPU, lalu memindahkan setiap prompt ke perangkat yang berbeda.

Pipeline dasar menggunakan pustaka diffusers mungkin terlihat seperti ini:

import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)

Kemudian dilanjutkan dengan melakukan inferensi berdasarkan prompt tertentu:

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    pipe.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    result = pipe(prompt).images[0]
    result.save(f"result_{rank}.png")

Akan terlihat bahwa kita harus memeriksa rank untuk mengetahui prompt apa yang harus dikirim, yang bisa sedikit merepotkan.

Pengguna mungkin juga berpikir bahwa dengan Accelerate, menggunakan Accelerator untuk menyiapkan dataloader untuk tugas semacam itu juga bisa menjadi cara sederhana untuk mengelolanya. (Untuk mempelajari lebih lanjut, lihat bagian terkait di Quick Tour)

Bisakah mengelolanya? Ya. Namun, apakah itu menambahkan kode tambahan yang tidak perlu: ya juga.

Dengan Accelerate, kita dapat menyederhanakan proses ini dengan menggunakan context manager Accelerator.split_between_processes() (yang juga tersedia di PartialState dan AcceleratorState). Fungsi ini akan secara otomatis membagi data apa pun yang Anda berikan (baik berupa prompt, kumpulan tensor, kamus data sebelumnya, dll.) ke semua proses (dengan potensi untuk di-pad) agar dapat Anda gunakan langsung.

Mari kita tulis ulang contoh di atas menggunakan context manager ini:

import torch
from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipe(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

Lalu, untuk menjalankan kode, kita dapat menggunakan Accelerate:

Jika Anda telah membuat file konfigurasi untuk digunakan dengan accelerate config:

accelerate launch distributed_inference.py

Jika Anda memiliki file konfigurasi tertentu yang ingin digunakan:

accelerate launch --config_file my_config.json distributed_inference.py

Atau jika Anda tidak ingin membuat file konfigurasi apa pun dan ingin meluncurkan di dua GPU:

Catatan: Anda akan menerima beberapa peringatan tentang nilai yang ditebak berdasarkan sistem Anda. Untuk menghapus peringatan ini, Anda dapat menjalankan accelerate config default atau melalui accelerate config untuk membuat file konfigurasi.

accelerate launch --num_processes 2 distributed_inference.py

Kami kini telah mengurangi kode boilerplate yang diperlukan untuk membagi data ini menjadi beberapa baris kode dengan mudah.

Bagaimana jika distribusi prompt ke GPU tidak merata? Misalnya, bagaimana jika kita memiliki 3 prompt, tetapi hanya 2 GPU?

Di bawah context manager, GPU pertama akan menerima dua prompt pertama dan GPU kedua menerima prompt ketiga, sehingga memastikan bahwa semua prompt terbagi dan tidak diperlukan overhead.

Meskipun demikian, bagaimana jika kita kemudian ingin melakukan sesuatu dengan hasil dari semua GPU? (Misalnya, mengumpulkan semuanya dan melakukan semacam pascapemrosesan) Anda dapat memasukkan apply_padding=True untuk memastikan bahwa daftar prompt dipadatkan hingga panjang yang sama, dengan data tambahan diambil dari sampel terakhir. Dengan cara ini, semua GPU akan memiliki jumlah prompt yang sama, dan Anda kemudian dapat mengumpulkan hasilnya.

Ini hanya diperlukan saat mencoba melakukan tindakan seperti mengumpulkan hasil, di mana data pada setiap perangkat harus memiliki panjang yang sama. Inferensi dasar tidak memerlukan ini.

Sebagai contoh:

import torch
from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat", "a chicken"], apply_padding=True) as prompt:
    result = pipe(prompt).images

Pada GPU pertama, prompt akan menjadi ["a dog", "a cat"], dan pada GPU kedua akan menjadi ["a chicken", "a chicken"]. Pastikan untuk membuang sampel terakhir, karena akan menjadi duplikat dari sampel sebelumnya.

Anda dapat menemukan contoh yang lebih kompleks di sini, seperti cara menggunakannya dengan LLM.