tpu-inference 0.0.1rc1__py3-none-any.whl → 0.11.1.dev202511130813__py3-none-any.whl

tpu_inference/models/jax/utils/quantization/quantization_utils.py

@@ -154,9 +154,12 @@ def qwix_quantize_nnx_model(model: nnx.Module, qwix_config: List[dict],
     logger.info(f"Memory usage before applying quantization of params: "
                 f"hbm={utils.hbm_usage_gb(jax.local_devices())}Gb")
 
-    if kv_cache_dtype != "auto":
-        kv_cache_jnp_dtype = utils.to_jax_dtype(kv_cache_dtype)
-    else:
+    # TODO (jacobplatin): we should refactor this to pass a dtype (or config) directly
+    kv_cache_jnp_dtype = utils.get_jax_dtype_from_str_dtype(kv_cache_dtype)
+
+    # Handle the case where kv_cache_dtype is "auto"
+    if kv_cache_jnp_dtype is None:
+        assert kv_cache_dtype == "auto", "kv_cache_dtype must be 'auto' if kv_cache_jnp_dtype is None"
         kv_cache_jnp_dtype = DEFAULT_KV_CACHE_DTYPE
 
     kv_caches = create_kv_caches(

tpu_inference/models/jax/utils/weight_utils.py

@@ -13,14 +13,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 import utils
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.utils import file_utils
 
@@ -199,11 +197,12 @@ def shard_put(x: jax.Array, shardings, mesh: jax.sharding.Mesh) -> jax.Array:
         return jax.device_put(x, shardings)
 
 
-def get_default_maps(model_config, mesh: Mesh,
+def get_default_maps(vllm_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"]
 
+    model_config = vllm_config.model_config
     hf_config = model_config.hf_config
 
     num_heads = hf_config.num_attention_heads
@@ -267,15 +266,14 @@ def get_default_maps(model_config, mesh: Mesh,
                        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_original_dtype_keys_regex: list[str]
-                           | None = None):
+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):
     name_map = metadata_map.name_map
     reshape_keys = metadata_map.reshape_map
     bias_reshape_keys = metadata_map.bias_reshape_map
@@ -292,118 +290,6 @@ def _load_and_shard_weight(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)
-
-    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 _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,
-):
-    """Loads weights from a single weights file."""
     try:
         shardings = nnx.get_named_sharding(params, mesh)
     except TypeError:
@@ -411,88 +297,147 @@ def _load_hf_weights_on_thread(
 
     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,
-        )
 
+        # 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
 
-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,
-):
-    """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):
+        # 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"
+                )
                 continue
+            model_key = name_map.get(hf_key, hf_key)
+        model_weight, model_sharding = get_param_and_sharding(
+            params, shardings, model_key)
 
-            # 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)
+        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}"
+        )
 
-            _load_and_shard_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):
+    """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 os.environ.get("TPU_MULTIHOST_BACKEND", "").lower() == "ray":
+        max_workers = 1
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [
+            executor.submit(
+                _load_hf_weights_on_thread,
                 vllm_config,
                 params,
-                shardings,
                 metadata_map,
                 mesh,
-                hf_key,
-                hf_weight_jax,
-                keep_original_dtype_keys_regex,
-            )
-    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,
-                ) for weights_file in weights_files
-            ]
-            for future in futures:
-                future.result()
-
+                weights_file,
+                filter_regex=filter_regex,
+                keep_original_dtype_keys_regex=keep_original_dtype_keys_regex)
+            for weights_file in weights_files
+        ]
+        for future in futures:
+            future.result()
     check_all_loaded(params)
     nnx.update(model, params)
 

tpu_inference/models/vllm/vllm_model_wrapper.py

@@ -9,7 +9,6 @@ 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
@@ -26,8 +25,6 @@ from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 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
 from tpu_inference.models.vllm.vllm_model_wrapper_context import (
     get_vllm_model_wrapper_context, set_vllm_model_wrapper_context)
 from tpu_inference.runner.lora_utils import replace_lora_metadata
@@ -92,14 +89,13 @@ class VllmModelWrapper:
             slice_config = self.vllm_config.device_config.slice
             modified_slice_config = True
             self.vllm_config.device_config.slice = None
