Supervised Fine-Tuning in NeMo RL#

This document explains how to perform SFT within NeMo RL. It outlines key operations, including initiating SFT runs, managing experiment configurations using YAML, and integrating custom datasets that conform to the required structure and attributes.

Launch an SFT Run#

The script, examples/run_sft.py, can be used to launch an experiment. This script can be launched either locally or via Slurm. For details on how to set up Ray and launch a job using Slurm, refer to the cluster documentation.

Be sure to launch the job using uv. The command to launch an SFT job is as follows:

uv run examples/run_sft.py --config <PATH TO YAML CONFIG> <OVERRIDES>

If not specified, config will default to examples/configs/sft.yaml.

Example Configuration File#

NeMo RL allows users to configure experiments using yaml config files. An example SFT configuration file can be found here.

To override a value in the config, either update the value in the yaml file directly, or pass the override via the command line. For example:

uv run examples/run_sft.py \
    cluster.gpus_per_node=1 \
    logger.wandb.name="sft-dev-1-gpu"

Reminder: Don’t forget to set your HF_HOME, WANDB_API_KEY, and HF_DATASETS_CACHE (if needed). You’ll need to do a huggingface-cli login as well for Llama models.

Datasets#

SFT datasets in NeMo RL are encapsulated using classes. Each SFT data class is expected to have the following attributes:

  1. dataset: A dictionary containing the formatted datasets. Each example in the dataset must conform to the format described below.

  2. task_name: A string identifier that uniquely identifies the dataset.

SFT datasets are expected to follow the HuggingFace chat format. Refer to the chat dataset document for details. If your data is not in the correct format, simply write a preprocessing script to convert the data into this format. response_datasets/squad.py has an example:

Note: The task_name field is required in each formatted example.

def format_data(self, data: dict[str, Any]) -> dict[str, Any]:
    return {
        "messages": [
            {
                "role": "system",
                "content": data["context"],
            },
            {
                "role": "user",
                "content": data["question"],
            },
            {
                "role": "assistant",
                "content": data["answers"]["text"][0],
            },
        ],
        "task_name": self.task_name,
    }

NeMo RL SFT uses Hugging Face chat templates to format the individual examples. Three types of chat templates are supported, which can be configured using the tokenizer.chat_template in your YAML config (see sft.yaml for an example):

  1. Apply the tokenizer’s default chat template. To use the tokenizer’s default, either omit tokenizer.chat_template from the config altogether, or set tokenizer.chat_template="default".

  2. Use a “passthrough” template which simply concatenates all messages. This is desirable if the chat template has been applied to your dataset as an offline preprocessing step. In this case, you should set tokenizer.chat_template to None as follows:

    tokenizer:
  chat_template: NULL

  3. Use a custom template: If you would like to use a custom template, create a string template in Jinja format, and add that string to the config. For example,

    tokenizer:
  chat_template: "{% for message in messages %}{%- if message['role'] == 'system'  %}{{'Context: ' + message['content'].strip()}}{%- elif message['role'] == 'user'  %}{{' Question: ' + message['content'].strip() + ' Answer: '}}{%- elif message['role'] == 'assistant'  %}{{message['content'].strip()}}{%- endif %}{% endfor %}"

By default, NeMo RL has some built-in supported datasets (e.g., OpenAssistant, NuminaMath-1.5, OpenMathInstruct-2, Squad, etc.), you can see the full list here. All of these datasets are downloaded from HuggingFace and preprocessed on-the-fly, so there’s no need to provide a path to any datasets on disk.

We provide a ResponseDataset class that is compatible with JSONL-formatted response datasets for loading datasets from local path or Hugging Face. You can use input_key, output_key to specify which fields in your data correspond to the question and answer respectively. Here’s an example configuration:

data:
  # other data settings, see `examples/configs/sft.yaml` for more details
  ...
  # dataset settings
  train:
    # this dataset will override input_key and use the default values for other vars
    data_path: /path/to/local/train_dataset.jsonl  # local file or hf_org/hf_dataset_name (HuggingFace)
    input_key: question
    subset: null  # used for HuggingFace datasets
    split: train  # used for HuggingFace datasets
    split_validation_size: 0.05  # use 5% of the training data as validation data
    seed: 42  # seed for train/validation split when split_validation_size > 0
  validation:
    # this dataset will use the default values for other vars except data_path
    data_path: /path/to/local/val_dataset.jsonl
  default:
    # will use below vars as default values if dataset doesn't specify it
    dataset_name: ResponseDataset
    input_key: input
    output_key: output
    prompt_file: null
    system_prompt_file: null
    processor: "sft_processor"

