gptoss120b-qlora-mathreasoning

KrackHack 3.0 submission — Domain: Gen AI | PS: Open Innovation — GPT-OSS-120B QLoRA finetuning using Unsloth for mathematical reasoning

A parameter-efficient approach to enhancing mathematical reasoning in large language models through QLoRA fine-tuning. This repository contains the complete pipeline for fine-tuning GPT-OSS-120B on Olympiad-level mathematics problems.

🎯 Key Results

• 40/50 accuracy on rigorous IMO-level evaluation set (+11% over baseline)
• Perfect consistency across evaluation runs (0.0 std dev)
• 3.4x improvement in computational tool usage
• 5.7x dataset efficiency through sequence length filtering

📋 Table of Contents

• Overview
• Installation
• Quick Start
• Training Pipeline
• Evaluation
• Results
• Project Structure
• Citation
• License

🔬 Overview

This project addresses the performance gap between closed-source and open-source models in mathematical reasoning. Using Quantized Low-Rank Adaptation (QLoRA) with Unsloth optimization, we fine-tune the 120-billion parameter GPT-OSS model on carefully curated mathematics datasets.

Key Innovation: Sequence Length Filtering

We discovered that dataset quality, measured by sequence length, significantly impacts fine-tuning effectiveness:

Dataset Size    Avg Length    Accuracy
1.7M examples   ~8000 tokens  32/50 (degraded)
1.1M examples   ~6000 tokens  32/50
300K examples   ~4000 tokens  33/50
295K examples   ~3800 tokens  40/50 ✅

🚀 Installation

Prerequisites

• Python 3.10 or higher
• CUDA 12.1+ with compatible NVIDIA GPU (recommended: A100 80GB)
• 160GB+ GPU memory (for dual-GPU setup) OR cloud environment

Environment Setup

# Clone the repository
git clone https://github.com/tomoeOOseven/gptoss120b-qlora-mathreasoning.git
cd gptoss120b-qlora-mathreasoning

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install core dependencies
pip install torch==2.5.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Install Unsloth (optimized for QLoRA training)
pip install "unsloth[cu121-torch251] @ git+https://github.com/unslothai/unsloth.git"

# Install remaining requirements
pip install -r requirements.txt

Requirements.txt

transformers==4.46.1
datasets==3.2.0
trl==0.12.1
bitsandbytes==0.44.0
accelerate==1.2.1
peft==0.13.2
scipy==1.14.1
jupyter-client==8.6.3
openai==1.57.4
polars==1.18.0
pandas==2.2.3
numpy==1.26.4
sympy==1.13.1
matplotlib==3.9.2

⚡ Quick Start

1. Download the Base Model

from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/gpt-oss-120b-unsloth-bnb-4bit",
    max_seq_length=16384,
    dtype=None,
    load_in_4bit=True,
)

2. Download Training Dataset

# Using Kaggle API
kaggle datasets download -d jeannkouagou/aimo3-tool-integrated-reasoning
unzip aimo3-tool-integrated-reasoning.zip -d data/

# Or download manually from:
# https://www.kaggle.com/datasets/jeannkouagou/aimo3-tool-integrated-reasoning

3. Run Training

# Option A: Using provided notebook
jupyter notebook notebooks/training_qlora.ipynb

# Option B: Using Python script
python scripts/train.py --config configs/qlora_config.yaml

4. Evaluate Model

# Run evaluation on 50-problem test set
python scripts/evaluate.py \
    --model_path outputs/checkpoint-30 \
    --eval_dataset data/evaluation_set.csv \
    --output_file results/evaluation_results.json

🎓 Training Pipeline

Step 1: Configure QLoRA Parameters

Edit configs/qlora_config.yaml:

# Model Configuration
model_name: "unsloth/gpt-oss-120b-unsloth-bnb-4bit"
max_seq_length: 16384
load_in_4bit: true

# LoRA Configuration
lora_r: 8
lora_alpha: 16
lora_dropout: 0.0
target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj
  - gate_proj
  - up_proj
  - down_proj

# Training Arguments
per_device_train_batch_size: 1
gradient_accumulation_steps: 4
learning_rate: 2e-4
max_steps: 30
warmup_steps: 5
optimizer: adamw_8bit
weight_decay: 0.01
lr_scheduler_type: linear

