feat: add Energon-owned SFT sequence packing

docs/design-docs/sequence-packing-and-dynamic-batching.md

@@ -88,7 +88,9 @@ We have the policy backends perform the actual packing because implementations c
 
 #### 2. Packing Algorithms (`nemo_rl/data/packing/algorithms.py`)
 
-Four packing algorithms are implemented, but we recommend you just use Modified First Fit Decreasing for the best packing efficiency:
+Six packing algorithms are implemented. Modified First Fit Decreasing is the
+default recommendation, and Energon-owned SFT packing supports all six through
+the same interface.
 
 ##### Concatenative Packer 
 - Sequential concatenation until bin capacity is reached
@@ -107,6 +109,12 @@ Four packing algorithms are implemented, but we recommend you just use Modified
   5. Greedy fit remaining items
   6. Apply FFD to leftovers
 
+##### Greedy Knapsack
+- Repeatedly selects the largest remaining sequence that fits in the current bin
+
+##### Balanced Greedy Knapsack
+- Places descending sequences into the least-full available bin
+
 ##### First Fit Decreasing (FFD)
 - Sort sequences by length (descending), place each in first fitting bin
 - O(n log n + n*m) where m = number of bins

docs/guides/sft.md

@@ -241,7 +241,17 @@ The processor runs inside Energon loader workers and returns the same tokenized
 
 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.
+Set `data.energon.packing_buffer_size` and enable fused
+`policy.sequence_packing` with any supported packing algorithm to let Energon
+form model-ready multimodal packs.
+Without an Energon packing buffer, SFTv2 currently requires fixed batching.
+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.
 

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

@@ -0,0 +1,14 @@
+defaults: ./vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.yaml
+sft:
+  max_num_steps: 50
+policy:
+  make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size},
+    ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}}
+  megatron_cfg:
+    tensor_model_parallel_size: 4
+    context_parallel_size: 2
+logger:
+  wandb_enabled: true
+  log_dir: logs/sft-nemotron-omni-30b-clevr-energon-packing-tp4ep8cp2
+  wandb:
+    name: sft-nemotron-omni-30b-clevr-energon-packing-tp4ep8cp2

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

@@ -0,0 +1,17 @@
+defaults: ./vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.yaml
+
+policy:
+  sequence_packing:
+    enabled: true
+    fuse_loss: true
+    algorithm: balanced_greedy_knapsack
+
+data:
+  energon:
+    packing_buffer_size: 64
+    max_samples_per_sequence: 16
+
+logger:
+  log_dir: logs/sft-nemotron-omni-30b-clevr-energon-packing
+  wandb:
+    name: sft-nemotron-omni-30b-clevr-energon-packing

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

@@ -0,0 +1,20 @@
+defaults: ./vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.yaml
+
+policy:
+  sequence_packing:
+    enabled: true
+    fuse_loss: true
+    algorithm: balanced_greedy_knapsack
+
+data:
+  energon:
+    packing_buffer_size: 64
+    max_samples_per_sequence: 16
+
+checkpointing:
+  checkpoint_dir: results/sft_${policy.model_name}_clevr_energon_packing
+
+logger:
+  log_dir: logs/sft-qwen2.5-vl-3b-clevr-energon-packing
+  wandb:
+    name: sft-qwen2.5-vl-3b-clevr-energon-packing

nemo_rl/algorithms/sft_v2.py

@@ -40,6 +40,7 @@
     DataLoaderPlacementPlan,
     resolve_topology_mapper,
 )
