feat(sft): add SFT v2 stage 1 with the Energon multimodal data path

docs/guides/sft.md

@@ -168,6 +168,83 @@ self.val_dataset = None
 self.split_train_validation(split_validation_size, seed)
 ```
 
+### 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`:
+
+```bash
+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:
+
+```json
+{
+  "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:
+
+```yaml
+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
+    packing_buffer_size: null
+  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, including model inputs such as Qwen3-VL grid metadata or Nano Omni image sizes and frame counts. `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.
+
+Sequence packing is unavailable in this path, on both sides: `packing_buffer_size` and `max_samples_per_sequence` are typed null-only, and `policy.sequence_packing` (like `policy.dynamic_batching`) is rejected at startup with `SFTv2 requires fixed NeMo-RL batching.` Packing is deferred to a later stage of the Energon integration. Energon does not provide a separate offline sequence-packing pipeline either; offline preparation may store length and media-cost metadata, but should not pre-concatenate multimodal conversations.
+
+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.
+
 ### 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.
@@ -297,4 +374,4 @@ policy:
 
 - This optimization applies to SFT training with `NLLLoss` and DPO training. See the [DPO guide](dpo.md#chunked-fused-linear-logprobs) 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.
+- 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.

examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.yaml

@@ -0,0 +1,47 @@
+defaults: vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.yaml
+sft:
+  max_num_steps: 100
+policy:
+  model_name: /data/models/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
+  megatron_cfg:
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+    bias_activation_fusion: false
+    tensor_model_parallel_size: 8
+    expert_model_parallel_size: 8
+    sequence_parallel: true
+    freeze_vision_model: true
+    freeze_vision_projection: false
+    freeze_sound_encoder: true
+    freeze_sound_projection: true
+    moe_router_dtype: fp32
+    moe_router_load_balancing_type: none
+    moe_aux_loss_coeff: 0.0
+    moe_router_enable_expert_bias: true
+    defer_fp32_logits: true
+    radio_force_cpe_eval_mode: true
+    optimizer:
+      optimizer_cpu_offload: true
+      optimizer_offload_fraction: 1.0
+    distributed_data_parallel_config:
+      overlap_param_gather: false
+checkpointing:
+  enabled: false
+data:
+  add_bos: false
+  add_eos: false
+  energon:
+    model_family: nemotron
+logger:
+  wandb_enabled: false
+  log_dir: logs/sft-nemotron-omni-30b-clevr-energon-tp4etp4-v2
+  wandb:
+    name: sft-nemotron-omni-30b-clevr-energon-tp4etp4-v2
+  tensorboard:
+    log_dir: tb_logs-sft-nemotron-omni-30b-clevr-energon-tp4etp4-v2
+  swanlab:
+    name: sft-nemotron-omni-30b-clevr-energon-tp4etp4-v2
+  mlflow:
+    run_name: sft-nemotron-omni-30b-clevr-energon-tp4etp4-v2
+cluster:
+  gpus_per_node: 8

examples/configs/recipes/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.yaml

@@ -0,0 +1,60 @@
+defaults: ../../sft_vlm_3B.yaml
+sft:
+  max_num_steps: 1000
+  val_period: 0
+  val_batches: 0
+  val_at_start: false
+policy:
+  dtensor_cfg:
+    enabled: false
+  megatron_cfg:
+    enabled: true
+    apply_rope_fusion: false
+    distributed_data_parallel_config:
+      overlap_grad_reduce: false
+  make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size}
+  optimizer: null
+checkpointing:
+  checkpoint_dir: results/sft_${policy.model_name}_clevr_energon_v2
+  metric_name: null
+data:
+  _override_: true
+  backend: energon
+  max_input_seq_length: ${policy.max_total_sequence_length}
+  add_bos: true
+  add_eos: true
+  add_generation_prompt: false
+  shuffle: true
+  num_workers: 0
+  energon:
+    model_family: qwen
+    num_workers: 8
+    shuffle_buffer_size: 1000
+    processor_adapter: hf_multimodal
+    topology_mapper: default
+    task_encoder:
+      name: generic_sft
+      options: {}
+    cookers:
+    - name: generic_conversation
+      options: {}
+  train:
+    path: /data/nemorl-datasets/clevr
+    split: train
+    virtual_epoch_length: 4375
+  validation: null
+data_plane:
+  enabled: true
+  impl: local
+  max_partitions: 2
+logger:
+  wandb_enabled: true
+  log_dir: logs/sft-qwen2.5-vl-3b-clevr-energon-v2
+  wandb:
+    name: sft-qwen2.5-vl-3b-clevr-energon-v2
+  tensorboard:
+    log_dir: tb_logs-sft-qwen2.5-vl-3b-clevr-energon-v2
+  swanlab:
+    name: sft-qwen2.5-vl-3b-clevr-energon-v2
+  mlflow:
+    run_name: sft-qwen2.5-vl-3b-clevr-energon-v2

examples/configs/sft.yaml

@@ -237,6 +237,7 @@ policy:
       fp8_param: false
 
 data:
+  backend: hf
   max_input_seq_length: ${policy.max_total_sequence_length}
   add_bos: true
   add_eos: true

examples/run_sft.py

@@ -202,7 +202,12 @@ def main(is_vlm: bool = False):
         tokenizer = get_tokenizer(config.policy["tokenizer"], get_processor=is_vlm)
 
         # setup data
-        dataset, val_dataset = setup_data(tokenizer, config.data)
+        if config.data.get("backend") == "energon":
+            raise ValueError(
+                "Energon backend is not supported for V1, use `run_sft_v2.py` instead."
+            )
+        else:
+            dataset, val_dataset = setup_data(tokenizer, config.data)
 
         (
             policy,

examples/run_sft_v2.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Run single-controller SFT with colocated Energon loaders."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import pprint
+import warnings
+
+import ray
+from omegaconf import OmegaConf
+
+from nemo_rl.algorithms.sft_v2 import (
+    MasterConfig,
+    SFTSingleControllerActor,
+    setup_sft_v2,
+)
+from nemo_rl.algorithms.utils import get_tokenizer
+from nemo_rl.distributed.virtual_cluster import init_ray
+from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
+from nemo_rl.utils.logger import get_next_experiment_dir
+
+
+def parse_args() -> tuple[argparse.Namespace, list[str]]:
+    """Parse the config path and Hydra-style overrides."""
+    parser = argparse.ArgumentParser(description="Run SFTv2 training")
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=os.path.join(
+            os.path.dirname(__file__),
+            "configs",
+            "recipes",
+            "vlm",
+            "vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.yaml",
+        ),
+    )
+    return parser.parse_known_args()
+
+
+def main() -> None:
+    """Load configuration and start the SFT single controller."""
+    register_omegaconf_resolvers()
+    args, overrides = parse_args()
+    config = load_config(args.config)
+    if overrides:
+        config = parse_hydra_overrides(config, overrides)
+    resolved = OmegaConf.to_container(config, resolve=True)
+    master_config = MasterConfig.model_validate(resolved)
+    master_config.logger["log_dir"] = get_next_experiment_dir(
+        master_config.logger["log_dir"]
+    )
+    pprint.pprint(master_config.model_dump())
+
+    # NEMO_RL_OTEL_* env is snapshotted into the Ray runtime_env and inherited
+    # by every worker, so this has to run before init_ray().
+    init_telemetry_driver(master_config, algorithm="sft_v2")
+
+    init_ray()
+    processor = get_tokenizer(master_config.policy["tokenizer"], get_processor=True)
+    actor_args = setup_sft_v2(master_config, processor)
+    controller = SFTSingleControllerActor.remote(master_config, actor_args)
+    try:
+        result = ray.get(controller.run.remote())
+        pprint.pprint(result)
+    finally:
+        try:
+            actor_args.trainer.shutdown()
+        except Exception as error:  # teardown must preserve the controller failure
+            warnings.warn(f"SFTv2 trainer shutdown failed: {error}", stacklevel=2)
+        shutdown_telemetry()
+
+
+if __name__ == "__main__":
+    main()

nemo_rl/algorithms/loss/utils.py

@@ -445,6 +445,7 @@ def prepare_packed_loss_input(
 
     input_ids = data["input_ids"]
     unpacked_seqlen = input_ids.shape[1]
+    input_is_prepacked = "cu_seqlens" in data
     cp_size = (
         1
         if context_parallel_group is None
@@ -456,14 +457,20 @@ def prepare_packed_loss_input(
         else torch.distributed.get_rank(context_parallel_group)
     )
 
-    packed_rolled_targets = _pack_input_ids(
-        input_ids,
-        cu_seqlens_q,
-        cu_seqlens_q_padded,
-        cp_rank=cp_rank,
-        cp_size=cp_size,
-        roll_shift=-1,
-    )
+    if input_is_prepacked:
+        # Energon already produced one physical row. The distributed helper
+        # shifts each source independently before CP slicing, so source N never
+        # predicts the first token of source N+1.
+        packed_targets = input_ids
+    else:
+        packed_targets = _pack_input_ids(
+            input_ids,
+            cu_seqlens_q,
+            cu_seqlens_q_padded,
+            cp_rank=cp_rank,
+            cp_size=cp_size,
+            roll_shift=-1,
+        )
 
     # With chunking, keep logits in their original dtype: the chunked logprob
     # kernel casts each chunk to float32 internally.
@@ -474,7 +481,7 @@ def prepare_packed_loss_input(
 
     logprobs = from_parallel_logits_to_logprobs_packed_sequences(
         logits_for_logprobs,
-        packed_rolled_targets,
+        packed_targets,
         cu_seqlens_q_padded,
         unpacked_seqlen,
         vocab_start_index=vocab_parallel_rank * logits.shape[-1],
@@ -484,7 +491,8 @@ def prepare_packed_loss_input(
         cp_group=context_parallel_group,
         sampling_params=sampling_params,
         chunk_size=chunk_size if use_chunking else None,
-        target_is_pre_rolled=True,
+        target_is_pre_rolled=not input_is_prepacked,
+        return_packed_layout=input_is_prepacked,
     )
 
     # Match prepare_loss_input behavior for top-k/top-p filtered training:
@@ -500,7 +508,7 @@ def prepare_packed_loss_input(
             data["curr_logprobs_unfiltered"] = (
                 from_parallel_logits_to_logprobs_packed_sequences(
                     logits_for_logprobs,
-                    packed_rolled_targets,
+                    packed_targets,
                     cu_seqlens_q_padded,
                     unpacked_seqlen,
                     vocab_start_index=vocab_parallel_rank * logits.shape[-1],
@@ -510,7 +518,8 @@ def prepare_packed_loss_input(
                     cp_group=context_parallel_group,
                     sampling_params=None,
                     chunk_size=chunk_size if use_chunking else None,
-                    target_is_pre_rolled=True,
+                    target_is_pre_rolled=not input_is_prepacked,
+                    return_packed_layout=input_is_prepacked,
                 )
             )