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
  • Llama-3.1 8B 및 70B
  • Llama-3.2 1B 및 3B
  • Llama-3.3 70B
  • Llama-Guard-3 1B 및 8B
  • Llama-4 Scout-17B-16E
  • Mistral 7B
  • Mixtral 8x7B 및 8x22B
  • Phi-3 mini 및 medium
  • Phi-4, Phi-4 reasoning 및 reasoning plus
  • Qwen-2 0.5B, 1.5B, 7B
  • Qwen-2.5 0.5B, 1.5B, 7B, 14B, 32B

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

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

고급 기능

Hex-LLM은 다음과 같은 고급 기능을 지원합니다.

  • 멀티 호스트 서빙
  • 분리형 서빙 [실험용]
  • 프리픽스 캐싱
  • 4비트 양자화 지원

멀티 호스트 서빙

이제 Hex-LLM은 멀티 호스트 TPU 슬라이스를 사용한 모델 서빙을 지원합니다. 이 기능을 사용하면 최대 8개의 v5e 코어만 포함된 단일 호스트 TPU VM에 로드할 수 없는 대규모 모델을 서빙할 수 있습니다.

이 기능을 사용 설정하려면 Hex-LLM 컨테이너 인수에 --num_hosts를 설정하고 Vertex AI SDK 모델 업로드 요청에서 --tpu_topology를 설정하세요. 다음 예시에서는 Llama 3.1 70B bfloat16 모델을 서빙하기 위해 TPU 4x4 v5e 토폴로지를 사용하여 Hex-LLM 컨테이너를 배포하는 방법을 보여줍니다.

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Meta-Llama-3.1-70B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=16",
    "--num_hosts=4",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    tpu_topology="4x4",
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

멀티 호스트 TPU 토폴로지로 Hex-LLM 컨테이너를 배포하는 엔드 투 엔드 튜토리얼은 Vertex AI Model Garden - Llama 3.1(배포) 노트북을 참조하세요.

일반적으로 멀티 호스트 서빙을 사용 설정하는 데 필요한 변경사항은 다음과 같습니다.

  1. --tensor_parallel_size 인수를 TPU 토폴로지 내 전체 코어 수로 설정합니다.
  2. --num_hosts 인수를 TPU 토폴로지 내 호스트 수로 설정합니다.
  3. Vertex AI SDK 모델 업로드 API에서 --tpu_topology를 설정합니다.

분리형 서빙 [실험용]

Hex-LLM은 이제 실험적 기능으로 분리형 서빙을 지원합니다. 이 기능은 단일 호스트 설정에서만 사용 설정할 수 있으며, 현재 성능 조정 중입니다.

분리형 서빙은 요청별 최초 토큰 생성 시간(TTFT)과 토큰당 출력 시간(TPOT), 그리고 전체 서빙 처리량을 균형 있게 조정하는 효과적인 방법입니다. 이 방식은 사전 입력 단계와 디코딩 단계를 서로 간섭하지 않도록 별도의 워크로드로 분리합니다. 이 방법은 특히 엄격한 지연 시간 요구사항이 설정된 시나리오에서 유용합니다.

이 기능을 사용 설정하려면 Hex-LLM 컨테이너 인수에서 --disagg_topo를 설정하세요. 다음 예시는 TPU v5e-8에서 Llama 3.1 8B bfloat16 모델을 서빙하도록 Hex-LLM 컨테이너를 배포하는 방법을 보여줍니다.

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=2",
    "--disagg_topo=3,1",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

--disagg_topo 인수는 "number_of_prefill_workers,number_of_decode_workers" 형식의 문자열을 허용합니다. 앞선 예시에서는 "3,1"로 설정하여 사전 입력 작업자 3개와 디코딩 작업자 1개를 구성했습니다. 각 작업자는 TPU v5e 코어 2개를 사용합니다.

프리픽스 캐싱

프리픽스 캐싱은 회사 전체 공통 서문, 일반적인 시스템 지침, 다중 턴 대화 기록과 같이 프롬프트 시작 부분의 콘텐츠가 동일한 경우에 최초 토큰 생성 시간(TTFT)을 줄여 줍니다. 동일한 입력 토큰을 반복적으로 처리하는 대신 Hex-LLM은 처리된 입력 토큰 계산 결과를 임시 캐시에 보관하여 TTFT를 개선할 수 있습니다.

