This Model is a first test to combine
Jamba architecture with bf16 bits linear layers, mixture of attention head and mixture of depth.
The goal is to developpe and test if this kind of architectures have not too much quality loss for a fast inference.
This model has a generation problem because of a softmax application in the mod process
If you want to test this model please look at this repo at this
commit
The data fit the model max length i.e. 512 tokens
Please look at the wandb metadata to see the hyperparameters or the train.py file in the repo