tpu-inference 0.11.1.dev202511220812__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tpu_inference/models/jax/utils/weight_utils.py

@@ -1,3 +1,16 @@
+# 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
@@ -13,10 +26,12 @@ from typing import Any, Optional
 import jax
 import jax.numpy as jnp
 import torch
+import torchax
 from flax import nnx
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 from safetensors import safe_open
+from vllm.config import VllmConfig
 
 from tpu_inference import envs, utils
 from tpu_inference.logger import init_logger
@@ -65,7 +80,13 @@ def transpose_params(param_key: str, param_tensor: jax.Array, transpose_map):
 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:
-            return jnp.reshape(param_tensor, new_shape)
+            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
 
 
@@ -265,15 +286,16 @@ def get_default_maps(model_config, mesh: Mesh,
                        bias_pad_map=bias_pad_keys)
 
 
-def _load_hf_weights_on_thread(vllm_config,
-                               params: nnx.State,
-                               metadata_map: MetadataMap,
-                               mesh: Mesh,
-                               weights_file: str,
-                               filter_regex: str | None = None,
-                               keep_original_dtype_keys_regex: list[str]
-                               | None = None,
-                               exclude_regex: list[str] | None = None):
+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_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
@@ -290,6 +312,131 @@ def _load_hf_weights_on_thread(vllm_config,
     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)
+
+    if hf_key.endswith(".weight"):
+        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[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[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_weight.sharding.spec if isinstance(
+        model_weight.sharding, NamedSharding) else model_weight.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,
+    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:
@@ -297,160 +444,92 @@ def _load_hf_weights_on_thread(vllm_config,
 
     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,
+            pp_missing_layers,
+        )
 
-        # Check if the key should be excluded
-        if exclude_regex:
-            should_exclude = False
-            for pattern in exclude_regex:
-                if re.search(pattern, hf_key):
-                    logger.info(
-                        f"Excluding {hf_key} based on pattern {pattern}")
-                    should_exclude = True
-                    break
-            if should_exclude:
-                continue
-
-        # 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)
-
-        if hf_key.endswith(".weight"):
-            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[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[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")
-                continue
-            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"
-                )
+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,
+):
+    """Load weights into a JAX model from either an iterator or files."""
+    params = nnx.state(model)
+    try:
+        shardings = nnx.get_named_sharding(params, mesh)
+    except TypeError:
+        shardings = params
+    weights_iterator = None
+    if hasattr(vllm_config.model_config, "model_weights_iterator"):
+        weights_iterator = vllm_config.model_config.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
-            model_key = name_map.get(hf_key, hf_key)
-        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}"
-        )
+            # 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)
 
-        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_weight.sharding.spec if isinstance(
-            model_weight.sharding, NamedSharding) else model_weight.sharding
-        model_weight.value = shard(hf_weight, spec)
-
-
-def load_hf_weights(vllm_config,
-                    model: nnx.Module,
-                    metadata_map: MetadataMap,
-                    mesh: Mesh,
-                    filter_regex: str | None = None,
-                    is_draft_model: bool = False,
-                    keep_original_dtype_keys_regex: list[str] | None = None,
-                    exclude_regex: list[str] | None = None):
-    """Load weights from all model weights files to the model, run in multi threads."""
-    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)
-    params = nnx.state(model)
-    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,
+            _load_and_shard_weight(
                 vllm_config,
                 params,
+                shardings,
                 metadata_map,
                 mesh,
-                weights_file,
-                filter_regex=filter_regex,
-                keep_original_dtype_keys_regex=keep_original_dtype_keys_regex,
-                exclude_regex=exclude_regex) for weights_file in weights_files
-        ]
-        for future in futures:
-            future.result()
+                hf_key,
+                hf_weight_jax,
+                keep_original_dtype_keys_regex,
+                pp_missing_layers=pp_missing_layers,
+            )
+    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,
+                ) for weights_file in weights_files
+            ]
+            for future in futures:
+                future.result()
+
     check_all_loaded(params)
     nnx.update(model, params)
 

