[Model] Changes to MLPSpeculator to support tie_weights and input_scale

vllm/model_executor/models/mlp_speculator.py

@@ -32,18 +32,27 @@ def __init__(
         self,
         normalized_shape,
         eps=1e-06,
+        elementwise_scale=True,
+        elementwise_shift=False,
     ):
         super(MLPSpeculatorLayerNorm, self).__init__()
-        self.weight = nn.Parameter(torch.empty(normalized_shape))
-        self.bias = nn.Parameter(torch.empty(normalized_shape))
+        self.elementwise_scale = elementwise_scale
+        self.elementwise_shift = elementwise_shift
+        if self.elementwise_scale:
+            self.weight = nn.Parameter(torch.empty(normalized_shape))
+        if self.elementwise_shift:
+            self.bias = nn.Parameter(torch.empty(normalized_shape))
         self.eps = eps
 
+
     def forward(self, x):
         xf = x
         xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
         x = xf.type_as(x)
-        x = self.weight * x
-        x = x + self.bias
+        if self.elementwise_scale:
+            x = self.weight * x
+        if self.elementwise_shift:
+            x = x + self.bias
         return x
 
 
@@ -59,32 +68,63 @@ def __init__(self, config: MLPSpeculatorConfig, **kwargs) -> None:
 
         self.max_speculative_tokens = config.num_lookahead_tokens
 
-        self.emb = nn.ModuleList([
-            VocabParallelEmbedding(config.vocab_size,
-                                   self.inner_dim,
-                                   org_num_embeddings=config.vocab_size)
-            for _ in range(self.max_speculative_tokens)
-        ])
-
-        self.proj = nn.ModuleList([
-            nn.Linear((self.emb_dim if i == 0 else self.inner_dim),
-                      self.inner_dim,
-                      bias=False) for i in range(self.max_speculative_tokens)
-        ])
-
-        self.head = nn.ModuleList([
-            nn.Linear(self.inner_dim, self.vocab_size, bias=False)
-            for _ in range(self.max_speculative_tokens)
-        ])
-        self.ln = nn.ModuleList([
-            MLPSpeculatorLayerNorm(self.inner_dim)
-            for _ in range(self.max_speculative_tokens)
-        ])
+        self.tie_weights = config.tie_weights
+        self.scale_input = config.scale_input
+
+        if self.tie_weights:
+            assert (self.n_predict > 1), "You cannot tie weights between stages when only 1 exists"
+            embedding = VocabParallelEmbedding(config.vocab_size, self.inner_dim, org_num_embeddings=config.vocab_size)
+            self.emb = nn.ModuleList([
+                embedding
+                for _ in range(self.max_speculative_tokens)
+            ])
+
+            # the initial projection from the base model may have a different size, so that stays separate.
+            proj_first = nn.Linear(self.emb_dim, self.inner_dim, bias=False)
+            proj_tied = nn.Linear(self.inner_dim, self.inner_dim, bias=False)
+            self.proj = nn.ModuleList([proj_first] + [proj_tied for _ in range(self.max_speculative_tokens - 1)])
+
+            head = nn.Linear(self.inner_dim, self.vocab_size, bias=False)
+            self.head = nn.ModuleList([
+                head
+                for _ in range(self.max_speculative_tokens)
+            ])
+
+            ln = MLPSpeculatorLayerNorm(self.inner_dim, elementwise_shift=True, elementwise_scale=True)
+            self.ln = nn.ModuleList([
+                ln
+                for _ in range(self.max_speculative_tokens)
+            ])
+        else:
+            self.emb = nn.ModuleList([
+                VocabParallelEmbedding(config.vocab_size,
+                                       self.inner_dim,
+                                       org_num_embeddings=config.vocab_size)
+                for _ in range(self.max_speculative_tokens)
+            ])
+
+            self.proj = nn.ModuleList([
+                nn.Linear((self.emb_dim if i == 0 else self.inner_dim),
+                          self.inner_dim,
+                          bias=False) for i in range(self.max_speculative_tokens)
+            ])
+
+            self.head = nn.ModuleList([
+                nn.Linear(self.inner_dim, self.vocab_size, bias=False)
+                for _ in range(self.max_speculative_tokens)
+            ])
+            self.ln = nn.ModuleList([
+                MLPSpeculatorLayerNorm(self.inner_dim, elementwise_shift=True, elementwise_scale=True)
+                for _ in range(self.max_speculative_tokens)
+            ])
+        if self.scale_input:
+            self.ln0 = MLPSpeculatorLayerNorm(self.emb_dim, elementwise_shift=False, elementwise_scale=False)
 
         self.state_weight = 0.5**(0.5 / config.n_predict)
         self.emb_weight = math.sqrt(
             (1 - self.state_weight**2) * (self.inner_dim / 2))
         self.activation = nn.GELU()
+
         self.config = config
         self.logits_processor = LogitsProcessor(config.vocab_size,
                                                 config.vocab_size, 1.0)
@@ -105,6 +145,9 @@ def generate_proposals(
         # b x 1 x d
         previous_hidden_states = previous_hidden_states.unsqueeze(1)
 
+        if self.scale_input:
+            previous_hidden_states = self.ln0(previous_hidden_states) / (2**0.5)
+
         # b x 1
         last_tokens = input_ids.unsqueeze(1)
 
@@ -137,8 +180,10 @@ def generate_proposals(
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         params_dict = dict(self.named_parameters())
         for name, loaded_weight in weights:
-            param = params_dict.get(name.replace("speculator.", ""))
-            if param is not None:
+            try:
+                param = params_dict[name.replace("speculator.", "")]
                 weight_loader = getattr(param, "weight_loader",
                                         default_weight_loader)
                 weight_loader(param, loaded_weight)
+            except:
+                pass

vllm/transformers_utils/configs/mlp_speculator.py

@@ -17,6 +17,8 @@ def __init__(self,
                  n_predict: int = 3,
                  top_k_tokens_per_head: Optional[List[int]] = None,
                  n_candidates: int = 5,
+                 tie_weights: bool = False,
+                 scale_input: bool = False,
                  **kwargs):
         """
         Initialize an MLPSpeculatorConfig
@@ -38,6 +40,11 @@ def __init__(self,
                 NOTE: This parameter is currently unused.
             n_candidates: int
                 number of child candidates to create per sequence
+            tie_weights: bool
+                If true, use a single set of weights for every model head/stage after the first. The initial projection
+                from the base model may have a different size, so that stays separate.
+            scale_input: bool
+                if True, will scale the initial hidden states from the base model
         """
         if top_k_tokens_per_head is None:
             top_k_tokens_per_head = [5, 4, 3]
@@ -49,5 +56,7 @@ def __init__(self,
         self.top_k_tokens_per_head = top_k_tokens_per_head
         self.n_candidates = n_candidates
         self.num_lookahead_tokens = n_predict
+        self.tie_weights = tie_weights
+        self.scale_input = scale_input
 
         super().__init__(**kwargs)