diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 554c7b9324..c9473b1361 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 554c7b9324225aa863eee52e8b8fdde7abced2b1 +Subproject commit c9473b136167434e687e826ee84b67e5818a5fa8 diff --git a/docs/guides/nemotron-3-nano-omni.md b/docs/guides/nemotron-3-nano-omni.md index 0a35ba3cfe..306d57dcb5 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/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 0000000000..b738cf17a1 --- /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-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 0000000000..844917c1fc --- /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 6a16e18fec..e1aefce2b3 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -23,6 +23,7 @@ import decord 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 @@ -82,6 +83,8 @@ def __init__( self, tensors: Union[torch.Tensor, list[Optional[torch.Tensor]], list[None]], dim_to_pack: int, + *, + pad_to_max_shape: bool = False, ) -> None: assert tensors is not None, "Input tensors to PackedTensor cannot be None" @@ -97,6 +100,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 @@ -109,8 +113,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 @@ -125,12 +172,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": @@ -158,12 +213,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( @@ -195,8 +258,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 734ab6cf0e..64a716fde8 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 model-owned 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/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 025707c700..57ae6e6928 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -126,6 +126,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 44eb5d5d89..a38ee75ac1 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -28,6 +28,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 1d7fbfbdbb..61e896f162 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -282,6 +282,167 @@ def _patch_vllm_hermes_tool_parser_thread_safety(logger) -> None: logger.info("Successfully patched hermes tool parser for thread-safety.") +def _patch_vllm_bad_words_tokenization_cache(logger) -> None: + """Cache vLLM ``bad_words`` tokenization under a process-wide lock. + + vLLM 0.20.0 tokenizes the same string list for every request. Besides the + avoidable cost, concurrent calls into a Hugging Face fast tokenizer can + fail with ``RuntimeError: Already borrowed``. The Omni vLLM fork carries + this cache; apply the same narrow behavior to the supported stock wheel. + + The patch is deliberately fail-closed on source shape. If vLLM changes the + method, leave the installation untouched and require the new version to be + audited instead of applying a partial textual rewrite. + """ + try: + file_to_patch = _get_vllm_file("sampling_params.py") + except RuntimeError: + logger.warning("Could not locate sampling_params.py for bad_words patch.") + return + + old_import = "import json as json_mod\nfrom dataclasses import field" + new_import = ( + "import json as json_mod\nimport threading\nfrom dataclasses import field" + ) + constants_anchor = "_SAMPLING_EPS = 1e-5\n_MAX_TEMP = 1e-2\n" + cache_definition = ( + "_SAMPLING_EPS = 1e-5\n" + "_MAX_TEMP = 1e-2\n\n" + "# Cache tokenized bad_words across requests. Fast tokenizers are not\n" + "# safe for concurrent encode calls on the same tokenizer instance.\n" + "_BAD_WORDS_TOKEN_IDS_CACHE: dict[\n" + " tuple[int, tuple[str, ...]], list[list[int]]\n" + "] = {}\n" + "_BAD_WORDS_TOKEN_IDS_CACHE_LOCK = threading.Lock()\n" + "_BAD_WORDS_TOKEN_IDS_CACHE_MAX_ENTRIES = 1024\n" + ) + old_method = """ def update_from_tokenizer(self, tokenizer: TokenizerLike) -> None: + if not self.bad_words: + return + self._bad_words_token_ids = [] + for bad_word in self.bad_words: + # To prohibit words both at the beginning + # and in the middle of text + # (related to add_prefix_space tokenizer parameter) + for add_prefix_space in [False, True]: + prefix = " " if add_prefix_space else "" + prompt = prefix + bad_word.lstrip() + prompt_token_ids = tokenizer.encode( + text=prompt, add_special_tokens=False + ) + + # If no space at the beginning + # or if prefix space produces a new word token + if (not add_prefix_space) or ( + add_prefix_space + and prompt_token_ids[0] != self._bad_words_token_ids[-1][0] + and len(prompt_token_ids) == len(self._bad_words_token_ids[-1]) + ): + self._bad_words_token_ids.append(prompt_token_ids) + + invalid_token_ids = [ + token_id + for bad_words_token_ids in self._bad_words_token_ids + for token_id in bad_words_token_ids + if token_id < 0 or token_id > tokenizer.max_token_id + ] + if len(invalid_token_ids) > 0: + raise VLLMValidationError( + f"The model vocabulary size is {tokenizer.max_token_id + 1}," + f" but the following tokens" + f" were specified as bad: {invalid_token_ids}." + f" All token id values should be integers satisfying:" + f" 0 <= token_id <= {tokenizer.max_token_id}.", + parameter="bad_words", + value=self.bad_words, + ) +""" + new_method = """ def update_from_tokenizer(self, tokenizer: TokenizerLike) -> None: + if not self.bad_words: + return + cache_key = (id(tokenizer), tuple(self.bad_words)) + cached = _BAD_WORDS_TOKEN_IDS_CACHE.get(cache_key) + if cached is not None: + self._bad_words_token_ids = cached + return + with _BAD_WORDS_TOKEN_IDS_CACHE_LOCK: + cached = _BAD_WORDS_TOKEN_IDS_CACHE.get(cache_key) + if cached is not None: + self._bad_words_token_ids = cached + return + self._bad_words_token_ids = self._tokenize_bad_words(tokenizer) + if ( + len(_BAD_WORDS_TOKEN_IDS_CACHE) + >= _BAD_WORDS_TOKEN_IDS_CACHE_MAX_ENTRIES + ): + _BAD_WORDS_TOKEN_IDS_CACHE.clear() + _BAD_WORDS_TOKEN_IDS_CACHE[cache_key] = self._bad_words_token_ids + + def _tokenize_bad_words(self, tokenizer: TokenizerLike) -> list[list[int]]: + bad_words_token_ids: list[list[int]] = [] + for bad_word in self.bad_words: + # To prohibit words both at the beginning + # and in the middle of text + # (related to add_prefix_space tokenizer parameter) + for add_prefix_space in [False, True]: + prefix = " " if add_prefix_space else "" + prompt = prefix + bad_word.lstrip() + prompt_token_ids = tokenizer.encode( + text=prompt, add_special_tokens=False + ) + + # If no space at the beginning + # or if prefix space produces a new word token + if (not add_prefix_space) or ( + add_prefix_space + and prompt_token_ids[0] != bad_words_token_ids[-1][0] + and len(prompt_token_ids) == len(bad_words_token_ids[-1]) + ): + bad_words_token_ids.append(prompt_token_ids) + + invalid_token_ids = [ + token_id + for token_ids in bad_words_token_ids + for token_id in token_ids + if token_id < 0 or token_id > tokenizer.max_token_id + ] + if len(invalid_token_ids) > 0: + raise VLLMValidationError( + f"The model vocabulary size is {tokenizer.max_token_id + 1}," + f" but the following tokens" + f" were specified as bad: {invalid_token_ids}." + f" All token id values should be integers satisfying:" + f" 0 <= token_id <= {tokenizer.max_token_id}.", + parameter="bad_words", + value=self.bad_words, + ) + return bad_words_token_ids +""" + + with _locked_file_patch(file_to_patch) as (content, write_back): + if "_BAD_WORDS_TOKEN_IDS_CACHE_LOCK" in content: + logger.info("vLLM bad_words tokenization cache already applied.") + return + if ( + old_import not in content + or constants_anchor not in content + or old_method not in content + ): + logger.warning( + "Could not apply vLLM bad_words tokenization cache: expected " + "vLLM 0.20.0 source shape was not found in %s.", + file_to_patch, + ) + return + + content = content.replace(old_import, new_import, 1) + content = content.replace(constants_anchor, cache_definition, 1) + content = content.replace(old_method, new_method, 1) + write_back(content) + + logger.info("Successfully patched vLLM bad_words tokenization cache.") + + def _apply_vllm_patches( py_executable: str, *, extra_env_vars: list[str] | None = None ) -> None: @@ -295,3 +456,4 @@ def _apply_vllm_patches( _patch_vllm_llama_eagle3_own_lm_head(patch_logger) _patch_vllm_hermes_tool_parser_thread_safety(patch_logger) + _patch_vllm_bad_words_tokenization_cache(patch_logger) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index f6b2e5adb9..de9d172b9d 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -165,6 +165,78 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: e.g. {tensor_name: (shape, dtype)} """ self.state_dict_info = state_dict_info # pyrefly: ignore[implicitly-defined-attribute] This class does not define __init__ so assignments like this should be ignored + self._validate_nemotron_omni_radio_layerscale_refit(state_dict_info) + + def _is_nemotron_omni(self) -> bool: + architectures = set( + self.model_runner.vllm_config.model_config.architectures or [] + ) + return not architectures.isdisjoint( + { + "NemotronH_Nano_Omni_Reasoning_V3", + "NemotronH_Super_Omni_Reasoning_V3", + } + ) + + def _validate_nemotron_omni_radio_layerscale_refit( + self, state_dict_info: dict[str, Any] + ) -> None: + """Reject explicit LayerScale state that stock vLLM would ignore. + + This method can run while a colocated vLLM engine is asleep. It must not + read or mutate model tensors because vLLM's level-1 sleep allocator has + temporarily released their CUDA storage. + """ + if not self._is_nemotron_omni(): + return + + explicit_layerscale = [ + name + for name in state_dict_info + if name.startswith("vision_model.radio_model.model.blocks.") + and name.endswith((".ls1", ".ls2")) + ] + if explicit_layerscale: + raise RuntimeError( + "Nemotron Omni refit contains explicit RADIO LayerScale tensors, " + "but the stock vLLM 0.20 loader ignores them. Refusing to replace " + "checkpoint values with the folded-checkpoint identity behavior." + ) + + def _initialize_nemotron_omni_radio_layerscale(self) -> int: + """Set folded RADIO LayerScale parameters to identity while vLLM is awake. + + Nano/Super Omni checkpoints fold RADIO LayerScale into the adjacent + projection weights and therefore do not export ``ls1``/``ls2``. Stock + vLLM 0.20 leaves those parameters at dummy-initialized values during + direct load or refit, corrupting image inference. Initialize the + vLLM-only parameters once, immediately after engine creation and before + colocated sleep can release their CUDA storage. + """ + if not self._is_nemotron_omni(): + return 0 + + model = self.model_runner.model + vision_model = getattr(model, "vision_model", None) + if vision_model is None: + raise RuntimeError( + "Nemotron Omni vLLM model has no vision_model during initialization." + ) + + initializer_factor = getattr(vision_model.config, "initializer_factor", 1.0) + initialized = 0 + with torch.no_grad(): + for name, parameter in vision_model.named_parameters(): + if name.rsplit(".", 1)[-1] in {"ls1", "ls2"}: + parameter.fill_(initializer_factor) + initialized += 1 + + if initialized == 0: + raise RuntimeError( + "Nemotron Omni vLLM model exposes no RADIO ls1/ls2 parameters; " + "the expected vLLM 0.20 model layout may have changed." + ) + return initialized def _maybe_process_fp8_kv_cache(self) -> None: """Process weights after loading for FP8 KV cache (static scales).""" diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index a2a59292ca..dd90a668e1 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -48,6 +48,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: @@ -471,11 +484,21 @@ 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", + # Keep the main default, while allowing an RL recipe to request + # raw model logprobs. This is required when policy logprobs are + # compared with vLLM generation under a logits processor. + logprobs_mode=self.cfg["vllm_cfg"].get( + "logprobs_mode", "processed_logprobs" + ), **vllm_kwargs, ) self._create_engine(llm_kwargs) + # Nemotron Omni checkpoints fold RADIO LayerScale into adjacent weights, + # while stock vLLM still allocates ls1/ls2 parameters. Initialize those + # parameters before colocated level-1 sleep releases their CUDA storage; + # mutating them later from prepare_refit_info corrupts sleep/wake state. + self.llm.collective_rpc("_initialize_nemotron_omni_radio_layerscale") log_gpu_memory_diagnostics( label="after_engine_create", worker_type="VllmGenerationWorker", device_id=0 ) @@ -532,6 +555,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") or None, ignore_eos=self.cfg.get("ignore_eos", False), ) @@ -682,10 +706,6 @@ 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, - ) # verify inputs have correct padding verify_right_padding(data, pad_value=self.cfg["_pad_token_id"]) @@ -693,13 +713,31 @@ def generate( # 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" ) + if self.cfg["vllm_cfg"].get("cap_max_tokens_to_context", False): + max_model_len = int(self.llm.llm_engine.model_config.max_model_len) + sampling_params = [ + self._build_sampling_params( + greedy=greedy, + stop_strings=stop_strings, + max_new_tokens=_context_capped_max_new_tokens( + configured_max_new_tokens=int(self.cfg["max_new_tokens"]), + input_length=int(input_length), + max_model_len=max_model_len, + ), + ) + for input_length in input_lengths.tolist() + ] + 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/megatron/data.py b/nemo_rl/models/megatron/data.py index bb216c8e29..df56fc3147 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -91,6 +91,7 @@ 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, ) -> Iterator[ProcessedMicrobatch]: """Wrap a raw microbatch iterator to yield processed microbatches. @@ -124,6 +125,7 @@ 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, straggler_timer=straggler_timer, ) @@ -148,6 +150,7 @@ 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, ) -> Tuple[Iterator[ProcessedMicrobatch], int, int, int, int]: """Create a processed microbatch iterator from a batch of data. @@ -212,6 +215,7 @@ 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, ) # Compute padded sequence length for pipeline parallelism @@ -246,6 +250,7 @@ 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, straggler_timer: Optional[StragglerDetector] = None, ) -> ProcessedInputs: """Process a microbatch for Megatron model forward pass.""" @@ -286,11 +291,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,6 +324,25 @@ 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: token_identity = None diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index a88cd33411..ec31588775 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -68,6 +68,11 @@ _HF_CONFIG_PATCHED = False +_NEMOTRON_OMNI_ARCHITECTURES = { + "NemotronH_Nano_Omni_Reasoning_V3", + "NemotronH_Super_Omni_Reasoning_V3", +} + def _patch_hf_config_double_instantiation(): """Patch HF config classes whose __post_init__ fails with Megatron's recursive instantiation. @@ -323,6 +328,41 @@ def _get_hf_config_overrides_hash(overrides: dict[str, Any]) -> str: return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:12] +def _validate_nemotron_omni_megatron_lm_layout( + hf_config: Any, + checkpoint_dir: str, +) -> None: + """Reject legacy LLaVA checkpoints before building the canonical Omni model. + + Raw ``megatron_lm`` checkpoints have no Bridge ``run_config.yaml`` in which + to persist the versioned model contract. Their tensor namespace is therefore + the only reliable way to distinguish the old LLaVA wrapper from the canonical + model-owned-packing implementation. + """ + architectures = set(getattr(hf_config, "architectures", None) or ()) + if not architectures.intersection(_NEMOTRON_OMNI_ARCHITECTURES): + return + + from megatron.core import dist_checkpointing + + tensor_metadata = dist_checkpointing.load_tensors_metadata(checkpoint_dir) + legacy_keys = [ + str(key) for key in tensor_metadata if "llava_model" in str(key).split(".") + ] + if not legacy_keys: + return + + examples = ", ".join(repr(key) for key in legacy_keys[:3]) + raise RuntimeError( + "Cannot load this Nemotron Omni megatron_lm checkpoint as NemotronOmniModel: " + "its tensor namespace contains the legacy 'llava_model' wrapper " + f"({examples}). In this branch NemotronOmniModel uses model-owned packing and a " + "different top-level checkpoint layout. Reconvert the original HF checkpoint with " + "this branch, or use the explicit NemotronOmniLlavaModelProvider/" + "NemotronOmniLlavaBridge with the compatible legacy Megatron-LM pin." + ) + + def _resolve_iter_dir_from_root(path: str, not_found_msg: str) -> str: """Resolve the latest iteration directory under ``path``. @@ -485,6 +525,7 @@ def setup_model_config( hf_cfg = AutoConfig.from_pretrained( hf_model_name, trust_remote_code=True, **hf_config_overrides ) + _validate_nemotron_omni_megatron_lm_layout(hf_cfg, pretrained_path) bridge_obj = AutoBridge.from_hf_config(hf_cfg) model_cfg = bridge_obj.to_megatron_provider(load_weights=False) else: @@ -539,6 +580,14 @@ def setup_model_config( model_cfg = cfg_from_pretrained.model cfg_from_pretrained.logger = LoggerConfig() + # New and legacy Nemotron Omni checkpoints once serialized the same + # provider class name with different semantics. Providers that implement + # this hook must prove their explicit, versioned contract before setup + # applies runtime parallelism or constructs any model layers. + validate_model_contract = getattr(model_cfg, "validate_model_contract", None) + if callable(validate_model_contract): + validate_model_contract() + # Apply parallelism settings _apply_parallelism_config(model_cfg, config) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index cba78d3acd..5f615c400a 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -128,16 +128,53 @@ 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 _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 # Classes with @ray.remote can't be inherited from, so we split the implementation out. @@ -441,6 +478,12 @@ 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" # vars used for refit ## will be initialized in prepare_refit_info @@ -641,6 +684,7 @@ 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, ) # Track total microbatches for MoE aux-loss averaging total_num_microbatches += int(num_microbatches) @@ -1456,6 +1500,7 @@ 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, ) use_fused_linear_logprobs = self.cfg["megatron_cfg"].get( @@ -1673,6 +1718,7 @@ 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, ) list_of_outputs = megatron_forward_backward( @@ -1817,17 +1863,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, - } - 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/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 8f1dce267c..74822398d7 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 a94412222a..a8cf28a3ea 100644 --- a/tests/unit/data/test_multimodal_dict.py +++ b/tests/unit/data/test_multimodal_dict.py @@ -349,3 +349,44 @@ 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.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) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 2fd8327811..254a9c2710 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -259,3 +259,82 @@ def test_load_mtp_weights_from_disk_raises_when_mtp_weights_missing( with pytest.raises(ValueError, match="No MTP layer weights"): ext.load_mtp_weights_from_disk(str(model_dir)) ext._load_draft_weights.assert_not_called() + + +class _TinyRadioVisionModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(initializer_factor=1.0) + self.ls1 = torch.nn.Parameter(torch.full((4,), 0.001)) + self.ls2 = torch.nn.Parameter(torch.full((4,), -0.001)) + self.projection = torch.nn.Parameter(torch.full((4,), 0.25)) + + +def _make_extension_with_radio(architecture): + from nemo_rl.models.generation.vllm.vllm_backend import ( + VllmInternalWorkerExtension, + ) + + ext = VllmInternalWorkerExtension.__new__(VllmInternalWorkerExtension) + ext.model_runner = SimpleNamespace( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace(architectures=[architecture]) + ), + model=SimpleNamespace(vision_model=_TinyRadioVisionModel()), + ) + return ext + + +@pytest.mark.vllm +def test_initialize_folded_nemotron_radio_layerscale_while_awake(): + ext = _make_extension_with_radio("NemotronH_Nano_Omni_Reasoning_V3") + + initialized = ext._initialize_nemotron_omni_radio_layerscale() + + vision_model = ext.model_runner.model.vision_model + assert initialized == 2 + assert torch.equal(vision_model.ls1, torch.ones_like(vision_model.ls1)) + assert torch.equal(vision_model.ls2, torch.ones_like(vision_model.ls2)) + assert torch.equal( + vision_model.projection, torch.full_like(vision_model.projection, 0.25) + ) + + +@pytest.mark.vllm +def test_prepare_refit_info_does_not_mutate_folded_layerscale(): + ext = _make_extension_with_radio("NemotronH_Nano_Omni_Reasoning_V3") + + state_dict_info = {"language_model.weight": ((4, 4), torch.bfloat16)} + ext.prepare_refit_info(state_dict_info) + + vision_model = ext.model_runner.model.vision_model + assert ext.state_dict_info is state_dict_info + assert torch.equal(vision_model.ls1, torch.full_like(vision_model.ls1, 0.001)) + assert torch.equal(vision_model.ls2, torch.full_like(vision_model.ls2, -0.001)) + + +@pytest.mark.vllm +def test_prepare_refit_info_rejects_explicit_nemotron_radio_layerscale(): + ext = _make_extension_with_radio("NemotronH_Super_Omni_Reasoning_V3") + + with pytest.raises(RuntimeError, match="explicit RADIO LayerScale"): + ext.prepare_refit_info( + { + "vision_model.radio_model.model.blocks.0.ls1": ( + (4,), + torch.bfloat16, + ) + } + ) + + +@pytest.mark.vllm +def test_prepare_refit_info_leaves_non_nemotron_model_unchanged(): + ext = _make_extension_with_radio("SomeOtherArchitecture") + + ext.prepare_refit_info({}) + assert ext._initialize_nemotron_omni_radio_layerscale() == 0 + + vision_model = ext.model_runner.model.vision_model + assert torch.equal(vision_model.ls1, torch.full_like(vision_model.ls1, 0.001)) + assert torch.equal(vision_model.ls2, torch.full_like(vision_model.ls2, -0.001)) diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index bb4f72feda..4582aec55c 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -37,6 +37,8 @@ ) 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 new file mode 100644 index 0000000000..78803028f5 --- /dev/null +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -0,0 +1,116 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +from nemo_rl.models.generation.vllm import patches + + +class _Logger: + def __init__(self): + self.info_messages = [] + self.warning_messages = [] + + def info(self, message, *args): + self.info_messages.append(message % args if args else message) + + def warning(self, message, *args): + self.warning_messages.append(message % args if args else message) + + +def _stock_sampling_params_source() -> str: + return """import copy +import json as json_mod +from dataclasses import field + +_SAMPLING_EPS = 1e-5 +_MAX_TEMP = 1e-2 + +class SamplingParams: + def update_from_tokenizer(self, tokenizer: TokenizerLike) -> None: + if not self.bad_words: + return + self._bad_words_token_ids = [] + for bad_word in self.bad_words: + # To prohibit words both at the beginning + # and in the middle of text + # (related to add_prefix_space tokenizer parameter) + for add_prefix_space in [False, True]: + prefix = " " if add_prefix_space else "" + prompt = prefix + bad_word.lstrip() + prompt_token_ids = tokenizer.encode( + text=prompt, add_special_tokens=False + ) + + # If no space at the beginning + # or if prefix space produces a new word token + if (not add_prefix_space) or ( + add_prefix_space + and prompt_token_ids[0] != self._bad_words_token_ids[-1][0] + and len(prompt_token_ids) == len(self._bad_words_token_ids[-1]) + ): + self._bad_words_token_ids.append(prompt_token_ids) + + invalid_token_ids = [ + token_id + for bad_words_token_ids in self._bad_words_token_ids + for token_id in bad_words_token_ids + if token_id < 0 or token_id > tokenizer.max_token_id + ] + if len(invalid_token_ids) > 0: + raise VLLMValidationError( + f"The model vocabulary size is {tokenizer.max_token_id + 1}," + f" but the following tokens" + f" were specified as bad: {invalid_token_ids}." + f" All token id values should be integers satisfying:" + f" 0 <= token_id <= {tokenizer.max_token_id}.", + parameter="bad_words", + value=self.bad_words, + ) +""" + + +def test_bad_words_patch_adds_bounded_thread_safe_cache(tmp_path, monkeypatch): + sampling_params = tmp_path / "sampling_params.py" + sampling_params.write_text(_stock_sampling_params_source()) + monkeypatch.setattr(patches, "_get_vllm_file", lambda _: str(sampling_params)) + logger = _Logger() + + patches._patch_vllm_bad_words_tokenization_cache(logger) + patched = sampling_params.read_text() + + assert "import threading" in patched + assert "_BAD_WORDS_TOKEN_IDS_CACHE_MAX_ENTRIES = 1024" in patched + assert "with _BAD_WORDS_TOKEN_IDS_CACHE_LOCK:" in patched + assert "def _tokenize_bad_words" in patched + assert "return bad_words_token_ids" in patched + assert not logger.warning_messages + + patches._patch_vllm_bad_words_tokenization_cache(logger) + assert sampling_params.read_text() == patched + assert ( + logger.info_messages[-1] == "vLLM bad_words tokenization cache already applied." + ) + + +def test_bad_words_patch_leaves_unknown_vllm_source_untouched(tmp_path, monkeypatch): + sampling_params = Path(tmp_path) / "sampling_params.py" + sampling_params.write_text("# newer vLLM source\n") + monkeypatch.setattr(patches, "_get_vllm_file", lambda _: str(sampling_params)) + logger = _Logger() + + patches._patch_vllm_bad_words_tokenization_cache(logger) + + assert sampling_params.read_text() == "# newer vLLM source\n" + assert "expected vLLM 0.20.0 source shape" in logger.warning_messages[-1] diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py index e72431ba11..c74b0a149b 100644 --- a/tests/unit/models/megatron/test_megatron_data.py +++ b/tests/unit/models/megatron/test_megatron_data.py @@ -566,13 +566,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 +593,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 02c634e3fe..2652d635e8 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. @@ -32,6 +32,64 @@ import torch +@pytest.mark.mcore +class TestNemotronOmniCheckpointContract: + """Tests for distinguishing raw Option 2 and legacy LLaVA checkpoints.""" + + def test_non_omni_checkpoint_does_not_read_tensor_metadata(self): + from nemo_rl.models.megatron.setup import ( + _validate_nemotron_omni_megatron_lm_layout, + ) + + hf_config = SimpleNamespace(architectures=["LlamaForCausalLM"]) + with patch( + "megatron.core.dist_checkpointing.load_tensors_metadata" + ) as load_metadata: + _validate_nemotron_omni_megatron_lm_layout(hf_config, "/unused") + + load_metadata.assert_not_called() + + def test_canonical_omni_checkpoint_layout_is_accepted(self): + from nemo_rl.models.megatron.setup import ( + _validate_nemotron_omni_megatron_lm_layout, + ) + + hf_config = SimpleNamespace(architectures=["NemotronH_Nano_Omni_Reasoning_V3"]) + metadata = { + "language_model.embedding.word_embeddings.weight": object(), + "vision_model.decoder.layers.0.self_attention.linear_qkv.weight": object(), + } + with patch( + "megatron.core.dist_checkpointing.load_tensors_metadata", + return_value=metadata, + ): + _validate_nemotron_omni_megatron_lm_layout( + hf_config, "/checkpoints/iter_0000001" + ) + + def test_legacy_llava_checkpoint_layout_is_rejected(self): + from nemo_rl.models.megatron.setup import ( + _validate_nemotron_omni_megatron_lm_layout, + ) + + hf_config = SimpleNamespace(architectures=["NemotronH_Super_Omni_Reasoning_V3"]) + metadata = { + "model.llava_model.language_model.embedding.word_embeddings.weight": object(), + "model.llava_model.vision_model.decoder.layers.0.self_attention.linear_qkv.weight": object(), + } + with patch( + "megatron.core.dist_checkpointing.load_tensors_metadata", + return_value=metadata, + ): + with pytest.raises( + RuntimeError, + match="tensor namespace contains the legacy 'llava_model' wrapper", + ): + _validate_nemotron_omni_megatron_lm_layout( + hf_config, "/checkpoints/iter_0002821" + ) + + @pytest.mark.mcore class TestValidateModelPaths: """Tests for validate_model_paths function.""" 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 0000000000..b7c0e81a40 --- /dev/null +++ b/tests/unit/models/megatron/test_nemotron_omni_model.py @@ -0,0 +1,295 @@ +# 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 +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 nemo_rl.distributed.model_utils import ( + from_parallel_logits_to_logprobs_packed_sequences, +) +from nemo_rl.models.megatron.data import process_microbatch + +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(): + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=2, + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + provider = _TinyOmniProvider( + freeze_language_model=True, + tensor_model_parallel_size=1, + context_parallel_size=2, + sequence_parallel=False, + ) + 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, + delegate_pack_to_model=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=0, + vocab_end_index=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 + + +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) diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 793cc52bd9..b51834483f 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -45,6 +45,77 @@ 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_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/uv.lock b/uv.lock index 976bc8f50b..b2eb99c621 100644 --- a/uv.lock +++ b/uv.lock @@ -2599,7 +2599,7 @@ requires-dist = [ { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" }, { name = "hypercorn", marker = "extra == 'dev'" }, - { name = "mamba-ssm", marker = "extra == 'ssm'", specifier = "~=2.2" }, + { name = "mamba-ssm", marker = "extra == 'ssm'", git = "https://github.com/state-spaces/mamba.git?rev=0048fbf2e7b2f214dcbe703ea3dec2b9647595e1" }, { name = "megatron-energon", extras = ["av-decode"], marker = "extra == 'dev'", specifier = "~=7.0" }, { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.50" }, { name = "numpy" }, @@ -2620,7 +2620,7 @@ requires-dist = [ { name = "tiktoken", marker = "extra == 'training'" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", marker = "extra == 'dev'" }, - { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4220403e831d29e93868f7793693ea83f6b8b05b" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=b9d690e042b1c4e455214e7dab65d6d3512c05d6" }, { name = "transformers", marker = "extra == 'mlm'" }, { name = "transformers", marker = "extra == 'training'" }, { name = "wandb", marker = "extra == 'mlm'" }, @@ -2662,7 +2662,7 @@ linting = [ ] no-pypi-wheels = [ { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, - { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=9edee0c022cd0938148a18e334203b0aab43aa19" }, + { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] test = [ { name = "coverage" },