BananaMind-2-MoE
BananaMind-2-MoE is a sparse decoder-only causal language model trained from scratch by BananaMind on a 30B-token curriculum.
The model has 25,086,592 total parameters, approximately 1,985,152 active parameters per token, a 4,096-token context window, and a custom 8k-token digit-aware byte-level BPE tokenizer.
Model Details
| Field | Value |
|---|---|
| Total parameters | 25,086,592 |
| Active parameters per token | 1,985,152 |
| Architecture | BananaMind2MoE decoder-only Transformer |
| Routing | Sparse top-1, no capacity drops |
| Experts per layer | 48 |
| Active experts per token | 1 |
| Layers | 8 |
| Hidden size | 128 |
| Expert intermediate size | 160 |
| Attention heads | 4 |
| KV heads | 2 |
| Head dim | 32 |
| Attention style | Grouped-query attention with QK norm |
| Expert MLP | SwiGLU |
| Position embeddings | RoPE |
| RoPE theta | 100,000 |
| Normalization | RMSNorm |
| Vocabulary size | 8,192 |
| Context length | 4,096 |
| Embeddings | Tied input/output embeddings |
| Weight format | safetensors |
| HF architecture | BananaMind2MoEForCausalLM |
| HF model type | bananamind2_moe |
| Final checkpoint | runs/bananamind2-moe/final.pt |
| Final training step | 55,485 |
| Tokens seen | 29,999,726,592 |
Active parameters count the tied embedding/output matrix, every attention and router parameter, normalization parameters, and one selected expert per layer.
Tokenizer
BananaMind-2-MoE uses the same custom 8k byte-level BPE tokenizer as BananaMind-2-Mini. Digits are kept as separate tokens so numbers do not collapse into large number tokens.
| Special token | ID |
|---|---|
| `< | pad |
| `< | bos |
| `< | eos |
| `< | unk |
Training Data
| Dataset | Target Tokens | Share |
|---|---|---|
| FineWeb-Edu | 16.5B | 55% |
| DCLM | 9.0B | 30% |
| Cosmopedia-v2 | 3.0B | 10% |
| FineMath-4+ | 1.5B | 5% |
| Total | 30.0B | 100% |
The run used a progressive curriculum, beginning web-heavy and gradually increasing synthetic textbook and mathematics data.
Training Setup
| Field | Value |
|---|---|
| Sequence length | 4,096 |
| Micro batch | 12 |
| Gradient accumulation | 11 |
| Effective batch | 132 sequences |
| Tokens per optimizer step | 540,672 |
| Final optimizer step | 55,485 |
| Optimizer | AdamW |
| Betas | 0.9, 0.95 |
| Peak learning rate | 2.3e-3 |
| Warmup steps | 1,750 |
| LR schedule | Warmup-stable-decay with cosine decay |
| Weight decay | 0.1, then 0.01 after 12B tokens |
| Router balance loss | 0.01 |
| Router z-loss | 1e-3 |
| Gradient clipping | 1.0 |
| Compile | PyTorch compile enabled |
| Seed | 1337 |
Evaluation
These are self-evaluated scores produced with lm_eval. Scores may vary slightly depending on the evaluation harness version, runtime settings, dtype, and environment.
An independent Open SLM Leaderboard evaluation is coming soon.
| Benchmark | Score | Metric |
|---|---|---|
| Average | 34.90 | mean |
| ARC Easy | 34.64 | acc_norm,none |
| PIQA | 56.37 | acc_norm,none |
| ARC Challenge | 21.16 | acc_norm,none |
| HellaSwag | 27.45 | acc_norm,none |
The unrounded average is 0.349047. Unrounded task results are ARC Easy 0.346380, PIQA 0.563656, ARC Challenge 0.211604, and HellaSwag 0.274547.
Usage
This model uses custom architecture code, so load it with trust_remote_code=True.
pip install -U transformers safetensors torch
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "BananaMind/BananaMind-2-MoE"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = (
torch.bfloat16
if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
else torch.float32
)
model = AutoModelForCausalLM.from_pretrained(
model_id,
trust_remote_code=True,
torch_dtype=dtype,
).to(device).eval()
prompt = "The color of the sky is"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
output = model.generate(
**inputs,
max_new_tokens=96,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.1,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
)
print(tokenizer.decode(output[0], skip_special_tokens=True))
Notes
- This is a base model, not a chat-tuned model.
- Routing is top-1 and deterministic in evaluation mode.
- Autoregressive generation supports the Transformers
DynamicCacheKV cache. - Current benchmark scores are self-evaluated; independent leaderboard evaluation is pending.
License
Apache 2.0