Comments (7)
You can use whichever training script / library you'd like, e.g. Megatron, DeepSpeed, lightning, hf accelerate etc. Just have to replace the model definition.
Examples:
Lightning has lit-gpt: https://github.com/Lightning-AI/lit-gpt
FlashAttention has training code, you can swap the model: https://github.com/Dao-AILab/flash-attention/tree/main/training
from mamba.
We've managed to train mamba by modifying the Huggingface Trainer class. Here is our implementation, we were actually able to train a chat model that seems to perform quite well.
from mamba.
does not seem to be so straightforward with HF trainer, quite literally:
File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 1555, in train
return inner_training_loop(
File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 1860, in _inner_training_loop
tr_loss_step = self.training_step(model, inputs)
File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2725, in training_step
loss = self.compute_loss(model, inputs)
File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2748, in compute_loss
outputs = model(**inputs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/accelerate/utils/operations.py", line 680, in forward
return model_forward(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/accelerate/utils/operations.py", line 668, in __call__
return convert_to_fp32(self.model_forward(*args, **kwargs))
File "/usr/local/lib/python3.10/dist-packages/torch/amp/autocast_mode.py", line 14, in decorate_autocast
return func(*args, **kwargs)
TypeError: MambaLMHeadModel.forward() got an unexpected keyword argument 'labels'
no labels
in forward()
?
would be very nice if you could provide a simple, minimal example of how to use the models with HF trainer. thank you!
from mamba.
We've managed to train mamba by modifying the Huggingface Trainer class. Here is our implementation, we were actually able to train a chat model that seems to perform quite well.
Cool, nice work! Are you using fp32 for this finetuning work?
from mamba.
Hmm... Doesn't seem to work out of the box with lit-gpt.
Minimal example:
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from mamba_ssm.utils.hf import load_config_hf
from mamba_ssm.models.config_mamba import MambaConfig
with fabric.init_module(
empty_init=isinstance(fabric.strategy, DeepSpeedStrategy)
):
config = load_config_hf('state-spaces/mamba-2.8b')
model = MambaLMHeadModel(MambaConfig(**config))
This will give the following error
Traceback (most recent call last):
File "/home/me/lit-gpt/pretrain/mamba.py", line 782, in <module>
setup(run_config)
File "/home/me/lit-gpt/pretrain/mamba.py", line 306, in setup
main(fabric, run_config)
File "/home/me/lit-gpt/pretrain/mamba.py", line 342, in main
model = MambaLMHeadModel(MambaConfig(**model_config))
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 385, in wrapper
f(module, *args, **kwargs)
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/mamba_ssm/models/mixer_seq_simple.py", line 199, in __init__
self.backbone = MixerModel(
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 385, in wrapper
f(module, *args, **kwargs)
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/mamba_ssm/models/mixer_seq_simple.py", line 118, in __init__
[
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/mamba_ssm/models/mixer_seq_simple.py", line 119, in <listcomp>
create_block(
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/mamba_ssm/models/mixer_seq_simple.py", line 42, in create_block
block = Block(
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 385, in wrapper
f(module, *args, **kwargs)
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/mamba_ssm/modules/mamba_simple.py", line 316, in __init__
self.mixer = mixer_cls(dim)
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 385, in wrapper
f(module, *args, **kwargs)
File "/home/me/lit-gpt/mamba-venv/lib/python3.9/site-packages/mamba_ssm/modules/mamba_simple.py", line 99, in __init__
self.dt_proj.bias.copy_(inv_dt)
RuntimeError: The size of tensor a (0) must match the size of tensor b (5120) at non-singleton dimension 0
from mamba.
if that's true, how the heck am I to pass attention?
from mamba.
Stage 2 works, but not stage 3. I don't have a fix at the moment. Problem is this line
mamba/mamba_ssm/modules/mamba_simple.py
Line 98 in eb2f7a5
from mamba.
Related Issues (20)
- RuntimeError: CUDA error: unknown error
- How to calculate Flops for mamba_chunk_scan_combined and mamba_split_conv1d_scan_combined function HOT 1
- Replace mamba1 with mamba2 and training becomes very slow! HOT 2
- Implement a "crossmamba" block for different length sequences HOT 3
- How to use sequence parallelism in Mamba2?
- Changing expansion_factor resulted in 'c10::DistBackendError HOT 1
- Mamba for encoder-decoder HOT 2
- mamba_split_conv1d_scan_combined error HOT 1
- torch.cuda.amp.custom_fwd(args...) Deprecated
- Is it possible to get ssm_states at specific seq len instead of just the last state? HOT 1
- How to modify mambaV2 to set causal=False? HOT 3
- Reducing _chunk_scan_bwd_kernel computation HOT 4
- Exploding gradients if ngroups is higher than 1. HOT 3
- clarification on how to interpret kernel size for conv1d HOT 1
- Some questions about the shape of A,B,C,D HOT 5
- How can I avoid using causal-conv1d? HOT 11
- Is it okay to put S6(MAMBA) and gated MLP blocks like a transformer? (Also please open the discussion tab)
- Vanishing gradient problem with more layer HOT 1
- Results vary greatly across experiments
- Gradient explosion in Mamba2 training, norm and loss divergence HOT 3
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from mamba.