diff --git a/.github/workflows/config/.secrets.baseline b/.github/workflows/config/.secrets.baseline index ae1e36e5fe..b889b24a82 100644 --- a/.github/workflows/config/.secrets.baseline +++ b/.github/workflows/config/.secrets.baseline @@ -354,6 +354,29 @@ "is_verified": false, "line_number": 670 } + ], + "nemo_curator/models/nemotron_h_vl.py": [ + { + "type": "Hex High Entropy String", + "filename": "nemo_curator/models/nemotron_h_vl.py", + "hashed_secret": "7e44041c8ec081054af6ccf0758efa3771e99123", + "is_verified": false, + "line_number": 52 + }, + { + "type": "Hex High Entropy String", + "filename": "nemo_curator/models/nemotron_h_vl.py", + "hashed_secret": "f1a89dd2b19e1d30ac4fc1a8d1f4993796c9d1e6", + "is_verified": false, + "line_number": 54 + }, + { + "type": "Hex High Entropy String", + "filename": "nemo_curator/models/nemotron_h_vl.py", + "hashed_secret": "d7f11e2581dd9d25522592e7cfb9346a749f51f3", + "is_verified": false, + "line_number": 55 + } ] }, "generated_at": "2026-03-17T21:09:47Z" diff --git a/nemo_curator/models/cosmos_embed1.py b/nemo_curator/models/cosmos_embed1.py index 758c0021f3..c3771944bc 100644 --- a/nemo_curator/models/cosmos_embed1.py +++ b/nemo_curator/models/cosmos_embed1.py @@ -32,9 +32,9 @@ } COSMOS_EMBED1_MODEL_REVISION_INFO: Final = { - "224p": "85f5627", - "336p": "5d8309d", - "448p": "9f4ff4d", + "224p": "787e0b9", + "336p": "0e8a28f", + "448p": "f60ec73", } @@ -201,7 +201,9 @@ def download_weights_on_node(cls, model_dir: str, variant: Literal["224p", "336p logger.info(f"CosmosEmbed1 {variant} weights downloaded to: {model_dir_path}") @classmethod - def download_processor_config_on_node(cls, model_dir: str, variant: Literal["224p", "336p", "448p"] = "336p") -> None: + def download_processor_config_on_node( + cls, model_dir: str, variant: Literal["224p", "336p", "448p"] = "336p" + ) -> None: """Download the processor config for the CosmosEmbed1 model on the node.""" model_dir_path = Path(model_dir) / _COSMOS_EMBED1_VARIANTS_INFO[variant] model_dir_path.mkdir(parents=True, exist_ok=True) diff --git a/nemo_curator/models/nemotron_h_vl.py b/nemo_curator/models/nemotron_h_vl.py new file mode 100644 index 0000000000..fb9d5bced1 --- /dev/null +++ b/nemo_curator/models/nemotron_h_vl.py @@ -0,0 +1,274 @@ +# Copyright (c) 2025, 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 typing import Any, Final, Literal + +from loguru import logger + +from nemo_curator.models.base import ModelInterface +from nemo_curator.utils import grouping +from nemo_curator.utils.hf_download_utils import download_model_from_hf + +# Constants for prompt processing +VIDEO_TAG_SPLIT_MAX = 1 +EXPECTED_VIDEO_TAG_PARTS = 2 + +try: + from vllm import LLM, SamplingParams + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + + class LLM: + pass + + class SamplingParams: + pass + + +# HuggingFace model IDs for Nemotron Nano V2 VL variants +# Available variants: BF16 (default), FP8, NVFP4-QAD +_NEMOTRON_VARIANTS_INFO: Final = { + "nemotron": "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", # Default BF16 variant + "nemotron-bf16": "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", + "nemotron-fp8": "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-FP8", + "nemotron-nvfp4": "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-NVFP4-QAD", +} + +_NEMOTRON_REVISION_INFO: Final = { + "nemotron": "5d250e2e111dc5e1434131bdf3d590c27a878ade", # BF16 default + "nemotron-bf16": "5d250e2e111dc5e1434131bdf3d590c27a878ade", + "nemotron-fp8": "7394488badb786e1decc0e00e308de1cab9560e6", + "nemotron-nvfp4": "b8d3c170d9ee3a078917ef9bfd508eff988d6de7", +} + +NemotronVariant = Literal["nemotron", "nemotron-bf16", "nemotron-fp8", "nemotron-nvfp4"] + + +class NemotronHVL(ModelInterface): + """NemotronH hybrid Mamba-Attention VLM for video captioning. + + Supports multiple checkpoint variants from HuggingFace: + - nemotron / nemotron-bf16: BF16 precision (default) + - nemotron-fp8: FP8 quantized + - nemotron-nvfp4: NVFP4 quantized + + Models are automatically downloaded from HuggingFace on first use. + """ + + def __init__( # noqa: PLR0913 + self, + model_dir: str, + model_variant: NemotronVariant = "nemotron", + caption_batch_size: int = 8, + max_output_tokens: int = 512, + stage2_prompt_text: str | None = None, + verbose: bool = False, + ): + """Initialize NemotronHVL model. + + Args: + model_dir: Base directory for model weights. Models will be downloaded + to subdirectories named after the HuggingFace model ID. + model_variant: Model variant to use. Options: + - "nemotron" or "nemotron-bf16": BF16 precision (default) + - "nemotron-fp8": FP8 quantized + - "nemotron-nvfp4": NVFP4 quantized + caption_batch_size: Batch size for caption generation. + max_output_tokens: Maximum number of tokens to generate. + stage2_prompt_text: Optional prompt text for stage 2 caption refinement. + verbose: Whether to enable verbose logging. + """ + # Normalize variant name - treat "nemotron" as "nemotron-bf16" + if model_variant == "nemotron": + self._normalized_variant: NemotronVariant = "nemotron-bf16" + else: + self._normalized_variant = model_variant # type: ignore[assignment] + + if self._normalized_variant not in _NEMOTRON_VARIANTS_INFO: + valid_variants = ", ".join(_NEMOTRON_VARIANTS_INFO.keys()) + msg = f"Invalid model_variant: {model_variant}. Valid options are: {valid_variants}" + raise ValueError(msg) + + self.model_dir = model_dir + self.model_variant = model_variant + self.caption_batch_size = caption_batch_size + self.max_output_tokens = max_output_tokens + self.stage2_prompt = stage2_prompt_text if stage2_prompt_text else "Please refine this caption: " + self.verbose = verbose + + # Set weight file path using HuggingFace model ID + self._hf_model_id = _NEMOTRON_VARIANTS_INFO[self._normalized_variant] + self.weight_file = str(Path(model_dir) / self._hf_model_id) + + @property + def model_id_names(self) -> list[str]: + """Return HuggingFace model ID for the selected variant.""" + return [self._hf_model_id] + + def setup(self) -> None: + if not VLLM_AVAILABLE: + msg = "vllm is required for NemotronHVL but is not installed. Please install vllm: pip install vllm" + raise ImportError(msg) + + # Determine quantization and dtype based on variant + # See: https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-FP8 + # https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-NVFP4-QAD + quantization = None + dtype = "bfloat16" # BF16 variant requires explicit dtype + if self._normalized_variant == "nemotron-fp8": + quantization = "modelopt" # vllm serve uses: --quantization modelopt + dtype = "auto" # FP8 determines dtype from quantization + elif self._normalized_variant == "nemotron-nvfp4": + quantization = "modelopt_fp4" # vllm serve uses: --quantization modelopt_fp4 + dtype = "auto" # FP4 determines dtype from quantization + + self.model = LLM( + model=self.weight_file, + trust_remote_code=True, + dtype=dtype, + tensor_parallel_size=1, + gpu_memory_utilization=0.9, + max_model_len=32768, + limit_mm_per_prompt={"video": 1}, + quantization=quantization, + video_pruning_rate=0, # Disable video pruning + ) + + self.sampling_params = SamplingParams( + temperature=0.6, + max_tokens=self.max_output_tokens, + top_p=0.95, + stop=["", "<|endoftext|>", "", ""], + ) + + logger.info( + f"NemotronHVL initialized: variant={self.model_variant}, " + f"quantization={quantization}, TP=1, GPU_util=0.9, max_len=32768" + ) + + def _refine_caption_prompt(self, original_prompt: str, refinement_text: str) -> str: + """Create a refined prompt for stage 2 captioning.""" + if "