Your JSONL files should contain one JSON object per line with the following structure:

{
  "input": "Hello",     // <input_key>: <input_content>
  "output": "Hi there!" // <output_key>: <output_content>
}

We support using multiple datasets for train and validation. You can refer to examples/configs/grpo_multiple_datasets.yaml for a full configuration example. Here’s an example configuration:

data:
  _override_: true # override the data config instead of merging with it
  # other data settings, see `examples/configs/sft.yaml` for more details
  ...
  # dataset settings
  train:
    # train dataset 1
    - dataset_name: OpenMathInstruct-2
      split_validation_size: 0.05 # use 5% of the training data as validation data
      seed: 42  # seed for train/validation split when split_validation_size > 0
    # train dataset 2
    - dataset_name: DeepScaler
  validation:
    # validation dataset 1
    - dataset_name: AIME2024
      repeat: 16
    # validation dataset 2
    - dataset_name: DAPOMathAIME2024
  # default settings for all datasets
  default:
    ...

AIME2025 and AIME2026 are registered alongside AIME2024 and accept the same config keys (e.g. repeat), so any of them can drop into the validation: list above.

Custom datasets defined outside NeMo RL#

If you want to plug in a dataset class that lives outside the nemo_rl package (so you don’t have to edit the built-in registry), set dataset_name to a fully qualified dotted import path. The dispatcher will import the module and resolve the class. The class must accept the same kwargs as the built-in datasets (i.e. the full data config) and implement set_task_spec and set_processor.

data:
  default:
    dataset_name: my_pkg.my_module.MyDataset  # importable from PYTHONPATH

The class must be importable — install it as a package or add its parent directory to PYTHONPATH before launching training.

We support using a single dataset for both train and validation by using split_validation_size to set the ratio of validation. This works for any dataset class that calls split_train_validation in its __init__ — which today includes most built-in datasets, among them OpenAssistant, OpenMathInstruct-2, OpenR1-Math-220k, ResponseDataset, and Tulu3SftMixtureDataset. A dataset class that does not call it ignores split_validation_size; the dataset dispatcher emits a warning in that case, so a misconfigured run is not silent. If you want to support this feature for your custom datasets or other built-in datasets, you can simply add the code to the dataset like ResponseDataset.

# `self.val_dataset` is used (not None) only when current dataset is used for both training and validation
self.val_dataset = None
self.split_train_validation(split_validation_size, seed)

Megatron-LM offline-packed datasets#

Use megatron_sft_packed when each JSONL record groups one or more conversations that must remain in one training row and you want to bypass NeMo RL’s generic online sequence-packing path. This dataset mode is available only with the Megatron backend.

Each line must be a JSON object whose messages field is an ordered list of OpenAI-style messages. A packed row must start with a system message and end with an assistant message. To place multiple conversations in one row, start each additional conversation with another system message:

{"messages":[{"role":"system","content":"You are helpful."},{"role":"user","content":"2+2?"},{"role":"assistant","content":"4"},{"role":"system","content":"You are concise."},{"role":"user","content":"Capital of France?"},{"role":"assistant","content":"Paris."}]}

A .jsonl.packed suffix is also accepted and is read as JSON Lines, so an offline packer can mark its output without the loader treating it as an unknown format.

This path is text-only. Every message content must be a string; multimodal records are not supported.

Producing a packed dataset#

Packing happens offline, before training. The producer decides how many conversations fit in one row; the loader only re-tokenizes what it is given.

  1. Pick the pack length and use it for both sides. The packer must target the same data.max_input_seq_length the training config sets.

  2. Pick the tokenizer and the prompt format and use them for both sides. The loader re-tokenizes each record with megatron_sft.prompt_format, so a more verbose tokenization than the packer’s overflows the row. The loader warns and drops the overflowing tail rather than failing, so check the logs after changing either.

  3. Greedily append whole conversations to a row until the next one would not fit, then start a new row. Each conversation begins with its own system message, which is what marks the segment boundary.

  4. For context parallelism, keep every conversation a multiple of 2 * context_parallel_size tokens, or let the loader pad each segment up to that multiple.

The processor tokenizes each record, then pads or right-truncates it to data.max_input_seq_length. A reference packer uses the loader’s token-count helper, so its row boundaries match the loader for the same tokenizer and prompt format. It expects raw JSONL records with a messages array; every conversation must start with system and end with assistant:

uv run examples/converters/pack_megatron_sft_jsonl.py \
  --input /path/to/raw.jsonl \
  --output /path/to/train.jsonl.packed \
  --tokenizer /path/to/tokenizer \
  --max-seq-length 262144 \
  --prompt-format identity \
  --context-parallel-size 16

Use the same tokenizer, prompt format, sequence length, and CP size for packing and training. The packer warns if one conversation exceeds the row length; the loader truncates that conversation, so inspect warnings before training.

Configure the dataset and the direct packed path as follows:

sft:
  only_unmask_final: false
  val_micro_batch_size: 1

policy:
  train_micro_batch_size: 1
  automodel_cfg:
    enabled: false
  dynamic_batching:
    enabled: false
  megatron_cfg:
    enabled: true
    context_parallel_size: 1
    use_fused_linear_logprobs: false

data:
  train:
    dataset_name: megatron_sft_packed
    data_path: /path/to/train.jsonl.packed
    megatron_sft:
      prompt_format: identity
      override_pad_token: null
  validation:
    dataset_name: megatron_sft_packed
    data_path: /path/to/validation.jsonl.packed
    megatron_sft:
      prompt_format: identity
      override_pad_token: null
  default:
    chat_key: messages

megatron_sft.prompt_format accepts identity, nemotron-nano-v2, or nemotron-h-aligned. The optional override_pad_token replaces the preset’s pad token; assistant-prefix masking is defined by the prompt-format preset.

For context parallelism, every conversation segment is padded to a multiple of 2 * context_parallel_size. The size comes from policy.megatron_cfg.context_parallel_size, not the dataset entry. A direct-packed training or validation split cannot be mixed with regular datasets, and direct-packed SFT does not support dynamic batching, draft training, router replay, sft.only_unmask_final=true, policy.sequence_packing.fuse_loss=true, or fused linear log-probability loss. The relevant training or validation micro batch size must be 1. At context-parallel size 1, online policy.sequence_packing.enabled is not required. For context-parallel size greater than 1, set it to true as required by the MCore context-parallel path.

Energon Multimodal Datasets#

The optional Energon SFT backend reads prepared WebDataset shards while the existing Hugging Face backend remains the default. megatron-energon ships in the mcore extra, and the Megatron policy workers pick that environment up from ACTOR_ENVIRONMENT_REGISTRY, so the driver runs under a plain uv run:

uv run examples/run_sft_v2.py \
  --config examples/configs/recipes/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.yaml

Each shard sample contains a JSON payload and its media members. The JSON payload uses one complete conversation as the sampling unit:

{
  "messages": [
    {"role": "user", "content": [
      {"type": "image", "media_index": 0},
      {"type": "text", "text": "Compare with the next image."},
      {"type": "image", "media_index": 1}
    ]},
    {"role": "assistant", "content": "The second image is brighter."}
  ],
  "media": [
    {"type": "image", "member": "000001.first.jpg"},
    {"type": "image", "member": "000001.second.jpg"}
  ],
  "tools": null
}

Media references must appear exactly once and in manifest order. Multi-turn conversations, assistant tool calls, matching tool results, and the final assistant response remain in the same sample. Tool results must reference an earlier tool-call ID.

Configure the prepared dataset path and split as follows:

sft:
  # SFTv2 has no validation loop; the exemplar defaults (val_period: 10,
  # val_at_start: true) are rejected at startup.
  val_period: 0
  val_at_start: false
  val_at_end: false
data:
  _override_: true              # replace the exemplar's HF data block wholesale
  backend: energon
  max_input_seq_length: ${policy.max_total_sequence_length}
  shuffle: true
  energon:
    model_family: qwen          # required, no default: "qwen" or "nemotron"
    num_workers: 8
    shuffle_buffer_size: 1000
    processor_adapter: hf_multimodal
    task_encoder:
      name: generic_sft
  train:
    path: /path/to/prepared/energon/dataset
    split: train
    virtual_epoch_length: 1000  # batches per virtual epoch
  validation: null              # SFTv2 builds a train loader only
data_plane:                     # required by run_sft_v2.py
  enabled: true
  impl: local
  max_partitions: 2

model_family and the top-level data_plane block have no defaults and are not supplied by any exemplar config, so both must be set explicitly. The sft overrides are needed for a different reason: SFTv2 runs no validation pass, so it rejects val_period, val_at_start or val_at_end left at their exemplar values, and data.validation must be null — the loader is always built with split_role="train".

The processor runs inside Energon loader workers and returns the same tokenized message_log representation as the Hugging Face path. prepare_sft_batch creates the assistant loss mask, flattens the messages, and pads the batch without checking which loader produced it.

