Skip to content

Mamba update#254

Open
jmercat wants to merge 71 commits intomlfoundations:mainfrom
jmercat:mamba_update
Open

Mamba update#254
jmercat wants to merge 71 commits intomlfoundations:mainfrom
jmercat:mamba_update

Conversation

@jmercat
Copy link
Collaborator

@jmercat jmercat commented Apr 18, 2024

Mamba is now using a class for its input, this updates OpenLM accordingly.

@jmercat jmercat force-pushed the mamba_update branch 2 times, most recently from 10ba930 to 6277b33 Compare May 4, 2024 00:36
@arjunsesh
Copy link
Contributor

arjunsesh commented Jun 12, 2024

High level comment that these changes are up to date as of mamba-ssm 1.2.1 where I have been able to successfully run the code (however gradient_checkpointing does not work). There are now more features with mamba-ssm>=2.0.0, and submodules have been renamed, so there are some changes to be made to make it compatible.

@jmercat
Copy link
Collaborator Author

jmercat commented Jun 12, 2024

I had the same problem with gradient checkpointing and it prevented some interesting experiments about long-context mamba... I didn’t have time to look into this yet but maybe updating to Mamba2 first would make sense.

jmercat added 24 commits June 26, 2024 11:19
@jmercat
allow inputs_embeds in model input args
32c596d
@jmercat
linted
9e8abdf
@jmercat
few fixes
7aa33c3 
llm-foundry removed from requirements and added in requirements_test
mosaicml added in requirements
version bump
try except around llm-foundry import
@jmercat
update Mamba compatibility
7739486
@jmercat
MambaParams default values
d468181
@jmercat
add mamba dataclass args
b14ef50
@jmercat
small fix in hf_config
35e2488
@jmercat
set the correct kind of model in OpenLMModel
762f0a9
@jmercat
get_input_embeddings for mamba
7fb747f
@jmercat
[wip] add inputs_embeddings in mamba class
aebb9d5
@jmercat
[wip] add inputs_embeddings in mamba class
3ca3986
@jmercat
[wip] add inputs_embeddings in mamba class
c21129e
@jmercat
[wip] Mamba kwargs
f5ae4d4
@jmercat
[wip] Mamba optional inputs
f89e36c
@jmercat
[wip] Mamba inputs propagated
a8fda98
@jmercat
[wip] Mamba output fix
8a88e0d
@jmercat
Filter keys keywords in load model (remove some layers from the state…
813cf72 
… dict based on names)
@jmercat
Mamba import optional and its dependencies too
81b6aac
@jmercat
Mamba import optional and its dependencies too
8b5d641
@jmercat
check for state_dict in checkpoint when no epoch
9525c5c
@jmercat
rotary inv_freq not a buffer anymore
7c62e96
@jmercat
rotary inv_freq not a buffer anymore -> remove it from state dicts wh…
bc64b0f 
…en loading
@jmercat
loosen panda dep
d2d08b6
@jmercat
avoid masking if the attention_mask is masking only the right padding
3e0ce49
jmercat added 30 commits July 23, 2024 18:52
@jmercat
update Mamba compatibility
bef9e8d
@jmercat
MambaParams default values
a8be17c
@jmercat
add mamba dataclass args
a286021
@jmercat
small fix in hf_config
b7fdd10
@jmercat
set the correct kind of model in OpenLMModel
47105c8
@jmercat
get_input_embeddings for mamba
a940f32
@jmercat
[wip] add inputs_embeddings in mamba class
d99bcec
@jmercat
[wip] add inputs_embeddings in mamba class
fb4613d
@jmercat
[wip] add inputs_embeddings in mamba class
046134d
@jmercat
[wip] Mamba kwargs
7b9ad58
@jmercat
[wip] Mamba optional inputs
a3f86a5
@jmercat
[wip] Mamba inputs propagated
e59c73e
@jmercat
[wip] Mamba output fix
75b983f
@jmercat
Filter keys keywords in load model (remove some layers from the state…
9ceeaa2 
… dict based on names)
@jmercat
Mamba import optional and its dependencies too
9c41b76
@jmercat
Mamba import optional and its dependencies too
1f46f1f
@jmercat
check for state_dict in checkpoint when no epoch
ed27e51
@jmercat
rotary inv_freq not a buffer anymore
5fd16fd
@jmercat
rotary inv_freq not a buffer anymore -> remove it from state dicts wh…
eba0181 
…en loading
@jmercat
loosen panda dep
e79644c
@jmercat
avoid masking if the attention_mask is masking only the right padding
f1226d9
@jmercat
handle none mask in hf_model
4f6d45f
@jmercat
Filter keys keywords in load model (remove some layers from the state…
80199f7 
… dict based on names)
@jmercat
Fix hf_model loss
412896b
@jmercat
Prevent wrong config to fail
1ce3442
@jmercat
fix mask shift
b401e75
@jmercat
black formatting and rebase on main
a1d6394
@jmercat
Merge remote-tracking branch 'origin/mamba_update' into mamba_update
dd97111
@jmercat
change model_norm to norm_type
1d6efab
@jmercat
more flexible OpenLM model config input
97cd5d4
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@jmercat @arjunsesh @sedrick-keh-tri @Vaishaal @ysharma1126