ofidaner

aimo3-gpt-oss-120b-sft-v08a

Key Idea: SFT as Behavioral Regularization

GPT-OSS-120B already solves most competition problems. Our SFT doesn't teach new solution methods — it regularizes the model's reasoning behavior, pushing hard problems from ~2/8 to ~4/8 correct attempts while reducing run-to-run variance (SD 1.0 vs 2.0 for the base model).

Files

Model

Table with columns: File, Description
File	Description
`adapter_model.safetensors`	LoRA adapter weights (91 MB)
`adapter_config.json`	LoRA configuration (rank 32, alpha 32, attention-only)
`merge_lora.py`	Sparse merge script — merges adapter into base model
`aimo_solver.py`	Full inference harness (TIR loop, voting, time management)
`RUNPOD_HOWTO.md`	Step-by-step RunPod setup and reproduction guide
`colab_reproduce.ipynb`	Google Colab notebook — one-click reproduction

Data

Table with columns: File, Description
File	Description
`data/tir_traces_v08a_balanced.jsonl`	Training data — 1,788 balanced TIR traces
`data/hard_devtest_50.csv`	PolyMath-Val-50 validation set (50 problems)
`data/hard_devtest_50_answers.csv`	Validation set answer key
`data/polymath_train_1000.csv`	Problem pool for data generation (1,000 problems)
`data/polymath_train_200.csv`	Initial problem pool (200 problems)

Scripts (full pipeline reproduction)

Table with columns: File, Description
File	Description
`scripts/run_hardtest_runpod.py`	RunPod inference — evaluate on PolyMath-Val-50
`scripts/datagen_cloud.py`	Data generation — run base model on problems
`scripts/build_training_from_traces.py`	Build balanced training data from raw traces
`scripts/train_unsloth.py`	Unsloth QLoRA training script
`scripts/patch_attention.py`	SDPA attention patch for Unsloth training

Analysis (figure generation)

Table with columns: File, Description
File	Description
`analysis/parse_hardtest.py`	Reusable parser for validation output files
`analysis/fig01_variance_comparison.py`	Figure 1: Competition score variance
`analysis/fig02_data_balance.py`	Figure 4: Balanced vs unbalanced sampling
`analysis/fig03_problem_stability_heatmap.py`	Figure 5a/5b: Problem stability heatmaps
`analysis/fig04_05_ablation_charts.py`	Figures 6, 8: Data composition and inference ablations

How to Use

Google Colab (recommended for reviewers)

Requirements: Colab Pro/Pro+ with A100 80GB runtime. The model requires 80GB VRAM — A100 40GB or smaller GPUs will not work.

The Colab notebook handles the full setup automatically — all downloads come from HuggingFace, no Kaggle credentials needed:

Installs Python 3.12 (required for vLLM wheels)
Downloads base model (ofidaner/gpt-oss-120b-mxfp4, ~65GB)
Downloads LoRA adapter, vLLM wheels, and tiktoken encodings from HuggingFace
Merges adapter into base model (~14 seconds)
Runs evaluation on the PolyMath-Val-50 validation set (50 problems, ~3.5 hours)

Manual setup (RunPod / any cloud GPU)

See RUNPOD_HOWTO.md for detailed step-by-step instructions.

Quick start (any machine with 80GB+ VRAM)

Prerequisites:

GPU with 80GB+ VRAM (H100 or A100 80GB)
Python 3.12
Base model: hf download ofidaner/gpt-oss-120b-mxfp4 (~65GB)
vLLM 0.16.0 (gpt-oss fork): hf download ofidaner/aimo3-vllm-wheels --repo-type dataset
Tiktoken encodings: hf download ofidaner/aimo3-tiktoken --repo-type dataset

Step 1: Merge adapter into base model

bash
python merge_lora.py \
  --base /path/to/gpt-oss-120b-kaggle \
  --adapter . \
  --output /path/to/merged_output

This produces a sparse merged shard (~1.78 GB) containing only the 144 modified attention tensors. Takes ~14 seconds on CPU.

Step 2: Reconstruct full model directory

