Cloud GPU에서 xDiT 컨테이너를 사용하여 Diffusion Transformer 모델 서빙

xDiT는 병렬 처리 및 최적화 기법을 사용하여 DiT(Diffusion Transformer) 모델의 추론을 가속화하는 오픈소스 라이브러리입니다. 이러한 기술을 사용하면 까다로운 워크로드에 확장 가능한 멀티 GPU 설정을 사용할 수 있습니다. 이 페이지에서는 Vertex AI에서 xDiT 및 Cloud GPU를 사용하여 DiT 모델을 배포하는 방법을 보여줍니다.

xDiT에 관한 자세한 내용은 xDiT GitHub 프로젝트를 참조하세요.

이점

다음 목록에서는 xDiT를 사용하여 Vertex AI에서 DiT 모델을 서빙할 경우의 주요 이점을 설명합니다.

  • 최대 3배 빠른 생성: 다른 서빙 솔루션에 비해 훨씬 짧은 시간 내에 고해상도 이미지와 동영상을 생성할 수 있습니다.
  • 확장 가능한 멀티 GPU 지원: 최적의 성능을 위해 여러 GPU에 워크로드를 효율적으로 분산합니다.
    • 하이브리드 병렬 처리: xDiT는 통합 시퀀스 병렬 처리, PipeFusion, CFG 병렬 처리, 데이터 병렬 처리와 같은 다양한 병렬 처리 접근 방식을 지원합니다. 이러한 메서드를 고유한 레시피로 결합하여 성능을 최적화할 수 있습니다.
  • 최적화된 단일 GPU 성능: xDiT는 단일 GPU에서도 더 빠른 추론을 제공합니다.
    • GPU 가속: xDiT는 여러 커널 가속 방법을 통합하고 DiTFastAttn의 기술을 사용하여 단일 GPU에서 추론 속도를 높입니다.
  • 간편한 배포: Vertex AI Model Garden에서 원클릭 배포 또는 Colab Enterprise 노트북을 사용하여 빠르게 시작할 수 있습니다.

지원되는 모델

xDiT는 Vertex AI Model Garden의 특정 DiT 모델 아키텍처(예: Flux.1 Schnell, CogVideoX-2b, Wan2.1 텍스트-동영상 모델 변형)에서 사용 가능합니다. DiT 모델에서 Model Garden의 xDiT를 지원하는지 확인하려면 Model Garden에서 모델 카드를 보세요.

다중 GPU 성능을 위한 하이브리드 병렬 처리:

xDiT는 다중 GPU 설정에서 성능을 극대화하기 위해 여러 병렬 처리 기법을 조합하여 사용합니다. 다음 기법이 연동되어 워크로드를 분산하고 리소스 사용률을 최적화합니다.

  • 통합 시퀀스 병렬 처리: 이 기법은 입력 데이터를 여러 GPU로 분할(예: 이미지를 패치로 분할)하여 메모리 사용량을 줄이고 확장성을 개선합니다.
  • PipeFusion: PipeFusion은 DiT 모델을 여러 단계로 나누고 각 단계를 다른 GPU에 할당하여 모델의 여러 부분을 병렬 처리할 수 있도록 합니다.
  • CFG 병렬 처리: 이 기법은 생성된 이미지의 스타일과 콘텐츠를 제어하는 일반적인 방법인 분류기 없는 안내를 사용하여 모델을 최적화합니다. 조건부 및 무조건부 브랜치의 계산을 병렬 처리하여 더 빠른 추론을 실행합니다.
  • 데이터 병렬 처리: 이 메서드는 각 GPU에 전체 모델을 복제하며, 각 GPU는 서로 다른 입력 데이터 배치를 처리하여 시스템의 전반적인 처리량을 늘립니다.

성능 개선에 관한 자세한 내용은 xDiT의 Flux.1 Schnell 또는 CogVideoX-2b에 관한 보고서를 참조하세요. Google은 Vertex AI Model Garden에서 이러한 결과를 재현할 수 있었습니다.

단일 GPU 가속

xDiT 라이브러리는 torch.compile 및 onediff를 사용하여 GPU의 런타임 속도를 향상시켜 단일 GPU 서빙의 이점을 제공합니다. 이러한 기법은 하이브리드 병렬 처리와 함께 사용할 수도 있습니다.

또한 xDiT에는 DiT의 계산 병목 현상을 해결하기 위한 효율적인 어텐션 계산 기법인 DiTFastAttn이 있습니다. 현재 이 기법은 단일 GPU 설정 또는 데이터 병렬 처리와 함께 사용할 수 있습니다.