tpu_inference/models/vllm/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tpu_inference/models/vllm/vllm_model_wrapper.py

@@ -1,3 +1,17 @@
+# 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.
+
 import copy
 import functools
 from collections.abc import Sequence
@@ -9,6 +23,7 @@ import jax
 import torch
 import torch.nn
 import torchax
+import vllm.envs as vllm_envs
 from flax.typing import PRNGKey
 from jax.sharding import Mesh, NamedSharding, PartitionSpec
 from torchax.interop import jax_view, torch_view
@@ -22,8 +37,10 @@ from vllm.model_executor.models import supports_lora, supports_multimodal
 from vllm.sequence import IntermediateTensors
 
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.vllm.process_weights.cleanup_sharding import \
+    shard_model_to_tpu
 from tpu_inference.layers.vllm.quantization import get_tpu_quantization_config
-from tpu_inference.layers.vllm.sharding import shard_model_to_tpu
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.jax_intermediate_tensor import \
     JaxIntermediateTensors
@@ -118,9 +135,16 @@ class VllmModelWrapper:
             "torch._sync",
             return_value=None) if use_random_weights else nullcontext()
 
+        # By default load weights to the CPU device first. If we are running
+        # under Pathways, this would cause weights to be loaded on a CPU-only
+        # node, so we'll need to remove this context.
+        jax_context = jax.default_device(
+            jax.devices("cpu")
+            [0]) if not vllm_envs.VLLM_TPU_USING_PATHWAYS else nullcontext()
+
         # Load the vLLM model and wrap it into a new model whose forward
         # function can calculate the hidden_state and logits.
-        with load_context, jax.default_device(jax.devices("cpu")[0]):
+        with load_context, jax_context:
             vllm_model = vllm_get_model(vllm_config=vllm_config_for_load)
         lora_manager = None
         if vllm_config_for_load.lora_config is not None:
@@ -189,7 +213,7 @@ class VllmModelWrapper:
                     kwargs={
                         "input_ids": torch_view(input_ids),
                         "positions": torch_view(input_positions),
-                        "intermediate_tensors": None,
+                        "intermediate_tensors": intermediate_tensors,
                         "inputs_embeds": None,
                     },
                     tie_weights=False,
@@ -212,8 +236,10 @@ class VllmModelWrapper:
 
         @functools.partial(
             jax.jit,
-            out_shardings=(NamedSharding(self.mesh,
-                                         PartitionSpec(None, "model"))),
+            out_shardings=(NamedSharding(
+                self.mesh,
+                PartitionSpec(ShardingAxisName.MLP_DATA,
+                              ShardingAxisName.MLP_TENSOR))),
         )
         def compute_logits_func(
             params_and_buffers: Any,
@@ -255,7 +281,6 @@ def load_lora_model(model: torch.nn.Module, vllm_config: VllmConfig,
         vllm_config,
         device,
         model.embedding_modules,
-        model.embedding_padding_modules,
     )
     return lora_manager, lora_manager.create_lora_manager(model)
 
@@ -269,10 +294,9 @@ def replace_set_lora(model):
         index: int,
         lora_a: torch.Tensor,
         lora_b: torch.Tensor,
-        embeddings_tensor: Optional[torch.Tensor],
     ):
         with torchax.default_env():
-            self._original_set_lora(index, lora_a, lora_b, embeddings_tensor)
+            self._original_set_lora(index, lora_a, lora_b)
 
     def _tpu_reset_lora(self, index: int):
         with torchax.default_env():

tpu_inference/models/vllm/vllm_model_wrapper_context.py

@@ -1,3 +1,17 @@
+# 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 contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Dict, List, Optional

tpu_inference/platforms/__init__.py

@@ -1,2 +1,16 @@
+# 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.
+
 # ruff: noqa
 from tpu_inference.platforms.tpu_platform import TpuPlatform