Cloud TPU에서 Hex-LLM 프리미엄 컨테이너를 사용하여 개방형 모델 서빙

XLA를 사용하는 고효율 대규모 언어 모델(LLM) 서빙인 Hex-LLM은 Cloud TPU 하드웨어용으로 설계되고 최적화된 Vertex AI LLM 서빙 프레임워크입니다. Hex-LLM은 지속적 일괄 처리 및 PagedAttention과 같은 LLM 서빙 기술을 XLA 및 Cloud TPU에 맞게 조정된 Vertex AI 최적화와 결합합니다. 오픈소스 모델용 Cloud TPU에 서빙되는 고효율 및 저비용 LLM입니다.

Hex-LLM은 모델 플레이그라운드, 클릭 한 번으로 배포, 노트북을 통해 Model Garden에서 사용할 수 있습니다.

기능

Hex-LLM은 XLA 및 Cloud TPU에 대한 Google의 자체 최적화를 지원하는 오픈소스 프로젝트를 기반으로 합니다. Hex-LLM은 자주 사용되는 LLM을 서빙할 때 높은 처리량과 짧은 지연 시간을 달성합니다.

Hex-LLM에는 다음 최적화가 포함됩니다.

  • 모델이 많은 수의 동시 요청으로 하드웨어를 최대한 활용할 수 있도록 지원하는 토큰 기반 연속 일괄 처리 알고리즘입니다.
  • XLA에 최적화된 어텐션 커널의 완전 재작성
  • 여러 Cloud TPU 칩에서 LLM을 효율적으로 실행하기 위해 고도로 최적화된 가중치 샤딩 메서드를 사용하여 유연하고 구성 가능한 데이터 동시 로드 및 텐서 동시 로드 전략입니다.

Hex-LLM은 다양한 밀도 및 희소 LLM을 지원합니다.

  • Gemma 2B 및 7B
  • Gemma 2 9B 및 27B
  • Llama 2 7B, 13B, 70B
  • Llama 3 8B 및 70B
  • Mistral 7B 및 Mixtral 8x7B

Hex-LLM은 다음과 같은 다양한 기능도 제공합니다.

  • Hex-LLM은 단일 컨테이너에 포함되어 있습니다. Hex-LLM은 API 서버, 추론 엔진, 지원되는 모델을 배포할 단일 Docker 이미지로 패키징합니다.
  • Hugging Face 모델 형식과 호환됩니다. Hex-LLM은 로컬 디스크, Hugging Face Hub, Cloud Storage 버킷에서 Hugging Face 모델을 로드할 수 있습니다.
  • 비트 및 바이트AWQ를 사용한 양자화
  • 동적 LoRA 로드. Hex-LLM은 서빙 중에 요청 인수를 읽어 LoRA 가중치를 로드할 수 있습니다.

Model Garden 시작하기

Hex-LLM Cloud TPU 서빙 컨테이너는 Model Garden에 통합되어 있습니다. 플레이그라운드, 클릭 한 번으로 배포, 다양한 모델의 Colab Enterprise 노트북 예시를 통해 이 서빙 기술에 액세스할 수 있습니다.

플레이그라운드 사용

Model Garden 플레이그라운드는 모델 카드에서 요청을 전송하여 연결할 수 있는 사전 배포된 Vertex AI 엔드포인트입니다.

  1. 프롬프트를 입력하고 원하는 경우 요청에 대한 인수를 포함합니다.

  2. 제출을 클릭하여 모델 응답을 빠르게 가져옵니다.

Gemma와 함께 사용해 보세요!

클릭 한 번으로 배포 사용

모델 카드를 사용하여 Hex-LLM으로 커스텀 Vertex AI 엔드포인트를 배포할 수 있습니다.

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

  2. 사용할 모델 변형에 대해 배포할 Cloud TPU v5e 머신 유형을 선택합니다.

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

Colab Enterprise 노트북 사용

유연성과 맞춤설정을 위해 Colab Enterprise 노트북 예시를 통해 Vertex AI SDK for Python을 사용하여 Hex-LLM으로 Vertex AI 엔드포인트를 배포할 수 있습니다.

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

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

  3. 노트북을 실행하여 Hex-LLM을 사용해 모델을 배포하고 엔드포인트에 예측 요청을 전송합니다. 배포의 코드 스니펫은 다음과 같습니다.

hexllm_args = [
    f"--model=google/gemma-2-9b-it",
    f"--tensor_parallel_size=4",
    f"--hbm_utilization_factor=0.8",
    f"--max_running_seqs=512",
]
hexllm_envs = {
    "PJRT_DEVICE": "TPU",
    "MODEL_ID": "google/gemma-2-9b-it",
    "DEPLOY_SOURCE": "notebook",
}
model = aiplatform.Model.upload(
    display_name="gemma-2-9b-it",
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=[
        "python", "-m", "hex_llm.server.api_server"
    ],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=hexllm_envs,
    serving_container_shared_memory_size_mb=(16 * 1024),
    serving_container_deployment_timeout=7200,
)

endpoint = aiplatform.Endpoint.create(display_name="gemma-2-9b-it-endpoint")
model.deploy(
    endpoint=endpoint,
    machine_type="ct5lp-hightpu-4t",
    deploy_request_timeout=1800,
    service_account="<your-service-account>",
    min_replica_count=1,
    max_replica_count=1,
)

커스텀 서빙을 위해 다음 Hex-LLM 서버 실행 인수를 수정할 수 있습니다.

  • --model: 로드할 모델입니다. Hugging Face 모델 ID, 로컬 절대 경로 또는 Cloud Storage 버킷 경로를 지정할 수 있습니다.
  • --tokenizer: 로드할 토크나이저입니다. Hugging Face 모델 ID, 로컬 절대 경로 또는 Cloud Storage 버킷 경로일 수 있습니다. 기본값은 --model과 동일합니다.
  • --enable_jit: JIT 모드를 사용 설정할지 여부입니다. 기본값은 True입니다.
  • --enable_lora: LoRA 로딩 모드를 사용 설정할지 여부입니다. 기본값은 False입니다.
  • --max_lora_rank: 요청에 정의된 LoRA 어댑터에 지원되는 최대 LoRA 순위입니다. 기본값은 16입니다.
  • --data_parallel_size: 데이터 동시 복제본 수입니다. 기본값은 1입니다.
  • --tensor_parallel_size: 텐서 동시 복제본 수입니다. 기본값은 1입니다.
  • --max_running_seqs: 서버가 동시에 처리할 수 있는 최대 요청 수입니다. 이 인수가 클수록 서버가 달성할 수 있는 처리량이 높아지지만 지연 시간에 부정적인 영향을 미칠 수 있습니다. 기본값은 256입니다.
  • --hbm_utilization_factor: 모델 가중치가 로드된 후 KV-캐시에 할당할 수 있는 무료 Cloud TPU HBM의 비율입니다. 이 인수를 더 낮은 값으로 설정하면 Cloud TPU HBM의 메모리 부족을 효과적으로 방지할 수 있습니다. 기본값은 0.9입니다.
  • --seed: 모든 난수 생성기를 초기화하기 위한 시드입니다. 이 인수를 변경하면 동일한 프롬프트에 대해 생성된 출력에 영향을 미칠 수 있습니다. 기본값은 0입니다.

Cloud TPU 할당량 요청

Model Garden에서 기본 할당량은 us-west1 리전의 Cloud TPU v5e 칩 4개입니다. 이 할당량은 클릭 한 번으로 배포 및 Colab Enterprise 노트북 배포에 적용됩니다. 추가 할당량을 요청하려면 할당량 상향 요청을 참조하세요.