Step 2: Prepare Dataset

The training script expects CSV format with columns:

• prompt: The mathematical problem
• completion: The solution with tool-integrated reasoning

from datasets import load_dataset

dataset = load_dataset(
    "csv", 
    data_files="data/aimo3_tir.csv",
    split="train",
    streaming=True
)

Step 3: Fine-Tune Model

from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig

# Load model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/gpt-oss-120b-unsloth-bnb-4bit",
    max_seq_length=16384,
    dtype=None,
    load_in_4bit=True,
)

# Apply LoRA
model = FastLanguageModel.get_peft_model(
    model,
    r=8,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    lora_dropout=0,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=3407,
)

# Train
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    max_seq_length=16384,
    args=SFTConfig(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=5,
        max_steps=30,
        learning_rate=2e-4,
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        output_dir="outputs",
    ),
)

trainer.train()

Step 4: Save and Merge Adapters

# Save LoRA adapters
model.save_pretrained("adapters/final")
tokenizer.save_pretrained("adapters/final")

# Merge adapters with base model (optional, for deployment)
model.save_pretrained_merged(
    "models/merged_model",
    tokenizer,
    save_method="merged_16bit",
)

📊 Evaluation

Running Evaluation

from evaluation.evaluator import AIMO3Evaluator

# Initialize evaluator
evaluator = AIMO3Evaluator(
    model_path="adapters/final",
    eval_dataset="data/evaluation_50_problems.csv",
    timeout=300,  # 5 minutes per problem
    num_workers=8,
)

# Run evaluation
results = evaluator.evaluate()

print(f"Accuracy: {results['accuracy']}/50")
print(f"Avg. solving time: {results['avg_time']:.2f}s")
print(f"Tool usage rate: {results['tool_usage_rate']:.1%}")

Evaluation Metrics

The evaluation computes:

• Accuracy: Correct answers / Total problems
• Consistency: Standard deviation across multiple runs
• Tool Usage: Average Python calls per problem
• Solving Time: Time to generate solution per problem

Note: Our evaluation set is a custom 50-problem benchmark curated by Lokesh comprising problems from AIMO reference competitions and IMO shortlist.

Sandbox Environment

Problems are evaluated in isolated Jupyter kernels with:

• Scientific computing libraries: NumPy, SciPy, SymPy, Matplotlib
• 5-minute timeout per problem
• Automatic result extraction and verification

📈 Results

Quantitative Results

Key Findings

1. Dataset Quality > Quantity: 295K curated examples outperform 1.7M unfiltered examples
2. Sequence Length Matters: Shorter, focused solutions (≤4K tokens) train better than lengthy derivations
3. Tool Integration is Critical: 3.4 Python calls/problem enables systematic verification
4. Perfect Reliability: Zero variance across evaluation runs demonstrates consistent reasoning

Comparison with Baseline

Our evaluation was conducted on a custom 50-problem test set (30 from AIMO reference problems + 20 from IMO shortlist), so direct comparison with SOTA systems on different benchmarks isn't appropriate. However, we can compare against the baseline:

Note: SOTA systems like Gemini Deep Think (Gold Medal at IMO 2025), GPT-5 (65.6%), o3 (61.1%), and DeepSeek R1 (60.8%) were evaluated on IMO-AnswerBench, a different benchmark. Our focus was on improving consistency and tool-use on our specific evaluation set rather than direct SOTA comparison.

📁 Project Structure

gptoss120b-qlora-mathreasoning/
├── README.md
├── LICENSE
├── requirements.txt
├── setup.py
│
├── configs/
│   ├── qlora_config.yaml          # Training hyperparameters
│   └── eval_config.yaml            # Evaluation settings
│
├── notebooks/
│   ├── training_qlora.ipynb        # Full training pipeline
│   ├── evaluation.ipynb            # Evaluation notebook
│   └── analysis.ipynb              # Results analysis
│
├── scripts/
│   ├── train.py                    # Training script
│   ├── evaluate.py                 # Evaluation script
│   ├── merge_adapters.py           # Merge LoRA with base model
│   └── filter_dataset.py           # Sequence length filtering
│
├── src/
│   ├── __init__.py
│   ├── model/
│   │   ├── __init__.py
│   │   ├── loader.py               # Model loading utilities
│   │   └── lora_config.py          # LoRA configuration
│   │
│   ├── data/
│   │   ├── __init__.py
│   │   ├── dataset.py              # Dataset loading & preprocessing
│   │   └── filtering.py            # Sequence length filters
│   │
│   ├── training/
│   │   ├── __init__.py
│   │   ├── trainer.py              # Training loop
│   │   └── callbacks.py            # Training callbacks
│   │
│   └── evaluation/
│       ├── __init__.py
│       ├── evaluator.py            # Evaluation pipeline
│       ├── sandbox.py              # Jupyter kernel sandbox
│       └── metrics.py              # Metric computation
│
├── data/                           # (Not tracked in git)
│   ├── aimo3_tir.csv
│   ├── evaluation_50_problems.csv
│   └── filtered_datasets/
│
├── outputs/                        # (Not tracked in git)
│   ├── checkpoint-10/
│   ├── checkpoint-20/
│   └── checkpoint-30/              # Final checkpoint
│
├── adapters/                       # (Not tracked in git)
│   └── final/                      # Saved LoRA adapters
│
├── models/                         # (Not tracked in git)
│   └── merged_model/               # Merged model (optional)
│
└── results/                        # (Not tracked in git)
    ├── evaluation_results.json
    ├── training_logs.txt
    └── analysis_plots/

🔧 Advanced Usage

Custom Dataset Filtering

Filter datasets by sequence length:

from src.data.filtering import filter_by_length

filtered_dataset = filter_by_length(
    input_path="data/openmathReasoning.csv",
    output_path="data/filtered_4k.csv",
    max_length=4096,
    tokenizer=tokenizer
)

Hyperparameter Tuning

Sweep over LoRA ranks:

for rank in 4 8 16 32; do
    python scripts/train.py \
        --config configs/qlora_config.yaml \
        --lora_r $rank \
        --output_dir outputs/rank_$rank
done

Multi-GPU Training

Modify training script for DDP:

from accelerate import Accelerator

accelerator = Accelerator()
model, optimizer, training_dataloader = accelerator.prepare(
    model, optimizer, training_dataloader
)

Deployment

Serve the merged model with vLLM:

python -m vllm.entrypoints.openai.api_server \
    --model models/merged_model \
    --tensor-parallel-size 2 \
    --max-model-len 16384

🐛 Troubleshooting

Out of Memory (OOM)

# Reduce batch size
per_device_train_batch_size = 1
gradient_accumulation_steps = 8  # Increase to maintain effective batch size

# Enable gradient checkpointing
use_gradient_checkpointing = "unsloth"

# Reduce max sequence length
max_seq_length = 8192

Slow Training

# Ensure Flash Attention 2 is installed
pip install flash-attn --no-build-isolation

# Use Unsloth optimizations
from unsloth import FastLanguageModel
FastLanguageModel.for_training(model)

# Increase number of dataloader workers
dataset_num_proc = 32

Port Conflicts (Jupyter Kernels)

# Clear hanging kernels
import os
os.system("lsof -ti:39001-39016 | xargs kill -9")

📚 Additional Resources

Documentation

• Unsloth Documentation
• QLoRA Paper
• IMO-Bench Paper

Related Projects

• OpenMathReasoning Dataset
• GPT-OSS Base Model

Datasets

• AIMO3-TIR
• GSM8K

📄 License

This project is licensed under the Apache License 2.0 - see the LICENSE file for details.

📝 Citation

If you use this work in your research, please cite:

@article{teamehe2025qlora,
  title={Efficient QLoRA Fine-Tuning of Large Language Models for Mathematical Reasoning: An Open Innovation Approach},
  author={Team EHE},
  journal={KrackHack 3.0 Submission},
  year={2025}
}

👥 Team

Team EHE

• Shrestha
• Pramukto
• Saurabh
• Anurag

🙏 Acknowledgments

• Lokesh (Kaggle) for curating the evaluation problem set
• KrackHack 3.0 organizers for this competition
• Unsloth team for optimization library
• Jean Kouagou for curating the AIMO3-TIR dataset
• NVIDIA for OpenMathReasoning dataset
• Open-source community for foundational tools

Note: This is a hackathon submission for the Open Innovation domain. All code and models are released for educational and research purposes.