이 기능을 사용 설정하려면 Hex-LLM 컨테이너 인수에서 --enable_prefix_cache_hbm을 설정하세요. 다음 예시는 TPU v5e-8에서 Llama 3.1 8B bfloat16 모델을 서빙하도록 Hex-LLM 컨테이너를 배포하는 방법을 보여줍니다.

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--hbm_utilization_factor=0.9",
    "--enable_prefix_cache_hbm",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

Hex-LLM은 일정 길이(기본값: 512 토큰, prefill_len_padding으로 구성 가능)를 초과하는 프롬프트에 대해 성능을 최적화하기 위해 프리픽스 캐싱을 사용합니다. 캐시 적중은 이 값의 배수 단위로 발생하며, 캐시된 토큰 수는 항상 prefill_len_padding의 배수가 됩니다. 채팅 완성 API 응답의 usage.prompt_tokens_detailscached_tokens 필드는 프롬프트 토큰 중 캐시 적중이 몇 개인지 나타냅니다.

"usage": {
  "prompt_tokens": 643,
  "total_tokens": 743,
  "completion_tokens": 100,
  "prompt_tokens_details": {
    "cached_tokens": 512
  }
}

청크 분할 사전 입력

청크 분할 사전 입력은 요청 사전 입력을 더 작은 청크로 분할하고 사전 입력과 디코딩을 하나의 배치 단계로 혼합합니다. Hex-LLM은 청크 분할 사전 입력을 구현하여 최초 토큰 생성 시간(TTFT)과 토큰당 출력 시간(TPOT)을 균형 있게 조정하고 처리량을 개선합니다.

이 기능을 사용 설정하려면 Hex-LLM 컨테이너 인수에서 --enable_chunked_prefill을 설정하세요. 다음 예시는 TPU v5e-8에서 Llama 3.1 8B 모델을 서빙하도록 Hex-LLM 컨테이너를 배포하는 방법을 보여줍니다.

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--hbm_utilization_factor=0.9",
    "--enable_chunked_prefill",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

4비트 양자화 지원

양자화는 추론 실행 시 일반적으로 가중치 또는 활성화 값을 사용하는 BF16 또는 FP32 대신 INT8이나 INT4 같은 저정밀도 데이터 유형으로 표현하여 계산 및 메모리 비용을 줄이는 기술입니다.

Hex-LLM은 INT8 가중치 전용 양자화를 지원합니다. 또한 AWQ 제로 포인트 양자화를 사용하여 INT4 가중치로 양자화된 모델도 지원합니다. Hex-LLM은 Mistral, Mixtral, Llama 모델 제품군의 INT4 변형 모델을 지원합니다.

양자화된 모델을 서빙할 때는 추가적인 플래그가 필요하지 않습니다.

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

Colab Enterprise 노트북 예시에는 다음이 포함됩니다.

서버 인수 및 환경 변수 구성

다음 인수를 설정하여 Hex-LLM 서버를 실행할 수 있습니다. 인수는 사용자가 의도한 목적과 요구사항에 가장 잘 맞도록 조정할 수 있습니다. 인수는 원클릭 배포용으로 미리 정의되어 있어 가장 간단한 배포 경험을 제공합니다. 인수를 맞춤설정하려면 노트북 예시를 참조하여 그에 따라 인수를 설정하면 됩니다.

