This Model is a first test to combine
Jamba architecture with mixture of attention head and mixture of depth.
Mamba and attention layers are in bf16 precision and the rest is in 1.58bits precision
107M over a total of 1025M parameters are in bf16 precision ~ 10% of the parameters are in bf16
The goal is to developpe and test if this kind of architectures have not too much quality loss for a fast inference.
If you want to test this model please look at this repo at this
commitThe data fit the model max length i.e. 512 tokens
Please look at the wandb meta data or the train.py in the repo to see the hyperparameters
