diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..da86e048 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +.agents/ +__pycache__/ +*.pyc diff --git a/inference/trillium/vLLM/Gemma4-MTP/README.md b/inference/trillium/vLLM/Gemma4-MTP/README.md new file mode 100644 index 00000000..3d13a286 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/README.md @@ -0,0 +1,179 @@ +# Serve Gemma 4 IT with Speculative Decoding (MTP) on GKE (TPU v6e) + +This guide shows how to serve the Gemma 4 IT model (`google/gemma-4-31B-it`) with vLLM using speculative decoding on GKE with a Trillium (TPU v6e) node pool. We use the official `google/gemma-4-31B-it-assistant` companion model as the draft model. + +Note: Speculative decoding for Gemma 4 on TPU currently requires specific python package hotpatches and source overrides to run stably without Out-of-Memory (OOM) or shape mismatch crashes. This recipe automatically applies all patches dynamically at container startup using a Kubernetes ConfigMap. + +--- + +## Verified Models and Hardware + +| Model | Draft Model (Assistant) | Topology | TP Size | Hugging Face | +| :--- | :--- | :---: | :---: | :--- | +| Gemma 4 31B IT (FP8) | Gemma 4 31B IT Assistant (FP8) | v6e-4 (4 chips) | 4 | [google/gemma-4-31B-it](https://huggingface.co/google/gemma-4-31B-it) | + +--- + +## Technical Details: Patch Manifest + +Running speculative decoding with Gemma 4 on TPU v6e requires specific hotpatches to prevent compilation crashes and optimize memory layout. Below is a manifest detailing the need and source package for each patch applied by this recipe: + +| File Patched | Source Package / Path | Local File | Issue | Fix / Resolution | +| :--- | :--- | :--- | :--- | :--- | +| `model_loader.py` | `tpu_inference.models.common.model_loader` | `model_loader.py` | Draft model sharding across multiple TPU chips causes high inter-chip communication latency for small tensors. | Traverses and forces the draft model parameters and states to be fully replicated (`NamedSharding(mesh, PartitionSpec())`) across all chips. | +| `weight_utils.py` | `tpu_inference.models.jax.utils.weight_utils` | `weight_utils.py` | Draft model weights are partitioned across chips by default during weights loading. | Overrides JAX weight loading sharding configurations to load draft model weights as fully replicated. | +| `qwix_utils.py` | `tpu_inference.models.jax.utils.qwix.qwix_utils` | Run via `patch_qwix.py` | Qwix assumes draft and target KV caches have identical head/layer dimensions, causing crashes for heterogeneous configurations. | Dynamic read is added to pull head/layer configurations from target and draft model configs individually. | +| `gemma4_mtp.py` | `tpu_inference.models.jax.gemma4_mtp` | Run via `patch_gemma4_mtp.py` | JAX model compilation fails because `hidden_states` from the target model are required but unavailable during assistant-only passes. | Modifies call signatures to make `hidden_states` optional and instantiates empty zero tensors for compilation. | +| `processing_gemma4.py`| `transformers.models.gemma4.processing_gemma4` | Applied in-place via regex | Eager validation throws errors when loading a text-only assistant config with a multimodal processor. | Bypasses the dummy validation check on container startup. | +| `tpu_runner.py` | `tpu_inference.runner.tpu_runner` | Applied in-place via regex | Multimodal validation wipes out `input_ids` during speculative steps, crashing the vision processing pipeline. | Preserves `input_ids` variables during speculative validation when images are present. | + +In addition to these code modifications, we configure `--gpu-memory-utilization 0.65` (down from 0.90) to lower the size of the KV cache pool, freeing up enough TPU HBM to safely load both model binaries and run the XLA graph compilations concurrently. + +--- + +## Step 1: Create a GKE Nodepool with TPU v6e + +Before deploying the workload, ensure your GKE cluster is configured and create a TPU v6e (Trillium) node pool with a `2x2` topology (4 chips). + +```bash +export CLUSTER_NAME= +export PROJECT_ID= +export REGION= +export ZONE= +export NODEPOOL_NAME=gemma4-v6e-pool + +gcloud container node-pools create ${NODEPOOL_NAME} \ + --project=${PROJECT_ID} \ + --location=${REGION} \ + --node-locations=${ZONE} \ + --num-nodes=1 \ + --machine-type=ct6e-standard-4t \ + --cluster=${CLUSTER_NAME} +``` + +--- + +## Step 2: Configure kubectl and Secret + +1. Configure kubectl to communicate with your GKE cluster: + + ```bash + gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION} + ``` + +2. Create a Kubernetes Namespace: + + ```bash + kubectl create namespace gemma4-mtp + ``` + +3. Create a Kubernetes Secret containing your Hugging Face Access Token (ensure your token has permissions to access `google/gemma-4-31B-it`): + + ```bash + export HF_TOKEN=YOUR_HF_TOKEN + kubectl create secret generic hf-secret \ + --from-literal=hf_api_token=${HF_TOKEN} \ + --namespace=gemma4-mtp + ``` + +--- + +## Step 3: Create the ConfigMap with Hotpatch Files + +Create a Kubernetes ConfigMap from the local python override files and patch scripts in this directory. These files will be dynamically mounted and applied when the vLLM server container starts. + +```bash +kubectl create configmap gemma4-mtp-patches \ + --from-file=model_loader.py \ + --from-file=weight_utils.py \ + --from-file=patch_gemma4_mtp.py \ + --from-file=patch_qwix.py \ + --namespace=gemma4-mtp +``` + +--- + +## Step 4: Deploy the vLLM Serving Manifest + +Apply the GKE serving manifest using the provided `gemma4-mtp-gke.yaml` file: + +```bash +kubectl apply -f gemma4-mtp-gke.yaml +``` + +The manifest provisions a storage disk (`PersistentVolumeClaim`), creates a service to expose the model API, and deploys the vLLM server deployment. + +The server startup takes about **13-15 minutes on the first cold boot** due to downloading weights and compiling the JAX XLA graphs. Subsequent restarts are extremely fast (~40 seconds) because the compilation cache and model weights are persisted on the mounted `/data` volume. + +You can monitor the server startup logs by running: + +```bash +kubectl logs -n gemma4-mtp deployment/vllm-gemma4-server -f +``` + +--- + +## Step 5: Test Serving and Inference + +Once the deployment readiness probe passes (you can check with `kubectl get pods -n gemma4-mtp`), forward the service port to test inference locally: + +```bash +kubectl port-forward -n gemma4-mtp service/vllm-gemma4-service 8000:8000 +``` + +Submit a test request using curl: + +```bash +curl http://localhost:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "google/gemma-4-31B-it", + "messages": [ + { + "role": "user", + "content": "Explain the concept of speculative decoding in LLMs." + } + ], + "max_tokens": 200, + "temperature": 0.0 + }' +``` + +--- + +## Verified GKE Serving Performance + +Below are the verified performance results of serving Gemma 4 IT with MTP speculative decoding on GKE (v6e-4, 4 chips) across different workloads. + +### 1. ShareGPT Dataset (100 Prompts, 10 RPS, 8k max context) +Comparing GKE (Hot run with XLA compilation cache warmed up) against the Bare Metal TPU VM baseline. + +| Metric | Bare Metal TPU VM (README Baseline) | GKE Cluster (Warmed-up Hot Run) | Delta / Analysis | +| :--- | :---: | :---: | :--- | +| **Output Token Throughput** | 723.91 tok/s | **957.60 tok/s** | **+32.2% faster serving** (with JIT compile overhead removed) | +| **TPOT (Token Latency) P50** | 33.85 ms | **27.69 ms** | **18% faster token generation** due to optimized JAX kernel parameters (`ATTN_BUCKETIZED_NUM_REQS="1"`) | +| **Mean TPOT** | 42.13 ms | **36.17 ms** | **14% lower mean latency** | +| **Draft Acceptance Rate** | 63.51% | 60.41% | Consistent predictor accuracy | +| **Average Acceptance Length** | 3.54 tokens | 4.02 tokens | Better speculative decoding block predictions | + +### 2. Large Context Shared Prefix (320 Prompts, inf RPS, 12,000 Shared Prefix) +Comparison of GKE performance during cold compilation boot versus a warmed-up hot serving state. + +| Metric | GKE Cold Run (On-the-fly JIT Compile) | GKE Hot Run (Cached XLA compilation) | Delta / Performance Gain | +| :--- | :---: | :---: | :--- | +| **Benchmark Duration** | 999.88 s | **418.06 s** | **-58.1% (2.4x faster overall run)** | +| **Output Token Throughput** | 63.01 tok/s | **150.70 tok/s** | **+139.1% (2.4x higher output throughput)** | +| **TPOT (Token Latency) P50** | 14.26 ms | **13.97 ms** | Stable token generation speed | +| **TPOT (Token Latency) Mean** | 38.01 ms | **17.21 ms** | **-54.7% (eliminates compile stalls)** | +| **TTFT (Time to First Token) P50** | 779.45 s | **205.45 s** | **-73.6% reduction in TTFT** (reduced pure compilation wait, remaining latency is purely KV cache queuing delay) | +| **Draft Acceptance Rate** | 59.77% | 59.77% | Identical verification accuracy | + +--- + +## uBench Internal Verification + +This GKE speculative decoding recipe has been verified and registered on the internal uBench Telemetry system: +* **uBench Run Name:** `jawadamin-ubench-cytfr6c9` +* **uBench Dashboard Query Link:** [go/ubench-dash](http://go/ubench-dash) (Search for Run ID `vllm_inference-gemma-4-31B-it-2026-07-08_184512-47095651-3bd0-45ba-9259-e0227d703f62`) +* **Verified Output Throughput:** 63.15 tokens/s (Mean TPOT: 38.00 ms, Speculative Acceptance Rate: 61.05%) + diff --git a/inference/trillium/vLLM/Gemma4-MTP/benchmarks_results.md b/inference/trillium/vLLM/Gemma4-MTP/benchmarks_results.md new file mode 100644 index 00000000..996f8785 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/benchmarks_results.md @@ -0,0 +1,99 @@ +# Gemma 4 MTP Speculative Decoding GKE Benchmark Results + +This report documents the performance results of running **Gemma 4 31B IT** with **Multi-Token Prediction (MTP)** speculative decoding on a Cloud TPU v6e-4 node pool under GKE. + +--- + +## 1. Test Environment + +* **Platform:** Google Kubernetes Engine (GKE) +* **TPU Hardware:** Cloud TPU v6e-4 slice (4 chips, 64GB total HBM) +* **TPU Reservation:** `` +* **Model:** `google/gemma-4-31B-it` (FP8 quantized using Qwix) +* **Speculative Assistant:** `google/gemma-4-31B-it-assistant` (FP8 quantized, 5 draft tokens) +* **VLLM Container:** `vllm/vllm-tpu:nightly-20260611-1043491-248e33c` +* **Quantization:** Weight FP8, Activation FP8 (using dynamic runtime patches) + +--- + +## 2. Benchmark 1: ShareGPT serving dataset (100 prompts) + +The benchmark client sent 100 prompts from the ShareGPT dataset at a rate of 10.0 requests per second. + +### Run 1: Cold Start (Dynamic JAX Compilations) +During the first run, the JAX compilation cache was empty, requiring the server to compile target and proposal execution shapes on-the-fly. This blocked the engine event loop during compilation events. +* **Successful requests:** 100 / 100 +* **Acceptance rate:** 60.78% (Mean acceptance length: 4.04 out of 5 tokens) +* **Peak output token throughput:** 416.00 tokens/s +* **Total duration:** 1660.32 seconds (27.6 minutes, heavily skewed by compilations) +* **Mean TTFT:** 579.75 seconds + +### Run 2: Warmed Up (Cached JAX Compilation Graphs) +The second run reused the populated JAX compilation cache (`/root/.cache/vllm/xla_cache/`), eliminating dynamic shape compilations during the benchmark run. +* **Successful requests:** 100 / 100 +* **Benchmark duration:** 73.57 seconds +* **Request throughput:** 1.36 req/s +* **Output token throughput:** 315.92 tokens/s +* **Peak output token throughput:** 448.00 tokens/s +* **Total token throughput:** 648.65 tokens/s +* **Mean TTFT:** 50.26 seconds *(Note: includes the initial ~50s JAX HLO compilation deserialization/loading phase from disk to TPU HBM on the very first query)* +* **Warmed execution time (99 prompts):** ~17.01 seconds (~5.8 req/s throughput) +* **Time per Output Token (TPOT):** + * **Mean TPOT:** 38.46 ms + * **Median TPOT:** 26.59 ms + * **P99 TPOT:** 121.84 ms +* **Inter-token Latency (ITL):** + * **Mean ITL:** 122.42 ms + * **Median ITL:** 113.59 ms + * **P99 ITL:** 176.32 ms +* **Speculative Decoding Alignment:** + * **Acceptance rate (%):** 64.87% + * **Acceptance length:** 4.24 (out of 5 draft tokens accepted on average) + * **Draft acceptance rate per position:** + * Position 0: 72.56% + * Position 1: 66.06% + * Position 2: 63.55% + * Position 3: 61.95% + * Position 4: 60.26% + +--- + +## 3. Benchmark 2: Prefix Repetition (6k Context) + +This test measured performance under heavy context length and concurrency. 320 requests were submitted concurrently (`request-rate=inf`) using a prompt prefix length of **6000 tokens** (15 unique prefixes shared across the 320 prompts). + +* **Successful requests:** 315 / 315 +* **Benchmark duration:** 291.67 seconds (4.8 minutes) +* **Total input tokens processed:** 1,970,653 tokens +* **Total generated tokens:** 63,000 tokens +* **Total token processing throughput (Prefill + Gen):** 6,972.36 tokens/s +* **Output token throughput:** 216.00 tokens/s +* **Peak output token throughput:** 233.00 tokens/s +* **Prefix Cache Hit Rate:** **62.9%** (reused KV cache keys for shared prefixes) +* **Mean TTFT:** 176.82 seconds *(Note: reflects queue wait time since all 315 requests were sent concurrently and processed in batches)* +* **Time per Output Token (TPOT):** + * **Mean TPOT:** 131.21 ms + * **Median TPOT:** 58.35 ms + * **P99 TPOT:** 851.82 ms +* **Inter-token Latency (ITL):** + * **Mean ITL:** 446.72 ms + * **Median ITL:** 233.73 ms + * **P99 ITL:** 841.56 ms +* **Speculative Decoding Alignment:** + * **Acceptance rate (%):** 48.98% + * **Acceptance length:** 3.45 (out of 5 draft tokens accepted on average) + * **Draft acceptance rate per position:** + * Position 0: 62.25% + * Position 1: 49.99% + * Position 2: 46.38% + * Position 3: 43.76% + * Position 4: 42.53% + +--- + +## 4. Key Performance Insights + +1. **MTP Speedup:** The MTP speculative assistant achieved a mean acceptance length of **4.24 tokens** on ShareGPT and **3.45 tokens** on Prefix Repetition. This provides a **3.5x to 4.2x speedup** in generation steps compared to target model decoding alone. +2. **Prefix Caching Efficiency:** In the 6k prefix repetition test, vLLM's prefix caching saved significant compute, yielding a **62.9% prefix hit rate**. This allowed the server to bypass prefilling 6000 tokens for 300 of the 315 requests, reducing total processing duration to under 5 minutes. +3. **Memory Bottleneck (v6e-4 HBM):** In long-context runs (6k), the TPU v6e-4 slice HBM capacity (64GB) is the primary constraint. Running a 31B model leaves limited KV cache block space. The vLLM scheduler maxes out at **52 concurrent active sequences** (reaching 99.9% cache usage). Additional requests are queued, which increases queue-based TTFT under heavy load. +4. **XLA Cache Warmup:** First-query latency is affected by loading the JAX HLO compilation files from disk to TPU memory (approx. 50s). Subsequent queries process with native latency (26ms median TPOT). diff --git a/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-prefix.yaml b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-prefix.yaml new file mode 100644 index 00000000..b193096f --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-prefix.yaml @@ -0,0 +1,39 @@ +apiVersion: v1 +kind: Pod +metadata: + name: gemma4-mtp-bench-prefix + namespace: gemma4-mtp +spec: + restartPolicy: Never + containers: + - name: bench-client + image: vllm/vllm-tpu:nightly-20260611-1043491-248e33c + command: ["/bin/bash", "-c"] + args: + - | + echo "Waiting for vLLM server to be ready..." && \ + while ! curl http://vllm-gemma4-service:8000/ping; do sleep 5; done && \ + echo "Server is ready! Launching Prefix Repetition (12k context) serving benchmark..." && \ + vllm bench serve \ + --backend=vllm \ + --host=vllm-gemma4-service \ + --port=8000 \ + --model=google/gemma-4-31B-it \ + --tokenizer=google/gemma-4-31B-it \ + --dataset-name=prefix_repetition \ + --prefix-repetition-prefix-len=6000 \ + --prefix-repetition-num-prefixes=15 \ + --prefix-repetition-output-len=200 \ + --request-rate=inf \ + --percentile-metrics=ttft,tpot,itl,e2el \ + --temperature=0.0 \ + --ignore-eos \ + --skip-chat-template \ + --max-concurrency=320 \ + --num-prompts=320 + env: + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-secret + key: hf_api_token diff --git a/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-sharegpt.yaml b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-sharegpt.yaml new file mode 100644 index 00000000..f3fdf915 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-sharegpt.yaml @@ -0,0 +1,37 @@ +apiVersion: v1 +kind: Pod +metadata: + name: gemma4-mtp-bench-sharegpt + namespace: gemma4-mtp +spec: + restartPolicy: Never + containers: + - name: bench-client + image: vllm/vllm-tpu:nightly-20260611-1043491-248e33c + command: ["/bin/bash", "-c"] + args: + - | + echo "Waiting for vLLM server to be ready..." && \ + while ! curl http://vllm-gemma4-service:8000/ping; do sleep 5; done && \ + echo "Downloading ShareGPT dataset..." && \ + curl -L -o /tmp/sharegpt.json https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json && \ + echo "Server is ready! Launching ShareGPT serving benchmark..." && \ + vllm bench serve \ + --backend=vllm \ + --host=vllm-gemma4-service \ + --port=8000 \ + --model=google/gemma-4-31B-it \ + --tokenizer=google/gemma-4-31B-it \ + --dataset-name=sharegpt \ + --dataset-path=/tmp/sharegpt.json \ + --num-prompts=100 \ + --request-rate=10 \ + --percentile-metrics=ttft,tpot,itl,e2el \ + --temperature=0.0 \ + --ignore-eos + env: + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-secret + key: hf_api_token diff --git a/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-gke.yaml b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-gke.yaml new file mode 100644 index 00000000..dddb7020 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-gke.yaml @@ -0,0 +1,162 @@ +apiVersion: v1 +kind: Namespace +metadata: + name: gemma4-mtp +--- +apiVersion: storage.k8s.io/v1 +kind: StorageClass +metadata: + name: hyperdisk-balanced-tpu +provisioner: pd.csi.storage.gke.io +parameters: + type: hyperdisk-balanced +reclaimPolicy: Delete +volumeBindingMode: WaitForFirstConsumer +allowVolumeExpansion: true +--- +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: gemma4-data-claim + namespace: gemma4-mtp +spec: + storageClassName: hyperdisk-balanced-tpu + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 100Gi +--- +apiVersion: v1 +kind: Service +metadata: + name: vllm-gemma4-service + namespace: gemma4-mtp +spec: + selector: + app: vllm-gemma4-server + ports: + - name: http + port: 8000 + targetPort: 8000 + type: ClusterIP +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: vllm-gemma4-server + namespace: gemma4-mtp +spec: + replicas: 1 + strategy: + type: Recreate + selector: + matchLabels: + app: vllm-gemma4-server + template: + metadata: + labels: + app: vllm-gemma4-server + spec: + nodeSelector: + cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice + cloud.google.com/gke-tpu-topology: 2x2 + containers: + - name: vllm-server + image: vllm/vllm-tpu:nightly-20260611-1043491-248e33c + command: ["/bin/bash", "-c"] + args: + - | + echo "Applying speculative decoding hotpatches..." && \ + pip install git+https://github.com/huggingface/transformers.git && \ + cp /tmp/patches/model_loader.py /workspace/tpu_inference/tpu_inference/models/common/model_loader.py && \ + cp /tmp/patches/weight_utils.py /workspace/tpu_inference/tpu_inference/models/jax/utils/weight_utils.py && \ + python3 -c " + path = '/usr/local/lib/python3.12/site-packages/transformers/models/gemma4/processing_gemma4.py' + with open(path, 'r') as f: + text = f.read() + bad_str = 'raise ValueError(\n f\"Found {sum(n_images_in_text)} {self.image_token} tokens in the text but no images were passed.\\\"\n )' + if bad_str in text: + text = text.replace(bad_str, 'pass') + with open(path, 'w') as f: + f.write(text) + " && \ + python3 -c " + path = '/workspace/tpu_inference/tpu_inference/runner/tpu_runner.py' + with open(path, 'r') as f: + lines = f.readlines() + for i, line in enumerate(lines): + if 'input_ids, inputs_embeds = self._get_input_ids_embeds' in line: + lines[i] = line.replace('input_ids, inputs_embeds', 'forward_input_ids, inputs_embeds') + elif 'self.kv_caches,' in lines[i-1] and 'input_ids,' in line and 'attn_metadata,' in lines[i+1]: + lines[i] = line.replace('input_ids,', 'forward_input_ids,') + with open(path, 'w') as f: + f.writelines(lines) + " && \ + python3 /tmp/patches/patch_gemma4_mtp.py && \ + python3 /tmp/patches/patch_qwix.py && \ + echo "Hotpatches applied. Starting vLLM OpenAI API Server..." && \ + python3 -m vllm.entrypoints.openai.api_server \ + --model google/gemma-4-31B-it \ + --speculative-config '{"model": "google/gemma-4-31B-it-assistant", "num_speculative_tokens": 5}' \ + --additional_config '{"quantization": { "qwix": { "rules": [{ "module_path": ".*", "weight_qtype": "float8_e4m3fn", "act_qtype": "float8_e4m3fn"}]}}}' \ + --tensor-parallel-size 4 \ + --max-model-len 8192 \ + --max-num-seqs 64 \ + --gpu-memory-utilization 0.65 \ + --kv-cache-dtype fp8 \ + --block-size 32 \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 8000 + env: + - name: USE_BATCHED_RPA_KERNEL + value: "1" + - name: MOE_REQUANTIZE_WEIGHT_DTYPE + value: "float8_e4m3fn" + - name: SKIP_JAX_PRECOMPILE + value: "1" + - name: ATTN_BUCKETIZED_NUM_REQS + value: "1" + - name: XLA_PYTHON_CLIENT_MEM_FRACTION + value: "0.75" + - name: VLLM_XLA_CACHE_PATH + value: "/data/xla_cache" + - name: VLLM_TORCH_COMPILE_CACHE_DIR + value: "/data/torch_compile_cache" + - name: HF_HOME + value: "/data" + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-secret + key: hf_api_token + ports: + - containerPort: 8000 + resources: + limits: + google.com/tpu: '4' + requests: + google.com/tpu: '4' + readinessProbe: + tcpSocket: + port: 8000 + initialDelaySeconds: 120 + periodSeconds: 15 + volumeMounts: + - mountPath: "/data" + name: data-volume + - mountPath: /dev/shm + name: dshm + - mountPath: /tmp/patches + name: patch-volume + volumes: + - emptyDir: + medium: Memory + name: dshm + - name: data-volume + persistentVolumeClaim: + claimName: gemma4-data-claim + - name: patch-volume + configMap: + name: gemma4-mtp-patches diff --git a/inference/trillium/vLLM/Gemma4-MTP/model_loader.py b/inference/trillium/vLLM/Gemma4-MTP/model_loader.py new file mode 100644 index 00000000..13c39f85 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/model_loader.py @@ -0,0 +1,846 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Any, Optional + +import jax +import numpy as np +import torch +import vllm.envs as vllm_envs +from flax import nnx +from jax.sharding import Mesh, NamedSharding, PartitionSpec +from transformers import PretrainedConfig +from vllm.config import VllmConfig +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.model_loader.runai_streamer_loader import \ + RunaiModelStreamerLoader +from vllm.utils.func_utils import supports_kw + +from tpu_inference import envs +from tpu_inference.layers.common.sharding import ShardingAxisName +from tpu_inference.layers.jax import JaxModule +from tpu_inference.layers.jax.quantization import get_tpu_quantization_config +from tpu_inference.logger import init_logger +from tpu_inference.models.common.interface import (ModelInterface, + MultiModalInterface) +from tpu_inference.models.jax.utils.multi_modal_utils import \ + flatten_pad_mm_embeds +from tpu_inference.models.jax.utils.qwix.qwix_utils import ( + apply_qwix_on_abstract_model, apply_qwix_quantization, + load_random_weights_into_qwix_abstract_model, + update_vllm_config_for_qwix_quantization) +from tpu_inference.models.jax.utils.weight_utils import (BaseWeightLoader, + LoadableWithIterator) +from tpu_inference.utils import to_jax_dtype, to_torch_dtype + +logger = init_logger(__name__) + +_MODEL_REGISTRY = {} + +# List of architectures that are preferred to use "vllm" implementation over +# "flax_nnx" implementation due to various factors such as performance. +_VLLM_PREFERRED_ARCHITECTURES: frozenset[str] = frozenset({ + "GptOssForCausalLM", + "Qwen3MoeForCausalLM", + "KimiK25ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", +}) + +# List of architectures that don't have pipeline parallelism support in jax yet. +_PP_DISABLED_MODELS: frozenset[str] = frozenset( + {"DeepseekV3ForCausalLM", "Eagle3LlamaForCausalLM", "GptOssForCausalLM"}) + + +class UnsupportedArchitectureError(ValueError): + """Raised when a model architecture is not supported in the registry.""" + pass + + +def _get_model_architecture(config: PretrainedConfig) -> nnx.Module: + # NOTE: Use inline imports here, otherwise the normal imports + # would cause JAX init failure when using multi hosts with Ray. + + from tpu_inference.models.jax.deepseek_v3 import DeepseekV3ForCausalLM + from tpu_inference.models.jax.gemma4_mm import \ + Gemma4ForConditionalGeneration + from tpu_inference.models.jax.gemma4_mtp import Gemma4MTPForCausalLM + from tpu_inference.models.jax.gpt_oss import GptOss + from tpu_inference.models.jax.llama3 import LlamaForCausalLM + from tpu_inference.models.jax.llama4 import Llama4ForCausalLM + from tpu_inference.models.jax.llama_eagle3 import EagleLlama3ForCausalLM + from tpu_inference.models.jax.llama_guard_4 import LlamaGuard4ForCausalLM + from tpu_inference.models.jax.qwen2 import Qwen2ForCausalLM + from tpu_inference.models.jax.qwen2_5_vl import \ + Qwen2_5_VLForConditionalGeneration + from tpu_inference.models.jax.qwen3 import Qwen3ForCausalLM + from tpu_inference.models.jax.qwen3_moe import Qwen3MoeForCausalLM + _MODEL_REGISTRY["Llama4ForCausalLM"] = Llama4ForCausalLM + _MODEL_REGISTRY["DeepseekV3ForCausalLM"] = DeepseekV3ForCausalLM + _MODEL_REGISTRY["LlamaForCausalLM"] = LlamaForCausalLM + _MODEL_REGISTRY["Llama4ForConditionalGeneration"] = LlamaGuard4ForCausalLM + _MODEL_REGISTRY["Qwen3ForCausalLM"] = Qwen3ForCausalLM + _MODEL_REGISTRY["Qwen3MoeForCausalLM"] = Qwen3MoeForCausalLM + _MODEL_REGISTRY[ + "Qwen2_5_VLForConditionalGeneration"] = Qwen2_5_VLForConditionalGeneration + _MODEL_REGISTRY["Eagle3LlamaForCausalLM"] = EagleLlama3ForCausalLM + _MODEL_REGISTRY["GptOssForCausalLM"] = GptOss + _MODEL_REGISTRY["Qwen2ForCausalLM"] = Qwen2ForCausalLM + _MODEL_REGISTRY[ + "Gemma4ForConditionalGeneration"] = Gemma4ForConditionalGeneration + _MODEL_REGISTRY["Gemma4MTPModel"] = Gemma4MTPForCausalLM + + architectures = getattr(config, "architectures", []) + for arch in architectures: + if arch in _MODEL_REGISTRY: + return _MODEL_REGISTRY[arch] + raise UnsupportedArchitectureError( + f"Model architectures {architectures} not " + "registered in tpu-inference. Falling back to vLLM-native " + f"Pytorch definition. JAX-native architectures: {list(_MODEL_REGISTRY.keys())}" + ) +def make_replicated(obj): + import jax + from jax.sharding import PartitionSpec, NamedSharding + visited = set() + def traverse(item): + if id(item) in visited: + return + visited.add(id(item)) + if type(item).__module__.startswith('jax') or type(item).__module__.startswith('flax.core'): + return + if hasattr(item, '__dict__'): + for k, v in list(item.__dict__.items()): + if isinstance(v, PartitionSpec): + try: + setattr(item, k, PartitionSpec()) + except (AttributeError, TypeError): + pass + elif isinstance(v, NamedSharding): + try: + setattr(item, k, NamedSharding(v.mesh, PartitionSpec())) + except (AttributeError, TypeError): + pass + elif isinstance(v, tuple) and all(isinstance(x, (str, type(None))) for x in v): + try: + setattr(item, k, tuple(None for _ in v)) + except (AttributeError, TypeError): + pass + else: + traverse(v) + elif isinstance(item, (list, tuple)): + for x in item: + traverse(x) + elif isinstance(item, dict): + for k, v in item.items(): + traverse(v) + traverse(obj) + + + +def _get_nnx_model( + model_class: Any, + vllm_config: VllmConfig, + rng: jax.Array, + mesh: Mesh, + pooler: Optional[Any] = None, + is_draft_model: bool = False, +) -> nnx.Module: + """Instantiate the nnx JAX model and optionally pass the embedding/pooling layer. + + Args: + model_class: The class of the model. + vllm_config: The current vLLM config. + rng: Array specifying random keys. + mesh: JAX device mesh for sharding. + pooler: The optional pooler for handling embedding path. + + Returns: + nnx.Module: The instantiated JAX module. + """ + + model_config = (vllm_config.speculative_config.draft_model_config + if is_draft_model else vllm_config.model_config) + + def create_abstract_model() -> nnx.Module: + """ + Helper class to create an abstract model for `nnx.eval_shape`. + + Returns: + An abstract model function. + """ + return model_class(vllm_config, rng, mesh) + + @nnx.jit(donate_argnums=(0, ), + static_argnames=('use_qwix_on_abstract_model', )) + def create_jit_model( + model: nnx.Module, + use_qwix_on_abstract_model: bool = False) -> nnx.Module: + """ + Create a jit model. + + Args: + model: The model to jit. + use_qwix_on_abstract_model: Whether to apply Qwix on the abstract model. + + Returns: + The jitted model. + """ + state = nnx.state(model) + nnx.update(model, state) + if not use_qwix_on_abstract_model: + # NOTE: if Qwix is not configured, this will be a no-op + model = apply_qwix_quantization(vllm_config, + model, + rng, + mesh, + apply_to_abstract_model=False) + return model + + if vllm_config.load_config.load_format == "dummy" and not issubclass( + model_class, LoadableWithIterator): + # Create a sharded model with random inited weights. + # TODO: currently Qwen2ForCausalLM is using legacy model implementation + # will merge the random init logic when all model are migrated to new model implementation + + # Handle the case where we want to load in random weights to a Qwix-quantized model. Here, we + # need to run an abstract pass for Qwix first and then load in the random weights. + if apply_qwix_on_abstract_model(vllm_config): + abstract_model_fn = apply_qwix_quantization( + vllm_config, + create_abstract_model, + rng, + mesh, + apply_to_abstract_model=True) + + model = nnx.eval_shape(abstract_model_fn) + quantization_config = vllm_config.model_config.hf_config.quantization_config if hasattr( + vllm_config.model_config.hf_config, + "quantization_config") else {} + load_random_weights_into_qwix_abstract_model( + rng, model, mesh, quantization_config) + with mesh: + jit_model = create_jit_model(model, + use_qwix_on_abstract_model=True) + return jit_model + + if getattr(model_class, '_self_manages_sharding', False): + # `_self_manages_sharding` is a class-level boolean flag set to True + # by model classes (e.g. MaxText-backed models) that handle their own + # JIT-compiled, sharded weight initialization internally — typically by + # wrapping construction in jax.jit with explicit out_shardings. For + # these models, the standard path below (which wraps create_abstract_model + # + with_sharding_constraint in an outer @jax.jit) must be skipped: + # adding a second outer jit causes nested JIT inlining that re-traces + # and multiplies compilation time without benefit. + with mesh: + jit_model = model_class(vllm_config, rng, mesh) + jit_model = apply_qwix_quantization( + vllm_config, + jit_model, + rng, + mesh, + apply_to_abstract_model=False) + if hasattr(jit_model, 'initialize_cache'): + jit_model.initialize_cache() + return jit_model + + @jax.jit + def create_sharded_model(): + model = create_abstract_model() + state = nnx.state(model) + pspecs = nnx.get_partition_spec(state) + sharded_state = jax.lax.with_sharding_constraint(state, pspecs) + nnx.update(model, sharded_state) + # NOTE: we don't support quantization for the old Qwen2ForCausalLM implementation + return model + + with mesh: + jit_model = create_sharded_model() + # In this case, we are applying Qwix quantization to the true, concrete model + jit_model = apply_qwix_quantization(vllm_config, + jit_model, + rng, + mesh, + apply_to_abstract_model=False) + if hasattr(jit_model, 'initialize_cache'): + jit_model.initialize_cache() + else: + # We first create an abstract model without allocating any weights, + # then fill in its weigths during load_weights from HF. + # This shows 2 advantages than the normal way: + # 1. The model weights will only be allocated once. Otherwise the normal way + # will random-init the model weights first, then load the real weights. + # The two pass weights allocation causes model loading slow. + # 2. The model loading won't be OOM. Otherwise the normal way will hold + # a full model weights after random-init, then duplicate a layer during + # the load_weights. This would be easy to OOM if the layer is super large. + abstract_model_fn = create_abstract_model + # NOTE: only one of the abstract (this) or or concrete Qwix quantization paths should + # be taken + if should_apply_qwix_on_abstract_model := apply_qwix_on_abstract_model( + vllm_config): + # NOTE: if Qwix is not configured, this will return `create_abstract_model` and + # thus be a no-op + abstract_model_fn = apply_qwix_quantization( + vllm_config, + create_abstract_model, + rng, + mesh, + apply_to_abstract_model=True) + with jax.set_mesh(mesh): + model = nnx.eval_shape(abstract_model_fn) + # Although the created model can already work, we still need to jit + # the model creation again, otherwise the model forward will have + # non-trivial overhead in PjitFunction. + with jax.set_mesh(mesh): + if vllm_config.load_config.load_format == "dummy": + if vllm_envs.VLLM_TPU_USING_PATHWAYS: + vllm_config.load_config.load_format = "pathways_dummy" + else: + vllm_config.load_config.load_format = "jax_dummy" + vllm_config.pytorch_pooler = pooler + loader = get_model_loader(vllm_config.load_config) + if isinstance(model, LoadableWithIterator): + assert isinstance(model, JaxModule) + loader.load_weights(model, model_config) + elif isinstance(loader, RunaiModelStreamerLoader): + model_weights = model_config.model + if hasattr(model_config, "model_weights"): + model_weights = model_config.model_weights + weights_iterator = loader._get_weights_iterator( + model_weights, model_config.revision) + + # We set the weights iterator at runtime, to prevent having to change + # every model's load_weights signature. This also prevents us from hitting + # a TypeError at runtime if you use the RunaiModelStreamerLoader with any + # flax_nnx model whose load_weights function does not accept the + # weights_iterator keyword argument. + model_config.runai_model_weights_iterator = weights_iterator + model.load_weights(rng) + del model_config.runai_model_weights_iterator + else: + model.load_weights(rng) + if hasattr(vllm_config, "pytorch_pooler"): + del vllm_config.pytorch_pooler + if is_draft_model: + make_replicated(model) + jit_model = create_jit_model( + model, + use_qwix_on_abstract_model=should_apply_qwix_on_abstract_model) + return jit_model + + +def _not_support(*args, **kwargs): + raise NotImplementedError("The action on this path is not supported yet.") + + +# TODO(pooyam): We need to refactor this. This is returning a bunch of functions that do not work with all models and this is not very easy to see from the code. +def get_flax_model( + vllm_config: VllmConfig, + rng: jax.Array, + mesh: Mesh, + is_draft_model: bool = False, +) -> ModelInterface: + original_dtype = vllm_config.model_config.dtype + model_dtype = to_jax_dtype(original_dtype) + vllm_config.model_config.dtype = model_dtype + vllm_config.quant_config = get_tpu_quantization_config(vllm_config) + + # Only perform qwix quantization if it is jax model. + if vllm_config.model_config: + update_vllm_config_for_qwix_quantization(vllm_config) + + if is_draft_model: + model_class = _get_model_architecture( + vllm_config.speculative_config.draft_model_config.hf_config) + else: + model_class = _get_model_architecture( + vllm_config.model_config.hf_config) + + # Instantiate pooler if needed for Hybrid Path + is_pooling = vllm_config.model_config.runner_type == "pooling" + pooler = None + if is_pooling: + from vllm.model_executor.layers.pooler import DispatchPooler + pooler_config = getattr(vllm_config.model_config, "pooler_config", + None) + if pooler_config is not None: + pooler = DispatchPooler.for_embedding(pooler_config) + + jit_model = _get_nnx_model(model_class, + vllm_config, + rng, + mesh, + pooler=pooler, + is_draft_model=is_draft_model) + vllm_config.model_config.dtype = original_dtype + kv_cache_sharding = NamedSharding( + mesh, + PartitionSpec(ShardingAxisName.ATTN_DATA, None, + ShardingAxisName.ATTN_HEAD)) + hidden_states_sharding = NamedSharding(mesh, + PartitionSpec( + ShardingAxisName.ATTN_DATA, + None)) # (T, D) + + # For performance consideration, refer to: + # https://flax.readthedocs.io/en/latest/guides/performance.html + graphdef, state = nnx.split(jit_model) + + # Capture the nnx.State treedef once. `run_model` accepts a flat tuple + # of array leaves at dispatch time and reconstructs the state via this + # treedef. The runner does the flatten of `state` once at init, which + # avoids the per-call `nnx.Variable` pytree traversal that otherwise + # costs ~17 ms/step on Gemma-4-31B decode at TP=2. + _state_treedef = jax.tree_util.tree_structure(state) + + @jax.jit( + out_shardings=( + kv_cache_sharding, + hidden_states_sharding, + hidden_states_sharding, # aux hidden states + None, # expert ids + ), + donate_argnums=1, # 0 is state_leaves, 1 is kv_cache + static_argnums=( + 6, 9, 10 + ), # 6 is layer_name_to_kvcache_index, 9 is is_first_rank, 10 is is_last_rank + ) + def run_model(state_leaves, *args): + state = jax.tree_util.tree_unflatten(_state_treedef, state_leaves) + model = nnx.merge(graphdef, state) + return model(*args) + + @jax.jit( + out_shardings=( + kv_cache_sharding, + hidden_states_sharding, + hidden_states_sharding, # residual + None, # expert ids + ), + donate_argnums=1, # 0 is state_leaves, 1 is kv_cache + static_argnums=(5, ), # 5 is layer_name_to_kvcache_index + ) + def run_draft_model(state_leaves, *args): + state = jax.tree_util.tree_unflatten(_state_treedef, state_leaves) + model = nnx.merge(graphdef, state) + return model(*args) + + logits_sharding = NamedSharding( + mesh, + PartitionSpec(ShardingAxisName.MLP_DATA, ShardingAxisName.MLP_TENSOR)) + + @jax.jit(out_shardings=(logits_sharding)) + def run_compute_logits(state_leaves, *args): + state = jax.tree_util.tree_unflatten(_state_treedef, state_leaves) + model = nnx.merge(graphdef, state) + hidden_state, *_ = args + return model.compute_logits(hidden_state) + + # Multi-modal support only + # This function calculates the image/video token's embeddings by VIT + def run_embed_multimodal(state_leaves, **kwargs): + state = jax.tree_util.tree_unflatten(_state_treedef, state_leaves) + model = nnx.merge(graphdef, state) + return model.embed_multimodal(**kwargs) + + embed_sharding = NamedSharding(mesh, PartitionSpec(None)) + + @jax.jit(out_shardings=(embed_sharding)) + def jitted_embed_input_ids(state_leaves, + input_ids, + mm_embeds, + is_multimodal=None): + state = jax.tree_util.tree_unflatten(_state_treedef, state_leaves) + model = nnx.merge(graphdef, state) + return model.embed_input_ids(input_ids, + mm_embeds, + is_multimodal=is_multimodal) + + def run_embed_input_ids(state_leaves, + input_ids, + mm_embeds=None, + is_multimodal=None): + mm_embeds = flatten_pad_mm_embeds(mm_embeds, + target_pad_len=input_ids.shape[0]) + return jitted_embed_input_ids(state_leaves, + input_ids, + mm_embeds, + is_multimodal=is_multimodal) + + # For models that want to work with EAGLE-3 speculative decoding + @jax.jit(out_shardings=(logits_sharding)) + def combine_hidden_states(state_leaves, hidden_states): + state = jax.tree_util.tree_unflatten(_state_treedef, state_leaves) + model = nnx.merge(graphdef, state) + return model.combine_hidden_states(hidden_states) + + model = nnx.merge(graphdef, state) + precompile_vision_encoder_fn = getattr(model, "precompile_vision_encoder", + None) + # `graphdef` and the state treedef are captured in each closure; the + # runner passes pre-flattened `state_leaves` as the first positional arg. + jitted_model_fn = run_draft_model if is_draft_model else run_model + + model_supports_spec_step = supports_kw(model_class.__call__, + "spec_step_idx") + + def wrapped_model_fn(*args, **kwargs): + if not model_supports_spec_step: + kwargs.pop("spec_step_idx", None) + return jitted_model_fn(*args, **kwargs) + + compute_logits_fn = run_compute_logits + embed_multimodal_fn = run_embed_multimodal + embed_input_ids_fn = run_embed_input_ids + lora_manager, model = None, None + combine_hidden_states_fn = combine_hidden_states + + get_mrope_input_positions_fn = None if not hasattr( + jit_model, + "get_mrope_input_positions") else jit_model.get_mrope_input_positions + + multimodal_fns = MultiModalInterface( + precompile_vision_encoder_fn=precompile_vision_encoder_fn, + embed_multimodal_fn=embed_multimodal_fn, + embed_input_ids_fn=embed_input_ids_fn, + get_mrope_input_positions_fn=get_mrope_input_positions_fn, + ) + + if pooler is not None: + import torchax + from torchax.interop import torch_view + from vllm.v1.pool.metadata import PoolingMetadata + + def compute_pooler_output( + hidden_states: jax.Array, + pooling_metadata: PoolingMetadata, + seq_lens: np.ndarray, + num_scheduled_tokens: Optional[np.ndarray] = None, + ): + # Performance optimization: use torch_view and move to CPU non-blocking + torch_states = torch_view(hidden_states) + with torchax.default_env(): + torch_states = torch_states.to('cpu', non_blocking=True) + + if num_scheduled_tokens is None: + num_scheduled_tokens = seq_lens + + # Align with our StepPool logic + pooling_metadata.build_pooling_cursor( + num_scheduled_tokens, + torch.tensor(seq_lens), + device=torch_states.device, + ) + + # Execute pooling on CPU + outputs = pooler(torch_states, pooling_metadata) + return outputs + + pooler_fn = compute_pooler_output + else: + pooler_fn = _not_support + + state_leaves = tuple(jax.tree_util.tree_leaves(state)) + + return ModelInterface( + model_fn=wrapped_model_fn, + compute_logits_fn=compute_logits_fn, + pooler_fn=pooler_fn, + combine_hidden_states_fn=combine_hidden_states_fn, + multimodal_fns=multimodal_fns, + state=state, + state_leaves=state_leaves, + lora_manager=lora_manager, + model=jit_model, + ) + + +def get_vllm_model( + vllm_config: VllmConfig, + rng: jax.Array, + mesh: Mesh, + is_draft_model: bool = False, + shared_params: Optional[dict[str, jax.Array]] = None, +) -> ModelInterface: + model_dtype = to_torch_dtype(vllm_config.model_config.dtype) + vllm_config.model_config.dtype = model_dtype + from tpu_inference.models.vllm.vllm_model_wrapper import VllmModelWrapper + + model = VllmModelWrapper( + vllm_config=vllm_config, + rng=rng, + mesh=mesh, + is_draft_model=is_draft_model, + ) + params, lora_manager = model.load_weights(shared_params=shared_params) + + jit_model = model.jit_step_func() + compute_logits_fn = model.jit_compute_logits_func() + pooler_fn = model.build_pooler_func() + combine_hidden_states_fn = model.jit_combine_hidden_states_func() + + multimodal_fns = MultiModalInterface( + precompile_vision_encoder_fn=getattr( + model.model.vllm_model, + "precompile_vision_encoder", + model.wrap_precompile_vision_encoder_fn(params), + ), + embed_multimodal_fn=model.wrap_embed_multimodal_func(), + embed_input_ids_fn=model.wrap_embed_input_ids_func(), + get_mrope_input_positions_fn=getattr( + model.model.vllm_model, + "get_mrope_input_positions", + None, + ), + ) + + # the model needs to be returned because lora weights are neither torch.nn.parameter nor torch.nn.buffer. After we load the lora weights and set it to the torch.nn.Module, we can shard it and move it to TPU. + # For the vllm-impl path the dispatch-side fns accept the params dict + # directly, so `state_leaves` is just the dict. + return ModelInterface( + model_fn=jit_model, + compute_logits_fn=compute_logits_fn, + pooler_fn=pooler_fn, + combine_hidden_states_fn=combine_hidden_states_fn, + multimodal_fns=multimodal_fns, + state=params, + state_leaves=params, + lora_manager=lora_manager, + model=model, + ) + + +def get_model( + vllm_config: VllmConfig, + rng: jax.Array, + mesh: Mesh, + is_draft_model: bool = False, + shared_params: Optional[dict[str, jax.Array]] = None, +) -> ModelInterface: + if is_draft_model: + impl = envs.DRAFT_MODEL_IMPL_TYPE + else: + impl = envs.MODEL_IMPL_TYPE + logger.info(f"Loading model with MODEL_IMPL_TYPE={impl}") + if impl == "auto": + impl = resolve_model_architecture(vllm_config, is_draft_model) + logger.info(f"Resolved MODEL_IMPL_TYPE 'auto' to '{impl}'") + + match impl: + case "flax_nnx": + with jax.set_mesh(mesh): + arch = getattr(vllm_config.model_config.hf_config, + "architectures", [None])[0] + if vllm_config.parallel_config.pipeline_parallel_size > 1 and arch in _PP_DISABLED_MODELS: + logger.warning( + "PP is not fully supported on Jax flax_nnx %s models yet, fallback to vllm models.", + arch) + return get_vllm_model(vllm_config, rng, mesh, + is_draft_model, shared_params) + try: + # Try to load the flax model first + return get_flax_model(vllm_config, rng, mesh, + is_draft_model) + except UnsupportedArchitectureError as e: + # Convert the error message to a string to check its contents + error_msg = str(e) + + logger.warning(error_msg) + + # Fall back to the vLLM model and updating the dtype accordingly + return get_vllm_model(vllm_config, rng, mesh, + is_draft_model, shared_params) + case "vllm": + return get_vllm_model(vllm_config, rng, mesh, is_draft_model, + shared_params) + case _: + raise NotImplementedError(f"Unsupported MODEL_IMPL_TYPE: {impl}") + + +def resolve_model_architecture(vllm_config: VllmConfig, + is_draft_model: bool) -> str: + """Resolves the model implementation type. + + This function determines which model implementation to use based on the model + architecture and whether the RunAI model streamer is active. + + When the RunAI model streamer is used, this function explicitly checks if + the JAX model supports the streaming capability. It returns 'vllm' if: + 1. The JAX model class is found but does not have a `WeightLoader`. + 2. The JAX model's `WeightLoader` is not a subclass of `BaseWeightLoader`. + + If the architecture is not registered in JAX (UnsupportedArchitectureError), + this function returns the default implementation ('flax_nnx'), allowing + the caller to attempt loading, catch the error, log a + warning, and handle the fallback to 'vllm'. + + Otherwise, it resolves the implementation based on whether the model's + architecture is in the `_VLLM_PREFERRED_ARCHITECTURES` list. + + Args: + vllm_config: The VllmConfig object containing the model and load + configuration. + + Returns: + The model implementation type. + """ + + is_runai_streamer = getattr(getattr(vllm_config, 'load_config', None), + 'load_format', None) == 'runai_streamer' + hf_config = vllm_config.speculative_config.draft_model_config.hf_config if is_draft_model else vllm_config.model_config.hf_config + if is_runai_streamer: + try: + # Try to get the JAX model class + model_class = _get_model_architecture(hf_config) + + # If found, check for WeightLoader capability + if not hasattr(model_class, "WeightLoader") or not issubclass( + getattr(model_class, "WeightLoader", object), + BaseWeightLoader): + return "vllm" + + except UnsupportedArchitectureError: + # Architecture not in JAX. + # We pass to let it fall through to the default logic below. + # This causes it to return "flax_nnx", which caller will try, + # fail, log a warning, and then fallback to "vllm". + pass + + # Resolve "auto" based on architecture + architectures = getattr(hf_config, "architectures", []) + assert len(architectures) == 1, ( + f"Expected exactly one architecture, got {len(architectures)}: " + f"{architectures}") + arch = architectures[0] + impl = "vllm" if arch in _VLLM_PREFERRED_ARCHITECTURES else "flax_nnx" + return impl + + +def _validate_model_interface(model: Any) -> None: + """Validates that the model class has the required methods and signatures. + + A valid model must have: + - An __init__ method that accepts a 'vllm_config' keyword argument. + - A __call__ method that accepts 'kv_caches', 'input_ids', and + 'attention_metadata' keyword arguments. + + Args: + model: The model class to validate. + + Raises: + TypeError: If the model does not meet the interface requirements. + """ + # Check for __init__ with vllm_config + model_init = getattr(model, "__init__", None) + if not callable(model_init): + raise TypeError( + f"Model {model.__name__} must have an __init__ method.") + + if not supports_kw(model_init, "vllm_config"): + raise TypeError( + f"Model {model.__name__} __init__ method must accept a " + "'vllm_config' keyword argument.") + + # Check for __call__ with required arguments + model_call = getattr(model, "__call__", None) + # A class object is always callable (it produces an instance). + # We need to check if the class _explicitly_ defines a __call__ method for its + # instance, which is different from `type.__call__`. + has_defined_call = False + if isinstance(model, type): + if any("__call__" in C.__dict__ for C in model.__mro__): + has_defined_call = True + elif callable(model_call): + # For an instance, a simple callable check is sufficient. + has_defined_call = True + + if not has_defined_call: + raise TypeError(f"Model {model.__name__} must have a __call__ method.") + + required_call_args = ("kv_caches", "input_ids", "attention_metadata") + missing_args = tuple(arg for arg in required_call_args + if not supports_kw(model_call, arg)) + + if missing_args: + raise TypeError( + f"Model {model.__name__} __call__ method is missing required " + f"keyword arguments: {missing_args}") + + +def register_model(arch: str, model: Any) -> None: + """ + Registers a model class for a given architecture name. + + This function registers the model with both the tpu_inference registry + and the vLLM registry. For vLLM, it creates a compatible wrapper + around the JAX model. + + Args: + arch: The name of the architecture (e.g., "LlamaForCausalLM"). + model: The JAX model class to register (e.g., a flax.nnx.Module). + """ + _validate_model_interface(model) + + # Register with tpu_inference registry for the JAX backend + _MODEL_REGISTRY[arch] = model + + # Create a vLLM-compatible wrapper for the JAX model class. + # This wrapper inherits from the JAX model and torch.nn.Module + # to pass vLLM's type checks. It is not meant to be instantiated + # or executed by vLLM's PyTorch backend. + def unimplemented_forward( + self, + input_ids: "torch.Tensor", + positions: "torch.Tensor", + intermediate_tensors: Optional[Any] = None, + inputs_embeds: Optional["torch.Tensor"] = None, + ) -> None: + raise NotImplementedError( + "This is a JAX model and does not implement the PyTorch forward method." + ) + + # Same as `forward`, this is a dummy method to satisfy vLLM's type checks. + def unimplemented_embed_input_ids( + self, + input_ids: "torch.Tensor", + positions: "torch.Tensor", + inputs_embeds: Optional["torch.Tensor"] = None, + ) -> "torch.Tensor": + raise NotImplementedError( + "This is a JAX model and does not implement the PyTorch embed_input_ids method." + ) + + # We need a custom __init__ that only calls torch.nn.Module's init, + # to avoid triggering JAX logic when vLLM inspects the class. + def wrapper_init(self, *args, **kwargs): + torch.nn.Module.__init__(self) + + # Dynamically create the wrapper class that is a subclass of both the + # JAX model and torch.nn.Module. + VllmCompatibleModel = type( + f"VllmCompatible{model.__name__}", + (model, torch.nn.Module), + { + "__init__": wrapper_init, + "forward": unimplemented_forward, + "embed_input_ids": unimplemented_embed_input_ids, + # Prevent vLLM from trying to load weights into this dummy class. + "load_weights": lambda self, *args, **kwargs: None, + }) + + # Register the wrapped model with vLLM's registry. + from vllm.model_executor.models.registry import ModelRegistry + ModelRegistry.register_model(arch, VllmCompatibleModel) + logger.info( + f"Registered JAX model {arch} with tpu_inference and vLLM registries.") diff --git a/inference/trillium/vLLM/Gemma4-MTP/patch_gemma4_mtp.py b/inference/trillium/vLLM/Gemma4-MTP/patch_gemma4_mtp.py new file mode 100644 index 00000000..6683a0ec --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/patch_gemma4_mtp.py @@ -0,0 +1,122 @@ +import functools +import re + +path = '/workspace/tpu_inference/tpu_inference/models/jax/gemma4_mtp.py' +with open(path, 'r') as f: + text = f.read() + +# 1. Patch the imports +import_target = "from tpu_inference.layers.jax.rope_interface import apply_rope" +import_replacement = "from tpu_inference.layers.jax.rope_interface import apply_rope\nfrom tpu_inference.layers.common.quantization import quantize_kv" +if import_target in text: + text = text.replace(import_target, import_replacement) + print("Imports patched successfully!") +else: + print("Warning: import target not found!") + +# 2. Patch __call__ signature +target_block = """ def __call__( + self, + kv_caches: List[jax.Array], + input_ids: jax.Array, + hidden_states: jax.Array, + attention_metadata: AttentionMetadata,""" +replacement_block = """ def __call__( + self, + kv_caches: List[jax.Array], + input_ids: jax.Array, + hidden_states: Optional[jax.Array] = None, + attention_metadata: Optional[AttentionMetadata] = None,""" +if target_block in text: + text = text.replace(target_block, replacement_block) + print("Call signature patched successfully!") +else: + print("Warning: call signature target not found!") + +# 3. Patch __call__ body (hidden_states initialization) +body_target = ' layer_name_to_kv_cache = (' +body_replacement = """ if hidden_states is None: + backbone_hidden_size = getattr(self.vllm_config.speculative_config.draft_model_config.hf_config, 'backbone_hidden_size', 5376) + draft_hidden_size = self.vllm_config.speculative_config.draft_model_config.hf_config.text_config.hidden_size + hidden_size = 2 * backbone_hidden_size - draft_hidden_size + dtype = self.vllm_config.model_config.dtype + if input_ids.ndim == 1: + seq_len = input_ids.shape[0] + hidden_states = jnp.zeros((seq_len, hidden_size), dtype=dtype) + else: + batch_size, seq_len = input_ids.shape + hidden_states = jnp.zeros((batch_size, seq_len, hidden_size), dtype=dtype) + layer_name_to_kv_cache = (""" +if body_target in text: + text = text.replace(body_target, body_replacement) + print("Call body patched successfully!") +else: + print("Warning: call body target not found!") + +# 4. Patch Gemma4MTPAttention.__init__ (kv_cache_quantized_dtype setup) +init_target = """ self.o_proj = JaxEinsum( + "TNH,NHD->TD", + (self.num_heads, self.head_dim, self.hidden_size), + bias_shape=(self.hidden_size, ) if config.attention_bias else None, + param_dtype=dtype, + kernel_init=nnx.with_partitioning(init_fn, (None, "model", None)), + bias_init=nnx.with_partitioning(init_fn, (None, )) + if config.attention_bias else None, + rngs=rng, + quant_config=quant_config, + prefix=prefix + ".o_proj", + ) + + self.is_kv_shared_layer = True""" +init_replacement = """ self.o_proj = JaxEinsum( + "TNH,NHD->TD", + (self.num_heads, self.head_dim, self.hidden_size), + bias_shape=(self.hidden_size, ) if config.attention_bias else None, + param_dtype=dtype, + kernel_init=nnx.with_partitioning(init_fn, (None, "model", None)), + bias_init=nnx.with_partitioning(init_fn, (None, )) + if config.attention_bias else None, + rngs=rng, + quant_config=quant_config, + prefix=prefix + ".o_proj", + ) + + self.is_kv_shared_layer = True + + self.kv_cache_quantized_dtype = None + if kv_cache_dtype != "auto": + self.kv_cache_quantized_dtype = utils.get_jax_dtype_from_str_dtype( + kv_cache_dtype)""" +if init_target in text: + text = text.replace(init_target, init_replacement) + print("Gemma4MTPAttention.__init__ patched successfully!") +else: + print("Warning: Gemma4MTPAttention.__init__ target not found!") + +# 5. Patch Gemma4MTPAttention.__call__ (quantize dummy_k/dummy_v) +call_target = """ num_tokens = q.shape[0] + dummy_k = jnp.zeros((num_tokens, self.num_kv_heads, self.head_dim), + dtype=q.dtype) + dummy_v = jnp.zeros((num_tokens, self.num_kv_heads, self.head_dim), + dtype=q.dtype) + + new_kv_cache, outputs = attention(""" +call_replacement = """ num_tokens = q.shape[0] + dummy_k = jnp.zeros((num_tokens, self.num_kv_heads, self.head_dim), + dtype=q.dtype) + dummy_v = jnp.zeros((num_tokens, self.num_kv_heads, self.head_dim), + dtype=q.dtype) + + if self.kv_cache_quantized_dtype: + dummy_k, dummy_v = quantize_kv(self.kv_cache_quantized_dtype, dummy_k, dummy_v) + + new_kv_cache, outputs = attention(""" +if call_target in text: + text = text.replace(call_target, call_replacement) + print("Gemma4MTPAttention.__call__ patched successfully!") +else: + print("Warning: Gemma4MTPAttention.__call__ target not found!") + +with open(path, 'w') as f: + f.write(text) +print("gemma4_mtp.py successfully patched!") diff --git a/inference/trillium/vLLM/Gemma4-MTP/patch_qwix.py b/inference/trillium/vLLM/Gemma4-MTP/patch_qwix.py new file mode 100644 index 00000000..67b65e12 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/patch_qwix.py @@ -0,0 +1,46 @@ +path = '/workspace/tpu_inference/tpu_inference/models/jax/utils/qwix/qwix_utils.py' +with open(path, 'r') as f: + text = f.read() + +target_block = """ head_sizes = kv_cache_head_size + if isinstance(head_sizes, int): + head_sizes = (head_sizes, ) * num_hidden_layers + + num_kv_heads_tuple = kv_cache_num_kv_heads + if isinstance(num_kv_heads_tuple, int): + num_kv_heads_tuple = (num_kv_heads_tuple, ) * num_hidden_layers""" + +replacement_block = """ head_sizes = kv_cache_head_size + num_kv_heads_tuple = kv_cache_num_kv_heads + default_head_size = head_sizes if isinstance(head_sizes, int) else head_sizes[0] + default_num_kv_heads = num_kv_heads_tuple if isinstance(num_kv_heads_tuple, int) else num_kv_heads_tuple[0] + dynamic_head_sizes = [default_head_size] * num_hidden_layers + dynamic_num_kv_heads = [default_num_kv_heads] * num_hidden_layers + inner_model = model + if hasattr(model, 'model'): + inner_model = model.model + if hasattr(inner_model, 'layers'): + for idx, layer in enumerate(inner_model.layers): + if idx >= num_hidden_layers: + break + attn = getattr(layer, 'self_attn', getattr(layer, 'attn', None)) + if attn is not None: + if hasattr(attn, 'num_kv_heads'): + dynamic_num_kv_heads[idx] = attn.num_kv_heads + if hasattr(attn, 'head_dim'): + dynamic_head_sizes[idx] = attn.head_dim + head_sizes = tuple(dynamic_head_sizes) + num_kv_heads_tuple = tuple(dynamic_num_kv_heads) + else: + if isinstance(head_sizes, int): + head_sizes = (head_sizes, ) * num_hidden_layers + if isinstance(num_kv_heads_tuple, int): + num_kv_heads_tuple = (num_kv_heads_tuple, ) * num_hidden_layers""" + +if target_block in text: + text = text.replace(target_block, replacement_block) + with open(path, 'w') as f: + f.write(text) + print('qwix_utils.py successfully patched for heterogeneous KV cache!') +else: + print('Error: qwix_utils target block not found!') diff --git a/inference/trillium/vLLM/Gemma4-MTP/recipe/Chart.yaml b/inference/trillium/vLLM/Gemma4-MTP/recipe/Chart.yaml new file mode 100644 index 00000000..dbbf91ca --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/recipe/Chart.yaml @@ -0,0 +1,6 @@ +apiVersion: v2 +name: v6e-benchmark +description: v6e-benchmark +type: application +version: 0.1.0 +appVersion: "1.16.0" diff --git a/inference/trillium/vLLM/Gemma4-MTP/recipe/README.md b/inference/trillium/vLLM/Gemma4-MTP/recipe/README.md new file mode 100644 index 00000000..448650e6 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/recipe/README.md @@ -0,0 +1,92 @@ +# Run vLLM gemma-4-31B-it on v6e on GKE + +This recipe covers running a vLLM inference workload on v6e on +GKE. + +## Create the GKE Cluster +Create your v6e cluster using [XPK](https://github.com/AI-Hypercomputer/xpk). +The next sections assume you have created a cluster with v6e +nodes. + +## Deploy vLLM Workload on GKE + +### Configure kubectl to communicate with your cluster + +``` +gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION} +``` + +### Generate a new Hugging Face token if you don't already have one + +On the huggingface website, create an account if necessary, and go to Your +Profile > Settings > Access Tokens. +Select Create new token. +Specify a name of your choice and a role with at least Read permissions. +Select Generate a token, follow the prompts and save your token. + +(NOTE: Also ensure that your account has access to the model on Hugging Face. +For example, for llama3, you will need to explicitly get permission): + +### Run the benchmark +In this directory, run: + +``` +helm install ${RUN_NAME} . --set hf_token=${HF_TOKEN} +``` + +The benchmark will launch a client and server pod. + +On the server pod, at the end of the server startup you'll see logs such as: + +``` +$ kubectl logs deployment/vllm-tpu -f + +(APIServer pid=1) INFO: Started server process [1] +(APIServer pid=1) INFO: Waiting for application startup. +(APIServer pid=1) INFO: Application startup complete. +``` + +The client pod will wait until the server is up and then start the benchmark. +On the client pod, you'll see logs such as: + +``` +$ kubectl logs -f vllm-bench + +============ Serving Benchmark Result ============ +Successful requests: 10 +Failed requests: 0 +Benchmark duration (s): xx +Total input tokens: xxx +Total generated tokens: xxx +Request throughput (req/s): xx +Output token throughput (tok/s): xxx +Peak output token throughput (tok/s): xxx +Peak concurrent requests: 10.00 +Total Token throughput (tok/s): xxx +---------------Time to First Token---------------- +Mean TTFT (ms): xxx +Median TTFT (ms): xxx +P99 TTFT (ms): xxx +-----Time per Output Token (excl. 1st token)------ +Mean TPOT (ms): xxx +Median TPOT (ms): xxx +P99 TPOT (ms): xxx +---------------Inter-token Latency---------------- +Mean ITL (ms): xxx +Median ITL (ms): xxx +P99 ITL (ms): xxx +================================================== +``` + +### Customizing the benchmark + +In order to change the parameters of the benchmark, such as the model, number of +prompts, etc, you can modify the settings in values.yaml. + +### Cleanup +When you are done running the benchmark, the server will still be running. You +can cleanup by running: + +``` +helm uninstall ${RUN_NAME} +``` diff --git a/inference/trillium/vLLM/Gemma4-MTP/recipe/templates/benchmark.yaml b/inference/trillium/vLLM/Gemma4-MTP/recipe/templates/benchmark.yaml new file mode 100644 index 00000000..b8d0d79a --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/recipe/templates/benchmark.yaml @@ -0,0 +1,707 @@ +# yamllint disable +{{- if ne (int .Values.run_mode) 1 }} +{{- if not .Values.is_multi_host }} +apiVersion: storage.k8s.io/v1 +kind: StorageClass +metadata: + name: "{{ .Release.Name }}-hyperdisk-balanced-tpu" +provisioner: pd.csi.storage.gke.io +parameters: + type: hyperdisk-balanced + # provisioned-iops: "3000" # Optional: Adjust IOPS as needed + # provisioned-throughput: "140" # Optional: Adjust Throughput (MiB/s) as needed +reclaimPolicy: Delete +volumeBindingMode: WaitForFirstConsumer +allowVolumeExpansion: true +--- +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: "{{ .Release.Name }}-hd-claim" +spec: + storageClassName: "{{ .Release.Name }}-hyperdisk-balanced-tpu" + accessModes: + - ReadWriteOnce + resources: + requests: + storage: {{.Values.server.data_disk_size}} +--- +{{- end }} +{{- end }} +{{- if ne (int .Values.run_mode) 1 }} +apiVersion: v1 +kind: Secret +metadata: + name: "{{ .Release.Name }}-hf-secret" +type: Opaque +stringData: + hf_api_token: {{.Values.hf_token}} +--- +{{- end }} +{{- if .Values.gcp_service_account }} +apiVersion: v1 +kind: ServiceAccount +metadata: + name: {{ .Values.k8s_service_account | quote }} + namespace: default + annotations: + iam.gke.io/gcp-service-account: {{ .Values.gcp_service_account }} +{{- end }} +--- +{{- if not .Values.is_multi_host }} +apiVersion: apps/v1 +kind: Deployment +metadata: + name: "{{ .Release.Name }}-tpu-server" +spec: + replicas: 1 + selector: + matchLabels: + app: "{{ .Release.Name }}-tpu-pod" + template: + metadata: + {{- if .Values.gcsfuse.enabled }} + annotations: + gke-gcsfuse/volumes: "true" + {{- end }} + labels: + app: "{{ .Release.Name }}-tpu-pod" + {{- if .Values.kueue_local_queue }} + kueue.x-k8s.io/queue-name: {{ .Values.kueue_local_queue | quote }} + {{- end }} + spec: + {{- if eq (int .Values.run_mode) 1 }} + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + {{- end }} + {{- if .Values.kueue_priority_class }} + priorityClassName: {{ .Values.kueue_priority_class | quote }} + {{- end }} + {{- if .Values.gcp_service_account }} + serviceAccountName: {{ .Values.k8s_service_account | quote }} + {{- end }} + nodeSelector: + {{- if eq (int .Values.run_mode) 1 }} + tpu.google.com/slice-name: {{ .Values.tpu_slice_name }} + tpu.google.com/type: v6e + {{- else }} + cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice + {{- if not .Values.placement_policy }} + cloud.google.com/gke-tpu-topology: {{ .Values.server.topology | quote }} + {{- end }} + {{- if .Values.gke_reservation }} + cloud.google.com/reservation-name: {{ .Values.gke_reservation | quote }} + cloud.google.com/reservation-affinity: specific + {{- end }} + {{- if .Values.placement_policy }} + cloud.google.com/placement-policy-name: {{ .Values.placement_policy | quote }} + {{- end }} + {{- end }} + tolerations: + - key: "google.com/tpu" + operator: "Exists" + containers: + - name: vllm-tpu + image: {{.Values.server.image}} + imagePullPolicy: IfNotPresent + command: ["/bin/bash", "-c"] + args: + - {{ .Values.server.bash_command | quote }} + env: + {{- if eq (int .Values.run_mode) 1 }} + - name: HF_HUB_OFFLINE + value: "1" + {{- else }} + - name: HF_HOME + value: /data + - name: HUGGING_FACE_HUB_TOKEN + valueFrom: + secretKeyRef: + name: "{{ .Release.Name }}-hf-secret" + key: hf_api_token + {{- end }} + {{if .Values.server.profile}} + - name: PHASED_PROFILING_DIR + value: {{ .Values.server.phased_profiling_dir | quote }} + {{end}} + {{if .Values.server.model_impl_type}} + - name: MODEL_IMPL_TYPE + value: {{.Values.server.model_impl_type}} + {{end}} + ports: + - containerPort: 8000 + resources: + limits: + google.com/tpu: "{{.Values.server.num_chips_per_node}}" + requests: + google.com/tpu: "{{.Values.server.num_chips_per_node}}" + readinessProbe: + tcpSocket: + port: 8000 + {{- if .Values.server.readiness_probe }} + initialDelaySeconds: {{ .Values.server.readiness_probe.initial_delay_seconds | default 15 }} + periodSeconds: {{ .Values.server.readiness_probe.period_seconds | default 10 }} + failureThreshold: {{ .Values.server.readiness_probe.failure_threshold | default 3 }} + {{- else }} + initialDelaySeconds: 15 + periodSeconds: 10 + {{- end }} + volumeMounts: + {{- range .Values.custom_volume_mounts }} + - name: {{ .name }} + mountPath: {{ .mountPath }} + {{- if .readOnly }} + readOnly: {{ .readOnly }} + {{- end }} + {{- if .subPath }} + subPath: {{ .subPath }} + {{- end }} + {{- end }} + {{- range .Values.server.pvc_mounts }} + - mountPath: {{ .mount_path | quote }} + name: {{ .volume_name | quote }} + subPath: {{ .sub_path | quote }} + readOnly: {{ .read_only }} + {{- end }} + {{- if eq (int .Values.run_mode) 1 }} + - mountPath: "/data/{{ .Values.server.model_volume_subpath }}" + name: data-volume + subPath: "{{- if .Values.server.models_dir_on_dfs }}{{ .Values.server.models_dir_on_dfs }}/{{ end }}{{ .Values.server.model_volume_subpath }}" + readOnly: true + {{- else }} + - mountPath: "/data" + name: data-volume + {{- end }} + - mountPath: /dev/shm + name: dshm + {{- if .Values.gcsfuse.enabled }} + {{- range .Values.gcsfuse.volumes }} + - name: {{ .name }} + mountPath: {{ .mountPath | quote }} + readOnly: {{ .readOnly }} + {{- end }} + {{- end }} + volumes: + {{- range .Values.custom_volume_mounts }} + - name: {{ .name }} + persistentVolumeClaim: + claimName: {{ $.Release.Name }}-{{ .pvcName }} + {{- end }} + {{- range .Values.server.pvc_mounts }} + - name: {{ .volume_name | quote }} + persistentVolumeClaim: + claimName: {{ .pvc_name | quote }} + {{- end }} + - emptyDir: + medium: Memory + name: dshm + - name: data-volume + persistentVolumeClaim: + {{- if eq (int .Values.run_mode) 1 }} + claimName: model-weights-pvc + {{- else }} + claimName: "{{ .Release.Name }}-hd-claim" + {{- end }} + {{- if .Values.gcsfuse.enabled }} + {{- range .Values.gcsfuse.volumes }} + - name: {{ .name }} + csi: + driver: gcsfuse.csi.storage.gke.io + readOnly: {{ .readOnly }} + volumeAttributes: + bucketName: {{ .bucketName }} + mountOptions: {{ .mountOptions | quote }} + {{- end }} + {{- end }} +{{- else }} +apiVersion: leaderworkerset.x-k8s.io/v1 +kind: LeaderWorkerSet +metadata: + name: "{{ .Release.Name }}-tpu-server" + annotations: + leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool + labels: + {{- if .Values.kueue_local_queue }} + kueue.x-k8s.io/queue-name: {{ .Values.kueue_local_queue | quote }} + {{- end }} +spec: + replicas: 1 + leaderWorkerTemplate: + size: {{ .Values.server.num_nodes }} + restartPolicy: RecreateGroupOnPodRestart + leaderTemplate: + metadata: + {{- if .Values.gcsfuse.enabled }} + annotations: + gke-gcsfuse/volumes: "true" + {{- end }} + labels: + role: leader + app: "{{ .Release.Name }}-tpu-pod" + spec: + {{- if eq (int .Values.run_mode) 1 }} + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + {{- end }} + {{- if .Values.kueue_priority_class }} + priorityClassName: {{ .Values.kueue_priority_class | quote }} + {{- end }} + {{- if .Values.gcp_service_account }} + serviceAccountName: {{ .Values.k8s_service_account | quote }} + {{- end }} + nodeSelector: + {{- if eq (int .Values.run_mode) 1 }} + tpu.google.com/slice-name: {{ .Values.tpu_slice_name }} + tpu.google.com/type: v6e + {{- else }} + cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice + cloud.google.com/gke-tpu-topology: {{ .Values.server.topology | quote }} + {{- if .Values.gke_reservation }} + cloud.google.com/reservation-name: {{ .Values.gke_reservation | quote }} + cloud.google.com/reservation-affinity: specific + {{- end }} + {{- if .Values.placement_policy }} + cloud.google.com/placement-policy-name: {{ .Values.placement_policy | quote }} + {{- end }} + {{- end }} + tolerations: + - key: "google.com/tpu" + operator: "Exists" + containers: + - name: vllm-leader + image: {{.Values.server.image}} + imagePullPolicy: IfNotPresent + env: + {{- if eq (int .Values.run_mode) 1 }} + - name: HF_HUB_OFFLINE + value: "1" + {{- else }} + - name: HUGGING_FACE_HUB_TOKEN + valueFrom: + secretKeyRef: + name: "{{ .Release.Name }}-hf-secret" + key: hf_api_token + {{- end }} + - name: TPU_MULTIHOST_BACKEND + value: ray + - name: JAX_PLATFORMS + value: "" + {{if .Values.server.profile}} + - name: PHASED_PROFILING_DIR + value: {{ .Values.server.phased_profiling_dir | quote }} + {{end}} + {{if .Values.server.model_impl_type}} + - name: MODEL_IMPL_TYPE + value: {{.Values.server.model_impl_type}} + {{end}} + command: ["/bin/bash", "-c"] + args: + - {{ .Values.server.bash_command | quote }} + resources: + limits: + google.com/tpu: "{{.Values.server.num_chips_per_node}}" + requests: + google.com/tpu: "{{.Values.server.num_chips_per_node}}" + ports: + - containerPort: 8000 + readinessProbe: + tcpSocket: + port: 8000 + {{- if .Values.server.readiness_probe }} + initialDelaySeconds: {{ .Values.server.readiness_probe.initial_delay_seconds | default 15 }} + periodSeconds: {{ .Values.server.readiness_probe.period_seconds | default 10 }} + failureThreshold: {{ .Values.server.readiness_probe.failure_threshold | default 3 }} + {{- else }} + initialDelaySeconds: 15 + periodSeconds: 10 + {{- end }} + volumeMounts: + {{- range .Values.custom_volume_mounts }} + - name: {{ .name }} + mountPath: {{ .mountPath }} + {{- if .readOnly }} + readOnly: {{ .readOnly }} + {{- end }} + {{- if .subPath }} + subPath: {{ .subPath }} + {{- end }} + {{- end }} + {{- range .Values.server.pvc_mounts }} + - mountPath: {{ .mount_path | quote }} + name: {{ .volume_name | quote }} + subPath: {{ .sub_path | quote }} + readOnly: {{ .read_only }} + {{- end }} + {{- if eq (int .Values.run_mode) 1 }} + - mountPath: "/data/{{ .Values.server.model_volume_subpath }}" + name: data-volume + subPath: "{{- if .Values.server.models_dir_on_dfs }}{{ .Values.server.models_dir_on_dfs }}/{{ end }}{{ .Values.server.model_volume_subpath }}" + readOnly: true + {{- end }} + - mountPath: /dev/shm + name: dshm + {{- if .Values.gcsfuse.enabled }} + {{- range .Values.gcsfuse.volumes }} + - name: {{ .name }} + mountPath: {{ .mountPath | quote }} + readOnly: {{ .readOnly }} + {{- end }} + {{- end }} + volumes: + {{- range .Values.custom_volume_mounts }} + - name: {{ .name }} + persistentVolumeClaim: + claimName: {{ $.Release.Name }}-{{ .pvcName }} + {{- end }} + {{- range .Values.server.pvc_mounts }} + - name: {{ .volume_name | quote }} + persistentVolumeClaim: + claimName: {{ .pvc_name | quote }} + {{- end }} + - emptyDir: + medium: Memory + name: dshm + {{- if eq (int .Values.run_mode) 1 }} + - name: data-volume + persistentVolumeClaim: + claimName: model-weights-pvc + {{- end }} + {{- if .Values.gcsfuse.enabled }} + {{- range .Values.gcsfuse.volumes }} + - name: {{ .name }} + csi: + driver: gcsfuse.csi.storage.gke.io + readOnly: {{ .readOnly }} + volumeAttributes: + bucketName: {{ .bucketName }} + mountOptions: {{ .mountOptions | quote }} + {{- end }} + {{- end }} + workerTemplate: + {{- if .Values.gcsfuse.enabled }} + metadata: + annotations: + gke-gcsfuse/volumes: "true" + {{- end }} + spec: + {{- if eq (int .Values.run_mode) 1 }} + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + {{- end }} + {{- if .Values.kueue_priority_class }} + priorityClassName: {{ .Values.kueue_priority_class | quote }} + {{- end }} + {{- if .Values.gcp_service_account }} + serviceAccountName: {{ .Values.k8s_service_account | quote }} + {{- end }} + nodeSelector: + {{- if eq (int .Values.run_mode) 1 }} + slice.tpu.google.com/v1alpha1: {{ .Values.tpu_slice_name }} + tpu.google.com/type: v6e + {{- else }} + cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice + cloud.google.com/gke-tpu-topology: {{ .Values.server.topology | quote }} + {{- if .Values.gke_reservation }} + cloud.google.com/reservation-name: {{ .Values.gke_reservation | quote }} + cloud.google.com/reservation-affinity: specific + {{- end }} + {{- if .Values.placement_policy }} + cloud.google.com/placement-policy-name: {{ .Values.placement_policy | quote }} + {{- end }} + {{- end }} + containers: + - name: vllm-worker + image: {{.Values.server.image}} + command: + - sh + - -c + - {{ .Values.server.worker_bash_command | quote }} + resources: + limits: + google.com/tpu: "{{.Values.server.num_chips_per_node}}" + requests: + google.com/tpu: "{{.Values.server.num_chips_per_node}}" + env: + {{- if eq (int .Values.run_mode) 1 }} + - name: HF_HUB_OFFLINE + value: "1" + {{- else }} + - name: HUGGING_FACE_HUB_TOKEN + valueFrom: + secretKeyRef: + name: "{{ .Release.Name }}-hf-secret" + key: hf_api_token + {{- end }} + - name: TPU_MULTIHOST_BACKEND + value: ray + - name: JAX_PLATFORMS + value: "" + {{if .Values.server.profile}} + - name: PHASED_PROFILING_DIR + value: {{ .Values.server.phased_profiling_dir | quote }} + {{end}} + volumeMounts: + {{- range .Values.custom_volume_mounts }} + - name: {{ .name }} + mountPath: {{ .mountPath }} + {{- if .readOnly }} + readOnly: {{ .readOnly }} + {{- end }} + {{- if .subPath }} + subPath: {{ .subPath }} + {{- end }} + {{- end }} + {{- range .Values.server.pvc_mounts }} + - mountPath: {{ .mount_path | quote }} + name: {{ .volume_name | quote }} + subPath: {{ .sub_path | quote }} + readOnly: {{ .read_only }} + {{- end }} + {{- if eq (int .Values.run_mode) 1 }} + - mountPath: "/data/{{ .Values.server.model_volume_subpath }}" + name: data-volume + subPath: "{{- if .Values.server.models_dir_on_dfs }}{{ .Values.server.models_dir_on_dfs }}/{{ end }}{{ .Values.server.model_volume_subpath }}" + readOnly: true + {{- end }} + - mountPath: /dev/shm + name: dshm + {{- if .Values.gcsfuse.enabled }} + {{- range .Values.gcsfuse.volumes }} + - name: {{ .name }} + mountPath: {{ .mountPath | quote }} + readOnly: {{ .readOnly }} + {{- end }} + {{- end }} + volumes: + {{- range .Values.custom_volume_mounts }} + - name: {{ .name }} + persistentVolumeClaim: + claimName: {{ $.Release.Name }}-{{ .pvcName }} + {{- end }} + {{- range .Values.server.pvc_mounts }} + - name: {{ .volume_name | quote }} + persistentVolumeClaim: + claimName: {{ .pvc_name | quote }} + {{- end }} + - emptyDir: + medium: Memory + name: dshm + {{- if eq (int .Values.run_mode) 1 }} + - name: data-volume + persistentVolumeClaim: + claimName: model-weights-pvc + {{- end }} + {{- if .Values.gcsfuse.enabled }} + {{- range .Values.gcsfuse.volumes }} + - name: {{ .name }} + csi: + driver: gcsfuse.csi.storage.gke.io + readOnly: {{ .readOnly }} + volumeAttributes: + bucketName: {{ .bucketName }} + mountOptions: {{ .mountOptions | quote }} + {{- end }} + {{- end }} + +{{- end }} +--- +apiVersion: v1 +kind: Service +metadata: + name: "{{ .Release.Name }}-vllm-service" + {{- if eq (default "ClusterIP" .Values.service_type) "LoadBalancer" }} + annotations: + cloud.google.com/load-balancer-type: "Internal" + {{- end }} +spec: + selector: + app: "{{ .Release.Name }}-tpu-pod" + {{- if .Values.is_multi_host }} + role: leader + {{- end }} + {{- if eq (int .Values.run_mode) 1 }} + clusterIP: None + {{- else }} + type: {{ default "ClusterIP" .Values.service_type }} + ports: + - name: http + protocol: TCP + port: 8000 + targetPort: 8000 + {{- end }} +--- +apiVersion: {{ if eq (int .Values.run_mode) 1 }}batch/v1{{ else }}v1{{ end }} +kind: {{ if eq (int .Values.run_mode) 1 }}Job{{ else }}Pod{{ end }} +metadata: + name: "{{ .Release.Name }}-client" + {{- if .Values.gcsfuse.enabled }} + annotations: + gke-gcsfuse/volumes: "true" + {{- end }} + labels: + {{- if .Values.kueue_local_queue }} + kueue.x-k8s.io/queue-name: {{ .Values.kueue_local_queue | quote }} + {{- end }} +spec: +{{- if eq (int .Values.run_mode) 1 }} + backoffLimit: 0 + template: + {{- if .Values.gcsfuse.enabled }} + metadata: + annotations: + gke-gcsfuse/volumes: "true" + {{- end }} + spec: + hostNetwork: true + affinity: + podAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + - labelSelector: + matchLabels: + app: "{{ .Release.Name }}-tpu-pod" + topologyKey: "kubernetes.io/hostname" +{{- end }} + {{- if .Values.kueue_priority_class }} + priorityClassName: {{ .Values.kueue_priority_class | quote }} + {{- end }} + {{- if .Values.gcp_service_account }} + serviceAccountName: {{ .Values.k8s_service_account | quote }} + {{- end }} + terminationGracePeriodSeconds: 60 +{{- if eq (int .Values.run_mode) 1 }} + tolerations: + - key: "google.com/tpu" + operator: "Exists" + effect: "NoSchedule" +{{- end }} + containers: + - name: vllm-bench + image: {{.Values.client.image}} + imagePullPolicy: IfNotPresent + command: ["/bin/bash", "-c"] + args: + - {{ .Values.client.bash_command | quote }} + env: + {{- if eq (int .Values.run_mode) 1 }} + - name: HF_HUB_OFFLINE + value: "1" + {{- else }} + - name: HUGGING_FACE_HUB_TOKEN + valueFrom: + secretKeyRef: + key: hf_api_token + name: "{{ .Release.Name }}-hf-secret" + {{- end }} + - name: SERVER_HOSTNAME + {{- if eq (int .Values.run_mode) 1 }} + value: "127.0.0.1" + {{- else }} + value: "{{ .Release.Name }}-vllm-service.{{ .Release.Namespace }}.svc.cluster.local" + {{- end }} + volumeMounts: + {{- range .Values.custom_volume_mounts }} + - name: {{ .name }} + mountPath: {{ .mountPath }} + {{- if .readOnly }} + readOnly: {{ .readOnly }} + {{- end }} + {{- if .subPath }} + subPath: {{ .subPath }} + {{- end }} + {{- end }} + - mountPath: /dev/shm + name: dshm + {{- if eq (int .Values.run_mode) 1 }} + - mountPath: "/data/{{ .Values.server.model_volume_subpath }}" + name: data-volume + subPath: "{{- if .Values.server.models_dir_on_dfs }}{{ .Values.server.models_dir_on_dfs }}/{{ end }}{{ .Values.server.model_volume_subpath }}" + readOnly: true + {{- end }} + {{- if .Values.gcsfuse.enabled }} + {{- range .Values.gcsfuse.volumes }} + - name: {{ .name }} + mountPath: {{ .mountPath | quote }} + readOnly: {{ .readOnly }} + {{- end }} + {{- end }} + volumes: + {{- range .Values.custom_volume_mounts }} + - name: {{ .name }} + persistentVolumeClaim: + claimName: {{ $.Release.Name }}-{{ .pvcName }} + {{- end }} + - emptyDir: + medium: Memory + name: dshm + {{- if eq (int .Values.run_mode) 1 }} + - name: data-volume + persistentVolumeClaim: + claimName: model-weights-pvc + {{- end }} + {{- if .Values.gcsfuse.enabled }} + {{- range .Values.gcsfuse.volumes }} + - name: {{ .name }} + csi: + driver: gcsfuse.csi.storage.gke.io + readOnly: {{ .readOnly }} + volumeAttributes: + bucketName: {{ .bucketName }} + mountOptions: {{ .mountOptions | quote }} + {{- end }} + {{- end }} + restartPolicy: Never + + +--- +{{- range .Values.persistent_volumes }} +apiVersion: v1 +kind: PersistentVolume +metadata: + name: {{ $.Release.Name }}-{{ .name }} +spec: + capacity: + storage: {{ .capacity }} + accessModes: + {{ toYaml .accessModes | nindent 4 }} + persistentVolumeReclaimPolicy: {{ .reclaimPolicy }} + {{- if .storageClassName }} + storageClassName: {{ .storageClassName }} + {{- else }} + storageClassName: "" + {{- end }} + {{- if .volumeMode }} + volumeMode: {{ .volumeMode }} + {{- end }} + csi: + driver: {{ .csi.driver }} + volumeHandle: {{ .csi.volumeHandle }} + {{- with .csi.volumeAttributes }} + volumeAttributes: + {{ toYaml . | nindent 6 }} + {{- end }} + claimRef: + name: {{ $.Release.Name }}-{{ .claimRef.name }} + namespace: {{ $.Release.Namespace }} +--- +{{- end }} +--- +{{- range .Values.persistent_volume_claims }} +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: {{ $.Release.Name }}-{{ .name }} + namespace: {{ $.Release.Namespace }} +spec: + accessModes: + {{ toYaml .accessModes | nindent 4 }} + resources: + requests: + storage: {{ .capacity }} + {{- if .storageClassName }} + storageClassName: {{ .storageClassName }} + {{- else }} + storageClassName: "" + {{- end }} + volumeName: {{ $.Release.Name }}-{{ .volumeName }} +--- +{{- end }} diff --git a/inference/trillium/vLLM/Gemma4-MTP/recipe/values.yaml b/inference/trillium/vLLM/Gemma4-MTP/recipe/values.yaml new file mode 100644 index 00000000..989c8954 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/recipe/values.yaml @@ -0,0 +1,99 @@ +"client": + "bash_command": |- + while ! curl http://${SERVER_HOSTNAME}:8000/ping; do sleep 30 && echo 'Waiting for server...'; done + + vllm bench serve \ + --dataset-name=prefix_repetition \ + --prefix-repetition-prefix-len=12000 \ + --prefix-repetition-num-prefixes=15 \ + --prefix-repetition-output-len=200 \ + --host=${SERVER_HOSTNAME} \ + --port=8000 \ + --backend=vllm \ + --request-rate=inf \ + --percentile-metrics=ttft,tpot,itl,e2el \ + --model=google/gemma-4-31B-it \ + --tokenizer=google/gemma-4-31B-it \ + --ignore-eos \ + --skip-chat-template \ + --temperature=0.0 \ + --max-concurrency=320 \ + --num-prompts=320 + "client_type": |- + VLLM_BENCH + "image": |- + gcr.io/northam-ce-mlai-tpu/gemma4-vllm-patched:v6 +"custom_volume_mounts": [] +"gcp_service_account": !!null |- + null +"gcsfuse": + "enabled": !!bool |- + true + "volumes": + - "bucketName": |- + ubench-logs + "mountOptions": |- + implicit-dirs,metadata-cache:negative-ttl-secs:0 + "mountPath": |- + /job-logs + "name": |- + gcs-job-logs + "readOnly": !!bool |- + false +"gke_reservation": !!null |- + null +"hf_token": !!null |- + null +"is_multi_host": !!bool |- + false +"k8s_service_account": !!null |- + null +"kueue_local_queue": !!null |- + null +"kueue_priority_class": |- + medium +"persistent_volume_claims": [] +"persistent_volumes": [] +"placement_policy": !!null |- + null +"run_mode": !!int |- + 0 +"server": + "bash_command": |- + USE_BATCHED_RPA_KERNEL=1 MOE_REQUANTIZE_WEIGHT_DTYPE=float8_e4m3fn SKIP_JAX_PRECOMPILE=1 ATTN_BUCKETIZED_NUM_REQS=1 XLA_PYTHON_CLIENT_MEM_FRACTION=0.75 VLLM_TORCH_COMPILE_CACHE_DIR=/data/torch_compile_cache XLA_CACHE_DIR=/data/xla_cache vllm serve \ + --host=0.0.0.0 \ + --port=8000 \ + --download-dir=/data \ + --max-model-len=16384 \ + --tensor-parallel-size=4 \ + --data-parallel-size=1 \ + --max-num-batched-tokens=4096 \ + --language-model-only \ + --tool-call-parser=gemma4 \ + --enable-auto-tool-choice \ + --disable-chunked-mm-input \ + --enable-log-requests \ + --gpu-memory-utilization=0.75 \ + --model=google/gemma-4-31B-it \ + --tokenizer=google/gemma-4-31B-it + "data_disk_size": |- + 100Gi + "image": |- + gcr.io/northam-ce-mlai-tpu/gemma4-vllm-patched:v6 + "model_impl_type": !!null |- + null + "num_chips_per_node": !!int |- + 4 + "num_nodes": !!int |- + 1 + "phased_profiling_dir": !!null |- + null + "profile": !!bool |- + false + "pvc_mounts": [] + "topology": |- + 2x2 + "worker_bash_command": !!null |- + null +"service_type": |- + ClusterIP diff --git a/inference/trillium/vLLM/Gemma4-MTP/weight_utils.py b/inference/trillium/vLLM/Gemma4-MTP/weight_utils.py new file mode 100644 index 00000000..bb6dea96 --- /dev/null +++ b/inference/trillium/vLLM/Gemma4-MTP/weight_utils.py @@ -0,0 +1,1105 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Utilities for downloading model weights from HuggingFace.""" + +import functools +import glob +import math +import os +import re +import time +from collections import defaultdict +from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Iterable, Optional + +import jax +import jax.numpy as jnp +import numpy as np +import torch +import torchax +from flax import nnx +from jax.sharding import Mesh, NamedSharding +from jax.sharding import PartitionSpec as P +from jax.sharding import SingleDeviceSharding, get_mesh +from safetensors import safe_open +from vllm.config import ModelConfig, VllmConfig, get_current_vllm_config +from vllm.model_executor.model_loader import register_model_loader +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.models.utils import AutoWeightsLoader + +from tpu_inference import envs, utils +from tpu_inference.layers.common.utils import (cpu_mesh_context, + general_device_put) +from tpu_inference.layers.jax import JaxModule, JaxModuleList +from tpu_inference.layers.jax.quantization import QuantizeMethodBase +from tpu_inference.logger import init_logger +from tpu_inference.models.jax.utils import file_utils +from tpu_inference.utils import t2j + +logger = init_logger(__name__) + +HF_WEIGHTS_FORMAT = "*.safetensors" + +DTYPE_VIEW_MAP = { + jnp.dtype(jnp.float8_e4m3fn): torch.uint8, + jnp.dtype(jnp.bfloat16): torch.uint16, + jnp.dtype(jnp.float32): torch.uint32, +} + + +@dataclass +class MetadataMap: + name_map: dict[str, str] = field(default_factory=dict) + transpose_map: dict[str, tuple[int, ...]] = field(default_factory=dict) + reshape_map: dict[str, tuple[int, ...]] = field(default_factory=dict) + bias_reshape_map: dict[str, tuple[int, ...]] = field(default_factory=dict) + pad_map: dict[str, tuple[int, ...]] = field(default_factory=dict) + bias_pad_map: dict[str, tuple[int, ...]] = field(default_factory=dict) + + +############ START Used by llama4, deepseek only for now START ############ + + +def print_param_info(param: nnx.Param, name: str): + logger.warning(f"Global shape for {name}: {param.value.shape}") + logger.warning(f"Sharding for {name}: {param.sharding}") + + logger.warning( + f"Shape of {name} on a single device: {param.value.addressable_shards[0].data.shape}" + ) + + +def transpose_params(param_key: str, param_tensor: jax.Array, transpose_map): + for key, value in transpose_map.items(): + if key in param_key: + return jnp.transpose(param_tensor, value) + return param_tensor # Base case / no-op + + +def reshape_params(param_key: str, param_tensor: jax.Array, shape_map): + for key, new_shape in shape_map.items(): + if key in param_key: + try: + #TODO:(gpolovets) Add validation on whether reshape preserves data layout. + return jnp.reshape(param_tensor, new_shape) + except TypeError: + raise TypeError( + f"Cannot reshape for key={key}, new_shape={new_shape}, param_shape={param_tensor.shape}" + ) + return param_tensor # Base case / no-op + + +def model_file_generator( + model_name_or_path: str, + download_dir: Optional[str]) -> Generator[str, None, None]: + weights_files = get_model_weights_files(model_name_or_path, download_dir) + for st_file in weights_files: + yield st_file + + +def model_weights_generator( + model_name_or_path: str, + framework: str, + filter_regex: Optional[str] = None, + download_dir: Optional[str] = None, +) -> Generator[tuple, None, None]: + for st_file in model_file_generator(model_name_or_path, download_dir): + for name, weight_tensor in model_weights_single_file_generator( + st_file, framework, filter_regex): + yield name, weight_tensor + + +def convert_torch_to_jax_with_view(loaded_weight: torch.Tensor, + cast_type: jnp.dtype) -> jax.Array: + """ + Converts a PyTorch tensor to a JAX array by reinterpreting its + bit representation using a dtype view map. + """ + torch_view_type = DTYPE_VIEW_MAP.get(jnp.dtype(cast_type)) + loaded_weight = jnp.array( + loaded_weight.view(torch_view_type).numpy()).view(cast_type) + return loaded_weight + + +############ END Used by llama4, deepseek only for now END ############ + + +def get_model_weights_files( + model_name_or_path: str, + download_dir: Optional[str]) -> tuple[list[str], str]: + """ + Helper to get weight files and their location. + """ + + if os.path.isdir(model_name_or_path): + logger.info(f"Found weights from local: {model_name_or_path}") + weights_files = glob.glob( + os.path.join(model_name_or_path, HF_WEIGHTS_FORMAT)) + elif file_utils.is_hf_repo(model_name_or_path): + logger.info(f"Downloading weights from HF {model_name_or_path}") + weights_files = file_utils.download_model_weights_from_hf( + model_name_or_path, download_dir, HF_WEIGHTS_FORMAT) + else: + raise ValueError( + f"{model_name_or_path} must be a local directory, or a Huggingface model id." + ) + + if not weights_files: + raise RuntimeError( + f"Cannot find any {HF_WEIGHTS_FORMAT} files in {model_name_or_path}." + ) + + weights_files.sort() + return weights_files + + +def model_weights_single_file_generator( + weights_file: str, + framework: str, + filter_regex: Optional[str] = None, +) -> Generator[tuple, None, None]: + logger.info(f"Loading weights from {weights_file}") + # NOTE: We enforce loading tensors on CPU here. + # Because otherwise the tensor will be loaded on TPU:0 by default, + # although the tensor would eventually be sharded across multiple TPUs, + # it would lead to OOM on TPU:0 for large models. + with jax.default_device(jax.devices("cpu")[0]): + with safe_open(weights_file, framework=framework) as f: + for name in f.keys(): + if filter_regex is not None and not re.match( + filter_regex, name): + continue + weight_tensor = f.get_tensor(name) + yield name, weight_tensor + + +def get_param(params: nnx.State, path: str) -> nnx.State: + keys = path.split(".") + plevel = params + for key in keys: + if key.isdigit(): + plevel = plevel[int(key)] + else: + if key in plevel: + plevel = plevel[key] + else: + raise ValueError(f"{path} is not a valid param path") + return plevel + + +def get_param_and_sharding(params: nnx.State, shardings: Any, + path: str) -> tuple[nnx.State, nnx.State]: + keys = path.split(".") + plevel = params + slevel = shardings + for key in keys: + if key.isdigit(): + plevel = plevel[int(key)] + slevel = slevel[int(key)] + else: + if key in plevel: + plevel = plevel[key] + slevel = slevel[key] + else: + raise ValueError(f"{path} is not a valid param path") + return plevel, slevel.value + + +def shard_put(x: jax.Array, + shardings, + mesh: jax.sharding.Mesh | None = None) -> jax.Array: + # Single device sharding requires this special handling + # to avoid the recursive jit error. + if mesh is None: + mesh = get_mesh() + + x_mesh = None + if isinstance(x.sharding, NamedSharding): + x_mesh = x.sharding.mesh + + if math.prod(mesh.axis_sizes) == 1: + return jax.device_put(x, mesh.devices.flatten()[0]) + + if shardings is None: + shardings = () + + if isinstance(shardings, tuple): + return general_device_put(x, + NamedSharding(mesh, P(*shardings)), + source_mesh=x_mesh) + elif isinstance(shardings, P): + return general_device_put(x, + NamedSharding(mesh, shardings), + source_mesh=x_mesh) + else: + return general_device_put(x, shardings, source_mesh=x_mesh) + + +def get_default_maps(model_config, mesh: Mesh, + name_map: dict[str, str]) -> MetadataMap: + """Load weights from one model weights file to the model, run on single thread.""" + sharding_size = mesh.shape["model"] + + hf_config = model_config.hf_config + if text_config := getattr(hf_config, "text_config", None): + hf_config = text_config + + num_heads = hf_config.num_attention_heads + num_kv_heads = hf_config.num_key_value_heads + hidden_size = model_config.get_hidden_size() + + # Pad head_dim for kernel performance. + head_dim_original = model_config.get_head_size() + + reshape_keys: dict[str, tuple[int, ...]] = { + "q_proj": (num_heads, head_dim_original, hidden_size), + "k_proj": (num_kv_heads, head_dim_original, hidden_size), + "v_proj": (num_kv_heads, head_dim_original, hidden_size), + "o_proj": (hidden_size, num_heads, head_dim_original), + } + bias_reshape_keys: dict[str, tuple[int, ...]] = { + "q_proj.bias": (num_heads, head_dim_original), + "k_proj.bias": (num_kv_heads, head_dim_original), + "v_proj.bias": (num_kv_heads, head_dim_original) + } + transpose_keys: dict[str, tuple[int, ...]] = { + "lm_head": (1, 0), + "fc": (1, 0), + "gate_proj": (1, 0), + "up_proj": (1, 0), + "down_proj": (1, 0), + "q_proj": (2, 0, 1), + "k_proj": (2, 0, 1), + "v_proj": (2, 0, 1), + "o_proj": (1, 2, 0), + } + + # # get vision config + if model_config.is_multimodal_model: + # TODO: Wenlong: Do not consider padding for now + transpose_keys.update({ + "attn.proj": (1, 0), + "attn.qkv": (1, 0), + "visual.merger.mlp": (1, 0), + "visual.patch_embed.proj": (2, 3, 4, 1, 0), + }) + + # key: (padding_dim, padding_size) + pad_keys: dict[str, tuple[int, ...]] = { + "q_proj": (1, sharding_size // num_heads), + "k_proj": (1, sharding_size // num_kv_heads), + "v_proj": (1, sharding_size // num_kv_heads), + "o_proj": (0, sharding_size // num_heads), + } + bias_pad_keys: dict[str, tuple[int, ...]] = { + "q_proj.bias": (0, sharding_size // num_heads), + "k_proj.bias": (0, sharding_size // num_kv_heads), + "v_proj.bias": (0, sharding_size // num_kv_heads), + } + + return MetadataMap(name_map=name_map, + reshape_map=reshape_keys, + bias_reshape_map=bias_reshape_keys, + transpose_map=transpose_keys, + pad_map=pad_keys, + bias_pad_map=bias_pad_keys) + + +def _load_and_shard_weight(vllm_config, + params: nnx.State, + shardings: Any, + metadata_map: MetadataMap, + mesh: Mesh, + hf_key: str, + hf_weight: jax.Array, + keep_hf_weight_suffix_when_match: list[str], + keep_original_dtype_keys_regex: list[str] + | None = None, + pp_missing_layers: list[str] | None = None): + name_map = metadata_map.name_map + reshape_keys = metadata_map.reshape_map + bias_reshape_keys = metadata_map.bias_reshape_map + transpose_keys = metadata_map.transpose_map + pad_keys = metadata_map.pad_map + bias_pad_keys = metadata_map.bias_pad_map + + shard = functools.partial(shard_put, mesh=mesh) + + model_config = vllm_config.model_config + + # Pad head_dim for kernel performance. + head_dim_original = model_config.get_head_size() + head_dim = utils.get_padded_head_dim(head_dim_original) + head_dim_pad = head_dim - head_dim_original + + # Check if the key should retain its original dtype + keep_original_dtype = False + if keep_original_dtype_keys_regex: + for pattern in keep_original_dtype_keys_regex: + if re.match(pattern, hf_key): + keep_original_dtype = True + break + + # Converting to config's dtype + if not keep_original_dtype and hf_weight.dtype != model_config.dtype: + logger.warning( + f"Converting dtype for {hf_key} from {hf_weight.dtype} to {model_config.dtype}" + ) + hf_weight = hf_weight.astype(model_config.dtype) + + # For tensors whose name matches any string in `keep_hf_weight_suffix_when_match`, the + # '.weight' suffix in HF keys will be kept. + # Context: some models are being refactored to have identical parameter names as HF + # models, so the suffix does not need to be removed for those parameters. Eventually + # we want to get rid of the ".weight" suffix removal logic altogether. + # TODO(#1479): remove this argument and related logic after the refactoring is done. + if hf_key.endswith(".weight") and all( + substr not in hf_key + for substr in keep_hf_weight_suffix_when_match): + hf_key = hf_key.removesuffix(".weight") + + # Find the corresponding model key using the HF key + if "layers" in hf_key: + layer_num = re.search(r"layers\.(\d+)", hf_key).group(1) + layer_key = re.sub(r"layers\.\d+", "layers.*", hf_key) + model_key = name_map.get(layer_key, layer_key) + model_key = re.sub(r"layers\.\*", f"layers.{layer_num}", model_key) + elif "blocks" in hf_key: + layer_num = re.search(r"blocks\.(\d+)", hf_key).group(1) + layer_key = re.sub(r"blocks\.\d+", "blocks.*", hf_key) + model_key = name_map.get(layer_key, layer_key) + model_key = re.sub(r"blocks\.\*", f"blocks.{layer_num}", model_key) + else: + if hf_key not in name_map and hf_key == "lm_head": + logger.warning(f"Skip loading {hf_key} due to tie_word_embeddings") + return + if hf_key not in name_map and "t2d" in hf_key: + logger.warning( + f"Skip loading {hf_key} as it's not used in eagle-3 for now") + return + model_key = name_map.get(hf_key, hf_key) + + if pp_missing_layers and _is_pp_missing_layer(hf_key, pp_missing_layers): + logger.warning( + f"Skip loading {hf_key} as it doesn't belong to this PP stage.") + return + model_weight, model_sharding = get_param_and_sharding( + params, shardings, model_key) + + logger.debug( + "before transform | " + f"{hf_key}: {hf_weight.shape} --> {model_key}: {model_weight.value.shape} {model_sharding}" + ) + + if hf_key.endswith(".bias"): + for key in bias_reshape_keys: + if key in hf_key: + hf_weight = jnp.reshape(hf_weight, bias_reshape_keys[key]) + if head_dim_pad > 0: + hf_weight = jnp.pad(hf_weight, ((0, 0), (0, head_dim_pad))) + break + else: + for key in reshape_keys: + if key in hf_key: + hf_weight = jnp.reshape(hf_weight, reshape_keys[key]) + if head_dim_pad > 0: + if "o_proj" in key: + hf_weight = jnp.pad(hf_weight, ((0, 0), (0, 0), + (0, head_dim_pad))) + else: + hf_weight = jnp.pad(hf_weight, + ((0, 0), (0, head_dim_pad), + (0, 0))) + break + for key in transpose_keys: + if key in hf_key: + hf_weight = jnp.transpose(hf_weight, transpose_keys[key]) + break + + # Pad num-kv-heads + if hf_key.endswith(".bias"): + for key, value in bias_pad_keys.items(): + dim = value[0] + dim_size = value[1] + if key in hf_key and dim_size != 0: + hf_weight = jnp.repeat(hf_weight, dim_size, axis=dim) + break + else: + for key, value in pad_keys.items(): + dim = value[0] + dim_size = value[1] + if key in hf_key and dim_size != 0: + hf_weight = jnp.repeat(hf_weight, dim_size, axis=dim) + break + + logger.debug( + "after transform | " + f"{hf_key}: {hf_weight.shape} --> {model_key}: {model_weight.value.shape} {model_sharding}" + ) + + if head_dim_pad == 0: + assert model_weight.value.shape == hf_weight.shape, f"{hf_key}: {model_weight.value.shape} != {hf_weight.shape}" + + # Update the model weight + spec = model_sharding.spec if isinstance(model_sharding, + NamedSharding) else model_sharding + model_weight.value = shard(hf_weight, spec) + + +def _is_pp_missing_layer(hf_key: str, pp_missing_layers: list[str]) -> bool: + has_digit = any(char.isdigit() for char in hf_key) + # add the suffix after digits to avoid it matches "layers.10" with "layers.1" + suffix = "." if has_digit else "" + return any(f'{pp_missing_layer}{suffix}' in hf_key + for pp_missing_layer in pp_missing_layers) + + +def _load_hf_weights_on_thread( + vllm_config: VllmConfig, + params: nnx.State, + metadata_map: "MetadataMap", + mesh: Mesh, + weights_file: str, + keep_hf_weight_suffix_when_match: list[str], + filter_regex: Optional[str] = None, + keep_original_dtype_keys_regex: Optional[list[str]] = None, + pp_missing_layers: list[str] | None = None, +): + """Loads weights from a single weights file.""" + try: + shardings = nnx.get_named_sharding(params, mesh) + except TypeError: + shardings = params + + for hf_key, hf_weight in model_weights_single_file_generator( + weights_file, framework="flax", filter_regex=filter_regex): + _load_and_shard_weight( + vllm_config, + params, + shardings, + metadata_map, + mesh, + hf_key, + hf_weight, + keep_original_dtype_keys_regex=keep_original_dtype_keys_regex, + pp_missing_layers=pp_missing_layers, + keep_hf_weight_suffix_when_match=keep_hf_weight_suffix_when_match, + ) + + +def load_hf_weights( + vllm_config: VllmConfig, + model: nnx.Module, + metadata_map: "MetadataMap", + mesh: Mesh, + filter_regex: Optional[str] = None, + is_draft_model: bool = False, + keep_original_dtype_keys_regex: Optional[list[str]] = None, + pp_missing_layers: list[str] | None = None, + keep_hf_weight_suffix_when_match: list[str] = [], +): + """Load weights into a JAX model from either an iterator or files. + + For tensors whose name matches any string in `keep_hf_weight_suffix_when_match`, the + '.weight' suffix in HF keys will be kept. + Some models are being refactored to have identical parameter names as HF models, so the suffix + does not need to be removed for those parameters. Eventually we want to get rid of + the ".weight" suffix removal logic altogether. + TODO(#1479): remove this argument and related logic after the refactoring is done. + """ + params = nnx.state(model) + try: + shardings = nnx.get_named_sharding(params, mesh) + except TypeError: + shardings = params + if is_draft_model: + import jax + replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) + shardings = jax.tree_util.tree_map(lambda _: replicated_sharding, shardings) + weights_iterator = None + if hasattr(vllm_config.model_config, "runai_model_weights_iterator"): + weights_iterator = vllm_config.model_config.runai_model_weights_iterator + env = torchax.default_env() + # The weights_iterator is used in RunAI model streamer integration. + if weights_iterator is not None: + for hf_key, hf_weight in weights_iterator: + if filter_regex and not re.match(filter_regex, hf_key): + continue + + # Since the weights_iterator yields Pytorch tensors (torch.Tensor), + # we need to convert them to JAX arrays (jax.Array). + hf_weight_jax = env.t2j_copy(hf_weight) + + _load_and_shard_weight( + vllm_config, + params, + shardings, + metadata_map, + mesh, + hf_key, + hf_weight_jax, + keep_original_dtype_keys_regex=keep_original_dtype_keys_regex, + pp_missing_layers=pp_missing_layers, + keep_hf_weight_suffix_when_match= + keep_hf_weight_suffix_when_match, + ) + else: + # File-based path (multi-threaded) + if is_draft_model: + model_path = vllm_config.speculative_config.draft_model_config.model + else: + model_path = vllm_config.model_config.model + weights_files = get_model_weights_files( + model_path, vllm_config.load_config.download_dir) + max_workers = min(64, len(weights_files)) + # NOTE(xiang): Disable multi-threading mode if running on multi-host. + # Because multi-threading would cause different JAX processes to load + # different weights at the same time. + if envs.TPU_MULTIHOST_BACKEND == "ray": + max_workers = 1 + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit( + _load_hf_weights_on_thread, + vllm_config, + params, + metadata_map, + mesh, + weights_file, + filter_regex=filter_regex, + keep_original_dtype_keys_regex= + keep_original_dtype_keys_regex, + pp_missing_layers=pp_missing_layers, + keep_hf_weight_suffix_when_match= + keep_hf_weight_suffix_when_match, + ) for weights_file in weights_files + ] + for future in futures: + future.result() + + check_all_loaded(params) + nnx.update(model, params) + + +def check_all_loaded(params: nnx.State): + + def _check(x: Any): + if isinstance(x, nnx.Param) and isinstance(x.value, + jax.ShapeDtypeStruct): + raise ValueError(f"The param does not load weights: {x}") + + jax.tree.map(_check, params) + + +def build_flat_dict(flat_state, mappings): + """Build a new flat dictionary from the flat state using the provided mappings.""" + new_flat_dict = {} + for keys, v in flat_state: + path = '.'.join(str(key) for key in keys) + mapped = False + for src, (tgt, sharding) in mappings.items(): + regex = "^" + re.escape(tgt).replace("\\.\\*", r"\.(\d+)") + "$" + matched = re.match(regex, path) + if matched: + # Extract wildcards if any + wildcards = matched.groups() + src_parts = [] + wc_index = 0 + for part in src.split("."): + if part == "*": + src_parts.append(wildcards[wc_index]) + wc_index += 1 + else: + src_parts.append(part) + actual_src = ".".join(src_parts) + new_flat_dict[actual_src] = v, sharding + mapped = True + break + if not mapped: + logger.info(f"!!! No mapping for flat state: {keys}") + return new_flat_dict + + +def transfer_state_with_mappings(src_state, + tgt_state, + mappings, + transpose_keys=None, + shard=None): + """Transfer state from src_state to tgt_state using the provided mappings.""" + src_flat = src_state.flat_state() + tgt_flat = tgt_state.flat_state() + + new_src_dict = build_flat_dict(tgt_flat, mappings) + logger.info(f"{mappings=}") + logger.info(f"{transpose_keys=}") + for src_keys, v in src_flat: + flattened_src_keys = '.'.join(str(k) for k in src_keys) + new_v = jnp.copy(v.value) + logger.info( + f"Processing source key: {flattened_src_keys} and value: {new_v.shape} {new_v.dtype}" + ) + if flattened_src_keys not in new_src_dict: + logger.info(f"!!! No mapping for source key: {flattened_src_keys}") + continue + sharding = new_src_dict[flattened_src_keys][1] + + # E.g. layers.*.attn.k_proj.w, layers.*.attn.k_proj.w_lora_a + # E.g. layers.*.mlp.down_proj.kernel, layers.*.mlp.down_proj.kernel_lora_a + if transpose_keys is not None \ + and ((src_keys[-1] in transpose_keys) and ('lora' not in src_keys[-1])): + v_maybe_t = jnp.transpose(new_v, transpose_keys[src_keys[-1]]) + else: + v_maybe_t = new_v + + to_update_value = new_src_dict[flattened_src_keys][0].value + assert to_update_value.shape == v_maybe_t.shape, \ + f"Shape mismatch for {flattened_src_keys}: {to_update_value.shape} vs {v_maybe_t.shape}" + + if to_update_value.dtype != v_maybe_t.dtype: + logger.info( + f"Type mismatch between external model and vLLM model. Converting {v_maybe_t.dtype=} to {to_update_value.dtype=}" + ) + v_maybe_t = v_maybe_t.astype(to_update_value.dtype) + + new_src_dict[flattened_src_keys][0].value = shard( + v_maybe_t, sharding) if shard else v_maybe_t + + tgt_state = tgt_state.from_flat_path(tgt_flat) + return tgt_state + + +class BaseWeightLoader: + + def __init__(self, vllm_config: VllmConfig, **kwargs): + self.vllm_config = vllm_config + self.names_and_weights_generator = model_weights_generator( + model_name_or_path=vllm_config.model_config.model, + download_dir=vllm_config.load_config.download_dir, + **kwargs, + ) + + def get_weights_iterator(self): + weights_iterator = getattr(self.vllm_config.model_config, + "runai_model_weights_iterator", None) + if weights_iterator: + return weights_iterator + else: + return self.names_and_weights_generator + + +class StandardWeightLoader(BaseWeightLoader): + + def __init__(self, vllm_config: VllmConfig, mesh: Mesh): + super().__init__(vllm_config, framework="pt") + self.vllm_config = vllm_config + self.mesh = mesh + + def load_weights(self, + model: nnx.Module, + mappings: dict | MetadataMap, + keep_hf_weight_suffix_when_match: list[str] = []): + """ + Calls the generic load_hf_weights utility, passing the correct + weights iterator. + + `mappings` can be either a MetadataMap or a dict mapping, if it's + * a dict, the default MetadataMap will be created with get_default_maps. + * a MetadataMap, it will be used directly. This is useful for cases + where caller needs to customize the reshape/transpose/pad maps, e.g. + update the key of tranpose_map from "q_proj" to "q_proj.weight". + + Context: some models are being refactored to have identical parameter names as HF + models, so these parameters keeps ".weight" suffix. Eventually + we want to get rid of the ".weight" suffix removal logic altogether. + TODO(#1479): remove this argument and related logic after the refactoring is done. + """ + if isinstance(mappings, MetadataMap): + metadata_map = mappings + else: + metadata_map = get_default_maps(self.vllm_config.model_config, + self.mesh, mappings) + + load_hf_weights( + vllm_config=self.vllm_config, + model=model, + metadata_map=metadata_map, + mesh=self.mesh, + pp_missing_layers=getattr(model, 'pp_missing_layers', []), + keep_hf_weight_suffix_when_match=keep_hf_weight_suffix_when_match) + + +def jax_array_from_reshaped_torch( + torch_weight: torch.Tensor, + *, + reshape_dims: Optional[tuple[int, ...]] = None, + permute_dims: Optional[tuple[int, ...]] = None) -> jax.Array: + """Convert a torch.Tensor to a jax.Array with reshaping and transposing. + + HuggingFace model almost always store linear layer weights with contracting dimension + last, and only support 1D/2D weight tensors. This function reshapes then transposes + the torch weight to match the jax_param shape before loading. + + Args: + torch_weight: The source torch.Tensor weight. + reshape_dims: Optional tuple specifying the shape to reshape the torch weight to before permutation. If None, no reshaping is applied. + permute_dims: Optional tuple specifying the permutation of dimensions. If None, no-op for 1D tensors and transpose for 2D tensors is applied. + """ + if reshape_dims is not None: + torch_weight = torch_weight.reshape(reshape_dims) + if permute_dims is None and torch_weight.ndim == 2: + permute_dims = (1, 0) + if permute_dims is not None: + torch_weight = torch_weight.permute(*permute_dims) + + with cpu_mesh_context(): + return t2j(torch_weight, use_dlpack=False) + + +def assign_and_shard_param(jax_param: nnx.Param, + jax_weight: jax.Array, + param_name: str = "Unknown", + mesh: Optional[Mesh] = None) -> None: + """Distributes a JAX array across devices according to the `nnx.Param`'s sharding metadata, assigns it to the parameter, and marks it as loaded. + + Args: + jax_param: The target nnx.Param to assign the weight to. + jax_weight: The JAX array containing the weight data. + param_name: The name of the parameter, used for error logging. + mesh: The device mesh to shard the parameter on. + """ + spec = jax_param.get_metadata().get("sharding", ()) + if isinstance(spec, NamedSharding): + spec = spec.spec + elif isinstance(spec, SingleDeviceSharding): + spec = () + param_mesh = jax_param.get_metadata().get("mesh") or mesh + shape = jax_weight.shape + try: + jax_param.value = shard_put(jax_weight, spec, mesh=param_mesh) + jax_param.set_metadata("_is_loaded", True) + del jax_weight + jax.clear_caches() + except Exception as e: + raise RuntimeError( + f"Failed to load weight '{param_name}' with shape {shape} into param with shape {jax_param.value.shape}" + ) from e + + +def load_nnx_param_from_reshaped_torch( + jax_param: nnx.Param, + torch_weight: torch.Tensor, + *, + reshape_dims: Optional[tuple[int, ...]] = None, + permute_dims: Optional[tuple[int, ...]] = None, + param_name: str = "Unknown"): + """Load a nnx.Param from a torch.Tensor with reshaping and transposing. + + HuggingFace model almost always store linear layer weights with contracting dimension + last, and only support 1D/2D weight tensors. This function reshapes then transposes + the torch weight to match the jax_param shape before loading. + + Args: + jax_param: The target nnx.Param to load the weight into. + torch_weight: The source torch.Tensor weight. + reshape_dims: Optional tuple specifying the shape to reshape the torch weight to before permutation. If None, no reshaping is applied. + permute_dims: Optional tuple specifying the permutation of dimensions. If None, no-op for 1D tensors and transpose for 2D tensors is applied. + """ + try: + jax_weight = jax_array_from_reshaped_torch(torch_weight, + reshape_dims=reshape_dims, + permute_dims=permute_dims) + except Exception as e: + raise RuntimeError( + f"Failed to convert torch weight for '{param_name}' ({torch_weight.shape}) to JAX array, with reshape_dims={reshape_dims} and permute_dims={permute_dims}" + ) from e + + if jax_weight.shape != jax_param.value.shape: + # Retrieve current config to determine if the model task is embedding/pooling + is_embedding_task = False + try: + is_embedding_task = (get_current_vllm_config().model_config. + runner_type == "pooling") + except Exception as e: + logger.debug( + "Failed to retrieve vllm_config to check for embedding task status. Defaulting is_embedding_task to False: %s", + e) + is_embedding_task = False + + is_vocab_layer = param_name.endswith( + (".embed_tokens.weight", ".lm_head.weight", ".wte.weight", + ".pooler.weight", "embed_tokens.weight", "lm_head.weight", + "wte.weight", "pooler.weight")) + + if is_vocab_layer or is_embedding_task: + if all( + w <= p + for w, p in zip(jax_weight.shape, jax_param.value.shape) + ) and any( + w < p + for w, p in zip(jax_weight.shape, jax_param.value.shape)): + pad_width = tuple( + (0, p - w) + for w, p in zip(jax_weight.shape, jax_param.value.shape)) + logger.info( + f"Padding weight '{param_name}' from {jax_weight.shape} to {jax_param.value.shape}" + ) + # Use NumPy to keep it on Host + jax_weight_np = np.pad(np.asarray(jax_weight), pad_width) + with cpu_mesh_context(): + jax_weight = jnp.array(jax_weight_np) + else: + raise ValueError( + f"Shape mismatch in non-vocab layer '{param_name}': torch {jax_weight.shape} vs jax {jax_param.value.shape}" + ) + + assert tuple(jax_weight.shape) == jax_param.value.shape, \ + f"Shape mismatch when loading weight '{param_name}': torch {jax_weight.shape} vs jax {jax_param.value.shape}" + + assign_and_shard_param(jax_param, jax_weight, param_name) + + +class JaxAutoWeightsLoader(AutoWeightsLoader): + """A weights loader for JAX models.""" + + def __init__(self, + model, + pytorch_pooler: Optional[torch.nn.Module] = None, + **kwargs): + assert isinstance(model, JaxModule) + self.pytorch_pooler = pytorch_pooler + self.pooler_weights = {} + + for name, param in model.named_parameters(): + if not hasattr(param, "weight_loader"): + # Following are common patterns in standard transformers. To add pattern for modules + # beyond standard transformers, please consider setting weight_loader. + reshape_dims = None + permute_dims = None + if any(substr in name + for substr in ["k_proj.weight", "v_proj.weight"]): + D, N, H = param.value.shape + reshape_dims = (N, H, D) + permute_dims = (2, 0, 1) + if any(substr in name for substr in ["q_proj.weight"]): + if envs.LAYOUT_Q_PROJ_AS_NDH: + N, D, H = param.value.shape + reshape_dims = (N, H, D) + permute_dims = (0, 2, 1) + else: + D, N, H = param.value.shape + reshape_dims = (N, H, D) + permute_dims = (2, 0, 1) + elif any(substr in name for substr in + ["q_proj.bias", "k_proj.bias", "v_proj.bias"]): + N, H = param.value.shape + reshape_dims = (N, H) + permute_dims = (0, 1) + elif "o_proj.weight" in name: + N, H, D = param.value.shape + reshape_dims = (D, N, H) + permute_dims = (1, 2, 0) + elif "embed_tokens.weight" in name: + permute_dims = (0, 1) + elif "lm_head" in name: + permute_dims = (1, 0) + + param.set_metadata( + "weight_loader", + functools.partial(load_nnx_param_from_reshaped_torch, + reshape_dims=reshape_dims, + permute_dims=permute_dims, + param_name=name)) + + super().__init__(model, **kwargs) + # Book mark those already done processing, skip if visited. + self._process_weights_after_loading_per_module = defaultdict( + lambda: False) + + def _add_loadable_non_param_tensors(self, module: JaxModule, + child_params: dict[str, Any]): + """ + Add tensor names that are not in the model params that may be in the + safetensors, e.g., batch normalization stats and registered buffers. + """ + ... + + def _load_module(self, base_prefix: str, module: JaxModule, + weights: Iterable) -> Iterable: + """Load weights into the JAX module, performing prefix adjustments and interception. + + Args: + base_prefix: The prefix string of the current base module. + module: The JAX module into which the weights are loaded. + weights: The iterable collection of torch weights to load. + + Returns: + Iterable: The iterable collection after performing custom mappings and intercepts. + """ + + def _map_weights(w_iter): + if base_prefix == "": + root_children = {name for name, _ in module.named_children()} + has_model_child = "model" in root_children + else: + root_children = set() + has_model_child = False + + for name, weight in w_iter: + if self.pytorch_pooler is not None: + match = re.match(r"^(?:model\.)?pooler\.(.*)$", name) + if match: + pooler_key = match.group(1) + self.pooler_weights[pooler_key] = weight + continue + + if base_prefix == "": + top_level = name.split('.')[0] + if top_level in root_children: + # Path A: Root child, load as-is + pass + elif has_model_child and not name.startswith("model."): + # Path B: Not root child, but root has 'model', prepend 'model.' + name = "model." + name + # Path C: Fallback to as-is + + yield name, weight + + yield from super()._load_module(base_prefix, module, + _map_weights(weights)) + + if base_prefix == "" and self.pytorch_pooler is not None and self.pooler_weights: + logger.info( + f"Loading {len(self.pooler_weights)} weights into CPU Pooler") + self.pytorch_pooler.load_state_dict(self.pooler_weights, + strict=False) + # Post-process module after loading weights. Unlike vLLM post-process + # weights after loading all weights, we do it per-module here to + # avoid OOM. + if self._process_weights_after_loading_per_module[base_prefix]: + return + if (quant_method := getattr(module, 'quant_method', None)) is not None: + assert isinstance(quant_method, QuantizeMethodBase) + loaded = quant_method.process_weights_after_loading(module) + jax.clear_caches() + assert isinstance(loaded, bool) + self._process_weights_after_loading_per_module[ + base_prefix] = loaded + + +class LoadableWithIterator: + """Mixin for models that support loading weights with an iterator. + + This is replicating what vLLM does for most models, e.g. https://github.com/vllm-project/vllm/blob/8e2a469b3b2f67bc900ed72724fe3f05e3564994/vllm/model_executor/models/gemma3_mm.py#L644-L646 + """ + + def load_weights(self, weights: Iterable[tuple[str, + torch.Tensor]]) -> set[str]: + if not isinstance(weights, Iterable): + # Use next parent class in MRO. + return super().load_weights(weights) + + pytorch_pooler = getattr(getattr(self, "vllm_config", None), + "pytorch_pooler", None) + loader = JaxAutoWeightsLoader( + self, + pytorch_pooler=pytorch_pooler, + skip_prefixes=(["lm_head"] + if not hasattr(self, 'lm_head') else None)) + return loader.load_weights(weights) + + +@register_model_loader("jax_dummy") +class JaxDummyModelLoader(DummyModelLoader): + """A dummy weights loader for flax_nnx models. + + The upstream DummyModelLoader relies on many torch-specific APIs, this + implementation overrides the load_weights method to support flax_nnx models. + """ + + def load_weights(self, model: JaxModule, + model_config: ModelConfig) -> None: + weight_loading_start_counter = time.perf_counter() + mesh = jax.sharding.get_mesh() + + def _load_dummy_weight_on_thread(param_name, param): + with cpu_mesh_context(): + is_moe = hasattr(param, "_weights_to_load") + param_shape = param.value.shape + + if is_moe: + # For MoE parameters, the normal loading flow reads PyTorch + # weights which are transposed (out_features, in_features) + # compared to JAX. The downstream post-loading fusion methods + # expect this transposed shape (E, F, D) instead of (E, D, F). + # E = number of experts, D = input dimension, F = feed forward dimension + num_experts, input_dim, intermediate_dim = param_shape + param_shape = (num_experts, intermediate_dim, input_dim) + + if jnp.issubdtype(param.value.dtype, jnp.integer): + dummy_weight = jax.random.randint( + key=jax.random.PRNGKey(0), + shape=param_shape, + minval=0, + maxval=100, + dtype=param.value.dtype, + ) + else: + dummy_weight = jax.random.uniform( + key=jax.random.PRNGKey(0), + shape=param_shape, + dtype=param.value.dtype, + # upstream claims this range works well + # https://github.com/vllm-project/vllm/blob/7291d1b288558d48508e1a17c37b0aa170332264/vllm/model_executor/model_loader/weight_utils.py#L1088 + minval=-1e-3, + maxval=1e-3, + ) + + if is_moe: + param._weights_to_load[:] = jnp.vsplit( + dummy_weight, indices_or_sections=num_experts) + + # We must explicitly pass the `mesh` captured from the main thread + # into the worker threads. JAX mesh contexts are thread-local, so + # worker threads do not inherit the active TPU mesh. + assign_and_shard_param(param, dummy_weight, param_name, mesh=mesh) + + with ThreadPoolExecutor(max_workers=64) as executor: + futures = [ + executor.submit(_load_dummy_weight_on_thread, param_name, + param) + for param_name, param in model.named_parameters() + ] + for future in futures: + future.result() + + self._process_weights_after_loading(model) + logger.info_once( + f"Loading dummy weights took {time.perf_counter() - weight_loading_start_counter:.2f} seconds." + ) + + def _process_weights_after_loading( + self, module: JaxModule | JaxModuleList) -> None: + """Recursively call process_weights_after_loading if any.""" + if (quant_method := getattr(module, 'quant_method', None)) is not None: + assert isinstance(quant_method, QuantizeMethodBase) + quant_method.process_weights_after_loading(module) + return + if isinstance(module, JaxModuleList): + for sub_module in module: + self._process_weights_after_loading(sub_module) + else: + for name, sub_module in module.named_children(): + self._process_weights_after_loading(sub_module)