kaggle-llm-124m-chat

A 124M-parameter chat model trained from scratch on free Kaggle GPUs (2x NVIDIA T4): pretrained on 0.98B tokens of FineWeb-Edu, then supervised fine-tuned on 226K conversations from smol-smoltalk.

It is a learning project, not a production assistant: it follows the chat format and writes simple code, but it often states wrong facts confidently.

Usage

The weights are a standard LlamaForCausalLM with the GPT-2 tokenizer, so plain transformers works:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "somdevsheel/kaggle-llm-124m-chat"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype=torch.float16).to("cuda")

messages = [{"role": "user", "content": "Write a Python function that reverses a string."}]
ids = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to("cuda")
out = model.generate(ids, attention_mask=torch.ones_like(ids), max_new_tokens=200,
                     do_sample=True, temperature=0.7, top_k=50, repetition_penalty=1.15)
print(tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True))

Recommended sampling: temperature 0.7, top-k 50, repetition penalty 1.15. Context length is 1024 tokens.

Chat format (role tags are plain text; each assistant turn ends with <|endoftext|>):

<|system|>
...
<|user|>
...
<|assistant|>
...<|endoftext|>

Model

Architecture	Decoder-only Transformer, LLaMA-style: RoPE, RMSNorm (pre-norm), SwiGLU, no biases, tied embeddings
Size	123.6M parameters (85.0M non-embedding)
Shape	12 layers, hidden 768, 12 heads, FFN 2048, context 1024
Tokenizer	GPT-2 BPE (vocab padded to 50,304)

Training

Stage	Data	Compute	Result
Pretraining	0.98B tokens of FineWeb-Edu (sample-10BT)	1,865 steps x 524K tokens, 10.9 h on 2x T4	val loss 3.43 (perplexity 31)
Chat SFT	226K conversations from smol-smoltalk, packed into 1024-token blocks	2,442 steps x 65K tokens, 1.7 h on 2x T4	val loss on assistant tokens 2.79 -> 1.61

• PyTorch DDP across 2 GPUs, fp16 mixed precision with loss scaling, AdamW (0.9, 0.95), weight decay 0.1, cosine LR schedule.
• Pretraining LR 1.2e-3; the trainer measured its own speed and fitted the cosine schedule into the 12-hour Kaggle session limit.
• SFT LR 2e-4; loss computed on assistant tokens only; whole conversations first-fit packed (92% fill).

Limitations

• Hallucination. About 8 training tokens per parameter (Chinchilla-optimal is ~20), so its knowledge is thin. Example: asked "What is DSA in computer science?", it invents a definition.
• English only; no safety tuning or preference training (no RLHF/DPO).
• Short code is often correct (string reversal, prime check); longer or vaguely specified code usually is not.
• Do not rely on it for factual, medical, legal or financial answers.

Base model

Pretrained-only weights: somdevsheel/kaggle-llm-124m-base.