srikarkashyap/telugu-gpt-moe-sangraha-10k

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.