add support torchax embedding

Author: carlesoctav

tpu_inference/models/jax/adapters.py

@@ -1,7 +1,9 @@
 import typing as tp
 
+import torch
 import jax
 from flax import nnx
+from flax.typing import PRNGKey
 from jax.sharding import Mesh
 
 from tpu_inference.layers.jax.pool.pooler import Pooler
@@ -84,3 +86,31 @@ def _init_pooler(self, vllm_config: VllmConfig) -> None:
         "ForEmbedding",
     )
     return ModelForEmbedding  # type: ignore[return-value]
+
+
+
+def init_pooler_from_vllm_model(
+        vllm_model: torch.nn.Module,
+        vllm_config: VllmConfig,
+        rng_key: PRNGKey, 
+        mesh: Mesh,
+):
+    class DummyModule:
+        def __init__(self, vllm_config, rng_key, mesh):
+            pass
+
+    for suffix in _GENERATE_SUFFIXES:
+        if suffix in vllm_model.__class__.__name__:
+            return None
+
+    if "ForEmbedding" in vllm_model.__class__.__name__:
+        EmbedModel = as_embedding_model(DummyModule)
+
+        embed_model = EmbedModel(vllm_config=vllm_config, rng_key=rng_key, mesh=mesh,)
+        embed_model._init_pooler(vllm_config)
+        return embed_model.pooler 
+    else:
+        raise NotImplementedError(
+            f"Pooling initialization for {vllm_model.__class__.__name__} is not implemented."
+        )
+

tpu_inference/models/vllm/vllm_model_wrapper.py

@@ -30,6 +30,8 @@
 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
+from tpu_inference.layers.jax.pool.pooler import Pooler
+from tpu_inference.models.jax.adapters import init_pooler_from_vllm_model
 
 logger = init_logger(__name__)
 
@@ -72,6 +74,7 @@ class VllmModelWrapper:
     rng: PRNGKey
     mesh: Mesh
     model: _VllmRunner
+    pooler: Pooler
 
     def __init__(self, vllm_config: VllmConfig, rng: PRNGKey, mesh: Mesh):
         self.vllm_config = vllm_config
@@ -137,6 +140,10 @@ def load_weights(self):
         self.model = _VllmRunner(vllm_model)
         params_and_buffers = shard_model_to_tpu(self.model, self.mesh)
 
+
+        # TODO: need to seperate this params_and_buffer for pooler (some pooler is not stateless)
+        self.pooler = init_pooler_from_vllm_model(vllm_model, self.vllm_config, self.rng, self.mesh)
+
         # Returning to the jax land, so we need to wrap it into a JaxValue.
         return jax_view(params_and_buffers), lora_manager
 

tpu_inference/runner/compilation_manager.py

@@ -21,6 +21,7 @@
 )
 from tpu_inference.logger import init_logger
 from tpu_inference.utils import device_array
+from torchax.ops.mappings import t2j_dtype
 
 if TYPE_CHECKING:
     from tpu_inference.runner.tpu_runner import TPUModelRunner
@@ -114,31 +115,17 @@ def _precompile_pooling(self) -> None:
 
         for num_tokens in self.runner.num_tokens_paddings:
             hidden_states = self._create_dummy_tensor(
-                (num_tokens, hidden_size), dtype, sharding=hidden_sharding)
+                (num_tokens, hidden_size), t2j_dtype(dtype), sharding=hidden_sharding)
 
             for num_reqs in self.runner.num_reqs_paddings:
                 if num_reqs == 0 or num_reqs > num_tokens:
                     continue
 
-                prompt_lens = np.ones(num_reqs, dtype=np.int32)
-                first_token_indices = np.arange(num_reqs, dtype=np.int32)
-                last_token_indices = first_token_indices.copy()
-                normalize = np.ones(num_reqs, dtype=np.int8)
-
-                (
-                    prompt_lens,
-                    normalize,
-                    first_token_indices,
-                    last_token_indices,
-                ) = device_array(
-                    self.runner.mesh,
-                    (
-                        prompt_lens,
-                        normalize,
-                        first_token_indices,
-                        last_token_indices,
-                    ),
-                )
+                prompt_lens = self._create_dummy_tensor(num_reqs, dtype = jnp.int32) 
+                first_token_indices = self._create_dummy_tensor(num_reqs, dtype = jnp.int32) 
+                last_token_indices = self._create_dummy_tensor(num_reqs, dtype = jnp.int32) 
+                normalize = self._create_dummy_tensor(num_reqs, dtype = jnp.int32) 
+
 
                 pooling_metadata = TPUSupportedPoolingMetadata(
                     prompt_lens=prompt_lens,