nanoGPT_hw_t4

NanoGPT T4 Arena — 165M Story Generation Model

A 165M parameter LLaMA-style nanoGPT model trained for five-sentence commonsense story completion on ROCStories, developed for the COMP4680/8650 Advanced Topics in Machine Learning assignment at the Australian National University (2026 S1).

Model Details

Model Description

• Developed by: Adithya Rama (ANU Master of Computing, Advanced)
• Model type: Autoregressive transformer (LLaMA-style nanoGPT)
• Language: English
• License: MIT
• Architecture: 12 layers, 12 heads, 768 embedding dim — 123.59M parameters
• Base: Trained from scratch (no external pretrained weights)

Architecture

Modern LLaMA-style GPT with all four architectural improvements enabled:

Training — Two Stage

Stage 1 — Pretraining (25,000 steps)

• Dataset: Combined corpus — ROCStories (5× upsampled) + TinyStories (~200M tokens total)
• lr = 3e-4, cosine decay, batch = 32, block = 512
• Output: out-t4-pretrain/ckpt_best.pt

Stage 2 — Fine-tuning (30,000 additional steps, total 55,000)

• Dataset: rocstories_synthetic — 408,722 synthetic ROCStories generated by llama-3.3-70b (Groq API) and llama3:8b (local Ollama), covering 28 culturally diverse categories (South Asian, East Asian, African, Latin American, Middle Eastern, workplace narratives, etc.)
  • original 70,676 ROCStories training stories = 32.65M tokens
• lr = 1e-4, cosine decay, batch = 16, gradient accumulation = 8, effective batch = 65,536 tokens/step
• always_save_checkpoint = True (val loss rises during Stage 2 by design — this reflects distribution shift from TinyStories to ROCStories, not overfitting)
• Final checkpoint: step 55,001

Sampling Parameters

Included in sample_params.json:

{
  "temperature": 0.80,
  "top_k": 50,
  "top_p": 0.92,
  "repetition_penalty": 1.03,
  "stop_token": 50256
}

Load and use with the provided sample_batch.py.

Example Outputs

Prompt: Emily forgot her umbrella before leaving for work.

Emily forgot her umbrella before leaving for work. She walked quickly to the office building, trying not to get wet. On her way, she stopped at a store to buy a cheap umbrella. The new umbrella turned out to be very sturdy and kept her dry. Emily arrived at the office, feeling relieved and grateful for her new umbrella.

Prompt: Tom decided to cook dinner for his friends.

Tom decided to cook dinner for his friends. He spent the evening following a complicated recipe he found online. Halfway through cooking, he realized he had forgotten to buy one of the key ingredients. Tom improvised with what he had and hoped for the best flavor combination. His friends enjoyed the meal and appreciated the effort he put into cooking for them.

Prompt: Sarah lost her keys on the way to the grocery store.

Sarah lost her keys on the way to the grocery store. She had been shopping for hours and still couldn't find them anywhere. Sarah started to feel frustrated and worried about being late for her dinner plans. She decided to retrace her steps from the previous day to see if she dropped them. After a short walk, she finally found the keys on a bench and continued shopping.

Evaluation

Note: Stage 2 val PPL rising is intentional and expected. The model is shifting its distribution from TinyStories (childlike) toward adult ROCStories style. The stored best_val_loss in the checkpoint reflects Stage 1 metrics and is not meaningful for Stage 2 story quality.

How to Use

# Load checkpoint
import torch
from model import GPT, GPTConfig
import json

ckpt = torch.load('ckpt.pt', map_location='cpu', weights_only=False)
cfg = GPTConfig(**ckpt['model_args'])
model = GPT(cfg)
state_dict = {k.replace('_orig_mod.', ''): v for k, v in ckpt['model'].items()}
model.load_state_dict(state_dict)
model.eval()

# Or just use the provided scripts:
# python sample_batch.py --init_from=resume --out_dir=. \
#     --start=FILE:data/rocstories/eval_prompts.txt \
#     --batch_prompts=True --max_new_tokens=120

Training Infrastructure

• Hardware: NVIDIA A100 (Google Colab)
• Stage 1 training time: ~2.5 hours
• Stage 2 training time: ~3.5 hours
• MFU (Stage 2): ~43.7%
• Framework: PyTorch 2.x, bfloat16

Limitations

• Trained for five-sentence ROCStories-style story completion only
• Occasionally loses track of character references in longer generations
• Out-of-distribution names (non-Western) may produce less coherent output
• Not instruction-tuned — expects a story opening sentence as prompt, not a directive like "Write a story about..."

Citation

@misc{karpathy2022nanogpt,
  author = {Karpathy, Andrej},
  title  = {nanoGPT},
  year   = {2022},
  url    = {https://github.com/karpathy/nanoGPT}
}