Add support for ibm granite #758
Conversation
…into pc/granite
|
The KVcache change broke my original implementation but I found a work arround it :) Ready for reaview @awni ✅ |
|
@Blaizzy this diff includes some previous PRs. Can you rebase on origin's main and force push to your branch? |
…into pc/granite
|
Hey @awni, It's done ✅ |
| attention_bias: bool = False | ||
| mlp_bias: bool = False |
There was a problem hiding this comment.
Was there a reason you added these / changed this file?
There was a problem hiding this comment.
Granite 3B and 8B require them, and its different from the base llama models.
There was a problem hiding this comment.
Oh I see, I did not realize those models had model_type = llama. Makes sense.
llms/mlx_lm/models/gpt_bigcode.py
Outdated
| def create_additive_causal_mask(N: int, offset: int = 0): | ||
| rinds = mx.arange(offset + N) | ||
| linds = mx.arange(offset, offset + N) if offset else rinds | ||
| mask = linds[:, None] < rinds[None] | ||
| return mask * -1e9 |
There was a problem hiding this comment.
would you mind refactoring this into base.py and calling that from this model and the llama.py model? Ultimately we'd like all the models to use this same version but we can start with that and then update from there rather than duplicate the function everywhere.
llms/mlx_lm/models/gpt_bigcode.py
Outdated
| attn_pdrop: float = 0.1 | ||
| embd_pdrop: float = 0.1 | ||
| resid_pdrop: float = 0.1 |
There was a problem hiding this comment.
They are being used on the attention, embedding and mlp blocks :)
There was a problem hiding this comment.
Right let's remove those Dropout layers as well.. we don't add dropout layers to model files here since we don't use them for trianing.
|
Very nicely done! I just left a couple minor comments. Please check them and then we can merge it! |
|
Thank you very much, @awni! I addressed all comments, Let me know if there anything else |
This PR adds support for openELM.
Twitter: @Prince_Canuma
Todo: