diff --git a/docs/design-docs/sequence-packing-and-dynamic-batching.md b/docs/design-docs/sequence-packing-and-dynamic-batching.md index bd3b229410c..94a1cf1d01e 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, and Energon-owned SFT packing supports all six through +the same interface. ##### Concatenative Packer - Sequential concatenation until bin capacity is reached @@ -107,6 +109,12 @@ Four packing algorithms are implemented, but we recommend you just use Modified 5. Greedy fit remaining items 6. Apply FFD to leftovers +##### Greedy Knapsack +- Repeatedly selects the largest remaining sequence that fits in the current bin + +##### Balanced Greedy Knapsack +- Places descending sequences into the least-full available bin + ##### First Fit Decreasing (FFD) - Sort sequences by length (descending), place each in first fitting bin - O(n log n + n*m) where m = number of bins diff --git a/docs/guides/sft.md b/docs/guides/sft.md index d6653abebab..a87a3f8df13 100644 --- a/docs/guides/sft.md +++ b/docs/guides/sft.md @@ -241,7 +241,17 @@ The processor runs inside Energon loader workers and returns the same tokenized The v1 `SFTProcessorAdapter` and `HFMultimodalSFTProcessorAdapter` are narrow integration interfaces. They are planned to be replaced by a more comprehensive modular processor implementation; dataset loading and the policy-facing batch shape should remain stable through that change. -Sequence packing is unavailable in this path, on both sides: `packing_buffer_size` and `max_samples_per_sequence` are typed null-only, and `policy.sequence_packing` (like `policy.dynamic_batching`) is rejected at startup with `SFTv2 requires fixed NeMo-RL batching.` Packing is deferred to a later stage of the Energon integration. Energon does not provide a separate offline sequence-packing pipeline either; offline preparation may store length and media-cost metadata, but should not pre-concatenate multimodal conversations. +Set `data.energon.packing_buffer_size` and enable fused +`policy.sequence_packing` with any supported packing algorithm to let Energon +form model-ready multimodal packs. +Without an Energon packing buffer, SFTv2 currently requires fixed batching. +Dynamic batching and HybridEP flex dispatch are not supported with +Energon-owned packs. + +With Energon-owned packing, each `sample_mask` entry represents one physical +pack, so `num_valid_samples` counts non-empty packs rather than source +conversations. NLL loss scaling is unchanged because it is normalized by +`global_valid_toks`. Training dataloader checkpoints include the Energon worker state plus a fingerprint of the source, loader, and processor settings. Restore must occur before the first iteration, and a changed fingerprint fails instead of silently continuing with a different stream. SFTv2 accepts a single train source; use an Energon metadataset to blend prepared sources. 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..7ce9f7bf6df --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp4ep8cp2-energon.v1.packing.yaml @@ -0,0 +1,14 @@ +defaults: ./vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.yaml +sft: + max_num_steps: 50 +policy: + make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size}, + ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}} + megatron_cfg: + tensor_model_parallel_size: 4 + context_parallel_size: 2 +logger: + wandb_enabled: true + log_dir: logs/sft-nemotron-omni-30b-clevr-energon-packing-tp4ep8cp2 + wandb: + name: sft-nemotron-omni-30b-clevr-energon-packing-tp4ep8cp2 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..d9280f70e11 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 ( @@ -196,6 +197,18 @@ def _setup_loaders(self) -> None: // self._placement_plan.logical_world_size, "max_sequence_length": config.data["max_input_seq_length"], "placement_fingerprint": self._placement_plan.placement_hash, + "packing_algorithm": config.policy["sequence_packing"]["algorithm"] + if config.data["energon"].packing_buffer_size is not None + else None, + # This caps sources per physical pack. Energon's similarly named + # max_samples_per_sequence instead controls sequential shard reads. + "max_sequences_per_bin": config.policy["sequence_packing"].get( + "max_sequences_per_bin" + ), + "sequence_length_pad_multiple": config.policy[ + "make_sequence_length_divisible_by" + ], + "only_unmask_final": config.sft.only_unmask_final, } if self._loader_states is None: futures = self._trainer.worker_group.run_all_workers_single_data( @@ -385,11 +398,25 @@ def setup_sft_v2( raise ValueError("SFTv2 requires data.backend=energon.") if not master_config.policy["megatron_cfg"]["enabled"]: raise ValueError("SFTv2 supports only the Megatron policy backend.") - if ( - master_config.policy["sequence_packing"]["enabled"] - or master_config.policy["dynamic_batching"]["enabled"] - ): - raise ValueError("SFTv2 requires fixed NeMo-RL batching.") + sequence_packing = master_config.policy["sequence_packing"] + dynamic_batching = master_config.policy["dynamic_batching"] + energon_packing = master_config.data["energon"].packing_buffer_size is not None + if not energon_packing: + if sequence_packing["enabled"] or dynamic_batching["enabled"]: + raise ValueError("SFTv2 without Energon packing requires fixed batching.") + else: + if not sequence_packing["enabled"] or not sequence_packing.get( + "fuse_loss", False + ): + raise ValueError( + "Energon packing requires sequence_packing enabled with fuse_loss." + ) + if sequence_packing.get("algorithm") not in { + algorithm.value for algorithm in PackingAlgorithm + }: + raise ValueError("Energon SFT requires a supported packing algorithm.") + if dynamic_batching["enabled"]: + raise ValueError("Energon packing does not support dynamic batching.") # SFTConfig carries validation knobs that default to on (val_period=10, # val_at_start=True) and this loop has no validation path, so reject them # rather than accepting a config whose validation silently never runs. @@ -418,6 +445,43 @@ def setup_sft_v2( max_sequence_length = master_config.data["max_input_seq_length"] if max_sequence_length is None: raise ValueError("SFTv2 requires data.max_input_seq_length.") + if energon_packing: + megatron_cfg = master_config.policy["megatron_cfg"] + if ( + megatron_cfg.get("moe_token_dispatcher_type") == "flex" + and megatron_cfg.get("moe_flex_dispatcher_backend") == "hybridep" + ): + raise ValueError("Energon packing does not support HybridEP flex dispatch.") + + cp_size = megatron_cfg["context_parallel_size"] + tp_size = megatron_cfg["tensor_model_parallel_size"] + pad_multiple = master_config.policy["make_sequence_length_divisible_by"] + parallel_multiple = (2 * cp_size if cp_size > 1 else 1) * ( + tp_size if tp_size > 1 and megatron_cfg["sequence_parallel"] else 1 + ) + if pad_multiple % parallel_multiple != 0: + raise ValueError( + "Energon packing requires make_sequence_length_divisible_by to " + f"be a multiple of {parallel_multiple}." + ) + if max_sequence_length % pad_multiple != 0: + raise ValueError( + "Energon packing requires max_input_seq_length to be divisible by " + "make_sequence_length_divisible_by." + ) + + fp8_cfg = megatron_cfg.get("fp8_cfg") or {} + if fp8_cfg.get("enabled", False): + fp8_multiple = { + "blockwise": 128, + "mxfp8": 32, + }.get(fp8_cfg["fp8_recipe"], 16) + fp8_multiple *= parallel_multiple + if max_sequence_length % fp8_multiple != 0: + raise ValueError( + "Energon packing requires max_input_seq_length to be divisible " + f"by the FP8 packed-token alignment ({fp8_multiple})." + ) processor = None tokenizer = tokenizer_or_processor diff --git a/nemo_rl/data/energon/config.py b/nemo_rl/data/energon/config.py index 9ffb52422f7..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: None = None - packing_buffer_size: None = None + max_samples_per_sequence: ( + Annotated[ + int, + Field( + ge=1, + description="Maximum sequential sample run used when sharding a dataset.", + ), + ] + | None + ) = None + packing_buffer_size: ( + Annotated[ + int, + Field( + ge=1, + description="Samples buffered by Energon for packing; None disables packing.", + ), + ] + | None + ) = None batch_grouping: Literal["auto"] = "auto" processor_adapter: Literal["hf_multimodal"] = "hf_multimodal" topology_mapper: Literal["default"] = "default" diff --git a/nemo_rl/data/energon/multimodal/packing.py b/nemo_rl/data/energon/multimodal/packing.py new file mode 100644 index 00000000000..f086a0e18b2 --- /dev/null +++ b/nemo_rl/data/energon/multimodal/packing.py @@ -0,0 +1,199 @@ +# 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.""" + +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.multimodal_utils import PackedTensor +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.") + packed_logs: list[list[dict[str, Any]]] = [] + boundaries: list[torch.Tensor | None] = [] + padded_boundaries: list[torch.Tensor | None] = [] + 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 + ) + 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 + 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) + 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) + 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_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(), + # 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, + } + ) + 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..14f8c5d17ec 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 @@ -226,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 @@ -251,7 +259,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 +268,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 +290,46 @@ 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 @stateless - def batch(self, samples: list[EncodedSFTSample]) -> BatchedDataDict[Any]: + 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, + ) + + @stateless + 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 | 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..34cbf1f0564 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,13 @@ def _loader_identity( batch_size: int, 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.""" - return { + identity = { "source": source.model_dump(mode="json"), "loader": loader_config.model_dump(mode="json"), "adapter": adapter_fingerprint, @@ -311,6 +316,14 @@ def _loader_identity( ), "topology": topology, } + 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 def _worker_config( @@ -338,6 +351,12 @@ def _task_encoder( loader_config: EnergonLoaderConfig, 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, + only_unmask_final: bool, ) -> BaseSFTTaskEncoder: cooker_functions = [ Cooker( @@ -353,12 +372,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=max_sequences_per_bin, + ) + 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 +408,10 @@ def build_energon_sft_loader( logical_rank: int, logical_world_size: int, placement_fingerprint: str, + 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: @@ -388,6 +425,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 +439,12 @@ def build_energon_sft_loader( loader_config=loader_config, 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, + only_unmask_final=only_unmask_final, ) worker_config = _worker_config( loader_config, @@ -421,9 +466,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 +480,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 +513,10 @@ def build_energon_sft_loader( logical_rank=logical_rank, 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 8af6ccb7af9..c0703a4b5f2 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,10 @@ def setup_sft_dataloader( batch_size: int, 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, ) -> bool: """Build the train loader on the TP0/PP0/CP0 rank of this DP replica.""" @@ -86,6 +91,10 @@ def setup_sft_dataloader( logical_rank=logical_rank, 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, ) if restored_state is not None: self._sft_loader.load_state_dict(restored_state) @@ -164,6 +173,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/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 4f047bb21e9..5124b9e452c 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -69,6 +69,7 @@ { "NemotronNanoVLV2Processor", "NemotronH_Nano_Omni_Reasoning_V3Processor", + "NemotronH_Omni_Reasoning_V3Processor", } ) @@ -165,7 +166,8 @@ def row_shapes_key(field: str) -> str: # Keys inside the :func:`row_shapes_key` tag value. A plain dict rather than a # record because ``tags`` rides TQ's own serializer. ROW_GEOMETRY_SHAPES = "shapes" -ROW_GEOMETRY_PAD = "pad" +ROW_GEOMETRY_PREPROCESS_MODE = "preprocess_mode" +ROW_GEOMETRY_PREPROCESS_KWARGS = "preprocess_kwargs" # Include-list of multimodal fields every forward-running dispatch (logprob @@ -199,7 +201,7 @@ def multimodal_row_tags( contents. Carries ``shapes`` (per-row, and unrecoverable once ``to_wire`` flattens) - and ``pad`` (the field's policy flag). Deliberately *not* a pad target: the + and the field's preprocessing settings. Deliberately *not* a pad target: the width padding lands at is scratch that the model discards -- mcore crops it via ``imgs_sizes`` before patchification, and the AutoModel path rejects mixed-resolution batches outright -- so each consumer pads to its own view @@ -223,11 +225,11 @@ def multimodal_row_tags( "sample_ids, so a disagreement here would pair one sample's " "pixels with another's shapes." ) - pad = value.pad_to_max_shape for row, row_shapes in enumerate(shapes): tags[row][row_shapes_key(key)] = { ROW_GEOMETRY_SHAPES: row_shapes, - ROW_GEOMETRY_PAD: pad, + ROW_GEOMETRY_PREPROCESS_MODE: value.preprocess_mode, + ROW_GEOMETRY_PREPROCESS_KWARGS: dict(value.preprocess_kwargs), } # ``None`` rather than ``B`` empty dicts: a text-only run has no packed # field, and an all-empty tags list would still be pickled on every @@ -272,16 +274,80 @@ def reassemble_packed_multimodal( + ". to_wire flattens each row, so without the companion the " "true per-segment shapes are unrecoverable." ) - # Indexed, not ``.get``-with-default: a producer-side rename of either - # key must fail here rather than silently restore ``pad=False``, which - # changes what ``as_tensor`` hands the vision encoder. + # Indexed, not ``.get``-with-default: a producer-side rename of any key + # must fail here rather than silently change what ``as_tensor`` hands + # the vision encoder. fields[key] = PackedTensor.from_wire( value, [[] if r is None else r[ROW_GEOMETRY_SHAPES] for r in rows], # type: ignore[union-attr] - pad_to_max_shape=bool(present[0][ROW_GEOMETRY_PAD]), + preprocess_mode=present[0][ROW_GEOMETRY_PREPROCESS_MODE], + preprocess_kwargs=present[0][ROW_GEOMETRY_PREPROCESS_KWARGS], ) +def _patchify_segments(segments: list[torch.Tensor], *, patch_dim: int) -> torch.Tensor: + """Cut pixel segments into vision patches and pack them into one sequence. + + Each ``[N, channels, H, W]`` segment is processed at its native resolution + into a ``[C_i, P²]`` block, where ``C_i`` is its spatial patch count and + ``P²`` is the flattened patch width (``channels * patch_dim**2``). Blocks + are packed along dimension zero, then a batch dimension is added to produce + ``[1, total_C, P²]``. Already-patchified segments in that final layout are + accepted so repeated materialization is safe. + """ + if patch_dim <= 0: + raise ValueError(f"patch_dim must be positive, got {patch_dim}") + + flattened: list[torch.Tensor] = [] + for segment in segments: + if segment.ndim == 3: + if segment.shape[0] != 1: + raise ValueError( + "Pre-patchified segments must be [1, total_C, P²], " + f"got shape {tuple(segment.shape)}" + ) + flattened.append(segment[0]) + continue + if segment.ndim != 4: + raise ValueError( + "patchify expects [N, C, H, W] pixel segments or " + "[1, total_C, P²] pre-patchified segments, got shape " + f"{tuple(segment.shape)}" + ) + count, channels, height, width = segment.shape + if height % patch_dim or width % patch_dim: + raise ValueError( + f"Image size {(height, width)} is not divisible by " + f"patch_dim={patch_dim}" + ) + rows = height // patch_dim + columns = width // patch_dim + flattened.append( + segment.reshape(count, channels, rows, patch_dim, columns, patch_dim) + .permute(0, 2, 4, 1, 3, 5) + .reshape(count * rows * columns, channels * patch_dim * patch_dim) + ) + + widths = {tensor.shape[-1] for tensor in flattened} + if len(widths) != 1: + raise ValueError( + f"patchify produced mismatched P² widths {sorted(widths)}; " + "the segments do not share a channel count" + ) + return torch.cat(flattened, dim=0).unsqueeze(0).contiguous() + + +def _shared_preprocess_spec( + from_packed_tensors: list["PackedTensor"], +) -> dict[str, Any]: + """Return the preprocessing setting shared by every input.""" + first = from_packed_tensors[0]._preprocess_spec + assert all( + packed_tensor._preprocess_spec == first for packed_tensor in from_packed_tensors + ), "All packed tensors must have the same preprocess setting" + return first + + class PackedTensor: """A logical batch of rows backed by packable tensor segments. @@ -325,7 +391,8 @@ def __init__( tensors: Union[torch.Tensor, list[Optional[torch.Tensor]], list[None]], dim_to_pack: int, *, - pad_to_max_shape: bool = False, + preprocess_mode: Optional[str] = None, + preprocess_kwargs: Optional[dict[str, Any]] = None, _row_offsets: Optional[list[int]] = None, _segment_indices: Optional[list[int]] = None, _segment_provenance: Optional[list[bytes]] = None, @@ -336,8 +403,10 @@ def __init__( tensors: A tensor or list of per-item tensors. List entries may be ``None`` for items without this modality. dim_to_pack: Dimension along which ``as_tensor`` concatenates. - pad_to_max_shape: Pad every non-packing dimension to its batch-wide - maximum before concatenating. All tensors must have the same rank. + preprocess_mode: Optional preprocessing applied by ``as_tensor``. + Supported values are ``pad_to_max_shape`` and ``patchify``. + preprocess_kwargs: Extra arguments for ``preprocess_mode``. Patchify + accepts ``patch_dim``. """ assert tensors is not None, "Input tensors to PackedTensor cannot be None" @@ -354,7 +423,13 @@ def __init__( f"Unsupported type for input tensors to PackedTensor: {type(tensors)}" ) self.dim_to_pack = dim_to_pack - self.pad_to_max_shape = pad_to_max_shape + if preprocess_mode not in (None, "pad_to_max_shape", "patchify"): + raise ValueError( + f"Unknown preprocess_mode {preprocess_mode!r}; expected None, " + "'pad_to_max_shape', or 'patchify'" + ) + self.preprocess_mode = preprocess_mode + self.preprocess_kwargs: dict[str, Any] = dict(preprocess_kwargs or {}) if (_row_offsets is None) != (_segment_indices is None): raise ValueError( "_row_offsets and _segment_indices must either both be set or both be None" @@ -395,6 +470,14 @@ def __setstate__(self, state: dict[str, Any]) -> None: self.__dict__.setdefault("_segment_indices", None) self.__dict__.setdefault("_segment_provenance", None) + @property + def _preprocess_spec(self) -> dict[str, Any]: + """Return keyword arguments that preserve preprocessing in a copy.""" + return { + "preprocess_mode": self.preprocess_mode, + "preprocess_kwargs": self.preprocess_kwargs, + } + @property def deduplication_enabled(self) -> bool: """Whether this value carries stable physical-segment provenance.""" @@ -445,7 +528,7 @@ def __deepcopy__(self, memo: dict[int, Any]) -> "PackedTensor": copied = PackedTensor( [deepcopy(item, memo) for item in self.tensors], self.dim_to_pack, - pad_to_max_shape=self.pad_to_max_shape, + **self._preprocess_spec, ) else: copied = PackedTensor( @@ -455,7 +538,7 @@ def __deepcopy__(self, memo: dict[int, Any]) -> "PackedTensor": else [deepcopy(item, memo) for item in self.tensors] ), self.dim_to_pack, - pad_to_max_shape=self.pad_to_max_shape, + **self._preprocess_spec, _row_offsets=( list(self._row_offsets) if self._row_offsets is not None else None ), @@ -488,12 +571,21 @@ def as_tensor( if len(non_none_tensors) == 0: return None + if self.preprocess_mode == "patchify": + if self.dim_to_pack != 0: + raise ValueError( + f"patchify requires dim_to_pack=0, got {self.dim_to_pack}" + ) + return _patchify_segments(non_none_tensors, **self.preprocess_kwargs).to( + device + ) + # Some multimodal processors produce a different shape per prompt, # such as dynamic-resolution images, variable-frame videos, or audio # feature sequences. Concatenation already permits the packing # dimension to vary; when explicitly requested, pad every other # dimension to the largest size in the batch. - if self.pad_to_max_shape: + if self.preprocess_mode == "pad_to_max_shape": ranks = {tensor.ndim for tensor in non_none_tensors} if len(ranks) != 1: raise ValueError( @@ -599,7 +691,7 @@ def converted(item: Optional[torch.Tensor]) -> Optional[torch.Tensor]: else list(self.tensors) ), self.dim_to_pack, - pad_to_max_shape=self.pad_to_max_shape, + **self._preprocess_spec, _row_offsets=( list(self._row_offsets) if self._row_offsets is not None else None ), @@ -626,7 +718,7 @@ def slice(self, indices: Union[list[int], torch.Tensor]) -> "PackedTensor": return PackedTensor( tensors, self.dim_to_pack, - pad_to_max_shape=self.pad_to_max_shape, + **self._preprocess_spec, ) physical_remap: dict[int, int] = {} @@ -650,7 +742,7 @@ def slice(self, indices: Union[list[int], torch.Tensor]) -> "PackedTensor": return PackedTensor( tensors, self.dim_to_pack, - pad_to_max_shape=self.pad_to_max_shape, + **self._preprocess_spec, _row_offsets=row_offsets, _segment_indices=segment_indices, _segment_provenance=( @@ -672,7 +764,7 @@ def empty_rows_like(cls, other: "PackedTensor", num_rows: int) -> "PackedTensor" return cls( [], other.dim_to_pack, - pad_to_max_shape=other.pad_to_max_shape, + **other._preprocess_spec, _row_offsets=[0] * (num_rows + 1), _segment_indices=[], _segment_provenance=[], @@ -681,7 +773,7 @@ def empty_rows_like(cls, other: "PackedTensor", num_rows: int) -> "PackedTensor" return cls( [], other.dim_to_pack, - pad_to_max_shape=other.pad_to_max_shape, + **other._preprocess_spec, _row_offsets=[0], _segment_indices=[], _segment_provenance=None, @@ -689,7 +781,7 @@ def empty_rows_like(cls, other: "PackedTensor", num_rows: int) -> "PackedTensor" return cls( [None] * num_rows, other.dim_to_pack, - pad_to_max_shape=other.pad_to_max_shape, + **other._preprocess_spec, ) @classmethod @@ -718,10 +810,7 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor": assert len(set(dim_to_packs)) == 1, ( "All packed tensors must have the same dim_to_pack" ) - pad_to_max_shapes = [batch.pad_to_max_shape for batch in from_packed_tensors] - assert len(set(pad_to_max_shapes)) == 1, ( - "All packed tensors must have the same pad_to_max_shape setting" - ) + preprocess_spec = _shared_preprocess_spec(from_packed_tensors) if any( packed_tensor.deduplication_enabled or packed_tensor._row_offsets is not None @@ -762,7 +851,7 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor": return cls( tensors, dim_to_packs[0], - pad_to_max_shape=pad_to_max_shapes[0], + **preprocess_spec, _row_offsets=row_offsets, _segment_indices=segment_indices, _segment_provenance=provenances, @@ -776,7 +865,7 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor": return cls( tensors, dim_to_pack, - pad_to_max_shape=pad_to_max_shapes[0], + **preprocess_spec, ) @classmethod @@ -800,7 +889,7 @@ def merge_segments( return cls( concatenated.tensors, concatenated.dim_to_pack, - pad_to_max_shape=concatenated.pad_to_max_shape, + **concatenated._preprocess_spec, _row_offsets=[0, len(concatenated._segment_indices)], _segment_indices=concatenated._segment_indices, _segment_provenance=concatenated._segment_provenance, @@ -836,10 +925,7 @@ def flattened_concat( assert len(set(dim_to_packs)) == 1, ( "All packed tensors must have the same dim_to_pack" ) - pad_to_max_shapes = [batch.pad_to_max_shape for batch in from_packed_tensors] - assert len(set(pad_to_max_shapes)) == 1, ( - "All packed tensors must have the same pad_to_max_shape setting" - ) + preprocess_spec = _shared_preprocess_spec(from_packed_tensors) if any( packed_tensor.deduplication_enabled or packed_tensor._row_offsets is not None @@ -854,7 +940,7 @@ def flattened_concat( return cls( tensors, from_packed_tensors[0].dim_to_pack, - pad_to_max_shape=pad_to_max_shapes[0], + **preprocess_spec, ) # ── Wire encoding (data-plane roundtrip) ───────────────────────── @@ -948,12 +1034,12 @@ def to_wire( # storage, and TQ never falls back to the deprecated strided layout. # * The per-row concat is 1-D, so it cannot raise on segments whose # trailing dims differ -- which is what previously forced - # ``pad_to_max_shape`` to pad *before* the concat. + # preprocessing to pad *before* the concat. # # The true shapes travel beside the payload (see the returned # ``shapes``) because TQ derives ``per_sample_shapes`` from what it is # handed: give it flat rows and it records flat lengths. Padding still - # happens for ``pad_to_max_shape`` values, but in worker memory at use + # happens for values that need it, but in worker memory at use # time via :meth:`as_tensor`, not on the wire. shapes = self._shapes_of(row_segments) # ``reshape(-1)`` on contiguous processor output is a view, so the @@ -990,7 +1076,8 @@ def from_wire( nested: torch.Tensor, shapes: list[list[list[int]]], *, - pad_to_max_shape: bool = False, + preprocess_mode: Optional[str] = None, + preprocess_kwargs: Optional[dict[str, Any]] = None, ) -> Optional["PackedTensor"]: """Reconstruct from the value produced by :meth:`to_wire`. @@ -1009,10 +1096,9 @@ def from_wire( reconstructs as legacy does: ``as_tensor`` returns ``None`` and ``logical_segment_counts_by_row`` reports 0 rather than 1. - ``pad_to_max_shape`` is restored onto the rebuilt value as a flag, not + The preprocessing settings are restored onto the rebuilt value, not materialized. Segments come back at their true shapes and stay separate - via the CSR row map, so nothing is padded or concatenated here; - :meth:`as_tensor` pads at use time. + via the CSR row map; :meth:`as_tensor` preprocesses at use time. Mirrors :meth:`to_wire`; both assume ``dim_to_pack=0``. """ @@ -1053,7 +1139,8 @@ def from_wire( return cls( segments_flat, # type: ignore[arg-type] dim_to_pack=0, - pad_to_max_shape=pad_to_max_shape, + preprocess_mode=preprocess_mode, + preprocess_kwargs=preprocess_kwargs, _row_offsets=row_offsets, _segment_indices=list(range(len(segments_flat))), ) @@ -1071,7 +1158,7 @@ def encode_multimodal_for_wire( Payload only. Per-token fields ride rectangular; packed fields ride as one flattened ``torch.jagged`` value. The geometry :meth:`PackedTensor.from_wire` needs to undo that flattening -- per-row segment shapes plus the - ``pad_to_max_shape`` flag -- is minted separately by + preprocessing settings -- is minted separately by :func:`multimodal_row_tags` and shipped on ``KVBatchMeta.tags``. TQ cannot derive it, because it reads ``per_sample_shapes`` off the flattened rows it is handed. @@ -1181,9 +1268,14 @@ def get_dim_to_pack_along(processor, key: str) -> int: return 0 -def get_pad_to_max_shape(processor: Any, key: str) -> bool: - """Return whether a processor input must pad non-packing dimensions.""" - return uses_image_placeholder(processor) and key == "pixel_values" +def get_preprocess(processor: Any, key: str) -> dict[str, Any]: + """Return materialization preprocessing for one processor input.""" + if uses_image_placeholder(processor) and key == "pixel_values": + return { + "preprocess_mode": "patchify", + "preprocess_kwargs": {"patch_dim": 16}, + } + return {"preprocess_mode": None, "preprocess_kwargs": {}} def extract_multimodal_model_inputs( @@ -1245,7 +1337,7 @@ def extract_multimodal_model_inputs( extracted[key] = PackedTensor( value, dim_to_pack=get_dim_to_pack_along(processor, key), - pad_to_max_shape=get_pad_to_max_shape(processor, key), + **get_preprocess(processor, key), ) for key in ("token_type_ids", "mm_token_type_ids"): @@ -1385,13 +1477,13 @@ def media_sources_equal( def _materialize_ragged_pixel_values( processed: dict[str, Any], processor: Any ) -> dict[str, Any]: - """Fold a ragged per-image ``pixel_values`` list into one padded tensor. + """Fold a ragged per-image ``pixel_values`` list into one patch sequence. Processors with dynamic per-image resolution return a list of CHW tensors rather than a stacked batch. ``imgs_sizes`` is derived from the *unpadded* shapes first, since those exact sizes are what the projector slices with; - padding happens afterwards so downstream sees the single tensor its - torch.Tensor contract expects. + patchification happens afterwards so downstream sees the single tensor its + ``torch.Tensor`` contract expects. """ processed = dict(processed) pixel_values = processed.get("pixel_values") @@ -1417,7 +1509,7 @@ def _materialize_ragged_pixel_values( def _stack_ragged_pixel_values( processed: dict[str, Any], tiles: list[torch.Tensor], processor: Any ) -> None: - """Derive imgs_sizes from unpadded shapes, then pad into one tensor.""" + """Derive image sizes, then patchify native-shape tiles into one tensor.""" if uses_image_placeholder(processor) and "imgs_sizes" not in processed: processed["imgs_sizes"] = torch.tensor( [[int(item.shape[-2]), int(item.shape[-1])] for item in tiles], @@ -1426,7 +1518,8 @@ def _stack_ragged_pixel_values( stacked = PackedTensor( [item.unsqueeze(0) for item in tiles], dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, ).as_tensor() assert stacked is not None processed["pixel_values"] = stacked 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..cd1f4b19ec6 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 = 0, + ) -> 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/data/processors.py b/nemo_rl/data/processors.py index 488a6fcd3e8..7e266e775d2 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -389,6 +389,7 @@ def vlm_preference_preprocessor( placeholder_style_processors = { "NemotronNanoVLV2Processor", "NemotronH_Nano_Omni_Reasoning_V3Processor", + "NemotronH_Omni_Reasoning_V3Processor", } message_processor = ( _NemotronOmniPreferenceProcessorProxy(processor) @@ -404,25 +405,33 @@ def _format_branch(completion: dict[str, Any]) -> VLMMessageLogType: task_data_spec, ) - # Mirror the canonical Nemotron Omni metadata contract. Dynamic-resolution - # image batches may differ spatially across rows, while imgs_sizes - # preserves the true crop consumed by model-owned patchification. + # Mirror the canonical Nemotron Omni metadata. Record native image sizes + # before patchification removes the spatial dimensions. for raw_message in message_log: message = cast(Any, raw_message) pixel_values = message.get("pixel_values") if not isinstance(pixel_values, PackedTensor): continue - pixel_values.pad_to_max_shape = True - pixels = pixel_values.as_tensor() - if pixels is not None and pixels.ndim == 4 and "imgs_sizes" not in message: - num_images, _, height, width = pixels.shape + if "imgs_sizes" not in message: + image_sizes: list[list[int]] = [] + for pixels in pixel_values.iter_logical_segments(): + if pixels is None: + continue + if pixels.ndim != 4: + raise ValueError( + "Nemotron Omni pixel values must be [N, C, H, W] " + f"before patchification, got {tuple(pixels.shape)}" + ) + image_sizes.extend( + [[int(pixels.shape[-2]), int(pixels.shape[-1])]] + * int(pixels.shape[0]) + ) message["imgs_sizes"] = PackedTensor( - torch.tensor( - [[height, width]] * num_images, - dtype=torch.long, - ), + torch.tensor(image_sizes, dtype=torch.long), dim_to_pack=0, ) + pixel_values.preprocess_mode = "patchify" + pixel_values.preprocess_kwargs = {"patch_dim": 16} imgs_sizes = message.get("imgs_sizes") if isinstance(imgs_sizes, PackedTensor) and "num_frames" not in message: sizes = imgs_sizes.as_tensor() diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index a9b7ab95932..f6fc3017d8a 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -116,7 +116,8 @@ def _broadcast_batched_data_dict( "empty_packed", len(v), v.dim_to_pack, - v.pad_to_max_shape, + v.preprocess_mode, + v.preprocess_kwargs, ) ) continue @@ -130,7 +131,8 @@ def _broadcast_batched_data_dict( str(values.device), nested.offsets().tolist(), shapes, - v.pad_to_max_shape, + v.preprocess_mode, + v.preprocess_kwargs, ) ) elif ( @@ -206,7 +208,14 @@ def _broadcast_batched_data_dict( ): out[key] = tensor.to(src_device) elif kind == "packed_wire": - dtype_str, src_device, offsets, shapes, pad_to_max_shape = entry[2:] + ( + dtype_str, + src_device, + offsets, + shapes, + preprocess_mode, + preprocess_kwargs, + ) = entry[2:] if is_leader: flat = leader_flat[key].to(bcast_device) else: @@ -224,17 +233,21 @@ def _broadcast_batched_data_dict( if torch.device(src_device).type != torch.device(bcast_device).type: nested = nested.to(src_device) out[key] = PackedTensor.from_wire( - nested, shapes, pad_to_max_shape=pad_to_max_shape + nested, + shapes, + preprocess_mode=preprocess_mode, + preprocess_kwargs=preprocess_kwargs, ) elif kind == "empty_packed": # Structural only: no payload, so followers rebuild from the # geometry and land on the leader's key set. - n_rows, dim_to_pack, pad_to_max_shape = entry[2:] + n_rows, dim_to_pack, preprocess_mode, preprocess_kwargs = entry[2:] if not is_leader: out[key] = PackedTensor( [None] * n_rows, dim_to_pack, - pad_to_max_shape=pad_to_max_shape, + preprocess_mode=preprocess_mode, + preprocess_kwargs=preprocess_kwargs, ) else: if not is_leader: diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index dfb03b24b98..639890e062c 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 @@ -283,10 +283,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"]: @@ -359,6 +374,104 @@ 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, 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] + 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 _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: + 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] + pack_length = int(padded[-1]) + if ( + cu.shape != padded.shape + or cu.numel() < 2 + or int(cu[0]) != 0 + or int(padded[0]) != 0 + or pack_length > input_ids.shape[1] + or bool((source_lengths <= 0).any()) + or bool((source_lengths > padded_lengths).any()) + ): + raise ValueError("Invalid prepacked source boundaries.") + batch_size, sequence_length = input_ids.shape[:2] + for key, value in list(data.items()): + if ( + key in {"cu_seqlens", "cu_seqlens_padded"} + or not torch.is_tensor(value) + or value.ndim < 2 + or value.shape[0] != batch_size + or value.shape[1] != sequence_length + ): + continue + data[key] = value[:, :pack_length].contiguous() + input_ids = data["input_ids"] + 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, + 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_cp_sharded.shape[1], + ) + return input_ids, input_ids_cp_sharded, params, padded + + def process_microbatch( data_dict: BatchedDataDict[Any], seq_length_key: Optional[str] = None, @@ -425,7 +538,46 @@ 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, + model_slices_context_parallel_inputs=( + model_slices_context_parallel_inputs + ), + ) + original_seq_length = input_ids.shape[1] + routed_experts = data_dict.get("routed_experts") + routed_experts_cp_sharded = routed_experts + if ( + routed_experts is not None + and not model_slices_context_parallel_inputs + ): + routed_experts_cp_sharded = _slice_prepacked_for_cp( + routed_experts, cu_seqlens_padded + ) + 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: 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 ae674a17c53..2b34e0f21b4 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -48,6 +48,8 @@ DP_TRAIN_FIELDS, GLOBAL_FORWARD_PAD_SEQLEN, LP_SEED_FIELDS, + MICRO_BATCH_INDICES, + MICRO_BATCH_LENGTHS, ROUTE_PASSTHROUGH_FLAG, ROUTE_PLAN_TAG, fields_with_optional_opd_full, @@ -576,10 +578,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") diff --git a/pyrefly.toml b/pyrefly.toml index 50b46c69694..1e919b9d9fc 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -123,6 +123,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", 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-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 "$@" 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/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index d4e0da4b969..f5a4adab36d 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -241,7 +241,8 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path): assert result["task_name"] == "mmpr-tiny" user_message = result["message_log"][0] assert torch.equal(user_message["num_frames"].as_tensor(), torch.tensor([1])) - assert user_message["pixel_values"].pad_to_max_shape is True + assert user_message["pixel_values"].preprocess_mode == "patchify" + assert user_message["pixel_values"].preprocess_kwargs == {"patch_dim": 16} assert user_message["pixel_values"].as_tensor().dtype == torch.float32 def test_text_only_row_preserves_formatted_vllm_content(self): 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/packing/test_knapsack.py b/tests/unit/data/packing/test_knapsack.py new file mode 100644 index 00000000000..413c8c1171e --- /dev/null +++ b/tests/unit/data/packing/test_knapsack.py @@ -0,0 +1,58 @@ +# 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 ( + 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..9b371b95cf4 --- /dev/null +++ b/tests/unit/data/test_energon_packing.py @@ -0,0 +1,147 @@ +# 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 +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 PackingAlgorithm, get_packer # 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, + ) + + +@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), + _sample("s2", 3, group="image"), + ] + + selected = select_samples_to_pack( + samples, + 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"}, + ] + + +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, + ) + + second_pack = pack_selected_samples( + [_sample("s2", 4)], + pack_capacity=12, + sequence_length_pad_multiple=4, + ) + prepared = prepare_packed_sft_batch( + [packed, second_pack], tokenizer=_Tokenizer(), only_unmask_final=False + ) + + 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, 4] + assert prepared["source_ids"] == [["s0", "s1"], ["s2"]] + assert isinstance(prepared["cu_seqlens"], PackedTensor) + assert isinstance(prepared["cu_seqlens_padded"], PackedTensor) + 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, 4] + + +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]] + + +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, + ) 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/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py index c88883316a2..cf7f71e2d76 100644 --- a/tests/unit/data/test_multimodal_dict.py +++ b/tests/unit/data/test_multimodal_dict.py @@ -23,8 +23,10 @@ PER_TOKEN_MULTIMODAL_FIELDS, PackedTensor, encode_multimodal_for_wire, + get_preprocess, multimodal_row_tags, reassemble_packed_multimodal, + uses_image_placeholder, ) from nemo_rl.distributed.batched_data_dict import ( BatchedDataDict, @@ -55,6 +57,28 @@ def test_packed_data_basic(): assert torch.equal(batch.as_tensor(), expected_tensor) +@pytest.mark.parametrize( + "processor_name", + [ + "NemotronNanoVLV2Processor", + "NemotronH_Nano_Omni_Reasoning_V3Processor", + "NemotronH_Omni_Reasoning_V3Processor", + ], +) +def test_placeholder_processors_use_patchify(processor_name): + processor = type(processor_name, (), {})() + + assert uses_image_placeholder(processor) + assert get_preprocess(processor, "pixel_values") == { + "preprocess_mode": "patchify", + "preprocess_kwargs": {"patch_dim": 16}, + } + assert get_preprocess(processor, "imgs_sizes") == { + "preprocess_mode": None, + "preprocess_kwargs": {}, + } + + def test_shard_by_batch_size_with_packed_data(): """Test shard_by_batch_size with packed multimodal data.""" # Create sample data @@ -387,7 +411,7 @@ def test_packedtensor_pads_mixed_dynamic_resolution_images(): second = 2 * torch.ones(1, 3, 4, 2) packed = PackedTensor( - [first, second], dim_to_pack=0, pad_to_max_shape=True + [first, second], dim_to_pack=0, preprocess_mode="pad_to_max_shape" ).as_tensor() assert packed.shape == (2, 3, 4, 4) @@ -412,7 +436,7 @@ def test_dynamic_resolution_padding_is_cropped_before_radio_patchification(): padded = PackedTensor( [small, large], dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", ).as_tensor() # Use nonzero garbage so this test cannot pass merely because F.pad uses zero. padded[0, :, 32:, :] = 123 @@ -461,7 +485,7 @@ def test_packedtensor_pad_to_max_shape_supports_audio_and_video( second = 2 * torch.ones(second_shape) packed = PackedTensor( - [first, second], dim_to_pack=0, pad_to_max_shape=True + [first, second], dim_to_pack=0, preprocess_mode="pad_to_max_shape" ).as_tensor() assert packed.shape == expected_shape @@ -476,7 +500,7 @@ def test_pad_to_max_shape_rejects_mismatched_ranks(): PackedTensor( [torch.ones(1, 3, 4), torch.ones(1, 3)], dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", ).as_tensor() @@ -485,7 +509,7 @@ def test_pad_to_max_shape_rejects_out_of_range_dim(): PackedTensor( [torch.ones(1, 3, 4), torch.ones(2, 3, 4)], dim_to_pack=3, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", ).as_tensor() @@ -493,25 +517,128 @@ def test_pad_to_max_shape_supports_negative_pack_dim(): packed = PackedTensor( [torch.ones(2, 3, 1), 2 * torch.ones(4, 3, 1)], dim_to_pack=-3, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", ).as_tensor() assert packed.shape == (6, 3, 1) -def test_slice_preserves_pad_to_max_shape_flag(): +def test_slice_preserves_preprocess_spec(): packed = PackedTensor( [torch.ones(1, 3, 2, 4), 2 * torch.ones(1, 3, 4, 2)], dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", + preprocess_kwargs={}, ) sliced = packed.slice([0, 1]) - assert sliced.pad_to_max_shape is True + assert sliced.preprocess_mode == "pad_to_max_shape" + assert sliced.preprocess_kwargs == {} assert sliced.as_tensor().shape == (2, 3, 4, 4) +def test_packedtensor_rejects_unknown_preprocess_mode(): + with pytest.raises(ValueError, match="Unknown preprocess_mode"): + PackedTensor( + torch.ones(1, 3, 4, 4), + dim_to_pack=0, + preprocess_mode="jagged", + ) + + +def test_patchify_packs_mixed_resolutions_without_padding(): + packed = PackedTensor( + [torch.ones(1, 3, 32, 32), 2 * torch.ones(1, 3, 64, 32)], + dim_to_pack=0, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, + ).as_tensor() + + assert packed.shape == (1, 12, 768) + assert torch.all(packed[0, :4] == 1) + assert torch.all(packed[0, 4:] == 2) + + +def test_patchify_preserves_pixel_order_within_a_patch(): + image = torch.arange(3 * 2 * 2, dtype=torch.float32).reshape(1, 3, 2, 2) + + packed = PackedTensor( + [image], + dim_to_pack=0, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 2}, + ).as_tensor() + + assert packed.shape == (1, 1, 12) + torch.testing.assert_close(packed[0, 0], image.reshape(12)) + + +def test_patchify_accepts_already_patchified_segments(): + raw = PackedTensor( + [torch.ones(1, 3, 32, 32)], + dim_to_pack=0, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, + ) + + once = raw.as_tensor() + assert once is not None + twice = PackedTensor( + [once], + dim_to_pack=0, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, + ).as_tensor() + + torch.testing.assert_close(once, twice) + + +def test_patchify_survives_flattened_concat(): + first = PackedTensor( + [torch.ones(1, 3, 32, 32)], + dim_to_pack=0, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, + ) + second = PackedTensor( + [2 * torch.ones(1, 3, 64, 32)], + dim_to_pack=0, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, + ) + + flattened = PackedTensor.flattened_concat([first, second]) + + assert len(flattened) == 2 + assert flattened.as_tensor().shape == (1, 12, 768) + torch.testing.assert_close( + flattened.as_tensor(), PackedTensor.concat([first, second]).as_tensor() + ) + + +def test_patchify_rejects_indivisible_image_size(): + with pytest.raises(ValueError, match="not divisible by patch_dim=16"): + PackedTensor( + [torch.ones(1, 3, 30, 32)], + dim_to_pack=0, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, + ).as_tensor() + + +def test_concat_rejects_mixed_preprocess_settings(): + padded = PackedTensor( + torch.ones(1, 3, 4, 4), + dim_to_pack=0, + preprocess_mode="pad_to_max_shape", + ) + plain = PackedTensor(torch.ones(1, 3, 4, 4), dim_to_pack=0) + + with pytest.raises(AssertionError, match="same preprocess setting"): + PackedTensor.concat([padded, plain]) + + def test_packedtensor_dedup_uses_provenance_not_prompt_position(): """Only segments descended from the same physical media are compacted.""" shared = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) @@ -559,12 +686,12 @@ def test_packedtensor_dedup_expands_before_dynamic_shape_padding(): first = PackedTensor( torch.ones(1, 1, 2), dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", ).enable_deduplication() second = PackedTensor( 2 * torch.ones(1, 2, 1), dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", ).enable_deduplication() packed = PackedTensor.concat([first, deepcopy(first), second]) @@ -632,25 +759,6 @@ def test_packedtensor_compact_dim_one_slice_empty_and_cloudpickle_roundtrip(): assert empty.as_tensor() is None -def test_packedtensor_unpickles_pre_deduplication_state(): - tensor = torch.tensor([[1.0], [2.0]]) - legacy = PackedTensor.__new__(PackedTensor) - legacy.__dict__ = { - "tensors": [tensor], - "dim_to_pack": 0, - "pad_to_max_shape": False, - } - - restored = cloudpickle.loads(cloudpickle.dumps(legacy, protocol=5)) - - assert not restored.deduplication_enabled - assert len(restored) == 1 - assert sum(restored.logical_segment_counts_by_row()) == 1 - torch.testing.assert_close(restored.as_tensor(), tensor) - restored.enable_deduplication() - assert restored.deduplication_enabled - - def test_packedtensor_empty_legacy_rows_survive_copy_pickle_and_slice(): legacy = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) empty = PackedTensor.empty_rows_like(legacy, 0) @@ -708,7 +816,7 @@ def test_to_wire_does_not_pad_segments_before_concat_under_dedup(): packed = PackedTensor( [torch.ones(1, 3, 2, 4), 2 * torch.ones(1, 3, 4, 2)], dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", _row_offsets=[0, 2], _segment_indices=[0, 1], ) @@ -721,7 +829,9 @@ def test_to_wire_does_not_pad_segments_before_concat_under_dedup(): assert [t.numel() for t in nested.unbind()] == [48] assert shapes == [[[1, 3, 2, 4], [1, 3, 4, 2]]] - restored = PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True).as_tensor() + restored = PackedTensor.from_wire( + nested, shapes, preprocess_mode="pad_to_max_shape" + ).as_tensor() assert torch.equal(restored, expected) @@ -758,7 +868,9 @@ def test_to_wire_does_not_materialize_pad_to_max_shape(): # Same rank, different trailing dims — nemotron-omni style tiles. first = torch.ones(1, 3, 2, 4) second = 2 * torch.ones(2, 3, 4, 2) - packed = PackedTensor([first, second], dim_to_pack=0, pad_to_max_shape=True) + packed = PackedTensor( + [first, second], dim_to_pack=0, preprocess_mode="pad_to_max_shape" + ) nested, shapes = packed.to_wire() rows = list(nested.unbind()) @@ -769,7 +881,9 @@ def test_to_wire_does_not_materialize_pad_to_max_shape(): # Padding is reapplied on read, reproducing the pre-wire as_tensor(). assert torch.equal( - PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True).as_tensor(), + PackedTensor.from_wire( + nested, shapes, preprocess_mode="pad_to_max_shape" + ).as_tensor(), packed.as_tensor(), ) @@ -836,7 +950,8 @@ def test_encode_multimodal_for_wire_packed_emits_single_nested_entry(): [[3, 4]], [[1, 4]], ] - assert tags[0]["pixel_values__row_shapes"]["pad"] is False + assert tags[0]["pixel_values__row_shapes"]["preprocess_mode"] is None + assert tags[0]["pixel_values__row_shapes"]["preprocess_kwargs"] == {} def test_multimodal_row_tags_does_not_encode_the_payload(): @@ -882,13 +997,20 @@ def test_reassemble_packed_multimodal_raises_without_companion(): def test_reassemble_packed_multimodal_round_trips_with_companion(): - packed = PackedTensor([torch.ones(3, 4), torch.ones(1, 4)], dim_to_pack=0) + packed = PackedTensor( + [torch.ones(1, 3, 32, 32), torch.ones(1, 3, 16, 32)], + dim_to_pack=0, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, + ) nested, _ = packed.to_wire() tags = multimodal_row_tags({"pixel_values": packed}, len(packed)) fields = {"pixel_values": nested} reassemble_packed_multimodal(fields, tags) + assert fields["pixel_values"].preprocess_mode == "patchify" + assert fields["pixel_values"].preprocess_kwargs == {"patch_dim": 16} assert torch.equal(fields["pixel_values"].as_tensor(), packed.as_tensor()) @@ -1034,11 +1156,13 @@ def test_to_wire_carries_mixed_rank_rows(): load-bearing. Reshaping on read restores the original ranks. """ rows = [torch.ones(1, 3, 2), torch.ones(2, 3)] - packed = PackedTensor(list(rows), dim_to_pack=0, pad_to_max_shape=True) + packed = PackedTensor(list(rows), dim_to_pack=0, preprocess_mode="pad_to_max_shape") nested, shapes = packed.to_wire() assert [t.numel() for t in nested.unbind()] == [6, 6] assert shapes == [[[1, 3, 2]], [[2, 3]]] - restored = PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True) + restored = PackedTensor.from_wire( + nested, shapes, preprocess_mode="pad_to_max_shape" + ) assert [tuple(t.shape) for t in restored.tensors] == [(1, 3, 2), (2, 3)] diff --git a/tests/unit/data/test_vlm_preference_processor.py b/tests/unit/data/test_vlm_preference_processor.py index 308212df768..3f59df888f4 100644 --- a/tests/unit/data/test_vlm_preference_processor.py +++ b/tests/unit/data/test_vlm_preference_processor.py @@ -101,6 +101,7 @@ def test_vlm_preference_processor_adds_nemotron_omni_media_metadata(): result["message_log_rejected"], ): message = message_log[0] - assert message["pixel_values"].pad_to_max_shape + assert message["pixel_values"].preprocess_mode == "patchify" + assert message["pixel_values"].preprocess_kwargs == {"patch_dim": 16} assert message["imgs_sizes"].as_tensor().tolist() == [[15, 23]] assert message["num_frames"].as_tensor().tolist() == [1] diff --git a/tests/unit/data_plane/test_leader_broadcast.py b/tests/unit/data_plane/test_leader_broadcast.py index 4523ae93356..c48ecb9bee6 100644 --- a/tests/unit/data_plane/test_leader_broadcast.py +++ b/tests/unit/data_plane/test_leader_broadcast.py @@ -91,7 +91,7 @@ def _packed(rows): return PackedTensor( [r.clone() if r is not None else None for r in rows], dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", ) @@ -148,7 +148,9 @@ def _all_empty_body(rank: int): { "input_ids": torch.arange(8, dtype=torch.long).reshape(2, 4), "pixel_values": PackedTensor( - [None, None], dim_to_pack=0, pad_to_max_shape=True + [None, None], + dim_to_pack=0, + preprocess_mode="pad_to_max_shape", ), } ) @@ -165,7 +167,8 @@ def _all_empty_body(rank: int): assert isinstance(packed, PackedTensor), type(packed).__name__ assert packed.logical_segment_counts_by_row() == [0, 0] assert packed.as_tensor() is None - assert packed.pad_to_max_shape is True + assert packed.preprocess_mode == "pad_to_max_shape" + assert packed.preprocess_kwargs == {} def _unsupported_type_body(rank: int): diff --git a/tests/unit/data_plane/test_local_sft.py b/tests/unit/data_plane/test_local_sft.py index 8dd2e426739..d5937c4645b 100644 --- a/tests/unit/data_plane/test_local_sft.py +++ b/tests/unit/data_plane/test_local_sft.py @@ -47,7 +47,7 @@ def _put_multimodal_batch( pixels = PackedTensor( [torch.full((1, 2), 1.0), torch.full((2, 2), 2.0)], dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="pad_to_max_shape", ).enable_deduplication() fields = local_batch_to_tensordict( { @@ -80,7 +80,7 @@ def test_local_round_trip_preserves_tensor_and_packed_tensor_fields() -> None: pixels = batch["pixel_values"] assert isinstance(pixels, PackedTensor) assert pixels.dim_to_pack == 0 - assert pixels.pad_to_max_shape + assert pixels.preprocess_mode == "pad_to_max_shape" assert pixels.deduplication_enabled assert pixels.logical_segment_counts_by_row() == [1, 1] assert torch.equal(pixels.tensors[0], torch.full((1, 2), 1.0)) diff --git a/tests/unit/environments/test_nemo_gym_image_placeholders.py b/tests/unit/environments/test_nemo_gym_image_placeholders.py index 5a05ddf366d..a5bf369acb5 100644 --- a/tests/unit/environments/test_nemo_gym_image_placeholders.py +++ b/tests/unit/environments/test_nemo_gym_image_placeholders.py @@ -60,7 +60,9 @@ def __call__(self, *, text, images, return_tensors): def _ragged(*shapes: tuple[int, ...]) -> NemotronNanoVLV2Processor: return NemotronNanoVLV2Processor( [torch.ones(*shape) for shape in shapes], - imgs_sizes=torch.tensor([[4, 4]] * len(shapes), dtype=torch.long), + imgs_sizes=torch.tensor( + [[shape[-2], shape[-1]] for shape in shapes], dtype=torch.long + ), ) @@ -83,9 +85,9 @@ def test_ragged_output_requested_only_for_multi_image_turns(): ) -def test_ragged_pixel_values_are_padded_to_one_tensor(): - """Heterogeneous CHW tensors become a single padded tensor for the message.""" - processor = _ragged((3, 2, 4), (3, 6, 4)) +def test_ragged_pixel_values_are_patchified_to_one_tensor(): + """Heterogeneous CHW tensors become one packed patch sequence.""" + processor = _ragged((3, 16, 32), (3, 32, 16)) user_message: dict = {} attach_image_model_inputs_to_message( user_message, @@ -94,10 +96,7 @@ def test_ragged_pixel_values_are_padded_to_one_tensor(): pad_dynamic_image_shapes=True, ) packed = user_message["pixel_values"].as_tensor() - # Two images, padded up to the tallest, channels preserved. - assert packed.shape[0] == 2 - assert packed.shape[-3] == 3 - assert packed.shape[-2] == 6 + assert packed.shape == (1, 4, 768) def test_ragged_pixel_values_reject_non_chw_entries(): diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py index 03f009dc491..21b95611feb 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,81 @@ 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=1 + ) + def test_process_microbatch_trims_prepacked_batch_padding( + self, mock_cp_world, mock_cp_rank + ): + from nemo_rl.models.megatron.data import process_microbatch + + data = self._prepacked_batch() + data["input_ids"] = torch.nn.functional.pad(data["input_ids"], (0, 4)) + data["token_mask"] = torch.nn.functional.pad(data["token_mask"], (0, 4)) + data["mtp_loss_mask"] = data["token_mask"].clone() + data["media_token_validity_mask"] = data["token_mask"].bool() + + result = process_microbatch( + data, + seq_length_key="input_lengths", + pack_sequences=True, + ) + + assert result.original_seq_length == 8 + assert result.input_ids.shape == (1, 8) + assert data["input_ids"].shape == (1, 8) + assert data["token_mask"].shape == (1, 8) + assert result.mtp_loss_mask.shape == (1, 8) + assert result.media_token_validity_mask.shape == (1, 8) + assert result.packed_seq_params.total_tokens == 8 + + @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/megatron/test_nemotron_omni_model.py b/tests/unit/models/megatron/test_nemotron_omni_model.py index 67990577c87..8ee53222459 100644 --- a/tests/unit/models/megatron/test_nemotron_omni_model.py +++ b/tests/unit/models/megatron/test_nemotron_omni_model.py @@ -246,7 +246,8 @@ def _deduplicated_expanded_fixture(device: torch.device): "pixel_values": PackedTensor( [image.clone(), image.clone()], dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, ), "imgs_sizes": PackedTensor( [image_size.clone(), image_size.clone()], @@ -257,7 +258,8 @@ def _deduplicated_expanded_fixture(device: torch.device): pixel_row = PackedTensor( image, dim_to_pack=0, - pad_to_max_shape=True, + preprocess_mode="patchify", + preprocess_kwargs={"patch_dim": 16}, ).enable_deduplication() image_size_row = PackedTensor( image_size, diff --git a/tests/unit/models/policy/test_tq_policy_placed.py b/tests/unit/models/policy/test_tq_policy_placed.py index 713aac7d2e9..00d886b2a31 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 @@ -92,15 +96,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()