Song song dữ liệu

Data parallelism nhân bản cùng một mô hình trên mỗi GPU và cấp cho mỗi GPU các prompt khác nhau. Nó tăng thông lượng khi mô hình đã vừa trên một GPU, và giữ nguyên độ trễ của một prompt đơn lẻ. Nếu bạn chưa làm, hãy thử Batch inference trên một GPU trước. Tham khảo Distributed inference để so sánh với các chiến lược khác.

Bạn có thể thiết lập song song hóa dữ liệu bằng Accelerate hoặc trực tiếp bằng PyTorch Distributed.

Accelerate

Accelerate thiết lập các tiến trình và gán một GPU cho mỗi tiến trình.

Cài đặt Accelerate bằng lệnh sau.

uv pip install accelerate

Tạo một PartialState trong tệp có tên run_distributed.py để thiết lập môi trường phân tán. Di chuyển DiffusionPipeline sang PartialState.device để gán một GPU cho mỗi tiến trình.

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

Trong cùng một tệp, hãy sử dụng split_between_processes như một trình quản lý ngữ cảnh để chia các prompt giữa các tiến trình.

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

Chạy script bằng accelerate launch, và đặt --num_processes bằng số lượng GPU.

accelerate launch --num_processes=2 run_distributed.py

Để xem một ví dụ hoàn chỉnh, hãy xem script suy luận đa GPU tối giản này. Để tìm hiểu thêm, hãy xem hướng dẫn Suy luận phân tán với Accelerate.

Phân tán PyTorch

Bạn cũng có thể thiết lập song song hóa dữ liệu trực tiếp bằng torch.distributed, không cần Accelerate.

Gọi init_process_group để tạo môi trường phân tán. torchrun thiết lập rank của từng tiến trình và tổng số tiến trình (world size) thông qua các biến môi trường, vì vậy bạn không cần tự truyền chúng.

Mỗi tiến trình cần một GPU riêng và một prompt riêng. torchrun đặt biến môi trường LOCAL_RANK thành chỉ số GPU của tiến trình trên node của nó, vì vậy hãy dùng nó để chọn GPU. Sử dụng get_rank, đếm số tiến trình trên tất cả các node, để chọn prompt.

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

def main():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")

    pipeline = DiffusionPipeline.from_pretrained(
        "Qwen/Qwen-Image", dtype=torch.bfloat16,
    ).to(device)

    prompts = ["a dog", "a cat"]
    prompt = prompts[rank]

    image = pipeline(prompt).images[0]
    image.save(f"{prompt.replace(' ', '_')}.png")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Lưu script dưới tên run_distributed.py và chạy bằng torchrun. Đặt --nproc-per-node bằng số lượng GPU. Danh sách prompts có một prompt cho mỗi tiến trình, vì vậy hãy thêm một prompt cho mỗi GPU bạn sử dụng.

torchrun --nproc-per-node 2 run_distributed.py

Update on GitHub

Diffusers

Song song dữ liệu

Accelerate

Phân tán PyTorch