Cost-Effective Development of Educational Chatbots: Training and Inference Optimization

Abstract

Large Language Models (LLMs) have revolutionized natural language processing tasks and show promise for interactive language learning applications. However, training these models and deploying them on resource-constrained devices (like smartphones or Raspberry Pi) is challenging due to their size and computational demands. This thesis presents an accelerated training pipeline for the GPT-2 124M model – the small- est GPT-2 – enabling it to be trained to competitive performance in under 3 minutes (using 0.8 billion tokens) on high-end hardware, a dramatic reduction compared to the original GPT-2 which required about 40 GB of text and multiple days of training. We achieve this speed-up through a combination of architectural optimizations (e.g. Rotary Positional Embeddings, ReLU² activations, QK-normalization, FlexAttention mechanism, and etc.), a custom Muon optimizer for faster convergence, and strate- gic hyperparameter tuning. The model was trained on the cleaned and deduplicated FineWeb dataset. The resulting model was then quantized to 8-bit, reducing memory by a factor of four with minimal loss in language modeling accuracy. We demonstrate the model’s effectiveness in an educational context, focusing on tasks like interac- tive conversation practice and grammar correction, and deployed on edge devices, including an 8GB Raspberry Pi. Benchmarks show that the optimized 124M model maintains language modelling performance (HellaSwag benchmark) while achieving real-time inference speeds within the limited memory ( 1 GB) e of these devices. This work highlights that with the right optimizations, even moderately sized LLMs can be trained and run off-line on personal devices for language learning, bridging the gap between large cloud-based models and accessible on-device AI tutors.