{ "id": 693113, "title": "Four-Phase Adaptation of a 120B Model for Tool-Augmented AIMO3 Reasoning", "content": "

\n

Abstract

\n

We present a four-phase fine-tuning pipeline that adapts the frozen gpt-oss-120b base model for competition mathematic, AIMO3. We apply sequential Supervised Fine-Tuning (SFT),Group Relative Policy Optimisation (DrGRPO), and Odds-Ratio PreferenceOptimisation (ORPO), followed by MXFP4 quantisation and merging for Kaggle deployment.\nEach phase addresses a specific failure mode identified through leaderboard results and internal hard-50 evaluation. The final MXFP4-baked checkpoint at 700 ORPO steps achieved a best public leaderboard score of 42 across 85 total submissions.

\n

1. Introduction

\n

The AI Mathematical Olympiad — Progress Prize 3 (AIMO3) requires systems to solve 50 competition-style integer-answer mathematics problems per submission. During the public leaderboard phase, 50 of these 100 problems are scored and contribute to the public rank; the remaining 50 are held out and scored only in the final private evaluation run. This split means the public LB score is noisy — it reflects performance on a fixed 50-problem subset, not the full competition distribution. Correct inference demands multi-step symbolic reasoning, often interleaved with verified Python computation via a tool-use interface (Harmony). A model that can reason symbolically but fails to call tools correctly wastes the computation window; a model that calls tools but cannot reason about hard problems produces wrong answers regardless.\nOur approach decomposes the adaptation problem into four sequentially targeted phases, each designed to fix a specific failure mode of the previous phase:

\n
    \n
  1. SFT warm-start — forces the model to follow the Harmony tool-use dialogue format on problems it cannot yet solve, before any reward signal is introduced.
    2. \n
  2. DrGRPO — reinforcement learning on sparse, verified oracle traces teaches the model to reason correctly through hard competition problems.
    3. \n
  3. ORPO — preference alignment on 1,415 competition problem pairs prevents the model from drifting to plausible but incorrect answers on problems it can already solve.
    4. \n
  4. MXFP4 baking — quantises and merges the full adapter stack into a ~61 GB checkpoint that fits on a single H100 80 GB GPU within the Kaggle runtime constraint.\nBefore any training, we conducted an initial exploration phase of five leaderboard submissions to estimate sensitivity to prompt format, inference temperature, and generation count. This established that the public LB varied by ±2–3 points under different hyperparameter settings, providing the\nsignal needed to calibrate subsequent phases.\nAll code is publicly available at\ngithub.com/codesrepo/aimo3-training-pipeline.\nThe sections below document each step with sufficient detail for full reproduction.
    5. \n
\n
\n

Figure 1 — Full training and inference pipeline (10 stages).\n \">\n

2. Method

\n

2.1 Dataset and Baseline Evaluation (Step 1)

\n

Dataset. We use the tir split of nvidia/OpenMathReasoning, filtered to examples whose\n expected answer can be parsed as an integer using Python's int(str(answer).strip()). The filter yields 249,341 problems from the tir split. This split contains tool-interleaved reasoning traces generated by NVIDIA's Nemotron models, making it the most suitable public source for Harmony-format fine-tuning.\n OpenMathReasoning subset Baseline: We evaluated the unmodified gpt-oss-120b base model under greedy decoding (temperature T=0.0, seed=42) on a 662-randomly selected problem subset using Harmony tool-augmented inference. The model answered 351 problems correctly (53.0%)

\n

2.2 Difficulty-Stratified Data Selection (Step 2)

\n

Fine-tuning on problems the model already solves wastes capacity and may induce regression on hard problems. We therefore stratify the scored subset by difficulty before selecting training data.\n Embedding and clustering. All 662 scored problems were embedded using the\n all-MiniLM-L6-v2 sentence transformer (384-dimensional embeddings). K-Means clustering with k=15 was applied to partition the embedding space into regions of similar semantic difficulty. The choice of k=15 was informed by both empirical validation and the elbow criterion — we ran multiple cluster configurations to identify a value that yielded a sizable subset with a high error rate, and k=15 struck the right balance. For completeness, the elbow criterion applied to Within-Cluster Sum of Squares (WCSS) over k ∈ {5, 10, …, 100} is presented in Figure 2.\n Figure 2 — K-Means elbow plot used to select k=15.\n \"\"\n Figure 3 — 2D PCA projection of the 662 scored problems coloured by cluster mismatch rate.\n Cluster 3 (68.12% mismatch, top right) is the selected high-difficulty region.\n \"\"\n Mismatch threshold. For each cluster we computed the mismatch rate — the fraction of problems the base model answered incorrectly. A threshold of ≥0.591 identified highest-difficulty clusters. A single cluster — Cluster 3, mismatch rate 68.12% — exceeded this threshold. The threshold was chosen to select a compact, maximally hard training set while retaining enough problems for meaningful fine-tuning.\n Figure 4 — Distribution of per-cluster mismatch rates (k=15). The red dashed line marks\n the 0.591 threshold; only Cluster 3 (right tail) exceeds it.\n \"\"\n The fitted K-Means model was then applied to the full 249,341-example dataset to assign every problem a cluster label and mismatch rate. From the full dataset, 2,710 hard examples belonging to Cluster 3 were selected as the SFT and DrGRPO training pool.

\n

2.3 SFT Warm-start (Step 2b)

\n

Motivation. Without format conditioning, the base model frequently violates the Harmony multi-turn tool-use protocol during early DrGRPO rollouts — misformatting tool calls, omitting protocol headers, or failing to emit \\boxed{answer}. These format errors dominate the DrGRPO reward signal in early steps and prevent the policy from learning reasoning strategies. The SFT warm-start fixes the format before any reward signal is introduced.\n Data preparation. Each of the 2,710 high-mismatch examples includes an oracle solution from the OpenMathReasoning dataset. We parsed each oracle solution to extract (i) Python tool call blocks, (ii) tool output blocks, and (iii) the final assistant reasoning turn, then re-encoded the full sequence into the Harmony multi-turn message format:

\n
\n
system prompt → user problem → [assistant turn → tool call → tool output]* → \\boxed{answer}\n
\n

Training. We trained a LoRA adapter on gpt-oss-120b (4-bit BnB NF4 quantisation) using the 2,710 Harmony-formatted examples. LoRA targeted the query, key, value, output, gate, up, and down projections (rank=4, α=8).

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
HyperparameterValue
Base modelgpt-oss-120b (BnB NF4)
LoRA rank / α4 / 8
Target modulesq, k, v, o, gate, up, down
Training examples2,710
HardwareRTX PRO 6000 Blackwell (96 GB VRAM), local WSL2
\n

2.4 Oracle Trace Collection (Step 3)

\n

DrGRPO is an on-policy RL method: the model samples its own rollouts during training and a reward function scores them group-relatively. The 207 verified oracle traces served not as supervised positive examples but as a confirmed-answer corpus — each trace's integer answer was used as the ground truth for the exact-match reward component, ensuring the reward signal is non-degenerate on hard problems. We collected these 207 samples iteratively by running gpt-oss-120b via a local vLLM server on every problem where the base model had previously failed, allowing up to 8 tool-augmented attempts per problem. Only traces where the predicted answer matched the expected answer were retained.\nEach retained trace was annotated with the cluster ID and mismatch rate of its source problem (using the saved K-Means model), enabling downstream cluster-aware reward scaling.\nTraces accumulated across training rounds:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
RoundCumulative correct traces
1151
2173
3207 ← used for DrGRPO
4219
5243
6283
The round-3 snapshot of 207 traces was used for DrGRPO training, as it represented the best available signal at the time of the reinforcement learning run.
An optional post-processing pass (aimo_gen_cert.py) extracted structured fields from each trace (key_idea, proof_skeleton, sanity_checks, solved_attempt), producing compact certificates for rag stage but it was not used in the final submission.
\n

2.5 DrGRPO Reinforcement Learning (Step 4)

\n

Motivation. After SFT the model reliably follows the Harmony format, but still fails on hard reasoning. DrGRPO provides a direct reward signal: sample candidate solutions from the current policy, compare each to the verified oracle answer, and reinforce correct reasoning chains via group-relative advantage normalisation.\nWhy sparse training data. Only ~30 oracle traces per iteration were used. This is a deliberate choice: the DrGRPO reward signal on verified hard problems is already strong, and overfitting to a small set forces the policy to generalise. Larger batches would require either more oracle data (expensive to collect) or accepting lower-quality traces. Each iteration runs on a single sample (batch=1).\nTwo-phase reward. Training was split into two phases to avoid reward collapse from competing\nobjectives:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
PhaseStepsDominant reward components
A — Certificate-first10\\boxed{} format (0.5), protocol headers (0.3), exact answer (5.0), tool consistency (0.7), length penalty (−0.05)
B — Answer-first20Exact answer (5.0, dominant), tool ok bonus (0.7), tool consistency (0.7), length penalty (−0.05)
Phase A reinforces structural correctness before Phase B shifts emphasis to answer accuracy.
\n
\n

Figure 5 — DrGRPO two-phase reward structure.\n \"\"

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
HyperparameterValue
BaseSFT-warmed gpt-oss-120b (BnB NF4)
LoRA rank / α4 / 8
LR5 × 10⁻⁶
KL β0.001
Generations per problem (G)2
Batch size1 (single sample per step)
Total steps30 (Phase A: 10, Phase B: 20)
scale_rewardsFalse
HardwareRTX PRO 6000 Blackwell (96 GB VRAM), local WSL2
\n

2.6 Community Dataset Aggregation (Step 5)

\n

ORPO requires diverse preference pairs covering the specific error types the model makes when it can solve a problem but selects the wrong answer from competing candidates. We assembled a community dataset from 11 public competition mathematics sources, filtered to integer answers in [0, 99999] and deduplicated on problem text.\n CrystalMath-Preview (Yi-Chia Chen, ycchen,\n ycchen/Crystal-Math-Preview)\n contributed 2,129 problems in raw form. After filtering to integer answers with non-null solutions, answers in [0, 99999], deduplication, and removal of problems with Chinese characters, 634 problems survived into the pool. The hard-50 benchmark (50 problems) used for checkpoint evaluation is from ZFTurbo\n (zfturbo,\n hard-math-problems-for-aimo-3).\n The full 2,664-problem pool was scored with gpt-oss-120b (MAX_ATTEMPTS=3 per problem) via a parallel vLLM server, recording both correct and wrong attempt text for every problem.

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
SourceProblems
ODA-Math-460k659
CrystalMath-Preview (filtered from 2,129)634
DAPO-17K561
olympiads-ref-base258
Nemotron-Math-V2228
PaCoRe-Train-8k133
IMO-AnswerBench67
Omni-MATH50
OlymMATH25
BeyondAIME22
MathArena18
Putnam-Axiom9
Total2,664
Figure 6 — Community dataset composition (2,664 problems from 11 sources).
\"\"
\n

2.7 Preference Pair Construction (Step 6)

\n

From the scored community dataset we constructed ORPO training pairs. Early versions of the dataset used full Harmony multi-turn traces as chosen/rejected sequences (382 pairs, Feb 2026). After experimentation, we found that answer-only pairs produced more stable ORPO training\n and better LB scores. The final format is:

\n

\n \n \n \n \n

2.8 ORPO Preference Fine-tuning (Step 7)

\n \n \n \n \n \n

Motivation. The purpose of DrGRPO was to improve the model’s reasoning behavior under a reward signal. However, given our limited compute budget, the DrGRPO stage was necessarily small: only 30 steps were run, so we treated the resulting checkpoint as a directionally improved but under-optimized policy rather than a fully converged RL model. We then applied ORPO as a lightweight preference-calibration stage on top of this checkpoint.\n This ordering was deliberate. The ORPO dataset contained 1,415 answer-only preference pairs from 11 competition sources, where the preferred response was the correct final answer and the rejected response was a plausible incorrect answer. These pairs were designed to address cases where the correct answer already appeared in the candidate pool but was not ranked highest. In other words, ORPO was used to shift the model’s final-answer ranking toward correct candidates, rather than to teach new reasoning procedures.\n To reduce the risk of overwriting reasoning behavior learned during DrGRPO, we used answer-only pairs rather than full solution traces. This format exposes the model to a broad range of error types, such as off-by-one mistakes in modular arithmetic, incorrect counting, and wrong base cases, without prescribing alternative reasoning chains. We do not claim this eliminates forgetting, since ORPO still updates the policy, but it reduces the most direct mechanism by which offline preference optimization could overwrite DrGRPO-induced reasoning patterns.\n ORPO, or Odds-Ratio Preference Optimization, combines a supervised loss on the preferred response with a log-odds preference loss, and does not require a separate reference model. This made it computationally practical for our setting: a 120B-parameter model trained under 4-bit quantization on a single-GPU budget.\n Incremental training. Training did not run from scratch to 1,000 steps in a single session. The i500 checkpoint was saved first; training then continued for 200 additional steps to produce i700, with checkpoints saved every 100 steps. This incremental approach allowed public leaderboard evaluation between checkpoints and avoided committing to a full training run before seeing LB signal.

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
HyperparameterValue
LR5 × 10⁻⁶
β (ORPO)0.001
Max prompt length2,048 tokens
Max total length4,224 tokens
Batch size1 (single sample per step)
Total steps1,000
LoRA rank / α / dropout4 / 8 / 0.05
OptimiserPaged AdamW 8-bit
HardwareRTX PRO 6000 Blackwell (96 GB VRAM), local WSL2
\n

2.9 Checkpoint Selection via Hard-50 Elimination (Step 7, continued)

\n

We evaluated five model ORPO variants on the 50-problem hard-50 benchmark (ZFTurbo), each run 2–3 times to account for stochastic generation. Checkpoint selection followed a three-stage elimination process:

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
ModelHard-50 meanRuntime meanDecision
i50031.5194.5 min① Eliminated — public LB score lower than i700
v23032026 (1.2k)25.3186.7 min② Eliminated — slowest remaining runtime
v25032026 (1.4k)28.0182.0 min② Eliminated — slower runtime
i1k-gold28.5172.5 min② Eliminated — public LB lower than i700
i70027.0175.0 minSelected — best public LB + acceptable runtime
Elimination rationale (two criteria).
\n
    \n
  1. i500 eliminated first (criterion 1: public LB score). Despite achieving the highest hard-50 mean (31.5 / 50), i500's public LB score on the 50 evaluated problems was lower than i700's. The public LB — scored on a fixed 50-problem subset — was the primary selection signal. i500's hard-50 runtime (194.5 min for 50 problems) also left the least headroom against the 5-hour Kaggle GPU notebook limit.
    2. \n
  2. Remaining candidates eliminated (criterion 2: runtime). Among the remaining checkpoints, v23032026 (186.7 min) and v25032026 (182.0 min) were eliminated on runtime; i1k-gold (172.5 min) was eliminated on public LB score, which was lower than i700's despite its faster runtime. i700 showed the best combination of public LB performance and acceptable runtime.
    3. \n
  3. i700 selected. At a mean hard-50 runtime of 175 min for 50 problems (~3.6 min/problem), the projected runtime for the full 100-problem notebook (~357 min on local hardware); actual Kaggle H100 submissions completed within the 5-hour GPU limit. i700 achieved the best public LB across 45 final submissions (mean 38.3, max 42).\n Figure 7 — Checkpoint evaluation: hard-50 scores and runtimes by model variant.\n \"\"\n Left panel: mean scores with individual run dots; right panel: mean runtimes.
    4. \n
\n

2.10 MXFP4 Baking (Step 8)

\n

The final ORPO LoRA adapter (i700) was merged into the base model weights using Unsloth's save_pretrained_merged with save_method=\"mxfp4\" on a node with 8× H100 80 GB GPUs (Kaggle notebook).

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
FormatSize
Naïve BF16 merge~218 GB
MXFP4 baked merge~61 GB
The ~61 GB checkpoint fits within the single H100 80 GB memory budget with headroom for the vLLM KV cache and Harmony tool interface.
\n

2.11 Kaggle Inference (Step 9–10)

\n

The inference notebook is based on v8/23 of @andreasbis's public notebook (scriptVersionId=290379281). Differences are limited to: model path (MXFP4 i700 checkpoint), notebook_limit increased from 17,400 to 17,520, and a use_rag flag set to 0 in the final submission.\n The merged MXFP4 checkpoint was uploaded as a Kaggle model and served via vLLM inside the submission notebook. Inference used Harmony-encoded prompts with tool-augmented self-consistency across multiple generation attempts. Each problem is allocated a timeout equal to the largest of: 300 s (minimum per problem), 900 s (hard-problem cap), or the remaining time budget after reserving 300 s for every unsolved problem — ensuring accumulated surplus is fully utilised.
\n For answer selection, majority voting is applied across up to 8 attempts with early stopping once 4 solutions agree. Otherwise, entropy weighing is performed to determine the rank — the generation the model was most confident about. Both the dynamic timeout strategy and majority voting with entropy weighing are widely used in public AIMO notebooks.

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
SettingValue
Serving frameworkvLLM
Hardware1× H100 80 GB
Tensor-parallel size1
Prompt encodingHarmony multi-turn
RAG at inferenceNone (despite notebook name aimo_rag_baked_c283)
The notebook name refers to the 283-certificate training set used for DrGRPO, not to retrieval augmentation at inference time. RAG was dropped in the final configuration after submissions with RAG (mean 37.2–37.9) consistently underperformed plain inference with the baked i700 model (mean 38.3).
\n

3. Results

\n

3.1 Leaderboard Progression

\n

We made 85 total submissions across seven distinct configuration phases (6 failed attempts excluded from analysis). The table and chart below show the progression from baseline exploration to the final baked i700 configuration.

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
PhaseValid submissionsMean LBMax LBNotes
Initial exploration536.640Hyperparameter sensitivity probing
LoRA adapter (not baked)534.437Slower inference hurt utilisation
RAG only (base model)437.541Retrieved solutions at inference
LoRA adapter + RAG536.438DrGRPO adapter loaded separately
Baked LoRA + RAG1037.940Seed DrGRPO + DPO iterations merged
DPO baked + RAG537.240
DPO baked, no RAG (final)4538.342MXFP4 i700, best configuration
Key findings from the progression:
\n \n

3.2 Ablation: Checkpoint Selection

\n

The hard-50 benchmark (Figure 7) revealed a counterintuitive result: i500 scored highest on hard-50 (mean 31.5/50) but worst on the public leaderboard, suggesting it may have overfit to the local benchmark distribution. Since the private LB score is the mean of 2 runs, we used public LB performance rather than hard-50 as the primary signal for checkpoint selection. At decision time the evidence was:

\n
\n\n

i700 was the clear choice on that basis. We also noticed ambiguities in several hard-50 problems that reinforced this decision — for example, N01418 implicitly assumes non-negative integers (unstated), and N01455's answer changes depending on whether (a,b) ranges over positive or all integers, exactly the kind of under-specification that IMO problems avoid. This made us more confident that the public LB distribution is closer to the private set than hard-50 is.

\n

4. Related Work

\n

Tool-augmented reasoning. Our inference pipeline builds on the Harmony tool-use framework, which interleaves Python execution with language model reasoning. Related approaches include PAL (Gao et al., 2023) and ToRA (Gou et al., 2024), which similarly augment LLMs with\nprogram synthesis for mathematical reasoning.\nReinforcement learning for reasoning. DrGRPO extends Group Relative Policy Optimisation (DrGRPO, Shao et al., 2024) with per-group advantage normalisation to reduce reward variance. Our two-phase reward schedule is closest in spirit to staged RLHF curricula (Ouyang et al., 2022),\nwhere format correctness is established before optimising for task performance.\nPreference optimisation. ORPO (Hong et al., 2024) eliminates the need for a reference model by incorporating the preference loss directly into the supervised fine-tuning objective. We apply it in answer-only mode — a design not present in the original paper — to avoid trace-level reward\nhacking while retaining the benefit of preference calibration.\nAIMO precedent. The AIMO2 winners (arXiv:2504.16891) used a similar pipeline combining SFT, RL, and inference-time scaling. Our work differs primarily in (i) the explicit difficulty stratification via K-Means before data selection, (ii) the answer-only ORPO format, and (iii)

\n

MXFP4 quantisation as the deployment path.

\n

5. Reproducibility

\n

All scripts, the Kaggle inference notebook, and requirements.txt are available at:
\ngithub.com/codesrepo/aimo3-training-pipeline
\nlink to kaggle notebook
\nlink to HF model
\nlink to google colab notebook

\n
\n

Note — Google Colab H100 Amendments
\n The following changes were made to run this notebook on Google Colab H100. They do not impact model accuracy but may affect runtime:

\n

os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True' # needed for Google Colab H100
\n batch_size = 32 # default: 256
\n gpu_memory_utilization = 0.90 # default: 0.96
\n To reproduce the full pipeline:

\n
\n
# Stage 0: baseline scoring\npython data_selection.py\n# Stage 1: clustering + SFT\npython cluster_high_mismatch.py\npython convert_to_harmony.py\npython sft_lora.py\n# Stage 2: oracle traces\npython oracle_traces_aimo.py\n# Stage 3: DrGRPO\npython aimo_drgrpo_lora_r4_copy.py   # submission version\n# : python grpo_lora.py            # for reference\n# Stage 4: evaluation\npython run_evaluation.py\n# Stage 5–6: community data + preference pairs\npython get_kaggle_datasets.py\npython data_positive_negative_math.py\npython preprocess_training_samples_multi.py\n# Stage 7: ORPO\npython train_dpo.py\n# Stage 8: bake (requires 8× H100)\npython merge_lora_v2.py\n
\n

Hardware summary:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
StageHardware
Stages 0–8 (data + training)RTX PRO 6000 Blackwell (96 GB VRAM), local WSL2
Stage 8b (MXFP4 baking)8× H100 80 GB (public cloud)
Inference (Kaggle submission)1× H100 80 GB, tensor-parallel-size=1
Key package versions (see requirements.txt):
unsloth, trl, peft, transformers>=4.46, vllm, sentence-transformers,
\n

scikit-learn, datasets

\n

6. Conclusion

\n

We described a four-phase adaptation pipeline for competition mathematics that sequentially addresses the distinct failure modes of a 120B base model: format compliance (SFT), hard-problem reasoning (DrGRPO), answer distribution alignment (ORPO), and memory-efficient deployment (MXFP4 baking). Key design choices — answer-only preference pairs, two-phase DrGRPO reward, and cluster-stratified data selection — were each motivated by specific\nempirical observations from leaderboard scores and internal evaluation.\nThe final baked i700 checkpoint achieved a best public score of 42 with a 45-submission mean of 38.3 (σ=1.67), ranking 1,799 / 4,138. The most impactful single decision was dropping RAG at inference (+1.1 pts mean), while the most impactful training decision was the DrGRPO phase (+24.3 pp on hard training problems).\nLimitations and future work. The hard-50 benchmark alone was insufficient to distinguish among checkpoint candidates — public LB score was the primary signal, with runtime as secondary. A larger, competition-representative local benchmark would reduce dependence on LB submissions. A larger, competition-representative local benchmark would reduce dependence on LB submissions for checkpoint selection. The DrGRPO phase was limited to 30 total steps by runtime constraints (≤30 min/iter on H100); extending to more steps or using a larger oracle trace corpus may yield further gains.

\n

Acknowledgements

\n\n

hard-math-problems-for-aimo-3: the 50-problem hard-50 benchmark used throughout for checkpoint evaluation.

\n

Figure Reference

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
FigureDescriptionFile
1Full 10-stage training & inference pipelinecharts/pipeline_diagram.png
2K-Means elbow plot (k=15 selection)imo/openmath_data/kmeans_elbow_plot.png
32D PCA cluster distribution coloured by mismatch rateimo/openmath_data/cluster_pca_distribution.png
4Distribution of per-cluster mismatch rates, threshold=0.59imo/openmath_data/mismatch_rate_distribution.png
5DrGRPO two-phase reward structurecharts/chart6_drgrpo_phases.png
6Community dataset composition pie (11 sources, 2,664 problems)charts/chart4_community_datasets.png
7Checkpoint hard-50 scores and runtimes (model selection)charts/chart10_hard50_ablation.png
8Phase-by-phase leaderboard progressioncharts/chart9_submission_scores.png
9Final submission (v4) score distributioncharts/chart11_v4_score_distribution.png
*9 figures total. Figures 2–4 are generated by cluster_high_mismatch.py; Figures 5–9 by
writeup/generate_charts.py and writeup/generate_pipeline_diagram.py.*
", "comments": [ { "id": 3458597, "author_name": "Sam Bealing", "author_url": "", "post_date": "2026-05-16T08:31:28.810000", "content": "

Following up on the questions from @philipvonderlind and @flamingice2801 (and thanks for the very thorough response), I only have one further question:

\n
    \n
  1. K-Means cluster validation: Your data selection uses K-Means on OpenMathReasoning to identify Cluster 3 as the most informative training distribution. Was the cluster choice validated on held-out problems (e.g. by checking that Cluster 3 problems also have high mismatch rates on a different problem reference set), or could the selection have overfit to the development sample you used to define it?
    2. \n
\n

Look forward to reading your response on this.

", "votes": 0, "replies": [ { "id": 3458636, "author_name": "Mohammad Shadab Alam", "author_url": "", "post_date": "2026-05-16T10:12:11.197000", "content": "

Thank you for the question @sambealing. We did not validate the Cluster 3 choice on a separate held-out reference set, so we agree that the cluster-selection step may have some dependence on the development sample used to identify it.

\n

As a post-selection and post-training observation from partial logging, the final model appeared to solve BYTE66 from the hard-50 set more frequently than the baseline: 2/3 successful runs for the final selected model versus 1/6 successful runs for the baseline. However, given the high inherent variance of these runs, we do not treat this as conclusive validation, and chance cannot be ruled out.

\n

The motivation for selecting Cluster 3 was practical rather than claiming it was globally optimal. Our goal was to concentrate the limited LoRA training budget on a narrow semantic region where the base model showed the largest mismatch or weakness, rather than spread updates thinly across the full OpenMathReasoning distribution. Since we were training only LoRA adapters rather than the full model, we expected the update to be localized and less likely to broadly disrupt existing capabilities.

\n

The intended best-case outcome was that the model would acquire a useful problem-solving pattern for a difficult region while retaining most of its existing behavior. The worst-case risk was overfitting to that narrow region, but even in that case the selection could still be useful if it helped solve a previously inaccessible high-complexity hard problem. We therefore treated the clustering step as a targeted, compute-constrained heuristic for data selection, not as a fully validated claim about the universally most informative training distribution.

", "votes": 0, "replies": [] } ] }, { "id": 3456815, "author_name": "Philip Vonderlind", "author_url": "", "post_date": "2026-05-12T16:16:08.390000", "content": "

Hi, thank you for your submission to the Writeup Prize. @sambealing and I will post some questions below that help us in assessing your submission more closely, as described in the updated discussion post.

\n

You already have a very extensive submission. I especially like the usage of plots to convey the architecture and experiments to the readers. I have a few questions regarding some stages in your approach:

\n\n

Looking forward to your response :)

