Kaggle Winners Walkthroughs: AI Mathematical Olympiad - Progress Prize 2 with Team NemoSkills

What is discussed, and when

1. 0:00 – 5:00: Art of Problem Solving, Nemo Skills, DeepC Carvan
2. 5:00 – 10:00: Lima 32B
3. 10:00 – 15:00: Nemo Skills, H100 GPUs
4. 20:00 – 22:32: Nemo Skills, Open Mass Reasoning

Notable quotes from "Kaggle Winners Walkthroughs: AI Mathematical Olympiad - Progress Prize 2 with Team NemoSkills"

• Side, what we did is we collected and processed about 300k mathematical problems from art of problem solving forums. We generated about 2.2 million chain of thought solutions for those problems using deepc carvan. We then fine-tuned quen.

  about Art of Problem Solving · at 1:07

• Each other. So we did it to maximize throughput and be able to cancel generations as needed. We used FP8 quantizations for keeping weights in FP8 and activations in Bflat 16.…

  about FP8 · at 2:28

• Simple API to do experiments that combine data generation training evaluation. It's very easy to switch from quick experiments that you're running on local machine to large…

  about Nemo Skills · at 3:39

• Show examples and so on. It didn't really work for us. So what we did instead is we took an instruction tuned model that was trained on a very small reasoning data set. In this case, it was our reproduction of Lima 32B.

  about Lima 32B · at 8:51

• Checkpoints to produce the final model which we find to improve results a little bit. So totally we use 512 H100 GPUs over 48 hours. So quite a bit of compute. But as you can…

  about H100 GPUs · at 11:38

• Activations for the weights and BF-16 for the activations. So you can see in this table that it basically has no accuracy degradation compared to the first row and reddrafter it's you know 2.6 six times.

  about Bfloat16 · at 14:26

• Things that we released. So we released the data set that we call open mass reasoning. So it contains all of the data that we generated and collected during this competition. So it includes about 540,000 problems extracted from these.

  about Open Mass Reasoning · at 21:14

Who and what is mentioned in this episode?

Companies & orgs

• Art of Problem Solving — Company · Data source · 1 mention · first at 1:12
  Used as a source for collecting mathematical problems for the competition.

Products & tech

• Nemo Skills — Product · Tool · 3 mentions · first at 3:22
  A codebase used by the team for fast experimentation and model improvement.
• DeepC Carvan — Product · Model · 1 mention · first at 1:20
  Used to generate chain of thought solutions for the mathematical problems.
  Processed about 300k mathematical problems from art of problem solving forums. We generated about 2.2 million chain of thought solutions for those problems using deepc carvan. We then fine-tuned quen 2.514b on this data to get what we'll.
• Lima 32B — Product · Model · 1 mention · first at 9:04
  A lightweight tuned reasoning model used to generate long reasoning solutions.
• Open Mass Reasoning — Product · Dataset · 1 mention · first at 21:15
  A dataset released by the team containing problems and solutions generated during the competition.
• H100 GPUs — Product · Hardware · 1 mention · first at 11:41
  GPUs used for training the models in the competition.

Concepts & topics

• Bfloat16 — Concept · Technique · 4 mentions
  A data type used for activations in the inference process.
• FP8 — Concept · Technique · 4 mentions
  A quantization technique used for keeping weights in FP8 to speed up inference.

Full transcript with timestamps

0:02 Awesome. Uh yeah. So today I'm going to talk about our first place solution to the AMO competition and uh we'll cover the following topics today. So first I'll introduce our team then summarize the solution briefly. Then we'll talk about details of uh of our solution. So we'll talk about how we collected a new set of problems, how we generated chain of thought solutions to integrated traing solutions. We'll talk about training and inference details. Discuss briefly what didn't work and uh finally show all the things that we released.

0:43 So let's go ahead and uh show the whole team. So we had seven people on the team. Ivan, Dra, another Ivan, Shupham, Crystal, Benedict, and myself. We and we are all from Nvidia. All right. So here is a summary of our solution. So on the training side, what we did is we collected and processed about 300k mathematical problems from art of problem solving forums. We generated uh about 2.2 million chain of thought solutions for those problems using deepc carvan. We then fine-tuned quen 2.514b on this data to get what we'll call a chain of thought model. Then we prompted this limo 32B which is a very lightweight tuned reasoning model to generate 1.2 million long reasoning solutions that integrate natural uh reasoning with Python code executions.