Model Garden 시작하기

xDiT 최적화 Cloud GPU 서빙 컨테이너는 Vertex AI Model Garden 내에 제공됩니다. 지원되는 모델의 경우 원클릭 배포 또는 Colab Enterprise 노트북 예를 사용할 때 배포에서 이 컨테이너를 사용합니다.

다음 예시에서는 Flux.1-schnell 모델을 사용하여 xDiT 컨테이너에 DiT 모델을 배포하는 방법을 보여줍니다.

클릭 한 번으로 배포 사용

모델 카드를 사용하여 xDiT 컨테이너로 커스텀 Vertex AI 엔드포인트를 배포할 수 있습니다.

  1. 모델 카드 페이지로 이동하고 배포를 클릭합니다.

  2. 사용할 모델 옵션에 대해 배포에 사용할 머신 유형을 선택합니다.

  3. 배포를 클릭하여 배포 프로세스를 시작합니다. 두 가지 이메일 알림이 수신됩니다. 하나는 모델이 업로드될 때 그리고 다른 하나는 엔드포인트가 준비될 때 전송됩니다.

Colab Enterprise 노트북 사용

유연성과 맞춤설정을 위해 Colab Enterprise 노트북 예를 사용하여 Vertex AI SDK for Python을 통해 xDiT 컨테이너로 Vertex AI 엔드포인트를 배포합니다.

  1. 모델 카드 페이지로 이동하고 노트북 열기를 클릭합니다.

  2. Vertex Serving 노트북을 선택합니다. Colab Enterprise에서 노트북이 열립니다.

  3. 노트북을 실행하여 xDiT 컨테이너를 사용하여 모델을 배포하고 엔드포인트에 예측 요청을 전송합니다. 배포의 코드 스니펫은 다음과 같습니다.

import vertexai
from vertexai import model_garden

vertexai.init(project=<YOUR_PROJECT_ID>, location=<REGION>)

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
endpoint = model.deploy()

xDiT 인수

xDiT는 특정 사용 사례에 맞게 성능을 최적화하도록 구성할 수 있는 다양한 서버 인수를 제공합니다. 이러한 인수는 배포 중에 환경 변수로 설정됩니다. 다음 목록은 구성해야 할 수 있는 주요 인수입니다.

모델 구성
  • MODEL_ID(문자열): 로드할 모델 식별자를 지정합니다. 레지스트리 또는 경로의 모델 이름과 일치해야 합니다.
런타임 최적화 인수
  • N_GPUS(정수): 추론에 사용할 GPU 수를 지정합니다. 기본값은 1입니다.
  • WARMUP_STEPS(정수): 추론이 시작되기 전에 필요한 준비 단계 수입니다. 이는 특히 PipeFusion이 사용 설정된 경우 안정적인 성능을 위해 중요합니다. 기본값은 1입니다.
  • USE_PARALLEL_VAE(불리언): 기기 간에 VAE 구성요소를 동시에 로드하여 고해상도 이미지(2,048픽셀 이상)를 효율적으로 처리할 수 있습니다. 이렇게 하면 대형 이미지의 OOM 문제가 방지됩니다. 기본값은 false입니다.
  • USE_TORCH_COMPILE(불리언): torch.compile을 통해 단일 GPU 가속을 사용 설정하여 성능 향상을 위한 커널 수준 최적화를 제공합니다. 기본값은 false입니다.
  • USE_ONEDIFF(불리언): GPU 커널 실행 속도가 최적화되도록 OneDiff 컴파일 가속 기술을 사용 설정합니다. 기본값은 false입니다.
데이터 병렬 인수
  • DATA_PARALLEL_DEGREE(정수): 데이터 병렬 처리 정도를 설정합니다. 사용 중지하려면 비워 둡니다. 또는 선택한 병렬 수준으로 설정합니다.
  • USE_CFG_PARALLEL(불리언): 분할 배치라고도 하는 분류기 없는 안내(CFG)에 병렬 처리 계산을 사용 설정합니다. 사용 설정 시 상수 병렬 처리 수준은 2입니다. 출력 스타일과 콘텐츠를 제어하기 위해 CFG를 사용하는 경우에는 true로 설정합니다. 기본값은 false입니다.
