device_map를 사용하여 모델 구성 요소를 여러 GPU에 분할

device_map 인수는 모델 구성 요소를 개별 GPU에 배치합니다. 모델이 단일 GPU에 맞지 않고 GPU 간에 빠른 인터커넥트가 없는 경우 사용하세요. 다른 전략과의 비교는 분산 추론을, device_map의 추가 옵션은 디바이스 배치를 참고하세요.

아래 예제는 각 구성 요소를 필요한 단계에서만 로드하고, device_map을 사용하여 큰 구성 요소를 두 개의 GPU에 분산 배치합니다.

프롬프트 인코딩

device_map="balanced"을 설정하여 텍스트 인코더를 개별 GPU에 배치합니다. max_memory를 사용하여 각 GPU가 사용하는 메모리 양을 제한합니다. transformer와 vae를 None으로 설정하여 텍스트 인코더만 로드되도록 합니다.

from diffusers import FluxPipeline
import torch

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=None,
    vae=None,
    device_map="balanced",
    max_memory={0: "16GB", 1: "16GB"},
    dtype=torch.bfloat16
)
with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=512
    )

텍스트 임베딩을 확보한 후에는 트랜스포머를 위한 공간을 확보하기 위해 GPU에서 텍스트 인코더를 제거합니다.

import gc

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline

flush()

sharded transformer로 노이즈 제거

device_map="auto"를 설정하면 Accelerate가 transformer를 배치합니다. 첫 번째 GPU가 가득 차면 다음 GPU로 이동하며, GPU 메모리가 부족할 때만 CPU 또는 디스크를 사용합니다. 두 개의 작은 GPU를 사용하는 경우 transformer는 두 GPU에 걸쳐 분할 배치됩니다. CPU 또는 디스크에 있는 가중치는 추론 속도를 늦춥니다.

from diffusers import AutoModel
import torch

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    device_map="auto",
    dtype=torch.bfloat16
)

transformer.hf_device_map을 확인하여 transformer가 여러 장치에 어떻게 분배되어 있는지 살펴볼 수 있습니다. device_map으로 로드된 pipeline에도 각 구성 요소가 어디에 배치되어 있는지 보여주는 hf_device_map 속성이 있습니다.

transformer만 포함된 pipeline을 로드하고, output_type="latent"를 설정하여 이미지 대신 latents를 반환하도록 합니다.

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    transformer=transformer,
    dtype=torch.bfloat16
)

print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images

잠재 변수 디코딩

메모리에서 pipeline과 transformer를 제거하고, latents를 디코딩하기 위해 VAE를 로드합니다. VAE는 하나의 GPU에 충분히 들어갈 만큼 작습니다.

import torch
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor

del pipeline.transformer
del pipeline
del transformer

flush()

vae = AutoencoderKL.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder="vae", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

with torch.no_grad():
    print("Running decoding.")
    latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

    image = vae.decode(latents, return_dict=False)[0]
    image = image_processor.postprocess(image, output_type="pil")
    image[0].save("split_transformer.png")

주어진 단계에서 필요한 모델을 선택적으로 로드하고 언로드하며, 가장 큰 모델을 여러 GPU에 걸쳐 샤딩하면 소비자용 GPU에서도 대규모 모델로 추론을 실행할 수 있습니다.

Update on GitHub

Diffusers

device_map를 사용하여 모델 구성 요소를 여러 GPU에 분할

프롬프트 인코딩

sharded transformer로 노이즈 제거

잠재 변수 디코딩