diff --git a/docs/guides/nemotron-3-nano-omni.md b/docs/guides/nemotron-3-nano-omni.md index 0a35ba3cfee..306d57dcb59 100644 --- a/docs/guides/nemotron-3-nano-omni.md +++ b/docs/guides/nemotron-3-nano-omni.md @@ -1,6 +1,8 @@ # Nemotron 3 Nano Omni -This guide explains how to post-train the Nemotron 3 Nano Omni vision-language model with GRPO using NeMo RL on the AutoModel backend. +This guide explains how to post-train the Nemotron 3 Nano Omni vision-language model with GRPO using NeMo RL. Both the AutoModel and Megatron backends are supported for image-and-text training. + +## AutoModel backend It covers two recipes: @@ -9,7 +11,7 @@ It covers two recipes: Both share the same checkpoint, model code, and reward pipeline; they differ only in the dataset, reward functions, and node count. -## Recipe 1 — CLEVR-CoGenT (single-node) +### Recipe 1 — CLEVR-CoGenT (single-node) The CLEVR-CoGenT recipe uses [`examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml`](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml). It expects 8 GPUs on a single node, EP=8 across the experts, and TP=8 in vLLM. @@ -47,7 +49,7 @@ uv run examples/run_vlm_grpo.py --config examples/configs/recipes/vlm/vlm_grpo-n cluster.gpus_per_node=8 cluster.num_nodes=1 ``` -## Recipe 2 — MMPR-Tiny (4-node Slurm) +### Recipe 2 — MMPR-Tiny (4-node Slurm) The MMPR-Tiny recipe uses [`examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml`](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml). Differences vs. the CLEVR recipe: @@ -109,3 +111,37 @@ sbatch \ ``` To run on a different node count, change `NUM_NODES` and the `--nodes` flag. + +## Megatron backend + +The Megatron backend uses a dedicated `NemotronOmniModel` supplied by Megatron Bridge. The Hugging Face processor expands each image placeholder into the complete media-token sequence before the batch reaches the model. NeMo RL passes that expanded sequence and the image tensors to the model; `NemotronOmniModel` replaces the media-token positions with RADIO encoder outputs and then performs sequence packing and context-parallel sharding. + +This is the same model-owned packing boundary used by maintained Megatron VLM integrations. It differs from the historical Nemotron Omni `LLaVAModel` path, which collapsed the expanded media-token sequence before packing and expanded it again inside the model. The dedicated model removes that extra representation change and allows the integration to use Megatron Bridge and Megatron-LM from their maintained main branches. + +The current Megatron recipes cover Nano image-and-text GRPO. Super, video, and audio training are follow-up work and are not enabled by these recipes. + +### Checkpoint compatibility + +Use the `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16` Hugging Face checkpoint or a checkpoint converted with the dedicated `NemotronOmniModel` integration. Legacy Megatron checkpoints whose parameter names use an `llava_model` prefix are not compatible with this model definition. Reconvert those checkpoints from the original Hugging Face checkpoint instead of loading them directly. + +### Maintained recipes + +| Workload | Recipe | Topology | +|---|---|---| +| CLEVR-CoGenT | [`vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml`](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml) | 1 node, 8 GPUs, TP=8, EP=8 | +| MMPR-Tiny | [`vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml`](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml) | 4 nodes, 8 GPUs per node, TP=8, EP=16, vLLM TP=2 | + +Launch the single-node Megatron recipe from inside the container on an 8-GPU node: + +```bash +uv run examples/run_vlm_grpo.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml +``` + +For a four-node Slurm run, use the `ray.sub` example above with the following configuration path and omit the AutoModel-specific `PYTHONPATH` addition: + +```bash +CONFIG_PATH=examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml +``` + +The recipes keep sequence packing enabled because the model owns the packing step after multimodal embedding insertion. They also request raw generation log probabilities so that vLLM and the Megatron policy compare the same pre-processor probability values when generation constraints such as `bad_words` are active. The generation context cap prevents the processor-expanded image prompt plus generated response from exceeding the configured 8192-token context length. diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml index 2fd9486bd9b..90e16000aa8 100644 --- a/examples/configs/distillation_math.yaml +++ b/examples/configs/distillation_math.yaml @@ -207,6 +207,7 @@ policy: &POLICY_BASE async_engine: false precision: ${...precision} kv_cache_dtype: "auto" + logprobs_mode: processed_logprobs tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 # When EP > 1, EP must be a multiple of TP since vLLM's EP = DP * TP diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 3abafb713b8..c4682a2c976 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -382,6 +382,7 @@ policy: async_engine: false precision: ${policy.precision} kv_cache_dtype: "auto" + logprobs_mode: processed_logprobs tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 # When EP > 1, EP must be a multiple of TP since vLLM's EP = DP * TP diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index c1e79ad6655..1d79963534b 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -249,6 +249,7 @@ policy: async_engine: false precision: ${policy.precision} kv_cache_dtype: "auto" + logprobs_mode: processed_logprobs tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml index 4844cbf96ea..c2ed8cbdb39 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml @@ -46,6 +46,8 @@ policy: tensor_parallel_size: 8 enforce_eager: true max_model_len: 8192 + cap_max_tokens_to_context: true + logprobs_mode: raw_logprobs gpu_memory_utilization: 0.5 enable_prefix_caching: false vllm_kwargs: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml new file mode 100644 index 00000000000..b738cf17a13 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml @@ -0,0 +1,55 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml +loss_fn: + reference_policy_kl_penalty: 0.0 +checkpointing: + checkpoint_dir: results/vlm_grpo_nemotron_omni_megatron +policy: + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + train_global_batch_size: 8 + logprob_batch_size: 1 + max_total_sequence_length: 8192 + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: '9.0' + tensor_model_parallel_size: 8 + expert_model_parallel_size: 8 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + generation: + max_new_tokens: 4096 + bad_words: + - + - + - + - + - + - + vllm_cfg: + tensor_parallel_size: 8 + enforce_eager: true + max_model_len: 8192 + cap_max_tokens_to_context: true + gpu_memory_utilization: 0.5 + enable_prefix_caching: false + logprobs_mode: raw_logprobs + vllm_kwargs: + limit_mm_per_prompt: + image: 2 + max_num_batched_tokens: 16384 + mamba_ssm_cache_dtype: float32 + skip_mm_profiling: true + kernel_config: + enable_flashinfer_autotune: false + moe_backend: triton +data: + default: + prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt +logger: + wandb: + project: grpo-vlm + name: nemotron-omni-megatron +cluster: + gpus_per_node: 8 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml index e3aef97a410..3a0f328a36a 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml @@ -78,6 +78,8 @@ policy: vllm_cfg: tensor_parallel_size: 8 enforce_eager: true + cap_max_tokens_to_context: true + logprobs_mode: raw_logprobs gpu_memory_utilization: 0.5 enable_prefix_caching: false vllm_kwargs: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml new file mode 100644 index 00000000000..844917c1fc2 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml @@ -0,0 +1,65 @@ +defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml +grpo: + num_prompts_per_step: 512 + overlong_filtering: true + zero_variance_prompt_filtering: false + deduplicate_multimodal_data: false +loss_fn: + ratio_clip_max: 0.28 + use_on_policy_kl_approximation: true + sequence_level_importance_ratios: true + token_level_loss: false +checkpointing: + checkpoint_dir: results/vlm_grpo_nemotron_omni_mmpr_megatron + keep_top_k: 4 + checkpoint_must_save_by: 00:03:45:00 +policy: + train_global_batch_size: 2048 + logprob_chunk_size: 1024 + megatron_cfg: + empty_unused_memory_level: 2 + expert_model_parallel_size: 16 + optimizer: + lr: 3.0e-06 + min_lr: 2.0e-09 + weight_decay: 0.0 + adam_beta2: 0.99 + scheduler: + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3.0e-08 + generation: + max_new_tokens: ${policy.max_total_sequence_length} + vllm_cfg: + tensor_parallel_size: 2 + load_format: auto + enforce_eager: false + max_model_len: ${policy.max_total_sequence_length} + gpu_memory_utilization: 0.75 + vllm_kwargs: + max_num_batched_tokens: 32768 + max_num_seqs: 512 +data: + train: + dataset_name: mmpr-tiny + download_dir: results/mmpr_tiny_processed + split_validation_size: 0.008 + seed: 42 + default: + prompt_file: null + env_name: mmpr-tiny +env: + mmpr-tiny: + num_workers: 8 + reward_functions: + - name: geo3k + weight: 1.0 + kwargs: + format_score: 0.1 +logger: + wandb_enabled: true + wandb: + project: nemotron-omni-main-migration + name: nemotron-omni-mmpr-megatron +cluster: + num_nodes: 4 diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 160d6a0923c..470618d1e74 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -22,6 +22,7 @@ import requests import torch +import torch.nn.functional as F from PIL import Image from transformers import PreTrainedTokenizerBase from transformers.audio_utils import load_audio @@ -81,7 +82,18 @@ def __init__( self, tensors: Union[torch.Tensor, list[Optional[torch.Tensor]], list[None]], dim_to_pack: int, + *, + pad_to_max_shape: bool = False, ) -> None: + """Wrap per-item tensors for concatenation along ``dim_to_pack``. + + Args: + 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. + """ assert tensors is not None, "Input tensors to PackedTensor cannot be None" if isinstance(tensors, torch.Tensor): @@ -96,6 +108,7 @@ 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 def as_tensor( self, device: Optional[torch.device] = None @@ -108,8 +121,51 @@ def as_tensor( non_none_tensors = [t for t in self.tensors if t is not None] if len(non_none_tensors) == 0: return None - else: - return torch.cat(non_none_tensors, dim=self.dim_to_pack).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: + ranks = {tensor.ndim for tensor in non_none_tensors} + if len(ranks) != 1: + raise ValueError( + "pad_to_max_shape requires tensors with the same rank, " + f"but received ranks {sorted(ranks)}" + ) + + rank = ranks.pop() + pack_dim = ( + self.dim_to_pack if self.dim_to_pack >= 0 else rank + self.dim_to_pack + ) + if not 0 <= pack_dim < rank: + raise IndexError( + f"dim_to_pack={self.dim_to_pack} is invalid for tensors with rank {rank}" + ) + max_shape = [ + max(tensor.shape[dim] for tensor in non_none_tensors) + for dim in range(rank) + ] + + def pad_to_batch_shape(tensor: torch.Tensor) -> torch.Tensor: + padding = [] + for dim in reversed(range(rank)): + padding.extend( + ( + 0, + 0 + if dim == pack_dim + else max_shape[dim] - tensor.shape[dim], + ) + ) + return F.pad(tensor, padding) + + non_none_tensors = [ + pad_to_batch_shape(tensor) for tensor in non_none_tensors + ] + + return torch.cat(non_none_tensors, dim=self.dim_to_pack).to(device) def __len__(self) -> int: # this is the number of tensors in this data wrapper @@ -124,12 +180,20 @@ def to(self, device: str | torch.device) -> "PackedTensor": def slice(self, indices: Union[list[int], torch.Tensor]) -> "PackedTensor": idx = indices.tolist() if isinstance(indices, torch.Tensor) else indices tensors = [self.tensors[i] for i in idx] - return PackedTensor(tensors, self.dim_to_pack) + return PackedTensor( + tensors, + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + ) @classmethod def empty_like(cls, other: "PackedTensor") -> "PackedTensor": """Return a new PackedTensor with same length and dim_to_pack as `other`, with all entries None.""" - return cls([None] * len(other.tensors), other.dim_to_pack) + return cls( + [None] * len(other.tensors), + other.dim_to_pack, + pad_to_max_shape=other.pad_to_max_shape, + ) @classmethod def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor": @@ -157,12 +221,20 @@ 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" + ) # concatenate the tensors tensors = [] for packed_tensor in from_packed_tensors: tensors.extend(packed_tensor.tensors) dim_to_pack = dim_to_packs[0] - return cls(tensors, dim_to_pack) + return cls( + tensors, + dim_to_pack, + pad_to_max_shape=pad_to_max_shapes[0], + ) @classmethod def flattened_concat( @@ -194,8 +266,16 @@ 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" + ) tensors = [p.as_tensor() for p in from_packed_tensors] - return cls(tensors, from_packed_tensors[0].dim_to_pack) + return cls( + tensors, + from_packed_tensors[0].dim_to_pack, + pad_to_max_shape=pad_to_max_shapes[0], + ) def get_multimodal_keys_from_processor(processor) -> list[str]: diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 39aae573ae8..a6a072b0357 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -614,15 +614,39 @@ def vlm_hf_data_processor( user_message["token_ids"] = message["input_ids"][0] # add all keys and values to the user message, and the list of keys multimodal_keys = list(get_multimodal_keys_from_processor(processor)) - # imgs_sizes is not declared in model_input_names by the NemotronOmni - # checkpoint's bundled image_processor, so append it explicitly when - # present. It packs along dim=0 (per-image). + # Current Nemotron Omni processors emit imgs_sizes. Historical MMPR + # checkpoints instead emit a batch of fixed-size image tiles and only + # declare pixel_values. Treat each tile as one dynamic-resolution image so + # the Nemotron Omni path can patchify it and preserve the processor's exact + # placeholder count. + if ( + _uses_image_placeholder + and "pixel_values" in message + and "imgs_sizes" not in message + and message["pixel_values"].ndim == 4 + ): + pixel_values = message["pixel_values"] + num_tiles, _, height, width = pixel_values.shape + message["imgs_sizes"] = torch.tensor( + [[height, width]] * num_tiles, dtype=torch.long + ) + + # imgs_sizes is not always declared in model_input_names by bundled image + # processors, so append it explicitly when present. RADIO uses temporal + # patching even for still images and requires one num_frames=1 entry per + # image/tile. if "imgs_sizes" in message and "imgs_sizes" not in multimodal_keys: multimodal_keys.append("imgs_sizes") + if "imgs_sizes" in message and "num_frames" not in message: + message["num_frames"] = torch.ones(len(message["imgs_sizes"]), dtype=torch.long) + if "num_frames" in message and "num_frames" not in multimodal_keys: + multimodal_keys.append("num_frames") for key in multimodal_keys: if key in message: user_message[key] = PackedTensor( - message[key], dim_to_pack=get_dim_to_pack_along(processor, key) + message[key], + dim_to_pack=get_dim_to_pack_along(processor, key), + pad_to_max_shape=_uses_image_placeholder and key == "pixel_values", ) # specifically for gemma, we need to add token_type_ids to the user message as a sequence-type value diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index d663ac1f611..1125245c98a 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -134,6 +134,12 @@ def setup_data_plane(self, cfg: DataPlaneConfig) -> None: Called once by the driver after worker construction. Idempotent. """ + if getattr(self, "model_slices_context_parallel_inputs", False): + raise NotImplementedError( + "TransferQueue/SingleController does not yet support models that " + "insert media before context-parallel input selection. Use the " + "synchronous NeMo-RL policy path for Nemotron Omni." + ) if self._dp_client is not None: return from nemo_rl.data_plane import build_data_plane_client diff --git a/nemo_rl/models/automodel/data.py b/nemo_rl/models/automodel/data.py index 29e798a3861..973fc7fff90 100644 --- a/nemo_rl/models/automodel/data.py +++ b/nemo_rl/models/automodel/data.py @@ -14,11 +14,14 @@ """Data processing utilities for automodel training and inference.""" +import inspect import itertools from dataclasses import dataclass, field +from functools import cache from typing import Any, Iterable, Iterator, Optional, Tuple import torch +from torch import nn from transformers import AutoTokenizer from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType @@ -29,6 +32,60 @@ ) +@cache +def _accepted_forward_kwargs( + model_type: type[nn.Module], +) -> Optional[frozenset[str]]: + """Return explicit ``forward`` kwargs, or ``None`` when all kwargs are accepted.""" + try: + parameters = inspect.signature(model_type.forward).parameters.values() + except (AttributeError, TypeError, ValueError): + return None + + if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters): + return None + return frozenset( + parameter.name for parameter in parameters if parameter.name != "self" + ) + + +def _all_image_sizes_equal(imgs_sizes: torch.Tensor) -> bool: + """Return True if every image in the batch has the same (height, width).""" + if imgs_sizes.ndim != 2 or imgs_sizes.shape[0] <= 1: + return True + return bool((imgs_sizes == imgs_sizes[0]).all()) + + +def filter_multimodal_kwargs_for_model( + model: nn.Module, multimodal_kwargs: dict[str, Any] +) -> dict[str, Any]: + """Drop processor metadata that is not accepted by an AutoModel forward.""" + accepted_kwargs = _accepted_forward_kwargs(type(model)) + if accepted_kwargs is None: + return multimodal_kwargs + # A forward that cannot consume imgs_sizes also cannot crop the per-image + # pad_to_max_shape padding, so mixed-resolution batches would feed padded + # pixels to the vision encoder and mismatch the placeholder count. This is + # the AutoModel Nemotron Omni path (nvidia/Nemotron-3-Nano-Omni-30B-A3B- + # Reasoning-BF16), whose HF forward takes pixel_values but not imgs_sizes, + # unlike the mcore NemotronOmniModel which crops via imgs_sizes. + imgs_sizes = multimodal_kwargs.get("imgs_sizes") + if ( + imgs_sizes is not None + and "imgs_sizes" not in accepted_kwargs + and not _all_image_sizes_equal(imgs_sizes) + ): + raise ValueError( + "This AutoModel does not accept `imgs_sizes` and cannot crop padded " + "pixel_values, but the batch contains mixed-resolution images. The " + "AutoModel backend only supports equal-resolution images/tiles; use " + "the Megatron backend for mixed-resolution inputs." + ) + return { + key: value for key, value in multimodal_kwargs.items() if key in accepted_kwargs + } + + @dataclass class ProcessedInputs: """Processed microbatch inputs ready for model forward pass. diff --git a/nemo_rl/models/automodel/train.py b/nemo_rl/models/automodel/train.py index d23f7f670a4..d6a3d1b3047 100644 --- a/nemo_rl/models/automodel/train.py +++ b/nemo_rl/models/automodel/train.py @@ -51,7 +51,11 @@ distributed_vocab_topk, get_logprobs_from_vocab_parallel_logits, ) -from nemo_rl.models.automodel.data import ProcessedInputs, ProcessedMicrobatch +from nemo_rl.models.automodel.data import ( + ProcessedInputs, + ProcessedMicrobatch, + filter_multimodal_kwargs_for_model, +) from nemo_rl.models.policy import PolicyConfig # Union type for any post-processing function @@ -118,7 +122,9 @@ def model_forward( # Add VLM kwargs if applicable if processed_inputs.is_multimodal: - model_args.update(processed_inputs.vlm_kwargs) + model_args.update( + filter_multimodal_kwargs_for_model(model, processed_inputs.vlm_kwargs) + ) # flash_attn_kwargs is not supported for multimodal if "flash_attn_kwargs" in model_args: del model_args["flash_attn_kwargs"] diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index ec7b3dd54ba..a757ad3a8ff 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -204,6 +204,7 @@ class GenerationConfig(TypedDict): model_name: NotRequired[str] # Not Required b/c GRPO writes this stop_token_ids: list[int] | None stop_strings: list[str] | None + bad_words: NotRequired[list[str] | None] colocated: NotRequired[ColocationConfig] port_range_low: NotRequired[int] port_range_high: NotRequired[int] diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index dc539dc5a5a..6281bb2af18 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -34,6 +34,15 @@ class VllmSpecificArgs(TypedDict): async_engine: bool load_format: NotRequired[str] precision: NotRequired[str] + # Whether vLLM returns logprobs before or after generation-time logit + # processors. RL policy recomputation uses raw model logits, so recipes + # with generation-time processors should request ``raw_logprobs`` when + # comparing generation and policy logprobs. + logprobs_mode: NotRequired[Literal["processed_logprobs", "raw_logprobs"]] + # Cap each request's generated tokens so the training prompt plus response + # fits within max_model_len. This is needed when multimodal processing makes + # the training prompt longer than its text-only representation. + cap_max_tokens_to_context: NotRequired[bool] # Use ModelOpt MXFP8 quantization when precision is fp8. is_mx: NotRequired[bool] kv_cache_dtype: Literal["auto", "fp8", "fp8_e4m3"] diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py index fb915a55eec..d5b16097ed9 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -491,6 +491,81 @@ def _patch_vllm_shm_broadcast_bind_retry(logger) -> None: ) +def _patch_vllm_radio_layerscale_loader(logger) -> None: + """Load explicit RADIO LayerScale weights and initialize folded weights. + + vLLM 0.25.1 uses ``ls1`` and ``ls2`` in ``RadioVisionEncoderLayer`` but + skips them in ``RadioModel.load_weights``. Explicit checkpoint values are + therefore ignored, while folded checkpoints leave the parameters at dummy + initialization. Patch the loader so explicit values are loaded and absent + values are initialized to RADIO's configured identity factor. + """ + try: + file_to_patch = _get_vllm_file("model_executor/models/radio.py") + except RuntimeError: + logger.warning("Could not locate radio.py for the LayerScale loader patch.") + return + + old_snippet = """ elif sub.startswith("model.blocks."): + # Encoder blocks: HF 'model.blocks.{i}.' -> + # vLLM 'model.encoder.layers.{i}.' + parts = sub.split(".") + if len(parts) >= 4: + layer_idx = parts[2] + suffix = ".".join(parts[3:]) + # Skip layer-scale entries that vLLM doesn't use + if suffix in {"ls1", "ls2"} or suffix.startswith(("ls1.", "ls2.")): + continue + vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}" + + if vllm_key and vllm_key in params_dict: + param = params_dict[vllm_key] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, weight) + loaded_params.add(vllm_key) + + return loaded_params +""" + new_snippet = """ elif sub.startswith("model.blocks."): + # Encoder blocks: HF 'model.blocks.{i}.' -> + # vLLM 'model.encoder.layers.{i}.' + parts = sub.split(".") + if len(parts) >= 4: + layer_idx = parts[2] + suffix = ".".join(parts[3:]) + vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}" + + if vllm_key and vllm_key in params_dict: + param = params_dict[vllm_key] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, weight) + loaded_params.add(vllm_key) + + initializer_factor = self.config.initializer_factor + for name, param in params_dict.items(): + if name.endswith((".ls1", ".ls2")) and name not in loaded_params: + param.data.fill_(initializer_factor) + loaded_params.add(name) + + return loaded_params +""" + + with _locked_file_patch(file_to_patch) as (content, write_back): + if new_snippet in content: + logger.info("vLLM RADIO LayerScale loader patch already applied.") + return + if old_snippet not in content: + logger.warning( + "Could not apply vLLM RADIO LayerScale loader patch: expected " + "vLLM 0.25.1 source shape was not found in %s.", + file_to_patch, + ) + return + write_back(content.replace(old_snippet, new_snippet, 1)) + + logger.info("Successfully patched vLLM RADIO LayerScale loading.") + + def ensure_vllm_source_compat() -> None: """Apply interpreter-independent vLLM source-compat patches. @@ -502,7 +577,9 @@ def ensure_vllm_source_compat() -> None: """ from vllm.logger import init_logger - _patch_vllm_tool_parser_namespace_tool(init_logger("vllm_patch")) + patch_logger = init_logger("vllm_patch") + _patch_vllm_tool_parser_namespace_tool(patch_logger) + _patch_vllm_radio_layerscale_loader(patch_logger) def _apply_vllm_patches( @@ -556,3 +633,4 @@ def _apply_vllm_patches( _patch_vllm_tool_parser_namespace_tool(patch_logger) _patch_vllm_ray_executor_v2_tcpstore_port(patch_logger) _patch_vllm_shm_broadcast_bind_retry(patch_logger) + _patch_vllm_radio_layerscale_loader(patch_logger) diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index 47c37d0ce84..ed3f045274e 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -64,6 +64,19 @@ logger = logging.getLogger(__name__) +def _context_capped_max_new_tokens( + *, configured_max_new_tokens: int, input_length: int, max_model_len: int +) -> int: + """Cap generation so the training prompt and response fit the context.""" + remaining_context = max_model_len - input_length + if remaining_context <= 0: + raise ValueError( + "Cannot generate from an input whose training length exhausts the " + f"model context: input_length={input_length}, max_model_len={max_model_len}." + ) + return min(configured_max_new_tokens, remaining_context) + + def _resolve_enable_prefix_caching(vllm_cfg: dict[str, Any]) -> bool: enable_prefix_caching = vllm_cfg.get("enable_prefix_caching", None) if enable_prefix_caching is None: @@ -576,10 +589,13 @@ def _load_model(self, bundle_indices, seed): enable_sleep_mode=True, # Set disable_log_stats=False so that self.llm.get_metrics() works. disable_log_stats=False, - logprobs_mode="processed_logprobs", **vllm_kwargs, ) + logprobs_mode = self.cfg["vllm_cfg"].get("logprobs_mode") + if logprobs_mode is not None: + llm_kwargs["logprobs_mode"] = logprobs_mode + self._create_engine(llm_kwargs) log_gpu_memory_diagnostics( label="after_engine_create", worker_type="VllmGenerationWorker", device_id=0 @@ -637,6 +653,7 @@ def _build_sampling_params( stop_token_ids=self.cfg["stop_token_ids"], stop=stop_strings, include_stop_str_in_output=True, + bad_words=self.cfg.get("bad_words"), ignore_eos=self.cfg.get("ignore_eos", False), ) @@ -668,6 +685,33 @@ def _spec_decode_max_tokens( 1, min(base_max_tokens, max_model_len - input_len - (spec_lookahead + 1)) ) + @classmethod + def _request_max_new_tokens( + cls, + *, + configured_max_new_tokens: int, + input_length: int, + max_model_len: int, + cap_to_context: bool, + spec_lookahead: int, + ) -> int: + """Apply context and speculative-decoding limits to one request.""" + max_new_tokens = configured_max_new_tokens + if cap_to_context: + max_new_tokens = _context_capped_max_new_tokens( + configured_max_new_tokens=max_new_tokens, + input_length=input_length, + max_model_len=max_model_len, + ) + if spec_lookahead > 0: + max_new_tokens = cls._spec_decode_max_tokens( + max_new_tokens, + input_length, + max_model_len, + spec_lookahead, + ) + return max_new_tokens + @staticmethod def _patch_vllm_nsight_config() -> None: """Override vLLM's nsight config for internal TP workers to use deferred capture. @@ -826,44 +870,54 @@ def generate( input_lengths = data["input_lengths"] batch_stop_strings: list[list[str]] = data.get("stop_strings", []) stop_strings = self._merge_stop_strings(batch_stop_strings) - sampling_params = self._build_sampling_params( - greedy=greedy, - stop_strings=stop_strings, - ) - # vLLM 0.20 eagle3 spec decode hits a CUDA illegal memory access when a + # vLLM Eagle3 spec decode hits a CUDA illegal memory access when a # request's total length reaches max_model_len (the drafter looks ahead # past the boundary). Clamp per-request max_tokens so speculative # requests stop short of the boundary by the drafter lookahead. spec_cfg = self.cfg.get("vllm_kwargs", {}).get("speculative_config") or {} spec_lookahead = int(spec_cfg.get("num_speculative_tokens", 0)) - if spec_lookahead > 0: - max_model_len = self.cfg["vllm_cfg"]["max_model_len"] - base_max_tokens = sampling_params.max_tokens - sampling_params = [ - self._build_sampling_params( - greedy=greedy, - stop_strings=stop_strings, - max_new_tokens=self._spec_decode_max_tokens( - base_max_tokens, int(input_len), max_model_len, spec_lookahead - ), - ) - for input_len in data["input_lengths"].tolist() - ] - # verify inputs have correct padding verify_right_padding(data, pad_value=self.cfg["_pad_token_id"]) # Original input length with padding padded_input_length = input_ids.size(1) - # Convert inputs to vLLM format - prompts = format_prompt_for_vllm_generation(data) - - # Generate outputs assert self.llm is not None, ( "Attempting to generate with either an uninitialized vLLM or non-model-owner" ) + cap_to_context = bool(self.cfg["vllm_cfg"].get("cap_max_tokens_to_context")) + if cap_to_context or spec_lookahead > 0: + max_model_len = int(self.cfg["vllm_cfg"]["max_model_len"]) + configured_max_new_tokens = int(self.cfg["max_new_tokens"]) + per_request_max_new_tokens = [] + for input_length in input_lengths.tolist(): + per_request_max_new_tokens.append( + self._request_max_new_tokens( + configured_max_new_tokens=configured_max_new_tokens, + input_length=int(input_length), + max_model_len=max_model_len, + cap_to_context=cap_to_context, + spec_lookahead=spec_lookahead, + ) + ) + + sampling_params = [ + self._build_sampling_params( + greedy=greedy, + stop_strings=stop_strings, + max_new_tokens=max_new_tokens, + ) + for max_new_tokens in per_request_max_new_tokens + ] + else: + sampling_params = self._build_sampling_params( + greedy=greedy, + stop_strings=stop_strings, + ) + + # Convert inputs to vLLM format and generate outputs. + prompts = format_prompt_for_vllm_generation(data) use_tqdm = self.cfg["vllm_cfg"].get("use_tqdm", True) outputs = self.llm.generate(prompts, sampling_params, use_tqdm=use_tqdm) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 3983a1029e7..201d0790044 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -954,11 +954,21 @@ async def process_single_sample(sample_idx): [per_sample_stop_strings] if per_sample_stop_strings else None ) - remaining_ctx = ( - self.cfg["vllm_cfg"]["max_model_len"] - current_input_actual_length - ) + max_model_len = int(self.cfg["vllm_cfg"]["max_model_len"]) + remaining_ctx = max_model_len - current_input_actual_length allowed_new_tokens = max(0, min(self.cfg["max_new_tokens"], remaining_ctx)) + spec_cfg = self.cfg.get("vllm_kwargs", {}).get("speculative_config") or {} + spec_lookahead = int(spec_cfg.get("num_speculative_tokens", 0)) + if allowed_new_tokens > 0 and spec_lookahead > 0: + allowed_new_tokens = self._request_max_new_tokens( + configured_max_new_tokens=allowed_new_tokens, + input_length=current_input_actual_length, + max_model_len=max_model_len, + cap_to_context=False, + spec_lookahead=spec_lookahead, + ) + # Handle case where no tokens can be generated due to length constraints if allowed_new_tokens == 0: # Access the input data directly from the function parameters diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index bb216c8e29c..f7b0230a9da 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -17,6 +17,9 @@ from typing import Any, Iterator, Optional, Tuple import torch +from megatron.bridge.training.utils.packed_seq_utils import ( + get_packed_seq_cp_partition_indices, +) from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( get_context_parallel_rank, @@ -59,7 +62,8 @@ class ProcessedMicrobatch: Attributes: data_dict: The original BatchedDataDict containing raw batch data input_ids: Processed input token IDs (may be packed for sequence packing) - input_ids_cp_sharded: Context-parallel sharded input token IDs + input_ids_cp_sharded: Model-forward token IDs. Usually CP-sharded; models + that insert media before CP selection receive the full packed THD row. attention_mask: Attention mask tensor (None for packed sequences) position_ids: Position IDs tensor (None for packed sequences) packed_seq_params: PackedSeqParams for sequence packing (None if not packing) @@ -91,6 +95,8 @@ def make_processed_microbatch_iterator( straggler_timer: StragglerDetector, pad_full_seq_to: Optional[int], delegate_pack_to_model: bool = False, + delegate_mtp_loss_mask_to_model: bool = False, + model_slices_context_parallel_inputs: bool = False, ) -> Iterator[ProcessedMicrobatch]: """Wrap a raw microbatch iterator to yield processed microbatches. @@ -124,6 +130,8 @@ def make_processed_microbatch_iterator( pad_full_seq_to=pad_full_seq_to, pack_sequences=pack_sequences, delegate_pack_to_model=delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=delegate_mtp_loss_mask_to_model, + model_slices_context_parallel_inputs=model_slices_context_parallel_inputs, straggler_timer=straggler_timer, ) @@ -148,6 +156,8 @@ def get_microbatch_iterator( straggler_timer: StragglerDetector, seq_length_key: Optional[str] = None, delegate_pack_to_model: bool = False, + delegate_mtp_loss_mask_to_model: bool = False, + model_slices_context_parallel_inputs: bool = False, ) -> Tuple[Iterator[ProcessedMicrobatch], int, int, int, int]: """Create a processed microbatch iterator from a batch of data. @@ -212,6 +222,8 @@ def get_microbatch_iterator( pad_full_seq_to=pad_full_seq_to, straggler_timer=straggler_timer, delegate_pack_to_model=delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=delegate_mtp_loss_mask_to_model, + model_slices_context_parallel_inputs=model_slices_context_parallel_inputs, ) # Compute padded sequence length for pipeline parallelism @@ -246,6 +258,8 @@ def process_microbatch( pad_full_seq_to: Optional[int] = None, pack_sequences: bool = False, delegate_pack_to_model: bool = False, + delegate_mtp_loss_mask_to_model: bool = False, + model_slices_context_parallel_inputs: bool = False, straggler_timer: Optional[StragglerDetector] = None, ) -> ProcessedInputs: """Process a microbatch for Megatron model forward pass.""" @@ -286,11 +300,10 @@ def process_microbatch( seq_lengths = data_dict[seq_length_key] if delegate_pack_to_model: - # The VLM packing path does not pack or propagate mtp_loss_mask, - # so MTP training would be silently dropped here. Fail loudly - # instead of producing wrong results. - assert "mtp_loss_mask" not in data_dict, ( - "MTP training is not supported with VLM sequence packing" + 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 " + "model_owns_mtp_loss_mask_packing" ) # VLM path: model (e.g. mbridge Qwen3VL) does its own # preprocess_packed_seqs; NeMo-RL must NOT pre-pack + CP-shard, @@ -320,8 +333,35 @@ def process_microbatch( pad_individual_seqs_to_multiple_of, pad_full_seq_to=pad_full_seq_to, ) + if has_mtp_loss_mask: + source_mtp_loss_mask = data_dict["mtp_loss_mask"] + assert source_mtp_loss_mask.ndim == 2 + assert ( + source_mtp_loss_mask.shape[0] == input_ids_cp_sharded.shape[0] + ) + mtp_loss_mask = source_mtp_loss_mask.new_zeros( + input_ids_cp_sharded.shape + ) + copied_length = min( + source_mtp_loss_mask.shape[1], + input_ids_cp_sharded.shape[1], + ) + mtp_loss_mask[:, :copied_length] = source_mtp_loss_mask[ + :, :copied_length + ] + mtp_loss_mask = mtp_loss_mask * attention_mask.to( + dtype=mtp_loss_mask.dtype + ) position_ids = None else: + if ( + model_slices_context_parallel_inputs + and "mtp_loss_mask" in data_dict + ): + raise NotImplementedError( + "Nemotron Omni caller-packed THD inputs do not yet support MTP. " + "Disable MTP for the Nano image/text path." + ) token_identity = None if routed_experts is not None and r3_trace_verify_forward_enabled(): token_identity = _make_r3_trace_token_identity( @@ -331,7 +371,7 @@ def process_microbatch( # Pack sequences on main's per-sequence zigzag CP layout. ( input_ids, - input_ids_cp_sharded, + local_input_ids, packed_seq_params, cu_seqlens, cu_seqlens_padded, @@ -344,6 +384,35 @@ def process_microbatch( cp_rank=get_context_parallel_rank(), cp_size=get_context_parallel_world_size(), ) + if model_slices_context_parallel_inputs: + packed_seq_params = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=int( + (cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]) + .max() + .item() + ), + max_seqlen_kv=int( + (cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]) + .max() + .item() + ), + # TE's default inference excludes the final boundary, so + # it misses trailing-only padding for a single sequence. + # CP zigzag can move that padding to a rank-local seam. + pad_between_seqs=not torch.equal(cu_seqlens, cu_seqlens_padded), + qkv_format="thd", + total_tokens=input_ids.shape[1], + ) + # This field is the model-forward input. For this capability + # the model needs the full THD row so it can insert media + # before selecting its CP-owned embeddings. + input_ids_cp_sharded = input_ids + else: + input_ids_cp_sharded = local_input_ids # routed_experts and the R3 trace token identity ride the SAME # per-seq zigzag CP sharding as input_ids, re-derived from # cu_seqlens_padded. @@ -362,6 +431,23 @@ def process_microbatch( get_context_parallel_rank(), get_context_parallel_world_size(), ) + if model_slices_context_parallel_inputs: + cp_partition_indices = get_packed_seq_cp_partition_indices( + packed_seq_params, + total_tokens=input_ids.shape[1], + cp_size=get_context_parallel_world_size(), + cp_rank=get_context_parallel_rank(), + device=input_ids.device, + ) + routed_experts_cp_sharded = routed_experts.index_select( + 1, cp_partition_indices + ).contiguous() + if _token_identity_packed is not None: + token_identity_cp_sharded = ( + _token_identity_packed.index_select( + 1, cp_partition_indices + ).contiguous() + ) if ( routed_experts_cp_sharded is not None and routed_experts_cp_sharded.dim() != 4 @@ -374,14 +460,22 @@ def process_microbatch( verified_token_count = _verify_r3_trace_cp_token_alignment( source_input_ids=data_dict["input_ids"], source_routed_experts=data_dict.get("routed_experts"), - input_ids_cp_sharded=input_ids_cp_sharded, + input_ids_cp_sharded=( + local_input_ids + if model_slices_context_parallel_inputs + else input_ids_cp_sharded + ), routed_experts_cp_sharded=routed_experts_cp_sharded, token_identity_cp_sharded=token_identity_cp_sharded, ) trace_cp_routed_experts( routed_experts_cp_sharded=routed_experts_cp_sharded, token_identity_cp_sharded=token_identity_cp_sharded, - input_ids_cp_sharded=input_ids_cp_sharded, + input_ids_cp_sharded=( + local_input_ids + if model_slices_context_parallel_inputs + else input_ids_cp_sharded + ), cp_token_identity_verified_count=verified_token_count, cp_rank=get_context_parallel_rank(), cp_size=get_context_parallel_world_size(), diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index d5aaae2d6fc..4b888cd9ea5 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -68,6 +68,8 @@ _HF_CONFIG_PATCHED = False +_NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT = "expanded_sequence_v1" + def _patch_hf_config_double_instantiation(): """Patch HF config classes whose __post_init__ fails with Megatron's recursive instantiation. @@ -920,9 +922,26 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.use_fused_weighted_squared_relu = config["megatron_cfg"][ "use_fused_weighted_squared_relu" ] - # Optional explicit attention backend override for environments where - # TE auto backend probing is unstable. + # NeMo-RL can pack multiple expanded Omni examples into one THD tensor. + # Flash attention does not support the resulting padded multi-row layout, + # so the canonical expanded-sequence contract must use backend dispatch. attention_backend = config["megatron_cfg"].get("attention_backend") + if ( + getattr(model_cfg, "nemotron_omni_contract", None) + == _NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT + ): + if attention_backend == "flash": + raise ValueError( + "Nemotron Omni's expanded-sequence contract does not support " + "attention_backend='flash' in NeMo-RL because packed batches can " + "contain multiple padded THD rows. Use attention_backend='auto' " + "or omit the setting." + ) + if attention_backend is None: + attention_backend = "auto" + + # Optional explicit attention backend override for other models, and the + # required auto selection for canonical Nemotron Omni. if attention_backend is not None: for _nvte_var in ("NVTE_FUSED_ATTN", "NVTE_FLASH_ATTN", "NVTE_UNFUSED_ATTN"): os.environ.pop(_nvte_var, None) diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 1b9e3c03982..beb02a89ad0 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -87,7 +87,8 @@ def model_forward( Args: model: The model to run forward pass on data_dict: Dictionary containing batch data - input_ids_cp_sharded: Context-parallel sharded input token IDs + input_ids_cp_sharded: Model-forward token IDs. Usually CP-sharded; models + that insert media before CP selection receive the full packed THD row. position_ids: Position IDs for tokens attention_mask: Attention mask for the sequence packed_seq_params: Parameters for packed sequences (optional) diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index d132d1c111f..e7e80f99f52 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -65,6 +65,7 @@ distributed_vocab_topk, get_logprobs_from_vocab_parallel_logits, ) +from nemo_rl.models.automodel.data import filter_multimodal_kwargs_for_model from nemo_rl.models.dtensor.parallelize import ( _parallelize_model, clip_grad_by_total_norm_, @@ -783,6 +784,9 @@ def train( vlm_kwargs = mb.get_multimodal_dict( as_tensors=True, device=input_ids.device ) + vlm_kwargs = filter_multimodal_kwargs_for_model( + self.model, vlm_kwargs + ) if len(vlm_kwargs) > 0: position_ids = None assert not self.cfg["dtensor_cfg"]["sequence_parallel"], ( @@ -1088,6 +1092,7 @@ def get_logprobs( vlm_kwargs = lp_batch.get_multimodal_dict( as_tensors=True, device=input_ids.device ) + vlm_kwargs = filter_multimodal_kwargs_for_model(self.model, vlm_kwargs) batch_size, seq_len = input_ids.shape if self.enable_seq_packing: @@ -1532,6 +1537,7 @@ def get_topk_logits( vlm_kwargs = lp_batch.get_multimodal_dict( as_tensors=True, device=input_ids.device ) + vlm_kwargs = filter_multimodal_kwargs_for_model(self.model, vlm_kwargs) batch_size, seq_len = input_ids.shape # Store original shapes for unpacking later diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 4052fe03774..52fefad7cdf 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -139,16 +139,65 @@ def _model_self_packs_for_cp(model: Any) -> bool: Such models (mbridge VLM wrappers) call ``preprocess_packed_seqs`` in their forward, so NeMo-RL must hand them an unpacked ``[B, S]`` batch instead of - pre-packing + CP-sharding itself. The only such model today is mbridge's - Qwen3VL, which is also the only mbridge VLM that supports context - parallelism; classic mcore GPTModel and other VLMs do not self-pack. + pre-packing + CP-sharding itself. New wrappers advertise the capability + through ``model_owns_packing``. The Qwen3VL type check remains as a + compatibility fallback until that upstream model exposes the capability. """ from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel from megatron.core.utils import unwrap_model unwrapped = unwrap_model(model) chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] - return any(isinstance(chunk, Qwen3VLModel) for chunk in chunks) + return any( + bool(getattr(chunk, "model_owns_packing", False)) + or isinstance(chunk, Qwen3VLModel) + for chunk in chunks + ) + + +def _model_self_packs_mtp_loss_mask(model: Any) -> bool: + """Whether a self-packing model also aligns and CP-shards MTP masks.""" + from megatron.core.utils import unwrap_model + + unwrapped = unwrap_model(model) + chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] + return any( + bool(getattr(chunk, "model_owns_mtp_loss_mask_packing", False)) + for chunk in chunks + ) + + +def _model_slices_context_parallel_inputs(model: Any) -> bool: + """Whether the model consumes full THD input and slices CP after embedding.""" + from megatron.core.utils import unwrap_model + + unwrapped = unwrap_model(model) + chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] + return any( + bool(getattr(chunk, "model_slices_context_parallel_inputs", False)) + for chunk in chunks + ) + + +def _estimate_refit_tensor_size_in_bytes( + param: torch.Tensor, + *, + export_dtype: torch.dtype, + tp_size: int, + ep_size: int, +) -> int: + """Estimate the gathered tensor size produced by Bridge export. + + Floating-point model weights are exported at the policy dtype. Integral + state (for example BatchNorm ``num_batches_tracked`` buffers) keeps its + original dtype and must not be looked up in a floating-point-only table. + """ + element_size = ( + torch.empty((), dtype=export_dtype).element_size() + if param.is_floating_point() + else param.element_size() + ) + return param.numel() * tp_size * ep_size * element_size @contextmanager @@ -501,6 +550,41 @@ def __init__( # (mbridge VLM wrappers like Qwen3VL). If so, NeMo-RL must hand it an # unpacked [B, S] batch rather than pre-packing + CP-sharding itself. self.delegate_pack_to_model = _model_self_packs_for_cp(self.model) + self.delegate_mtp_loss_mask_to_model = _model_self_packs_mtp_loss_mask( + self.model + ) + assert ( + not self.delegate_mtp_loss_mask_to_model or self.delegate_pack_to_model + ), "A model cannot own MTP-mask packing without owning sequence packing" + self.model_slices_context_parallel_inputs = ( + _model_slices_context_parallel_inputs(self.model) + ) + if self.model_slices_context_parallel_inputs: + if self.delegate_pack_to_model: + raise RuntimeError( + "A model cannot both own sequence packing and consume caller-packed " + "full THD inputs." + ) + model_config = self._get_model_config() + mtp_num_layers = getattr(model_config, "mtp_num_layers", None) + if mtp_num_layers is not None and mtp_num_layers > 0: + raise NotImplementedError( + "Nemotron Omni caller-packed THD inputs do not yet support MTP. " + "Disable MTP for the Nano image/text path." + ) + if self.cfg["megatron_cfg"].get("use_fused_linear_logprobs", False): + raise NotImplementedError( + "Nemotron Omni caller-packed THD inputs do not support " + "use_fused_linear_logprobs=true." + ) + virtual_pipeline_size = self.cfg["megatron_cfg"].get( + "virtual_pipeline_model_parallel_size" + ) + if virtual_pipeline_size not in (None, 1): + raise NotImplementedError( + "Nemotron Omni caller-packed THD inputs do not yet support " + "virtual pipeline parallelism." + ) # vars used for refit ## will be initialized in prepare_refit_info @@ -703,6 +787,8 @@ def train( mbs, straggler_timer=self.mcore_state.straggler_timer, delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, + model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, ) # Track total microbatches for MoE aux-loss averaging total_num_microbatches += int(num_microbatches) @@ -1535,6 +1621,8 @@ def get_logprobs( logprob_batch_size, straggler_timer=self.mcore_state.straggler_timer, delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, + model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, ) use_fused_linear_logprobs = self.cfg["megatron_cfg"].get( @@ -1752,6 +1840,8 @@ def get_topk_logits( logprob_batch_size, straggler_timer=self.mcore_state.straggler_timer, delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, + model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, ) list_of_outputs = megatron_forward_backward( @@ -1974,18 +2064,11 @@ def calculate_size_in_bytes(param, tp_size, ep_size): # need to broadcast for other pp ranks size_in_bytes = None else: - # Calculate size for this parameter - prec_to_bytes = { - torch.bfloat16: 2, - torch.float16: 2, - torch.float32: 4, - torch.float8_e4m3fn: 1, - torch.float8_e5m2: 1, - torch.uint8: 1, - } - scale = prec_to_bytes[self.dtype] / prec_to_bytes[param.dtype] - size_in_bytes = ( - param.element_size() * param.numel() * tp_size * ep_size * scale + size_in_bytes = _estimate_refit_tensor_size_in_bytes( + param, + export_dtype=self.dtype, + tp_size=tp_size, + ep_size=ep_size, ) # Broadcast size_in_bytes across pipeline parallel ranks diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 8acb0503dda..dea2565dbcc 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -49,6 +49,10 @@ tests/test_suites/vlm/vlm_grpo-qwen3-omni-30ba3b-audiomcq-4n8g-megatron.v1.sh tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.sh tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.sh +# Functional Nemotron-Omni 30B-A3B VLM GRPO runs (Megatron) +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.sh + # Functional Qwen3.5-35B VLM GRPO run # The AutoModel variant is re-enabled with the vLLM 0.25.1 bump (no longer hits # https://github.com/vllm-project/vllm/issues/36237). The Megatron variant still diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.sh new file mode 100755 index 00000000000..81c7e5624a9 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.sh @@ -0,0 +1,42 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +GPUS_PER_NODE=8 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +# Run the experiment +cd $PROJECT_ROOT +uv run examples/run_vlm_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/reward"]) > 0.5' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.sh new file mode 100755 index 00000000000..318b6a9649c --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.sh @@ -0,0 +1,42 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=8 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +# Run the experiment +cd $PROJECT_ROOT +uv run examples/run_vlm_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/reward"]) > 0.5' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/unit/algorithms/test_sequence_packing_fusion.py b/tests/unit/algorithms/test_sequence_packing_fusion.py index 59869639808..a979cbff5da 100644 --- a/tests/unit/algorithms/test_sequence_packing_fusion.py +++ b/tests/unit/algorithms/test_sequence_packing_fusion.py @@ -360,6 +360,7 @@ def _run_compare_sequence_packing_wrappers_with_sampling( ) +@pytest.mark.mcore @pytest.mark.parametrize( "cp_tp", [ @@ -391,6 +392,7 @@ def test_sequence_packing_fusion_vs_baseline(distributed_test_runner, cp_tp): distributed_test_runner(test_fn, world_size=world_size) +@pytest.mark.mcore @pytest.mark.parametrize( "cp_tp", [ diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 8f1dce267c7..74822398d7b 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -128,7 +128,7 @@ def test_missing_download_dir_raises_value_error(self): MMPRTinyDataset(download_dir="") -def _make_stub_nemotron_processor(): +def _make_stub_nemotron_processor(*, include_imgs_sizes=True, num_tiles=1): """Build a minimal stub whose class name is NemotronNanoVLV2Processor. The stub implements just enough of the AutoProcessor interface for @@ -164,10 +164,13 @@ def apply_chat_template(self, messages, **kwargs): def __call__(self, text=None, images=None, **kwargs): self.captured_call_text = text - return { + result = { "input_ids": fake_input_ids, - "pixel_values": torch.randn(1, 3, 224, 224), + "pixel_values": torch.randn(num_tiles, 3, 224, 224), } + if include_imgs_sizes: + result["imgs_sizes"] = torch.tensor([[224, 224]] * num_tiles) + return result return NemotronNanoVLV2Processor() @@ -233,6 +236,37 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path): assert "vllm_images" in result assert len(result["vllm_images"]) == 1 assert result["task_name"] == "mmpr-tiny" + user_message = result["message_log"][0] + assert torch.equal(user_message["num_frames"].as_tensor(), torch.tensor([1])) + + def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): + from nemo_rl.data.interfaces import TaskDataSpec + from nemo_rl.data.processors import vlm_hf_data_processor + + task_data_spec = TaskDataSpec(task_name="mmpr-tiny") + task_data_spec.prompt = _TEST_PROMPT_TEMPLATE + processor = _make_stub_nemotron_processor(include_imgs_sizes=False, num_tiles=3) + result = vlm_hf_data_processor( + datum_dict={ + "images": [tiny_image_path], + "question": _RAW_QUESTION, + "answer": "A", + "task_name": "mmpr-tiny", + }, + task_data_spec=task_data_spec, + processor=processor, + max_seq_length=8192, + idx=0, + ) + + user_message = result["message_log"][0] + assert torch.equal( + user_message["imgs_sizes"].as_tensor(), + torch.tensor([[224, 224], [224, 224], [224, 224]]), + ) + assert torch.equal( + user_message["num_frames"].as_tensor(), torch.ones(3, dtype=torch.long) + ) def test_prompted_text_contains_boxed_literal_and_no_raw_dataset_string( self, tiny_image_path diff --git a/tests/unit/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py index a94412222a5..23b7cdacdbd 100644 --- a/tests/unit/data/test_multimodal_dict.py +++ b/tests/unit/data/test_multimodal_dict.py @@ -349,3 +349,134 @@ def test_packedtensor_as_tensor_with_mixed_none_and_tensors(): out = pt.as_tensor() expected = torch.cat([t1, t3], dim=0) assert torch.equal(out, expected) + + +def test_packedtensor_pads_mixed_dynamic_resolution_images(): + """Raw image batches pad spatial dimensions before packing on dim 0.""" + first = torch.ones(1, 3, 2, 4) + second = 2 * torch.ones(1, 3, 4, 2) + + packed = PackedTensor( + [first, second], dim_to_pack=0, pad_to_max_shape=True + ).as_tensor() + + assert packed.shape == (2, 3, 4, 4) + torch.testing.assert_close(packed[0, :, :2, :4], first[0]) + torch.testing.assert_close(packed[0, :, 2:, :], torch.zeros(3, 2, 4)) + torch.testing.assert_close(packed[1, :, :4, :2], second[0]) + torch.testing.assert_close(packed[1, :, :, 2:], torch.zeros(3, 4, 2)) + + +@pytest.mark.mcore +def test_dynamic_resolution_padding_is_cropped_before_radio_patchification(): + """Batch-shape padding must not become RADIO image content.""" + from megatron.bridge.models.nemotron_omni.modeling_nemotron_omni import ( + NemotronOmniModel, + ) + + generator = torch.Generator().manual_seed(2026) + small = torch.randn(1, 3, 32, 32, generator=generator) + large = torch.randn(1, 3, 64, 64, generator=generator) + imgs_sizes = torch.tensor([[32, 32], [64, 64]], dtype=torch.long) + + padded = PackedTensor( + [small, large], + dim_to_pack=0, + pad_to_max_shape=True, + ).as_tensor() + # Use nonzero garbage so this test cannot pass merely because F.pad uses zero. + padded[0, :, 32:, :] = 123 + padded[0, :, :, 32:] = -456 + + class _Patchifier: + patch_dim = 16 + + patchifier = _Patchifier() + packed_patches = NemotronOmniModel._patchify_dynamic_images( + patchifier, + padded, + imgs_sizes, + ) + expected_patches = torch.cat( + [ + NemotronOmniModel._patchify_dynamic_images( + patchifier, + small, + imgs_sizes[:1], + ), + NemotronOmniModel._patchify_dynamic_images( + patchifier, + large, + imgs_sizes[1:], + ), + ], + dim=1, + ) + + torch.testing.assert_close(packed_patches, expected_patches) + + +@pytest.mark.parametrize( + ("first_shape", "second_shape", "expected_shape"), + [ + ((1, 2, 3), (2, 4, 3), (3, 4, 3)), + ((1, 2, 3, 2, 4), (2, 4, 3, 4, 2), (3, 4, 3, 4, 4)), + ], +) +def test_packedtensor_pad_to_max_shape_supports_audio_and_video( + first_shape, second_shape, expected_shape +): + """Padding is generic across non-packing dimensions and tensor ranks.""" + first = torch.ones(first_shape) + second = 2 * torch.ones(second_shape) + + packed = PackedTensor( + [first, second], dim_to_pack=0, pad_to_max_shape=True + ).as_tensor() + + assert packed.shape == expected_shape + slices = (slice(0, first_shape[0]),) + tuple( + slice(0, size) for size in first_shape[1:] + ) + torch.testing.assert_close(packed[slices], first) + + +def test_pad_to_max_shape_rejects_mismatched_ranks(): + with pytest.raises(ValueError, match="same rank"): + PackedTensor( + [torch.ones(1, 3, 4), torch.ones(1, 3)], + dim_to_pack=0, + pad_to_max_shape=True, + ).as_tensor() + + +def test_pad_to_max_shape_rejects_out_of_range_dim(): + with pytest.raises(IndexError, match="dim_to_pack=3 is invalid"): + PackedTensor( + [torch.ones(1, 3, 4), torch.ones(2, 3, 4)], + dim_to_pack=3, + pad_to_max_shape=True, + ).as_tensor() + + +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, + ).as_tensor() + + assert packed.shape == (6, 3, 1) + + +def test_slice_preserves_pad_to_max_shape_flag(): + packed = PackedTensor( + [torch.ones(1, 3, 2, 4), 2 * torch.ones(1, 3, 4, 2)], + dim_to_pack=0, + pad_to_max_shape=True, + ) + + sliced = packed.slice([0, 1]) + + assert sliced.pad_to_max_shape is True + assert sliced.as_tensor().shape == (2, 3, 4, 4) diff --git a/tests/unit/models/automodel/test_automodel_train.py b/tests/unit/models/automodel/test_automodel_train.py index d1f209152ff..4efd0eb59a2 100644 --- a/tests/unit/models/automodel/test_automodel_train.py +++ b/tests/unit/models/automodel/test_automodel_train.py @@ -181,6 +181,157 @@ def test_forward_with_multimodal(self, mock_model, processed_inputs_multimodal): # Flash attention should be removed for multimodal assert "flash_attn_kwargs" not in call_kwargs + def test_forward_filters_unsupported_multimodal_metadata( + self, processed_inputs_multimodal + ): + class ExplicitMultimodalModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.pixel_values = None + + def forward( + self, + input_ids, + attention_mask=None, + position_ids=None, + use_cache=False, + pixel_values=None, + ): + self.pixel_values = pixel_values + return MagicMock(logits=torch.randn(2, 64, 1000)) + + model = ExplicitMultimodalModel() + processed_inputs_multimodal.vlm_kwargs.update( + { + "imgs_sizes": torch.tensor([[224, 224]]), + "num_frames": torch.tensor([1]), + } + ) + + model_forward(model, processed_inputs_multimodal) + + assert ( + model.pixel_values is processed_inputs_multimodal.vlm_kwargs["pixel_values"] + ) + + def test_forward_rejects_mixed_resolution_without_imgs_sizes_support( + self, processed_inputs_multimodal + ): + class ExplicitMultimodalModel(torch.nn.Module): + def forward( + self, + input_ids, + attention_mask=None, + position_ids=None, + use_cache=False, + pixel_values=None, + ): + return MagicMock(logits=torch.randn(2, 64, 1000)) + + model = ExplicitMultimodalModel() + processed_inputs_multimodal.vlm_kwargs.update( + { + "imgs_sizes": torch.tensor([[224, 320], [256, 288]]), + "num_frames": torch.ones(2, dtype=torch.long), + } + ) + + with pytest.raises(ValueError, match="mixed-resolution"): + model_forward(model, processed_inputs_multimodal) + + def test_forward_allows_uniform_resolution_without_imgs_sizes_support( + self, processed_inputs_multimodal + ): + class ExplicitMultimodalModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.pixel_values = None + + def forward( + self, + input_ids, + attention_mask=None, + position_ids=None, + use_cache=False, + pixel_values=None, + ): + self.pixel_values = pixel_values + return MagicMock(logits=torch.randn(2, 64, 1000)) + + model = ExplicitMultimodalModel() + processed_inputs_multimodal.vlm_kwargs.update( + { + "imgs_sizes": torch.tensor([[224, 224], [224, 224]]), + "num_frames": torch.ones(2, dtype=torch.long), + } + ) + + # Uniform sizes (the shipped fixed-tile case): no raise, and imgs_sizes + # is filtered out for a model that cannot consume it. + model_forward(model, processed_inputs_multimodal) + + assert ( + model.pixel_values is processed_inputs_multimodal.vlm_kwargs["pixel_values"] + ) + + def test_forward_preserves_dynamic_resolution_omni_inputs( + self, processed_inputs_multimodal + ): + class DynamicResolutionOmniLikeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.forward_kwargs = {} + + def forward( + self, + input_ids, + attention_mask=None, + position_ids=None, + use_cache=False, + pixel_values=None, + imgs_sizes=None, + ): + self.forward_kwargs = { + "pixel_values": pixel_values, + "imgs_sizes": imgs_sizes, + } + return MagicMock(logits=torch.randn(2, 64, 1000)) + + model = DynamicResolutionOmniLikeModel() + padded_images = torch.randn(2, 3, 256, 320) + image_sizes = torch.tensor([[224, 320], [256, 288]]) + processed_inputs_multimodal.vlm_kwargs.update( + { + "pixel_values": padded_images, + "imgs_sizes": image_sizes, + "num_frames": torch.ones(2, dtype=torch.long), + } + ) + + model_forward(model, processed_inputs_multimodal) + + assert model.forward_kwargs["pixel_values"] is padded_images + assert model.forward_kwargs["imgs_sizes"] is image_sizes + + def test_forward_preserves_multimodal_metadata_for_kwargs_model( + self, processed_inputs_multimodal + ): + class KwargsModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.forward_kwargs = {} + + def forward(self, **kwargs): + self.forward_kwargs = kwargs + return MagicMock(logits=torch.randn(2, 64, 1000)) + + model = KwargsModel() + processed_inputs_multimodal.vlm_kwargs["num_frames"] = torch.tensor([1]) + + model_forward(model, processed_inputs_multimodal) + + assert "num_frames" in model.forward_kwargs + def test_forward_reward_model_removes_flash_attn( self, mock_model, processed_inputs_with_flash ): diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index ba8dd52df4f..0821f0865b6 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -38,6 +38,8 @@ from nemo_rl.models.generation.openai_server_utils import replace_prefix_tokens from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration from nemo_rl.models.generation.vllm.vllm_worker import ( + VllmGenerationWorkerImpl, + _context_capped_max_new_tokens, _resolve_enable_prefix_caching, ) from nemo_rl.models.generation.vllm.vllm_worker_async import ( @@ -138,6 +140,52 @@ } +def test_context_capped_max_new_tokens(): + assert ( + _context_capped_max_new_tokens( + configured_max_new_tokens=8192, + input_length=3058, + max_model_len=8192, + ) + == 5134 + ) + assert ( + _context_capped_max_new_tokens( + configured_max_new_tokens=256, + input_length=3058, + max_model_len=8192, + ) + == 256 + ) + with pytest.raises(ValueError, match="exhausts the model context"): + _context_capped_max_new_tokens( + configured_max_new_tokens=8192, + input_length=8192, + max_model_len=8192, + ) + + +def test_sampling_params_preserve_bad_words(): + worker = object.__new__(VllmGenerationWorkerImpl) + worker.cfg = { + "top_k": None, + "temperature": 1.0, + "top_p": 1.0, + "max_new_tokens": 128, + "stop_token_ids": None, + "bad_words": ["", ""], + "ignore_eos": False, + } + worker.SamplingParams = lambda **kwargs: kwargs + + sampling_params = worker._build_sampling_params( + greedy=False, + stop_strings=None, + ) + + assert sampling_params["bad_words"] == ["", ""] + + def test_resolve_enable_prefix_caching_respects_explicit_config(monkeypatch): def raise_if_called(): raise AssertionError("CUDA capability should not be queried") diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py index dc7fceae38f..0b1ed867e58 100644 --- a/tests/unit/models/generation/test_vllm_patches.py +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -42,6 +42,9 @@ _TOOL_PARSER_SOURCE = "tool_parsers/utils.py" _PATCH_FN = "_patch_vllm_tool_parser_namespace_tool" _MARKER = "except ImportError: # openai < 2.25.0 predates namespace tools" +_RADIO_SOURCE = "model_executor/models/radio.py" +_RADIO_PATCH_FN = "_patch_vllm_radio_layerscale_loader" +_RADIO_MARKER = "initializer_factor = self.config.initializer_factor" @pytest.fixture @@ -53,6 +56,15 @@ def patched_tool_parser_source(tmp_path, monkeypatch): return copied +@pytest.fixture +def patched_radio_source(tmp_path, monkeypatch): + """The installed vLLM RADIO loader, unpatched then patched in tmp.""" + copied = write_unpatched_copy(_RADIO_SOURCE, _RADIO_PATCH_FN, tmp_path / "radio.py") + monkeypatch.setattr(patches, "_get_vllm_file", lambda _relative: str(copied)) + patches._patch_vllm_radio_layerscale_loader(logging.getLogger(__name__)) + return copied + + @pytest.mark.vllm def test_namespace_tool_patch_anchor_still_matches_installed_vllm( patched_tool_parser_source, @@ -100,6 +112,52 @@ def test_namespace_tool_stub_never_matches(patched_tool_parser_source): assert not isinstance(value, stub_cls) +@pytest.mark.vllm +def test_radio_layerscale_patch_anchor_still_matches_installed_vllm( + patched_radio_source, +): + """Pin the vLLM 0.25.1 RADIO loader shape used by the source patch.""" + content = patched_radio_source.read_text() + assert _RADIO_MARKER in content + assert "Skip layer-scale entries that vLLM doesn't use" not in content + ast.parse(content) + + +@pytest.mark.vllm +def test_radio_layerscale_patch_loads_explicit_and_initializes_folded_weights( + patched_radio_source, +): + content = patched_radio_source.read_text() + assert 'vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}"' in content + assert 'name.endswith((".ls1", ".ls2"))' in content + assert "param.data.fill_(initializer_factor)" in content + assert "loaded_params.add(name)" in content + + +@pytest.mark.vllm +def test_radio_layerscale_patch_is_idempotent(patched_radio_source, monkeypatch): + before = patched_radio_source.read_text() + monkeypatch.setattr( + patches, "_get_vllm_file", lambda _relative: str(patched_radio_source) + ) + + patches._patch_vllm_radio_layerscale_loader(logging.getLogger(__name__)) + + assert patched_radio_source.read_text() == before + + +def test_radio_layerscale_patch_warns_on_unknown_source(monkeypatch, tmp_path, caplog): + radio_source = tmp_path / "radio.py" + radio_source.write_text("class RadioModel:\n pass\n") + monkeypatch.setattr(patches, "_get_vllm_file", lambda _relative: str(radio_source)) + + with caplog.at_level(logging.WARNING): + patches._patch_vllm_radio_layerscale_loader(logging.getLogger(__name__)) + + assert radio_source.read_text() == "class RadioModel:\n pass\n" + assert "vLLM 0.25.1 source shape was not found" in caplog.text + + @pytest.mark.parametrize( "existing,extra,expected", [ diff --git a/tests/unit/models/generation/test_vllm_sparse_refit.py b/tests/unit/models/generation/test_vllm_sparse_refit.py index 1bee0dfa7ea..e70c809ca6f 100644 --- a/tests/unit/models/generation/test_vllm_sparse_refit.py +++ b/tests/unit/models/generation/test_vllm_sparse_refit.py @@ -571,6 +571,21 @@ async def test_async_sparse_refit_post_init_records_worker_locality() -> None: ] +def test_sync_post_init_binds_numa() -> None: + worker = VllmGenerationWorkerImpl.__new__(VllmGenerationWorkerImpl) + worker._sparse_refit_receiver = None + worker._mtp_load_from_disk = False + worker.report_device_id = MagicMock(return_value=["0"]) + worker.llm = MagicMock() + + worker.post_init() + + assert worker.vllm_device_ids == ["0"] + assert worker.llm.collective_rpc.call_args_list == [ + call("bind_numa", args=()), + ] + + def test_async_sparse_refit_exposes_zmq_relay(monkeypatch) -> None: from nemo_rl.models.generation.vllm import vllm_sparse_refit as refit_module diff --git a/tests/unit/models/generation/test_vllm_spec_decode_clamp.py b/tests/unit/models/generation/test_vllm_spec_decode_clamp.py index 6baa9b416ac..f88ca65dd97 100644 --- a/tests/unit/models/generation/test_vllm_spec_decode_clamp.py +++ b/tests/unit/models/generation/test_vllm_spec_decode_clamp.py @@ -36,3 +36,27 @@ def test_spec_decode_max_tokens_clamp( ) == expected ) + + +@pytest.mark.parametrize( + ("cap_to_context", "spec_lookahead", "expected"), + [ + (False, 0, 400), + (True, 0, 300), + (False, 5, 294), + (True, 5, 294), + ], +) +def test_request_max_new_tokens_combines_context_and_spec_limits( + cap_to_context, spec_lookahead, expected +): + assert ( + BaseVllmGenerationWorker._request_max_new_tokens( + configured_max_new_tokens=400, + input_length=700, + max_model_len=1000, + cap_to_context=cap_to_context, + spec_lookahead=spec_lookahead, + ) + == expected + ) diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py index e72431ba114..834d49a1612 100644 --- a/tests/unit/models/megatron/test_megatron_data.py +++ b/tests/unit/models/megatron/test_megatron_data.py @@ -330,6 +330,158 @@ def test_process_microbatch_no_packing_propagates_mtp_loss_mask( assert result.mtp_loss_mask is not None assert torch.equal(result.mtp_loss_mask, mtp_loss_mask) + @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 + ) + @patch( + "nemo_rl.models.megatron.data.get_packed_seq_cp_partition_indices", + return_value=torch.tensor([0, 3, 4, 7]), + ) + @patch("nemo_rl.models.megatron.data._pack_sequences_for_megatron") + def test_process_microbatch_keeps_full_thd_for_model_cp_slicing( + self, mock_pack, mock_indices, mock_cp_world, mock_cp_rank + ): + """Full THD input does not calculate replay indices when routes are absent.""" + from nemo_rl.models.megatron.data import process_microbatch + + full_tokens = torch.tensor([[1, 2, 3, 0, 4, 5, 0, 0]]) + local_tokens = full_tokens[:, [0, 3, 4, 7]] + cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32) + cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32) + mock_pack.return_value = ( + full_tokens, + local_tokens, + MagicMock(), + cu_seqlens, + cu_seqlens_padded, + ) + input_ids = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]]) + + result = process_microbatch( + {"input_ids": input_ids, "input_lengths": torch.tensor([3, 2])}, + seq_length_key="input_lengths", + pack_sequences=True, + model_slices_context_parallel_inputs=True, + straggler_timer=MagicMock(), + ) + + assert torch.equal(result.input_ids, full_tokens) + assert torch.equal(result.input_ids_cp_sharded, full_tokens) + assert torch.equal(result.packed_seq_params.cu_seqlens_q, cu_seqlens) + assert torch.equal( + result.packed_seq_params.cu_seqlens_q_padded, cu_seqlens_padded + ) + assert result.packed_seq_params.total_tokens == 8 + mock_indices.assert_not_called() + + @pytest.mark.parametrize( + ("cu_seqlens", "cu_seqlens_padded", "expected_pad_between_seqs"), + [ + ([0, 3], [0, 4], True), + ([0, 4], [0, 4], False), + ], + ids=["trailing-padding", "no-padding"], + ) + @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 + ) + @patch("nemo_rl.models.megatron.data._pack_sequences_for_megatron") + def test_process_microbatch_marks_single_sequence_trailing_padding( + self, + mock_pack, + mock_cp_world, + mock_cp_rank, + cu_seqlens, + cu_seqlens_padded, + expected_pad_between_seqs, + ): + from nemo_rl.models.megatron.data import process_microbatch + + full_tokens = torch.tensor([[1, 2, 3, 0]]) + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32) + cu_seqlens_padded = torch.tensor(cu_seqlens_padded, dtype=torch.int32) + mock_pack.return_value = ( + full_tokens, + full_tokens, + MagicMock(), + cu_seqlens, + cu_seqlens_padded, + ) + + result = process_microbatch( + { + "input_ids": full_tokens, + "input_lengths": cu_seqlens[1:].clone(), + }, + seq_length_key="input_lengths", + pack_sequences=True, + model_slices_context_parallel_inputs=True, + straggler_timer=MagicMock(), + ) + + assert result.packed_seq_params.pad_between_seqs is expected_pad_between_seqs + assert torch.equal(result.packed_seq_params.cu_seqlens_q, cu_seqlens) + assert torch.equal( + result.packed_seq_params.cu_seqlens_q_padded, cu_seqlens_padded + ) + + def test_process_microbatch_rejects_mtp_with_model_cp_slicing(self): + from nemo_rl.models.megatron.data import process_microbatch + + with pytest.raises(NotImplementedError, match="do not yet support MTP"): + process_microbatch( + { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "input_lengths": torch.tensor([4]), + "mtp_loss_mask": torch.ones(1, 4), + }, + seq_length_key="input_lengths", + pack_sequences=True, + model_slices_context_parallel_inputs=True, + straggler_timer=MagicMock(), + ) + + def test_caller_packing_matches_mbridge_thd_contract(self): + from megatron.bridge.data.packing.in_batch import ( + pack_right_padded_sequence_batch_to_mcore_thd, + ) + + from nemo_rl.models.megatron.data import _pack_sequences_for_megatron + + input_ids = torch.tensor([[1, 2, 3, 0, 0], [4, 5, 0, 0, 0]]) + seq_lengths = torch.tensor([3, 2]) + ( + full_tokens, + _local_tokens, + _packed_seq_params, + cu_seqlens, + cu_seqlens_padded, + ) = _pack_sequences_for_megatron( + input_ids, + seq_lengths, + pad_individual_seqs_to_multiple_of=4, + cp_size=1, + ) + mbridge_batch = { + "input_ids": input_ids.clone(), + "position_ids": torch.arange(input_ids.shape[1]) + .unsqueeze(0) + .expand_as(input_ids) + .clone(), + "attention_mask": torch.arange(input_ids.shape[1]).unsqueeze(0) + < seq_lengths.unsqueeze(1), + } + pack_right_padded_sequence_batch_to_mcore_thd( + mbridge_batch, + pad_to_multiple_of=4, + ) + + assert torch.equal(full_tokens, mbridge_batch["input_ids"]) + assert torch.equal(cu_seqlens, mbridge_batch["cu_seqlens_q"]) + assert torch.equal(cu_seqlens_padded, mbridge_batch["cu_seqlens_q_padded"]) + @patch("nemo_rl.models.megatron.data.get_ltor_masks_and_position_ids") def test_process_microbatch_no_packing_mtp_loss_mask_absent(self, mock_get_masks): """mtp_loss_mask defaults to None when not provided.""" @@ -463,6 +615,69 @@ def test_process_microbatch_packs_routed_experts_with_tokens( assert torch.equal(result.routed_experts, packed_routed_experts) assert torch.equal(result.routed_experts_cp_sharded, cp_routed_experts) + @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 + ) + @patch( + "nemo_rl.models.megatron.data.get_packed_seq_cp_partition_indices", + return_value=torch.tensor([0, 3, 4, 7]), + ) + @patch("nemo_rl.models.megatron.data._shard_routed_experts_for_cp") + @patch("nemo_rl.models.megatron.data._pack_sequences_for_megatron") + def test_model_cp_slicing_uses_shared_indices_for_router_replay( + self, + mock_pack, + mock_shard, + mock_indices, + mock_cp_world, + mock_cp_rank, + ): + from nemo_rl.models.megatron.data import process_microbatch + + input_ids = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]]) + routed_experts = torch.arange(2 * 4 * 3 * 2, dtype=torch.int32).reshape( + 2, 4, 3, 2 + ) + packed_tokens = torch.tensor([[1, 2, 3, 0, 4, 5, 0, 0]]) + packed_routes = torch.arange(1 * 8 * 3 * 2, dtype=torch.int32).reshape( + 1, 8, 3, 2 + ) + cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32) + cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32) + mock_pack.return_value = ( + packed_tokens, + packed_tokens[:, [0, 3, 4, 7]], + MagicMock(), + cu_seqlens, + cu_seqlens_padded, + ) + mock_shard.return_value = ( + packed_routes, + torch.full_like(packed_routes[:, :4], -1), + None, + None, + ) + + result = process_microbatch( + { + "input_ids": input_ids, + "input_lengths": torch.tensor([3, 2]), + "routed_experts": routed_experts, + }, + seq_length_key="input_lengths", + pack_sequences=True, + model_slices_context_parallel_inputs=True, + straggler_timer=MagicMock(), + ) + + assert torch.equal(result.input_ids_cp_sharded, packed_tokens) + assert torch.equal( + result.routed_experts_cp_sharded, + packed_routes[:, [0, 3, 4, 7]], + ) + mock_indices.assert_called_once() + def test_process_microbatch_packing_requires_seq_length_key(self): """Test that packing requires seq_length_key.""" from nemo_rl.models.megatron.data import process_microbatch @@ -566,13 +781,10 @@ def test_process_microbatch_delegate_pack_to_model(self, mock_prepare, mock_pack assert torch.equal(result.cu_seqlens_padded, mock_cu_seqlens_padded) def test_process_microbatch_delegate_pack_rejects_mtp_loss_mask(self): - """delegate_pack_to_model must reject a pre-computed mtp_loss_mask. + """Self-packing models must explicitly advertise MTP-mask ownership. - The VLM self-packing path does not pack/propagate mtp_loss_mask, so MTP - training would be silently dropped. process_microbatch must fail loudly - rather than produce wrong results. Regression guard for issue #2869: the - worker now only creates mtp_loss_mask when MTP is enabled, but if a mask - ever reaches this path it must raise instead of being silently ignored. + Qwen3-VL and other wrappers that have not implemented this contract stay + fail-closed rather than receiving a full-batch mask for CP-sharded tokens. """ from nemo_rl.models.megatron.data import process_microbatch @@ -596,8 +808,31 @@ def test_process_microbatch_delegate_pack_rejects_mtp_loss_mask(self): straggler_timer=MagicMock(), ) - assert "MTP training is not supported with VLM sequence packing" in str( - exc_info.value + assert "model_owns_mtp_loss_mask_packing" in str(exc_info.value) + + def test_process_microbatch_delegates_padded_mtp_loss_mask(self): + """A capable wrapper receives a padded full mask to pack with its IDs.""" + from nemo_rl.models.megatron.data import process_microbatch + + input_ids = torch.tensor([[1, 2, 3, 0, 0], [4, 5, 0, 0, 0]]) + mtp_loss_mask = torch.tensor([[0, 0, 1, 0, 0], [0, 1, 0, 0, 0]]) + result = process_microbatch( + { + "input_ids": input_ids, + "input_lengths": torch.tensor([3, 2]), + "mtp_loss_mask": mtp_loss_mask, + }, + seq_length_key="input_lengths", + pad_individual_seqs_to_multiple_of=4, + pack_sequences=True, + delegate_pack_to_model=True, + delegate_mtp_loss_mask_to_model=True, + ) + + assert result.input_ids_cp_sharded.shape == (2, 4) + assert torch.equal( + result.mtp_loss_mask, + torch.tensor([[0, 0, 1, 0], [0, 1, 0, 0]]), ) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index dc57b56a3eb..f95e7e99e76 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -910,6 +910,19 @@ def test_pipeline_dtype_mapping(self): class TestApplyPerformanceConfig: """Tests for _apply_performance_config function.""" + @staticmethod + def _config(*, attention_backend=None): + megatron_cfg = { + "activation_checkpointing": False, + "apply_rope_fusion": False, + "bias_activation_fusion": False, + "gradient_accumulation_fusion": False, + "use_fused_weighted_squared_relu": False, + } + if attention_backend is not None: + megatron_cfg["attention_backend"] = attention_backend + return {"megatron_cfg": megatron_cfg} + def test_basic_performance_config(self): """Test applying basic performance configuration.""" from nemo_rl.models.megatron.setup import _apply_performance_config @@ -954,6 +967,94 @@ def test_activation_checkpointing_enabled(self): assert model_cfg.recompute_method == "uniform" assert model_cfg.recompute_num_layers == 1 + def test_expanded_omni_defaults_to_auto_attention(self, monkeypatch): + """Expanded Omni uses backend dispatch without relying on a recipe.""" + from megatron.core.transformer.enums import AttnBackend + + from nemo_rl.models.megatron.setup import _apply_performance_config + + for variable in ("NVTE_FUSED_ATTN", "NVTE_FLASH_ATTN", "NVTE_UNFUSED_ATTN"): + monkeypatch.setenv(variable, "1") + + model_cfg = SimpleNamespace( + gated_linear_unit=True, + attention_backend=AttnBackend.flash, + nemotron_omni_contract="expanded_sequence_v1", + ) + _apply_performance_config(model_cfg, self._config()) + + assert model_cfg.attention_backend is AttnBackend.auto + for variable in ("NVTE_FUSED_ATTN", "NVTE_FLASH_ATTN", "NVTE_UNFUSED_ATTN"): + assert variable not in os.environ + + @pytest.mark.parametrize("attention_backend", ["auto", "unfused"]) + def test_expanded_omni_preserves_supported_explicit_attention_backend( + self, attention_backend + ): + """Expanded Omni preserves an explicitly selected compatible backend.""" + from megatron.core.transformer.enums import AttnBackend + + from nemo_rl.models.megatron.setup import _apply_performance_config + + model_cfg = SimpleNamespace( + gated_linear_unit=True, + attention_backend=AttnBackend.flash, + nemotron_omni_contract="expanded_sequence_v1", + ) + _apply_performance_config( + model_cfg, self._config(attention_backend=attention_backend) + ) + + assert model_cfg.attention_backend is AttnBackend[attention_backend] + + def test_expanded_omni_rejects_flash_attention(self): + """Flash cannot represent expanded Omni's padded multi-row THD batches.""" + from nemo_rl.models.megatron.setup import _apply_performance_config + + model_cfg = SimpleNamespace( + gated_linear_unit=True, + nemotron_omni_contract="expanded_sequence_v1", + ) + with pytest.raises( + ValueError, + match="does not support attention_backend='flash'", + ): + _apply_performance_config( + model_cfg, self._config(attention_backend="flash") + ) + + @pytest.mark.parametrize( + "model_contract", + [None, "llava_collapse_expand_v1"], + ids=["non-omni", "legacy-llava"], + ) + def test_non_expanded_model_preserves_provider_attention_backend( + self, model_contract + ): + """Models outside the expanded Omni contract retain provider defaults.""" + from megatron.core.transformer.enums import AttnBackend + + from nemo_rl.models.megatron.setup import _apply_performance_config + + model_cfg = SimpleNamespace( + gated_linear_unit=True, + attention_backend=AttnBackend.flash, + nemotron_omni_contract=model_contract, + ) + _apply_performance_config(model_cfg, self._config()) + + assert model_cfg.attention_backend is AttnBackend.flash + + def test_invalid_attention_backend_raises(self): + """Invalid explicit backends retain the generic validation behavior.""" + from nemo_rl.models.megatron.setup import _apply_performance_config + + model_cfg = SimpleNamespace(gated_linear_unit=True) + with pytest.raises(ValueError, match="Invalid attention backend"): + _apply_performance_config( + model_cfg, self._config(attention_backend="invalid") + ) + def test_activation_func_required_when_not_gated(self): """Test that activation_func is required when not using gated_linear_unit.""" from nemo_rl.models.megatron.setup import _apply_performance_config diff --git a/tests/unit/models/megatron/test_nemotron_omni_model.py b/tests/unit/models/megatron/test_nemotron_omni_model.py new file mode 100644 index 00000000000..2cc5827e727 --- /dev/null +++ b/tests/unit/models/megatron/test_nemotron_omni_model.py @@ -0,0 +1,516 @@ +# 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. + +"""Distributed functional coverage for NeMo-RL's Nemotron Omni contract.""" + +import copy +import functools +import gc +import os +from dataclasses import dataclass + +import pytest +import torch + +# This module is collected by catch-all unit-test lanes that intentionally do +# not install the mcore extra. Skip before importing MBridge so those lanes can +# deselect the mcore-marked tests without failing during collection. +pytest.importorskip("megatron.bridge") + +from megatron.bridge.models.nemotron_omni.nemotron_omni_provider import ( + NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT, + NemotronOmniModelProvider, +) +from megatron.core import dist_checkpointing, parallel_state +from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnBackend + +from nemo_rl.data.multimodal_utils import PackedTensor +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.model_utils import ( + from_parallel_logits_to_logprobs_packed_sequences, +) +from nemo_rl.models.megatron.data import get_microbatch_iterator, process_microbatch +from nemo_rl.models.megatron.train import ( + LogprobsPostProcessor, + megatron_forward_backward, +) + +pytestmark = pytest.mark.mcore + +_IMAGE_TOKEN_ID = 18 + + +@dataclass +class _TinyOmniProvider(NemotronOmniModelProvider): + """Small real RADIO/NemotronH model for a two-rank functional test.""" + + has_sound: bool = False + language_model_type: str = "nemotron6-moe" + hidden_size: int = 128 + ffn_hidden_size: int = 256 + num_attention_heads: int = 4 + num_query_groups: int = 2 + kv_channels: int = 32 + mamba_num_heads: int = 4 + mamba_head_dim: int = 32 + mamba_num_groups: int = 2 + mamba_state_dim: int = 16 + hybrid_layer_pattern: str = "M" + vocab_size: int = 128 + seq_length: int = 32 + image_token_index: int = _IMAGE_TOKEN_ID + img_start_token_id: int = 21 + img_end_token_id: int = 22 + tokenizer_type: str = "nemotron6-moe" + dynamic_resolution: bool = True + use_vision_backbone_fp8_arch: bool = False + vision_proj_ffn_hidden_size: int = 256 + pipeline_model_parallel_size: int = 1 + use_cpu_initialization: bool = True + gradient_accumulation_fusion: bool = False + nemotron_omni_contract: str = NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT + + def _build_vision_config(self, language_cfg): + vision_cfg = copy.deepcopy(language_cfg) + vision_cfg.sequence_parallel = False + vision_cfg.context_parallel_size = 1 + vision_cfg.tp_comm_overlap = False + vision_cfg.recompute_granularity = None + vision_cfg.recompute_method = None + vision_cfg.recompute_num_layers = None + vision_cfg.mtp_num_layers = None + vision_cfg.num_layers = 1 + vision_cfg.pipeline_model_parallel_size = 1 + vision_cfg.num_attention_heads = 4 + vision_cfg.add_bias_linear = True + vision_cfg.add_qkv_bias = True + vision_cfg.hidden_size = 128 + vision_cfg.ffn_hidden_size = 256 + vision_cfg.gated_linear_unit = False + vision_cfg.kv_channels = 32 + vision_cfg.num_query_groups = 4 + vision_cfg.normalization = "LayerNorm" + vision_cfg.qk_layernorm = False + vision_cfg.layernorm_epsilon = 1e-6 + vision_cfg.class_token_len = 10 + return vision_cfg + + +def _build_distributed_model( + *, + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + context_parallel_size: int = 2, + sequence_parallel: bool = False, + language_layer_pattern: str = "M", + attention_backend: AttnBackend | None = None, +): + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=tensor_parallel_size, + pipeline_model_parallel_size=pipeline_parallel_size, + context_parallel_size=context_parallel_size, + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + provider_kwargs = { + "freeze_language_model": True, + "tensor_model_parallel_size": tensor_parallel_size, + "pipeline_model_parallel_size": pipeline_parallel_size, + "context_parallel_size": context_parallel_size, + "sequence_parallel": sequence_parallel, + "hybrid_layer_pattern": "|".join( + language_layer_pattern for _ in range(pipeline_parallel_size) + ), + } + if attention_backend is not None: + provider_kwargs["attention_backend"] = attention_backend + provider = _TinyOmniProvider( + **provider_kwargs, + ) + provider.finalize() + models = provider.provide_distributed_model( + ddp_config=DistributedDataParallelConfig( + grad_reduce_in_fp32=True, + overlap_grad_reduce=False, + use_distributed_optimizer=False, + check_for_nan_in_grad=True, + ), + wrap_with_ddp=True, + mixed_precision_wrapper=None, + ) + assert len(models) == 1 + return models[0] + + +def _expanded_fixture(device: torch.device): + input_ids = torch.tensor( + [ + [7, 21, 18, 18, 22, 9, 10, 0], + [11, 21, 18, 22, 12, 0, 0, 0], + ], + dtype=torch.long, + device=device, + ) + lengths = torch.tensor([7, 5], dtype=torch.long, device=device) + generator = torch.Generator(device=device) + generator.manual_seed(2026) + images = torch.randn(2, 3, 32, 64, generator=generator, device=device) + images[1, :, :, 32:] = 0 + image_sizes = torch.tensor([[32, 64], [32, 32]], dtype=torch.int32, device=device) + return input_ids, lengths, images, image_sizes + + +def _forward(model): + device = torch.device("cuda", torch.cuda.current_device()) + input_ids, lengths, images, image_sizes = _expanded_fixture(device) + processed = process_microbatch( + {"input_ids": input_ids, "input_lengths": lengths}, + seq_length_key="input_lengths", + pad_individual_seqs_to_multiple_of=4, + pack_sequences=True, + model_slices_context_parallel_inputs=True, + ) + output = model( + input_ids=processed.input_ids_cp_sharded, + attention_mask=processed.attention_mask, + packed_seq_params=processed.packed_seq_params, + pixel_values=images, + imgs_sizes=image_sizes, + ) + logprobs = from_parallel_logits_to_logprobs_packed_sequences( + output, + target=processed.input_ids, + cu_seqlens_padded=processed.cu_seqlens_padded, + unpacked_seqlen=input_ids.shape[1], + vocab_start_index=parallel_state.get_tensor_model_parallel_rank() + * output.shape[-1], + vocab_end_index=(parallel_state.get_tensor_model_parallel_rank() + 1) + * output.shape[-1], + group=parallel_state.get_tensor_model_parallel_group(), + inference_only=False, + cp_group=parallel_state.get_context_parallel_group(), + ) + prediction_mask = torch.arange(input_ids.shape[1] - 1, device=device).unsqueeze( + 0 + ) < (lengths - 1).unsqueeze(1) + loss = -(logprobs * prediction_mask).sum() / prediction_mask.sum() + return loss, output, logprobs, processed + + +def _run_training_checkpoint_roundtrip( + rank: int, + world_size: int, + *, + checkpoint_dir: str, +) -> None: + assert world_size == 2 + model = _build_distributed_model() + model.train() + model.zero_grad_buffer() + + loss, output, _, _ = _forward(model) + loss.backward() + model.finish_grad_sync() + + core_model = model.module + gradients = {} + before_update = {} + optimizer_parameters = [] + for name, parameter in core_model.named_parameters(): + if not parameter.requires_grad: + continue + assert name.startswith(("vision_model.", "vision_projection.")) + assert hasattr(parameter, "main_grad") + assert torch.isfinite(parameter.main_grad).all() + rank_zero_gradient = parameter.main_grad.detach().clone() + torch.distributed.broadcast(rank_zero_gradient, src=0) + torch.testing.assert_close( + parameter.main_grad, rank_zero_gradient, rtol=0, atol=0 + ) + gradients[name] = parameter.main_grad + before_update[name] = parameter.detach().clone() + parameter.grad = parameter.main_grad.to(parameter.dtype).clone() + optimizer_parameters.append(parameter) + assert gradients + + optimizer = torch.optim.SGD(optimizer_parameters, lr=1.0) + optimizer.step() + changed = { + name + for name, parameter in core_model.named_parameters() + if name in before_update and not torch.equal(parameter, before_update[name]) + } + assert any(name.startswith("vision_model.") for name in changed) + assert any(name.startswith("vision_projection.") for name in changed) + + model.eval() + with torch.no_grad(): + _, post_update_output, _, _ = _forward(model) + post_update_output = post_update_output.detach().clone() + + metadata = { + "dp_cp_group": parallel_state.get_data_parallel_group( + with_context_parallel=True + ) + } + sharded_state = core_model.sharded_state_dict(metadata=metadata) + assert changed <= sharded_state.keys() + if rank == 0: + os.makedirs(checkpoint_dir, exist_ok=True) + torch.distributed.barrier() + dist_checkpointing.save({"model": sharded_state}, checkpoint_dir) + + provider = _TinyOmniProvider( + freeze_language_model=True, + tensor_model_parallel_size=1, + context_parallel_size=2, + sequence_parallel=False, + ) + provider.finalize() + restored_model = provider.provide().cuda().eval() + restore_template = restored_model.sharded_state_dict(metadata=metadata) + loaded_state = dist_checkpointing.load({"model": restore_template}, checkpoint_dir) + incompatible = restored_model.load_state_dict(loaded_state["model"]) + assert not incompatible.missing_keys + assert not incompatible.unexpected_keys + + restored_parameters = dict(restored_model.named_parameters()) + original_parameters = dict(core_model.named_parameters()) + for name in changed: + torch.testing.assert_close( + restored_parameters[name], original_parameters[name], rtol=0, atol=0 + ) + with torch.no_grad(): + _, restored_output, _, _ = _forward(restored_model) + torch.testing.assert_close(restored_output, post_update_output, rtol=0, atol=0) + + if rank == 0: + print( + "NEMOTRON_OMNI_CP2_DCP_ROUNDTRIP " + f"loss={loss.item():.8f} changed_tensors={len(changed)} " + "post_restore_max_logit_abs_diff=0.00000000", + flush=True, + ) + + del output, post_update_output, restored_output, optimizer, model + del core_model, restored_model + gc.collect() + torch.cuda.empty_cache() + torch.distributed.barrier() + parallel_state.destroy_model_parallel() + + +def test_nemotron_omni_cp2_training_and_checkpoint_roundtrip( + distributed_test_runner, + tmp_path, +): + test_fn = functools.partial( + _run_training_checkpoint_roundtrip, + checkpoint_dir=str(tmp_path / "nemotron_omni_cp2_dcp"), + ) + distributed_test_runner(test_fn, world_size=2) + + +def _run_parallel_forward_contract( + rank: int, + world_size: int, + *, + tensor_parallel_size: int, + context_parallel_size: int, +) -> None: + assert world_size == tensor_parallel_size * context_parallel_size + model = _build_distributed_model( + tensor_parallel_size=tensor_parallel_size, + context_parallel_size=context_parallel_size, + sequence_parallel=True, + ) + model.eval() + + with torch.no_grad(): + loss, output, logprobs, _ = _forward(model) + + assert torch.isfinite(loss) + assert torch.isfinite(output).all() + assert torch.isfinite(logprobs).all() + assert logprobs.shape == (2, 7) + + reference = logprobs.clone() + torch.distributed.broadcast(reference, src=0) + torch.testing.assert_close(logprobs, reference, rtol=0, atol=0) + + del model, output, logprobs + gc.collect() + torch.cuda.empty_cache() + torch.distributed.barrier() + parallel_state.destroy_model_parallel() + + +@pytest.mark.parametrize( + ("tensor_parallel_size", "context_parallel_size", "world_size"), + [ + pytest.param(2, 1, 2, id="tp2-sp"), + pytest.param(2, 2, 4, id="tp2-cp2-sp"), + ], +) +def test_nemotron_omni_parallel_forward_and_logprob_contract( + distributed_test_runner, + tensor_parallel_size, + context_parallel_size, + world_size, +): + test_fn = functools.partial( + _run_parallel_forward_contract, + tensor_parallel_size=tensor_parallel_size, + context_parallel_size=context_parallel_size, + ) + distributed_test_runner(test_fn, world_size=world_size) + + +def _run_padded_multirow_attention_contract(rank: int, world_size: int) -> None: + assert world_size == 2 + for variable in ("NVTE_FUSED_ATTN", "NVTE_FLASH_ATTN", "NVTE_UNFUSED_ATTN"): + os.environ.pop(variable, None) + model = _build_distributed_model( + context_parallel_size=2, + language_layer_pattern="*", + attention_backend=AttnBackend.auto, + ) + model.train() + model.zero_grad_buffer() + + loss, output, logprobs, processed = _forward(model) + packed_seq_params = processed.packed_seq_params + assert packed_seq_params.qkv_format == "thd" + torch.testing.assert_close( + packed_seq_params.cu_seqlens_q, + torch.tensor([0, 7, 12], dtype=torch.int32, device=output.device), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + packed_seq_params.cu_seqlens_q_padded, + torch.tensor([0, 8, 16], dtype=torch.int32, device=output.device), + rtol=0, + atol=0, + ) + assert torch.isfinite(loss) + assert torch.isfinite(output).all() + assert torch.isfinite(logprobs).all() + + loss.backward() + model.finish_grad_sync() + trainable_gradients = [ + parameter.main_grad + for parameter in model.module.parameters() + if parameter.requires_grad and hasattr(parameter, "main_grad") + ] + assert trainable_gradients + assert all(torch.isfinite(gradient).all() for gradient in trainable_gradients) + + del model, output, logprobs + gc.collect() + torch.cuda.empty_cache() + torch.distributed.barrier() + parallel_state.destroy_model_parallel() + + +def test_nemotron_omni_padded_multirow_attention_forward_backward( + distributed_test_runner, +): + distributed_test_runner(_run_padded_multirow_attention_contract, world_size=2) + + +def _run_pipeline_forward_contract(rank: int, world_size: int) -> None: + assert world_size == 2 + model = _build_distributed_model( + pipeline_parallel_size=2, + context_parallel_size=1, + ) + model.eval() + + device = torch.device("cuda", torch.cuda.current_device()) + input_ids, lengths, images, image_sizes = _expanded_fixture(device) + data = BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": lengths, + "pixel_values": PackedTensor( + [images[0:1], images[1:2]], + dim_to_pack=0, + ), + "imgs_sizes": PackedTensor( + [image_sizes[0:1], image_sizes[1:2]], + dim_to_pack=0, + ), + } + ) + data.micro_batch_indices = [[[0, 2]]] + data.micro_batch_lengths = [[int(lengths.sum().item())]] + cfg = { + "dynamic_batching": {"enabled": False}, + "sequence_packing": {"enabled": True}, + "make_sequence_length_divisible_by": 1, + "megatron_cfg": { + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 2, + "context_parallel_size": 1, + "sequence_parallel": False, + }, + } + ( + data_iterator, + num_microbatches, + micro_batch_size, + _, + padded_seq_length, + ) = get_microbatch_iterator( + data, + cfg, + mbs=2, + straggler_timer=None, + model_slices_context_parallel_inputs=True, + ) + + results = megatron_forward_backward( + model=model, + data_iterator=data_iterator, + num_microbatches=num_microbatches, + seq_length=padded_seq_length, + mbs=micro_batch_size, + post_processing_fn=LogprobsPostProcessor(cfg), + forward_only=True, + ) + + if parallel_state.is_pipeline_last_stage(): + assert len(results) == num_microbatches + for result in results: + assert torch.isfinite(result["logprobs"]).all() + assert result["logprobs"].shape == input_ids.shape + else: + assert results == [] + + del model + gc.collect() + torch.cuda.empty_cache() + torch.distributed.barrier() + parallel_state.destroy_model_parallel() + + +def test_nemotron_omni_pp2_scheduled_forward_contract(distributed_test_runner): + distributed_test_runner(_run_pipeline_forward_contract, world_size=2) diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 759ff20b1c1..8cd4b715d1d 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -18,6 +18,7 @@ from pathlib import Path from types import SimpleNamespace from typing import Optional +from unittest.mock import MagicMock import numpy as np import pytest @@ -45,6 +46,103 @@ pytestmark = pytest.mark.mcore +def test_model_owned_packing_capability_is_detected(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_self_packs_for_cp, + ) + + class ModelOwnedPackingModel: + model_owns_packing = True + + assert _model_self_packs_for_cp(ModelOwnedPackingModel()) + + +def test_model_owned_mtp_loss_mask_packing_capability_is_detected(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_self_packs_mtp_loss_mask, + ) + + class ModelOwnedPackingModel: + model_owns_mtp_loss_mask_packing = True + + assert _model_self_packs_mtp_loss_mask(ModelOwnedPackingModel()) + assert not _model_self_packs_mtp_loss_mask(object()) + + +def test_regular_model_does_not_delegate_packing(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_self_packs_for_cp, + ) + + assert not _model_self_packs_for_cp(object()) + + +def test_model_cp_slicing_capability_is_detected(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_slices_context_parallel_inputs, + ) + + class ModelSlicesContextParallelInputs: + model_slices_context_parallel_inputs = True + + assert _model_slices_context_parallel_inputs(ModelSlicesContextParallelInputs()) + assert not _model_slices_context_parallel_inputs(object()) + + +def test_model_cp_slicing_rejects_transfer_queue_setup(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model_slices_context_parallel_inputs = True + + with pytest.raises( + NotImplementedError, match="TransferQueue/SingleController does not yet support" + ): + worker.setup_data_plane(MagicMock()) + + +def test_refit_size_estimate_preserves_integral_buffer_dtype(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _estimate_refit_tensor_size_in_bytes, + ) + + param = torch.zeros(3, dtype=torch.int64) + + assert ( + _estimate_refit_tensor_size_in_bytes( + param, export_dtype=torch.bfloat16, tp_size=2, ep_size=4 + ) + == 3 * 8 * 2 * 4 + ) + + +def test_refit_size_estimate_casts_floating_weight_to_export_dtype(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _estimate_refit_tensor_size_in_bytes, + ) + + param = torch.zeros(3, dtype=torch.float32) + + assert ( + _estimate_refit_tensor_size_in_bytes( + param, export_dtype=torch.bfloat16, tp_size=2, ep_size=4 + ) + == 3 * 2 * 2 * 4 + ) + + +def test_qwen3vl_type_fallback_still_delegates_packing(): + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel + + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_self_packs_for_cp, + ) + + assert _model_self_packs_for_cp(Qwen3VLModel.__new__(Qwen3VLModel)) + + class _FakeTrainableModel: def __init__(self): self.train_called = False diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml index 828396d9914..5b14bd9bb8b 100644 --- a/tests/unit/reference_configs/distillation_math.yaml +++ b/tests/unit/reference_configs/distillation_math.yaml @@ -197,6 +197,7 @@ policy: &POLICY_BASE async_engine: false precision: ${...precision} kv_cache_dtype: "auto" + logprobs_mode: processed_logprobs tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 # When EP > 1, EP must be a multiple of TP since vLLM's EP = DP * TP diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index ce1a40e7ee6..5f4ffb2270d 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -378,6 +378,7 @@ policy: async_engine: false precision: ${policy.precision} kv_cache_dtype: "auto" + logprobs_mode: processed_logprobs tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 # When EP > 1, EP must be a multiple of TP since vLLM's EP = DP * TP diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index bfa297f5cf2..3610bfd09e8 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -237,6 +237,7 @@ policy: async_engine: false precision: ${policy.precision} kv_cache_dtype: "auto" + logprobs_mode: processed_logprobs tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index 1b7990b78cd..a3826111de3 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -256,7 +256,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites( ) -def test_nightly_compute_stays_below_3680_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_3800_hours(nightly_test_suite, tracker): command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}" print(f"Running command: {command}") @@ -288,8 +288,8 @@ def test_nightly_compute_stays_below_3680_hours(nightly_test_suite, tracker): f"Last line of output was not as expected: '{last_line}'" ) total_gpu_hours = float(last_line.split(":")[-1].strip()) - assert total_gpu_hours <= 3680, ( - f"Total GPU hours exceeded 3680: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 3800, ( + f"Total GPU hours exceeded 3800: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours)