Fine-tune or distill from reviewed production data

Learning objective

Fine-tune or distill models from production data.

Core idea

Production data identifies tasks worth improving, but an observed answer is not automatically a correct training target. Separate the request, available context, attempted answer, and reviewed target before constructing a dataset. Supervised fine-tuning teaches selected responses; knowledge distillation transfers a teacher model’s behavior to a student. Both can reproduce mistakes in their supervision.

Use NeMo RL’s supervised fine-tuning interface for reviewed input-output examples. Megatron Bridge’s Model Optimizer integration provides a distinct logit-based distillation path. The latter trains against teacher outputs inside the training computation; it is more specific than collecting text from a larger model and calling the resulting dataset “distillation.”

The trigger that justifies training belongs to 8.4.1. This lesson concerns the dataset and training job that turn that decision into a reproducible candidate.

How it works

Freeze an eligible dataset

Create a training-data manifest that names the source release, collection period, source-case identifiers, permitted use, preprocessing revision, reviewer decisions, and train/validation membership. These are team-defined provenance fields, not mandatory NeMo API fields. The manifest lets the training owner reproduce selection and remove ineligible cases before they reach the trainer.

Retain the context the model actually needs. A corrected answer to a retrieval question is not a useful supervised target if the example omits the supporting passage. Conversely, copying all captured context can carry confidential or irrelevant material into the dataset. Apply the data quality and privacy checks covered in 6.4.2, then materialize only the approved fields.

Partition by the unit that can leak information: for example, keep an entire customer conversation or document family in one split. A row-level split can put near-identical turns in both training and validation. Deduplication mechanics belong to 6.4.1; here, preserve the resulting split membership in the manifest so a rerun cannot silently change it.

Keep a final release evaluation set separate from the data used to select training settings. Repeatedly choosing a candidate from its validation score makes that score part of the development process. The model owner must retain untouched cases for the eventual release decision in 8.4.3.

Construct supervised examples

NeMo RL SFT datasets expose dataset and task_name attributes. Formatted chat examples use messages with role and content, and each formatted example requires task_name. Use this structure when integrating a custom dataset class rather than inventing an incompatible trajectory format.

For local input-output data, the documented ResponseDataset loader accepts JSONL and maps source fields through input_key and output_key. Configure data.train.data_path and data.validation.data_path, with data.default.dataset_name: ResponseDataset and data.default.processor: sft_processor. This connects the approved dataset snapshots to the training job explicitly.

For production tool-use traces, select data.train.dataset_name: openai_format, set chat_key: "messages" and tool_key: "tools", and retain the conversation messages and per-record tool definitions in JSONL. This route preserves tool-use supervision rather than reducing each incident to a final answer pair. NeMo RL passes messages and tools to the tokenizer’s apply_chat_template(): before training, the owner checks that the model’s chat template renders that tool-calling format correctly. For heterogeneous tool argument schemas, set use_preserving_dataset: true to avoid data corruption.

The ResponseDataset route also documents split_validation_size and seed for constructing a split. When conversation-level separation matters, create separate files beforehand and use the explicit validation path. A random row split cannot infer which records belong to the same business case.

NeMo RL uses Hugging Face chat templates. tokenizer.chat_template can select the tokenizer’s default template; the documented passthrough mode uses YAML null when a template was already applied offline. Preserve the chosen template with the experiment. Applying formatting twice can make the model train on a different message structure from the one served.

Launch the approved YAML experiment with uv run examples/run_sft.py --config CONFIG_PATH, where CONFIG_PATH identifies that configuration. Archive the effective training inputs and produced candidate together. A completed job proves that training ran; task correctness still requires evaluation.

Distill with a compatible teacher

Megatron Bridge’s knowledge-distillation setup uses NVIDIA Model Optimizer. It loads teacher and student checkpoints, replaces the standard loss with KL divergence between output logits, optionally adds intermediate-state losses, runs forward passes through both models, and backpropagates only through the student. The teacher stays frozen, and only student checkpoints are saved.

The supported providers are GPT and Mamba-based models. Teacher and student must support the same parallelism strategy. With pipeline parallelism, intermediate-state KD losses are supported only on the final pipeline stage; virtual pipeline parallelism is unsupported. Do not assume that any pair of remotely hosted model endpoints can enter this training path.

Set the model provider’s teacher attribute to the teacher configuration, configure kd_config, and call convert_to_distillation_provider(). The page recommends NVIDIA Model Optimizer’s general Megatron Bridge distillation script for custom use cases. With the provided recipe script, KD settings go under model.kd_config in the configuration override file. logit_layers selects logit modules and intermediate_layer_pairs selects student-teacher state pairs; inspect the architectures before naming layers.

Supply the frozen production-derived dataset through the recipe’s training configuration, preserving the manifest and split membership established for supervised training. Review input eligibility and any response targets before submission. The teacher computes logits on these supplied inputs during training; production logs provide the cases, not precomputed teacher logits. Megatron Bridge recommends high-quality, diverse training data.

skip_lm_loss defaults to true; setting it to false adds the ordinary language-modeling loss and consumes more memory. kd_loss_scale defaults to 1.0 and controls the relative distillation scale when that loss is included; it is unused when skip_lm_loss is true. logit_kl_temperature, also defaulting to 1.0, controls the temperature for KL-divergence calculation.

The teacher adds a forward pass even though it receives no backward pass. Budget training for both loaded models and their computation. A smaller deployed student does not imply a small distillation job. Review the teacher’s task behavior on the production slice before paying to imitate it.

Key distinctions and trade-offs

Choose the learning signal and verify its implementation:

Teacher-generated text can also become reviewed input-output targets for SFT. That dataset route does not supply the teacher’s full output logits to the documented KD loss. Name which route the job uses so its compute estimate and reproducibility record describe the real experiment.

Architecture / decision scenario

A manufacturing maintenance assistant needs a smaller model for on-premises inference. On 600 held-out cases, the incumbent passes 576 and an untrained small candidate passes 510. The team requires at least 570 passes, zero unsupported shutdown instructions, and p95 latency below two seconds. These are scenario targets, not product guarantees.

The tempting alternative is to train directly on 20,000 captured answers. Review finds that 8% contain unsupported advice, so the model owner excludes those outputs as targets and commissions corrections. The data owner separates entire equipment incident threads between training and validation, then freezes the JSONL files and provenance manifest.

The team compares supervised training on corrected responses with logit-based distillation from an approved teacher. The platform owner first checks checkpoint and parallelism compatibility and reserves capacity for both teacher and student; the small inference footprint alone cannot justify scheduling KD.

Suppose the SFT candidate passes 572 cases with no prohibited instruction and reaches 1.8 seconds, while the distilled candidate passes 582 but emits two unsupported shutdown instructions. The release owner advances the SFT candidate for production regression checks. Higher imitation accuracy does not outweigh a mandatory operational safety condition.

Common exam traps

• Treating raw production answers as ground truth.
• Randomly splitting related conversation turns and calling the result independent validation.
• Confusing text-target SFT with teacher-logit distillation.
• Assuming a remote teacher API satisfies the documented Megatron Bridge provider requirements.
• Budgeting distillation from student inference memory alone.
• Promoting a trained checkpoint because its training loss decreased.

What to remember

Freeze eligible targets, provenance, formatting, and split membership before training. Use NeMo RL’s dataset and SFT interfaces for reviewed responses, or the compatible Megatron Bridge teacher-student path for logit-based distillation. Evaluate the saved candidate on task behavior and serving constraints before release.

Knowledge check

References

• NVIDIA, Supervised Fine-Tuning in NeMo RL
• NVIDIA, Knowledge Distillation