[Spec Decoding][Eagle3] Fix bug of eagle-3 not being compataible with non-8b models. (#1165)

Signed-off-by: Pooya Moradi <[email protected]>
Co-authored-by: Pooya Moradi <[email protected]>

Author: py4

tpu_inference/models/jax/llama_eagle3.py

@@ -304,15 +304,15 @@ def load_weights(self, rng_key: jax.Array):
             "fc": "model.fc.kernel",
             "lm_head": "lm_head.kernel",
             "d2t": "draft_id_to_target_id",
+            "embed_tokens":
+            "model.embed_tokens.embedding",  # Some checkpoints need this
         }
 
         # Define keys to keep in original dtype (e.g., float32 for stability)
         keep_original_dtype_keys_regex = [
             r".*d2t.*",
         ]
 
-        # `embed_tokens` is shared between target and draft.
-        exclude_regex = [r".*embed_tokens.*"]
         metadata_map = get_default_maps(
             self.vllm_config.speculative_config.draft_model_config, self.mesh,
             mappings)
@@ -325,10 +325,9 @@ def load_weights(self, rng_key: jax.Array):
             metadata_map=metadata_map,
             mesh=self.mesh,
             is_draft_model=True,
-            keep_original_dtype_keys_regex=keep_original_dtype_keys_regex,
-            exclude_regex=exclude_regex if exclude_regex else None)
+            keep_original_dtype_keys_regex=keep_original_dtype_keys_regex)
 
-        # If the embedding is not initialized, initialize it with a dummpy array here to pass jit compilation. The real weights will be shared from the target model in eagle3 class.
+        # If the embedding is not initialized, initialize it with a dummy array here to pass jit compilation. The real weights will be shared from the target model in eagle3 class.
         if isinstance(self.model.embed_tokens.embedding.value,
                       jax.ShapeDtypeStruct):
             self.model.embed_tokens.embedding.value = jnp.zeros(

tpu_inference/models/jax/utils/weight_utils.py

@@ -402,7 +402,6 @@ def _load_hf_weights_on_thread(
     weights_file: str,
     filter_regex: Optional[str] = None,
     keep_original_dtype_keys_regex: Optional[list[str]] = None,
-    exclude_regex: Optional[list[str]] = None,
 ):
     """Loads weights from a single weights file."""
     try:
@@ -412,17 +411,6 @@ 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):
-        # 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
         _load_and_shard_weight(
             vllm_config,
             params,
@@ -443,7 +431,6 @@ def load_hf_weights(
     filter_regex: Optional[str] = None,
     is_draft_model: bool = False,
     keep_original_dtype_keys_regex: Optional[list[str]] = None,
-    exclude_regex: Optional[list[str]] = None,
 ):
     """Load weights into a JAX model from either an iterator or files."""
     params = nnx.state(model)
@@ -491,17 +478,17 @@ def load_hf_weights(
             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,
-                                exclude_regex=exclude_regex)
-                for weights_file in weights_files
+                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()

tpu_inference/runner/compilation_manager.py

@@ -548,7 +548,9 @@ def _precompile_rejection_sampler(self) -> None:
     def _precompile_eagle3_helpers(self) -> None:
         logger.info(
             "Compiling eagle3 jitted helpers with different input shapes.")
-        hidden_size = self.runner.model_config.get_hidden_size()
+        target_hidden_size = self.runner.model_config.get_hidden_size()
+        draft_hidden_size = self.runner.speculative_config.draft_model_config.get_hidden_size(
+        )
         dtype = self.runner.model_config.dtype
 
         num_kv_cache_groups = len(self.runner.kv_cache_config.kv_cache_groups)
@@ -595,7 +597,7 @@ def _precompile_eagle3_helpers(self) -> None:
 
         for num_logits in self.runner.num_logits_paddings:
             hidden_states = self._create_dummy_tensor(
-                (num_logits, hidden_size), jnp.bfloat16)
+                (num_logits, draft_hidden_size), jnp.bfloat16)
             self._run_compilation(
                 "eagle3_get_draft_token_ids",
                 self.runner.drafter._get_draft_token_ids,
@@ -606,18 +608,21 @@ def _precompile_eagle3_helpers(self) -> None:
         input_ids_loop = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32,
             NamedSharding(self.runner.mesh, PartitionSpec()))