The v1 SFTProcessorAdapter and HFMultimodalSFTProcessorAdapter are narrow integration interfaces. They are planned to be replaced by a more comprehensive modular processor implementation; dataset loading and the policy-facing batch shape should remain stable through that change.

Optional controls under data.energon preserve Energon’s defaults when unset:

Field

Behavior

gc_collect_every_n_steps

Steps between garbage collections; 0 disables GC and null uses Energon’s default.

cache_pool_num_workers

Positive number of cache-copy threads; null uses Energon’s default.

cache_pool_max_gbytes

Positive cache-size limit in GiB; null uses Energon’s limit.

nvdataset_cache_dir

DSS cache root exported as NVDATASET_CACHE_DIR; null leaves the environment setting unchanged.

To let Energon form model-ready multimodal packs, configure packing on the task encoder and enable fused sequence packing:

policy:
  sequence_packing:
    enabled: true
    fuse_loss: true
data:
  energon:
    task_encoder:
      name: generic_sft
      packing:
        name: balanced_greedy_knapsack
        buffer_size: 64
        options:
          max_sequences_per_bin: 16   # optional cap on source conversations per pack
          balanced_knapsack_delta: 0 # optional extra initial bins for this algorithm
    max_samples_per_sequence: null  # optional shard read-order control

The task encoder owns the packer selection and buffer size; the policy block enables packed fused loss. max_samples_per_sequence only controls how many consecutive samples Energon reads from one shard; it does not affect pack layout. Without an Energon packing configuration, SFTv2 requires fixed batching. max_sequences_per_bin defaults to no cap. balanced_knapsack_delta defaults to zero and is supported only by balanced_greedy_knapsack. Pack capacity is taken from data.max_input_seq_length, and its padding multiple comes from policy.make_sequence_length_divisible_by. Dynamic batching and HybridEP flex dispatch are not supported with Energon-owned packs.

With Energon-owned packing, each sample_mask entry represents one physical pack, so num_valid_samples counts non-empty packs rather than source conversations. NLL loss scaling is unchanged because it is normalized by global_valid_toks.

Training dataloader checkpoints include the Energon worker state plus a fingerprint of the source, loader, and processor settings. Restore must occur before the first iteration, and a changed fingerprint fails instead of silently continuing with a different stream. SFTv2 accepts a single train source; use an Energon metadataset to blend prepared sources.

To use an application-specific task encoder or cooker without adding it to NeMo-RL, select an exported object from an absolute Python file path. The file must be available at the same path in every loader-owning worker environment. The task encoder must subclass BaseSFTTaskEncoder; cookers must be callable; both must declare supported model families with @supports_model_families.

data:
  energon:
    model_family: nemotron
    task_encoder:
      python_file: /workspace/plugins/nemotron_energon/__init__.py
      object_name: NemotronMultiModalTaskEncoder
      options:
        prompt_format: nemotron6-moe
    cookers:
      - python_file: /workspace/plugins/nemotron_energon/__init__.py
        object_name: cook_nemotron_conversation

Use object_name to select the exported class or function.

For GenericSFTTaskEncoder, loss_mask_mode=None (the default) creates masks from assistant roles. With loss_mask_mode="precomputed", each message must provide a binary, one-dimensional token_loss_mask tensor with the same shape as token_ids. This mode preserves the supplied mask for packed and unpacked conversations and rejects only_unmask_final=true. Packed batches also mask padding and the first token of each source, which has no preceding prediction within that source.

The selected file’s SHA-256 digest is included in the loader checkpoint identity. When the selected file is a package __init__.py, the digest covers every Python source file in that package, so resuming after changing plugin code fails safely.

Preformatted text sources#

For the generic Energon cooker, set the boolean dataset subflavor skip_chat_template: true when each message’s text already contains the chat format. The cooker maps this to CanonicalSFTSample.chat_template_preapplied, which the generic processor adapter forwards to get_formatted_message_log. This is a per-source setting: normal and preformatted sources can share an adapter and blend. Missing or false values use the normal template path; non-boolean values are rejected.

The preformatted path supports text-only turns and encodes each turn independently. The source owns the rendered template, BOS/EOS, assistant end-of-turn markers, and the division of text between roles for loss masking. Use the same format expected at generation time. In a blend with normal sources, include BOS exactly when data.add_bos is enabled. The path does not insert BOS/EOS and rejects tools, generation prompts, and task prompts. Independent encoding may differ from encoding the joined conversation at boundaries, including tokenizer-specific leading spaces.

OpenAI Format Datasets (with Tool Calling Support)#