bash
mkdir -p /path/to/merged_model
for f in /path/to/gpt-oss-120b-kaggle/*; do
  ln -s "$f" /path/to/merged_model/$(basename "$f")
done
for f in /path/to/merged_output/*; do
  ln -sf "$f" /path/to/merged_model/$(basename "$f")
done

Step 3: Run inference

bash
# The solver handles vLLM startup, Jupyter kernel pool, and the full TIR loop
python aimo_solver.py --model /path/to/merged_model --problems problems.csv

Training Details

Table with columns: Parameter, Value
Parameter	Value
Base model	GPT-OSS-120B (openai/gpt-oss-120b)
Training platform	Tinker (Thinking Machines fine-tuning API)
Training context	32K tokens
LoRA rank	32
LoRA alpha	32 (scaling = 1.0)
Target modules	q_proj, k_proj, v_proj, o_proj (attention only)
Epochs	1
Effective batch size	8
Learning rate

Data Generation

Training traces were generated by running the base GPT-OSS-120B model on 1,200 problems from the PolyMath dataset (8 attempts per problem, temperature 1.0). Correct traces were collected and balanced by difficulty:

Table with columns: Difficulty (correct/8), Problems, Sampling, Share
Difficulty (correct/8)	Problems	Sampling	Share
8/8 (easy)	~350	1 trace each	~27%
7/8	~150	2 traces each	~14%
5-6/8 (medium)	~280	all traces	~33%
1-4/8 (hard)	~260	all traces x2	~26%

Inference Configuration

vLLM 0.16.0 (gpt-oss fork) with openai_harmony encoding
8 parallel attempts per problem at temperature 1.0
Entropy-weighted majority voting using logprobs
16 persistent Jupyter kernels with numpy/sympy/mpmath
Multi-turn TIR loop: up to 128 turns per attempt
Time budget: 300s reserved per remaining problem, max 900s per problem

Citation

If you use this work, please cite:

markdown
@misc{fidaner2026aimo3sft,
  title={SFT as Behavioral Regularization: Fine-Tuning a 120B MoE for Mathematical Reasoning},
  author={Fidaner, Onur},
  year={2026},
  note={AIMO Progress Prize 3 Competition}
}

Acknowledgments

PolyMath dataset by Dekoninck, Hiss, Petrov, Rognvaldsson, and Sun
GPT-OSS-120B with Tools notebook by Andreas Bisiadis
Thinking Machines Tinker fine-tuning platform
XTX Markets, AIMOPrize, and Kaggle for organizing AIMO 3

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.

Learn more

Model Details

Model Provider

ofidaner

Model Tree

Base

openai/gpt-oss-120b

Adapter

this model

Input Modalities

Text

Output Modalities

Text

Supported Functionality

Dedicated EndpointsContainer

Explore FriendliAI today

Get started Talk to an engineer

Key Idea: SFT as Behavioral Regularization

Files

Model

Table with columns: File, Description
File	Description
`adapter_model.safetensors`	LoRA adapter weights (91 MB)
`adapter_config.json`	LoRA configuration (rank 32, alpha 32, attention-only)
`merge_lora.py`	Sparse merge script — merges adapter into base model
`aimo_solver.py`	Full inference harness (TIR loop, voting, time management)
`RUNPOD_HOWTO.md`	Step-by-step RunPod setup and reproduction guide
`colab_reproduce.ipynb`	Google Colab notebook — one-click reproduction

Data

Table with columns: File, Description
File	Description
`data/tir_traces_v08a_balanced.jsonl`	Training data — 1,788 balanced TIR traces
`data/hard_devtest_50.csv`	PolyMath-Val-50 validation set (50 problems)
`data/hard_devtest_50_answers.csv`	Validation set answer key
`data/polymath_train_1000.csv`	Problem pool for data generation (1,000 problems)
`data/polymath_train_200.csv`	Initial problem pool (200 problems)

Scripts (full pipeline reproduction)

Table with columns: File, Description
File	Description
`scripts/run_hardtest_runpod.py`	RunPod inference — evaluate on PolyMath-Val-50
`scripts/datagen_cloud.py`	Data generation — run base model on problems
`scripts/build_training_from_traces.py`	Build balanced training data from raw traces
`scripts/train_unsloth.py`	Unsloth QLoRA training script
`scripts/patch_attention.py`	SDPA attention patch for Unsloth training

Analysis (figure generation)

Table with columns: File, Description
File	Description
`analysis/parse_hardtest.py`	Reusable parser for validation output files
`analysis/fig01_variance_comparison.py`	Figure 1: Competition score variance
`analysis/fig02_data_balance.py`	Figure 4: Balanced vs unbalanced sampling
`analysis/fig03_problem_stability_heatmap.py`	Figure 5a/5b: Problem stability heatmaps
`analysis/fig04_05_ablation_charts.py`	Figures 6, 8: Data composition and inference ablations

How to Use

Google Colab (recommended for reviewers)

Requirements: Colab Pro/Pro+ with A100 80GB runtime. The model requires 80GB VRAM — A100 40GB or smaller GPUs will not work.

The Colab notebook handles the full setup automatically — all downloads come from HuggingFace, no Kaggle credentials needed:

Installs Python 3.12 (required for vLLM wheels)
Downloads base model (ofidaner/gpt-oss-120b-mxfp4, ~65GB)
Downloads LoRA adapter, vLLM wheels, and tiktoken encodings from HuggingFace
Merges adapter into base model (~14 seconds)
Runs evaluation on the PolyMath-Val-50 validation set (50 problems, ~3.5 hours)

Manual setup (RunPod / any cloud GPU)

See RUNPOD_HOWTO.md for detailed step-by-step instructions.

Quick start (any machine with 80GB+ VRAM)

Prerequisites:

GPU with 80GB+ VRAM (H100 or A100 80GB)
Python 3.12
Base model: hf download ofidaner/gpt-oss-120b-mxfp4 (~65GB)
vLLM 0.16.0 (gpt-oss fork): hf download ofidaner/aimo3-vllm-wheels --repo-type dataset
Tiktoken encodings: hf download ofidaner/aimo3-tiktoken --repo-type dataset

Step 1: Merge adapter into base model

bash
python merge_lora.py \
  --base /path/to/gpt-oss-120b-kaggle \
  --adapter . \
  --output /path/to/merged_output

This produces a sparse merged shard (~1.78 GB) containing only the 144 modified attention tensors. Takes ~14 seconds on CPU.

Step 2: Reconstruct full model directory

bash
mkdir -p /path/to/merged_model
for f in /path/to/gpt-oss-120b-kaggle/*; do
  ln -s "$f" /path/to/merged_model/$(basename "$f")
done
for f in /path/to/merged_output/*; do
  ln -sf "$f" /path/to/merged_model/$(basename "$f")
done

Step 3: Run inference

bash
# The solver handles vLLM startup, Jupyter kernel pool, and the full TIR loop
python aimo_solver.py --model /path/to/merged_model --problems problems.csv

Training Details

Table with columns: Parameter, Value
Parameter	Value
Base model	GPT-OSS-120B (openai/gpt-oss-120b)
Training platform	Tinker (Thinking Machines fine-tuning API)
Training context	32K tokens
LoRA rank	32
LoRA alpha	32 (scaling = 1.0)
Target modules	q_proj, k_proj, v_proj, o_proj (attention only)
Epochs	1
Effective batch size	8
Learning rate

Data Generation

Table with columns: Difficulty (correct/8), Problems, Sampling, Share
Difficulty (correct/8)	Problems	Sampling	Share
8/8 (easy)	~350	1 trace each	~27%
7/8	~150	2 traces each	~14%
5-6/8 (medium)	~280	all traces	~33%
1-4/8 (hard)	~260	all traces x2	~26%

Inference Configuration

vLLM 0.16.0 (gpt-oss fork) with openai_harmony encoding
8 parallel attempts per problem at temperature 1.0
Entropy-weighted majority voting using logprobs
16 persistent Jupyter kernels with numpy/sympy/mpmath
Multi-turn TIR loop: up to 128 turns per attempt
Time budget: 300s reserved per remaining problem, max 900s per problem

Citation

If you use this work, please cite:

markdown
@misc{fidaner2026aimo3sft,
  title={SFT as Behavioral Regularization: Fine-Tuning a 120B MoE for Mathematical Reasoning},
  author={Fidaner, Onur},
  year={2026},
  note={AIMO Progress Prize 3 Competition}
}

Acknowledgments

PolyMath dataset by Dekoninck, Hiss, Petrov, Rognvaldsson, and Sun
GPT-OSS-120B with Tools notebook by Andreas Bisiadis
Thinking Machines Tinker fine-tuning platform
XTX Markets, AIMOPrize, and Kaggle for organizing AIMO 3

aimo3-gpt-oss-120b-sft-v08a

README

Key Idea: SFT as Behavioral Regularization

Files

Model

Data

Scripts (full pipeline reproduction)

Analysis (figure generation)

How to Use

Google Colab (recommended for reviewers)

Manual setup (RunPod / any cloud GPU)

Quick start (any machine with 80GB+ VRAM)

Training Details

Data Generation

Inference Configuration

Citation

Acknowledgments

Explore FriendliAI today

README

Key Idea: SFT as Behavioral Regularization

Files

Model

Data

Scripts (full pipeline reproduction)

Analysis (figure generation)

How to Use

Google Colab (recommended for reviewers)

Manual setup (RunPod / any cloud GPU)

Quick start (any machine with 80GB+ VRAM)

Training Details

Data Generation

Inference Configuration

Citation

Acknowledgments