BananaMind-2-Nano
BananaMind-2-Nano is a compact decoder-only causal language model trained from scratch by BananaMind on a 30B-token curriculum.
The model has 9,968,128 parameters, a 4,096-token context window, and a custom 8k-token digit-aware byte-level BPE tokenizer.
Model Details
| Field | Value |
|---|---|
| Parameters | 9,968,128 |
| Architecture | BananaMind2Nano decoder-only Transformer |
| Layers | 10 |
| Hidden size | 256 |
| Intermediate size | 768 |
| Attention heads | 4 |
| KV heads | 2 |
| Head dim | 64 |
| Attention style | Grouped-query attention with QK norm |
| MLP | SwiGLU |
| Position embeddings | RoPE |
| RoPE theta | 100,000 |
| Normalization | RMSNorm |
| RMSNorm epsilon | 1e-06 |
| Vocabulary size | 8,192 |
| Context length | 4,096 |
| Embeddings | Tied input/output embeddings |
| Weight format | safetensors |
| HF architecture | BananaMind2NanoForCausalLM |
| HF model type | bananamind2_nano |
| Final checkpoint | runs/bananamind2-nano/final.pt |
| Final training step | 55,485 |
| Tokens seen | 29,999,726,592 |
Credits to AxiomicLabs and GPT X2 125M for the architecture inspiration.
Tokenizer
BananaMind-2-Nano 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|> | 0 |
<|bos|> | 1 |
<|eos|> | 2 |
<|unk|> | 3 |
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 | 0.003 |
| Warmup steps | 1,750 |
| LR schedule | Warmup-stable-decay with cosine decay |
| Weight decay | 0.1, then 0.01 after 12B tokens |
| Gradient clipping | 1 |
| Z-loss coefficient | 1e-4 until 12B tokens, then off |
| Compile | PyTorch compile enabled |
| Seed | 1337 |
Evaluation
These are self-reported scores produced with lm_eval. Scores may vary slightly depending on the evaluation harness version, runtime settings, dtype, and environment.
All task scores use acc_norm,none. The average is the mean of ARC Easy, PIQA, ARC Challenge, and HellaSwag.
| Benchmark | Score | Metric |
|---|---|---|
| Average | 35.77 | mean |
| ARC Easy | 36.20 | acc_norm,none |
| PIQA | 55.98 | acc_norm,none |
| ARC Challenge | 23.38 | acc_norm,none |
| HellaSwag | 27.50 | acc_norm,none |
The unrounded average is 0.357659. Available unrounded task results are ARC Easy 0.361953, PIQA 0.559848, and ARC Challenge 0.233788.
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-Nano"
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))
License
Apache 2.0