NeMo RL also supports datasets in the OpenAI conversation format, which is commonly used for chat models and function calling. This format is particularly useful for training models with tool-use capabilities.

Basic Usage#

To use an OpenAI format dataset, configure your YAML as follows:

data:
  train:
    dataset_name: openai_format
    data_path: <PathToTrainingDataset>       # Path to training data
    chat_key: "messages"                     # Key for messages in the data (default: "messages")
    system_key: null                         # Key for system message in the data (optional)
    system_prompt: null                      # Default system prompt if not in data (optional)
    tool_key: "tools"                        # Key for tools in the data (default: "tools")
    use_preserving_dataset: false            # Set to true for heterogeneous tool schemas (see below)
  validation:
    ...

Data Format#

Your JSONL files should contain one JSON object per line following the OpenAI Chat Completions function calling format:

{
  "messages": [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What's the weather in Paris?"},
    {"role": "assistant", "content": "I'll check the weather for you.", "tool_calls": [
      {
        "id": "call_123",
        "type": "function",
        "function": {
          "name": "get_weather",
          "arguments": {"city": "Paris", "unit": "celsius"}
        }
      }
    ]},
    {"role": "tool", "content": "22°C, sunny", "tool_call_id": "call_123"},
    {"role": "assistant", "content": "The weather in Paris is currently 22°C and sunny."}
  ],
  "tools": [
    {
      "type": "function",
      "name": "get_weather",
      "description": "Get current weather for a city",
      "parameters": {
        "type": "object",
        "properties": {
          "city": {"type": "string", "description": "City name"},
          "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
        },
        "required": ["city"]
      }
    }
  ]
}

Note

NeMo RL passes messages and tools directly to the tokenizer’s apply_chat_template(), so correct tool call rendering also depends on the model’s chat template supporting this format.

Tool Calling with Heterogeneous Schemas#

When your dataset contains tools with different argument structures (heterogeneous schemas), you should enable use_preserving_dataset: true to avoid data corruption:

data:
  dataset_name: openai_format
  ...
  use_preserving_dataset: true  # IMPORTANT: Enable this for tool calling datasets

Why this matters: Standard HuggingFace dataset loading enforces uniform schemas by adding None values for missing keys. For example:

  • Tool A has arguments: {"query": "search term"}

  • Tool B has arguments: {"expression": "2+2", "precision": 2}

Without use_preserving_dataset: true, the loader would incorrectly add:

  • Tool A becomes: {"query": "search term", "expression": None, "precision": None}

  • Tool B becomes: {"query": None, "expression": "2+2", "precision": 2}

This corrupts your training data and can lead to models generating invalid tool calls. The PreservingDataset mode maintains the exact structure of each tool call.

Evaluate the Trained Model#

Upon completion of the training process, see the evaluation overview to assess model capabilities.

LoRA Configuration#

NeMo RL supports LoRA (Low-Rank Adaptation) for parameter-efficient fine-tuning of SFT models, including Nano‑v3 models, on both the Automodel and Megatron backends. To enable LoRA for SFT on the default Automodel backend:

uv run examples/run_sft.py policy.automodel_cfg.lora_cfg.enabled=true

For the full reference — backend support, the Automodel vs Megatron schema comparison, config examples, parameter details, example recipes, and Hugging Face export — see the dedicated LoRA guide.

Optimizations#

Chunked Linear Cross-Entropy Fusion Loss#

During standard SFT training the model materializes a full logit tensor of shape [batch_size, seq_length, vocab_size], which can cause out-of-memory (OOM) errors for long sequences or large vocabularies. The chunked linear cross-entropy fusion loss avoids this by computing the loss directly from the hidden states: it chunks the sequence dimension, projects each chunk to logits on the fly, computes per-token log probabilities, and discards the logits before moving to the next chunk.

Benefits:

  • Extends the maximum trainable sequence length significantly (e.g. from <65K to >100K tokens) by eliminating the large logit tensor from GPU memory.

  • Produces numerically equivalent loss values to the standard path.

How to enable:

Add the following to your Megatron config in your YAML file:

policy:
  megatron_cfg:
    enabled: true
    use_fused_linear_logprobs: true
    fused_linear_logprobs_chunk_size: 256  # tokens per chunk; smaller = less memory, larger = more throughput

Notes:

  • This optimization applies to SFT training with NLLLoss and DPO training. See the DPO guide for DPO-specific details.

  • Context parallelism is not supported when fused linear logprobs are enabled.

  • The fused_linear_logprobs_chunk_size parameter controls the trade-off between memory savings and compute throughput. The default value of 256 is a good starting point.