모델

  • --model: 로드할 모델입니다. Hugging Face 모델 ID, Cloud Storage 버킷 경로(gs://my-bucket/my-model), 또는 로컬 경로를 지정할 수 있습니다. 모델 아티팩트는 Hugging Face 형식을 따라야 하며, 모델 가중치는 safetensors 파일을 사용해야 합니다. Llama, Gemma 2, Mistral/Mixtral에 대해서는 BitsAndBytes int8 및 AWQ 양자화 모델 아티팩트가 지원됩니다.
  • --tokenizer: 로드할 토크나이저입니다. Hugging Face 모델 ID, Cloud Storage 버킷 경로(gs://my-bucket/my-model), 또는 로컬 경로일 수 있습니다. 이 인수가 설정되지 않으면 기본적으로 --model의 값이 사용됩니다.
  • --tokenizer_mode: 토크나이저 모드입니다. 가능한 값은 ["auto", "slow"]입니다. 기본값은 "auto"입니다. "auto"로 설정하면 가능한 경우 빠른 토크나이저가 사용됩니다. 느린 토크나이저는 Python으로 작성되어 Transformers 라이브러리에서 제공되며, 빠른 토크나이저는 Rust로 작성되어 Tokenizers 라이브러리에서 제공되며 성능 향상을 지원합니다. 자세한 내용은 Hugging Face 문서를 참조하세요.
  • --trust_remote_code: Hugging Face 모델 저장소에 정의된 원격 코드 파일을 허용할지 여부입니다. 기본값은 False입니다.
  • --load_format: 로드할 모델 체크포인트의 형식입니다. 가능한 값은 ["auto", "dummy"]입니다. 기본값은 "auto"입니다. "auto"로 설정하면 모델 가중치가 safetensors 형식으로 로드됩니다. "dummy"로 설정하면 모델 가중치가 무작위로 초기화됩니다. "dummy" 설정은 실험 목적으로 유용합니다.
  • --max_model_len: 모델이 서빙할 최대 컨텍스트 길이(입력 길이 + 출력 길이)입니다. 기본값은 Hugging Face 형식의 모델 구성 파일(config.json)에서 읽어옵니다. 최대 컨텍스트 길이가 클수록 TPU 메모리를 더 많이 요구합니다.
  • --sliding_window: 설정 시, 이 인수는 모델이 기본적으로 사용하는 슬라이딩 윈도우 어텐션의 윈도우 크기를 재정의합니다. 이 인수를 더 큰 값으로 설정하면 어텐션 메커니즘이 더 많은 토큰을 포함하게 되어 표준 셀프 어텐션과 유사한 효과를 냅니다. 이 인수는 실험용으로만 사용됩니다. 일반적인 사용 사례에서는 모델의 원래 윈도우 크기를 사용하는 것이 좋습니다.
  • --seed: 모든 난수 생성기를 초기화하기 위한 시드입니다. 이 인수를 변경하면 동일한 프롬프트에 대해 샘플링되는 다음 토큰이 달라져 출력에 영향을 미칠 수 있습니다. 기본값은 0입니다.

추론 엔진

  • --num_hosts: 실행할 호스트 수입니다. 기본값은 1입니다. 자세한 내용은 TPU v5e 구성 문서를 참조하세요.
  • --disagg_topo: 실험적 기능인 분리형 서빙에서 사전 입력 작업자와 디코딩 작업자의 수를 정의합니다. 기본값은 None입니다. 인수는 "number_of_prefill_workers,number_of_decode_workers" 형식을 따릅니다.
  • --data_parallel_size: 데이터 동시 복제본 수입니다. 기본값은 1입니다. 이 값을 1에서 N으로 늘리면 지연 시간은 유지하면서 처리량은 대략 N배 향상됩니다.
  • --tensor_parallel_size: 텐서 병렬 복제본 수입니다. 기본값은 1입니다. 텐서 병렬 복제본 수를 늘리면 행렬 곱셈 크기를 줄여 속도를 높이므로 일반적으로 지연 시간이 개선됩니다.
  • --worker_distributed_method: 작업자 실행 시 사용할 분산 방식입니다. mp(멀티 프로세싱 모듈) 또는 ray(Ray 라이브러리) 중 선택할 수 있습니다. 기본값은 mp입니다.
  • --enable_jit: JIT(Just-in-Time Compilation) 모드를 사용 설정할지 여부입니다. 기본값은 True입니다. --no-enable_jit로 설정하면 사용 중지됩니다. JIT 모드로 사용 설정하면 초기 컴파일에 시간이 더 들지만, 전체적인 추론 성능이 향상됩니다. 일반적으로 성능 개선 효과가 초기 오버헤드보다 더 큽니다.
  • --warmup: 서버 초기화 시 샘플 요청으로 워밍업할지 여부입니다. 기본값은 True입니다. --no-warmup으로 설정하면 사용 중지됩니다. 초기 요청은 더 무거운 컴파일을 유발하여 느리기 때문에 워밍업을 권장합니다.
  • --max_prefill_seqs: 한 반복 작업에서 사전 입력으로 예약할 수 있는 최대 시퀀스 수입니다. 기본값은 1입니다. 이 값이 클수록 서버가 달성할 수 있는 처리량은 높아지지만 지연 시간에 부정적인 영향을 미칠 수 있습니다.
  • --prefill_seqs_padding: 서버가 사전 입력 배치 크기를 이 값의 배수로 패딩합니다. 기본값은 8입니다. 이 값을 늘리면 모델 재컴파일 횟수가 줄어들지만 불필요한 계산과 추론 오버헤드는 증가합니다. 최적 설정은 요청 트래픽에 따라 달라집니다.
  • --prefill_len_padding: 서버가 시퀀스 길이를 이 값의 배수로 패딩합니다. 기본값은 512입니다. 이 값을 늘리면 모델 재컴파일 횟수가 줄어들지만 불필요한 계산과 추론 오버헤드는 증가합니다. 최적 설정은 요청의 데이터 분포에 따라 달라집니다.
  • --max_decode_seqs/--max_running_seqs: 한 반복 작업에서 디코딩으로 예약할 수 있는 최대 시퀀스 수입니다. 기본값은 256입니다. 이 값이 클수록 서버가 달성할 수 있는 처리량은 높아지지만 지연 시간에 부정적인 영향을 미칠 수 있습니다.
  • --decode_seqs_padding: 서버가 디코딩 배치 크기를 이 값의 배수로 패딩합니다. 기본값은 8입니다. 이 값을 늘리면 모델 재컴파일 횟수가 줄어들지만 불필요한 계산과 추론 오버헤드는 증가합니다. 최적 설정은 요청 트래픽에 따라 달라집니다.
  • --decode_blocks_padding: 서버가 디코딩 중 시퀀스의 키-값 캐시(KV 캐시)에 사용되는 메모리 블록 수를 이 값의 배수로 패딩합니다. 기본값은 128입니다. 이 값을 늘리면 모델 재컴파일 횟수가 줄어들지만 불필요한 계산과 추론 오버헤드는 증가합니다. 최적 설정은 요청의 데이터 분포에 따라 달라집니다.
  • --enable_prefix_cache_hbm: HBM에서 프리픽스 캐싱을 사용 설정할지 여부입니다. 기본값은 False입니다. 이 인수를 설정하면 이전 요청에서 공유된 프리픽스 계산 결과를 재사용해 성능을 개선할 수 있습니다.
  • --enable_chunked_prefill: 청크 분할 사전 입력을 사용 설정할지 여부입니다. 기본값은 False입니다. 이 인수를 설정하면 더 긴 컨텍스트 길이를 지원하며 성능을 개선할 수 있습니다.

메모리 관리

  • --hbm_utilization_factor: 모델 가중치가 로드된 후 KV 캐시에 할당할 수 있는 Cloud TPU 고대역폭 메모리(HBM)의 여유 메모리 비율입니다. 기본값은 0.9입니다. 이 인수를 더 높은 값으로 설정하면 KV 캐시 크기가 증가하여 처리량이 개선될 수 있지만, 초기화 및 런타임 중 Cloud TPU HBM 부족 위험이 커집니다.
  • --num_blocks: KV 캐시에 할당할 기기 블록 수입니다. 이 인수가 설정되면 서버는 --hbm_utilization_factor를 무시합니다. 이 인수가 설정되지 않으면 서버는 HBM 사용량을 프로파일링하고 --hbm_utilization_factor에 따라 할당할 기기 블록 수를 계산합니다. 이 인수를 더 높은 값으로 설정하면 KV 캐시 크기가 증가하여 처리량이 개선될 수 있지만, 초기화 및 런타임 중 Cloud TPU HBM 부족 위험이 커집니다.
  • --block_size: 하나의 블록에 저장된 토큰 수입니다. 가능한 값은 [8, 16, 32, 2048, 8192]입니다. 기본값은 32입니다. 이 인수를 더 큰 값으로 설정하면 블록 관리 오버헤드는 줄어들지만, 메모리 낭비가 늘어납니다. 정확한 성능 영향은 실험적으로 결정해야 합니다.

동적 LoRA

  • --enable_lora: Cloud Storage에서 동적 LoRA 어댑터 로딩을 사용 설정할지 여부입니다. 기본값은 False입니다. 이 기능은 Llama 모델 제품군에서 지원됩니다.
  • --max_lora_rank: 요청에 정의된 LoRA 어댑터에 지원되는 최대 LoRA 순위입니다. 기본값은 16입니다. 이 인수를 더 높은 값으로 설정하면 서버에 사용할 수 있는 LoRA 어댑터의 유연성이 높아지지만, LoRA 가중치에 할당되는 Cloud TPU HBM 사용량이 늘어나고 처리량이 감소합니다.
  • --enable_lora_cache: 동적 LoRA 어댑터 캐싱을 사용 설정할지 여부입니다. 기본값은 True입니다. --no-enable_lora_cache로 설정하면 사용 중지됩니다. 캐싱을 활성화하면 이전에 사용한 LoRA 어댑터 파일을 다시 다운로드할 필요가 없어 성능이 개선됩니다.
  • --max_num_mem_cached_lora: TPU 메모리 캐시에 저장할 수 있는 LoRA 어댑터의 최대 개수입니다. 기본값은 16입니다. 이 인수를 더 큰 값으로 설정하면 캐시 적중 가능성이 높아지지만 Cloud TPU HBM 사용량이 늘어납니다.

다음 환경 변수를 사용하여 서버를 구성할 수도 있습니다.

  • HEX_LLM_LOG_LEVEL: 생성되는 로깅 정보의 양을 제어합니다. 기본값은 INFO입니다. 로깅 모듈에서 정의된 표준 Python 로깅 수준 중 하나로 설정할 수 있습니다.
  • HEX_LLM_VERBOSE_LOG: 상세 로깅 출력을 사용 설정할지 여부입니다. 허용되는 값은 true 또는 false입니다. 기본값은 false입니다.

서버 인수 조정

서버 인수는 서로 연관되어 있으며, 함께 서빙 성능에 영향을 줍니다. 예를 들어 --max_model_len=4096으로 더 큰 값을 설정하면 TPU 메모리 사용량이 증가하므로 더 큰 메모리 할당이 필요하고 작업 배치는 줄어듭니다. 또한 일부 인수는 사용 사례에 의해 결정되지만, 다른 인수들은 조정할 수 있습니다. 다음은 Hex-LLM 서버를 구성하기 위한 워크플로입니다.

  1. 관심 있는 모델 계열과 모델 변형을 결정합니다. 예를 들어 Llama 3.1 8B Instruct를 선택할 수 있습니다.
  2. 모델 크기와 정밀도를 기반으로 필요한 TPU 메모리의 하한선을 추정합니다(model_size * (num_bits / 8)). 예를 들어 8B 모델과 bfloat16 정밀도의 경우 필요한 TPU 메모리의 하한선은 8 * (16 / 8) = 16 GB입니다.
  3. 필요한 TPU v5e 칩 수를 추정합니다. 각 v5e 칩은 16GB를 제공합니다(tpu_memory / 16). 예를 들어 8B 모델과 bfloat16 정밀도에서는 1개 칩 이상이 필요합니다. 1칩, 4칩, 8칩 구성 중 1개를 초과하는 최소 구성은 4칩 구성인 ct5lp-hightpu-4t입니다. 이후 --tensor_parallel_size=4로 설정할 수 있습니다.
  4. 사용 사례에 맞는 최대 컨텍스트 길이(입력 길이 + 출력 길이)를 결정합니다. 예: 4096. 이후 --max_model_len=4096으로 설정할 수 있습니다.
  5. 모델, 하드웨어, 서버 구성을 고려하여 KV 캐시에 할당할 수 있는 TPU 여유 메모리를 최대치로 조정합니다(--hbm_utilization_factor). 우선 0.95로 시작합니다. Hex-LLM 서버를 배포하고 긴 프롬프트 및 높은 동시성을 사용해 서버를 테스트합니다. 서버에 메모리 부족이 발생하면 적절하게 사용률을 줄입니다.

Llama 3.1 8B Instruct 배포용 샘플 인수 집합은 다음과 같습니다.

python -m hex_llm.server.api_server \
    --model=meta-llama/Llama-3.1-8B-Instruct \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.95

ct5lp-hightpu-4t에서 Llama 3.1 70B Instruct AWQ를 배포하기 위한 샘플 인수 집합은 다음과 같습니다.

python -m hex_llm.server.api_server \
    --model=hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4 \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.45

Cloud TPU 할당량 요청

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