1:45 We then aggressively filtered those solutions down to just 15k high quality samples. Then we fine-tuned our chain of thought model and this data to get the TI model to integrate reasoning model and finally we merged the two models together to get the final one that we submitted. So these are the details on the training side and then on the inference side our system was doing 12 generations in parallel and we aggregated predictions using majority voting. We use tensor RTL um for inference. The uh inference implementation was fully asynchronous.

2:25 So different generations didn't need to communicate with each other or wait for each other. So we did it to maximize throughput and be able to cancel generations as needed. We used uh FP8 quantizations for keeping weights in FP8 and activations in Bflat 16. That was about 1.5 times faster than the original. And then we further used the reddrafttor for speculative decoding which made the total speed up to be about 2.6 uh times compared to the original. We implemented uh a fairly simple time buffer. So the current problem finished faster than we expect. We allocated more time to be spent on future problems. And we also had three different early stopping strategies that we'll cover in details later slides.

3:16 Okay. So the the final component of our solution was really our codebase which you can see from the name of our team. We use Nemo skills uh codebase from Nvidia. Uh so that was very important for us to be able to do fast experimentation. So this this library lets you orchestrate very complex pipelines for improving language models and do it easily. So there there is a simple API to do experiments that combine data generation training evaluation. It's very easy to switch from quick experiments that you're running on local machine to large scale learn jobs. It's easy to convert checkpoints to different formats from MIMO for training to TLM or hugging face for inference and uh it also integrates different uh inference and training libraries in a unified way. So you can easily switch between ton v and so on and uh you know different train back ends as well. Okay. Now let's uh dive a little bit deeper into the details. So first

4:17 let's talk about the problems collection. So as I was mentioning we use this art of problemsolving.com forums. So you can see here uh how one particular discussion looks like. So it starts with a problem definition then different people propose solutions. So this one's relatively clean. Sometimes there are multiple problems presented and people discuss both or or many. Sometimes there are no problems at all. Uh sometimes the solutions are correct. Sometimes people point out mistakes in them. So does require quite a bit of uh processing to get clean data out of this uh raw format.

4:59 So here is how how we structured our pipeline for extracting and processing problems. So we start with raw form post and about there are about 600,000 of those and then each block in this diagram represents a call to an LLM that does some kind of cleaning or or processing. So you start with a block that extracts all problems. You then do a series of classifications. Is this a proof question? Is it not proof question? Is it multiple choice? Is it binary question that's just yes or no answer? Is the question invalid or missing some context? We also uh try to extract the ground truth answers from the discussion so that we can later use them to filter out incorrect solutions.

5:41 The final block is the decontamination to make sure that our training set and validation set don't overlap. So here is an example how the prompt looks like for the very first block. So you can see that it's fairly detailed. we give a lot of instructions to the model and then we show some examples. So this prompt is not complete. Uh there are more examples and then it ends up is actually giving a a particular post to a model and uh asking it to complete the output field and uh each block here is a link. So we published all of the prompts. It's all open source inosklls. Okay.

6:22 So uh the only clarification here is that we actually for the final model that we submit we actually did not use this four blocks. So the pipeline was relatively uh simpler and we'll talk about why that's the case towards the end. So let's now talk about channel solution generation. So we ended up with about 5 uh,000 540,000 problems uh from the problem extraction pipeline and then we generate a lot of solutions for them using a mix of DC car 1 QWC2B for any solutions where we were able to extract the answer we use it directly to judge whether the solution is correct or not for any other solutions where we don't know the answer we just treat majority answer as as a ground truth. And for for the judge step, we don't just use symbolic comparison. We actually have an alarm call here and we use CL 2.5 to be instruct for that. And the model is shown a problem statement and two answers and is asked to compare

7:28 whether the answers are the same or not in the context of the problem. We do this because in this forums sometimes the answers are fairly complicated to compare in a rulebased system. So our final data is about 4 million samples and again a clarification that uh while we do have more data for the final model we submitted that took the first place actually used quite a bit fewer so only about 2.2 to a million solutions and only from deepc car one as the teacher. So let's now talk about how we integrated code execution into this long reasoning solutions. So here is an example of a trir solution where a model reasons in natural language and then at some point decides that it wants to write a python code to compute something. The output of this Python code is not produced by the model, but instead we extract the code, we execute in an in an actual Python interpreter, we get the output and we

8:30 put it back into the prompt, let the model continue. So it turns out that it's actually um not that simple to do this. So the reasoning models that we had R1 QW are really hard to just prompt into generating such solutions. So we did a lot of tries. tried to instruct them, show examples and so on. It didn't really work for us. So what we did instead is uh we took an instruction tuned model that was trained on a very small reasoning data set. In this case, it was our reproduction of Lima 32B model. And since that model is instruction tuned and is not overfeed to this reasoning long chain of flow solutions, it will we were actually able to just prompt it to produce solutions that integrate code in natural language.

