Multi-Task LLM Post-Training

Data and Reward Alignment for Reasoning, Code, and Instruction Following

Overview

We studied how to post-train one language model for three different capabilities: following verifiable instructions, solving mathematical problems, and completing Python functions. Improving any one benchmark in isolation was straightforward. The harder question was how to improve one behavior without erasing the others. We used Llama-3.2-3B [1] to test data mixtures and training settings, then used what we learned to train Llama-3.1-8B. All stages used LoRA adapters [2], so the base model remained frozen. For more details, read the full report.

Tasks and Evaluation

We evaluated the model on three held-out benchmarks:

Task Evaluation set What it measures
Instruction following IFEval [4] Whether a response satisfies verifiable constraints involving keywords, length, language, structure, and output format.
Mathematical reasoning GSM8K test split [5] Whether the model solves a multi-step grade-school word problem and returns the correct final numerical answer.
Code generation HumanEval test suite [6] Whether a generated Python function, completed from its signature and docstring, passes executable tests.

MATH [9] problems were also used during mixed RLVR, but were not part of the reported three-task average.

Balancing the Training Data

We first fine-tuned only on GSM8K. Math accuracy rose to 45.0%, but HumanEval fell to 0.0% and IFEval remained at 22.3%. Training for one task had made the model worse at the other two.

Mixing the tasks recovered the missing abilities. A small balanced set reached a 39.5% average, and a broad 120,000-example mixture reached 46.8%. A smaller 88,803-example mixture built from task-specific Tulu 3 subsets [3] did better, reaching 55.8%. It contained constrained instructions, worked math solutions, and direct code responses. The type of example mattered more than the total number of examples.

Selected supervised fine-tuning experiments comparing IFEval, GSM8K, and HumanEval accuracy across single-task, balanced, broad, and task-aligned data mixtures
The GSM8K-only run improved math but lost code generation. The task-aligned mixture gave the best balance across all three evaluations.

The code data showed the same pattern. Keeping the math and instruction data fixed but replacing PersonaHub Code with Evol-CodeAlpaca lowered HumanEval from 40.2% to 34.1%. PersonaHub used more direct code responses, while Evol-CodeAlpaca leaned toward competitive-programming-style solutions. HumanEval expects a short function completion, so the more similar training format transferred better.

We used the resulting mixture to train Llama-3.1-8B on 93,800 examples: 30,000 math, 28,803 instruction-following, and 34,997 code examples. This supervised model reached 74.6% on IFEval, 75.9% on GSM8K, and 53.0% on HumanEval, for a 67.8% average.

LoRA Rank and Batch Size

We compared LoRA ranks on the same Llama-3.2-3B data mixture. Rank 64 improved all three evaluations over rank 32, while rank 128 reduced the math and code scores.

LoRA rank IFEval GSM8K HumanEval Average
32 48.3% 50.5% 39.0% 45.9%
64 52.8% 52.8% 41.5% 49.0%
128 50.0% 45.5% 31.7% 42.4%

We therefore used rank 64 for the final model. The 8B supervised stage used a batch size of 32 examples. Mixed RLVR used 32 prompts per update, MBPP code RLVR used 8, and the final instruction-following stage used 32. Each RLVR prompt produced eight candidate answers.

Adding Verifiable Rewards

After supervised fine-tuning, we used reinforcement learning with verifiable rewards (RLVR), where answers are checked directly. GSM8K [5] used exact numerical answers, MATH [9] used normalized symbolic answers, and IFEval checked whether each requested constraint was satisfied [4].

For each prompt, the model produced eight answers. The GRPO-style update [8] compared their rewards. It could learn when some answers passed and others failed. If all eight received the same reward, there was no better answer to reinforce.

This stage raised IFEval from 74.6% to 77.7% and GSM8K from 75.9% to 78.1%. HumanEval changed only slightly, from 53.0% to 53.7%, since code was not part of this stage. The three-task average increased from 67.8% to 69.8%.

RLVR training curves showing reward by task and the fraction of answer groups that received identical rewards
GSM8K produced a useful mix of correct and incorrect answers. MATH reward stayed near zero because most groups contained no correct answer.

The difference between GSM8K and MATH was not the accuracy of the checker. Both had clear answer checks. MATH was simply too difficult for the model at this stage: when all eight answers were wrong, the reward could not show which direction was better. Easier problems, partial credit, or a gradual increase in difficulty could provide a stronger learning signal.

Matching Code Training to HumanEval

Code remained the weakest result, so we added a code RLVR stage with executable tests. Our first dataset asked for long, standalone solutions. The tests worked, but HumanEval [6] stayed at 53.7%. HumanEval instead gives the model a function signature and docstring and asks it to complete the function. The two formats did not match.

We replaced that dataset with MBPP [7], which also uses short function-writing problems and unit tests. We removed one problem whose function signature overlapped with HumanEval and trained on the remaining 169 problems. HumanEval then rose from 53.7% to 58.5%, and the three-task average reached 71.7%.

HumanEval accuracy before and after two code RLVR strategies alongside the MBPP unit-test reward curve
Executable rewards helped HumanEval only when the training problems asked for the same kind of short function completion.

Results

We finished with a short instruction-following RLVR stage. The final model scored 78.3% on IFEval, 78.2% on GSM8K, and 60.4% on HumanEval, for a 72.3% average. Across the 8B training sequence, the average rose from 67.8% after supervised fine-tuning to 69.8% after mixed RLVR, 71.7% after MBPP code training, and 72.3% at the end.

IFEval, GSM8K, HumanEval, and average accuracy across task-aligned supervised fine-tuning and three RLVR stages
Performance across the four stages of the 8B post-training sequence.

Evaluation Notes

Final results used greedy decoding and the complete benchmark suites. We scanned the training data for exact or substring matches with the test prompts and found none. This check does not detect paraphrased or semantic overlap. We did not repeat every experiment with multiple random seeds. The results describe performance on IFEval, GSM8K, and HumanEval rather than general instruction-following, math, or programming ability.

Takeaway

More data or more training stages did not automatically produce a better model. The data had to cover all three tasks and match the expected response format. RLVR helped when the model produced both successful and unsuccessful answers, giving the reward a useful comparison.

References

  1. A. Grattafiori et al., The Llama 3 Herd of Models , 2024.
  2. E. J. Hu et al., LoRA: Low-Rank Adaptation of Large Language Models , ICLR 2022.
  3. N. Lambert et al., Tülu 3: Pushing Frontiers in Open Language Model Post-Training , 2024.
  4. J. Zhou et al., Instruction-Following Evaluation for Large Language Models , 2023.
  5. K. Cobbe et al., Training Verifiers to Solve Math Word Problems , 2021.
  6. M. Chen et al., Evaluating Large Language Models Trained on Code , 2021.
  7. J. Austin et al., Program Synthesis with Large Language Models , 2021.
  8. Z. Shao et al., DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models , 2024.
  9. D. Hendrycks et al., Measuring Mathematical Problem Solving With the MATH Dataset , NeurIPS Datasets and Benchmarks 2021.