20231088/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] MLPSpeculator: Use ParallelLMHead in tie_weights=False case. (#6303)

Browse Source 
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
This commit is contained in:
Thomas Parnell 2024-07-10 15:04:07 +02:00 committed by GitHub
parent e72ae80b06
commit c38eba3046
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
vllm/model_executor/models/mlp_speculator.py
View File

@ -110,7 +110,7 @@ class MLPSpeculator(nn.Module):
]) ])
self.head = nn.ModuleList([ self.head = nn.ModuleList([
nn.Linear(self.inner_dim, self.vocab_size, bias=False) ParallelLMHead(self.vocab_size, self.inner_dim, bias=False)
for _ in range(self.max_speculative_tokens) for _ in range(self.max_speculative_tokens)
]) ])
self.ln = nn.ModuleList([ self.ln = nn.ModuleList([