9:20 So we use it to generate about 1.2 million solutions. But the problem is that they were actually low quality. And there there were mainly two problems with them that we tried to address with this uh filtering stop using alum as a judge. So the first problem was that sometimes the solutions did not add anything new. They just were used to verify existing calculation which might be fine but they kind of want code execution to drive towards you know to to be used outside of just verification. And then the second one sometimes uh code execution was done to do something new but that was fairly simple calculation that is just easier to do in chain of thought. So we designed this two prompts. So these are again just all prompts uh to filter any code block that doesn't uh contribute significant noity or is not significant enough on itself. So we filter this 1.2 2 million solutions down to just 15,000. And uh you can see that we

10:23 remove about 99% of the data. But what we end up with is fairly high quality samples that you can then use for training of our final models. And of course you can do a few more iterations here. So you can now train on this 15k samples. Do more generation, do more filtering, repeat this process. So we actually did try that for for you know our best models. But this is not used for the final keggo model. So the final keggo model is only trained on this 15k high quality samples. So let's now cover the training details. So for the training we started with quen 2.5 14bit base model and we did sftt for about eight epox on 2.2 million chain of thought solutions. So for the hyperparameters we used Adamw optimizer uh we used cosign learning create schedule fairly big batch size 16k maximum length uh we did leverage advanced paralization techniques like tensor and context parallel and sequence packing which helped us to speed up this

11:33 training quite a bit. We also averaged four equally space checkpoints to produce the final model which we find to improve results a little bit. So totally we use 512 H100 GPUs over 48 hours. So quite a bit of compute. But as you can see on the plot uh here on the right we actually could get about 90% of the results with you know four times uh less than that. And uh so but but we really wanted to push the models to to the limit. So we did spend all of it to train the best models possible. So this this is followed by a lightweight SFT on just 50 and thousand samples. So this doesn't take almost any compute compared to the first stage. So just do four 100 stops with a constant viewing rate no checkpoint averaging very quick. We use Nemo line our back end of Nemo skills for training and then we merge uh these two models is just a linear combination of weights and we use 0.3 weight for the chain fault model and

12:39 0.7 weight for the TI model and I should say that when data is more balanced so when we did generate more TI data when it's also millions you could you can actually train on all of these together so you don't need two stages and in fact the plot on the right shows the joint train when we just use different prompts for chain of thought and TR models and that works just as well. All right, so let's now cover the inference details. So for the inference we used uh reductor technique for speculative decoding. So the idea is fairly simple.

13:19 Whenever you want to predict the next token, instead of running generations for a big LM, use a very small RNN head that attached to the last layer and then you predict three tokens using that RNN and you try to verify them using an LM doing parallel predictions for all three tokens. If the LM accepts all three, we take them and move on. If it doesn't accept, we regenerate the first token with with an LM and repeat the process. So this itself uh gave about 1.5x speed up which was very nice and the total acceptance rate was about 65%. So we did train this on the data generated from our final model. So it's important to match the distribution and uh we suspect that there are other drafting techniques that might work better for example using toal lamps but we did not uh get it implemented during this competition.

14:13 So besides the speculative decoding we also used quantization. We experimented with different options but the one that uh showed no decrease in accuracy and gave the best speed was using FP8 activations for the weights and BF-16 for for the activations. So you can see in this table that it uh it basically has no accuracy degradation compared to the first row and reddrafter it's you know 2.6 six times faster. Uh now for for the generations as I was mentioning our generations fully as sync. So we using flight batching and you can see here an example of synchronous generations. So in this case even though different samples finish at different times we only get all outputs when the L one finishes. I think sync generation you get outputs as soon as any of them is finished. And you can also integrate this with code execution in which case you know whenever you need to execute code you stop the generation you take the Python

15:16 run it and then add it back to the problem then you continue the generation. This is again fully I think no dependency between different samples and we use this uh to be able to cancel different generations to implement our early stopping. So this is combined with a time buffer for each question. We allocate about three or exactly 350 seconds and then there is a variable buffer of 200 seconds that can be added if we finish some problems early. We did implement a different version that's a bit smoother and less has less variance depending on the uh random order of the problems but it's not used in our final solution. So for early stopping we use these three criteria. The first one is we look at the first five generations that are finished. If four of the five have the same answer, we just take that as an answer and we don't uh we cancel all the remaining generations. When the first 10 out of 12 are finished, we don't look at

