Curriculum SLM (PyTorch)

Complete, modular PyTorch project for training a Small Language Model (decoder-only Transformer) with curriculum learning and evaluating downstream classification on AG News.

Structure

• model/ GPT-style model + classification head
• tokenizer/ SentencePiece tokenizer training/loading
• data/ language modeling and downstream loaders
• training/ curriculum pretraining + downstream fine-tuning + baseline
• evaluation/ metrics plotting
• config/ experiment configuration
• main.py end-to-end orchestration

Key Features

• GPT-style decoder-only Transformer:
  • 8 layers, 512 hidden, 8 heads, 2048 FFN
  • learned positional embeddings
  • causal mask attention
• Tokenizer trained across all curriculum stages.
• Curriculum pretraining:
  • Stage 1: from scratch on simple corpus
  • Stage 2: 90% stage2 + 10% replay from stage1
  • Stage 3: 80% stage3 + 20% replay from stage2
• Downstream AG News:
  • fine-tune pretrained model with classification head
  • baseline model trained only on AG News (no pretraining)
• Metrics:
  • train/val LM loss
  • perplexity
  • downstream accuracy
  • TensorBoard logging
• Reproducibility:
  • global seed control
• Outputs:
  • checkpoints
  • JSON metrics
  • comparison plots

Installation

python -m venv .venv
source .venv/bin/activate
pip install torch datasets sentencepiece pyyaml matplotlib tensorboard

Data preparation

Provide stage corpora as plain text files (one sample per line):

• data/raw/stage1.txt
• data/raw/stage2.txt
• data/raw/stage3.txt (optional domain corpus; can be empty)

Update paths in config/default.yaml as needed.

Run

python main.py --config config/default.yaml

TensorBoard:

tensorboard --logdir outputs/tensorboard

Outputs

In outputs/:

• pretrained_final.pt
• stage checkpoints stage*_step_*.pt
• pretraining_metrics.json
• curriculum_downstream.json
• baseline_downstream.json
• pretraining_loss.png
• pretraining_perplexity.png
• downstream_comparison.png