View a markdown version of this page

Tiered KV-Cache einrichten - Amazon SageMaker KI
VoraussetzungenKonfigurieren Sie Ihre RayClusterKonfigurieren Sie Ihre Ray Serve-AnwendungCPU-Offloading aktivieren (optional)

Die vorliegende Übersetzung wurde maschinell erstellt. Im Falle eines Konflikts oder eines Widerspruchs zwischen dieser übersetzten Fassung und der englischen Fassung (einschließlich infolge von Verzögerungen bei der Übersetzung) ist die englische Fassung maßgeblich.

Tiered KV-Cache einrichten

Der verwaltete mehrstufige KV-Cache fügt Ray Serve-Bereitstellungen eine clusterweite Cache-Ebene hinzu, sodass ein Replikat ein Token-Präfix liest, das ein anderes Replikat bereits berechnet hat. Dies reduziert die Zeit bis zum ersten Token bei langen Dokumenten, Konversationen mit mehreren Durchgängen und gemeinsam genutzten Systemaufforderungen. Die clusterweite Ebene wird auf HyperPod Tiered Storage unter Verwendung von LMCache mit dem HyperPod Speicher-Backend in der LMCache-Dokumentation ausgeführt, das in Ray Serve und vLLM integriert ist, um verteiltes KV-Caching bereitzustellen, das durch Tiered Storage unterstützt wird. HyperPod

Der Cache hat zwei Ebenen:

  • Eine lokale Ebene im Speicher des Knotens, der die Anfrage bedient, für die schnellste Wiederverwendung.

  • Eine clusterweite Ebene auf HyperPod Tiered Storage, einer gepoolten Speicherebene, die sich über Clusterknoten erstreckt. Ein Replikat liest ein Präfix, das von einem anderen Replikat berechnet wurde, anstatt es erneut zu berechnen.

Wenn eine Anforderung ein Token-Präfix mit einer früheren Arbeit gemeinsam hat, liest die Bereitstellung den gecachten KV-Status von der lokalen Ebene und dann von der clusterweiten Ebene, bevor irgendetwas neu berechnet wird.

Voraussetzungen

  • Ein von Amazon EKS orchestrierter HyperPod Cluster, auf dem der Operator installiert ist. KubeRay

  • HyperPod Tiered Storage ist auf dem Cluster aktiviert. Anweisungen zur Einrichtung finden Sie unter Managed Tiered Checkpointing einrichten.

Konfigurieren Sie Ihre RayCluster

Fügen Sie Ihrer Worker-Gruppenspezifikation Folgendes hinzu, um den Tiered Storage-Endpunkt und den gemeinsam genutzten Speicher für die Ray Serve-Replikate verfügbar zu machen.

workerGroupSpecs:
- template:
    spec:
      volumes:
        - name: host-shm
          hostPath:
            path: /dev/shm
      containers:
      - name: ray-worker
        env:
          - name: NODE_IP
            valueFrom:
              fieldRef:
                fieldPath: status.hostIP
        volumeMounts:
          - name: host-shm
            mountPath: /dev/shm/ai_toolkit_cache
            subPath: ai_toolkit_cache

Die NODE_IP Umgebungsvariable stellt die Host-IP-Adresse bereit, sodass der LMCache-Client auf Port 9200 eine Verbindung zum Tiered Storage-Dienst herstellen kann. Die Bereitstellung eines gemeinsam genutzten Speicher-Volumes ermöglicht den CPU-Offloading-Pfad für den lokalen KV-Cache.

Konfigurieren Sie Ihre Ray Serve-Anwendung

Verwenden Sie die LLMConfig API mit dem LMCache-Connector, um KV-Caching in Ihrer Ray Serve-Bereitstellung zu aktivieren. Das folgende Beispiel dient einem Modell, bei dem mehrstufiges KV-Caching aktiviert ist:

from ray.serve.llm import LLMConfig, build_openai_app

llm_config = LLMConfig(
    model_loading_config={
        "model_id": "my-llm",
        "model_source": "Qwen/Qwen-7B-Chat",
    },
    engine_kwargs={
        "trust_remote_code": True,
        "kv_transfer_config": {
            "kv_connector": "LMCacheConnectorV1",
            "kv_role": "kv_both",
        },
    },
    runtime_env={
        "env_vars": {
            "LMCACHE_REMOTE_URL": "sagemaker-hyperpod://$(NODE_IP):9200",
            "LMCACHE_EXTRA_CONFIG": '{"sagemaker_hyperpod_bucket": "lmcache", "sagemaker_hyperpod_shared_memory_name": "ai_toolkit_cache"}',
            "LMCACHE_CHUNK_SIZE": "256",
        },
    },
    accelerator_type="A10G",
)

app = build_openai_app({"llm_configs": [llm_config]})

Die wichtigsten Konfigurationswerte sind:

  • kv_connector: Auf setzen, LMCacheConnectorV1 um die LMCache-Integration mit vLLM zu aktivieren.

  • kv_role: Legen Sie fest, kv_both dass jedes Replikat sowohl aus dem gemeinsamen Cache liest als auch in den gemeinsamen Cache schreibt.

  • LMCACHE_REMOTE_URL: Der Tiered Storage-Endpunkt auf dem lokalen Knoten. Verwendet die im RayCluster Manifest konfigurierte NODE_IP Umgebungsvariable.

  • LMCACHE_CHUNK_SIZE: Die Anzahl der Token pro Cache-Block. Kleinere Werte erhöhen die Cache-Trefferquote für gemeinsam genutzte Präfixe auf Kosten von mehr Cache-Einträgen.

CPU-Offloading aktivieren (optional)

Um eine lokale CPU-Speicherebene hinzuzufügen, die KV-Einträge zwischenspeichert, bevor sie in Tiered Storage geschrieben werden, fügen Sie die folgenden Umgebungsvariablen hinzu: runtime_env

"env_vars": {
    "LMCACHE_REMOTE_URL": "sagemaker-hyperpod://$(NODE_IP):9200",
    "LMCACHE_EXTRA_CONFIG": '{"sagemaker_hyperpod_bucket": "lmcache", "sagemaker_hyperpod_shared_memory_name": "ai_toolkit_cache"}',
    "LMCACHE_CHUNK_SIZE": "256",
    "LMCACHE_LOCAL_CPU": "True",
    "LMCACHE_MAX_LOCAL_CPU_SIZE": "100",
}

Wenn CPU-Offloading aktiviert ist, werden KV-Cache-Einträge im CPU-Speicher auf dem lokalen Knoten gespeichert, bevor sie in den clusterweiten Tiered Storage geschrieben werden. Dies ermöglicht schnellere Cache-Lesevorgänge für Replikate auf demselben Knoten.