MathLM

MathLM is a lightweight, from-scratch implementation of a Mathematical Language Model using the GPT-2 architecture. It demonstrates the complete pipeline of initializing a transformer model with random weights and performing Supervised Fine-Tuning (SFT) for domain-specific mathematical tasks.

---

🚀 Project Highlights

• Small Language Model (SLM): Optimized for local training with a custom GPT-2 config ($n_layer=4, n_head=4, n_embd=128$).
• Targeted Training: Specialized for arithmetic patterns (e.g., 2+3=5) using trl and assistant_only_loss.
• Hardware Optimized: Native support for Apple Silicon (MPS) and NVIDIA CUDA.

📂 Project Structure

File/Folder	Description
data/	Training datasets in .jsonl format
models/	Trained checkpoints (e.g., checkpoint-3300)
train.py	Core training script using SFTTrainer
test.py	Inference & evaluation script

🛠 Tech Stack & Config

• LLM Frameworks: transformers, trl, datasets
• Architecture: GPT-2 (Vocab: 50,257 | Context: 1024 | Layers: 4 | Heads: 4)
• Acceleration: MPS (Metal Performance Shaders) / CUDA

---

⚙️ Getting Started

1. Installation

git clone [https://github.com/guoxinyue1112/MathLM.git](https://github.com/guoxinyue1112/MathLM.git)
cd MathLM
pip install torch transformers trl datasets

2. Dataset Preparation

Place your training data in the data/ folder in .jsonl format. Each line should be a JSON object that the trl.SFTTrainer can parse.

3. Training

1. Place your .jsonl data in data/.

2. Configure train.py (default: 1000 epochs, batch size 4).

3. Run:

python train.py

4. Inference & Testing

After training, use test.py to evaluate the model's mathematical reasoning. Update the model_path variable in test.py to point to your latest checkpoint (e.g., checkpoint-3300).

python test.py

🧠 Logic & Design

The core philosophy of MathLM is to treat mathematical reasoning as a structured sequence-to-sequence task, specifically optimized for small-scale parameter efficiency.

1. Training Strategy & Optimization

To achieve high-precision results on a localized dataset, the training pipeline implements the following:

• Cold-Start Initialization: Unlike traditional fine-tuning, this project starts with a randomized GPT2Config. This ensures the model learns mathematical syntax and logic strictly from the ground up without pre-existing linguistic bias.
• Assistant-Only Loss (SFT): By utilizing the assistant_only_loss=True flag in the SFTTrainer, the model ignores the prompt tokens during backpropagation.
  • Mechanism: Gradients are only calculated based on the "completion" tokens. This focuses the model's capacity on generating the correct answer rather than reconstructing the question.
• High-Density Convergence: Given the small hidden dimension ($128$), we employ a high epoch count ($1000$) to allow the model to fully map the discrete numerical relationships within the training set.

2. Inference & Generation Pipeline

The inference engine in test.py is architected for cross-platform efficiency:

• Execution State: Invokes model.eval() and torch.no_grad() to freeze weights and disable dropout, ensuring stable outputs and reducing peak memory consumption.
• Hardware-Aware Backend:
  • MPS (Metal Performance Shaders): Directly utilizes the GPU cores of Apple Silicon for accelerated tensor operations.
  • Automatic Fallback: Implements a robust device-agnostic bridge: MPS → CUDA → CPU.
• Decoding Strategy:
  • Temperature Sampling ($T=0.8$): Introduces controlled stochasticity to prevent the model from getting stuck in repetitive loops.
  • Context Control: Uses max_length constraints to ensure the decoder terminates appropriately after generating the mathematical solution.

📊 Hardware Support

This project is specifically tested on:

• Apple Silicon (M1 Pro): Utilizing dataloader_pin_memory=False for stability.

• CUDA: Fully compatible with NVIDIA GPUs.

---

📄 License

This project is licensed under the MIT License.

Copyright (c) 2025 Guo Xinyue

---

✉️ Contact

Guo Xinyue GitHub: @guoxinyue1112