데이터 병렬

Data parallelism은 동일한 모델을 각 GPU에 복제하고 각 GPU에 서로 다른 프롬프트를 할당합니다. 모델이 이미 단일 GPU에 들어가는 경우 처리량을 높이며, 단일 프롬프트의 지연 시간은 변하지 않습니다. 아직 시도하지 않았다면, 먼저 단일 GPU에서 Batch inference를 사용해 보세요. 다른 전략과 비교하려면 Distributed inference를 참고하세요.

Accelerate 또는 PyTorch Distributed를 직접 사용하여 데이터 병렬성을 설정할 수 있습니다.

Accelerate

Accelerate는 프로세스를 설정하고 각 프로세스에 GPU를 할당합니다.

다음 명령어로 Accelerate를 설치합니다.

uv pip install accelerate

run_distributed.py라는 파일에서 PartialState를 생성하여 분산 환경을 설정합니다. DiffusionPipeline을 PartialState.device로 이동하여 각 프로세스에 GPU를 할당합니다.

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)

같은 파일에서 split_between_processes를 컨텍스트 매니저로 사용하여 프롬프트를 프로세스 간에 분할합니다.

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

accelerate launch로 스크립트를 실행하고 --num_processes를 GPU 수로 설정하세요.

accelerate launch --num_processes=2 run_distributed.py

완전한 예제는 이 최소 다중 GPU 추론 스크립트를 참고하세요. 자세한 내용은 Distributed inference with Accelerate 가이드를 참고하세요.

PyTorch Distributed

Accelerate 없이 torch.distributed를 직접 사용하여 데이터 병렬성을 설정할 수도 있습니다.

분산 환경을 생성하려면 init_process_group을 호출합니다. torchrun은 환경 변수를 통해 각 프로세스의 rank와 전체 프로세스 수(world size)를 설정하므로, 직접 전달할 필요가 없습니다.

각 프로세스는 자체 GPU와 자체 프롬프트가 필요합니다. torchrun은 LOCAL_RANK 환경 변수를 해당 노드에서 프로세스의 GPU 인덱스로 설정하므로, 이를 사용하여 GPU를 선택합니다. 모든 노드에 걸쳐 프로세스를 세는 get_rank를 사용하여 프롬프트를 선택합니다.

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

run_distributed.py로 스크립트를 저장하고 torchrun으로 실행합니다. --nproc-per-node를 GPU 수로 설정합니다. prompts 리스트에는 프로세스당 하나의 프롬프트가 포함되어 있으므로, 사용하는 GPU마다 프롬프트를 추가해야 합니다.

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

Update on GitHub

Diffusers

데이터 병렬

Accelerate

PyTorch Distributed