-        target_hidden_state_loop = self._create_dummy_tensor(
-            (self.runner.max_num_reqs, hidden_size), dtype,
+        draft_hidden_state_loop = self._create_dummy_tensor(
+            (self.runner.max_num_reqs, draft_hidden_size), dtype,
             NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
         next_token_ids = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32)
         last_token_indices = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32)
         for num_tokens in self.runner.num_tokens_paddings:
             aux_hidden_states = [
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
             ]
 
             positions = self._create_dummy_tensor((num_tokens, ), jnp.int32)
@@ -648,15 +653,15 @@ def filter_token_and_prepare_initial_inputs_wrapper(
             input_ids = self._create_dummy_tensor((num_tokens, ), jnp.int32)
             aux_hidden_states = [
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
             ]
@@ -688,17 +693,17 @@ def draft_model_fn_wrapper(
                 state,
                 kv_caches,
                 input_ids,
-                target_hidden_states,
+                draft_hidden_states,
                 attention_metadata,
             ):
                 kv_caches, hidden_states, _ = self.runner.drafter.model_fn(
-                    state, kv_caches, input_ids, target_hidden_states,
+                    state, kv_caches, input_ids, draft_hidden_states,
                     attention_metadata)
                 self.runner.kv_caches = kv_caches
                 return hidden_states
 
-            target_hidden_states = self._create_dummy_tensor(
-                (num_tokens, hidden_size), dtype,
+            draft_hidden_states = self._create_dummy_tensor(
+                (num_tokens, draft_hidden_size), dtype,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, "model")))
             input_ids = self._create_dummy_tensor(
                 (num_tokens, ), jnp.int32,
@@ -709,7 +714,7 @@ def draft_model_fn_wrapper(
                 self.runner.drafter.state,
                 self.runner.kv_caches,
                 input_ids,
-                target_hidden_states,
+                draft_hidden_states,
                 attention_metadata,
                 num_tokens=num_tokens,
             )
@@ -741,13 +746,13 @@ def draft_model_fn_wrapper(
                 self.runner.drafter.state,
                 self.runner.kv_caches,
                 input_ids_loop,
-                target_hidden_state_loop,
+                draft_hidden_state_loop,
                 attention_metadata,
                 num_tokens=num_tokens,
             )
 
             hidden_states = self._create_dummy_tensor(
-                (num_tokens, hidden_size), jnp.bfloat16,
+                (num_tokens, draft_hidden_size), jnp.bfloat16,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
 
             self._run_compilation(

tpu_inference/spec_decode/jax/eagle3.py

@@ -9,10 +9,13 @@
 from vllm.config import VllmConfig
 
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.logger import init_logger
 from tpu_inference.models.common.model_loader import get_model
 from tpu_inference.runner import utils as runner_utils
 from tpu_inference.utils import device_array
 
+logger = init_logger(__name__)
+
 
 class Eagle3Proposer:
     """A proposer for speculative decoding using the Eagle3 method.
@@ -51,9 +54,22 @@ def load_model(self, target_model: Any) -> None:
         """Loads the draft model."""
         self.model_fn, self.compute_logits_fn, self.combine_hidden_states_fn, _, self.state, _, _ = get_model(
             self.vllm_config, self.rng_key, self.mesh, is_draft_model=True)
-        if 'embed_tokens' in self.state.model:
-            del self.state.model['embed_tokens']
-        self.state.model.embed_tokens = target_model.model.embed
+
+        draft_embed_tokens = getattr(self.state.model, 'embed_tokens', None)
+        if draft_embed_tokens is None or ~jnp.any(
+                draft_embed_tokens.embedding):
+            logger.info(
+                "Draft model does not have embedding. Setting draft model's embed_tokens to target model's embed"
+            )
+            self.state.model.embed_tokens = target_model.model.embed
+        elif jnp.array_equal(draft_embed_tokens.embedding,
+                             target_model.model.embed.embedding):
+            logger.info(
+                "Draft model's embed_tokens is identical to target model's embed. Sharing the embedding."
+            )
+            self.state.model.embed_tokens = target_model.model.embed
+        else:
+            logger.info("Draft model has its own embed_tokens.")
 
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _prepare_input_ids(