Suy luận phân tán

Suy luận phân tán có thể được chia thành hai nhóm:

1. Nạp toàn bộ mô hình lên mỗi GPU và gửi các phần của một lô (batch) qua bản sao mô hình trên mỗi GPU một lần
2. Tải các phần của mô hình lên từng GPU và xử lý một đầu vào tại một thời điểm

Hướng dẫn này bao phủ phần đầu tiên. Đối với phần thứ hai, xem Big Model Inference.

Tự động gửi các phần của một lô đến từng mô hình đã tải

Đây là giải pháp tiêu tốn bộ nhớ nhất, vì nó yêu cầu mỗi GPU phải giữ một bản sao đầy đủ của mô hình trong bộ nhớ tại một thời điểm nhất định.

Thông thường khi thực hiện thao tác này, người dùng sẽ chuyển mô hình sang một thiết bị cụ thể để nạp từ CPU, sau đó di chuyển từng prompt sang một thiết bị khác.

Một pipeline cơ bản sử dụng thư viện diffusers có thể trông như sau:

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

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

Sau đó, thực hiện suy luận dựa trên prompt cụ thể:

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")

Người ta sẽ nhận thấy rằng chúng ta phải kiểm tra rank để biết prompt nào cần gửi, điều này có thể hơi phiền phức.

Người dùng cũng có thể nghĩ rằng với Accelerate, việc sử dụng Accelerator để chuẩn bị một dataloader cho một tác vụ như vậy cũng có thể là một cách đơn giản để quản lý điều này. (Để tìm hiểu thêm, hãy xem phần liên quan trong Quick Tour)

Nó có thể quản lý được không? Có. Tuy nhiên, nó có thêm mã không cần thiết không? Cũng có.

Với Accelerate, chúng ta có thể đơn giản hóa quy trình này bằng cách sử dụng trình quản lý ngữ cảnh Accelerator.split_between_processes() (cũng tồn tại trong PartialState và AcceleratorState). Hàm này sẽ tự động chia nhỏ bất kỳ dữ liệu nào bạn truyền vào (dù là một prompt, một tập tensor, một từ điển chứa dữ liệu trước đó, v.v.) trên tất cả các tiến trình (có thể được đệm) để bạn có thể sử dụng ngay lập tức.

Hãy viết lại ví dụ trên bằng trình quản lý ngữ cảnh này:

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")

Và sau đó để chạy mã, chúng ta có thể sử dụng Accelerate:

Nếu bạn đã tạo một tệp cấu hình để sử dụng bằng accelerate config:

accelerate launch distributed_inference.py

Nếu bạn có một tệp cấu hình cụ thể mà bạn muốn sử dụng:

accelerate launch --config_file my_config.json distributed_inference.py

Hoặc nếu bạn không muốn tạo bất kỳ tệp cấu hình nào và khởi chạy trên hai GPU:

Lưu ý: Bạn sẽ nhận được một số cảnh báo về việc các giá trị được suy đoán dựa trên hệ thống của bạn. Để loại bỏ chúng, bạn có thể chạy accelerate config default hoặc thực hiện accelerate config để tạo tệp cấu hình.

accelerate launch --num_processes 2 distributed_inference.py

Chúng ta giờ đã giảm mã mẫu cần thiết để tách dữ liệu này xuống chỉ còn vài dòng mã một cách dễ dàng.

Nhưng điều gì sẽ xảy ra nếu chúng ta có một phân bố prompt không đều lên các GPU? Ví dụ, nếu chúng ta có 3 prompt nhưng chỉ có 2 GPU?

Trong trình quản lý ngữ cảnh, GPU đầu tiên sẽ nhận hai prompt đầu tiên và GPU thứ hai nhận prompt thứ ba, đảm bảo rằng tất cả các prompt đều được chia đều và không cần bất kỳ chi phí phụ nào.

Tuy nhiên, nếu chúng ta muốn thực hiện một thao tác nào đó với kết quả từ tất cả các GPU thì sao? (Ví dụ như thu thập tất cả và thực hiện một dạng xử lý hậu kỳ nào đó) Bạn có thể truyền apply_padding=True để đảm bảo rằng các danh sách prompt được đệm (padding) đến cùng một độ dài, với dữ liệu bổ sung được lấy từ mẫu cuối cùng. Bằng cách này, tất cả các GPU sẽ có cùng số lượng prompt, và bạn có thể thu thập kết quả.

Điều này chỉ cần thiết khi thực hiện các thao tác như thu thập kết quả, trong đó dữ liệu trên mỗi thiết bị cần có cùng độ dài. Suy luận cơ bản không yêu cầu điều này.

Ví dụ:

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

Trên GPU đầu tiên, các prompt sẽ là ["a dog", "a cat"], và trên GPU thứ hai sẽ là ["a chicken", "a chicken"]. Hãy đảm bảo loại bỏ mẫu cuối cùng, vì nó sẽ là bản sao của mẫu trước đó.

Bạn có thể tìm thấy các ví dụ phức tạp hơn tại đây, chẳng hạn như cách sử dụng với LLM.