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 Guard 3 1B 및 8B
  • Mistral 7B
  • Mixtral 8x7B 및 8x22B
  • Phi-3 mini 및 medium
  • Qwen2 0.5B, 1.5B, 7B
  • Qwen2.5 0.5B, 1.5B, 7B, 14B, 32B AWQ

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를 설정합니다. 다음은 Llama 3.1 8B bfloat16 모델을 제공하는 TPU v5e-8에 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를 설정합니다. 다음은 Llama 3.1 8B bfloat16 모델을 제공하는 TPU v5e-8에 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",
    "--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
  }
}

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) 컴파일 모드를 사용 설정할지 여부입니다. 기본값은 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입니다. 이 인수를 설정하면 이전 요청의 공유 접두사 계산을 재사용하여 성능을 개선할 수 있습니다.

메모리 관리

  • --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. 모델 크기 및 정밀도(model_size * (num_bits / 8))에 따라 필요한 TPU 메모리의 하한선을 추정합니다. 8B 모델 및 bfloat16 정밀도의 경우 필요한 TPU 메모리의 하한은 8 * (16 / 8) = 16 GB입니다.
  3. 각 v5e 칩이 16GB를 제공하는 경우 필요한 TPU v5e 칩 수를 추정합니다. tpu_memory / 16 8B 모델과 bfloat16 정밀도의 경우 칩이 2개 이상 필요합니다. 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 칩 4개입니다. 이 할당량은 클릭 한 번으로 배포 및 Colab Enterprise 노트북 배포에 적용됩니다. 추가 할당량을 요청하려면 할당량 상향 요청을 참조하세요.