-        self.vllm_config.compilation_config.static_forward_context.clear()
-
         vllm_config_for_load = copy.deepcopy(self.vllm_config)
         if modified_slice_config:
             self.vllm_config.device_config.slice = slice_config
         assert self.vllm_config.model_config.dtype in TORCH_DTYPE_TO_JAX, "The model_config.dtype must be a PyTorch dtype."
         vllm_config_for_load.device_config.device = "cpu"
         # Clearing the cached compilation config, otherwise vllm model init will fail
+        vllm_config_for_load.compilation_config.static_forward_context.clear()
 
         # When expert parallelism is enabled, vLLM loads weight in sharding
         # aware manner. Since tpu-inference has its own sharding logic, this
@@ -119,16 +115,9 @@ 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_context:
+        with load_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:
@@ -160,8 +149,7 @@ class VllmModelWrapper:
                 "xla_tpu_reduce_scatter_collective_matmul_mode":
                 "post_spmd_conservative"
             },
-            static_argnames=("layer_name_to_kvcache_index", "is_first_rank",
-                             "is_last_rank"),
+            static_argnames=("layer_name_to_kvcache_index", ),
         )
         def step_fun(
             params_and_buffers,  # This has been wrapped into torchax TorchValue
@@ -169,12 +157,8 @@ class VllmModelWrapper:
             input_ids: jax.Array,
             attn_metadata: AttentionMetadata,
             input_embeds: jax.Array,
-            input_positions: jax.Array,
             layer_name_to_kvcache_index: Sequence[Tuple[str, int]],
             lora_metadata,
-            intermediate_tensors: JaxIntermediateTensors = None,
-            is_first_rank: bool = True,
-            is_last_rank: bool = True,
             *args,
         ) -> Tuple[List[jax.Array], jax.Array]:
             layer_name_to_kvcache_index = dict(layer_name_to_kvcache_index)
@@ -189,14 +173,12 @@ class VllmModelWrapper:
                 # torch_view in order to call the Torch function.
                 original_lora_metadata = replace_lora_metadata(
                     self.model, lora_metadata, self.vllm_config.lora_config)
-                if not is_first_rank:
-                    intermediate_tensors = intermediate_tensors.to_torch()
-                output_from_torch = torch.func.functional_call(
+                hidden_states = torch.func.functional_call(
                     self.model,
                     torch_view(params_and_buffers),
                     kwargs={
                         "input_ids": torch_view(input_ids),
-                        "positions": torch_view(input_positions),
+                        "positions": torch_view(attn_metadata.input_positions),
                         "intermediate_tensors": None,
                         "inputs_embeds": None,
                     },
@@ -206,13 +188,11 @@ class VllmModelWrapper:
                                       self.vllm_config.lora_config)
                 vllm_model_wrapper_context = get_vllm_model_wrapper_context()
                 new_kv_caches = vllm_model_wrapper_context.kv_caches
-            # Wrap the output(hidden states or intermediate tensor)
-            # from torch land into a JaxValue for the jax code to consume.
-            if not is_last_rank:
-                output = JaxIntermediateTensors.from_torch(output_from_torch)
-            else:
-                output = jax_view(output_from_torch)
-            return new_kv_caches, output, []
+            # Wrap the hidden_states from torch land into a JaxValue for the jax
+            # code to consume.
+            hidden_states = jax_view(hidden_states)
+
+            return new_kv_caches, hidden_states, []
 
         return step_fun
 
@@ -221,7 +201,7 @@ class VllmModelWrapper:
         @functools.partial(
             jax.jit,
             out_shardings=(NamedSharding(self.mesh,
-                                         PartitionSpec("data", "model"))),
+                                         PartitionSpec(None, "model"))),
         )
         def compute_logits_func(
             params_and_buffers: Any,
@@ -263,6 +243,7 @@ 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)
 
@@ -276,9 +257,10 @@ 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)
+            self._original_set_lora(index, lora_a, lora_b, embeddings_tensor)
 
     def _tpu_reset_lora(self, index: int):
         with torchax.default_env():