16:17 the last two. So we just cancel them to uh kind of work around edge cases. when some generations take a really long time and of course if you run out of time in our buffer we just cancel you know whenever that happens and take all the answers that we have. So these are all details of of our final submission and let's not now talk about the things that we tried but uh that that didn't work. So first we experimented with different methods to select one solution out of all the generations and the goal here is to improve upon majority voting. So we tried uh using different reward uh models for example outcome based reward models. We tried some heristics for example taking uh you know different methods to break tights in majority like taking the shortest or longest solutions and we even implemented a novel technique that we call J select which uh the idea here is that you show all of the solutions in

17:22 prompt you ask an all to pick one. You can see here that it works substantially better than majority voting especially at low sample count which was particularly relevant in Kaggle. So like we couldn't do more than 20 generations for sure. And our final method for J select actually is fairly efficient. So you can use a single model for both generations and selection and you only need about 10% extra compute compared to majority voting. But still even this uh even with this we could not really see improvement in scores in the leaderboard perhaps because the extra time that we need to spend here wasn't you know just like the score improvements weren't big enough to justify this extra time that we need to spend and uh the other thing we could not do is as I was mentioning we actually did have better models than what the model that we eventually submitted trained on more data Here is the table showing all of the model

18:23 operations that you trained. So you can see our uh kegle model over here. That's the actual one that that I submit. You can see that it's substantially worse than the best 14 bit model that we had. So like 8 point difference here and also on AM benchmarks as even comparable to the 7B model. Um and uh you know why didn't we use the bath models? Of course it's not because we didn't want to. We we tried to use them but ultimately just weren't able to make a high score in submission and we think that the likely reason is a combination of these four factors. So first as probably all people know who participated in this competition there is a fairly high variance in scores. So even if you submit exactly the same uh notebook the scores might vary by two or three points sometimes much more. You only have one submission per day. So sometimes it's just hard to measure uh different options on the leaderboard.

19:33 And we did have our local evalations of course, but as you saw the 7B model worked really well in our local, but somehow we never saw a high score with it in KGO even though we could do more generations and it was more stable. So we did spend a lot of time focusing on smaller models in particular 7B but that might have been a mistake um because maybe our local distribution in evaluation doesn't match what we had in Kaggle or maybe smaller models just can't solve you know a few particular hard problems or something else but yeah we did spend quite a bit of time without any benefit here and uh because of those two first points We also uh didn't really have enough time to address the fact that our latest models turned out that they produce about 10% more tokens on average. So we always controlled for the maximum lens but we only uh towards the end of the competition we noticed that even though the maximum length is the

20:35 same but for easy and medium problems uh the latest models were just consuming a little bit more tokens and this accumulated overall uh you know over five hours of the of the submission and uh ultimately seemed enough to affect the scores and uh you know we just did not really allocate ourselves enough time to fix all of this because we maximized the time we spent training the best models only left about one week to submit those final models and just didn't have time to address these issues. So let's uh let's see all the things that we released. So we released the data set that we call open mass reasoning. So it contains all of the data that we generated and collected during this competition. So it includes about 540,000 problems extracted from these forums. 4.3 million long reasoning solutions, 1.7 million T solutions and 560,000 gentle L traces. We released a series of open mass neatron models that are

21:47 state-of-the-art on mathematical reasoning tasks in their respective model sizes. And this includes our vin model in Kaggle. We also released all the code that's needed to reproduce our results in Nemo skills. So you can go to our documentation, see all of the commands and scripts you need to run to fully repeat what we did. We published a paper with more details. So all of the things are available on this links. Fully open sourced uh commercially permissive license. Welcome everyone to try try and use those. And that's the end of the talk.

22:28 Happy to answer any questions.

Questions about this episode

What is discussed in "Kaggle Winners Walkthroughs: AI Mathematical Olympiad - Progress Prize 2 with Team NemoSkills"?

In order: 0:00 Art of Problem Solving, Nemo Skills, DeepC Carvan; 5:00 Lima 32B; 10:00 Nemo Skills, H100 GPUs; 20:00 Nemo Skills, Open Mass Reasoning.

What is mentioned in this Kaggle episode?

8 people, works and ideas are mentioned, including Nemo Skills, Art of Problem Solving, DeepC Carvan, Lima 32B, Open Mass Reasoning, Bfloat16. Each one on this page links to its own page with every episode that mentions it.

How long is "Kaggle Winners Walkthroughs: AI Mathematical Olympiad - Progress Prize 2 with Team NemoSkills"?

The episode runs 23 min; the transcript is about 3,500 words in 25 timestamped sections.

More from Kaggle

• All analyzed episodes
• People mentioned on Kaggle
• Topics mentioned on Kaggle