Telugu GPT-MoE Sangraha 10k
This repository contains a 10,000-step base language-model checkpoint trained from scratch on Telugu text from the verified subset of ai4bharat/sangraha.
It is a research/personal pretraining checkpoint, not an instruction-tuned assistant. Expect Telugu text continuation behavior, not reliable chat or question answering.
Model Summary
- Architecture: decoder-only causal LM
- Scale: GPT-2-small-like hidden width/depth
- Layers: 12
- Hidden size: 768
- Attention heads: 12
- KV heads: 4, grouped-query attention
- Context length: 1024 tokens
- Normalization: PreNorm RMSNorm
- Positional encoding: RoPE
- FFN: SwiGLU MoE
- Experts: 4 experts, top-1 routing
- Attention backend: PyTorch SDPA / FlashAttention when available
- Vocabulary size: 68,096
- Unique parameters: 297,720,576 (~298M)
- Raw tied state-dict parameters: 350,018,304
- Weight format:
model.safetensors - Tokenizer: pre-existing Sarvam tokenizer files saved with the checkpoint
This is a custom PyTorch architecture. It loads through Hugging Face Auto classes with trust_remote_code=True.
Training Data
- Dataset:
ai4bharat/sangraha - Subset/config:
verified - Telugu split used locally:
tel - Packed block size: 1024 tokens
- Training tokens seen: 327,680,000
The Sangraha dataset card on Hugging Face lists the dataset license as cc-by-4.0.
Training Setup
- GPU: NVIDIA RTX PRO 2000 Blackwell laptop GPU
- Mixed precision: FP8 via Accelerate/TorchAO where available
- Optimizer: AdamW
- Schedule: warmup + cosine decay
- Batch shape: 2 sequences per microbatch, gradient accumulation 16
- Tokens per optimizer step: 32,768
- Total optimizer steps: 10,000
- Peak observed throughput: roughly 3k tokens/s after warmup
Training Result
Final train step:
step=10000
loss=2.9328
lm_loss=2.9224
router_aux_loss=1.0224
router_dropped_fraction=0.0151
router_entropy=1.0828
tokens_seen=327,680,000
Final validation result on the prepared local validation split:
eval_loss=2.9257
eval_lm_loss=2.9153
eval_perplexity=18.65
Eval perplexity trend:
500 84.64
1000 47.87
1500 36.22
2000 31.95
2500 29.52
3000 27.38
3500 25.95
4000 24.82
4500 23.57
5000 22.79
5500 22.06
6000 21.36
6500 20.81
7000 20.32
7500 19.87
8000 19.51
8500 19.24
9000 19.00
9500 18.75
10000 18.65
Global-MMLU Telugu
Evaluated on CohereForAI/Global-MMLU, config te, split test.
Four-pass full-option scoring with mean token log probability:
total=14,042
correct=3,400
accuracy=24.21%
This is essentially random-level for a 4-way benchmark. The model has learned some Telugu surface-form continuation, but it is not yet a useful knowledge or reasoning model.
Usage With Transformers Auto Classes
Because this is a custom architecture, pass trust_remote_code=True:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo_id = "srikarkashyap/telugu-gpt-moe-sangraha-10k"
tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
repo_id,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
device_map="cuda",
)
prompt = "తెలుగు భాష"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=96,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.15,
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
The model does not implement KV-cache yet, so generation works but is slower than stock GPT-style models with cached decoding.
Usage With Included Sampler
Clone this repository and run the included sampler from the repo root:
uv run python sample_checkpoint.py `
--checkpoint . `
--device cuda `
--dtype bfloat16 `
--prompt "తెలుగు భాష" `
--max_new_tokens 96 `
--temperature 0.7 `
--top_p 0.9 `
--repetition_penalty 1.15
CPU sampling also works, but is slower:
uv run python sample_checkpoint.py --checkpoint . --device cpu --prompt "తెలుగు భాష"
Limitations
- Not instruction tuned.
- Not aligned for chat.
- Not reliable for factual QA.
- Shows scraped-web/news-style continuation artifacts.
- Global-MMLU Telugu is random-level.
- Custom architecture requires the included code.
Intended Use
This checkpoint is best used as:
- a Telugu base-LM experiment,
- a continuation point for more pretraining,
- a starting point for Telugu SFT experiments,
- a reproducibility artifact for a small local-GPU training run.
It should not be used for production factual answering, medical/legal/financial advice, or safety-critical applications.
Citation / Data
Training data came from ai4bharat/sangraha. Evaluation used CohereForAI/Global-MMLU.