---
license: apache-2.0
language:
- en
library_name: transformers
pipeline_tag: text-generation
datasets:
- HuggingFaceFW/fineweb-edu
- mlfoundations/dclm-baseline-1.0
- HuggingFaceTB/smollm-corpus
- HuggingFaceTB/finemath
tags:
- causal-lm
- language-model
- base-model
- small-language-model
- bananamind
- bananamind2
- bananamind2-nano
- fineweb-edu
- dclm
- cosmopedia-v2
- finemath
- digit-tokenizer
- pytorch
- safetensors
- custom-code
- trust-remote-code
- custom-architecture
---

# BananaMind-2-Nano
![Banner](banner.png)

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`.

```bash
pip install -U transformers safetensors torch
```

```python
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