시퀀스 병렬 인수(USP - 통합 시퀀스 병렬 처리)
  • ULYSSES_DEGREE(정수): DeepSpeed-Ulysses 및 Ring-Attention을 결합하는 통합 시퀀스 병렬 방식의 Ulysses 수준을 설정합니다. all-to-all 통신 패턴을 제어합니다. 비워두면 기본값이 사용됩니다.
  • RING_DEGREE(정수): 시퀀스 병렬 처리에서 피어 투 피어 통신의 Ring 수준을 설정합니다. ULYSSES_DEGREE와 함께 작동하여 2D 프로세스 메시를 형성합니다. 비워두면 기본값이 사용됩니다.
텐서 병렬 인수
  • TENSOR_PARALLEL_DEGREE(정수): 텐서 병렬 처리 정도를 설정합니다. 텐서 병렬 처리는 기기당 메모리 비용을 줄이기 위해 기능 차원을 따라 기기 간에 모델 파라미터를 분할합니다. 사용 중지하려면 비워 둡니다.
  • SPLIT_SCHEME(문자열): 기기 간에 모델 텐서를 분할하는 방법을 정의합니다(예: 어텐션 헤드, 숨겨진 차원). 기본 분할 스킴을 사용하려면 비워 둡니다.
Ray 분산 인수
  • USE_RAY(불리언): 여러 노드에서 계산을 확장할 수 있도록 Ray 분산 실행 프레임워크를 사용 설정합니다. 기본값은 false입니다.
  • RAY_WORLD_SIZE(정수): Ray 클러스터의 총 프로세스 수입니다. 기본값은 1입니다.
  • VAE_PARALLEL_SIZE(정수): Ray를 사용할 때 VAE 병렬 처리에 전용으로 사용되는 프로세스 수입니다. 기본값은 0입니다.
  • DIT_PARALLEL_SIZE(정수): Ray를 사용할 때 DiT 백본 병렬 처리에만 사용되는 프로세스 수입니다. 기본값은 0입니다.
PipeFusion 병렬 인수
  • PIPEFUSION_PARALLEL_DEGREE(정수): 확산 모델의 입력 시간 중복 특성을 활용하는 시퀀스 수준 파이프라인 병렬 처리인 PipeFusion의 병렬 처리 정도를 설정합니다. 값이 클수록 병렬 처리가 늘어나지만 더 많은 메모리가 필요합니다. 기본값은 1입니다.
  • NUM_PIPELINE_PATCH(정수): 파이프라인 처리를 위해 시퀀스를 분할하는 패치 수입니다. 비워 두면 자동으로 결정됩니다.
  • ATTN_LAYER_NUM_FOR_PP(문자열): 파이프라인 병렬 처리에 사용할 어텐션 레이어를 지정합니다. 쉼표로 구분하거나(예: '10,9') 또는 공백으로 구분(예: '10 9')할 수 있습니다. 비워 두면 모든 레이어가 사용됩니다.
메모리 최적화 인수
  • ENABLE_MODEL_CPU_OFFLOAD(불리언): 사용하지 않을 때 모델 가중치를 CPU 메모리로 오프로드합니다. 그러면 지연 시간이 증가하는 대신 GPU 메모리 사용량이 줄어듭니다. 기본값은 false입��다.
  • ENABLE_SEQUENTIAL_CPU_OFFLOAD(불리언): 순방향 전달 중에 모델 레이어를 순차적으로 CPU로 오프로드하여 GPU 메모리보다 큰 모델의 추론을 지원합니다. 기본값은 false입니다.
  • ENABLE_TILING(불리언): VAE 구성요소를 한 번에 한 타일씩 디코딩하여 GPU 메모리 사용량을 줄���니다. 이 인수는 큰 이미지나 동영상에 유용하며 메모리 부족 오류를 방지하는 데도 도움이 됩니다. 기본값은 false입니다.
  • ENABLE_SLICING(불리언): 입력 텐서를 VAE 디코딩을 위한 슬라이스로 분할하여 GPU 메모리 사용량을 줄입니다. 기본값은 false입니다.
DiTFastAttn 인수(어텐션 최적화)
  • USE_FAST_ATTN(불리언): 단일 GPU 추론을 위해 DiTFastAttn 가속을 사용 설정하여 입력 시간 감소를 활용해 계산 복잡성을 줄입니다. 기본값은 false입니다.
  • N_CALIB(정수): DiTFastAttn 최적화를 위한 조정 샘플 수입니다. 기본값은 8입니다.
  • THRESHOLD(부동 소수점 수): DiTFastAttn의 시간 유사성 감소를 위한 유사성 기준점입니다. 기본값은 0.5입니다.
  • WINDOW_SIZE(정수): 공간 중복을 줄이기 위해 잔여 캐싱을 사용하는 윈도우 어텐션의 윈도우 크기입니다. 기본값은 64입니다.
  • COCO_PATH(문자열): DiTFastAttn 보정용 COCO 데이터 세트의 경로입니다. USE_FAST_ATTN이 true인 경우 필요합니다. 비워 두면 사용되지 않습니다.
