# 1st place solution - NemoSkills We would like to extend our thanks to the hosts from AIMOPrize, XTX Markets and Kaggle for this year's challenge which obviously had a lot of thought and preparation spent on the problems and the Kaggle environment. And a special shout out to Simon Frieder for being so accessible throughout the challenge. We were thrilled with the results - after developing our own TIR models we struggled to get Public LB to match the strong CV results but at the end private LB was more closely aligned. Resources [Paper](https://arxiv.org/abs/2504.16891) [Kaggle Submission script](https://www.kaggle.com/code/darraghdog/nemoskills-qwen2-5-14b-instruct-mix48-redrafter) [Training Dataset - nvidia/OpenMathReasoning](https://huggingface.co/datasets/nvidia/OpenMathReasoning) - 5.5M OpenMath solutions - CoT&TIR [Collection of Math Reasoning models](https://huggingface.co/collections/nvidia/openmathreasoning-68072c0154a5099573d2e730) - 1.5B/7B/14B/32B released - 1.5B beats R1 on AIME (Kaggle upload in-work) [Submitted model](https://huggingface.co/nvidia/OpenMath-Nemotron-14B-Kaggle) [Reproduce results in Nemo-Skills](https://nvidia.github.io/NeMo-Skills/openmathreasoning1/) TL;DR The dataset was built by collecting and filtering 540K math problems from AoPS forums. For our submission, 2.2M million Chain-of-Thought (CoT) responses were generated using models like DeepSeek-R1 with incorrect answers filtered out through verification. We also created a high quality Tool-Integrated Reasoning (TIR) dataset of 15K problems and generations. Training occurred in two phases: first a Qwen2.5-14B model was fine-tuned on CoT data, and second, the smaller TIR dataset was used to further refine, enhancing its problem-solving with structured reasoning. To retain the strength of TIR, along with the shorter responses of CoT for speed, we found mergekit worked well to merge both models into one, retaining the strength of TIR, while shortening the number of output tokens. For inference optimization, we implemented a pipeline combining TensorRT-LLM conversion with FP8 quantization, achieving 1.5x speedup over BF16. ReDrafter, a speculative decoding technique developed by Apple, was used for a further 1.8x speedup. Our serving system features dynamic batching with early stopping and a time-buffering strategy which allocates compute resources efficiently. The solution includes tool-integrated reasoning capabilities through a sandboxed execution environment. Training Data Large-scale long-reasoning dataset We started by collecting an extensive set of mathematical problems from the internet. We developed an LLM-based problem extraction and refinement pipeline to construct a dataset of 540K unique problems. Using this dataset, we then generated 3.2M long reasoning CoT solutions. The process generates multiple solution candidates (up to 32) per problem using models like DeepSeek-R1 and QwQ-32B, with temperature 0.7 and top-p 0.95. Harder problems—identified by Qwen2.5-72B-Math-Instruct's pass-rates—receive more candidates. Incorrect solutions are filtered by verifying answer equivalence with Qwen2.5-32B-Instruct. If no answer is found, the most frequent candidate answer is used. For our submission a 2.2M subset generated by DeepSeek-R1 was used. Tool-Integrated Reasoning For tool integrated reasoning the model prompts code for a calculation, and then the code is executed within a sandbox. We recognised code snippets with the special tokens and <\tool_call> . The code is then appended to the LLM output, between the text ``` and ```output . An example snippet from an output is below. , k must divide , k > So, k can be Therefore, b = k - = So, same result. Therefore, sum \n\nAlternatively, maybe I can a small program bases b > , compute b + b + , the latter divides the former, collect such bases. sum them. Let\s a Python code the :\n\n\n# Initialize a list store bases\nvalid_bases = []\n\n# bases upwards\nfor b range(, ): # Arbitrary upper \n num1 = * b + \n num2 = b + \n num1 % num2 == :\n valid_bases.append(b)\n print(f"Found base: {b}")\n\n# Sum the bases\nsum_bases = sum(valid_bases)\nprint(f"Sum: {sum_bases}")\n\n# sum , take modulo \nif sum_bases > :\n result = sum_bases % \nelse:\n result = sum_bases\n\nprint(f"Final Result: {result}")\n\n```output\nFound base: \nFound base: \nSum: \nFinal Result: \n```\nThe code confirms that the bases are , summing Therefore, the final answer Since less than , modulo operation needed. We developed a method for integrating code execution into long reasoning generations. Our initial attempts to elicit Tool-Integrated Reasoning (TIR) generations from DeepSeek-R1 through simple prompting proved unsuccessful. We hypothesize that these models struggle to deviate from their standard solution format due to extensive training on reasoning tasks and limited exposure to instruction-following. To overcome this challenge, we built a pipeline that starts with a small-scale reasoning finetuning of an instruction-following model, [LIMO](https://huggingface.co/GAIR/LIMO) . By prompting this model to generate long-reasoning TIR solutions followed by aggressive quality filtering, we established an initial dataset suitable for training. Through multiple iterations of training, generation, and filtering, we constructed a 1.7M TIR solution set that was crucial for improving the accuracy of our final models. For final stage TIR training, the 1.7M solutions are aggressively filtered down to 15K samples based on number of code executions and correctness of answers. Model Training For our Kaggle submission we used a slightly different training recipe than detailed in the paper. We found this different recipe resulted in fewer tokens than the models publicly released, and we were out of time to experiment more with reducing tokens on the final models even though they performed better. We first trained a Qwen2.5-14B-Base model with SFT on the 2.2M subset of CoT solutions for 8 epochs. We changed base RoPE to 500k to allow long reasoning. The model is trained with NVIDIA/Nemo-Skills for 8 epochs with 1e-4 LR, AdamW, weight decay 0.01 and 10% linear warmup decaying to a LR of 1e-7. We use batch size of 1024 samples and leverage sequence packing and context parallelization techniques from NVIDIA/NeMo-Aligner which significantly accelerates training on the long-reasoning data. The training lasted 48 hours on 512 H100 (yes, 512!). We could achieve most of the strength of the model with 20% of the compute but we wanted to scale up to see where learning saturates. See figure 3(b) from the linked publication below which shows metrics at different stages of training. Weight averaging is used from different stages for the final weights. ![](assets/99d7608b279b6259f493b8b9.png) This is followed by a light-weight TIR fine-tuning on the 15k TIR samples. We trained the TIR model for 400 steps with a constant learning rate of 1e-5 and use the last checkpoint without averaging. As detailed below, we then merge CoT and TIR checkpoints as it both improves accuracy and speeds up generation by reducing solution length and number of code executions. Evaluation Dataset During the competition we mostly use problems from the 2024 year for American Invitational Mathematics Examinations (AIME) and Harvard-MIT Mathematics Tournaments (HMMT). Later we added problems from 2025. Our final benchmark Comp-Math-24-25 includes 256 problems. Dataset Problems AIME24 30 AIME25 30 HMMT Nov 2024 62 HMMT Feb 2024 68 HMMT Feb 2025 66 Inference Model Merging In this competition, we explored various methods for merging two LLMs with CoT and TIR behaviors. Our primary goal was to effectively combine the distinct strengths of these two fine-tuning stages to enhance the model's performance. We experimented with several merging techniques from mergekit package. Surprisingly, the most effective approach turned out to be a simple linear combination of the two checkpoints: the CoT checkpoint used before TIR fine-tuning and the best TIR checkpoint attained thereafter. This strategy allowed us to control the extent to which each stage influenced the final model's behavior. In the following table, we show accuracy and generation statistics of merged models on Comp-Math-24-25 dataset. length indicates the average number of tokens per solution, while code refers to the average number of code executions per solution. Method maj@16 pass@16 length code CoT 62.9 76.2 11203 - TIR 66.8 80.1 15834 2.73 CoT * 0.3 + TIR * 0.7 69.1 81.3 12489 0.85 Model acceleration We prioritized Int8 weight-only (W8A16) and FP8 quantization, which delivered faster inference than BF16 with minimal accuracy loss. The reduced weight size also freed up memory for larger key-value caches. [ReDrafter](https://developer.nvidia.com/blog/nvidia-tensorrt-llm-now-supports-recurrent-drafting-for-optimizing-llm-inference/) is a speculative decoding technique developed by Apple, and implemented in TensorRT-LLM. We trained a ReDrafter head on a random subset of problems from OpenMathReasoning-1 dataset. Using these problems we generated 100k solutions with the target model. The resulting redrafter generated 3 tokens for every LLM step with 65% acceptance rate delivering a speed up of approximately 1.8X. The TensorRT-LLM repo’s example implementation has an example for a non quantised llama, so we needed some changes to make it work. We tried other drafting methods, but could not get them to work as well. Especially LLM drafting should work and may be better. Quantization Speed tok/sec L4x4 50 question eval time L40x4 Accuracy score % AIME24 Accuracy score % AIME25 bf16 210 2h 82.7 66.7 w8a16 (int8) 315 1h 45m 82.7 66.7 w4a16 (int4) 436 1h 35m 72.7 60.7 f8a16 310 1h 40m 83.3 68.7 f8a16 + Redrafter 554 1h 81.3 71.3 Accuracy score in the table is using maj@12 of the merged model averaged over 5 runs. TensorRT-LLM Inference Pretrained models were converted to TensorRT engines using TensorRT-LLM. TensorRT’s inflight batching boosts throughput by dynamically grouping inference requests, releasing each sample as soon as it completes—reducing latency and optimizing GPU utilization. See below some recent [benchmarks](https://x.com/vllm_project/status/1913513173342392596) by vLLM team. ![](assets/0e96bf0fd0440cf2e010aca6.jpeg) Since samples are processed independently, batches can mix different prompts or seeds seamlessly. TensorRT-LLM includes a number of other optimizations such as custom attention kernels and paged KV caching. ![](assets/b6728ff13779921ab11eace6.png) For each new problem we used different seeds to launch 12 async generations via in-flight batching in TensorRT. The streaming for each sample was monitored for code blocks, stop phrases, max tokens or timeouts. If the LLM generates code, LLM generation stops, and the code block is executed in a sandbox. The sandbox output (or partial error trace) is appended to the LLM and generation continues. The generation continues until another code block. When no other code blocks are encountered, the LLM stops at one of max tokens, timeout time, or stop phrase. We experimented with various sampling parameters but observed minimal differences in the results. However, our winning submission was based on an "almost" greedy search strategy. Despite setting the temperature to 0 and enabling the redrafter_greedy_search parameter, TensorRT-LLM still produced slightly varying outputs within a single batch. We did not investigate this behavior in detail, but speculate that it could be related to rounding or batching issues. Ultimately, we chose this approach because it provided more stable results at small batch sizes and offered a slight improvement in the speed of speculative decoding. Within the kaggle notebook, generations are monitored for completion. When initial answers are the same we early-stop, e.g. if 4 of the first 5 generations are the same we cancel the remaining generations in TensorRT. ![](assets/71b4d3fce772af251a647e8d.png) During monitoring of async generations, we also early stop after completing 10 of the 12 generations to avoid waiting excessively long for any stragglers. A buffering strategy was implemented, allocating 350 seconds per question as the base time limit. If a question completed early, the unused time was added to a shared buffer. The next question could then draw up to 210 extra seconds from this buffer, making a total of 560 seconds. ![](assets/85f1dcab95abd7d2ed9ddac9.png) Using this inference pipeline and deepseek-distill-14b, we scored in the range of 23 to 29 on public leaderboard using different settings. The score variability was concerning and our own trained models helped to stabilise the scores as well as increasing the accuracy. For the final chosen submissions, one 14b CoT model and the above MIX TIR model were chosen. The MIX TIR model scored much better on CV, and a little better on Public LB (public LB scores : 32,33,28). The one 28 scoring submission made us nervous and was perhaps due to a larger maj@ used. In the end, the private LB reflected more our CV results than LB results which was quite satisfying. We had other models as seen which were quite strong on CV (see the report) but the same size model produced more tokens than the previous versions, and we believe in public LB they were timing out before having a chance to properly attempt the questions. Ultimately we were out of time and submissions to properly narrow down the discrepancies seen between CV and public LB, especially given the variance with one submission per day and only 50 questions scored. What was not used A number of experiments were made but were not in the final submission choices, mostly due to limited time likely causing early exit before a good answer was made, Stronger models - as detailed in the paper we had stronger models trained on hard problems, however there were more tokens generated and we were not confident it would finish within the Kaggle constraints. Generative Solution Re-Ranking - this type of reward model requires additional memory and time resources. Increasing the batch size and max tokens proved to be a simpler alternative compared to serving an extra generative model for smaller sample sets. Outcome Reward Model - Initial experiments showed some advantages with this model. However, for long reasoning generations and TIR, these benefits diminished, particularly when evaluating smaller subsets such as 50 problems. 7b models (both TIR and CoT) with redrafter were very fast but had lower CV and LB despite finishing comfortably with batch size 32. 32b CoT models comfortably fit a batch size of 8 and 18k generation length but the score was close to 14b. A potentially better smoothed buffer strategy was developed which calculated buffer time based on time remaining from the 5 hours. However it was not used when we made the above submission.