", "votes": 0, "replies": [ { "id": 3457342, "author_name": "Mohammad Shadab Alam", "author_url": "", "post_date": "2026-05-13T15:36:15.310000", "content": "

Q1: ORPO after GRPO ordering\nThank you for the insightful question. We agree that the more common ordering is ORPO before GRPO, since applying offline preference optimization after online RL can risk moving the policy away from dynamic behaviors discovered during RL. In our case, however, the DrGRPO stage was intentionally small and compute-constrained: it was run for only 30 steps on 30 data points, so we treated it less as a fully converged RL-trained model and more as a reward-aligned initialization. The subsequent ORPO stage used a much larger preference set of 1,415 pairs, providing a broader and more diverse signal.

\n

The motivation for this ordering was incremental: DrGRPO helped move the model toward better candidate behavior, while ORPO was then used to improve final-answer ranking among candidates, especially in cases where the correct answer was already present but not ranked highest. We mitigated the risk of overwriting DrGRPO-induced reasoning behavior by using answer-only preference pairs rather than full solution traces. As a result, ORPO primarily calibrated the output distribution over final answers instead of explicitly supervising alternative reasoning paths. This does not eliminate forgetting risk, but it reduces the most direct way in which the offline stage could overwrite reasoning patterns discovered during DrGRPO. We further mitigated this risk by using a conservative learning rate of 5×10−6, so the ORPO stage acted more as a mild preference-calibration step than a large distribution-shifting update.

\n

Q2: GRPO requires verified correct trajectories\nThank you for pointing this out. You are right that the phrasing in the writeup is imprecise. GRPO/DrGRPO is an on-policy RL method: the model samples outputs during training, and a reward function scores those outputs relative to each other within each group. Thus, fully verified pre-collected trajectories are not required in principle.

\n

In our setup, the verified oracle traces collected during training were not used as supervised positive trajectories for GRPO, nor were they fed to the trainer for imitation learning. Instead, they served two limited purposes. First, they defined a corpus of hard problems with confirmed integer answers. Second, the confirmed answers supplied the expected answer string used by the exact-match reward component. During DrGRPO, the policy still generated its own rollouts online, and updates were based on rewards assigned to those generated outputs.

\n

Separately, the oracle traces were also used to generate certificates for a RAG component during development. However, RAG was not activated in the final submission, so these certificates did not affect the final DrGRPO training or inference pipeline. We will revise the writeup to avoid implying that GRPO requires verified positive traces.

\n

Q3: Inference setup\nThank you for the question. Our inference setup was based on v8/23 of @andreasbis' public notebook, aimo-3-gpt-oss-120b-with-tools (scriptVersionId=290379281), which we used as the baseline notebook. We did test our fine-tuned checkpoint integrated directly into this public notebook framework.

\n

Relative to the upstream notebook, the final submitted configuration made only minimal changes:

\n

Model path — replaced the base gpt-oss-120b checkpoint with our fine-tuned MXFP4-baked i700 checkpoint, outliar/gpt-oss-120b-merged-mxfp4-i700.

\n

notebook_limit — increased slightly from 17,400 to 17,520, only to marginally extend the usable runtime budget.

\n

use_rag flag and cert_path — added for RAG ablation runs, but use_rag=0 in the final submitted configuration.

\n

All other inference parameters were kept identical to the upstream notebook: attempts=8, early_stop=4, temperature=1.0, min_p=0.02, context_tokens=65536, and gpu_memory_utilization=0.96. Therefore, the final inference strategy, including dynamic timeout, majority voting, and entropy-weighted fallback, was directly inherited from @andreasbis' implementation.

\n

In this sense, yes: the model was evaluated inside the public notebook framework, and the final submitted run differs from the public baseline primarily by the fine-tuned model weights. We will clarify this in the writeup and include the exact inference configuration to make the comparison reproducible.

\n

kaggle notebook used for final submission link

", "votes": 0, "replies": [ { "id": 3457351, "author_name": "Mohammad Shadab Alam", "author_url": "", "post_date": "2026-05-13T15:59:15.413000", "content": "

Edited on 13th May 2026 for the rebuttal phase:

\n

DrGRPO terminology: Standardized the naming of this stage as “DrGRPO” throughout the writeup. Earlier drafts used inconsistent variants such as “DGRPO,” “Dr.GRPO,” and “GRPO.”

\n

GRPO phrasing correction: Corrected the description of DrGRPO training (section 2.4). DrGRPO is an on-policy RL method and does not require pre-collected correct trajectories as supervised training examples. In our setup, the oracle traces were used only to obtain verified ground-truth integer answers for the exact-match reward component.

\n

ORPO ordering rationale: Added a clearer explanation for applying ORPO after DrGRPO (section 2.8). Because the DrGRPO stage was compute-constrained and ran for only 30 steps, we treated it as a directionally improved but under-optimized checkpoint. ORPO was then applied as a lightweight answer re-ranking stage using 1,415 answer-only preference pairs.

\n

Inference setup: Clarified the exact final inference configuration (section 2.11). The submission notebook was based on @andreasbis' public notebook version with scriptVersionId=290379281. The final setup differed from the upstream notebook in only three ways: the model path was changed to our fine-tuned checkpoint, notebook_limit was increased from 17,400 to 17,520, and RAG was disabled in the submitted configuration with use_rag=0.

", "votes": 0, "replies": [] } ] } ] }, { "id": 3445573, "author_name": "Utkarsh Sharma", "author_url": "", "post_date": "2026-04-20T05:17:02.997000", "content": "

Great writeup @outliar - the 10-stage pipeline diagram is one of the clearest I've seen in this competition, and the fact that you actually got MXFP4 baking to work with save_method=\"mxfp4\" is a genuine contribution. We attempted GRPO fine-tuning as well and lost ~$300 to the merged_4bit_forced → BitsAndBytes NF4 deployment bug that you avoided. So first: thank you for documenting the correct path.

\n

Two things jumped out from your results that I'd love your perspective on:

\n
    \n

  1. Checkpoint selection with n=2. Your Figure 7 shows i700 selected despite having the lowest hard50 mean (27.0) of all 5 checkpoints. Looking at the individual run dots, each checkpoint was evaluated only 2-3 times on hard50. With that sample size, i500's mean of 31.5 and i700's mean of 27.0 could easily overlap - the 4.5-point gap is within a single run's variance. Do you think i500 might actually generalize better to the private set, given that the private distribution is closer to hard50's difficulty than to the public LB's fixed 50-problem subset?

    2. \n

  2. Variance comparison. Your Figure 9 shows n=35 final submissions with μ=38.57, σ=1.67 - very well documented. We ran the unmodified base model (no fine-tuning, just entropy-weighted voting) 9 times and got μ=40.78, σ=1.20. The fine-tuned model has both higher variance and lower mean than the base model with a better voting function. Does this suggest that ORPO might have flattened the model's confidence across competing answers rather than sharpening it? Or do you attribute the gap primarily to the MXFP4 quantization step?

    3. \n
\n

Really appreciate the open-sourced pipeline - the difficulty-stratified data selection via K-Means is an idea I wish I'd tried.

", "votes": 0, "replies": [ { "id": 3445762, "author_name": "Mohammad Shadab Alam", "author_url": "", "post_date": "2026-04-20T13:30:43.743000", "content": "

Thank you for your thorough review and kind words your report was a direct inspiration for mine.

\n

1.Checkpoint selection

\n

Since the final LB score is the mean of 2 runs, I used this as a guiding principle throughout, including checkpoint selection, where I relied on public LB performance rather than hard-50. At decision time the evidence was:

\n\n

On hard-50 reliability, I also noticed ambiguities in the benchmark problems themselves. For example, N01418 implicitly assumes non-negative integers (unstated), and N01455's answer changes depending on whether (a,b) ranges over positive or all integers, exactly the kind of under-specification IMO problems avoid. This made me more confident that the public LB distribution should be more closer to the private set than hard-50 is. That said, if hard-50 does turn out to be more representative, the existing evidence would favour i500.

\n

2.Variance comparison

\n

Looking at my own baseline (initial exploration, no fine-tuning): scores 36, 36, 40, 35, 36 give μ=36.6, σ=1.74. The final ORPO submissions (μ=38.57, σ=1.67) represent roughly 2 additional correct problems, consistent with ORPO sharpening the model's early attention toward the correct answer on a subset of problems, rather than flattening confidence generally.

\n

This aligns with @natnitarach finding in section 5 that with μ̂=39.3, σ̂=1.7 over 21 baseline runs, each submission is essentially a lottery ticket, and their verifier-aware re-ranking (which achieves better answer selection) outperforms the fine-tuned approach. The gap you observe between your base model (μ=40.78, σ=1.20) and my fine-tuned one is worth investigating further MXFP4 quantization could be a plausible contributor and I cannot fully rule it out.

", "votes": 0, "replies": [ { "id": 3448910, "author_name": "Utkarsh Sharma", "author_url": "", "post_date": "2026-04-27T02:07:37.260000", "content": "

Thank you for the detailed checkpoint data the expanded i700 runs (38, 39, 40, 40, 40, μ=39.4) make the selection much more defensible than the n=2 I was originally questioning.

\n

The point about hard50 problem ambiguities (N01418 assuming non-negative integers, N01455's range dependency) resonates strongly. We used the same hard50 dataset for offline evaluation and found similar issues - our Branch B multicell scored 27/50 on it, but we couldn't determine how much of the gap vs our leaderboard mean (40.78) was real difficulty vs dataset noise.

\n

Your baseline comparison (μ=36.6 pre-ORPO → μ=38.57 post-ORPO, Δ≈+2) is actually a cleaner measurement of fine-tuning value than most teams produced. The fact that our inference-only approach (μ=40.78) still exceeds the fine-tuned model (μ=38.57) is consistent with both our writeups' findings - but as you note, the MXFP4 quantization step is a confound that neither of us can fully rule out.

", "votes": 0, "replies": [] }, { "id": 3449038, "author_name": "Mohammad Shadab Alam", "author_url": "", "post_date": "2026-04-27T11:03:18.390000", "content": "

Based on these findings, it is highly likely that your private LB score will be at least 1.5 points higher than mine.

", "votes": 0, "replies": [] } ] } ] } ], "comment_count": 9, "messages": [ { "id": 3445347, "postDate": "2026-04-19T15:24:44.873Z", "content": "

\n

Abstract

\n

We present a four-phase fine-tuning pipeline that adapts the frozen gpt-oss-120b base model for competition mathematic, AIMO3. We apply sequential Supervised Fine-Tuning (SFT),Group Relative Policy Optimisation (DrGRPO), and Odds-Ratio PreferenceOptimisation (ORPO), followed by MXFP4 quantisation and merging for Kaggle deployment.\nEach phase addresses a specific failure mode identified through leaderboard results and internal hard-50 evaluation. The final MXFP4-baked checkpoint at 700 ORPO steps achieved a best public leaderboard score of 42 across 85 total submissions.

\n

1. Introduction

\n

The AI Mathematical Olympiad — Progress Prize 3 (AIMO3) requires systems to solve 50 competition-style integer-answer mathematics problems per submission. During the public leaderboard phase, 50 of these 100 problems are scored and contribute to the public rank; the remaining 50 are held out and scored only in the final private evaluation run. This split means the public LB score is noisy — it reflects performance on a fixed 50-problem subset, not the full competition distribution. Correct inference demands multi-step symbolic reasoning, often interleaved with verified Python computation via a tool-use interface (Harmony). A model that can reason symbolically but fails to call tools correctly wastes the computation window; a model that calls tools but cannot reason about hard problems produces wrong answers regardless.\nOur approach decomposes the adaptation problem into four sequentially targeted phases, each designed to fix a specific failure mode of the previous phase:

\n
    \n
  1. SFT warm-start — forces the model to follow the Harmony tool-use dialogue format on problems it cannot yet solve, before any reward signal is introduced.
    2. \n
  2. DrGRPO — reinforcement learning on sparse, verified oracle traces teaches the model to reason correctly through hard competition problems.
    3. \n
  3. ORPO — preference alignment on 1,415 competition problem pairs prevents the model from drifting to plausible but incorrect answers on problems it can already solve.
    4. \n
  4. MXFP4 baking — quantises and merges the full adapter stack into a ~61 GB checkpoint that fits on a single H100 80 GB GPU within the Kaggle runtime constraint.\nBefore any training, we conducted an initial exploration phase of five leaderboard submissions to estimate sensitivity to prompt format, inference temperature, and generation count. This established that the public LB varied by ±2–3 points under different hyperparameter settings, providing the\nsignal needed to calibrate subsequent phases.\nAll code is publicly available at\ngithub.com/codesrepo/aimo3-training-pipeline.\nThe sections below document each step with sufficient detail for full reproduction.
    5. \n
\n
\n

Figure 1 — Full training and inference pipeline (10 stages).\n \">\n

2. Method

\n

2.1 Dataset and Baseline Evaluation (Step 1)

\n

Dataset. We use the tir split of nvidia/OpenMathReasoning, filtered to examples whose\n expected answer can be parsed as an integer using Python's int(str(answer).strip()). The filter yields 249,341 problems from the tir split. This split contains tool-interleaved reasoning traces generated by NVIDIA's Nemotron models, making it the most suitable public source for Harmony-format fine-tuning.\n OpenMathReasoning subset Baseline: We evaluated the unmodified gpt-oss-120b base model under greedy decoding (temperature T=0.0, seed=42) on a 662-randomly selected problem subset using Harmony tool-augmented inference. The model answered 351 problems correctly (53.0%)

\n

2.2 Difficulty-Stratified Data Selection (Step 2)

\n

Fine-tuning on problems the model already solves wastes capacity and may induce regression on hard problems. We therefore stratify the scored subset by difficulty before selecting training data.\n Embedding and clustering. All 662 scored problems were embedded using the\n all-MiniLM-L6-v2 sentence transformer (384-dimensional embeddings). K-Means clustering with k=15 was applied to partition the embedding space into regions of similar semantic difficulty. The choice of k=15 was informed by both empirical validation and the elbow criterion — we ran multiple cluster configurations to identify a value that yielded a sizable subset with a high error rate, and k=15 struck the right balance. For completeness, the elbow criterion applied to Within-Cluster Sum of Squares (WCSS) over k ∈ {5, 10, …, 100} is presented in Figure 2.\n Figure 2 — K-Means elbow plot used to select k=15.\n \"\"\n Figure 3 — 2D PCA projection of the 662 scored problems coloured by cluster mismatch rate.\n Cluster 3 (68.12% mismatch, top right) is the selected high-difficulty region.\n \"\"\n Mismatch threshold. For each cluster we computed the mismatch rate — the fraction of problems the base model answered incorrectly. A threshold of ≥0.591 identified highest-difficulty clusters. A single cluster — Cluster 3, mismatch rate 68.12% — exceeded this threshold. The threshold was chosen to select a compact, maximally hard training set while retaining enough problems for meaningful fine-tuning.\n Figure 4 — Distribution of per-cluster mismatch rates (k=15). The red dashed line marks\n the 0.591 threshold; only Cluster 3 (right tail) exceeds it.\n \"\"\n The fitted K-Means model was then applied to the full 249,341-example dataset to assign every problem a cluster label and mismatch rate. From the full dataset, 2,710 hard examples belonging to Cluster 3 were selected as the SFT and DrGRPO training pool.

\n

2.3 SFT Warm-start (Step 2b)

\n

Motivation. Without format conditioning, the base model frequently violates the Harmony multi-turn tool-use protocol during early DrGRPO rollouts — misformatting tool calls, omitting protocol headers, or failing to emit \\boxed{answer}. These format errors dominate the DrGRPO reward signal in early steps and prevent the policy from learning reasoning strategies. The SFT warm-start fixes the format before any reward signal is introduced.\n Data preparation. Each of the 2,710 high-mismatch examples includes an oracle solution from the OpenMathReasoning dataset. We parsed each oracle solution to extract (i) Python tool call blocks, (ii) tool output blocks, and (iii) the final assistant reasoning turn, then re-encoded the full sequence into the Harmony multi-turn message format:

\n
\n
system prompt → user problem → [assistant turn → tool call → tool output]* → \\boxed{answer}\n
\n

Training. We trained a LoRA adapter on gpt-oss-120b (4-bit BnB NF4 quantisation) using the 2,710 Harmony-formatted examples. LoRA targeted the query, key, value, output, gate, up, and down projections (rank=4, α=8).

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
HyperparameterValue
Base modelgpt-oss-120b (BnB NF4)
LoRA rank / α4 / 8
Target modulesq, k, v, o, gate, up, down
Training examples2,710
HardwareRTX PRO 6000 Blackwell (96 GB VRAM), local WSL2
\n

2.4 Oracle Trace Collection (Step 3)

\n

DrGRPO is an on-policy RL method: the model samples its own rollouts during training and a reward function scores them group-relatively. The 207 verified oracle traces served not as supervised positive examples but as a confirmed-answer corpus — each trace's integer answer was used as the ground truth for the exact-match reward component, ensuring the reward signal is non-degenerate on hard problems. We collected these 207 samples iteratively by running gpt-oss-120b via a local vLLM server on every problem where the base model had previously failed, allowing up to 8 tool-augmented attempts per problem. Only traces where the predicted answer matched the expected answer were retained.\nEach retained trace was annotated with the cluster ID and mismatch rate of its source problem (using the saved K-Means model), enabling downstream cluster-aware reward scaling.\nTraces accumulated across training rounds:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
RoundCumulative correct traces
1151
2173
3207 ← used for DrGRPO
4219
5243
6283
The round-3 snapshot of 207 traces was used for DrGRPO training, as it represented the best available signal at the time of the reinforcement learning run.
An optional post-processing pass (aimo_gen_cert.py) extracted structured fields from each trace (key_idea, proof_skeleton, sanity_checks, solved_attempt), producing compact certificates for rag stage but it was not used in the final submission.
\n

2.5 DrGRPO Reinforcement Learning (Step 4)

\n

Motivation. After SFT the model reliably follows the Harmony format, but still fails on hard reasoning. DrGRPO provides a direct reward signal: sample candidate solutions from the current policy, compare each to the verified oracle answer, and reinforce correct reasoning chains via group-relative advantage normalisation.\nWhy sparse training data. Only ~30 oracle traces per iteration were used. This is a deliberate choice: the DrGRPO reward signal on verified hard problems is already strong, and overfitting to a small set forces the policy to generalise. Larger batches would require either more oracle data (expensive to collect) or accepting lower-quality traces. Each iteration runs on a single sample (batch=1).\nTwo-phase reward. Training was split into two phases to avoid reward collapse from competing\nobjectives:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
PhaseStepsDominant reward components
A — Certificate-first10\\boxed{} format (0.5), protocol headers (0.3), exact answer (5.0), tool consistency (0.7), length penalty (−0.05)
B — Answer-first20Exact answer (5.0, dominant), tool ok bonus (0.7), tool consistency (0.7), length penalty (−0.05)
Phase A reinforces structural correctness before Phase B shifts emphasis to answer accuracy.
\n
\n

Figure 5 — DrGRPO two-phase reward structure.\n \"\"

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
HyperparameterValue
BaseSFT-warmed gpt-oss-120b (BnB NF4)
LoRA rank / α4 / 8
LR5 × 10⁻⁶
KL β0.001
Generations per problem (G)2
Batch size1 (single sample per step)
Total steps30 (Phase A: 10, Phase B: 20)
scale_rewardsFalse
HardwareRTX PRO 6000 Blackwell (96 GB VRAM), local WSL2
\n

2.6 Community Dataset Aggregation (Step 5)

\n

ORPO requires diverse preference pairs covering the specific error types the model makes when it can solve a problem but selects the wrong answer from competing candidates. We assembled a community dataset from 11 public competition mathematics sources, filtered to integer answers in [0, 99999] and deduplicated on problem text.\n CrystalMath-Preview (Yi-Chia Chen, ycchen,\n ycchen/Crystal-Math-Preview)\n contributed 2,129 problems in raw form. After filtering to integer answers with non-null solutions, answers in [0, 99999], deduplication, and removal of problems with Chinese characters, 634 problems survived into the pool. The hard-50 benchmark (50 problems) used for checkpoint evaluation is from ZFTurbo\n (zfturbo,\n hard-math-problems-for-aimo-3).\n The full 2,664-problem pool was scored with gpt-oss-120b (MAX_ATTEMPTS=3 per problem) via a parallel vLLM server, recording both correct and wrong attempt text for every problem.

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
SourceProblems
ODA-Math-460k659
CrystalMath-Preview (filtered from 2,129)634
DAPO-17K561
olympiads-ref-base258
Nemotron-Math-V2228
PaCoRe-Train-8k133
IMO-AnswerBench67
Omni-MATH50
OlymMATH25
BeyondAIME22
MathArena18
Putnam-Axiom9
Total2,664
Figure 6 — Community dataset composition (2,664 problems from 11 sources).
\"\"
\n

2.7 Preference Pair Construction (Step 6)

\n

From the scored community dataset we constructed ORPO training pairs. Early versions of the dataset used full Harmony multi-turn traces as chosen/rejected sequences (382 pairs, Feb 2026). After experimentation, we found that answer-only pairs produced more stable ORPO training\n and better LB scores. The final format is:

\n

\n \n \n \n \n

2.8 ORPO Preference Fine-tuning (Step 7)

\n \n \n \n \n \n

Motivation. The purpose of DrGRPO was to improve the model’s reasoning behavior under a reward signal. However, given our limited compute budget, the DrGRPO stage was necessarily small: only 30 steps were run, so we treated the resulting checkpoint as a directionally improved but under-optimized policy rather than a fully converged RL model. We then applied ORPO as a lightweight preference-calibration stage on top of this checkpoint.\n This ordering was deliberate. The ORPO dataset contained 1,415 answer-only preference pairs from 11 competition sources, where the preferred response was the correct final answer and the rejected response was a plausible incorrect answer. These pairs were designed to address cases where the correct answer already appeared in the candidate pool but was not ranked highest. In other words, ORPO was used to shift the model’s final-answer ranking toward correct candidates, rather than to teach new reasoning procedures.\n To reduce the risk of overwriting reasoning behavior learned during DrGRPO, we used answer-only pairs rather than full solution traces. This format exposes the model to a broad range of error types, such as off-by-one mistakes in modular arithmetic, incorrect counting, and wrong base cases, without prescribing alternative reasoning chains. We do not claim this eliminates forgetting, since ORPO still updates the policy, but it reduces the most direct mechanism by which offline preference optimization could overwrite DrGRPO-induced reasoning patterns.\n ORPO, or Odds-Ratio Preference Optimization, combines a supervised loss on the preferred response with a log-odds preference loss, and does not require a separate reference model. This made it computationally practical for our setting: a 120B-parameter model trained under 4-bit quantization on a single-GPU budget.\n Incremental training. Training did not run from scratch to 1,000 steps in a single session. The i500 checkpoint was saved first; training then continued for 200 additional steps to produce i700, with checkpoints saved every 100 steps. This incremental approach allowed public leaderboard evaluation between checkpoints and avoided committing to a full training run before seeing LB signal.

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
HyperparameterValue
LR5 × 10⁻⁶
β (ORPO)0.001
Max prompt length2,048 tokens
Max total length4,224 tokens
Batch size1 (single sample per step)
Total steps1,000
LoRA rank / α / dropout4 / 8 / 0.05
OptimiserPaged AdamW 8-bit
HardwareRTX PRO 6000 Blackwell (96 GB VRAM), local WSL2
\n

2.9 Checkpoint Selection via Hard-50 Elimination (Step 7, continued)

\n

We evaluated five model ORPO variants on the 50-problem hard-50 benchmark (ZFTurbo), each run 2–3 times to account for stochastic generation. Checkpoint selection followed a three-stage elimination process:

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
ModelHard-50 meanRuntime meanDecision
i50031.5194.5 min① Eliminated — public LB score lower than i700
v23032026 (1.2k)25.3186.7 min② Eliminated — slowest remaining runtime
v25032026 (1.4k)28.0182.0 min② Eliminated — slower runtime
i1k-gold28.5172.5 min② Eliminated — public LB lower than i700
i70027.0175.0 minSelected — best public LB + acceptable runtime
Elimination rationale (two criteria).
\n
    \n
  1. i500 eliminated first (criterion 1: public LB score). Despite achieving the highest hard-50 mean (31.5 / 50), i500's public LB score on the 50 evaluated problems was lower than i700's. The public LB — scored on a fixed 50-problem subset — was the primary selection signal. i500's hard-50 runtime (194.5 min for 50 problems) also left the least headroom against the 5-hour Kaggle GPU notebook limit.
    2. \n
  2. Remaining candidates eliminated (criterion 2: runtime). Among the remaining checkpoints, v23032026 (186.7 min) and v25032026 (182.0 min) were eliminated on runtime; i1k-gold (172.5 min) was eliminated on public LB score, which was lower than i700's despite its faster runtime. i700 showed the best combination of public LB performance and acceptable runtime.
    3. \n
  3. i700 selected. At a mean hard-50 runtime of 175 min for 50 problems (~3.6 min/problem), the projected runtime for the full 100-problem notebook (~357 min on local hardware); actual Kaggle H100 submissions completed within the 5-hour GPU limit. i700 achieved the best public LB across 45 final submissions (mean 38.3, max 42).\n Figure 7 — Checkpoint evaluation: hard-50 scores and runtimes by model variant.\n \"\"\n Left panel: mean scores with individual run dots; right panel: mean runtimes.
    4. \n
\n

2.10 MXFP4 Baking (Step 8)

\n

The final ORPO LoRA adapter (i700) was merged into the base model weights using Unsloth's save_pretrained_merged with save_method=\"mxfp4\" on a node with 8× H100 80 GB GPUs (Kaggle notebook).

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
FormatSize
Naïve BF16 merge~218 GB
MXFP4 baked merge~61 GB
The ~61 GB checkpoint fits within the single H100 80 GB memory budget with headroom for the vLLM KV cache and Harmony tool interface.
\n

2.11 Kaggle Inference (Step 9–10)

\n

The inference notebook is based on v8/23 of @andreasbis's public notebook (scriptVersionId=290379281). Differences are limited to: model path (MXFP4 i700 checkpoint), notebook_limit increased from 17,400 to 17,520, and a use_rag flag set to 0 in the final submission.\n The merged MXFP4 checkpoint was uploaded as a Kaggle model and served via vLLM inside the submission notebook. Inference used Harmony-encoded prompts with tool-augmented self-consistency across multiple generation attempts. Each problem is allocated a timeout equal to the largest of: 300 s (minimum per problem), 900 s (hard-problem cap), or the remaining time budget after reserving 300 s for every unsolved problem — ensuring accumulated surplus is fully utilised.
\n For answer selection, majority voting is applied across up to 8 attempts with early stopping once 4 solutions agree. Otherwise, entropy weighing is performed to determine the rank — the generation the model was most confident about. Both the dynamic timeout strategy and majority voting with entropy weighing are widely used in public AIMO notebooks.

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
SettingValue
Serving frameworkvLLM
Hardware1× H100 80 GB
Tensor-parallel size1
Prompt encodingHarmony multi-turn
RAG at inferenceNone (despite notebook name aimo_rag_baked_c283)
The notebook name refers to the 283-certificate training set used for DrGRPO, not to retrieval augmentation at inference time. RAG was dropped in the final configuration after submissions with RAG (mean 37.2–37.9) consistently underperformed plain inference with the baked i700 model (mean 38.3).
\n

3. Results

\n

3.1 Leaderboard Progression

\n

We made 85 total submissions across seven distinct configuration phases (6 failed attempts excluded from analysis). The table and chart below show the progression from baseline exploration to the final baked i700 configuration.

\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
PhaseValid submissionsMean LBMax LBNotes
Initial exploration536.640Hyperparameter sensitivity probing
LoRA adapter (not baked)534.437Slower inference hurt utilisation
RAG only (base model)437.541Retrieved solutions at inference
LoRA adapter + RAG536.438DrGRPO adapter loaded separately
Baked LoRA + RAG1037.940Seed DrGRPO + DPO iterations merged
DPO baked + RAG537.240
DPO baked, no RAG (final)4538.342MXFP4 i700, best configuration
Key findings from the progression:
\n \n

3.2 Ablation: Checkpoint Selection

\n

The hard-50 benchmark (Figure 7) revealed a counterintuitive result: i500 scored highest on hard-50 (mean 31.5/50) but worst on the public leaderboard, suggesting it may have overfit to the local benchmark distribution. Since the private LB score is the mean of 2 runs, we used public LB performance rather than hard-50 as the primary signal for checkpoint selection. At decision time the evidence was:

\n
\n\n

i700 was the clear choice on that basis. We also noticed ambiguities in several hard-50 problems that reinforced this decision — for example, N01418 implicitly assumes non-negative integers (unstated), and N01455's answer changes depending on whether (a,b) ranges over positive or all integers, exactly the kind of under-specification that IMO problems avoid. This made us more confident that the public LB distribution is closer to the private set than hard-50 is.

\n

4. Related Work

\n

Tool-augmented reasoning. Our inference pipeline builds on the Harmony tool-use framework, which interleaves Python execution with language model reasoning. Related approaches include PAL (Gao et al., 2023) and ToRA (Gou et al., 2024), which similarly augment LLMs with\nprogram synthesis for mathematical reasoning.\nReinforcement learning for reasoning. DrGRPO extends Group Relative Policy Optimisation (DrGRPO, Shao et al., 2024) with per-group advantage normalisation to reduce reward variance. Our two-phase reward schedule is closest in spirit to staged RLHF curricula (Ouyang et al., 2022),\nwhere format correctness is established before optimising for task performance.\nPreference optimisation. ORPO (Hong et al., 2024) eliminates the need for a reference model by incorporating the preference loss directly into the supervised fine-tuning objective. We apply it in answer-only mode — a design not present in the original paper — to avoid trace-level reward\nhacking while retaining the benefit of preference calibration.\nAIMO precedent. The AIMO2 winners (arXiv:2504.16891) used a similar pipeline combining SFT, RL, and inference-time scaling. Our work differs primarily in (i) the explicit difficulty stratification via K-Means before data selection, (ii) the answer-only ORPO format, and (iii)

\n

MXFP4 quantisation as the deployment path.

\n

5. Reproducibility

\n

All scripts, the Kaggle inference notebook, and requirements.txt are available at:
\ngithub.com/codesrepo/aimo3-training-pipeline
\nlink to kaggle notebook
\nlink to HF model
\nlink to google colab notebook

\n
\n

Note — Google Colab H100 Amendments
\n The following changes were made to run this notebook on Google Colab H100. They do not impact model accuracy but may affect runtime:

\n

os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True' # needed for Google Colab H100
\n batch_size = 32 # default: 256
\n gpu_memory_utilization = 0.90 # default: 0.96
\n To reproduce the full pipeline:

\n
\n
# Stage 0: baseline scoring\npython data_selection.py\n# Stage 1: clustering + SFT\npython cluster_high_mismatch.py\npython convert_to_harmony.py\npython sft_lora.py\n# Stage 2: oracle traces\npython oracle_traces_aimo.py\n# Stage 3: DrGRPO\npython aimo_drgrpo_lora_r4_copy.py   # submission version\n# : python grpo_lora.py            # for reference\n# Stage 4: evaluation\npython run_evaluation.py\n# Stage 5–6: community data + preference pairs\npython get_kaggle_datasets.py\npython data_positive_negative_math.py\npython preprocess_training_samples_multi.py\n# Stage 7: ORPO\npython train_dpo.py\n# Stage 8: bake (requires 8× H100)\npython merge_lora_v2.py\n
\n

Hardware summary:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
StageHardware
Stages 0–8 (data + training)RTX PRO 6000 Blackwell (96 GB VRAM), local WSL2
Stage 8b (MXFP4 baking)8× H100 80 GB (public cloud)
Inference (Kaggle submission)1× H100 80 GB, tensor-parallel-size=1
Key package versions (see requirements.txt):
unsloth, trl, peft, transformers>=4.46, vllm, sentence-transformers,
\n

scikit-learn, datasets

\n

6. Conclusion

\n

We described a four-phase adaptation pipeline for competition mathematics that sequentially addresses the distinct failure modes of a 120B base model: format compliance (SFT), hard-problem reasoning (DrGRPO), answer distribution alignment (ORPO), and memory-efficient deployment (MXFP4 baking). Key design choices — answer-only preference pairs, two-phase DrGRPO reward, and cluster-stratified data selection — were each motivated by specific\nempirical observations from leaderboard scores and internal evaluation.\nThe final baked i700 checkpoint achieved a best public score of 42 with a 45-submission mean of 38.3 (σ=1.67), ranking 1,799 / 4,138. The most impactful single decision was dropping RAG at inference (+1.1 pts mean), while the most impactful training decision was the DrGRPO phase (+24.3 pp on hard training problems).\nLimitations and future work. The hard-50 benchmark alone was insufficient to distinguish among checkpoint candidates — public LB score was the primary signal, with runtime as secondary. A larger, competition-representative local benchmark would reduce dependence on LB submissions. A larger, competition-representative local benchmark would reduce dependence on LB submissions for checkpoint selection. The DrGRPO phase was limited to 30 total steps by runtime constraints (≤30 min/iter on H100); extending to more steps or using a larger oracle trace corpus may yield further gains.

\n

Acknowledgements

\n\n

hard-math-problems-for-aimo-3: the 50-problem hard-50 benchmark used throughout for checkpoint evaluation.

\n

Figure Reference

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
FigureDescriptionFile
1Full 10-stage training & inference pipelinecharts/pipeline_diagram.png
2K-Means elbow plot (k=15 selection)imo/openmath_data/kmeans_elbow_plot.png
32D PCA cluster distribution coloured by mismatch rateimo/openmath_data/cluster_pca_distribution.png
4Distribution of per-cluster mismatch rates, threshold=0.59imo/openmath_data/mismatch_rate_distribution.png
5DrGRPO two-phase reward structurecharts/chart6_drgrpo_phases.png
6Community dataset composition pie (11 sources, 2,664 problems)charts/chart4_community_datasets.png
7Checkpoint hard-50 scores and runtimes (model selection)charts/chart10_hard50_ablation.png
8Phase-by-phase leaderboard progressioncharts/chart9_submission_scores.png
9Final submission (v4) score distributioncharts/chart11_v4_score_distribution.png
*9 figures total. Figures 2–4 are generated by cluster_high_mismatch.py; Figures 5–9 by
writeup/generate_charts.py and writeup/generate_pipeline_diagram.py.*
", "rawMarkdown": "\n**Mohammad Shadab Alam** | **Kaggle**: [outliar](https://www.kaggle.com/outliar) | Public leaderboard score (best): **42** | Private leaderboard score (avg of 2): **39.5**| Rank **2008 / 4,138** \n\n## Abstract\n\nWe present a four-phase fine-tuning pipeline that adapts the frozen gpt-oss-120b base model for competition mathematic, AIMO3. We apply sequential Supervised Fine-Tuning (SFT),Group Relative Policy Optimisation (DrGRPO), and Odds-Ratio PreferenceOptimisation (ORPO), followed by MXFP4 quantisation and merging for Kaggle deployment.\nEach phase addresses a specific failure mode identified through leaderboard results and internal hard-50 evaluation. The final MXFP4-baked checkpoint at 700 ORPO steps achieved a best public leaderboard score of **42** across 85 total submissions.\n\n\n## 1. Introduction\n\nThe AI Mathematical Olympiad — Progress Prize 3 (AIMO3) requires systems to solve 50 competition-style integer-answer mathematics problems per submission. During the public leaderboard phase, **50 of these 100 problems are scored** and contribute to the public rank; the remaining 50 are held out and scored only in the final private evaluation run. This split means the public LB score is noisy — it reflects performance on a fixed 50-problem subset, not the full competition distribution. Correct inference demands multi-step symbolic reasoning, often interleaved with verified Python computation via a tool-use interface (Harmony). A model that can reason symbolically but fails to call tools correctly wastes the computation window; a model that calls tools but cannot reason about hard problems produces wrong answers regardless.\n\nOur approach decomposes the adaptation problem into four sequentially targeted phases, each designed to fix a specific failure mode of the previous phase:\n\n1. **SFT warm-start** — forces the model to follow the Harmony tool-use dialogue format on problems it cannot yet solve, before any reward signal is introduced.\n2. **DrGRPO** — reinforcement learning on sparse, verified oracle traces teaches the model to reason correctly through hard competition problems.\n3. **ORPO** — preference alignment on 1,415 competition problem pairs prevents the model from drifting to plausible but incorrect answers on problems it can already solve. \n4. **MXFP4 baking** — quantises and merges the full adapter stack into a ~61 GB checkpoint that fits on a single H100 80 GB GPU within the Kaggle runtime constraint.\n\nBefore any training, we conducted an **initial exploration phase** of five leaderboard submissions to estimate sensitivity to prompt format, inference temperature, and generation count. This established that the public LB varied by ±2–3 points under different hyperparameter settings, providing the\nsignal needed to calibrate subsequent phases.\n\nAll code is publicly available at\n[github.com/codesrepo/aimo3-training-pipeline](https://github.com/codesrepo/aimo3-training-pipeline).\nThe sections below document each step with sufficient detail for full reproduction.\n\n> **Figure 1** — Full training and inference pipeline (10 stages).\n![> `charts/pipeline_diagram.png`](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2Fa33430ea68af208fe41111a394d36e10%2Fpipeline_diagram.png?generation=1776602464559958&alt=media)\n\n\n## 2. Method\n\n### 2.1 Dataset and Baseline Evaluation (Step 1)\n\n**Dataset.** We use the `tir` split of `nvidia/OpenMathReasoning`, filtered to examples whose\nexpected answer can be parsed as an integer using Python's `int(str(answer).strip())`. The filter yields **249,341 problems** from the `tir` split. This split contains tool-interleaved reasoning traces generated by NVIDIA's Nemotron models, making it the most suitable public source for Harmony-format fine-tuning.\n\nOpenMathReasoning subset Baseline: We evaluated the unmodified gpt-oss-120b base model under greedy decoding (temperature T=0.0, seed=42) on a 662-randomly selected problem subset using Harmony tool-augmented inference. The model answered **351** problems correctly (53.0%)\n\n### 2.2 Difficulty-Stratified Data Selection (Step 2)\n\nFine-tuning on problems the model already solves wastes capacity and may induce regression on hard problems. We therefore stratify the scored subset by difficulty before selecting training data.\n\n\n\n**Embedding and clustering.** All 662 scored problems were embedded using the\n`all-MiniLM-L6-v2` sentence transformer (384-dimensional embeddings). K-Means clustering with k=15 was applied to partition the embedding space into regions of similar semantic difficulty. The choice of k=15 was informed by both empirical validation and the elbow criterion — we ran multiple cluster configurations to identify a value that yielded a sizable subset with a high error rate, and k=15 struck the right balance. For completeness, the elbow criterion applied to Within-Cluster Sum of Squares (WCSS) over k ∈ {5, 10, …, 100} is presented in Figure 2.\n\n> **Figure 2** — K-Means elbow plot used to select k=15.\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2Fb987e17bbd4cd4febff7e9c852215e16%2Fkmeans_elbow_plot.png?generation=1776602585406084&alt=media)\n\n> **Figure 3** — 2D PCA projection of the 662 scored problems coloured by cluster mismatch rate.\n> Cluster 3 (68.12% mismatch, top right) is the selected high-difficulty region.\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2Fe22cfdcc8cab428336bc57e9fca8876d%2Fcluster_pca_distribution.png?generation=1776602608035905&alt=media)\n\n**Mismatch threshold.** For each cluster we computed the *mismatch rate* — the fraction of problems the base model answered incorrectly. A threshold of **≥0.591** identified highest-difficulty clusters. A single cluster — **Cluster 3, mismatch rate 68.12%** — exceeded this threshold. The threshold was chosen to select a compact, maximally hard training set while retaining enough problems for meaningful fine-tuning.\n\n> **Figure 4** — Distribution of per-cluster mismatch rates (k=15). The red dashed line marks\n> the 0.591 threshold; only Cluster 3 (right tail) exceeds it.\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2F324c3649e92395605ca0405a58e5fe4c%2Fmismatch_rate_distribution.png?generation=1776602655512937&alt=media)\n\n\nThe fitted K-Means model was then applied to the full 249,341-example dataset to assign every problem a cluster label and mismatch rate. From the full dataset, 2,710 hard examples belonging to Cluster 3 were selected as the SFT and DrGRPO training pool.\n\n\n### 2.3 SFT Warm-start (Step 2b)\n\n**Motivation.** Without format conditioning, the base model frequently violates the Harmony multi-turn tool-use protocol during early DrGRPO rollouts — misformatting tool calls, omitting protocol headers, or failing to emit `\\boxed{answer}`. These format errors dominate the DrGRPO reward signal in early steps and prevent the policy from learning reasoning strategies. The SFT warm-start fixes the format before any reward signal is introduced.\n\n**Data preparation.** Each of the 2,710 high-mismatch examples includes an oracle solution from the OpenMathReasoning dataset. We parsed each oracle solution to extract (i) Python tool call blocks, (ii) tool output blocks, and (iii) the final assistant reasoning turn, then re-encoded the full sequence into the Harmony multi-turn message format:\n\n```\nsystem prompt → user problem → [assistant turn → tool call → tool output]* → \\boxed{answer}\n```\n\n**Training.** We trained a LoRA adapter on gpt-oss-120b (4-bit BnB NF4 quantisation) using the 2,710 Harmony-formatted examples. LoRA targeted the query, key, value, output, gate, up, and down projections (rank=4, α=8).\n\n| Hyperparameter | Value |\n|----------------|-------|\n| Base model | gpt-oss-120b (BnB NF4) |\n| LoRA rank / α | 4 / 8 |\n| Target modules | q, k, v, o, gate, up, down |\n| Training examples | 2,710 |\n| Hardware | RTX PRO 6000 Blackwell (96 GB VRAM), local WSL2 |\n\n\n### 2.4 Oracle Trace Collection (Step 3)\n\nDrGRPO is an on-policy RL method: the model samples its own rollouts during training and a reward function scores them group-relatively. The 207 verified oracle traces served not as supervised positive examples but as a confirmed-answer corpus — each trace's integer answer was used as the ground truth for the exact-match reward component, ensuring the reward signal is non-degenerate on hard problems. We collected these 207 samples iteratively by running gpt-oss-120b via a local vLLM server on every problem where the base model had previously failed, allowing up to 8 tool-augmented attempts per problem. Only traces where the predicted answer matched the expected answer were retained.\n\nEach retained trace was annotated with the cluster ID and mismatch rate of its source problem (using the saved K-Means model), enabling downstream cluster-aware reward scaling.\n\nTraces accumulated across training rounds:\n\n| Round | Cumulative correct traces |\n|-------|--------------------------|\n| 1 | 151 |\n| 2 | 173 |\n| 3 | **207** ← used for DrGRPO |\n| 4 | 219 |\n| 5 | 243 |\n| 6 | 283 |\n\nThe round-3 snapshot of **207 traces** was used for DrGRPO training, as it represented the best available signal at the time of the reinforcement learning run.\n\nAn optional post-processing pass (`aimo_gen_cert.py`) extracted structured fields from each trace (`key_idea`, `proof_skeleton`, `sanity_checks`, `solved_attempt`), producing compact certificates for rag stage but it was not used in the final submission.\n\n### 2.5 DrGRPO Reinforcement Learning (Step 4)\n\n**Motivation.** After SFT the model reliably follows the Harmony format, but still fails on hard reasoning. DrGRPO provides a direct reward signal: sample candidate solutions from the current policy, compare each to the verified oracle answer, and reinforce correct reasoning chains via group-relative advantage normalisation.\n\n**Why sparse training data.** Only **~30 oracle traces per iteration** were used. This is a deliberate choice: the DrGRPO reward signal on verified hard problems is already strong, and overfitting to a small set forces the policy to generalise. Larger batches would require either more oracle data (expensive to collect) or accepting lower-quality traces. Each iteration runs on a single sample (batch=1).\n\n**Two-phase reward.** Training was split into two phases to avoid reward collapse from competing\nobjectives:\n\n| Phase | Steps | Dominant reward components |\n|-------|-------|---------------------------|\n| A — Certificate-first | 10 | `\\boxed{}` format (0.5), protocol headers (0.3), exact answer (5.0), tool consistency (0.7), length penalty (−0.05) |\n| B — Answer-first | 20 | Exact answer (5.0, dominant), tool ok bonus (0.7), tool consistency (0.7), length penalty (−0.05) |\n\nPhase A reinforces structural correctness before Phase B shifts emphasis to answer accuracy.\n\n> **Figure 5** — DrGRPO two-phase reward structure.\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2F638a254b1379e7209fae62d0d3793b91%2Fchart6_drgrpo_phases.png?generation=1776603176696481&alt=media)\n| Hyperparameter | Value |\n|----------------|-------|\n| Base | SFT-warmed gpt-oss-120b (BnB NF4) |\n| LoRA rank / α | 4 / 8 |\n| LR | 5 × 10⁻⁶ |\n| KL β | 0.001 |\n| Generations per problem (G) | 2 |\n| Batch size | 1 (single sample per step) |\n| Total steps | 30 (Phase A: 10, Phase B: 20) |\n| scale_rewards | False |\n| Hardware | RTX PRO 6000 Blackwell (96 GB VRAM), local WSL2 |\n\n\n### 2.6 Community Dataset Aggregation (Step 5)\n\nORPO requires diverse preference pairs covering the specific error types the model makes when it *can* solve a problem but selects the wrong answer from competing candidates. We assembled a community dataset from **11 public competition mathematics sources**, filtered to integer answers in [0, 99999] and deduplicated on problem text.\n\n**CrystalMath-Preview** (Yi-Chia Chen, [ycchen](https://www.kaggle.com/threerabbits),\n[ycchen/Crystal-Math-Preview](https://huggingface.co/datasets/ycchen/Crystal-Math-Preview))\ncontributed 2,129 problems in raw form. After filtering to integer answers with non-null solutions, answers in [0, 99999], deduplication, and removal of problems with Chinese characters, **634 problems** survived into the pool. The **hard-50 benchmark** (50 problems) used for checkpoint evaluation is from ZFTurbo\n([zfturbo](https://www.kaggle.com/zfturbo),\n[hard-math-problems-for-aimo-3](https://www.kaggle.com/datasets/zfturbo/hard-math-problems-for-aimo-3)).\n\nThe full 2,664-problem pool was scored with gpt-oss-120b (MAX_ATTEMPTS=3 per problem) via a parallel vLLM server, recording both correct and wrong attempt text for every problem.\n\n| Source | Problems |\n|--------|---------|\n| ODA-Math-460k | 659 |\n| CrystalMath-Preview (filtered from 2,129) | 634 |\n| DAPO-17K | 561 |\n| olympiads-ref-base | 258 |\n| Nemotron-Math-V2 | 228 |\n| PaCoRe-Train-8k | 133 |\n| IMO-AnswerBench | 67 |\n| Omni-MATH | 50 |\n| OlymMATH | 25 |\n| BeyondAIME | 22 \n| MathArena | 18 |\n| Putnam-Axiom | 9 |\n| **Total** | **2,664** |\n\n> **Figure 6** — Community dataset composition (2,664 problems from 11 sources).\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2F8e2e3132a212892d2012df0738c3b72e%2Fchart4_community_datasets.png?generation=1776603215544557&alt=media)\n\n\n### 2.7 Preference Pair Construction (Step 6)\n\nFrom the scored community dataset we constructed ORPO training pairs. Early versions of the dataset used full Harmony multi-turn traces as chosen/rejected sequences (382 pairs, Feb 2026). After experimentation, we found that **answer-only pairs** produced more stable ORPO training\nand better LB scores. The final format is:\n\n- **Prompt**: problem statement\n- **Chosen**: `\\boxed{correct_answer}`\n- **Rejected**: `\\boxed{wrong_answer}`\n\nNo solution trace or reasoning chain is included. The model is calibrated on *which numerical answer is correct*, relying on the reasoning capabilities acquired during SFT and DrGRPO to produce solution paths. This design avoids introducing hallucinated reasoning in chosen traces, keeps the token budget minimal, and allows multiple rejected pairs per problem (one per wrong attempt).\n\nThe dataset grew iteratively as additional community problems were scored:\n\n| Version | Date | Pairs |\n|---------|------|-------|\n| v1 | Early Feb 2026 | 382 |\n| v2 | Mid Feb 2026 | 856 |\n| v3 | Mar 18, 2026 | 1,416 |\n| v4 (final) | Mar 25, 2026 | 1,415 |\n\n Of the 2,664 scored problems, 1,415 required more than one attempt, each yielding at least one wrong trace paired against a correct one.\n\n### 2.8 ORPO Preference Fine-tuning (Step 7)\n\n**Motivation.** The purpose of DrGRPO was to improve the model’s reasoning behavior under a reward signal. However, given our limited compute budget, the DrGRPO stage was necessarily small: only 30 steps were run, so we treated the resulting checkpoint as a directionally improved but under-optimized policy rather than a fully converged RL model. We then applied ORPO as a lightweight preference-calibration stage on top of this checkpoint.\n\nThis ordering was deliberate. The ORPO dataset contained 1,415 answer-only preference pairs from 11 competition sources, where the preferred response was the correct final answer and the rejected response was a plausible incorrect answer. These pairs were designed to address cases where the correct answer already appeared in the candidate pool but was not ranked highest. In other words, ORPO was used to shift the model’s final-answer ranking toward correct candidates, rather than to teach new reasoning procedures.\n\nTo reduce the risk of overwriting reasoning behavior learned during DrGRPO, we used answer-only pairs rather than full solution traces. This format exposes the model to a broad range of error types, such as off-by-one mistakes in modular arithmetic, incorrect counting, and wrong base cases, without prescribing alternative reasoning chains. We do not claim this eliminates forgetting, since ORPO still updates the policy, but it reduces the most direct mechanism by which offline preference optimization could overwrite DrGRPO-induced reasoning patterns.\n\nORPO, or Odds-Ratio Preference Optimization, combines a supervised loss on the preferred response with a log-odds preference loss, and does not require a separate reference model. This made it computationally practical for our setting: a 120B-parameter model trained under 4-bit quantization on a single-GPU budget.\n\n**Incremental training.** Training did not run from scratch to 1,000 steps in a single session. The i500 checkpoint was saved first; training then continued for 200 additional steps to produce i700, with checkpoints saved every 100 steps. This incremental approach allowed public leaderboard evaluation between checkpoints and avoided committing to a full training run before seeing LB signal.\n\n| Hyperparameter | Value |\n|----------------|-------|\n| LR | 5 × 10⁻⁶ |\n| β (ORPO) | 0.001 |\n| Max prompt length | 2,048 tokens |\n| Max total length | 4,224 tokens |\n| Batch size | 1 (single sample per step) |\n| Total steps | 1,000 |\n| LoRA rank / α / dropout | 4 / 8 / 0.05 |\n| Optimiser | Paged AdamW 8-bit |\n| Hardware | RTX PRO 6000 Blackwell (96 GB VRAM), local WSL2 |\n\n\n### 2.9 Checkpoint Selection via Hard-50 Elimination (Step 7, continued)\n\nWe evaluated five model ORPO variants on the 50-problem hard-50 benchmark (ZFTurbo), each run 2–3 times to account for stochastic generation. Checkpoint selection followed a three-stage elimination process:\n\n| Model | Hard-50 mean | Runtime mean | Decision |\n|-------|-------------|-------------|---------|\n| i500 | 31.5 | 194.5 min | ① Eliminated — public LB score lower than i700 |\n| v23032026 (1.2k) | 25.3 | 186.7 min | ② Eliminated — slowest remaining runtime |\n| v25032026 (1.4k) | 28.0 | 182.0 min | ② Eliminated — slower runtime |\n| i1k-gold | 28.5 | 172.5 min | ② Eliminated — public LB lower than i700 |\n| **i700** | **27.0** | **175.0 min** | **Selected — best public LB + acceptable runtime** |\n\n**Elimination rationale (two criteria).**\n\n1. **i500 eliminated first (criterion 1: public LB score).** Despite achieving the highest hard-50 mean (31.5 / 50), i500's public LB score on the 50 evaluated problems was lower than i700's. The public LB — scored on a fixed 50-problem subset — was the primary selection signal. i500's hard-50 runtime (194.5 min for 50 problems) also left the least headroom against the 5-hour Kaggle GPU notebook limit.\n\n2. **Remaining candidates eliminated (criterion 2: runtime).** Among the remaining checkpoints, v23032026 (186.7 min) and v25032026 (182.0 min) were eliminated on runtime; i1k-gold (172.5 min) was eliminated on public LB score, which was lower than i700's despite its faster runtime. i700 showed the best combination of public LB performance and acceptable runtime.\n\n3. **i700 selected.** At a mean hard-50 runtime of 175 min for 50 problems (~3.6 min/problem), the projected runtime for the full 100-problem notebook (~357 min on local hardware); actual Kaggle H100 submissions completed within the 5-hour GPU limit. i700 achieved the best public LB across 45 final submissions (mean 38.3, max 42).\n\n> **Figure 7** — Checkpoint evaluation: hard-50 scores and runtimes by model variant.\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2Fb1849543e1835c2217a8a8c226d9bb1b%2Fchart10_hard50_ablation.png?generation=1776604596757256&alt=media)\n> Left panel: mean scores with individual run dots; right panel: mean runtimes.\n\n### 2.10 MXFP4 Baking (Step 8)\n\nThe final ORPO LoRA adapter (i700) was merged into the base model weights using Unsloth's `save_pretrained_merged` with `save_method=\"mxfp4\"` on a node with 8× H100 80 GB GPUs (Kaggle notebook).\n\n| Format | Size |\n|--------|------|\n| Naïve BF16 merge | ~218 GB |\n| MXFP4 baked merge | ~61 GB |\n\nThe ~61 GB checkpoint fits within the single H100 80 GB memory budget with headroom for the vLLM KV cache and Harmony tool interface.\n\n### 2.11 Kaggle Inference (Step 9–10)\nThe inference notebook is based on v8/23 of @andreasbis's public notebook (scriptVersionId=290379281). Differences are limited to: model path (MXFP4 i700 checkpoint), `notebook_limit` increased from 17,400 to 17,520, and a `use_rag` flag set to 0 in the final submission.\nThe merged MXFP4 checkpoint was uploaded as a Kaggle model and served via vLLM inside the submission notebook. Inference used Harmony-encoded prompts with tool-augmented self-consistency across multiple generation attempts. Each problem is allocated a timeout equal to the largest of: 300 s (minimum per problem), 900 s (hard-problem cap), or the remaining time budget after reserving 300 s for every unsolved problem — ensuring accumulated surplus is fully utilised. \nFor answer selection, majority voting is applied across up to 8 attempts with early stopping once 4 solutions agree. Otherwise, entropy weighing is performed to determine the rank — the generation the model was most confident about. Both the dynamic timeout strategy and majority voting with entropy weighing are widely used in public AIMO notebooks. \n\n| Setting | Value |\n|---------|-------|\n| Serving framework | vLLM |\n| Hardware | 1× H100 80 GB |\n| Tensor-parallel size | 1 |\n| Prompt encoding | Harmony multi-turn |\n| RAG at inference | None (despite notebook name `aimo_rag_baked_c283`) |\n\nThe notebook name refers to the 283-certificate training set used for DrGRPO, not to retrieval augmentation at inference time. RAG was dropped in the final configuration after submissions with RAG (mean 37.2–37.9) consistently underperformed plain inference with the baked i700 model (mean 38.3).\n\n## 3. Results\n\n### 3.1 Leaderboard Progression\n\nWe made **85 total submissions** across seven distinct configuration phases (6 failed attempts excluded from analysis). The table and chart below show the progression from baseline exploration to the final baked i700 configuration.\n\n| Phase | Valid submissions | Mean LB | Max LB | Notes |\n|-------|-----------------|---------|--------|-------|\n| Initial exploration | 5 | 36.6 | 40 | Hyperparameter sensitivity probing |\n| LoRA adapter (not baked) | 5 | 34.4 | 37 | Slower inference hurt utilisation |\n| RAG only (base model) | 4 | 37.5 | 41 | Retrieved solutions at inference |\n| LoRA adapter + RAG | 5 | 36.4 | 38 | DrGRPO adapter loaded separately |\n| Baked LoRA + RAG | 10 | 37.9 | 40 | Seed DrGRPO + DPO iterations merged |\n| DPO baked + RAG | 5 | 37.2 | 40 | |\n| **DPO baked, no RAG (final)** | **45** | **38.3** | **42** | MXFP4 i700, best configuration |\n\n**Key findings from the progression:**\n- Dropping RAG at inference increased the mean from 37.2 to 38.3 (+1.1 pts). Retrieved solutions introduced conflicting answer candidates at inference time, degrading self-consistency.\n- Baking the LoRA adapter into MXFP4 (vs loading it separately at inference) eliminated adapter loading overhead, recovering ~3 pts versus the unbaked adapter configuration (34.4 → 37.9).\n- The final 45-submission DPO phase showed a mean of 38.3 with σ=1.67, indicating stable inference with occasional higher runs explaining the max of 42.\n\n> **Figure 8** — Phase-by-phase leaderboard progression (left: scatter over submission index; right: per-phase mean ± std).\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2F374337aefcc4871a8dbfa4914ab1c13c%2Fchart9_submission_scores.png?generation=1776603617928091&alt=media)\n> **Figure 9** — Score distribution of the 35 final DPO v4 submissions (mean=38.57, σ=1.67,\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F150708%2F0bba4be2216d02524df805adb93a2d6c%2Fchart11_v4_score_distribution.png?generation=1776603640035900&alt=media)\n\n\n### 3.2 Ablation: Checkpoint Selection\n\nThe hard-50 benchmark (Figure 7) revealed a counterintuitive result: i500 scored highest on hard-50 (mean 31.5/50) but worst on the public leaderboard, suggesting it may have overfit to the local benchmark distribution. Since the private LB score is the mean of 2 runs, we used public LB performance rather than hard-50 as the primary signal for checkpoint selection. At decision time the evidence was: \n \n - i500: 2 runs → 40, 35 (μ=37.5) \n - i700: 2 runs → 38, 39 (μ=38.5); later confirmed with 5 runs → 38, 39, 40, 40, 40 (μ=39.4) \n \n i700 was the clear choice on that basis. We also noticed ambiguities in several hard-50 problems that reinforced this decision — for example, N01418 implicitly assumes non-negative integers (unstated), and N01455's answer changes depending on whether (a,b) ranges over positive or all integers, exactly the kind of under-specification that IMO problems avoid. This made us more confident that the public LB distribution is closer to the private set than hard-50 is.\n\n## 4. Related Work\n\n**Tool-augmented reasoning.** Our inference pipeline builds on the Harmony tool-use framework, which interleaves Python execution with language model reasoning. Related approaches include PAL (Gao et al., 2023) and ToRA (Gou et al., 2024), which similarly augment LLMs with\nprogram synthesis for mathematical reasoning.\n\n**Reinforcement learning for reasoning.** DrGRPO extends Group Relative Policy Optimisation (DrGRPO, Shao et al., 2024) with per-group advantage normalisation to reduce reward variance. Our two-phase reward schedule is closest in spirit to staged RLHF curricula (Ouyang et al., 2022),\nwhere format correctness is established before optimising for task performance.\n\n**Preference optimisation.** ORPO (Hong et al., 2024) eliminates the need for a reference model by incorporating the preference loss directly into the supervised fine-tuning objective. We apply it in answer-only mode — a design not present in the original paper — to avoid trace-level reward\nhacking while retaining the benefit of preference calibration.\n\n**AIMO precedent.** The AIMO2 winners (arXiv:2504.16891) used a similar pipeline combining SFT, RL, and inference-time scaling. Our work differs primarily in (i) the explicit difficulty stratification via K-Means before data selection, (ii) the answer-only ORPO format, and (iii)\nMXFP4 quantisation as the deployment path.\n\n---\n\n## 5. Reproducibility\n\nAll scripts, the Kaggle inference notebook, and `requirements.txt` are available at: \n[github.com/codesrepo/aimo3-training-pipeline](https://github.com/codesrepo/aimo3-training-pipeline) \n\n\n[link to kaggle notebook](https://www.kaggle.com/code/outliar/aimo-rag-baked-c283?scriptVersionId=300694970) \n\n\n[link to HF model](https://huggingface.co/RecursiveRouge/gpt-oss-120b-merged-mxfp4-i700/tree/main) \n\n[link to google colab notebook](https://colab.research.google.com/github/codesrepo/aimo3-training-pipeline/blob/master/google_colab_notebook/outliar_aimo_inference.ipynb) \n\n > *Note — Google Colab H100 Amendments* \n *The following changes were made to run this notebook on Google Colab H100. They do not impact model accuracy but may affect runtime:*\n \n >os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True' # needed for Google Colab H100 \n\n >batch_size = 32 # default: 256 \n\n >gpu_memory_utilization = 0.90 # default: 0.96 \n\n\n**To reproduce the full pipeline:**\n\n```bash\n# Stage 0: baseline scoring\npython data_selection.py\n\n# Stage 1: clustering + SFT\npython cluster_high_mismatch.py\npython convert_to_harmony.py\npython sft_lora.py\n\n# Stage 2: oracle traces\npython oracle_traces_aimo.py\n\n# Stage 3: DrGRPO\npython aimo_drgrpo_lora_r4_copy.py # submission version\n# : python grpo_lora.py # for reference\n\n# Stage 4: evaluation\npython run_evaluation.py\n\n# Stage 5–6: community data + preference pairs\npython get_kaggle_datasets.py\npython data_positive_negative_math.py\npython preprocess_training_samples_multi.py\n\n# Stage 7: ORPO\npython train_dpo.py\n\n# Stage 8: bake (requires 8× H100)\npython merge_lora_v2.py\n```\n\n**Hardware summary:**\n\n| Stage | Hardware |\n|-------|----------|\n| Stages 0–8 (data + training) | RTX PRO 6000 Blackwell (96 GB VRAM), local WSL2 |\n| Stage 8b (MXFP4 baking) | 8× H100 80 GB (public cloud) |\n| Inference (Kaggle submission) | 1× H100 80 GB, tensor-parallel-size=1 |\n\n**Key package versions** (see `requirements.txt`):\n`unsloth`, `trl`, `peft`, `transformers>=4.46`, `vllm`, `sentence-transformers`,\n`scikit-learn`, `datasets`\n\n---\n\n## 6. Conclusion\n\nWe described a four-phase adaptation pipeline for competition mathematics that sequentially addresses the distinct failure modes of a 120B base model: format compliance (SFT), hard-problem reasoning (DrGRPO), answer distribution alignment (ORPO), and memory-efficient deployment (MXFP4 baking). Key design choices — answer-only preference pairs, two-phase DrGRPO reward, and cluster-stratified data selection — were each motivated by specific\nempirical observations from leaderboard scores and internal evaluation.\n\nThe final baked i700 checkpoint achieved a best public score of **42** with a 45-submission mean of 38.3 (σ=1.67), ranking **1,799 / 4,138**. The most impactful single decision was dropping RAG at inference (+1.1 pts mean), while the most impactful training decision was the DrGRPO phase (+24.3 pp on hard training problems).\n\n\n**Limitations and future work.** The hard-50 benchmark alone was insufficient to distinguish among checkpoint candidates — public LB score was the primary signal, with runtime as secondary. A larger, competition-representative local benchmark would reduce dependence on LB submissions. A larger, competition-representative local benchmark would reduce dependence on LB submissions for checkpoint selection. The DrGRPO phase was limited to 30 total steps by runtime constraints (≤30 min/iter on H100); extending to more steps or using a larger oracle trace corpus may yield further gains.\n\n\n\n## Acknowledgements\n\n- **Yi-Chia Chen** ([ycchen](https://www.kaggle.com/threerabbits)) —\n [CrystalMath-Preview](https://huggingface.co/datasets/ycchen/Crystal-Math-Preview):\n 2,129 verified, high-difficulty contest math problems for RLVR training. 634 problems from this dataset contributed to our ORPO training pool after filtering.\n\n- **ZFTurbo** ([zfturbo](https://www.kaggle.com/zfturbo)) —\n [hard-math-problems-for-aimo-3](https://www.kaggle.com/datasets/zfturbo/hard-math-problems-for-aimo-3): the 50-problem hard-50 benchmark used throughout for checkpoint evaluation.\n\n---\n\n## Figure Reference\n\n| Figure | Description | File |\n|--------|-------------|------|\n| 1 | Full 10-stage training & inference pipeline | `charts/pipeline_diagram.png` |\n| 2 | K-Means elbow plot (k=15 selection) | `imo/openmath_data/kmeans_elbow_plot.png` |\n| 3 | 2D PCA cluster distribution coloured by mismatch rate | `imo/openmath_data/cluster_pca_distribution.png` |\n| 4 | Distribution of per-cluster mismatch rates, threshold=0.59 | `imo/openmath_data/mismatch_rate_distribution.png` |\n| 5 | DrGRPO two-phase reward structure | `charts/chart6_drgrpo_phases.png` |\n| 6 | Community dataset composition pie (11 sources, 2,664 problems) | `charts/chart4_community_datasets.png` |\n| 7 | Checkpoint hard-50 scores and runtimes (model selection) | `charts/chart10_hard50_ablation.png` |\n| 8 | Phase-by-phase leaderboard progression | `charts/chart9_submission_scores.png` |\n| 9 | Final submission (v4) score distribution | `charts/chart11_v4_score_distribution.png` |\n\n*9 figures total. Figures 2–4 are generated by `cluster_high_mismatch.py`; Figures 5–9 by\n`writeup/generate_charts.py` and `writeup/generate_pipeline_diagram.py`.*\n", "votes": 13 }, { "id": 3458597, "postDate": "2026-05-16T08:31:28.810Z", "content": "

Following up on the questions from @philipvonderlind and @flamingice2801 (and thanks for the very thorough response), I only have one further question:

\n
    \n
  1. K-Means cluster validation: Your data selection uses K-Means on OpenMathReasoning to identify Cluster 3 as the most informative training distribution. Was the cluster choice validated on held-out problems (e.g. by checking that Cluster 3 problems also have high mismatch rates on a different problem reference set), or could the selection have overfit to the development sample you used to define it?
    2. \n
\n

Look forward to reading your response on this.

", "rawMarkdown": "Following up on the questions from @philipvonderlind and @flamingice2801 (and thanks for the very thorough response), I only have one further question:\n\n1. **K-Means cluster validation:** Your data selection uses K-Means on OpenMathReasoning to identify Cluster 3 as the most informative training distribution. Was the cluster choice validated on held-out problems (e.g. by checking that Cluster 3 problems also have high mismatch rates on a different problem reference set), or could the selection have overfit to the development sample you used to define it?\n\nLook forward to reading your response on this.", "replies": [ { "id": 3458636, "postDate": "2026-05-16T10:12:11.197Z", "content": "

Thank you for the question @sambealing. We did not validate the Cluster 3 choice on a separate held-out reference set, so we agree that the cluster-selection step may have some dependence on the development sample used to identify it.

\n

As a post-selection and post-training observation from partial logging, the final model appeared to solve BYTE66 from the hard-50 set more frequently than the baseline: 2/3 successful runs for the final selected model versus 1/6 successful runs for the baseline. However, given the high inherent variance of these runs, we do not treat this as conclusive validation, and chance cannot be ruled out.

\n

The motivation for selecting Cluster 3 was practical rather than claiming it was globally optimal. Our goal was to concentrate the limited LoRA training budget on a narrow semantic region where the base model showed the largest mismatch or weakness, rather than spread updates thinly across the full OpenMathReasoning distribution. Since we were training only LoRA adapters rather than the full model, we expected the update to be localized and less likely to broadly disrupt existing capabilities.

\n

The intended best-case outcome was that the model would acquire a useful problem-solving pattern for a difficult region while retaining most of its existing behavior. The worst-case risk was overfitting to that narrow region, but even in that case the selection could still be useful if it helped solve a previously inaccessible high-complexity hard problem. We therefore treated the clustering step as a targeted, compute-constrained heuristic for data selection, not as a fully validated claim about the universally most informative training distribution.

", "rawMarkdown": "Thank you for the question @sambealing. We did not validate the Cluster 3 choice on a separate held-out reference set, so we agree that the cluster-selection step may have some dependence on the development sample used to identify it. \n\nAs a post-selection and post-training observation from partial logging, the final model appeared to solve BYTE66 from the hard-50 set more frequently than the baseline: 2/3 successful runs for the final selected model versus 1/6 successful runs for the baseline. However, given the high inherent variance of these runs, we do not treat this as conclusive validation, and chance cannot be ruled out.\n\nThe motivation for selecting Cluster 3 was practical rather than claiming it was globally optimal. Our goal was to concentrate the limited LoRA training budget on a narrow semantic region where the base model showed the largest mismatch or weakness, rather than spread updates thinly across the full OpenMathReasoning distribution. Since we were training only LoRA adapters rather than the full model, we expected the update to be localized and less likely to broadly disrupt existing capabilities.\n\nThe intended best-case outcome was that the model would acquire a useful problem-solving pattern for a difficult region while retaining most of its existing behavior. The worst-case risk was overfitting to that narrow region, but even in that case the selection could still be useful if it helped solve a previously inaccessible high-complexity hard problem. We therefore treated the clustering step as a targeted, compute-constrained heuristic for data selection, not as a fully validated claim about the universally most informative training distribution." } ] }, { "id": 3456815, "postDate": "2026-05-12T16:16:08.390Z", "content": "

Hi, thank you for your submission to the Writeup Prize. @sambealing and I will post some questions below that help us in assessing your submission more closely, as described in the updated discussion post.

\n

You already have a very extensive submission. I especially like the usage of plots to convey the architecture and experiments to the readers. I have a few questions regarding some stages in your approach:

\n\n

Looking forward to your response :)

", "rawMarkdown": "Hi, thank you for your submission to the Writeup Prize. @sambealing and I will post some questions below that help us in assessing your submission more closely, as described in the [updated discussion post.](https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-3/discussion/689703)\n\nYou already have a very extensive submission. I especially like the usage of plots to convey the architecture and experiments to the readers. I have a few questions regarding some stages in your approach:\n\n* **ORPO after GRPO**: Usually, it is more common to do ORPO (offline) before GRPO (online), due to the risk of possibly unlearning the dynamic patterns discovered by GRPO with the static dataset used by ORPO. Was there a specific reason to do it in this order, and did you mitigate those risks?\n* **GRPO requirements**: You write 'GRPO requires verified correct trajectories to serve as positive training examples'.\nHowever, what you actually need is a *reward model*, not a trace. Because GRPO is an online/on-policy RL method, the model samples outputs during training and updates the policy based on their relative rewards. Could you clarify why fully verified traces are necessary here, rather than just a reward function over generated outputs?\n\n* **Inference**: I could not find the exact inference setup in the Kaggle writeup (maybe it's in the appended PDF) that got you your final results. Furthermore, did you test your model integrated into the [public notebook](https://www.kaggle.com/code/andreasbis/aimo-3-gpt-oss-120b-with-tools) by @andreasbis to test if it improves results or stability?\n\nLooking forward to your response :)\n", "replies": [ { "id": 3457342, "postDate": "2026-05-13T15:36:15.310Z", "content": "

Q1: ORPO after GRPO ordering\nThank you for the insightful question. We agree that the more common ordering is ORPO before GRPO, since applying offline preference optimization after online RL can risk moving the policy away from dynamic behaviors discovered during RL. In our case, however, the DrGRPO stage was intentionally small and compute-constrained: it was run for only 30 steps on 30 data points, so we treated it less as a fully converged RL-trained model and more as a reward-aligned initialization. The subsequent ORPO stage used a much larger preference set of 1,415 pairs, providing a broader and more diverse signal.

\n

The motivation for this ordering was incremental: DrGRPO helped move the model toward better candidate behavior, while ORPO was then used to improve final-answer ranking among candidates, especially in cases where the correct answer was already present but not ranked highest. We mitigated the risk of overwriting DrGRPO-induced reasoning behavior by using answer-only preference pairs rather than full solution traces. As a result, ORPO primarily calibrated the output distribution over final answers instead of explicitly supervising alternative reasoning paths. This does not eliminate forgetting risk, but it reduces the most direct way in which the offline stage could overwrite reasoning patterns discovered during DrGRPO. We further mitigated this risk by using a conservative learning rate of 5×10−6, so the ORPO stage acted more as a mild preference-calibration step than a large distribution-shifting update.

\n

Q2: GRPO requires verified correct trajectories\nThank you for pointing this out. You are right that the phrasing in the writeup is imprecise. GRPO/DrGRPO is an on-policy RL method: the model samples outputs during training, and a reward function scores those outputs relative to each other within each group. Thus, fully verified pre-collected trajectories are not required in principle.

\n

In our setup, the verified oracle traces collected during training were not used as supervised positive trajectories for GRPO, nor were they fed to the trainer for imitation learning. Instead, they served two limited purposes. First, they defined a corpus of hard problems with confirmed integer answers. Second, the confirmed answers supplied the expected answer string used by the exact-match reward component. During DrGRPO, the policy still generated its own rollouts online, and updates were based on rewards assigned to those generated outputs.

\n

Separately, the oracle traces were also used to generate certificates for a RAG component during development. However, RAG was not activated in the final submission, so these certificates did not affect the final DrGRPO training or inference pipeline. We will revise the writeup to avoid implying that GRPO requires verified positive traces.

\n

Q3: Inference setup\nThank you for the question. Our inference setup was based on v8/23 of @andreasbis' public notebook, aimo-3-gpt-oss-120b-with-tools (scriptVersionId=290379281), which we used as the baseline notebook. We did test our fine-tuned checkpoint integrated directly into this public notebook framework.

\n

Relative to the upstream notebook, the final submitted configuration made only minimal changes:

\n

Model path — replaced the base gpt-oss-120b checkpoint with our fine-tuned MXFP4-baked i700 checkpoint, outliar/gpt-oss-120b-merged-mxfp4-i700.

\n

notebook_limit — increased slightly from 17,400 to 17,520, only to marginally extend the usable runtime budget.

\n

use_rag flag and cert_path — added for RAG ablation runs, but use_rag=0 in the final submitted configuration.

\n

All other inference parameters were kept identical to the upstream notebook: attempts=8, early_stop=4, temperature=1.0, min_p=0.02, context_tokens=65536, and gpu_memory_utilization=0.96. Therefore, the final inference strategy, including dynamic timeout, majority voting, and entropy-weighted fallback, was directly inherited from @andreasbis' implementation.

\n

In this sense, yes: the model was evaluated inside the public notebook framework, and the final submitted run differs from the public baseline primarily by the fine-tuned model weights. We will clarify this in the writeup and include the exact inference configuration to make the comparison reproducible.

\n

kaggle notebook used for final submission link

", "rawMarkdown": "**Q1: ORPO after GRPO ordering**\nThank you for the insightful question. We agree that the more common ordering is ORPO before GRPO, since applying offline preference optimization after online RL can risk moving the policy away from dynamic behaviors discovered during RL. In our case, however, the DrGRPO stage was intentionally small and compute-constrained: it was run for only 30 steps on 30 data points, so we treated it less as a fully converged RL-trained model and more as a reward-aligned initialization. The subsequent ORPO stage used a much larger preference set of 1,415 pairs, providing a broader and more diverse signal.\n\nThe motivation for this ordering was incremental: DrGRPO helped move the model toward better candidate behavior, while ORPO was then used to improve final-answer ranking among candidates, especially in cases where the correct answer was already present but not ranked highest. We mitigated the risk of overwriting DrGRPO-induced reasoning behavior by using answer-only preference pairs rather than full solution traces. As a result, ORPO primarily calibrated the output distribution over final answers instead of explicitly supervising alternative reasoning paths. This does not eliminate forgetting risk, but it reduces the most direct way in which the offline stage could overwrite reasoning patterns discovered during DrGRPO. We further mitigated this risk by using a conservative learning rate of 5×10−6, so the ORPO stage acted more as a mild preference-calibration step than a large distribution-shifting update.\n\n**Q2: GRPO requires verified correct trajectories**\nThank you for pointing this out. You are right that the phrasing in the writeup is imprecise. GRPO/DrGRPO is an on-policy RL method: the model samples outputs during training, and a reward function scores those outputs relative to each other within each group. Thus, fully verified pre-collected trajectories are not required in principle.\n\nIn our setup, the verified oracle traces collected during training were not used as supervised positive trajectories for GRPO, nor were they fed to the trainer for imitation learning. Instead, they served two limited purposes. First, they defined a corpus of hard problems with confirmed integer answers. Second, the confirmed answers supplied the expected answer string used by the exact-match reward component. During DrGRPO, the policy still generated its own rollouts online, and updates were based on rewards assigned to those generated outputs.\n\nSeparately, the oracle traces were also used to generate certificates for a RAG component during development. However, RAG was not activated in the final submission, so these certificates did not affect the final DrGRPO training or inference pipeline. We will revise the writeup to avoid implying that GRPO requires verified positive traces.\n\n**Q3: Inference setup**\nThank you for the question. Our inference setup was based on v8/23 of @andreasbis' public [notebook](https://www.kaggle.com/code/andreasbis/aimo-3-gpt-oss-120b-with-tools?scriptVersionId=290379281), aimo-3-gpt-oss-120b-with-tools (scriptVersionId=290379281), which we used as the baseline notebook. We did test our fine-tuned checkpoint integrated directly into this public notebook framework.\n\nRelative to the upstream notebook, the final submitted configuration made only minimal changes:\n\nModel path — replaced the base gpt-oss-120b checkpoint with our fine-tuned MXFP4-baked i700 checkpoint, outliar/gpt-oss-120b-merged-mxfp4-i700.\n\nnotebook_limit — increased slightly from 17,400 to 17,520, only to marginally extend the usable runtime budget.\n\nuse_rag flag and cert_path — added for RAG ablation runs, but use_rag=0 in the final submitted configuration.\n\nAll other inference parameters were kept identical to the upstream notebook: attempts=8, early_stop=4, temperature=1.0, min_p=0.02, context_tokens=65536, and gpu_memory_utilization=0.96. Therefore, the final inference strategy, including dynamic timeout, majority voting, and entropy-weighted fallback, was directly inherited from @andreasbis' implementation.\n\nIn this sense, yes: the model was evaluated inside the public notebook framework, and the final submitted run differs from the public baseline primarily by the fine-tuned model weights. We will clarify this in the writeup and include the exact inference configuration to make the comparison reproducible. \n\nkaggle notebook used for final submission [link ](https://www.kaggle.com/code/outliar/aimo-rag-baked-c283?scriptVersionId=300694970)", "replies": [ { "id": 3457351, "postDate": "2026-05-13T15:59:15.413Z", "content": "

Edited on 13th May 2026 for the rebuttal phase:

\n

DrGRPO terminology: Standardized the naming of this stage as “DrGRPO” throughout the writeup. Earlier drafts used inconsistent variants such as “DGRPO,” “Dr.GRPO,” and “GRPO.”

\n

GRPO phrasing correction: Corrected the description of DrGRPO training (section 2.4). DrGRPO is an on-policy RL method and does not require pre-collected correct trajectories as supervised training examples. In our setup, the oracle traces were used only to obtain verified ground-truth integer answers for the exact-match reward component.

\n

ORPO ordering rationale: Added a clearer explanation for applying ORPO after DrGRPO (section 2.8). Because the DrGRPO stage was compute-constrained and ran for only 30 steps, we treated it as a directionally improved but under-optimized checkpoint. ORPO was then applied as a lightweight answer re-ranking stage using 1,415 answer-only preference pairs.

\n

Inference setup: Clarified the exact final inference configuration (section 2.11). The submission notebook was based on @andreasbis' public notebook version with scriptVersionId=290379281. The final setup differed from the upstream notebook in only three ways: the model path was changed to our fine-tuned checkpoint, notebook_limit was increased from 17,400 to 17,520, and RAG was disabled in the submitted configuration with use_rag=0.

", "rawMarkdown": "**Edited on 13th May 2026 for the rebuttal phase:**\n\n**DrGRPO terminology:** Standardized the naming of this stage as “DrGRPO” throughout the writeup. Earlier drafts used inconsistent variants such as “DGRPO,” “Dr.GRPO,” and “GRPO.”\n\n**GRPO phrasing correction:** Corrected the description of DrGRPO training (*section 2.4*). DrGRPO is an on-policy RL method and does not require pre-collected correct trajectories as supervised training examples. In our setup, the oracle traces were used only to obtain verified ground-truth integer answers for the exact-match reward component.\n\n**ORPO ordering rationale:** Added a clearer explanation for applying ORPO after DrGRPO (*section 2.8*). Because the DrGRPO stage was compute-constrained and ran for only 30 steps, we treated it as a directionally improved but under-optimized checkpoint. ORPO was then applied as a lightweight answer re-ranking stage using 1,415 answer-only preference pairs.\n\n**Inference setup:** Clarified the exact final inference configuration (*section 2.11*). The submission notebook was based on @andreasbis' public notebook version with scriptVersionId=290379281. The final setup differed from the upstream notebook in only three ways: the model path was changed to our fine-tuned checkpoint, notebook_limit was increased from 17,400 to 17,520, and RAG was disabled in the submitted configuration with use_rag=0." } ] } ] }, { "id": 3445573, "postDate": "2026-04-20T05:17:02.997Z", "content": "

Great writeup @outliar - the 10-stage pipeline diagram is one of the clearest I've seen in this competition, and the fact that you actually got MXFP4 baking to work with save_method=\"mxfp4\" is a genuine contribution. We attempted GRPO fine-tuning as well and lost ~$300 to the merged_4bit_forced → BitsAndBytes NF4 deployment bug that you avoided. So first: thank you for documenting the correct path.

\n

Two things jumped out from your results that I'd love your perspective on:

\n
    \n

  1. Checkpoint selection with n=2. Your Figure 7 shows i700 selected despite having the lowest hard50 mean (27.0) of all 5 checkpoints. Looking at the individual run dots, each checkpoint was evaluated only 2-3 times on hard50. With that sample size, i500's mean of 31.5 and i700's mean of 27.0 could easily overlap - the 4.5-point gap is within a single run's variance. Do you think i500 might actually generalize better to the private set, given that the private distribution is closer to hard50's difficulty than to the public LB's fixed 50-problem subset?

    2. \n

  2. Variance comparison. Your Figure 9 shows n=35 final submissions with μ=38.57, σ=1.67 - very well documented. We ran the unmodified base model (no fine-tuning, just entropy-weighted voting) 9 times and got μ=40.78, σ=1.20. The fine-tuned model has both higher variance and lower mean than the base model with a better voting function. Does this suggest that ORPO might have flattened the model's confidence across competing answers rather than sharpening it? Or do you attribute the gap primarily to the MXFP4 quantization step?

    3. \n
\n

Really appreciate the open-sourced pipeline - the difficulty-stratified data selection via K-Means is an idea I wish I'd tried.

", "rawMarkdown": "Great writeup @outliar - the 10-stage pipeline diagram is one of the clearest I've seen in this competition, and the fact that you actually got MXFP4 baking to work with `save_method=\"mxfp4\"` is a genuine contribution. We attempted GRPO fine-tuning as well and lost ~$300 to the `merged_4bit_forced` → BitsAndBytes NF4 deployment bug that you avoided. So first: thank you for documenting the correct path.\n\nTwo things jumped out from your results that I'd love your perspective on:\n\n1. **Checkpoint selection with n=2.** Your Figure 7 shows i700 selected despite having the lowest hard50 mean (27.0) of all 5 checkpoints. Looking at the individual run dots, each checkpoint was evaluated only 2-3 times on hard50. With that sample size, i500's mean of 31.5 and i700's mean of 27.0 could easily overlap - the 4.5-point gap is within a single run's variance. Do you think i500 might actually generalize better to the private set, given that the private distribution is closer to hard50's difficulty than to the public LB's fixed 50-problem subset?\n\n2. **Variance comparison.** Your Figure 9 shows n=35 final submissions with μ=38.57, σ=1.67 - very well documented. We ran the unmodified base model (no fine-tuning, just entropy-weighted voting) 9 times and got μ=40.78, σ=1.20. The fine-tuned model has both higher variance and lower mean than the base model with a better voting function. Does this suggest that ORPO might have flattened the model's confidence across competing answers rather than sharpening it? Or do you attribute the gap primarily to the MXFP4 quantization step?\n\nReally appreciate the open-sourced pipeline - the difficulty-stratified data selection via K-Means is an idea I wish I'd tried.", "replies": [ { "id": 3445762, "postDate": "2026-04-20T13:30:43.743Z", "content": "

Thank you for your thorough review and kind words your report was a direct inspiration for mine.

\n

1.Checkpoint selection

\n

Since the final LB score is the mean of 2 runs, I used this as a guiding principle throughout, including checkpoint selection, where I relied on public LB performance rather than hard-50. At decision time the evidence was:

\n\n

On hard-50 reliability, I also noticed ambiguities in the benchmark problems themselves. For example, N01418 implicitly assumes non-negative integers (unstated), and N01455's answer changes depending on whether (a,b) ranges over positive or all integers, exactly the kind of under-specification IMO problems avoid. This made me more confident that the public LB distribution should be more closer to the private set than hard-50 is. That said, if hard-50 does turn out to be more representative, the existing evidence would favour i500.

\n

2.Variance comparison

\n

Looking at my own baseline (initial exploration, no fine-tuning): scores 36, 36, 40, 35, 36 give μ=36.6, σ=1.74. The final ORPO submissions (μ=38.57, σ=1.67) represent roughly 2 additional correct problems, consistent with ORPO sharpening the model's early attention toward the correct answer on a subset of problems, rather than flattening confidence generally.

\n

This aligns with @natnitarach finding in section 5 that with μ̂=39.3, σ̂=1.7 over 21 baseline runs, each submission is essentially a lottery ticket, and their verifier-aware re-ranking (which achieves better answer selection) outperforms the fine-tuned approach. The gap you observe between your base model (μ=40.78, σ=1.20) and my fine-tuned one is worth investigating further MXFP4 quantization could be a plausible contributor and I cannot fully rule it out.

", "rawMarkdown": "Thank you for your thorough review and kind words your report was a direct inspiration for mine.\n \n **1.Checkpoint selection** \n \nSince the final LB score is the mean of 2 runs, I used this as a guiding principle throughout, including checkpoint selection, where I relied on public LB performance rather than hard-50. At decision time the evidence was: \n - i500: 2 runs → 40, 35 (μ=37.5) \n - i700: 2 runs → 38, 39 (μ=38.5); later confirmed with 5 runs → 38, 39, 40, 40, 40 (μ=39.4) \n ![](https://www.googleapis.com/download/storage/v1/b/kaggle-forum-message-attachments/o/inbox%2F150708%2F4f07fb97f42f1c0cd4404439cd1f6de7%2Fearly_evidence.jpg?generation=1776691366783193&alt=media) \n i700 was the clear choice on that basis. \n \nOn hard-50 reliability, I also noticed ambiguities in the benchmark problems themselves. For example, N01418 implicitly assumes non-negative integers (unstated), and N01455's answer changes depending on whether (a,b) ranges over positive or all integers, exactly the kind of under-specification IMO problems avoid. This made me more confident that the public LB distribution should be more closer to the private set than hard-50 is. That said, if hard-50 does turn out to be more representative, the existing evidence would favour i500. \n \n**2.Variance comparison**\n\n Looking at my own baseline (initial exploration, no fine-tuning): scores 36, 36, 40, 35, 36 give μ=36.6, σ=1.74. The final ORPO submissions (μ=38.57, σ=1.67) represent roughly 2 additional correct problems, consistent with ORPO sharpening the model's early attention toward the correct answer on a subset of problems, rather than flattening confidence generally.\n \nThis aligns with @natnitarach [finding in section 5 ](https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-3/writeups/model-capability-dominates) that with μ̂=39.3, σ̂=1.7 over 21 baseline runs, each submission is essentially a lottery ticket, and their verifier-aware re-ranking (which achieves better answer selection) outperforms the fine-tuned approach. The gap you observe between your base model (μ=40.78, σ=1.20) and my fine-tuned one is worth investigating further MXFP4 quantization could be a plausible contributor and I cannot fully rule it out. ", "replies": [ { "id": 3448910, "postDate": "2026-04-27T02:07:37.260Z", "content": "

Thank you for the detailed checkpoint data the expanded i700 runs (38, 39, 40, 40, 40, μ=39.4) make the selection much more defensible than the n=2 I was originally questioning.

\n

The point about hard50 problem ambiguities (N01418 assuming non-negative integers, N01455's range dependency) resonates strongly. We used the same hard50 dataset for offline evaluation and found similar issues - our Branch B multicell scored 27/50 on it, but we couldn't determine how much of the gap vs our leaderboard mean (40.78) was real difficulty vs dataset noise.

\n

Your baseline comparison (μ=36.6 pre-ORPO → μ=38.57 post-ORPO, Δ≈+2) is actually a cleaner measurement of fine-tuning value than most teams produced. The fact that our inference-only approach (μ=40.78) still exceeds the fine-tuned model (μ=38.57) is consistent with both our writeups' findings - but as you note, the MXFP4 quantization step is a confound that neither of us can fully rule out.

", "rawMarkdown": "Thank you for the detailed checkpoint data the expanded i700 runs (38, 39, 40, 40, 40, μ=39.4) make the selection much more defensible than the n=2 I was originally questioning.\n\nThe point about hard50 problem ambiguities (N01418 assuming non-negative integers, N01455's range dependency) resonates strongly. We used the same hard50 dataset for offline evaluation and found similar issues - our Branch B multicell scored 27/50 on it, but we couldn't determine how much of the gap vs our leaderboard mean (40.78) was real difficulty vs dataset noise.\n\nYour baseline comparison (μ=36.6 pre-ORPO → μ=38.57 post-ORPO, Δ≈+2) is actually a cleaner measurement of fine-tuning value than most teams produced. The fact that our inference-only approach (μ=40.78) still exceeds the fine-tuned model (μ=38.57) is consistent with both our writeups' findings - but as you note, the MXFP4 quantization step is a confound that neither of us can fully rule out." }, { "id": 3449038, "postDate": "2026-04-27T11:03:18.390Z", "content": "

Based on these findings, it is highly likely that your private LB score will be at least 1.5 points higher than mine.

", "rawMarkdown": "Based on these findings, it is highly likely that your private LB score will be at least 1.5 points higher than mine." } ] } ] } ], "complete": true, "source": "live", "fetched_at": "2026-10-20T04:21:27Z", "competition_id": 118448, "competition_slug": "ai-mathematical-olympiad-progress-prize-3", "team_id": 15101388, "solution_url": "/competitions/ai-mathematical-olympiad-progress-prize-3/writeups/four-phase-adaptation-of-a-120b-model-for-tool-aug" }