캐시 최적화 인수
  • USE_CACHE(불리언): 중복 계산을 줄이기 위해 일반 캐싱 메커니즘을 사용 설정합니다. 기본값은 false입니다.
  • USE_TEACACHE(불리언): 중간 결과를 캐싱하기 위한 TeaCache 최적화 방법을 사용 설정합니다. 기본값은 false입니다.
  • USE_FBCACHE(불리언): 첫 번째 블록 캐시 최적화 방법을 사용 설정합니다. 기본값은 false입니다.
정밀도 최적화 인수
  • USE_FP8_T5_ENCODER(불리언): T5 텍스트 인코더에 FP8(8비트 부동 소수점) 정밀도를 사용 설정합니다. 그러면 메모리 사용량이 줄어들고 품질에 미치는 영향이 최소화되면서 처리량이 개선됩니다. 기본값은 false입니다.

서빙 맞춤설정

Model Garden은 지원되는 모델에 기본 xDiT 병렬 처리 구성을 제공합니다. Vertex AI SDK for Python을 사용하여 이러한 기본 설정을 검사할 수 있습니다.

'black-forest-labs/FLUX.1-schnell'과 같은 모델의 기본 배포 구성을 보려면 다음 코드 스니펫을 실행하면 됩니다.

import vertexai
from vertexai import model_garden

vertexai.init(project=<YOUR_PROJECT_ID>, location=<REGION>)

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
deploy_options = model.list_deploy_options()


# Example Response
# ['black-forest-labs/flux1-schnell@flux.1-schnell']
# [model_display_name: "Flux1-schnell"
# container_spec {
#   image_uri: "us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-2.ubuntu2204.py310"
#  env {
#    name: "DEPLOY_SOURCE"
#    value: "UI_NATIVE_MODEL"
#  }
#  env {
#    name: "MODEL_ID"
#    value: "gs://vertex-model-garden-restricted-us/black-forest-labs/FLUX.1-schnell"
#  }
#  env {
#    name: "TASK"
#    value: "text-to-image"
#  }
#  env {
#    name: "N_GPUS"
#    value: "2"
#  }
#  env {
#    name: "USE_TORCH_COMPILE"
#    value: "true"
#  }
#  env {
#    name: "RING_DEGREE"
#    value: "2"
#  }
# ..........]

list_deploy_options() 메서드는 xDiT 구성을 정의하는 환경 변수(env)를 포함한 컨테이너 사양을 반환합니다.

병렬 처리 전략을 맞춤설정하려면 모델을 배포할 때 이러한 환경 변수를 재정의하면 됩니다. 다음 예시에서는 2-GPU 설정의 RING_DEGREE 및 ULYSSES_DEGREE를 수정하여 병렬 처리 방식을 변경하는 방법을 보여줍니다.

import vertexai
from vertexai import model_garden

# Replace with your project ID and region
vertexai.init(project="<YOUR_PROJECT_ID>", location="<REGION>")

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")

# Custom environment variables to override default settings
# This example sets N_GPUS as 2, so RING_DEGREE * ULYSSES_DEGREE must equal 2
container_env_vars = {
    "N_GPUS": "2",
    "RING_DEGREE": "1",
    "ULYSSES_DEGREE": "2"
    # Add other environment variables to customize here
}

machine_type = "a3-highgpu-2g"
accelerator_type = "NVIDIA_H100_80GB"
accelerator_count = 2

# Deploy the model with the custom environment variables
endpoint = model.deploy(
    machine_type=machine_type,
    accelerator_type=accelerator_type,
    accelerator_count=accelerator_count,
  container_env_vars=container_env_vars
)

각 환경 변수에 대한 자세한 내용은 'xDiT 관련 인수 이해' 섹션을 참조하세요. 병렬 처리 수준(예: PIPEFUSION_PARALLEL_DEGREE, ULYSSES_DEGREE, RING_DEGREE, USE_CFG_PARALLEL)의 곱이 총 GPU 수(N_GPUS)와 같은지 확인합니다.

다양한 모델의 서빙 레시피 및 구성에 대한 자세한 예시는 xDiT 공식 문서를 참조하세요. Model Garden SDK에 대한 자세한 내용은 문서를 참조하세요.