Supervised Fine-Tuning (SFT) for Code Generation #

A base language model emerges from pretraining as a broad next-token predictor over internet-scale text and code. It is fluent, but it is not yet a reliable assistant for any particular task. When prompted to produce a program to a specification, it may ramble, ignore the format, or drift off task. Supervised fine-tuning (SFT) is the first and most direct step that adapts such a base model into a task-specific generator. It is the foundation this tutorial builds on: Section 3.1 applies this recipe to kernel corpora, and the reinforcement-learning stage of Section 1.6 assumes an SFT model as its starting point. This section develops SFT generically, for code generation, before those later sections specialize it.

It helps to place SFT within the full training pipeline that turns raw text into a capable assistant:

  Pre-training  →  Mid-training  →  SFT  →  RL
                                    └────┬────┘
                                     post-training

Each stage plays a distinct role. Microsoft’s MAI-Thinking-1 technical report [1] is a useful concrete reference throughout, since it documents each stage of a from-scratch build, and we use it for illustrative figures below.

• Pre-training. Self-supervised next-token prediction over a broad, internet-scale corpus of text and code. It instills general language ability, world knowledge, and coding fluency, and is by far the most compute-heavy stage. Its output is the “base model.” MAI-Thinking-1, for instance, pre-trains on 30 trillion tokens of in-house-processed web data, public GitHub code, books, academic papers, news, and multilingual text [1].
• Mid-training. Continued next-token training on a deliberately reshaped data mixture: upweighting high-value, target-relevant data (code, math, reasoning), injecting domain-specific or long-context data, and annealing the learning rate. It sharpens the capabilities that post-training will rely on, bridging the generic base and the specialized assistant. (This is a newer, less-standardized term than the others.) In MAI-Thinking-1 this is 3.55 trillion tokens across two phases that emphasize STEM, math, and coding “to build a strong foundation for reasoning RL climbs,” and progressively extend the context window from 16K to 64K and then 256K tokens [1].
• Post-training. The umbrella for everything after pre-training. It converts the knowledgeable base model into a useful, instruction-following assistant, typically in two stages:
  • SFT (this section): imitate curated (prompt, target) demonstrations to teach instruction-following, output formatting, and task grounding, moving the prior toward good solutions. Its ceiling is the demonstrations themselves.
  • RL (Section 1.6): optimize a reward over the model’s own generations, pushing past imitation toward the outcome that actually matters (for kernels, a program that is correct and fast). It can surpass the demonstrations. MAI-Thinking-1’s RL climb rewards responses by executing code or by feedback from a prompted judge, and trains three domain specialists (STEM reasoning, agentic coding and tool use, and helpfulness/safety) [1].

The two post-training stages on one toy task: SFT matches the model's answer token by token against a given expert example, whereas RL samples several answers and scores them, reinforcing the one the checker rewards without any example to imitate.

The first two stages share one objective, next-token prediction on ever-more-curated data, while the two post-training stages differ sharply: SFT imitates fixed targets, whereas RL learns from the consequences of the model’s own outputs. (The ordering can vary: MAI-Thinking-1 notably drives its main RL climb directly from the mid-trained base with no prior SFT on reasoning traces, and only later uses a light SFT to consolidate its specialists [1]. The pre-training → mid-training → SFT → RL pipeline above is the common default, not a rigid law.) The rest of this section develops SFT.

The idea is imitation (behavior cloning): collect a corpus of high-quality (prompt, target) demonstrations and train the model to reproduce the targets by next-token prediction. For code generation, a prompt is a task specification (a natural-language description, a function signature, a docstring, perhaps tests or examples) and the target is a correct program that satisfies it. SFT does two things at once: it grounds the model in the syntax, APIs, and idioms of the target language, and it reshapes the model’s behavioral prior toward the region of the output space where good solutions live.

The SFT objective #

Conditional sequence modeling. We fine-tune an autoregressive model \( \pi_\theta \) with parameters \( \theta \) . Given a prompt \( x \) , the model assigns a probability to a target \( y \) by the chain rule, one token at a time:

where \( y_{\lt t} = (y_1, \dots, y_{t-1}) \) is the prefix of already-generated tokens. Each factor is a categorical distribution over the vocabulary \( \mathcal{V} \) , a softmax over the model’s logits \( z_t \in \mathbb{R}^{|\mathcal{V}|} \) at position \( t \) :

and because \( z_t \) depends on the whole context \( (x, y_{\lt t}) \) through the transformer, every token is predicted conditioned on the entire prompt and all previously generated tokens.

Maximum likelihood. Given a corpus \( \mathcal{D} = \{(x^{(i)}, y^{(i)})\}_{i=1}^{N} \) , SFT fits \( \theta \) by maximum likelihood estimation (MLE): choose the parameters that make the reference targets as probable as possible. Taking the log turns the product into a sum and negating turns maximization into minimization, giving the negative log-likelihood loss, which expands over tokens as:

Token-level cross-entropy. Each inner term is the cross-entropy between the one-hot target \( e_{y_t} \in \{0,1\}^{|\mathcal{V}|} \) and the model’s predicted distribution, \( -\sum_{v} e_{y_t}[v]\,\log \pi_\theta(v \mid x, y_{\lt t}) \) . So SFT is just next-token classification, supervised at every position of every target.

Conditioned on the teacher prefix \( y_{\lt t} \) (teacher forcing), the loss of SFT at each position is how far the model's predicted distribution sits from the target's next token.

Two properties of this objective are worth stating, because the later RL stage is motivated by the second:

1. Teacher forcing. At every position the model is conditioned on the ground-truth prefix \( y_{\lt t} \) , not its own past predictions. The whole target is scored in a single forward pass under a causal mask, with no autoregressive sampling during training, which is what lets the loss decompose into \( |y| \) independent, densely-supervised terms per example (training is cheap and stable).

2. Exposure bias. Because training always supplies the correct prefix while inference must consume the model’s own (possibly erroneous) generations, SFT is subject to the well-known train/inference exposure bias. One early mistake shifts the model into contexts it never saw during training. This mismatch, together with the fact that the loss rewards matching tokens rather than task success, is exactly what reinforcement learning (Section 1.6) is introduced to address.

Practical considerations for code #

Several choices specific to code generation shape how the objective above is applied in practice.

• Loss masking (prompt vs. completion). Each training example concatenates the prompt and the target into one sequence, but the loss is applied only to the target tokens. The prompt is attended to as context but never used as a supervised target (there is nothing to learn by predicting the user’s own request), so a loss mask zeroes out the prompt positions and the average is taken over completion tokens alone.

• Instruction and chat formatting. Code assistants are trained on a fixed template that wraps the task in role markers and delimits code blocks. Consistent formatting at training time is what lets the model reliably emit parseable, well-delimited code at inference time rather than prose mixed with snippets. This is the instruction-tuning recipe that turns a base model into a controllable assistant [2, 3].

• Fill-in-the-middle (FIM). Much real code editing is infilling, i.e., completing a span given both the code before and after it, not just left-to-right continuation. Code models are therefore commonly trained on a fill-in-the-middle transformation [5]: a document is split into prefix, middle, and suffix, reordered so the middle is predicted from both surrounding contexts. This teaches infilling without changing the next-token objective above.

• Reasoning distillation. As we will see in the kernel setting (Section 3.1), targets can be augmented with an explicit reasoning trace distilled from a stronger model, so the student learns to think through the approach before emitting code, not merely to reproduce final programs.

• Evaluation with pass@k. Code quality is judged by functional correctness, not token overlap with a reference. The standard metric is pass@k [4]: sample \( k \) programs per task and count the task solved if any one passes its tests. This directly reflects that many distinct programs can be correct, and it is the notion of success the RL reward later optimizes more directly.

Limitations, and what comes next #

SFT is powerful and usually indispensable, but imitation of a fixed corpus has intrinsic ceilings. Because the objective only maximizes the likelihood of reference targets, the model can at best reproduce the behavior it was shown; it cannot discover solutions better than its corpus, and its quality is capped by the strength of the demonstrations and the curation filters. And because the loss scores token overlap rather than whether the program actually runs correctly and efficiently, SFT has no direct signal for the outcome we care about. Closing that gap, by letting the model learn from executing its own generations against a reward, is the role of reinforcement-learning post-training, developed in Section 1.6 and applied to kernels in Section 3.2.

References #

1. The Microsoft AI Team. MAI-Thinking-1: Building a Hill-Climbing Machine. Microsoft AI, 2026. microsoft.ai/pdf/mai-thinking-1.pdf
2. J. Wei, M. Bosma, V. Y. Zhao, et al. Finetuned Language Models Are Zero-Shot Learners (instruction tuning / FLAN). arXiv:2109.01652, 2021. arxiv.org/abs/2109.01652
3. L. Ouyang, J. Wu, X. Jiang, et al. Training Language Models to Follow Instructions with Human Feedback (InstructGPT; SFT then RL). arXiv:2203.02155, 2022. arxiv.org/abs/2203.02155
4. M. Chen, J. Tworek, H. Jun, et al. Evaluating Large Language Models Trained on Code (Codex; the pass@k metric). arXiv:2107.03374, 2021. arxiv.org/abs/2107.03374
5. M. Bavarian, H. Jun, N. Tezak, et al. Efficient Training of Language Models to Fill in the Middle. arXiv:2207.14255, 2022. arxiv.org/abs/2207.14255