From ad2110b8d2acc5bf9a82d237656a93bd21bc26fc Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 9 Sep 2026 21:08:06 -0700 Subject: [PATCH 01/11] feat: add Energon-owned SFT sequence packing Signed-off-by: rohitrango --- ...8g-megatron-tp8ep8-energon.v1.packing.yaml | 17 ++ ...r-1n2g-megatrontp1-energon.v1.packing.yaml | 20 ++ nemo_rl/algorithms/sft_v2.py | 32 ++- nemo_rl/data/energon/config.py | 4 +- nemo_rl/data/energon/multimodal/packing.py | 196 ++++++++++++++++++ .../energon/multimodal/task_encoders/base.py | 4 +- .../multimodal/task_encoders/generic_sft.py | 52 ++++- nemo_rl/data/energon/multimodal/types.py | 11 + nemo_rl/data/energon/sft_dataloader.py | 46 +++- nemo_rl/data/energon/sft_worker.py | 12 ++ nemo_rl/data/packing/__init__.py | 4 + nemo_rl/data/packing/algorithms.py | 88 +++++++- nemo_rl/models/megatron/data.py | 82 +++++++- nemo_rl/models/policy/tq_policy.py | 13 +- 14 files changed, 560 insertions(+), 21 deletions(-) create mode 100644 examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.yaml create mode 100644 examples/configs/recipes/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.packing.yaml create mode 100644 nemo_rl/data/energon/multimodal/packing.py diff --git a/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.yaml b/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.yaml new file mode 100644 index 00000000000..f542c82dc9f --- /dev/null +++ b/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 diff --git a/examples/configs/recipes/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.packing.yaml b/examples/configs/recipes/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.packing.yaml new file mode 100644 index 00000000000..b2ee0b5d43a --- /dev/null +++ b/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 diff --git a/nemo_rl/algorithms/sft_v2.py b/nemo_rl/algorithms/sft_v2.py index bf59b037174..7b938c862e3 100644 --- a/nemo_rl/algorithms/sft_v2.py +++ b/nemo_rl/algorithms/sft_v2.py @@ -196,6 +196,13 @@ 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, + "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 +392,26 @@ 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["algorithm"] not in { + "greedy_knapsack", + "balanced_greedy_knapsack", + }: + raise ValueError("Energon SFT supports only the two knapsack packers.") + 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. diff --git a/nemo_rl/data/energon/config.py b/nemo_rl/data/energon/config.py index 9ffb52422f7..ee7f47d10e4 100644 --- a/nemo_rl/data/energon/config.py +++ b/nemo_rl/data/energon/config.py @@ -65,8 +65,8 @@ 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)] | None = None + packing_buffer_size: Annotated[int, Field(ge=1)] | None = None batch_grouping: Literal["auto"] = "auto" processor_adapter: Literal["hf_multimodal"] = "hf_multimodal" topology_mapper: Literal["default"] = "default" diff --git a/nemo_rl/data/energon/multimodal/packing.py b/nemo_rl/data/energon/multimodal/packing.py new file mode 100644 index 00000000000..d51354118a5 --- /dev/null +++ b/nemo_rl/data/energon/multimodal/packing.py @@ -0,0 +1,196 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Energon-owned selection and materialization of multimodal SFT packs.""" + +from __future__ import annotations + +from typing import Any + +import torch +from transformers import PreTrainedTokenizerBase + +from nemo_rl.data.energon.multimodal.types import EncodedSFTSample, PackedSFTSample +from nemo_rl.data.llm_message_utils import ( + add_loss_mask_to_message_log, + batched_message_log_to_flat_message, + message_log_to_flat_messages, +) +from nemo_rl.data.packing import SequencePacker +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +def _cost(sample: EncodedSFTSample, multiple: int) -> int: + if sample.packing_cost < sample.length: + raise ValueError(f"Invalid packing cost for sample {sample.sample_key!r}.") + return ((sample.packing_cost + multiple - 1) // multiple) * multiple + + +def select_samples_to_pack( + samples: list[EncodedSFTSample], + *, + packer: SequencePacker, + sequence_length_pad_multiple: int, +) -> list[list[EncodedSFTSample]]: + """Group compatible sources and run the configured packer.""" + if sequence_length_pad_multiple <= 0: + raise ValueError("Packing alignment must be positive.") + groups: dict[tuple[Any, ...], list[EncodedSFTSample]] = {} + for sample in samples: + groups.setdefault(sample.group_key, []).append(sample) + result: list[list[EncodedSFTSample]] = [] + for group in groups.values(): + bins = packer.pack( + [_cost(sample, sequence_length_pad_multiple) for sample in group] + ) + indexes = [index for bin_indexes in bins for index in bin_indexes] + if sorted(indexes) != list(range(len(group))): + raise RuntimeError("Packing must preserve every source exactly once.") + result.extend([[group[index] for index in bin_indexes] for bin_indexes in bins]) + return result + + +def pack_selected_samples( + samples: list[EncodedSFTSample], + *, + pack_capacity: int, + sequence_length_pad_multiple: int, +) -> PackedSFTSample: + """Turn one selected source group into a physical pack.""" + if not samples or any( + sample.group_key != samples[0].group_key for sample in samples + ): + raise ValueError("A physical pack needs compatible sources.") + padded_lengths = [_cost(sample, sequence_length_pad_multiple) for sample in samples] + if sum(padded_lengths) > pack_capacity: + raise ValueError("Selected sources exceed the pack capacity.") + return PackedSFTSample.derive_from( + samples[0], + __key__=",".join(sample.sample_key for sample in samples), + samples=list(samples), + source_padded_lengths=padded_lengths, + group_key=samples[0].group_key, + pack_capacity=pack_capacity, + ) + + +def prepare_packed_sft_batch( + packs: list[PackedSFTSample], + *, + tokenizer: PreTrainedTokenizerBase, + only_unmask_final: bool, +) -> BatchedDataDict[Any]: + """Create model tensors for a batch of physical Energon packs.""" + if not packs or tokenizer.pad_token_id is None: + raise ValueError("Packed SFT requires packs and a tokenizer pad token.") + capacities = {pack.pack_capacity for pack in packs} + if len(capacities) != 1: + raise ValueError("All physical packs in a batch need one capacity.") + capacity = capacities.pop() + packed_logs: list[list[dict[str, Any]]] = [] + boundaries: list[torch.Tensor] = [] + padded_boundaries: list[torch.Tensor] = [] + source_ids: list[list[str]] = [] + + for pack in packs: + logs = [ + [dict(message) for message in sample.message_log] for sample in pack.samples + ] + add_loss_mask_to_message_log( + logs, roles_to_train_on=["assistant"], only_unmask_final=only_unmask_final + ) + lengths: list[int] = [] + combined: list[dict[str, Any]] = [] + token_dtype = torch.long + for log, sample, padded_length in zip( + logs, pack.samples, pack.source_padded_lengths + ): + tokens = message_log_to_flat_messages(log).get("token_ids") + if not isinstance(tokens, torch.Tensor) or tokens.numel() == 0: + raise TypeError("Packed SFT sources require token tensors.") + token_dtype = tokens.dtype + length = tokens.shape[0] + lengths.append(length) + templates = { + key: value + for message in log + for key, value in message.items() + if key not in {"token_ids", "token_loss_mask"} + and isinstance(value, torch.Tensor) + } + for message in log: + message["token_loss_mask"] = ( + message["token_loss_mask"] * sample.loss_multiplier + ) + for key, template in templates.items(): + message.setdefault( + key, + torch.zeros( + (message["token_ids"].shape[0], *template.shape[1:]), + dtype=template.dtype, + ), + ) + log[0]["token_loss_mask"][0] = 0 + padding = padded_length - length + if padding < 0: + raise ValueError("A source exceeds its padded length.") + if padding: + pad_message = { + "role": "padding", + "token_ids": torch.full( + (padding,), tokenizer.pad_token_id, dtype=token_dtype + ), + "token_loss_mask": torch.zeros(padding, dtype=torch.float32), + } + pad_message.update( + { + key: torch.zeros((padding, *value.shape[1:]), dtype=value.dtype) + for key, value in templates.items() + } + ) + log.append(pad_message) + combined.extend(log) + tail = capacity - sum(pack.source_padded_lengths) + if tail: + combined.append( + { + "role": "padding", + "token_ids": torch.full( + (tail,), tokenizer.pad_token_id, dtype=token_dtype + ), + "token_loss_mask": torch.zeros(tail, dtype=torch.float32), + } + ) + packed_logs.append(combined) + boundaries.append( + torch.tensor( + [0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32 + ) + ) + padded = [0, *torch.tensor(pack.source_padded_lengths).cumsum(0).tolist()] + padded[-1] = capacity + padded_boundaries.append(torch.tensor(padded, dtype=torch.int32)) + source_ids.append([sample.sample_key for sample in pack.samples]) + + flat, input_lengths = batched_message_log_to_flat_message( + packed_logs, pad_value_dict={"token_ids": tokenizer.pad_token_id} + ) + prepared = BatchedDataDict( + { + "input_ids": flat["token_ids"], + "input_lengths": input_lengths, + "token_mask": flat["token_loss_mask"], + "sample_mask": flat["token_loss_mask"].bool().any(1).float(), + "cu_seqlens": boundaries, + "cu_seqlens_padded": padded_boundaries, + "source_ids": source_ids, + } + ) + prepared.update(flat.get_multimodal_dict(as_tensors=False)) + return prepared + + +__all__ = [ + "pack_selected_samples", + "prepare_packed_sft_batch", + "select_samples_to_pack", +] diff --git a/nemo_rl/data/energon/multimodal/task_encoders/base.py b/nemo_rl/data/energon/multimodal/task_encoders/base.py index f51455a4d26..837231f32e2 100644 --- a/nemo_rl/data/energon/multimodal/task_encoders/base.py +++ b/nemo_rl/data/energon/multimodal/task_encoders/base.py @@ -14,7 +14,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Sequence -from typing import Any, ClassVar, TypeAlias +from typing import Any, TypeAlias from megatron.energon import Cooker, CrudeSample, DefaultTaskEncoder @@ -40,8 +40,6 @@ class BaseSFTTaskEncoder( ): """Common SFT lifecycle shared by the Energon task encoders.""" - sample_schema: ClassVar[str] - def __init__( self, *, diff --git a/nemo_rl/data/energon/multimodal/task_encoders/generic_sft.py b/nemo_rl/data/energon/multimodal/task_encoders/generic_sft.py index 45cb2d54a3b..dd97220be4e 100644 --- a/nemo_rl/data/energon/multimodal/task_encoders/generic_sft.py +++ b/nemo_rl/data/energon/multimodal/task_encoders/generic_sft.py @@ -27,6 +27,11 @@ ALL_MODEL_FAMILIES, supports_model_families, ) +from nemo_rl.data.energon.multimodal.packing import ( + pack_selected_samples, + prepare_packed_sft_batch, + select_samples_to_pack, +) from nemo_rl.data.energon.multimodal.task_encoders.base import ( BaseSFTTaskEncoder, SFTCooker, @@ -37,10 +42,12 @@ from nemo_rl.data.energon.multimodal.types import ( CanonicalSFTSample, EncodedSFTSample, + PackedSFTSample, ) from nemo_rl.data.interfaces import TaskDataSpec from nemo_rl.data.llm_message_utils import get_formatted_message_log from nemo_rl.data.multimodal_utils import PackedTensor +from nemo_rl.data.packing import SequencePacker from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -251,7 +258,6 @@ class GenericSFTTaskEncoder(BaseSFTTaskEncoder): # which would let a systematically broken dataset retry forever. 1 fails on the # first bad sample; raise it to tolerate transient decode errors. __default_failure_tolerance__ = 1 - sample_schema = "nemo_rl.sft.encoded.v1" # Match the existing HF VLM path. Its processor expects PIL RGB images. decoder = SampleDecoder(image_decode="pilrgb") @@ -261,10 +267,18 @@ def __init__( adapter: SFTProcessorAdapter, cooker_functions: Sequence[SFTCooker], include_source_ids: bool, + packer: SequencePacker | None = None, + tokenizer: Any | None = None, + sequence_length_pad_multiple: int = 1, + only_unmask_final: bool = False, ) -> None: super().__init__(cooker_functions=cooker_functions) self.adapter = adapter self.include_source_ids = include_source_ids + self.packer = packer + self.tokenizer = tokenizer + self.sequence_length_pad_multiple = sequence_length_pad_multiple + self.only_unmask_final = only_unmask_final @stateless def preencode_sample(self, sample: CanonicalSFTSample) -> EncodedSFTSample: @@ -275,12 +289,44 @@ def postencode_sample(self, sample: EncodedSFTSample) -> EncodedSFTSample: return sample def batch_group_criterion( - self, sample: EncodedSFTSample + self, sample: EncodedSFTSample | PackedSFTSample ) -> tuple[tuple[Any, ...], None]: return sample.group_key, None + def select_samples_to_pack( + self, samples: list[EncodedSFTSample] + ) -> list[list[EncodedSFTSample]]: + if self.packer is None: + raise RuntimeError("Energon packing is not configured.") + return select_samples_to_pack( + samples, + packer=self.packer, + sequence_length_pad_multiple=self.sequence_length_pad_multiple, + ) + + def pack_selected_samples(self, samples: list[EncodedSFTSample]) -> PackedSFTSample: + if self.packer is None: + raise RuntimeError("Energon packing is not configured.") + return pack_selected_samples( + samples, + pack_capacity=self.packer.bin_capacity, + sequence_length_pad_multiple=self.sequence_length_pad_multiple, + ) + @stateless - def batch(self, samples: list[EncodedSFTSample]) -> BatchedDataDict[Any]: + def batch( + self, samples: list[EncodedSFTSample | PackedSFTSample] + ) -> BatchedDataDict[Any]: + if samples and isinstance(samples[0], PackedSFTSample): + if not all(isinstance(sample, PackedSFTSample) for sample in samples): + raise TypeError("Energon batches cannot mix packed and unpacked rows.") + if self.tokenizer is None: + raise RuntimeError("Packed SFT requires a tokenizer.") + return prepare_packed_sft_batch( + cast(list[PackedSFTSample], samples), + tokenizer=self.tokenizer, + only_unmask_final=self.only_unmask_final, + ) if not all(isinstance(sample, EncodedSFTSample) for sample in samples): raise TypeError("Energon SFT batches accept only encoded samples.") encoded_samples = cast(list[EncodedSFTSample], samples) diff --git a/nemo_rl/data/energon/multimodal/types.py b/nemo_rl/data/energon/multimodal/types.py index 80cb548bf57..3dbac336fd2 100644 --- a/nemo_rl/data/energon/multimodal/types.py +++ b/nemo_rl/data/energon/multimodal/types.py @@ -72,11 +72,22 @@ class EncodedSFTSample(Sample): pending_sample: CanonicalSFTSample | None = None +@edataclass +class PackedSFTSample(Sample): + """One physical pack of compatible encoded conversations.""" + + samples: list[EncodedSFTSample] + source_padded_lengths: list[int] + group_key: tuple[Any, ...] + pack_capacity: int + + __all__ = [ "CanonicalSFTSample", "EncodedSFTSample", "FrozenMediaMetadata", "MediaRef", "MediaMetadataValue", + "PackedSFTSample", "freeze_media_metadata", ] diff --git a/nemo_rl/data/energon/sft_dataloader.py b/nemo_rl/data/energon/sft_dataloader.py index 29aa6016797..8d423b1f781 100644 --- a/nemo_rl/data/energon/sft_dataloader.py +++ b/nemo_rl/data/energon/sft_dataloader.py @@ -46,6 +46,7 @@ build_processor_adapter, ) from nemo_rl.data.energon.multimodal.types import CanonicalSFTSample +from nemo_rl.data.packing import get_packer from nemo_rl.distributed.batched_data_dict import BatchedDataDict _V2_STATE_FORMAT_VERSION = 2 @@ -286,9 +287,11 @@ def _loader_identity( batch_size: int, shuffle: bool | None, topology: dict[str, Any], + packing_algorithm: str | None, + sequence_length_pad_multiple: int, ) -> dict[str, Any]: """Describe what a restored loader must still agree with.""" - return { + identity = { "source": source.model_dump(mode="json"), "loader": loader_config.model_dump(mode="json"), "adapter": adapter_fingerprint, @@ -311,6 +314,12 @@ def _loader_identity( ), "topology": topology, } + if packing_algorithm is not None: + identity.update( + packing_algorithm=packing_algorithm, + sequence_length_pad_multiple=sequence_length_pad_multiple, + ) + return identity def _worker_config( @@ -338,6 +347,11 @@ def _task_encoder( loader_config: EnergonLoaderConfig, adapter: Any, include_source_ids: bool, + packing_algorithm: str | None, + max_sequence_length: int, + sequence_length_pad_multiple: int, + tokenizer: Any, + only_unmask_final: bool, ) -> BaseSFTTaskEncoder: cooker_functions = [ Cooker( @@ -353,12 +367,26 @@ def _task_encoder( Any, TASK_ENCODER_REGISTRY.resolve(loader_config.task_encoder.name) ) encoder_options: dict[str, Any] = dict(loader_config.task_encoder.options) + packer = ( + get_packer( + packing_algorithm, + max_sequence_length, + max_sequences_per_bin=loader_config.max_samples_per_sequence, + ) + if loader_config.packing_buffer_size is not None + and packing_algorithm is not None + else None + ) return cast( BaseSFTTaskEncoder, encoder_type( adapter=adapter, cooker_functions=cooker_functions, include_source_ids=include_source_ids, + packer=packer, + tokenizer=tokenizer, + sequence_length_pad_multiple=sequence_length_pad_multiple, + only_unmask_final=only_unmask_final, **encoder_options, ), ) @@ -375,6 +403,9 @@ def build_energon_sft_loader( logical_rank: int, logical_world_size: int, placement_fingerprint: str, + packing_algorithm: str | None = None, + sequence_length_pad_multiple: int = 1, + only_unmask_final: bool = False, ) -> EnergonSFTDataLoader: """Build one loader for an explicit logical data shard and split.""" if "energon" not in data_config: @@ -388,6 +419,8 @@ def build_energon_sft_loader( resolved_source = _source_config(source, name=split_role) loader_config = _loader_config(data_config["energon"]) + if loader_config.packing_buffer_size is not None and packing_algorithm is None: + raise ValueError("Energon packing requires a packing algorithm.") adapter = build_processor_adapter( processor_adapter=loader_config.processor_adapter, processor=processor, @@ -400,6 +433,11 @@ def build_energon_sft_loader( loader_config=loader_config, adapter=adapter, include_source_ids=True, + packing_algorithm=packing_algorithm, + max_sequence_length=max_sequence_length, + sequence_length_pad_multiple=sequence_length_pad_multiple, + tokenizer=processor.tokenizer, + only_unmask_final=only_unmask_final, ) worker_config = _worker_config( loader_config, @@ -421,9 +459,10 @@ def build_energon_sft_loader( worker_config=worker_config, batch_size=batch_size, batch_drop_last=True, + packing_buffer_size=loader_config.packing_buffer_size, shuffle_buffer_size=(loader_config.shuffle_buffer_size), shuffle_over_epochs_multiplier=1, - max_samples_per_sequence=None, + max_samples_per_sequence=loader_config.max_samples_per_sequence, virtual_epoch_length=resolved_source.virtual_epoch_length, task_encoder=task_encoder, ) @@ -434,6 +473,7 @@ def build_energon_sft_loader( worker_config=worker_config, batch_size=batch_size, batch_drop_last=False, + packing_buffer_size=loader_config.packing_buffer_size, limit=resolved_source.limit, task_encoder=task_encoder, ) @@ -466,6 +506,8 @@ def build_energon_sft_loader( logical_rank=logical_rank, logical_world_size=logical_world_size, ), + packing_algorithm=packing_algorithm, + sequence_length_pad_multiple=sequence_length_pad_multiple, ), ) diff --git a/nemo_rl/data/energon/sft_worker.py b/nemo_rl/data/energon/sft_worker.py index 8af6ccb7af9..994bd3b24c8 100644 --- a/nemo_rl/data/energon/sft_worker.py +++ b/nemo_rl/data/energon/sft_worker.py @@ -31,6 +31,7 @@ ) from nemo_rl.data.energon.sft_types import StepEnvelope from nemo_rl.data_plane.adapters.local import local_batch_to_tensordict +from nemo_rl.data_plane.schema import MICRO_BATCH_INDICES, MICRO_BATCH_LENGTHS from nemo_rl.models.policy.utils import get_runtime_env_for_policy_worker from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, @@ -64,6 +65,9 @@ def setup_sft_dataloader( batch_size: int, max_sequence_length: int, placement_fingerprint: str, + packing_algorithm: str | None, + sequence_length_pad_multiple: int, + only_unmask_final: bool, restored_state: Optional[dict[str, Any]] = None, ) -> bool: """Build the train loader on the TP0/PP0/CP0 rank of this DP replica.""" @@ -86,6 +90,9 @@ def setup_sft_dataloader( logical_rank=logical_rank, logical_world_size=logical_world_size, placement_fingerprint=placement_fingerprint, + packing_algorithm=packing_algorithm, + sequence_length_pad_multiple=sequence_length_pad_multiple, + only_unmask_final=only_unmask_final, ) if restored_state is not None: self._sft_loader.load_state_dict(restored_state) @@ -164,6 +171,11 @@ def load_next_sft_batch( (sample_mask.unsqueeze(-1) * prepared["token_mask"][:, 1:]).sum().item() ) extra_info = dict(published_meta.extra_info) + if "cu_seqlens" in prepared: + extra_info[MICRO_BATCH_INDICES] = [ + [[index, index + 1] for index in range(batch_size)] + ] + extra_info[MICRO_BATCH_LENGTHS] = [list(lengths)] if make_sequence_length_divisible_by > 1: extra_info["pad_to_multiple"] = int(make_sequence_length_divisible_by) envelope = StepEnvelope( diff --git a/nemo_rl/data/packing/__init__.py b/nemo_rl/data/packing/__init__.py index a955f681cce..9e539be501b 100644 --- a/nemo_rl/data/packing/__init__.py +++ b/nemo_rl/data/packing/__init__.py @@ -13,9 +13,11 @@ # limitations under the License. from nemo_rl.data.packing.algorithms import ( + BalancedGreedyKnapsackPacker, ConcatenativePacker, FirstFitDecreasingPacker, FirstFitShufflePacker, + GreedyKnapsackPacker, ModifiedFirstFitDecreasingPacker, PackingAlgorithm, SequencePacker, @@ -24,11 +26,13 @@ from nemo_rl.data.packing.metrics import PackingMetrics __all__ = [ + "BalancedGreedyKnapsackPacker", "PackingAlgorithm", "SequencePacker", "ConcatenativePacker", "FirstFitDecreasingPacker", "FirstFitShufflePacker", + "GreedyKnapsackPacker", "ModifiedFirstFitDecreasingPacker", "get_packer", "PackingMetrics", diff --git a/nemo_rl/data/packing/algorithms.py b/nemo_rl/data/packing/algorithms.py index af36c9b947e..1a2c97edead 100644 --- a/nemo_rl/data/packing/algorithms.py +++ b/nemo_rl/data/packing/algorithms.py @@ -18,7 +18,7 @@ import math import random from abc import ABC, abstractmethod -from bisect import bisect +from bisect import bisect, bisect_right from typing import Dict, List, Optional, Tuple, Type, Union @@ -29,6 +29,8 @@ class PackingAlgorithm(enum.Enum): FIRST_FIT_DECREASING = "first_fit_decreasing" FIRST_FIT_SHUFFLE = "first_fit_shuffle" MODIFIED_FIRST_FIT_DECREASING = "modified_first_fit_decreasing" + GREEDY_KNAPSACK = "greedy_knapsack" + BALANCED_GREEDY_KNAPSACK = "balanced_greedy_knapsack" class SequencePacker(ABC): @@ -293,6 +295,88 @@ def _estimate_bins_needed(self, sequence_lengths: List[int]) -> int: return max(1, math.ceil(total_length / self.bin_capacity)) +class GreedyKnapsackPacker(SequencePacker): + """Repeatedly take the largest remaining sequence that fits.""" + + def _pack_implementation(self, sequence_lengths: List[int]) -> List[List[int]]: + self._validate_sequence_lengths(sequence_lengths) + remaining = sorted( + (length, -index, index) for index, length in enumerate(sequence_lengths) + ) + bins: List[List[int]] = [] + while remaining: + current: List[int] = [] + capacity = self.bin_capacity + while ( + self.max_sequences_per_bin is None + or len(current) < self.max_sequences_per_bin + ): + fit = bisect_right(remaining, (capacity, 1, len(sequence_lengths))) + if fit == 0: + break + length, _, index = remaining.pop(fit - 1) + capacity -= length + current.append(index) + bins.append(current) + return bins + + +class BalancedGreedyKnapsackPacker(SequencePacker): + """Place descending sequences into the least-full available bin.""" + + def __init__( + self, + bin_capacity: int, + collect_metrics: bool = False, + min_bin_count: Optional[int] = None, + bin_count_multiple: Optional[int] = None, + max_sequences_per_bin: Optional[int] = None, + balanced_knapsack_delta: int = 20, + ) -> None: + super().__init__( + bin_capacity, + collect_metrics, + min_bin_count, + bin_count_multiple, + max_sequences_per_bin, + ) + if balanced_knapsack_delta < 0: + raise ValueError("balanced_knapsack_delta must be nonnegative") + self.balanced_knapsack_delta = balanced_knapsack_delta + + def _pack_implementation(self, sequence_lengths: List[int]) -> List[List[int]]: + self._validate_sequence_lengths(sequence_lengths) + if not sequence_lengths: + return [] + count = math.ceil(sum(sequence_lengths) / self.bin_capacity) + bins: List[List[int]] = [ + [] for _ in range(count + self.balanced_knapsack_delta) + ] + loads = [0] * len(bins) + for index in sorted( + range(len(sequence_lengths)), + key=sequence_lengths.__getitem__, + reverse=True, + ): + candidates = [ + i + for i, load in enumerate(loads) + if load + sequence_lengths[index] <= self.bin_capacity + and ( + self.max_sequences_per_bin is None + or len(bins[i]) < self.max_sequences_per_bin + ) + ] + if not candidates: + bins.append([]) + loads.append(0) + candidates = [len(bins) - 1] + target = min(candidates, key=loads.__getitem__) + bins[target].append(index) + loads[target] += sequence_lengths[index] + return [bin_indexes for bin_indexes in bins if bin_indexes] + + class ConcatenativePacker(SequencePacker): """Concatenative packing algorithm. @@ -700,6 +784,8 @@ def get_packer( PackingAlgorithm.FIRST_FIT_DECREASING: FirstFitDecreasingPacker, PackingAlgorithm.FIRST_FIT_SHUFFLE: FirstFitShufflePacker, PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING: ModifiedFirstFitDecreasingPacker, + PackingAlgorithm.GREEDY_KNAPSACK: GreedyKnapsackPacker, + PackingAlgorithm.BALANCED_GREEDY_KNAPSACK: BalancedGreedyKnapsackPacker, } # Convert string to enum if needed diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index 6ef2eaa7f30..c1c72205165 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -275,10 +275,25 @@ def get_microbatch_iterator( if seq_length_key is None and cfg["sequence_packing"]["enabled"]: seq_length_key = "input_lengths" + prepacked = "cu_seqlens" in data or "cu_seqlens_padded" in data + if prepacked and not all( + key in data for key in ("cu_seqlens", "cu_seqlens_padded") + ): + raise ValueError("Prepacked input requires both cumulative boundary fields.") + if prepacked and ( + not cfg["sequence_packing"]["enabled"] + or not cfg["sequence_packing"].get("fuse_loss", False) + or cfg["dynamic_batching"]["enabled"] + ): + raise ValueError("Prepacked input requires fused sequence packing only.") if not cfg["sequence_packing"]["enabled"]: pad_factor = _get_non_packed_sequence_pad_factor(cfg) - if cfg["dynamic_batching"]["enabled"]: + if prepacked: + raw_iterator = data.make_microbatch_iterator(1) + data_iterator_len = data.size + micro_batch_size = 1 + elif cfg["dynamic_batching"]["enabled"]: raw_iterator = data.make_microbatch_iterator_with_dynamic_shapes() data_iterator_len = data.get_microbatch_iterator_dynamic_shapes_len() elif cfg["sequence_packing"]["enabled"]: @@ -350,6 +365,57 @@ def get_ltor_masks_and_position_ids(*args: Any, **kwargs: Any) -> Any: return _impl(*args, **kwargs) +def _prepacked_boundary( + data: BatchedDataDict[Any], key: str, device: torch.device +) -> torch.Tensor: + value = data[key] + if isinstance(value, list): + if len(value) != 1: + raise ValueError(f"{key} must describe one physical pack.") + value = value[0] + elif torch.is_tensor(value) and value.ndim == 2 and value.shape[0] == 1: + value = value[0] + if not torch.is_tensor(value) or value.ndim != 1: + raise ValueError(f"{key} must be a one-dimensional tensor.") + return value.to(device=device, dtype=torch.int32) + + +def _prepare_prepacked( + data: BatchedDataDict[Any], +) -> tuple[torch.Tensor, torch.Tensor, PackedSeqParams, torch.Tensor]: + input_ids = data["input_ids"] + if not torch.is_tensor(input_ids) or input_ids.shape[0] != 1: + raise ValueError("Prepacked input_ids must contain one physical row.") + cu = _prepacked_boundary(data, "cu_seqlens", input_ids.device) + padded = _prepacked_boundary(data, "cu_seqlens_padded", input_ids.device) + source_lengths = cu[1:] - cu[:-1] + padded_lengths = padded[1:] - padded[:-1] + if ( + cu.shape != padded.shape + or cu.numel() < 2 + or int(cu[0]) != 0 + or int(padded[0]) != 0 + or int(padded[-1]) != input_ids.shape[1] + or bool((source_lengths <= 0).any()) + or bool((source_lengths > padded_lengths).any()) + ): + raise ValueError("Invalid prepacked source boundaries.") + if get_context_parallel_world_size() != 1: + raise NotImplementedError("Energon-owned packing currently requires CP=1.") + params = PackedSeqParams( + cu_seqlens_q=padded, + cu_seqlens_kv=padded, + cu_seqlens_q_padded=padded, + cu_seqlens_kv_padded=padded, + max_seqlen_q=int(padded_lengths.max()), + max_seqlen_kv=int(padded_lengths.max()), + pad_between_seqs=False, + qkv_format="thd", + total_tokens=input_ids.shape[1], + ) + return input_ids, input_ids, params, padded + + def process_microbatch( data_dict: BatchedDataDict[Any], seq_length_key: Optional[str] = None, @@ -415,7 +481,19 @@ def process_microbatch( # Get sequence lengths and context parallel size seq_lengths = data_dict[seq_length_key] - if delegate_pack_to_model: + prepacked = "cu_seqlens" in data_dict + if prepacked: + if delegate_pack_to_model: + raise ValueError("Prepacked input cannot use model-owned packing.") + ( + input_ids, + input_ids_cp_sharded, + packed_seq_params, + cu_seqlens_padded, + ) = _prepare_prepacked(data_dict) + position_ids = None + attention_mask = None + elif delegate_pack_to_model: has_mtp_loss_mask = "mtp_loss_mask" in data_dict assert not has_mtp_loss_mask or delegate_mtp_loss_mask_to_model, ( "MTP training requires a self-packing VLM that advertises " diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index 47d4b2d8391..13be7340246 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -47,6 +47,8 @@ from nemo_rl.data_plane.schema import ( DP_TRAIN_FIELDS, GLOBAL_FORWARD_PAD_SEQLEN, + MICRO_BATCH_INDICES, + MICRO_BATCH_LENGTHS, LP_SEED_FIELDS, ROUTE_PASSTHROUGH_FLAG, ROUTE_PLAN_TAG, @@ -558,10 +560,15 @@ def train_placed_microbatches( f"got {len(dp_metas)} batches for dp_world={dp_world}." ) spa, dba = self._packing_args("train_mb_tokens") - if spa is not None or dba is not None: + if dba is not None: + raise ValueError("Placed metadata does not support dynamic batching.") + if spa is not None and any( + MICRO_BATCH_INDICES not in meta.extra_info + or MICRO_BATCH_LENGTHS not in meta.extra_info + for meta in dp_metas + ): raise ValueError( - "Placed metadata supports fixed batches only. Disable NeMo-RL " - "sequence packing and dynamic batching." + "Placed packed metadata requires producer microbatch shapes." ) train_metas = [ replace(meta, task_name="train") From 127785797a6c80be43f7b473cce6e1f7f60dedfd Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 9 Sep 2026 21:32:25 -0700 Subject: [PATCH 02/11] fix: transport packed boundaries across TP ranks Signed-off-by: rohitrango --- nemo_rl/data/energon/multimodal/packing.py | 7 +++++-- nemo_rl/models/megatron/data.py | 6 ++++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/nemo_rl/data/energon/multimodal/packing.py b/nemo_rl/data/energon/multimodal/packing.py index d51354118a5..67ac02c8c12 100644 --- a/nemo_rl/data/energon/multimodal/packing.py +++ b/nemo_rl/data/energon/multimodal/packing.py @@ -15,6 +15,7 @@ batched_message_log_to_flat_message, message_log_to_flat_messages, ) +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data.packing import SequencePacker from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -180,8 +181,10 @@ def prepare_packed_sft_batch( "input_lengths": input_lengths, "token_mask": flat["token_loss_mask"], "sample_mask": flat["token_loss_mask"].bool().any(1).float(), - "cu_seqlens": boundaries, - "cu_seqlens_padded": padded_boundaries, + # TP replica broadcast rejects tensor-bearing Python lists; + # PackedTensor carries these jagged per-pack boundaries over NCCL. + "cu_seqlens": PackedTensor(boundaries, dim_to_pack=0), + "cu_seqlens_padded": PackedTensor(padded_boundaries, dim_to_pack=0), "source_ids": source_ids, } ) diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index c1c72205165..43005ed7d9b 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -29,7 +29,7 @@ from megatron.core.utils import StragglerDetector from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType -from nemo_rl.data.multimodal_utils import PACKED_MULTIMODAL_FIELDS +from nemo_rl.data.multimodal_utils import PACKED_MULTIMODAL_FIELDS, PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank from nemo_rl.models.megatron.common import _round_up_to_multiple @@ -369,7 +369,9 @@ def _prepacked_boundary( data: BatchedDataDict[Any], key: str, device: torch.device ) -> torch.Tensor: value = data[key] - if isinstance(value, list): + if isinstance(value, PackedTensor): + value = value.as_tensor() + elif isinstance(value, list): if len(value) != 1: raise ValueError(f"{key} must describe one physical pack.") value = value[0] From 2615ba98df3bfdf8efb9854b4d51ec3e777e21d2 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 9 Sep 2026 21:44:25 -0700 Subject: [PATCH 03/11] test: cover Energon-owned sequence packing Signed-off-by: rohitrango --- tests/unit/data/packing/test_knapsack.py | 44 +++++++++ tests/unit/data/test_energon_packing.py | 109 +++++++++++++++++++++++ 2 files changed, 153 insertions(+) create mode 100644 tests/unit/data/packing/test_knapsack.py create mode 100644 tests/unit/data/test_energon_packing.py diff --git a/tests/unit/data/packing/test_knapsack.py b/tests/unit/data/packing/test_knapsack.py new file mode 100644 index 00000000000..fc9168a4e4c --- /dev/null +++ b/tests/unit/data/packing/test_knapsack.py @@ -0,0 +1,44 @@ +import pytest + +from nemo_rl.data.packing import ( + BalancedGreedyKnapsackPacker, + GreedyKnapsackPacker, + PackingAlgorithm, + get_packer, +) + + +@pytest.mark.parametrize( + ("algorithm", "packer_type"), + [ + (PackingAlgorithm.GREEDY_KNAPSACK, GreedyKnapsackPacker), + (PackingAlgorithm.BALANCED_GREEDY_KNAPSACK, BalancedGreedyKnapsackPacker), + ], +) +def test_factory_builds_knapsack_packers(algorithm, packer_type) -> None: + assert isinstance(get_packer(algorithm, 10), packer_type) + assert isinstance(get_packer(algorithm.value, 10), packer_type) + + +def test_greedy_knapsack_takes_largest_remaining_item_that_fits() -> None: + assert GreedyKnapsackPacker(10).pack([6, 5, 4, 3, 2]) == [ + [0, 2], + [1, 3, 4], + ] + + +def test_balanced_knapsack_spreads_equal_items_across_minimum_bins() -> None: + packer = BalancedGreedyKnapsackPacker(8, balanced_knapsack_delta=0) + + assert packer.pack([4, 4, 4, 4]) == [[0, 2], [1, 3]] + + +@pytest.mark.parametrize( + "packer", + [GreedyKnapsackPacker(10), BalancedGreedyKnapsackPacker(10)], +) +def test_knapsack_packers_keep_common_interface_constraints(packer) -> None: + packer.max_sequences_per_bin = 1 + assert packer.pack([4, 3, 2]) == [[0], [1], [2]] + with pytest.raises(ValueError, match="exceeds bin capacity"): + packer.pack([11]) diff --git a/tests/unit/data/test_energon_packing.py b/tests/unit/data/test_energon_packing.py new file mode 100644 index 00000000000..4050bc0b5de --- /dev/null +++ b/tests/unit/data/test_energon_packing.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import pytest +import torch + +pytest.importorskip("megatron.energon") +pytest.importorskip("megatron.core") + +pytestmark = pytest.mark.mcore + +from nemo_rl.data.energon.multimodal.packing import ( # noqa: E402 + pack_selected_samples, + prepare_packed_sft_batch, + select_samples_to_pack, +) +from nemo_rl.data.energon.multimodal.types import EncodedSFTSample # noqa: E402 +from nemo_rl.data.multimodal_utils import PackedTensor # noqa: E402 +from nemo_rl.data.packing import GreedyKnapsackPacker # noqa: E402 +from nemo_rl.models.megatron.data import _prepacked_boundary # noqa: E402 + + +class _Tokenizer: + pad_token_id = 0 + + +def _sample( + key: str, + length: int, + *, + group: str = "text", + packing_cost: int | None = None, +) -> EncodedSFTSample: + user_length = max(1, length - 2) + return EncodedSFTSample( + __key__=key, + __restore_key__=(key,), + message_log=[ + {"role": "user", "token_ids": torch.arange(1, user_length + 1)}, + { + "role": "assistant", + "token_ids": torch.arange(user_length + 1, length + 1), + }, + ], + length=length, + packing_cost=length if packing_cost is None else packing_cost, + loss_multiplier=1.0, + group_key=(group,), + sample_key=key, + ) + + +def test_selection_uses_aligned_costs_and_keeps_groups_separate() -> None: + samples = [ + _sample("s0", 5), + _sample("s1", 3), + _sample("s2", 3, group="image"), + ] + + selected = select_samples_to_pack( + samples, + packer=GreedyKnapsackPacker(12), + sequence_length_pad_multiple=4, + ) + + assert [[sample.sample_key for sample in pack] for pack in selected] == [ + ["s0", "s1"], + ["s2"], + ] + + +def test_preparation_builds_model_ready_pack_and_jagged_boundaries() -> None: + packed = pack_selected_samples( + [_sample("s0", 5), _sample("s1", 3)], + pack_capacity=12, + sequence_length_pad_multiple=4, + ) + + prepared = prepare_packed_sft_batch( + [packed], tokenizer=_Tokenizer(), only_unmask_final=False + ) + + assert prepared["input_ids"].tolist() == [[1, 2, 3, 4, 5, 0, 0, 0, 1, 2, 3, 0]] + assert prepared["token_mask"].tolist() == [[0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0]] + assert prepared["input_lengths"].tolist() == [12] + assert prepared["source_ids"] == [["s0", "s1"]] + assert "packed_schema_version" not in prepared + assert isinstance(prepared["cu_seqlens"], PackedTensor) + assert isinstance(prepared["cu_seqlens_padded"], PackedTensor) + assert prepared["cu_seqlens"].as_tensor().tolist() == [0, 5, 8] + assert prepared["cu_seqlens_padded"].as_tensor().tolist() == [0, 8, 12] + assert torch.equal( + _prepacked_boundary(prepared.slice(0, 1), "cu_seqlens", torch.device("cpu")), + torch.tensor([0, 5, 8], dtype=torch.int32), + ) + + +def test_physical_pack_rejects_incompatible_or_over_capacity_sources() -> None: + with pytest.raises(ValueError, match="compatible sources"): + pack_selected_samples( + [_sample("s0", 3), _sample("s1", 3, group="image")], + pack_capacity=8, + sequence_length_pad_multiple=1, + ) + with pytest.raises(ValueError, match="exceed the pack capacity"): + pack_selected_samples( + [_sample("s0", 5), _sample("s1", 4)], + pack_capacity=8, + sequence_length_pad_multiple=1, + ) From eb6221428faeea2caa81879f61886d5a815f04dd Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 9 Sep 2026 21:52:54 -0700 Subject: [PATCH 04/11] chore: add packing module to pyrefly Signed-off-by: rohitrango --- nemo_rl/algorithms/sft_v2.py | 2 +- nemo_rl/data/energon/multimodal/packing.py | 4 ++-- pyrefly.toml | 1 + 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/nemo_rl/algorithms/sft_v2.py b/nemo_rl/algorithms/sft_v2.py index 7b938c862e3..7330cafe6c7 100644 --- a/nemo_rl/algorithms/sft_v2.py +++ b/nemo_rl/algorithms/sft_v2.py @@ -405,7 +405,7 @@ def setup_sft_v2( raise ValueError( "Energon packing requires sequence_packing enabled with fuse_loss." ) - if sequence_packing["algorithm"] not in { + if sequence_packing.get("algorithm") not in { "greedy_knapsack", "balanced_greedy_knapsack", }: diff --git a/nemo_rl/data/energon/multimodal/packing.py b/nemo_rl/data/energon/multimodal/packing.py index 67ac02c8c12..3524373a624 100644 --- a/nemo_rl/data/energon/multimodal/packing.py +++ b/nemo_rl/data/energon/multimodal/packing.py @@ -88,8 +88,8 @@ def prepare_packed_sft_batch( raise ValueError("All physical packs in a batch need one capacity.") capacity = capacities.pop() packed_logs: list[list[dict[str, Any]]] = [] - boundaries: list[torch.Tensor] = [] - padded_boundaries: list[torch.Tensor] = [] + boundaries: list[torch.Tensor | None] = [] + padded_boundaries: list[torch.Tensor | None] = [] source_ids: list[list[str]] = [] for pack in packs: diff --git a/pyrefly.toml b/pyrefly.toml index 01788d1fb52..292bdce1a21 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -121,6 +121,7 @@ project-includes = [ "nemo_rl/data/energon/multimodal/cookers/__init__.py", "nemo_rl/data/energon/multimodal/cookers/generic.py", "nemo_rl/data/energon/multimodal/model_families.py", + "nemo_rl/data/energon/multimodal/packing.py", "nemo_rl/data/energon/multimodal/registry.py", "nemo_rl/data/energon/multimodal/task_encoders/__init__.py", "nemo_rl/data/energon/multimodal/task_encoders/base.py", From da8a26306c6cdf1bf6012a09afa9c7f895d3cfa3 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 9 Sep 2026 21:58:14 -0700 Subject: [PATCH 05/11] chore: add test copyright headers Signed-off-by: rohitrango --- tests/unit/data/packing/test_knapsack.py | 14 ++++++++++++++ tests/unit/data/test_energon_packing.py | 14 ++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/tests/unit/data/packing/test_knapsack.py b/tests/unit/data/packing/test_knapsack.py index fc9168a4e4c..413c8c1171e 100644 --- a/tests/unit/data/packing/test_knapsack.py +++ b/tests/unit/data/packing/test_knapsack.py @@ -1,3 +1,17 @@ +# 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. + import pytest from nemo_rl.data.packing import ( diff --git a/tests/unit/data/test_energon_packing.py b/tests/unit/data/test_energon_packing.py index 4050bc0b5de..6aeba6e78d0 100644 --- a/tests/unit/data/test_energon_packing.py +++ b/tests/unit/data/test_energon_packing.py @@ -1,3 +1,17 @@ +# 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. + from __future__ import annotations import pytest From ac2519de8f875ed168cf6f872271c2df72a5177f Mon Sep 17 00:00:00 2001 From: Rohit Jena Date: Thu, 10 Sep 2026 08:55:40 -0700 Subject: [PATCH 06/11] Update nemo_rl/data/packing/algorithms.py Signed-off-by: Rohit Jena --- nemo_rl/data/packing/algorithms.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_rl/data/packing/algorithms.py b/nemo_rl/data/packing/algorithms.py index 1a2c97edead..cd1f4b19ec6 100644 --- a/nemo_rl/data/packing/algorithms.py +++ b/nemo_rl/data/packing/algorithms.py @@ -331,7 +331,7 @@ def __init__( min_bin_count: Optional[int] = None, bin_count_multiple: Optional[int] = None, max_sequences_per_bin: Optional[int] = None, - balanced_knapsack_delta: int = 20, + balanced_knapsack_delta: int = 0, ) -> None: super().__init__( bin_capacity, From 51cdc68d25bb9f45b99155478d0378428f5b307d Mon Sep 17 00:00:00 2001 From: rohitrango Date: Thu, 10 Sep 2026 10:27:07 -0700 Subject: [PATCH 07/11] test(sft): add CP2 Energon packing recipe Signed-off-by: rohitrango --- ...-megatron-tp4ep8cp2-energon.v1.packing.yaml | 18 ++++++++++++++++++ ...8g-megatron-tp4ep8cp2-energon.v1.packing.sh | 4 ++++ 2 files changed, 22 insertions(+) create mode 100644 examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml create mode 100755 tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.sh diff --git a/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml b/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml new file mode 100644 index 00000000000..8b6e94633b3 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml @@ -0,0 +1,18 @@ +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: + project: sft-dev + name: sft-nemotron-omni-30b-clevr-energon-packing-tp4ep8cp2 diff --git a/tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.sh b/tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.sh new file mode 100755 index 00000000000..9bf8a7f783c --- /dev/null +++ b/tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.sh @@ -0,0 +1,4 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +export EXP_NAME=$(basename "$0" .sh) +source "$SCRIPT_DIR/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.sh" sft.max_num_steps=50 "$@" From 31d3111f927778f28643c0f65f9ea27a23d6e294 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Thu, 10 Sep 2026 10:31:58 -0700 Subject: [PATCH 08/11] fix(sft): address Energon packing review feedback Signed-off-by: rohitrango --- .../sequence-packing-and-dynamic-batching.md | 10 +- docs/guides/sft.md | 7 +- nemo_rl/algorithms/sft_v2.py | 47 ++++++- nemo_rl/data/energon/config.py | 22 +++- nemo_rl/data/energon/multimodal/packing.py | 46 +++++-- .../multimodal/task_encoders/generic_sft.py | 3 + nemo_rl/data/energon/sft_dataloader.py | 17 ++- nemo_rl/data/energon/sft_worker.py | 2 + nemo_rl/models/megatron/data.py | 59 ++++++++- tests/test_suites/disabled.txt | 2 + ...1n8g-megatron-tp8ep8-energon.v1.packing.sh | 4 + ...evr-1n2g-megatrontp1-energon.v1.packing.sh | 4 + tests/unit/algorithms/test_sft_v2.py | 57 +++++++- tests/unit/data/packing/test_algorithms.py | 123 +++--------------- tests/unit/data/test_energon_packing.py | 47 +++++-- tests/unit/data/test_energon_sft.py | 63 +++++++-- tests/unit/data/test_energon_sft_v2.py | 4 + .../models/megatron/test_megatron_data.py | 59 +++++++++ .../models/policy/test_tq_policy_placed.py | 39 +++++- 19 files changed, 454 insertions(+), 161 deletions(-) create mode 100755 tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.sh create mode 100755 tests/test_suites/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.packing.sh diff --git a/docs/design-docs/sequence-packing-and-dynamic-batching.md b/docs/design-docs/sequence-packing-and-dynamic-batching.md index bd3b229410c..3af17341a60 100644 --- a/docs/design-docs/sequence-packing-and-dynamic-batching.md +++ b/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 for NeMo-RL-owned packing; Energon-owned SFT packing uses +one of the two knapsack algorithms. ##### 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 diff --git a/docs/guides/sft.md b/docs/guides/sft.md index d6653abebab..d18235eb666 100644 --- a/docs/guides/sft.md +++ b/docs/guides/sft.md @@ -241,7 +241,12 @@ 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 `greedy_knapsack` or +`balanced_greedy_knapsack` 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. 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. diff --git a/nemo_rl/algorithms/sft_v2.py b/nemo_rl/algorithms/sft_v2.py index 7330cafe6c7..7e57222f6ce 100644 --- a/nemo_rl/algorithms/sft_v2.py +++ b/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 ( @@ -199,6 +200,11 @@ def _setup_loaders(self) -> None: "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" ], @@ -406,8 +412,8 @@ def setup_sft_v2( "Energon packing requires sequence_packing enabled with fuse_loss." ) if sequence_packing.get("algorithm") not in { - "greedy_knapsack", - "balanced_greedy_knapsack", + PackingAlgorithm.GREEDY_KNAPSACK.value, + PackingAlgorithm.BALANCED_GREEDY_KNAPSACK.value, }: raise ValueError("Energon SFT supports only the two knapsack packers.") if dynamic_batching["enabled"]: @@ -440,6 +446,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 diff --git a/nemo_rl/data/energon/config.py b/nemo_rl/data/energon/config.py index ee7f47d10e4..b7fb7d307f1 100644 --- a/nemo_rl/data/energon/config.py +++ b/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: Annotated[int, Field(ge=1)] | None = None - packing_buffer_size: Annotated[int, Field(ge=1)] | 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" diff --git a/nemo_rl/data/energon/multimodal/packing.py b/nemo_rl/data/energon/multimodal/packing.py index 3524373a624..81efb964eeb 100644 --- a/nemo_rl/data/energon/multimodal/packing.py +++ b/nemo_rl/data/energon/multimodal/packing.py @@ -1,4 +1,16 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# 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. """Energon-owned selection and materialization of multimodal SFT packs.""" @@ -99,6 +111,14 @@ def prepare_packed_sft_batch( add_loss_mask_to_message_log( logs, roles_to_train_on=["assistant"], only_unmask_final=only_unmask_final ) + templates = { + key: value + for log in logs + for message in log + for key, value in message.items() + if key not in {"token_ids", "token_loss_mask"} + and isinstance(value, torch.Tensor) + } lengths: list[int] = [] combined: list[dict[str, Any]] = [] token_dtype = torch.long @@ -111,13 +131,6 @@ def prepare_packed_sft_batch( token_dtype = tokens.dtype length = tokens.shape[0] lengths.append(length) - templates = { - key: value - for message in log - for key, value in message.items() - if key not in {"token_ids", "token_loss_mask"} - and isinstance(value, torch.Tensor) - } for message in log: message["token_loss_mask"] = ( message["token_loss_mask"] * sample.loss_multiplier @@ -152,15 +165,20 @@ def prepare_packed_sft_batch( combined.extend(log) tail = capacity - sum(pack.source_padded_lengths) if tail: - combined.append( + tail_message = { + "role": "padding", + "token_ids": torch.full( + (tail,), tokenizer.pad_token_id, dtype=token_dtype + ), + "token_loss_mask": torch.zeros(tail, dtype=torch.float32), + } + tail_message.update( { - "role": "padding", - "token_ids": torch.full( - (tail,), tokenizer.pad_token_id, dtype=token_dtype - ), - "token_loss_mask": torch.zeros(tail, dtype=torch.float32), + key: torch.zeros((tail, *value.shape[1:]), dtype=value.dtype) + for key, value in templates.items() } ) + combined.append(tail_message) packed_logs.append(combined) boundaries.append( torch.tensor( diff --git a/nemo_rl/data/energon/multimodal/task_encoders/generic_sft.py b/nemo_rl/data/energon/multimodal/task_encoders/generic_sft.py index dd97220be4e..14f8c5d17ec 100644 --- a/nemo_rl/data/energon/multimodal/task_encoders/generic_sft.py +++ b/nemo_rl/data/energon/multimodal/task_encoders/generic_sft.py @@ -233,6 +233,7 @@ def encode(self, sample: CanonicalSFTSample) -> EncodedSFTSample: for key, value in list(message.items()): if isinstance(value, PackedTensor): message[key] = PackedTensor.empty_like(value) + length = sum(len(message["token_ids"]) for message in message_log) loss_multiplier = 0.0 # group_key is the adapter fingerprint alone. Keying on the tensor names @@ -293,6 +294,7 @@ def batch_group_criterion( ) -> tuple[tuple[Any, ...], None]: return sample.group_key, None + @stateless def select_samples_to_pack( self, samples: list[EncodedSFTSample] ) -> list[list[EncodedSFTSample]]: @@ -304,6 +306,7 @@ def select_samples_to_pack( sequence_length_pad_multiple=self.sequence_length_pad_multiple, ) + @stateless def pack_selected_samples(self, samples: list[EncodedSFTSample]) -> PackedSFTSample: if self.packer is None: raise RuntimeError("Energon packing is not configured.") diff --git a/nemo_rl/data/energon/sft_dataloader.py b/nemo_rl/data/energon/sft_dataloader.py index 8d423b1f781..34cbf1f0564 100644 --- a/nemo_rl/data/energon/sft_dataloader.py +++ b/nemo_rl/data/energon/sft_dataloader.py @@ -288,7 +288,9 @@ def _loader_identity( shuffle: bool | None, topology: dict[str, Any], packing_algorithm: str | None, + max_sequences_per_bin: int | None, sequence_length_pad_multiple: int, + only_unmask_final: bool, ) -> dict[str, Any]: """Describe what a restored loader must still agree with.""" identity = { @@ -317,7 +319,9 @@ def _loader_identity( if packing_algorithm is not None: identity.update( packing_algorithm=packing_algorithm, + max_sequences_per_bin=max_sequences_per_bin, sequence_length_pad_multiple=sequence_length_pad_multiple, + only_unmask_final=only_unmask_final, ) return identity @@ -348,6 +352,7 @@ def _task_encoder( adapter: Any, include_source_ids: bool, packing_algorithm: str | None, + max_sequences_per_bin: int | None, max_sequence_length: int, sequence_length_pad_multiple: int, tokenizer: Any, @@ -371,7 +376,7 @@ def _task_encoder( get_packer( packing_algorithm, max_sequence_length, - max_sequences_per_bin=loader_config.max_samples_per_sequence, + max_sequences_per_bin=max_sequences_per_bin, ) if loader_config.packing_buffer_size is not None and packing_algorithm is not None @@ -403,9 +408,10 @@ def build_energon_sft_loader( logical_rank: int, logical_world_size: int, placement_fingerprint: str, - packing_algorithm: str | None = None, - sequence_length_pad_multiple: int = 1, - only_unmask_final: bool = False, + packing_algorithm: str | None, + max_sequences_per_bin: int | None, + sequence_length_pad_multiple: int, + only_unmask_final: bool, ) -> EnergonSFTDataLoader: """Build one loader for an explicit logical data shard and split.""" if "energon" not in data_config: @@ -434,6 +440,7 @@ def build_energon_sft_loader( adapter=adapter, include_source_ids=True, packing_algorithm=packing_algorithm, + max_sequences_per_bin=max_sequences_per_bin, max_sequence_length=max_sequence_length, sequence_length_pad_multiple=sequence_length_pad_multiple, tokenizer=processor.tokenizer, @@ -507,7 +514,9 @@ def build_energon_sft_loader( logical_world_size=logical_world_size, ), packing_algorithm=packing_algorithm, + max_sequences_per_bin=max_sequences_per_bin, sequence_length_pad_multiple=sequence_length_pad_multiple, + only_unmask_final=only_unmask_final, ), ) diff --git a/nemo_rl/data/energon/sft_worker.py b/nemo_rl/data/energon/sft_worker.py index 994bd3b24c8..c0703a4b5f2 100644 --- a/nemo_rl/data/energon/sft_worker.py +++ b/nemo_rl/data/energon/sft_worker.py @@ -66,6 +66,7 @@ def setup_sft_dataloader( max_sequence_length: int, placement_fingerprint: str, packing_algorithm: str | None, + max_sequences_per_bin: int | None, sequence_length_pad_multiple: int, only_unmask_final: bool, restored_state: Optional[dict[str, Any]] = None, @@ -91,6 +92,7 @@ def setup_sft_dataloader( logical_world_size=logical_world_size, placement_fingerprint=placement_fingerprint, packing_algorithm=packing_algorithm, + max_sequences_per_bin=max_sequences_per_bin, sequence_length_pad_multiple=sequence_length_pad_multiple, only_unmask_final=only_unmask_final, ) diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index 43005ed7d9b..29d285a635e 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -382,8 +382,29 @@ def _prepacked_boundary( return value.to(device=device, dtype=torch.int32) +def _slice_prepacked_for_cp(value: torch.Tensor, padded: torch.Tensor) -> torch.Tensor: + """Apply Megatron's per-source zigzag CP slicing to a packed row.""" + if value.ndim < 2 or value.shape[:2] != (1, int(padded[-1])): + raise ValueError( + "Prepacked token-aligned tensors must have shape [1, pack length, ...]." + ) + cp_rank = get_context_parallel_rank() + cp_size = get_context_parallel_world_size() + return torch.cat( + [ + _get_tokens_on_this_cp_rank( + value[:, int(start) : int(end)], cp_rank, cp_size, seq_dim=1 + ) + for start, end in zip(padded[:-1], padded[1:]) + ], + dim=1, + ).contiguous() + + def _prepare_prepacked( data: BatchedDataDict[Any], + *, + model_slices_context_parallel_inputs: bool, ) -> tuple[torch.Tensor, torch.Tensor, PackedSeqParams, torch.Tensor]: input_ids = data["input_ids"] if not torch.is_tensor(input_ids) or input_ids.shape[0] != 1: @@ -402,8 +423,19 @@ def _prepare_prepacked( or bool((source_lengths > padded_lengths).any()) ): raise ValueError("Invalid prepacked source boundaries.") - if get_context_parallel_world_size() != 1: - raise NotImplementedError("Energon-owned packing currently requires CP=1.") + cp_size = get_context_parallel_world_size() + if cp_size > 1 and bool((padded_lengths % (2 * cp_size) != 0).any()): + raise ValueError( + "Every prepacked padded source length must be divisible by 2 * " + f"context_parallel_size ({2 * cp_size})." + ) + local_input_ids = _slice_prepacked_for_cp(input_ids, padded) + input_ids_cp_sharded = ( + input_ids if model_slices_context_parallel_inputs else local_input_ids + ) + # Keep physical boundaries in cu_seqlens_q as well as cu_seqlens_q_padded. + # MTP loss rolling still has consumers that use cu_seqlens_q as the wrap + # boundary, so logical boundaries can roll into padding or the next source. params = PackedSeqParams( cu_seqlens_q=padded, cu_seqlens_kv=padded, @@ -413,9 +445,9 @@ def _prepare_prepacked( max_seqlen_kv=int(padded_lengths.max()), pad_between_seqs=False, qkv_format="thd", - total_tokens=input_ids.shape[1], + total_tokens=input_ids_cp_sharded.shape[1], ) - return input_ids, input_ids, params, padded + return input_ids, input_ids_cp_sharded, params, padded def process_microbatch( @@ -492,7 +524,24 @@ def process_microbatch( input_ids_cp_sharded, packed_seq_params, cu_seqlens_padded, - ) = _prepare_prepacked(data_dict) + ) = _prepare_prepacked( + data_dict, + model_slices_context_parallel_inputs=( + model_slices_context_parallel_inputs + ), + ) + if "mtp_loss_mask" in data_dict: + mtp_loss_mask = data_dict["mtp_loss_mask"] + if not model_slices_context_parallel_inputs: + mtp_loss_mask = _slice_prepacked_for_cp( + mtp_loss_mask, cu_seqlens_padded + ) + if "media_token_validity_mask" in data_dict: + media_token_validity_mask = data_dict["media_token_validity_mask"] + if not model_slices_context_parallel_inputs: + media_token_validity_mask = _slice_prepacked_for_cp( + media_token_validity_mask, cu_seqlens_padded + ) position_ids = None attention_mask = None elif delegate_pack_to_model: diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index 0eb54edaf77..4c263c23f0d 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -60,12 +60,14 @@ tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-ready-first-single-controll # no tier it can join (test_nightly_suites_match_gpus_per_node enforces this). # Covered by tests/functional/sft_v2_energon.sh in the L1 SFT suite. tests/test_suites/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.sh +tests/test_suites/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.packing.sh # The Nemotron-Omni 30B-A3B Energon SFTv2 run is self-contained, but its 1x8-GPU # 90-minute allocation adds 12 GPU-hours and the nightly suite is already at # ~4176 of its 4181 GPU-hour cap (test_nightly_compute_stays_below_4181_hours). # Run it manually until the nightly suite has room. tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.sh +tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.sh # OOMs on H100 2n8g during the colocated reshard: the colocated training state # offload is currently too slow to use; tracked by issue #3976. tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.sh diff --git a/tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.sh b/tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.sh new file mode 100755 index 00000000000..fea0a0f9791 --- /dev/null +++ b/tests/test_suites/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.sh @@ -0,0 +1,4 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +export EXP_NAME=$(basename "$0" .sh) +source "$SCRIPT_DIR/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.sh" "$@" diff --git a/tests/test_suites/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.packing.sh b/tests/test_suites/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.packing.sh new file mode 100755 index 00000000000..680b2e391ed --- /dev/null +++ b/tests/test_suites/vlm/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.packing.sh @@ -0,0 +1,4 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +export EXP_NAME=$(basename "$0" .sh) +source "$SCRIPT_DIR/vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.sh" "$@" diff --git a/tests/unit/algorithms/test_sft_v2.py b/tests/unit/algorithms/test_sft_v2.py index fcb72d9a0ce..9cce438a769 100644 --- a/tests/unit/algorithms/test_sft_v2.py +++ b/tests/unit/algorithms/test_sft_v2.py @@ -86,12 +86,19 @@ def _valid_setup_config( "backend": "energon", "validation": None, "max_input_seq_length": 128, + "energon": SimpleNamespace(packing_buffer_size=None), } data.update(data_overrides or {}) policy = { - "megatron_cfg": {"enabled": True}, + "megatron_cfg": { + "enabled": True, + "context_parallel_size": 1, + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + }, "sequence_packing": {"enabled": False}, "dynamic_batching": {"enabled": False}, + "make_sequence_length_divisible_by": 1, } for section, values in (policy_overrides or {}).items(): policy[section].update(values) @@ -220,11 +227,11 @@ def test_checkpoint_metric_rejects_a_metric_no_step_produces() -> None: ), ( {"policy_overrides": {"sequence_packing": {"enabled": True}}}, - "fixed NeMo-RL batching", + "fixed batching", ), ( {"policy_overrides": {"dynamic_batching": {"enabled": True}}}, - "fixed NeMo-RL batching", + "fixed batching", ), ({"sft_overrides": {"val_period": 10}}, "has no validation loop"), ( @@ -258,6 +265,50 @@ def test_setup_rejects_a_validation_checkpoint_metric() -> None: ) +@pytest.mark.parametrize( + ("megatron_overrides", "policy_multiple", "message"), + [ + ({"context_parallel_size": 2}, 2, "multiple of 4"), + ( + { + "moe_token_dispatcher_type": "flex", + "moe_flex_dispatcher_backend": "hybridep", + }, + 1, + "HybridEP", + ), + ( + {"fp8_cfg": {"enabled": True, "fp8_recipe": "blockwise"}}, + 1, + "FP8 packed-token alignment", + ), + ], +) +def test_setup_rejects_unsupported_energon_packing_layouts( + megatron_overrides: dict[str, Any], policy_multiple: int, message: str +) -> None: + from nemo_rl.algorithms.sft_v2 import setup_sft_v2 + + config = _valid_setup_config( + data_overrides={ + "max_input_seq_length": 130, + "energon": SimpleNamespace(packing_buffer_size=64), + }, + policy_overrides={ + "megatron_cfg": megatron_overrides, + "sequence_packing": { + "enabled": True, + "fuse_loss": True, + "algorithm": "greedy_knapsack", + }, + }, + ) + config.policy["make_sequence_length_divisible_by"] = policy_multiple + + with pytest.raises(ValueError, match=message): + setup_sft_v2(config, MagicMock()) + + def test_restore_rejects_changed_placement() -> None: from nemo_rl.algorithms.sft_v2 import _restore_save_state diff --git a/tests/unit/data/packing/test_algorithms.py b/tests/unit/data/packing/test_algorithms.py index f46fdacad6f..978c96636a2 100644 --- a/tests/unit/data/packing/test_algorithms.py +++ b/tests/unit/data/packing/test_algorithms.py @@ -25,6 +25,13 @@ get_packer, ) +ALL_ALGORITHMS = list(PackingAlgorithm) +DETERMINISTIC_ALGORITHMS = [ + algorithm + for algorithm in ALL_ALGORITHMS + if algorithm is not PackingAlgorithm.FIRST_FIT_SHUFFLE +] + def validate_solution( sequence_lengths: List[int], bins: List[List[int]], bin_capacity: int @@ -92,16 +99,10 @@ def edge_cases(self) -> Dict[str, List[int]]: "mixed_sizes": [10, 50, 100, 20, 80, 30, 70, 40, 60, 90], } - # TODO(ahmadki): use the function to specify all test algorithms ins tead of lists below @pytest.fixture def algorithms(self) -> List[PackingAlgorithm]: """Fixture for packing algorithms.""" - return [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ] + return ALL_ALGORITHMS def test_get_packer(self, bin_capacity: int, algorithms: List[PackingAlgorithm]): """Test the get_packer factory function.""" @@ -116,15 +117,7 @@ def test_get_packer(self, bin_capacity: int, algorithms: List[PackingAlgorithm]) invalid_algorithm = object() get_packer(invalid_algorithm, bin_capacity) # type: ignore - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", ALL_ALGORITHMS) def test_small_sequences( self, bin_capacity: int, @@ -141,15 +134,7 @@ def test_small_sequences( # Print the number of bins used (for information) print(f"{algorithm.name} used {len(bins)} bins for small sequences") - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", ALL_ALGORITHMS) def test_medium_sequences( self, bin_capacity: int, @@ -166,15 +151,7 @@ def test_medium_sequences( # Print the number of bins used (for information) print(f"{algorithm.name} used {len(bins)} bins for medium sequences") - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", ALL_ALGORITHMS) def test_large_sequences( self, bin_capacity: int, @@ -191,16 +168,7 @@ def test_large_sequences( # Print the number of bins used (for information) print(f"{algorithm.name} used {len(bins)} bins for large sequences") - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) - # TODO(ahmadki): use the function to specify all test algorithms instead of lists below + @pytest.mark.parametrize("algorithm", ALL_ALGORITHMS) @pytest.mark.parametrize( "case_name, sequence_lengths", [ @@ -228,15 +196,7 @@ def test_edge_cases( if case_name == "single_item": assert len(bins) == 1 - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", ALL_ALGORITHMS) def test_empty_list(self, bin_capacity: int, algorithm: PackingAlgorithm): """Test empty list with algorithms that can handle it.""" packer = get_packer(algorithm, bin_capacity) @@ -245,15 +205,7 @@ def test_empty_list(self, bin_capacity: int, algorithm: PackingAlgorithm): # For empty list, check that no bins are created assert len(bins) == 0 - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", ALL_ALGORITHMS) def test_error_cases(self, bin_capacity: int, algorithm: PackingAlgorithm): """Test error cases with all algorithms.""" # Test with a sequence length that exceeds bin capacity @@ -263,14 +215,7 @@ def test_error_cases(self, bin_capacity: int, algorithm: PackingAlgorithm): with pytest.raises(ValueError): packer.pack(sequence_lengths) - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", DETERMINISTIC_ALGORITHMS) def test_deterministic( self, bin_capacity: int, @@ -325,15 +270,7 @@ def test_randomized( f"Warning: {algorithm.name} produced the same result with different seeds" ) - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", ALL_ALGORITHMS) def test_min_bin_count( self, bin_capacity: int, @@ -366,15 +303,7 @@ def test_min_bin_count( for bin_contents in bins_more: assert len(bin_contents) > 0, "Found empty bin" - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.FIRST_FIT_SHUFFLE, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", ALL_ALGORITHMS) def test_bin_count_multiple( self, bin_capacity: int, @@ -418,14 +347,7 @@ def test_bin_count_multiple( for bin_contents in bins_force: assert len(bin_contents) > 0, "Found empty bin" - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", DETERMINISTIC_ALGORITHMS) def test_combined_constraints( self, bin_capacity: int, @@ -503,14 +425,7 @@ def test_insufficient_sequences_for_constraints(self, bin_capacity: int): ): packer.pack(sequence_lengths) - @pytest.mark.parametrize( - "algorithm", - [ - PackingAlgorithm.CONCATENATIVE, - PackingAlgorithm.FIRST_FIT_DECREASING, - PackingAlgorithm.MODIFIED_FIRST_FIT_DECREASING, - ], - ) + @pytest.mark.parametrize("algorithm", DETERMINISTIC_ALGORITHMS) def test_packing_preservation( self, bin_capacity: int, diff --git a/tests/unit/data/test_energon_packing.py b/tests/unit/data/test_energon_packing.py index 6aeba6e78d0..059d0268169 100644 --- a/tests/unit/data/test_energon_packing.py +++ b/tests/unit/data/test_energon_packing.py @@ -30,7 +30,6 @@ from nemo_rl.data.energon.multimodal.types import EncodedSFTSample # noqa: E402 from nemo_rl.data.multimodal_utils import PackedTensor # noqa: E402 from nemo_rl.data.packing import GreedyKnapsackPacker # noqa: E402 -from nemo_rl.models.megatron.data import _prepacked_boundary # noqa: E402 class _Tokenizer: @@ -89,24 +88,48 @@ def test_preparation_builds_model_ready_pack_and_jagged_boundaries() -> None: sequence_length_pad_multiple=4, ) + second_pack = pack_selected_samples( + [_sample("s2", 4)], + pack_capacity=12, + sequence_length_pad_multiple=4, + ) prepared = prepare_packed_sft_batch( - [packed], tokenizer=_Tokenizer(), only_unmask_final=False + [packed, second_pack], tokenizer=_Tokenizer(), only_unmask_final=False ) - assert prepared["input_ids"].tolist() == [[1, 2, 3, 4, 5, 0, 0, 0, 1, 2, 3, 0]] - assert prepared["token_mask"].tolist() == [[0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0]] - assert prepared["input_lengths"].tolist() == [12] - assert prepared["source_ids"] == [["s0", "s1"]] - assert "packed_schema_version" not in prepared + assert prepared["input_ids"][0].tolist() == [1, 2, 3, 4, 5, 0, 0, 0, 1, 2, 3, 0] + assert prepared["token_mask"][0].tolist() == [0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0] + assert prepared["input_lengths"].tolist() == [12, 12] + assert prepared["source_ids"] == [["s0", "s1"], ["s2"]] assert isinstance(prepared["cu_seqlens"], PackedTensor) assert isinstance(prepared["cu_seqlens_padded"], PackedTensor) - assert prepared["cu_seqlens"].as_tensor().tolist() == [0, 5, 8] - assert prepared["cu_seqlens_padded"].as_tensor().tolist() == [0, 8, 12] - assert torch.equal( - _prepacked_boundary(prepared.slice(0, 1), "cu_seqlens", torch.device("cpu")), - torch.tensor([0, 5, 8], dtype=torch.int32), + first = prepared.slice(0, 1) + assert first["cu_seqlens"].as_tensor().tolist() == [0, 5, 8] + assert first["cu_seqlens_padded"].as_tensor().tolist() == [0, 8, 12] + sliced = prepared.slice(1, 2) + assert sliced["cu_seqlens"].as_tensor().tolist() == [0, 4] + assert sliced["cu_seqlens_padded"].as_tensor().tolist() == [0, 12] + + +def test_preparation_backfills_multimodal_token_fields() -> None: + text_sample = _sample("text", 4) + multimodal_sample = _sample("image", 4) + for message in multimodal_sample.message_log: + message["mm_token_type_ids"] = torch.ones_like(message["token_ids"]) + packed = pack_selected_samples( + [text_sample, multimodal_sample], + pack_capacity=12, + sequence_length_pad_multiple=1, ) + prepared = prepare_packed_sft_batch( + [packed], tokenizer=_Tokenizer(), only_unmask_final=False + ) + + assert prepared["mm_token_type_ids"].tolist() == [ + [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0] + ] + def test_physical_pack_rejects_incompatible_or_over_capacity_sources() -> None: with pytest.raises(ValueError, match="compatible sources"): diff --git a/tests/unit/data/test_energon_sft.py b/tests/unit/data/test_energon_sft.py index cdd16bb7a09..decaab34c6d 100644 --- a/tests/unit/data/test_energon_sft.py +++ b/tests/unit/data/test_energon_sft.py @@ -271,6 +271,16 @@ def test_qwen_adapter_returns_tokenized_message_log_with_model_inputs(): ) +def test_adapter_uses_truncated_length_for_packing() -> None: + encoded = _adapter(_FakeQwenProcessor(), max_sequence_length=16).encode(_sample()) + + assert encoded.length == sum( + len(message["token_ids"]) for message in encoded.message_log + ) + assert encoded.packing_cost == encoded.length + assert encoded.packing_cost < 16 + + def test_hf_and_energon_backends_agree_on_the_same_conversation(): """Both backends feed prepare_sft_batch; the prepared tensors must match. @@ -385,11 +395,7 @@ def test_task_encoder_runs_split_encode_and_batch_lifecycle_methods(): assert encoder.encode_batch(batch) is batch assert batch["source_ids"] == ["sample-0"] - # Stage 1 does not override select_samples_to_pack, so this falls through to - # Energon's base implementation. - with pytest.raises( - NotImplementedError, match="Packing only effective when overridden" - ): + with pytest.raises(RuntimeError, match="packing is not configured"): encoder.select_samples_to_pack([preencoded]) @@ -455,7 +461,7 @@ def test_rejected_restore_names_the_settings_that_changed(): ).load_state_dict(state) -def test_energon_config_disables_sequence_packing(): +def test_energon_config_validates_sequence_packing(): config = EnergonLoaderConfig(model_family="qwen") assert config.model_family == "qwen" assert config.packing_buffer_size is None @@ -469,10 +475,15 @@ def test_energon_config_disables_sequence_packing(): path="/data/prepared", split="train", virtual_epoch_length=10 ) assert source.virtual_epoch_length == 10 - with pytest.raises(ValueError): - EnergonLoaderConfig(model_family="qwen", packing_buffer_size=10) - with pytest.raises(ValueError): - EnergonLoaderConfig(model_family="qwen", max_samples_per_sequence=2) + packed = EnergonLoaderConfig( + model_family="qwen", packing_buffer_size=10, max_samples_per_sequence=2 + ) + assert packed.packing_buffer_size == 10 + assert packed.max_samples_per_sequence == 2 + for field in ("packing_buffer_size", "max_samples_per_sequence"): + for value in (0, -1): + with pytest.raises(ValueError): + EnergonLoaderConfig(model_family="qwen", **{field: value}) with pytest.raises(ValueError): EnergonLoaderConfig.model_validate({}) with pytest.raises(ValueError): @@ -485,6 +496,9 @@ def _identity( batch_size: int = 8, shuffle: bool | None = True, logical_rank: int = 0, + packing_algorithm: str | None = None, + max_sequences_per_bin: int | None = None, + only_unmask_final: bool = False, ) -> dict: config = loader_config or EnergonLoaderConfig(model_family="qwen") return _loader_identity( @@ -502,6 +516,10 @@ def _identity( "logical_rank": logical_rank, "logical_world_size": 2, }, + packing_algorithm=packing_algorithm, + max_sequences_per_bin=max_sequences_per_bin, + sequence_length_pad_multiple=1, + only_unmask_final=only_unmask_final, ) @@ -542,6 +560,27 @@ def test_identity_refuses_a_changed_batch_size_or_shuffle(): assert _identity_fingerprint(changed) != _identity_fingerprint(baseline) +def test_identity_pins_packing_semantics(): + baseline = _identity( + packing_algorithm="greedy_knapsack", + max_sequences_per_bin=4, + ) + + for changed in ( + _identity( + packing_algorithm="balanced_greedy_knapsack", + max_sequences_per_bin=4, + ), + _identity(packing_algorithm="greedy_knapsack", max_sequences_per_bin=2), + _identity( + packing_algorithm="greedy_knapsack", + max_sequences_per_bin=4, + only_unmask_final=True, + ), + ): + assert _identity_fingerprint(changed) != _identity_fingerprint(baseline) + + def test_train_loader_rejects_shuffle_false(): # get_train_dataset shards by slice and Energon asserts a single slice # iterator when it does not shuffle over epochs, so shuffle=false is not a @@ -560,6 +599,10 @@ def test_train_loader_rejects_shuffle_false(): logical_rank=0, logical_world_size=1, placement_fingerprint="same-placement", + packing_algorithm=None, + max_sequences_per_bin=None, + sequence_length_pad_multiple=1, + only_unmask_final=False, ) diff --git a/tests/unit/data/test_energon_sft_v2.py b/tests/unit/data/test_energon_sft_v2.py index 70d04ff2c37..2821a2233de 100644 --- a/tests/unit/data/test_energon_sft_v2.py +++ b/tests/unit/data/test_energon_sft_v2.py @@ -71,6 +71,10 @@ def _v2_fingerprint( logical_rank=logical_rank, logical_world_size=logical_world_size, ), + packing_algorithm=None, + max_sequences_per_bin=None, + sequence_length_pad_multiple=1, + only_unmask_final=False, ) ) diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py index 6d99a3d2254..6a16b4b87f6 100644 --- a/tests/unit/models/megatron/test_megatron_data.py +++ b/tests/unit/models/megatron/test_megatron_data.py @@ -199,6 +199,19 @@ def test_get_and_validate_seqlen_still_checks_per_token_multimodal(self): class TestProcessMicrobatch: """Tests for process_microbatch function.""" + @staticmethod + def _prepacked_batch() -> BatchedDataDict: + return BatchedDataDict( + { + "input_ids": torch.tensor([[1, 2, 3, 0, 5, 6, 7, 0]]), + "input_lengths": torch.tensor([8]), + "token_mask": torch.tensor([[1, 1, 1, 0, 1, 1, 1, 0]]), + "sample_mask": torch.tensor([1.0]), + "cu_seqlens": [torch.tensor([0, 3, 6], dtype=torch.int32)], + "cu_seqlens_padded": [torch.tensor([0, 4, 8], dtype=torch.int32)], + } + ) + @patch("nemo_rl.models.megatron.data.get_ltor_masks_and_position_ids") def test_process_microbatch_no_packing(self, mock_get_masks): """Test process_microbatch without sequence packing.""" @@ -389,6 +402,52 @@ def test_process_microbatch_with_packing( # Verify pack was called mock_pack.assert_called_once() + @patch("nemo_rl.models.megatron.data.get_context_parallel_rank", return_value=0) + @patch( + "nemo_rl.models.megatron.data.get_context_parallel_world_size", return_value=1 + ) + @patch("nemo_rl.models.megatron.data._pack_sequences_for_megatron") + def test_process_microbatch_uses_prepacked_physical_boundaries( + self, mock_pack, mock_cp_world, mock_cp_rank + ): + from nemo_rl.models.megatron.data import process_microbatch + + data = self._prepacked_batch() + result = process_microbatch( + data, + seq_length_key="input_lengths", + pack_sequences=True, + ) + + mock_pack.assert_not_called() + assert torch.equal(result.input_ids_cp_sharded, data["input_ids"]) + assert torch.equal( + result.packed_seq_params.cu_seqlens_q, + torch.tensor([0, 4, 8], dtype=torch.int32), + ) + assert result.packed_seq_params.pad_between_seqs is False + + @patch("nemo_rl.models.megatron.data.get_context_parallel_rank", return_value=0) + @patch( + "nemo_rl.models.megatron.data.get_context_parallel_world_size", return_value=2 + ) + def test_process_microbatch_cp_slices_each_prepacked_source( + self, mock_cp_world, mock_cp_rank + ): + from nemo_rl.models.megatron.data import process_microbatch + + data = self._prepacked_batch() + data["mtp_loss_mask"] = data["token_mask"].clone() + result = process_microbatch( + data, + seq_length_key="input_lengths", + pack_sequences=True, + ) + + assert torch.equal(result.input_ids_cp_sharded, torch.tensor([[1, 0, 5, 0]])) + assert torch.equal(result.mtp_loss_mask, torch.tensor([[1, 0, 1, 0]])) + assert result.packed_seq_params.total_tokens == 4 + @patch("nemo_rl.models.megatron.data.get_ltor_masks_and_position_ids") def test_process_microbatch_no_packing_propagates_mtp_loss_mask( self, mock_get_masks diff --git a/tests/unit/models/policy/test_tq_policy_placed.py b/tests/unit/models/policy/test_tq_policy_placed.py index 1f9cc7ef8b0..20076ef17f1 100644 --- a/tests/unit/models/policy/test_tq_policy_placed.py +++ b/tests/unit/models/policy/test_tq_policy_placed.py @@ -19,7 +19,11 @@ import pytest from nemo_rl.data_plane import KVBatchMeta -from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN +from nemo_rl.data_plane.schema import ( + GLOBAL_FORWARD_PAD_SEQLEN, + MICRO_BATCH_INDICES, + MICRO_BATCH_LENGTHS, +) from nemo_rl.models.policy.tq_policy import TQPolicy @@ -91,15 +95,44 @@ def test_train_placed_microbatches_requires_one_batch_per_dp_rank() -> None: worker_group.run_all_workers_sharded_data.assert_not_called() -def test_train_placed_microbatches_rejects_sequence_packing() -> None: +def test_train_placed_microbatches_rejects_dynamic_batching() -> None: + policy, worker_group = _policy() + policy.use_dynamic_batches = True + policy.dynamic_batching_args = {} + policy.cfg["dynamic_batching"] = {"train_mb_tokens": 4096} + + with pytest.raises(ValueError, match="dynamic batching"): + policy.train_placed_microbatches( + [_meta(0, ["input_ids"]), _meta(1, ["input_ids"])] + ) + + worker_group.run_all_workers_sharded_data.assert_not_called() + + +def test_train_placed_microbatches_requires_producer_packing_shapes() -> None: policy, worker_group = _policy() policy.use_sequence_packing = True policy.sequence_packing_args = {"algorithm": "modified_first_fit_decreasing"} policy.cfg["sequence_packing"] = {"train_mb_tokens": 4096} - with pytest.raises(ValueError, match="fixed batches only"): + with pytest.raises(ValueError, match="producer microbatch shapes"): policy.train_placed_microbatches( [_meta(0, ["input_ids"]), _meta(1, ["input_ids"])] ) worker_group.run_all_workers_sharded_data.assert_not_called() + + +def test_train_placed_microbatches_accepts_producer_packing_shapes() -> None: + policy, worker_group = _policy() + policy.use_sequence_packing = True + policy.sequence_packing_args = {"algorithm": "modified_first_fit_decreasing"} + policy.cfg["sequence_packing"] = {"train_mb_tokens": 4096} + dp_metas = [_meta(0, ["input_ids"]), _meta(1, ["input_ids"])] + for meta in dp_metas: + meta.extra_info[MICRO_BATCH_INDICES] = [[[0, 1], [1, 2]]] + meta.extra_info[MICRO_BATCH_LENGTHS] = [[8, 16]] + + policy.train_placed_microbatches(dp_metas) + + worker_group.run_all_workers_sharded_data.assert_called_once() From 1fa103345129b8a33b7a481001a7f627c611cba7 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Thu, 10 Sep 2026 10:41:13 -0700 Subject: [PATCH 09/11] docs(sft): clarify packed sample metrics Signed-off-by: rohitrango --- docs/guides/sft.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/guides/sft.md b/docs/guides/sft.md index d18235eb666..1c1802c62fa 100644 --- a/docs/guides/sft.md +++ b/docs/guides/sft.md @@ -248,6 +248,11 @@ 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. ### OpenAI Format Datasets (with Tool Calling Support) From 295add73cb9d1ff9c49b17a330872c24c2fa08c5 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Thu, 10 Sep 2026 10:54:58 -0700 Subject: [PATCH 10/11] feat(sft): allow all Energon packing algorithms Signed-off-by: rohitrango --- .../sequence-packing-and-dynamic-batching.md | 4 ++-- docs/guides/sft.md | 4 ++-- nemo_rl/algorithms/sft_v2.py | 5 ++--- tests/unit/data/test_energon_packing.py | 15 +++++++++------ 4 files changed, 15 insertions(+), 13 deletions(-) diff --git a/docs/design-docs/sequence-packing-and-dynamic-batching.md b/docs/design-docs/sequence-packing-and-dynamic-batching.md index 3af17341a60..94a1cf1d01e 100644 --- a/docs/design-docs/sequence-packing-and-dynamic-batching.md +++ b/docs/design-docs/sequence-packing-and-dynamic-batching.md @@ -89,8 +89,8 @@ We have the policy backends perform the actual packing because implementations c #### 2. Packing Algorithms (`nemo_rl/data/packing/algorithms.py`) Six packing algorithms are implemented. Modified First Fit Decreasing is the -default recommendation for NeMo-RL-owned packing; Energon-owned SFT packing uses -one of the two knapsack algorithms. +default recommendation, and Energon-owned SFT packing supports all six through +the same interface. ##### Concatenative Packer - Sequential concatenation until bin capacity is reached diff --git a/docs/guides/sft.md b/docs/guides/sft.md index 1c1802c62fa..a87a3f8df13 100644 --- a/docs/guides/sft.md +++ b/docs/guides/sft.md @@ -242,8 +242,8 @@ 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. Set `data.energon.packing_buffer_size` and enable fused -`policy.sequence_packing` with `greedy_knapsack` or -`balanced_greedy_knapsack` to let Energon form model-ready multimodal packs. +`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. diff --git a/nemo_rl/algorithms/sft_v2.py b/nemo_rl/algorithms/sft_v2.py index 7e57222f6ce..d9280f70e11 100644 --- a/nemo_rl/algorithms/sft_v2.py +++ b/nemo_rl/algorithms/sft_v2.py @@ -412,10 +412,9 @@ def setup_sft_v2( "Energon packing requires sequence_packing enabled with fuse_loss." ) if sequence_packing.get("algorithm") not in { - PackingAlgorithm.GREEDY_KNAPSACK.value, - PackingAlgorithm.BALANCED_GREEDY_KNAPSACK.value, + algorithm.value for algorithm in PackingAlgorithm }: - raise ValueError("Energon SFT supports only the two knapsack packers.") + 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, diff --git a/tests/unit/data/test_energon_packing.py b/tests/unit/data/test_energon_packing.py index 059d0268169..664eab3140b 100644 --- a/tests/unit/data/test_energon_packing.py +++ b/tests/unit/data/test_energon_packing.py @@ -29,7 +29,7 @@ ) from nemo_rl.data.energon.multimodal.types import EncodedSFTSample # noqa: E402 from nemo_rl.data.multimodal_utils import PackedTensor # noqa: E402 -from nemo_rl.data.packing import GreedyKnapsackPacker # noqa: E402 +from nemo_rl.data.packing import PackingAlgorithm, get_packer # noqa: E402 class _Tokenizer: @@ -62,7 +62,10 @@ def _sample( ) -def test_selection_uses_aligned_costs_and_keeps_groups_separate() -> None: +@pytest.mark.parametrize("algorithm", list(PackingAlgorithm)) +def test_selection_uses_aligned_costs_and_keeps_groups_separate( + algorithm: PackingAlgorithm, +) -> None: samples = [ _sample("s0", 5), _sample("s1", 3), @@ -71,13 +74,13 @@ def test_selection_uses_aligned_costs_and_keeps_groups_separate() -> None: selected = select_samples_to_pack( samples, - packer=GreedyKnapsackPacker(12), + packer=get_packer(algorithm, 12), sequence_length_pad_multiple=4, ) - assert [[sample.sample_key for sample in pack] for pack in selected] == [ - ["s0", "s1"], - ["s2"], + assert [{sample.sample_key for sample in pack} for pack in selected] == [ + {"s0", "s1"}, + {"s2"}, ] From 619c3d6a2003bdc69a12e21ff2ff11bdbaa2aad6 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Thu, 10 Sep 2026 11:10:40 -0700 Subject: [PATCH 11/11] fix(sft): satisfy lint checks Signed-off-by: rohitrango --- ...ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml | 4 ---- nemo_rl/models/policy/tq_policy.py | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml b/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml index 8b6e94633b3..7ce9f7bf6df 100644 --- a/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml +++ b/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml @@ -1,18 +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: - project: sft-dev name: sft-nemotron-omni-30b-clevr-energon-packing-tp4ep8cp2 diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index 618c66a6eb4..2b34e0f21b4 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -47,9 +47,9 @@ from nemo_rl.data_plane.schema import ( DP_TRAIN_FIELDS, GLOBAL_FORWARD_PAD_SEQLEN, + LP_SEED_FIELDS, MICRO_BATCH_INDICES, MICRO_BATCH_LENGTHS, - LP_SEED_FIELDS, ROUTE_PASSTHROUGH_FLAG, ROUTE_PLAN_TAG, fields_with_optional_opd_full,