+from nemo_rl.data.packing import PackingAlgorithm
 from nemo_rl.data_plane.interfaces import LocalDataPlaneConfig
 from nemo_rl.distributed.named_sharding import REPLICATED_AXES
 from nemo_rl.distributed.virtual_cluster import (
@@ -196,6 +197,18 @@ def _setup_loaders(self) -> None:
             // self._placement_plan.logical_world_size,
             "max_sequence_length": config.data["max_input_seq_length"],
             "placement_fingerprint": self._placement_plan.placement_hash,
+            "packing_algorithm": config.policy["sequence_packing"]["algorithm"]
+            if config.data["energon"].packing_buffer_size is not None
+            else None,
+            # This caps sources per physical pack. Energon's similarly named
+            # max_samples_per_sequence instead controls sequential shard reads.
+            "max_sequences_per_bin": config.policy["sequence_packing"].get(
+                "max_sequences_per_bin"
+            ),
+            "sequence_length_pad_multiple": config.policy[
+                "make_sequence_length_divisible_by"
+            ],
+            "only_unmask_final": config.sft.only_unmask_final,
         }
         if self._loader_states is None:
             futures = self._trainer.worker_group.run_all_workers_single_data(
@@ -385,11 +398,25 @@ def setup_sft_v2(
         raise ValueError("SFTv2 requires data.backend=energon.")
     if not master_config.policy["megatron_cfg"]["enabled"]:
         raise ValueError("SFTv2 supports only the Megatron policy backend.")
-    if (
-        master_config.policy["sequence_packing"]["enabled"]
-        or master_config.policy["dynamic_batching"]["enabled"]
-    ):
-        raise ValueError("SFTv2 requires fixed NeMo-RL batching.")
+    sequence_packing = master_config.policy["sequence_packing"]
+    dynamic_batching = master_config.policy["dynamic_batching"]
+    energon_packing = master_config.data["energon"].packing_buffer_size is not None
+    if not energon_packing:
+        if sequence_packing["enabled"] or dynamic_batching["enabled"]:
+            raise ValueError("SFTv2 without Energon packing requires fixed batching.")
+    else:
+        if not sequence_packing["enabled"] or not sequence_packing.get(
+            "fuse_loss", False
+        ):
+            raise ValueError(
+                "Energon packing requires sequence_packing enabled with fuse_loss."
+            )
+        if sequence_packing.get("algorithm") not in {
+            algorithm.value for algorithm in PackingAlgorithm
+        }:
+            raise ValueError("Energon SFT requires a supported packing algorithm.")
+        if dynamic_batching["enabled"]:
+            raise ValueError("Energon packing does not support dynamic batching.")
     # SFTConfig carries validation knobs that default to on (val_period=10,
     # val_at_start=True) and this loop has no validation path, so reject them
     # rather than accepting a config whose validation silently never runs.
@@ -418,6 +445,43 @@ def setup_sft_v2(
     max_sequence_length = master_config.data["max_input_seq_length"]
     if max_sequence_length is None:
         raise ValueError("SFTv2 requires data.max_input_seq_length.")
+    if energon_packing:
+        megatron_cfg = master_config.policy["megatron_cfg"]
+        if (
+            megatron_cfg.get("moe_token_dispatcher_type") == "flex"
+            and megatron_cfg.get("moe_flex_dispatcher_backend") == "hybridep"
+        ):
+            raise ValueError("Energon packing does not support HybridEP flex dispatch.")
+
+        cp_size = megatron_cfg["context_parallel_size"]
+        tp_size = megatron_cfg["tensor_model_parallel_size"]
+        pad_multiple = master_config.policy["make_sequence_length_divisible_by"]
+        parallel_multiple = (2 * cp_size if cp_size > 1 else 1) * (
+            tp_size if tp_size > 1 and megatron_cfg["sequence_parallel"] else 1
+        )
+        if pad_multiple % parallel_multiple != 0:
+            raise ValueError(
+                "Energon packing requires make_sequence_length_divisible_by to "
+                f"be a multiple of {parallel_multiple}."
+            )
+        if max_sequence_length % pad_multiple != 0:
+            raise ValueError(
+                "Energon packing requires max_input_seq_length to be divisible by "
+                "make_sequence_length_divisible_by."
+            )
+
+        fp8_cfg = megatron_cfg.get("fp8_cfg") or {}
+        if fp8_cfg.get("enabled", False):
+            fp8_multiple = {
+                "blockwise": 128,
+                "mxfp8": 32,
+            }.get(fp8_cfg["fp8_recipe"], 16)
+            fp8_multiple *= parallel_multiple
+            if max_sequence_length % fp8_multiple != 0:
+                raise ValueError(
+                    "Energon packing requires max_input_seq_length to be divisible "
+                    f"by the FP8 packed-token alignment ({fp8_multiple})."
+                )
 
     processor = None
     tokenizer = tokenizer_or_processor

nemo_rl/data/energon/config.py

@@ -65,8 +65,26 @@ class EnergonLoaderConfig(BaseModel, extra="allow"):
     )
     num_workers: Annotated[int, Field(ge=0)] = 8
     shuffle_buffer_size: Annotated[int, Field(ge=0)] = 1000
-    max_samples_per_sequence: None = None
-    packing_buffer_size: None = None
+    max_samples_per_sequence: (
+        Annotated[
+            int,
+            Field(
+                ge=1,
+                description="Maximum sequential sample run used when sharding a dataset.",
+            ),
+        ]
+        | None
+    ) = None
+    packing_buffer_size: (
+        Annotated[
+            int,
+            Field(
+                ge=1,
+                description="Samples buffered by Energon for packing; None disables packing.",
+            ),
+        ]
+        | None
+    ) = None
     batch_grouping: Literal["auto"] = "auto"
     processor_adapter: Literal["hf_multimodal"] = "hf_multimodal"
     topology_mapper: Literal["default"] = "default"