---
language:
- ne
license: mit
tags:
- nepali
- causal-lm
- from-scratch
- llm
- cuda
- fyp
datasets:
- tonibirat/neBrahma-Nepali-Pretrain-Corpus
---

# neBrahma-base-58m: Nepali Causal Language Model

A 58M-parameter Nepali decoder-only language model trained **from scratch** on 1.16B tokens
of Nepali text. This is the P3 base model for the neBrahma FYP project - a unified
Nepali text+audio system targeting ASR and TTS.

The training engine is **hand-written in C++/CUDA** (llm.c style), targeting an RTX 3060
Mobile (6 GB VRAM). No HuggingFace Transformers used in training.

## Quick Stats

| | |
|---|---|
| Parameters | 57.7M |
| Vocabulary | 32,768 (SentencePiece BPE) |
| Training context | 256 tokens |
| Training tokens | 1.157B |
| Training steps | 70,624 |
| Final val loss | 3.8413 |
| Perplexity | 46.6 |
| Random baseline | 10.40 (ln 32768) |
| Hardware | RTX 3060 Mobile 6GB |
| Training time | ~4 days |

## Architecture

| | |
|---|---|
| Type | Decoder-only transformer |
| d_model | 512 |
| n_layers | 12 |
| n_heads | 8 (head_dim = 64) |
| FFN dim | 1,536 |
| Attention | Causal, no biases |
| Positional encoding | RoPE (theta=10000) |
| FFN activation | SwiGLU |
| Normalization | RMSNorm (eps=1e-5), pre-norm |
| Embeddings | Tied (tok_emb = LM head) |
| Dtype | FP32 (BF16 Tensor Cores via WMMA during training) |

## Training Details

- **Data**: 1.845B-token corpus, Chinchilla-optimal budget of 1.16B tokens
- **Tokenizer**: 32k SentencePiece BPE trained on Nepali corpus
- **Optimizer**: AdamW (beta1=0.9, beta2=0.95, weight_decay=0.1)
- **LR schedule**: Linear warmup (2000 steps) + cosine decay: 3e-4 -> 3e-5
- **Batch**: micro=8, grad_accum=8, seq_len=256 -> 16,384 tokens/step
- **Steps**: 70,625 (Chinchilla-optimal for 58M params / 1.16B token budget)
- **Throughput**: ~3,700 tok/s (WMMA TF32 Tensor Core matmuls)
- **Engine**: Hand-written C++/CUDA (custom training loop, no PyTorch in training path)

## Training Curve

![Training Curve](images/loss_curve.png)

![LR and Gradient Norm](images/lr_gnorm.png)

### Validation Loss by Step

| Step | Val Loss |
|------|----------|
| 0 | 10.5076 |
| 1,000 | 6.1994 |
| 2,000 | 5.2937 |
| 5,000 | 4.5601 |
| 10,000 | 4.2322 |
| 20,000 | 4.1278 |
| 30,000 | 3.9827 |
| 40,000 | 3.9185 |
| 50,000 | 3.8962 |
| 60,000 | 3.8378 |
| 70,000 | 3.8174 |
| 70,624 | 3.8413 |

## Dataset

Trained on [tonibirat/neBrahma-Nepali-Pretrain-Corpus](https://huggingface.co/datasets/tonibirat/neBrahma-Nepali-Pretrain-Corpus):

| | |
|---|---|
| Documents | 20,321,968 |
| Tokens | 1.845B |
| Sources | IndicCorp v2, FineWeb-2, IRIIS, Sagarmatha ASR |
| Quality | FIT TO TRAIN certified (zero warnings) |
| Keep rate | 77.7% (5.8M docs rejected by quality gates) |

## Usage

This model uses a custom architecture. Load with `engine/oracle/model.py` from the
[neBrahma-llm repo](https://github.com/ToniBirat7/neBrahma-llm):

```python
from safetensors import safe_open
import torch

# Load weights
tensors = {}
with safe_open("model.safetensors", framework="pt") as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)

# Build model (requires engine/oracle/model.py from neBrahma-llm repo)
from model import TinyLlama
cfg = {
    "d_model": 512, "n_layers": 12, "n_heads": 8, "head_dim": 64,
    "d_ff": 1536, "vocab_size": 32768, "rms_eps": 1e-5, "rope_theta": 10000.0,
}
model = TinyLlama(cfg)

# Load weights (ordered_params() gives canonical order matching safetensors keys)
param_keys = (
    ["model.tok_emb"] +
    [f"model.layers.{l}.{n}" for l in range(12)
     for n in ["attn_norm","wq","wk","wv","wo","ffn_norm","w_gate","w_up","w_down"]] +
    ["model.final_norm"]
)
for key, p in zip(param_keys, model.ordered_params()):
    p.data.copy_(tensors[key])

model.eval()

# Tokenize with sentencepiece
import sentencepiece as spm
sp = spm.SentencePieceProcessor()
sp.Load("tokenizer.model")

text = "नेपाली भाषा"
tokens = torch.tensor([sp.encode(text)], dtype=torch.long)
with torch.no_grad():
    logits = model(tokens)
```

## Limitations

- **Base model only**: decoder-only, not instruction-tuned or RLHF-aligned
- **Short context**: trained at seq_len=256; extrapolation beyond this is unreliable
- **Nepali only**: vocabulary and training data are Nepali-only
- **Custom architecture**: not compatible with `AutoModelForCausalLM`; requires the neBrahma engine
- **FYP artifact**: this is a research checkpoint, not a production model
- **Next phase**: P5 will fine-tune this on ASR/TTS data with a 45,056-token audio-extended vocabulary

## Citation

```bibtex
@misc{nebrahma2026,
  author    = {Birat Gautam},
  title     = {neBrahma: A Nepali Multimodal Language Model Trained from Scratch},
  year      = {2026},
  note      = {Final Year Project, Birmingham City University},
  url       = {https://github.com/ToniBirat7/neBrahma-llm},
}
```