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

xDiT는 병렬 처리 및 최적화 기법을 사용하여 확산 변환기 (DiT) 모델의 추론을 가속화하는 오픈소스 라이브러리입니다. 이러한 기법을 사용하면 까다로운 워크로드에 확장 가능한 멀티 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는 Flux.1 Schnell, CogVideoX-2b와 같은 Vertex AI Model Garden의 특정 DiT 모델 아키텍처에서 사용할 수 있습니다. Model Garden에서 DiT 모델이 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.compileonediff를 사용하여 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 노트북 예시를 통해 Python용 Vertex AI SDK를 사용하여 xDiT 컨테이너로 Vertex AI 엔드포인트를 배포합니다.

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

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

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

XDIT_DOCKER_URI=us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-1.ubuntu2204.py310

serving_env = {
    "MODEL_ID": "black-forest-labs/FLUX.1-schnell",
    "TASK": "text-to-image",
    "DEPLOY_SOURCE": "notebook",
    "N_GPUS": "2",
    "ULYSSES_DEGREE": "1",
    "RING_DEGREE": "2",
    "PIPEFUSION_PARALLEL_DEGREE": "1",
    "USE_TORCH_COMPILE": "true",
}

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=XDIT_DOCKER_URI,
    serving_container_ports=[7080],
    serving_container_predict_route="/predict",
    serving_container_health_route="/health",
    serving_container_environment_variables=serving_env,
)

model.deploy(
    endpoint=endpoint,
    machine_type="a3-highgpu-2g",
    accelerator_type="NVIDIA_H100_80GB",
    accelerator_count=2,
    deploy_request_timeout=1800,
    service_account=SERVICE_ACCOUNT,
)

환경 변수

  • MODEL_ID: 배포할 DiT 모델의 ID를 지정합니다(예: 'black-forest-labs/FLUX.1-schnell').
  • TASK: 모델이 실행하는 태스크를 정의합니다(예: 'text-to-image-flux-xdit').
  • N_GPUS: 추론에 사용할 GPU 수를 설정합니다.
  • ULYSSES_DEGREE, RING_DEGREE, PIPEFUSION_PARALLEL_DEGREE: xDiT에서 사용하는 동시 실행 기법을 제어합니다. 각 인수에 관한 자세한 내용은 xDiT 인수를 참고하세요.
  • USE_TORCH_COMPILE: torch.compile를 사용하여 단일 GPU 가속을 사용 설정합니다.

xDiT 인수

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

  • N_GPUS (정수): 추론에 사용할 GPU 수를 지정합니다. 기본값은 1입니다.
  • ENABLE_TILING (불리언): VAE 구성요소를 한 번에 하나씩 디코딩하여 GPU 메모리 사용량을 줄입니다. 이 인수는 큰 이미지나 동영상에 유용하며 메모리 부족 오류를 방지합니다. 기본값은 false입니다.
  • ENABLE_SLICING (불리언): 입력 텐서를 VAE 디코딩을 위한 슬라이스로 분할하여 GPU 메모리 사용량을 줄입니다. 기본값은 false입니다.
  • USE_TORCH_COMPILE (불리언): torch.compile를 통해 단일 GPU 가속을 사용 설정하여 컴파일 속도를 개선합니다. 기본값은 false입니다.
  • PIPEFUSION_PARALLEL_DEGREE (정수): PipeFusion의 동시 로드 정도를 설정합니다. 값이 클수록 동시 로드가 늘어나지만 메모리가 더 많이 필요할 수 있습니다. 기본값은 1입니다.
  • WARMUP_STEPS (정수): PipeFusion이 사용 설정된 경우 이 인수는 추론이 시작되기 전에 필요한 워밍업 단계 수를 지정합니다. 기본값은 0입니다.
  • ULYSSES_DEGREE (정수): Ulysses 각도를 설정합니다. 기본값은 1입니다.
  • RING_DEGREE (정수): 링 각도를 설정합니다. 기본값은 1입니다.
  • USE_CFG_PARALLEL (불리언): DiT 모델의 출력을 제어하는 데 사용되는 기법인 분류기 없는 안내 (CFG)에 병렬 계산을 사용 설정합니다. 사용 설정하면 상수 동시 로드 수준은 2입니다. CFG를 사용할 때 'true'로 설정합니다. 기본값은 false입니다.
  • USE_PARALLEL_VAE (불리언): VAE 구성요소를 병렬화하여 고해상도 이미지 (2,048픽셀 초과)를 효율적으로 처리할 수 있습니다. 기본값은 false입니다.

인수의 전체 목록은 xDiT GitHub 프로젝트의 xFuserArgs 클래스를 참고하세요.