Add Speculator Architecture #6
Conversation
fms_extras/modules/speculator.py
Outdated
| import torch.nn.functional as F | ||
| from fms.modules.layernorm import LayerNormParameterized | ||
|
|
||
| class Speculator(nn.Module): |
There was a problem hiding this comment.
I wonder if we might end up with multiple speculators, and if there's something more specific we might want to name this
There was a problem hiding this comment.
might be helpful to have an explanation of our speculation strategy in the comments too
There was a problem hiding this comment.
Good point, I'll add an explanation. We could call this MLP_Speculator
There was a problem hiding this comment.
FYI snake_case class names aren't convention in python: https://visualgit.readthedocs.io/en/latest/pages/naming_convention.html
we might want to rename the dataset classes at some point too
fms_extras/modules/speculator.py
Outdated
There was a problem hiding this comment.
should the path be models instead of modules? at least in the main fms repo, we distinguish between the two. I think the speculator runs standalone, not as a child of another model
There was a problem hiding this comment.
Haha I was using the exact opposite logic - I had this in modules a la main repo, because it's a standalone object and not built from any dedicated sub-modules. But I can move this into models - at least for now that would make the organization of this repo slightly simpler
fms_extras/models/speculator.py
Outdated
|
|
||
| class MLP_Speculator(nn.Module): | ||
| """ | ||
| This is a simple MLP-based speculator that functions similarly to Medusa, ingesting context via |
There was a problem hiding this comment.
might want to link to the medusa paper
fms_extras/models/speculator.py
Outdated
| The architecture is as flat and simple as possible: for each prediction head, the current | ||
| state vector is projected into a new latent space and added to the previous token's embedding. | ||
| This sum goes through layernorm and activation, forming the new state vector. This state predicts | ||
| the next token (or set of candidate tokens) for the current head, and then is passed on to the next. |
There was a problem hiding this comment.
nit: in diffs on github, the text will be easier to read if each lines is shorter (short enough to not need to wrap in a side-by-side window).
long enough that whole paragraphs are wrapped works too though then they become harder to comment on
fms_extras/models/speculator.py
Outdated
| import torch.nn.functional as F | ||
| from fms.modules.layernorm import LayerNormParameterized | ||
|
|
||
| class MLP_Speculator(nn.Module): |
There was a problem hiding this comment.
nit: snake_case for class names is unconventional in python
fms_extras/models/speculator.py
Outdated
| ) | ||
| # Weights ensure that state_0 accounts for 50% of state magnitude by final head in expectation | ||
| self.state_weight = .5**(.5/n_predict) | ||
| self.emb_weight = math.sqrt(1-self.state_weight**2) |
There was a problem hiding this comment.
will need to run black, expects spaces are -
|
I made the Blacking and Snake_Casing changes but it looks like a maintainer needs to approve the automated tests again |
fms_extras/models/speculator.py
Outdated
| m.weight.data.fill_(1) | ||
| m.bias.data.zero_() | ||
|
|
||
| def generate_suffixes(self, state, ind, topk=[5, 4, 3], n=5): |
There was a problem hiding this comment.
can we add type hints and docstrings for this
fms_extras/models/speculator.py
Outdated
| 1, best_guesses.unsqueeze(2).expand(-1, -1, self.n_predict) | ||
| ) # b n h | ||
|
|
||
| def forward(self, state, inds): |
There was a problem hiding this comment.
can we add type hints and docstrings for this
Add support for speculative decoding architecture, with implementations of distinct parallel pretraining and generative inference forward passes.