diff --git a/examples/models/nemotron/nemotron_3_omni/README.md b/examples/models/nemotron/nemotron_3_omni/README.md index c29306b118..b82a922a7c 100644 --- a/examples/models/nemotron/nemotron_3_omni/README.md +++ b/examples/models/nemotron/nemotron_3_omni/README.md @@ -96,7 +96,7 @@ the container), `uv` is resolving against a stale environment — re-run ## Checkpoint Conversion -[conversion.sh](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_omni/examples/models/vlm/nemotron_3_omni/conversion.sh) covers HF → Megatron import, Megatron → HF +[conversion.sh](conversion.sh) covers HF → Megatron import, Megatron → HF export, and a multi-GPU HF↔Megatron round-trip verification. - **Import** writes `iter_0000000/`, `latest_train_state.pt`, and @@ -107,6 +107,9 @@ export, and a multi-GPU HF↔Megatron round-trip verification. tensors (regenerated from config on the HF side): `sound_encoder.encoder.feature_extractor.featurizer.{fb,window}` and `vision_model.radio_model.input_conditioner.{norm_mean,norm_std}`. + `--trust-remote-code` is also required for export because the exporter + loads the HF config, which references the custom modeling module shipped + with `NemotronH_Nano_Omni_Reasoning_V3`. - **Round-trip** loads HF → Megatron (TP=2, EP=2) and re-exports back to HF, diffing every tensor; all weights should match (✅) and the same 4 expected-missing tensors are reported on re-export. The re-exported HF @@ -130,8 +133,8 @@ bash examples/models/nemotron/nemotron_3_omni/conversion.sh ## Inference -[inference.sh](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_omni/examples/models/vlm/nemotron_3_omni/inference.sh) drives -`examples/conversion/hf_to_megatron_generate_nemotron_omni.py` over the four +[inference.sh](inference.sh) drives +`examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py` over the four modality combinations exercised by the model: | # | Modality | GPUs | Parallelism | @@ -144,6 +147,14 @@ modality combinations exercised by the model: The default assets are pulled automatically from the public HF model card ([`nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16`](https://huggingface.co/nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16/tree/main/media)) on the first run — `curl` must be available. + +> **Video prerequisite:** the video paths (rows 2 and 4) sample frames via +> [`decord`](https://github.com/dmlc/decord), which is not pulled in by any +> pyproject extra. Install it before running those modes: +> +> ```bash +> uv pip install decord +> ``` Override `IMAGE_PATH` / `VIDEO_PATH` / `AUDIO_PATH` with your own assets to use different inputs; omit `--megatron_model_path` (set `MEGATRON_PATH=""`) to convert HF → Megatron on the fly instead of reusing the imported @@ -214,8 +225,8 @@ base config. Recipe base: `nemotron_omni_cord_v2_*_config` in | Mode | Script | Recipe | |---|---|---| -| Full SFT | [slurm_sft_cord_v2.sh](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_omni/examples/models/vlm/nemotron_3_omni/slurm_sft_cord_v2.sh) | `nemotron_omni_cord_v2_sft_config` | -| LoRA | [slurm_peft_cord_v2.sh](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_omni/examples/models/vlm/nemotron_3_omni/slurm_peft_cord_v2.sh) | `nemotron_omni_cord_v2_peft_config` | +| Full SFT | [slurm_sft_cord_v2.sh](slurm_sft_cord_v2.sh) | `nemotron_omni_cord_v2_sft_config` | +| LoRA | [slurm_peft_cord_v2.sh](slurm_peft_cord_v2.sh) | `nemotron_omni_cord_v2_peft_config` | Parallelism (both): TP=2, EP=8, CP=1, MBS=2, GBS=16, packed sequences, selective recompute. LoRA targets `linear_qkv`, `linear_proj`, `in_proj`, @@ -235,17 +246,18 @@ embedder: frames are fused in pairs (`temporal_patch_dim=2`, `separate_video_embedder=True`) and audio is fed through the Parakeet encoder. Recipe base: `nemotron_omni_valor32k_*_config`. -Prepare the Energon shards once: +Prepare the Energon shards once. For the full walkthrough, see +[`tutorials/data/valor32k-avqa/data-preparation.md`](../../../../tutorials/data/valor32k-avqa/data-preparation.md). ```bash -uv run python examples/models/nemotron/nemotron_3_omni/data/build_valor32k_avqa_shards.py \ +uv run python tutorials/data/valor32k-avqa/build_valor32k_avqa_shards.py \ --output_dir ${WORKSPACE}/datasets/valor32k_avqa ``` | Mode | Script | Recipe | |---|---|---| -| Full SFT | [slurm_sft_valor32k_avqa.sh](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_omni/examples/models/vlm/nemotron_3_omni/slurm_sft_valor32k_avqa.sh) | `nemotron_omni_valor32k_sft_config` | -| LoRA | [slurm_peft_valor32k_avqa.sh](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_omni/examples/models/vlm/nemotron_3_omni/slurm_peft_valor32k_avqa.sh) | `nemotron_omni_valor32k_peft_config` | +| Full SFT | [slurm_sft_valor32k_avqa.sh](slurm_sft_valor32k_avqa.sh) | `nemotron_omni_valor32k_sft_config` | +| LoRA | [slurm_peft_valor32k_avqa.sh](slurm_peft_valor32k_avqa.sh) | `nemotron_omni_valor32k_peft_config` | Parallelism (both): TP=2, EP=8, CP=1, MBS=2, packed sequences, selective recompute. SFT uses GBS=16 and the recipe-default LR; LoRA uses GBS=64 and @@ -268,8 +280,8 @@ finetuned Megatron checkpoint on the same datasets used for training: | Dataset | Script | Output | |---|---|---| -| CORD-V2 | [cord_v2_inference.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_omni/examples/models/vlm/nemotron_3_omni/cord_v2_inference.py) | JSON of `{prompt, gold, prediction}` per sample plus image bytes for eyeballing | -| VALOR32K-AVQA | [valor32k_avqa_inference.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_omni/examples/models/vlm/nemotron_3_omni/valor32k_avqa_inference.py) | Per-sample predictions and an aggregate multiple-choice accuracy | +| CORD-V2 | [cord_v2_inference.py](cord_v2_inference.py) | JSON of `{prompt, gold, prediction}` per sample plus image bytes for eyeballing | +| VALOR32K-AVQA | [valor32k_avqa_inference.py](valor32k_avqa_inference.py) | Per-sample predictions and an aggregate multiple-choice accuracy | Example invocations (8 GPUs, single node). The slurm scripts tag `OUTPUT_DIR` with the run config (`__`, where diff --git a/examples/models/nemotron/nemotron_3_omni/conversion.sh b/examples/models/nemotron/nemotron_3_omni/conversion.sh new file mode 100755 index 0000000000..bbfc8e6431 --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/conversion.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# 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. + +# ============================================================================== +# Nemotron-3 Nano Omni - Checkpoint conversion (HF <-> Megatron) + roundtrip +# ============================================================================== + +set -xeuo pipefail + +WORKSPACE=${WORKSPACE:-/workspace} +HF_MODEL_ID=${HF_MODEL_ID:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16} +MODEL_NAME=$(basename "$HF_MODEL_ID") + +MEGATRON_PATH=${WORKSPACE}/models/${MODEL_NAME} +HF_EXPORT_PATH=${WORKSPACE}/models/${MODEL_NAME}-hf-export + +# Import HF -> Megatron +uv run python examples/conversion/convert_checkpoints.py import \ + --hf-model "$HF_MODEL_ID" \ + --megatron-path "$MEGATRON_PATH" \ + --trust-remote-code + +# Export Megatron -> HF (--not-strict allows 4 expected-missing tensors that +# are regenerated from config on the HF side: sound_encoder featurizer fb/window +# and vision_model input_conditioner norm_mean/norm_std) +uv run python examples/conversion/convert_checkpoints.py export \ + --hf-model "$HF_MODEL_ID" \ + --megatron-path "$MEGATRON_PATH" \ + --hf-path "$HF_EXPORT_PATH" \ + --not-strict + +# Round-trip validation (multi-GPU): HF -> Megatron -> HF and compare weights +uv run python -m torch.distributed.run --nproc_per_node=4 \ + examples/conversion/hf_megatron_roundtrip_multi_gpu.py \ + --hf-model-id "$HF_MODEL_ID" \ + --trust-remote-code \ + --tp 2 --ep 2 --not-strict diff --git a/examples/models/nemotron/nemotron_3_omni/cord_v2_inference.py b/examples/models/nemotron/nemotron_3_omni/cord_v2_inference.py new file mode 100644 index 0000000000..2af4936024 --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/cord_v2_inference.py @@ -0,0 +1,384 @@ +# 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 + +"""CORD-V2 inference for a Nemotron Omni Megatron checkpoint. + +Loads N samples from the CORD-V2 HF dataset, runs the Megatron model on each +(image + "Describe this image." prompt), and writes a JSON file containing the +input prompt, the gold response (parsed receipt JSON as a token string), and +the model's generated prediction for each sample. Image bytes are also saved +to disk so the outputs can be eyeballed later. + +Vision backbone: dynamic resolution, temporal_patch_dim=1, separate_video_embedder=True +(matches `nemotron_omni_cord_v2_sft_config` with the updated dynamic-resolution training). + +Usage: + uv run torchrun --nproc-per-node=8 examples/models/nemotron/nemotron_3_omni/cord_v2_inference.py \ + --hf_model_path /chcui/pretrained_models/Nemotron-3-Nano-Omni-30B-A3B-Reasoning \ + --megatron_model_path /path/to/cord_v2/checkpoints \ + --tp 4 --ep 2 \ + --max_samples 100 \ + --output /chcui/mbridge_home/inference_results/cord_v2_sft_rerun.json +""" + +import argparse +import json +from pathlib import Path +from typing import Optional + +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.pipeline_parallel.schedules import get_forward_backward_func +from transformers import AutoProcessor, AutoTokenizer + +from megatron.bridge import AutoBridge +from megatron.bridge.data.vlm_datasets.hf_dataset_makers import make_cord_v2_dataset +from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import adjust_image_tokens +from megatron.bridge.utils.common_utils import get_last_rank, print_rank_0 + + +_VISION_PATCH_DIM = 16 + + +def _build_vision_packed_seq_params(imgs_sizes: Optional[torch.Tensor]) -> Optional[PackedSeqParams]: + """PackedSeqParams from per-image (H, W) sizes — mirrors nemotron_omni_step.""" + if imgs_sizes is None or imgs_sizes.numel() == 0: + return None + sizes = imgs_sizes.tolist() if torch.is_tensor(imgs_sizes) else list(imgs_sizes) + seq_lens = [(int(h) // _VISION_PATCH_DIM) * (int(w) // _VISION_PATCH_DIM) for h, w in sizes] + cu = [0] + for sl in seq_lens: + cu.append(cu[-1] + sl) + device = imgs_sizes.device if torch.is_tensor(imgs_sizes) else torch.device("cpu") + cu_tensor = torch.tensor(cu, dtype=torch.int32, device=device) + max_len = torch.tensor(max(seq_lens) if seq_lens else 0, dtype=torch.int32, device=device) + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_tensor, + cu_seqlens_kv=cu_tensor, + max_seqlen_q=max_len, + max_seqlen_kv=max_len, + ) + + +class SingleBatchIterator: + """Iterator that yields one prepared inference batch.""" + + def __init__(self, input_ids, position_ids, attention_mask, **kwargs): + self.batch = dict(tokens=input_ids, position_ids=position_ids, attention_mask=attention_mask) + for key in ("images", "imgs_sizes", "num_frames", "vision_packed_seq_params"): + if kwargs.get(key) is not None: + self.batch[key] = kwargs[key] + self._yielded = False + + def __iter__(self): + return self + + def __next__(self): + if self._yielded: + raise StopIteration + self._yielded = True + return self.batch + + +def vlm_forward_step(data_iterator, model, **_): + """Run one VLM forward pass for text generation.""" + + batch = next(data_iterator) + forward_args = { + "input_ids": batch["tokens"], + "position_ids": batch["position_ids"], + "attention_mask": batch.get("attention_mask", None), + } + if "images" in batch: + forward_args["images"] = batch["images"] + else: + forward_args["images"] = torch.tensor([], dtype=torch.bfloat16, device=batch["tokens"].device).reshape(0, 0, 0) + for key in ("imgs_sizes", "num_frames", "vision_packed_seq_params"): + if key in batch: + forward_args[key] = batch[key] + + def loss_func(x, **_): + return x + + output = model(**forward_args) + if isinstance(output, tuple): + output = output[0] + return output, loss_func + + +def prepare_image_sample(tokenizer, processor, image, prompt, system_prompt=None): + """Build input_ids and dynamic-resolution image tensors for a single image + text prompt.""" + text_content = f"\n{prompt}" + messages = [{"role": "user", "content": text_content}] + if system_prompt: + messages.insert(0, {"role": "system", "content": system_prompt}) + + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + inputs = processor(text=[text], images=[image], return_tensors="pt") + + input_ids = inputs.input_ids + + # Adjust image tokens: collapse ...... to single per tile. + img_start_id = tokenizer.convert_tokens_to_ids("") + img_end_id = tokenizer.convert_tokens_to_ids("") + num_patches = torch.ones(1, dtype=torch.long) # 1 tile per image + if img_start_id != tokenizer.unk_token_id and (input_ids == img_start_id).any(): + input_ids = adjust_image_tokens(input_ids, num_patches, img_start_id, img_end_id) + + # Patchify [1, 3, H, W] → [1, num_patches, 3*P*P] (mirrors collate.py dynamic-res path). + pv = inputs.pixel_values # [1, 3, H, W], float32, already normalized + P = _VISION_PATCH_DIM + _, C, H, W = pv.shape + py, px = H // P, W // P + pv_patched = ( + pv[0] + .reshape(3, py, P, px, P) + .permute(1, 3, 0, 2, 4) + .reshape(py * px, 3 * P * P) + .unsqueeze(0) + .contiguous() + .bfloat16() + ) + imgs_sizes = torch.tensor([[H, W]], dtype=torch.long) + num_frames = torch.tensor([1], dtype=torch.long) + + return input_ids, pv_patched, imgs_sizes, num_frames + + +@torch.no_grad() +def generate(model, tokenizer, input_ids, images, imgs_sizes, num_frames, max_new_tokens=200): + """Generate tokens for one CORD-V2 sample.""" + + prompt_len = input_ids.size(1) + input_ids = input_ids.cuda() + images = images.cuda() + imgs_sizes = imgs_sizes.cuda() + num_frames = num_frames.cuda() + + position_ids = ( + torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device).unsqueeze(0).expand_as(input_ids) + ) + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + generated_ids = input_ids.clone() + stop_tokens = {tokenizer.eos_token_id} + + fwd_bwd = get_forward_backward_func() + for _ in range(max_new_tokens): + # Rebuild each iteration: RADIO mutates cu_seqlens_q in-place when inserting class tokens, + # so reusing the same object would cause cu_seqlens to grow by class_token_len each step. + vision_packed_seq_params = _build_vision_packed_seq_params(imgs_sizes) + iterator = SingleBatchIterator( + input_ids, + position_ids, + attention_mask, + images=images, + imgs_sizes=imgs_sizes, + num_frames=num_frames, + vision_packed_seq_params=vision_packed_seq_params, + ) + output = fwd_bwd( + forward_step_func=vlm_forward_step, + data_iterator=iterator, + model=model, + num_microbatches=1, + forward_only=True, + seq_length=input_ids.size(1), + micro_batch_size=1, + collect_non_loss_data=True, + ) + if isinstance(output, list) and len(output) > 0: + output = output[0] + if isinstance(output, tuple): + output = output[0] + + if parallel_state.is_pipeline_last_stage(): + world_size = parallel_state.get_tensor_model_parallel_world_size() + gathered = [torch.zeros_like(output) for _ in range(world_size)] + dist.all_gather(gathered, output, group=parallel_state.get_tensor_model_parallel_group()) + full = torch.cat(gathered, dim=2) + next_token_ids = torch.argmax(full[:, -1], dim=-1, keepdim=True) + else: + next_token_ids = torch.ones((1, 1), device=generated_ids.device, dtype=generated_ids.dtype) + + dist.broadcast(next_token_ids, get_last_rank()) + generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1) + input_ids = generated_ids + position_ids = ( + torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device) + .unsqueeze(0) + .expand_as(input_ids) + ) + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + if int(next_token_ids.item()) in stop_tokens: + break + + cleaned = tokenizer.decode(generated_ids[0, prompt_len:].tolist(), skip_special_tokens=True).strip() + full_text = tokenizer.decode(generated_ids[0].tolist(), skip_special_tokens=False) + return cleaned, full_text + + +def extract_gt_text_from_conversation(conv): + """Extract the ground-truth assistant text from a CORD-V2 conversation.""" + for turn in conv: + if turn.get("role") == "assistant": + content = turn.get("content") + if isinstance(content, list): + for item in content: + if item.get("type") == "text": + return item.get("text", "") + elif isinstance(content, str): + return content + return "" + + +def main(): + """Run CORD-V2 inference.""" + + parser = argparse.ArgumentParser(description="CORD-V2 inference for Nemotron Omni") + parser.add_argument("--hf_model_path", type=str, required=True) + parser.add_argument("--megatron_model_path", type=str, default=None) + parser.add_argument("--split", type=str, default="test") + parser.add_argument("--max_samples", type=int, default=100) + parser.add_argument("--max_new_tokens", type=int, default=300) + parser.add_argument("--tp", type=int, default=4) + parser.add_argument("--pp", type=int, default=1) + parser.add_argument("--ep", type=int, default=2) + parser.add_argument("--etp", type=int, default=1) + parser.add_argument( + "--output", type=str, required=True, help="Output JSON path (images saved next to it under cord_v2_images/)" + ) + parser.add_argument("--prompt", type=str, default="Describe this image.") + args = parser.parse_args() + + bridge = AutoBridge.from_hf_pretrained(args.hf_model_path, trust_remote_code=True) + model_provider = bridge.to_megatron_provider(load_weights=(args.megatron_model_path is None)) + model_provider.tensor_model_parallel_size = args.tp + model_provider.pipeline_model_parallel_size = args.pp + model_provider.expert_model_parallel_size = args.ep + model_provider.expert_tensor_parallel_size = args.etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.dynamic_resolution = True + model_provider.temporal_patch_dim = 1 + model_provider.separate_video_embedder = True + model_provider.temporal_ckpt_compat = True + model_provider.vision_class_token_len = 10 + model_provider.initialize_model_parallel(seed=0) + + if args.megatron_model_path: + print_rank_0(f"Loading Megatron checkpoint from {args.megatron_model_path}") + model = bridge.load_megatron_model( + args.megatron_model_path, + mp_overrides={ + "tensor_model_parallel_size": args.tp, + "pipeline_model_parallel_size": args.pp, + "expert_model_parallel_size": args.ep, + "expert_tensor_parallel_size": args.etp, + "pipeline_dtype": torch.bfloat16, + "dynamic_resolution": True, + "temporal_patch_dim": 1, + "separate_video_embedder": True, + "temporal_ckpt_compat": True, + "vision_class_token_len": 10, + }, + wrap_with_ddp=False, + ) + model = [m.cuda().eval() for m in model] + for m in model: + inner = m.module if hasattr(m, "module") else m + if hasattr(inner, "config"): + inner.config.grad_scale_func = None + if hasattr(inner, "llava_model") and hasattr(inner.llava_model, "config"): + inner.llava_model.config.grad_scale_func = None + else: + print_rank_0(f"Converting HF from {args.hf_model_path} on the fly") + model_provider.finalize() + model = model_provider.provide_distributed_model(wrap_with_ddp=False) + model = [m.cuda().bfloat16().eval() for m in model] + + tokenizer = AutoTokenizer.from_pretrained(args.hf_model_path, trust_remote_code=True) + processor = AutoProcessor.from_pretrained(args.hf_model_path, trust_remote_code=True) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + print_rank_0(f"Loading CORD-V2 split={args.split} ...") + examples = make_cord_v2_dataset(split=args.split) + n = min(args.max_samples, len(examples)) + print_rank_0(f"Running inference on {n}/{len(examples)} samples") + + output_path = Path(args.output) + images_dir = output_path.parent / "cord_v2_images" + if dist.get_rank() == 0: + output_path.parent.mkdir(parents=True, exist_ok=True) + images_dir.mkdir(parents=True, exist_ok=True) + + results = [] + for i in range(n): + ex = examples[i] + conv = ex["conversation"] + user_content = conv[0]["content"] + image = None + for item in user_content: + if item.get("type") == "image": + image = item.get("image") + gt_text = extract_gt_text_from_conversation(conv) + + if dist.get_rank() == 0: + img_path = images_dir / f"sample_{i:03d}.png" + try: + image.save(img_path) + except Exception as e: + print_rank_0(f"WARN: could not save sample image {i}: {e}") + + input_ids, pv_patched, imgs_sizes, num_frames = prepare_image_sample(tokenizer, processor, image, args.prompt) + + cleaned, prediction_full = generate( + model, + tokenizer, + input_ids, + pv_patched, + imgs_sizes, + num_frames, + max_new_tokens=args.max_new_tokens, + ) + + record = { + "sample_index": i, + "image_path": str(images_dir / f"sample_{i:03d}.png"), + "prompt": args.prompt, + "ground_truth": gt_text, + "prediction": cleaned, + "prediction_full_decode": prediction_full, + } + results.append(record) + print_rank_0(f"[{i + 1}/{n}] done (gt_len={len(gt_text)}, pred_len={len(cleaned)})") + + if dist.get_rank() == 0: + with open(output_path, "w") as f: + json.dump( + { + "checkpoint": args.megatron_model_path, + "hf_model": args.hf_model_path, + "split": args.split, + "num_samples": len(results), + "results": results, + }, + f, + indent=2, + default=str, + ) + print_rank_0(f"Saved results to {output_path}") + + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py b/examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py new file mode 100644 index 0000000000..23afdae1b7 --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py @@ -0,0 +1,965 @@ +# 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. + +""" +Nemotron Omni VL Generation Script. + +This script demonstrates inference with Nemotron Omni model +using Megatron-Bridge. Unlike the InternVL-based Nemotron VL models that rely on +qwen_vl_utils, this script uses the model's native HF processor with tokens. + +Vision backbone config is modality-dependent: + * Image: dynamic_resolution=True, temporal_patch_dim=1, + separate_video_embedder=False. Each HF-processor tile is pre-patchified + into [1, total_patches, 3*P*P] and passed through RADIO's packed + dynamic-resolution path (is_packed_dynamic_res=True in LlavaModel). The + ``imgs_sizes`` / ``vision_packed_seq_params`` tensors are built from the + per-tile shapes, and ``num_image_tiles`` is recomputed by LlavaModel from + RADIO output (256 tokens/tile after pixel_shuffle). + * Audio / text-only: dynamic_resolution=False, temporal_patch_dim=1. + * Video (and video+audio): dynamic_resolution=True, temporal_patch_dim=2, + separate_video_embedder=True, temporal_ckpt_compat=True so RADIO ViT + exercises the trained `video_embedder`. The video preprocessing mirrors + the SFT data pipeline (see `NemotronOmniTaskEncoder` with + `use_temporal_video_embedder=True`): frames are grouped in pairs, all frames + are pre-patchified into [1, total_patches, 3*P*P], and `imgs_sizes` + / `num_frames` / `vision_packed_seq_params` are plumbed through to LLaVAModel. + +Examples: + # Single image: + uv run python examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py \ + --hf_model_path="nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" \ + --image_path="https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/table.png" \ + --prompt="Describe this image." \ + --max_new_tokens 300 + + # Multiple images: + uv run python examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py \ + --hf_model_path="nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" \ + --image_path="https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/example1a.jpeg,https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/example1b.jpeg" \ + --prompt="Describe the two images in detail." \ + --max_new_tokens 300 + + # Video description: + uv run python examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py \ + --hf_model_path="nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" \ + --video_path="https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/demo.mp4" \ + --prompt="Describe what you see." \ + --max_new_tokens 300 + + # Audio transcription: + uv run python examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py \ + --hf_model_path="nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" \ + --audio_path="/path/to/audio.wav" \ + --prompt="Transcribe the audio." \ + --max_new_tokens 300 +""" + +import argparse +import math +from typing import Optional + +import requests +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.pipeline_parallel.schedules import get_forward_backward_func +from PIL import Image +from transformers import AutoProcessor, AutoTokenizer + +from megatron.bridge import AutoBridge +from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import adjust_image_tokens +from megatron.bridge.utils.common_utils import get_last_rank, print_rank_0 + + +# Must stay in sync with nemotron_omni_provider.py and the temporal SFT recipe +# `nemotron_omni_valor32k_sft_config` (patch_dim=16, 512x512 frames, +# tps=2). Reused across video / video+audio preprocessing for temporal inference. +_VIDEO_TEMPORAL_PATCH_SIZE = 2 +_VIDEO_FRAME_H = 512 +_VIDEO_FRAME_W = 512 +_VISION_PATCH_DIM = 16 +_VIDEO_FPS = 1 +_VIDEO_NFRAMES = 8 + +# CLIP / RADIO normalization constants (mirrors NemotronOmniTaskEncoder._patchify_frame) +_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) +_CLIP_STD = (0.26862954, 0.26130258, 0.27577711) + + +def _patchify_frame( + pil_img: Image.Image, target_h: int = _VIDEO_FRAME_H, target_w: int = _VIDEO_FRAME_W +) -> torch.Tensor: + """Resize + normalize a PIL frame and pack into [num_patches, 3*P*P] patches. + + Mirrors ``NemotronOmniTaskEncoder._patchify_frame`` exactly so inference-time + tensors have the same shape/distribution as SFT training tensors. + """ + from torchvision import transforms + + img = pil_img.convert("RGB").resize((target_w, target_h)) + tensor = transforms.ToTensor()(img) + mean = torch.tensor(_CLIP_MEAN).view(3, 1, 1) + std = torch.tensor(_CLIP_STD).view(3, 1, 1) + tensor = (tensor - mean) / std + P = _VISION_PATCH_DIM + py, px = target_h // P, target_w // P + patches = tensor.reshape(3, py, P, px, P).permute(1, 3, 0, 2, 4).reshape(py * px, 3 * P * P) + return patches + + +def _build_vision_packed_seq_params(imgs_sizes: Optional[torch.Tensor]) -> Optional[PackedSeqParams]: + """Build vision PackedSeqParams from pre-grouping per-frame (H, W). + + RADIO's dynamic-resolution + class-token path reads ``packed_seq_params.cu_seqlens_q`` + to insert class tokens at per-image boundaries, and ``_apply_temporal_grouping`` + rebuilds cu_seqlens after tubelet fusion. Copied from + ``megatron.bridge.training.nemotron_omni_step._build_vision_packed_seq_params``. + """ + if imgs_sizes is None or imgs_sizes.numel() == 0: + return None + sizes = imgs_sizes.tolist() if torch.is_tensor(imgs_sizes) else list(imgs_sizes) + seq_lens = [(int(h) // _VISION_PATCH_DIM) * (int(w) // _VISION_PATCH_DIM) for h, w in sizes] + cu = [0] + for sl in seq_lens: + cu.append(cu[-1] + sl) + device = imgs_sizes.device if torch.is_tensor(imgs_sizes) else torch.device("cpu") + cu_tensor = torch.tensor(cu, dtype=torch.int32, device=device) + max_len = torch.tensor(max(seq_lens) if seq_lens else 0, dtype=torch.int32, device=device) + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_tensor, + cu_seqlens_kv=cu_tensor, + max_seqlen_q=max_len, + max_seqlen_kv=max_len, + ) + + +def _fastconformer_output_length(mel_length: int, subsampling_factor: int = 8) -> int: + """Mirror BridgeSoundEncoder._compute_output_lengths for a scalar input. + + The FastConformer subsampling stack does Conv2D(kernel=3, stride=2, padding=1, + floor mode) repeated ``log2(subsampling_factor)`` times. The recurrence + reduces to ``L -> (L - 1) // 2 + 1`` per layer. + """ + num_layers = int(math.log2(subsampling_factor)) + length = int(mel_length) + for _ in range(num_layers): + length = (length - 1) // 2 + 1 + return length + + +def _align_sound_tokens(input_ids: torch.Tensor, sound_token_id: int, desired_count: int) -> torch.Tensor: + """Rewrite ``input_ids`` so the contiguous run of sound tokens has ``desired_count`` entries. + + The HF processor estimates sound-token count from ``len(waveform) // hop_length``, but + ``ParakeetFeatureExtractor`` (STFT with ``center=True``) produces one extra mel frame, + so ``BridgeSoundEncoder``'s output length can exceed the HF-inserted count by 1. We + realign here using the true encoder output length computed from ``sound_length``. + """ + assert input_ids.dim() == 2 and input_ids.size(0) == 1, "Expected a single-sample batch" + mask = input_ids[0] == sound_token_id + positions = torch.where(mask)[0] + if positions.numel() == 0: + raise ValueError(f"No sound tokens (id={sound_token_id}) found in input_ids") + current_count = int(positions.numel()) + if current_count == desired_count: + return input_ids + start = int(positions[0].item()) + end = int(positions[-1].item()) + 1 + if (end - start) != current_count: + raise ValueError("Sound tokens are not contiguous; cannot safely re-align") + prefix = input_ids[:, :start] + suffix = input_ids[:, end:] + middle = torch.full((1, desired_count), sound_token_id, dtype=input_ids.dtype, device=input_ids.device) + return torch.cat([prefix, middle, suffix], dim=1) + + +class SingleBatchIterator: + """Iterator that yields a single batch of data for text generation. + Required by the forward_backward_func function. + + This class creates an iterator that yields exactly one batch containing + input tokens, position IDs, attention mask, and optional vision inputs, + then raises StopIteration. Used for single-step inference in the forward pass. + """ + + def __init__(self, input_ids, position_ids, attention_mask, **kwargs): + self.batch = dict( + tokens=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + ) + + if kwargs.get("images", None) is not None: + self.batch["images"] = kwargs.get("images", None) + elif kwargs.get("pixel_values", None) is not None: + self.batch["pixel_values"] = kwargs.get("pixel_values", None) + + # Sound inputs + if kwargs.get("sound_clips", None) is not None: + self.batch["sound_clips"] = kwargs["sound_clips"] + if kwargs.get("sound_length", None) is not None: + self.batch["sound_length"] = kwargs["sound_length"] + + # Temporal video embedder inputs (dynamic-resolution pre-patchified path) + if kwargs.get("imgs_sizes", None) is not None: + self.batch["imgs_sizes"] = kwargs["imgs_sizes"] + if kwargs.get("num_frames", None) is not None: + self.batch["num_frames"] = kwargs["num_frames"] + if kwargs.get("vision_packed_seq_params", None) is not None: + self.batch["vision_packed_seq_params"] = kwargs["vision_packed_seq_params"] + + self._yielded = False + + def __iter__(self): + return self + + def __next__(self): + if self._yielded: + raise StopIteration + self._yielded = True + return self.batch + + +def vlm_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: + """Forward step function for vision-language generation. + Required by the forward_backward_func function. + + Extracts a batch from the data iterator and runs the model forward pass + with the provided input tokens, position IDs, attention mask, and vision inputs. + + Args: + data_iterator: Iterator providing batches of input data + model: The Megatron model to run forward pass on + **kwargs: Additional keyword arguments (unused) + + Returns: + Tuple of (model_output, loss_function) + """ + batch = next(data_iterator) + forward_args = { + "input_ids": batch["tokens"], + "position_ids": batch["position_ids"], + "attention_mask": batch.get("attention_mask", None), + } + + if "images" in batch: + forward_args["images"] = batch["images"] + elif "pixel_values" in batch: + forward_args["pixel_values"] = batch["pixel_values"] + + # LLaVAModel.forward() requires `images` even for audio-only inference + if "images" not in forward_args and "pixel_values" not in forward_args: + forward_args["images"] = torch.tensor([], dtype=torch.bfloat16, device=batch["tokens"].device).reshape(0, 0, 0) + + if "sound_clips" in batch: + forward_args["sound_clips"] = batch["sound_clips"] + if "sound_length" in batch: + forward_args["sound_length"] = batch["sound_length"] + + # Temporal video embedder plumbing (RADIO dynamic-resolution + separate + # video embedder path). Matches nemotron_omni_step.forward_step kwargs. + if "imgs_sizes" in batch: + forward_args["imgs_sizes"] = batch["imgs_sizes"] + if "num_frames" in batch: + forward_args["num_frames"] = batch["num_frames"] + if "vision_packed_seq_params" in batch: + forward_args["vision_packed_seq_params"] = batch["vision_packed_seq_params"] + + def loss_func(x, **kwargs): + return x + + output = model(**forward_args) + # LlavaModel returns (logits, loss_mask) tuple; pipeline expects a single tensor + if isinstance(output, tuple): + output = output[0] + return output, loss_func + + +def load_image(image_path: str) -> Image.Image: + """Load an image from URL or file path. + + Args: + image_path: URL or local file path to the image + + Returns: + PIL Image object + """ + if image_path.startswith(("http://", "https://")): + response = requests.get(image_path, stream=True) + response.raise_for_status() + return Image.open(response.raw) + else: + return Image.open(image_path) + + +def _patchify_pixel_values(pv: torch.Tensor, patch_dim: int = _VISION_PATCH_DIM): + """Pack [N, 3, H, W] image-tiles into [1, total_patches, 3*P*P] patches. + + ``N`` is typically 1 image * 1 tile (single-tile inference). When the HF + processor returns multiple rows (multi-image), they're concatenated along + the patch dim so RADIO's dynamic-resolution path sees a single packed + sequence that matches ``imgs_sizes``. + """ + P = patch_dim + patches_list = [] + sizes = [] + for i in range(pv.shape[0]): + _, H, W = pv[i].shape + py, px = H // P, W // P + p = pv[i : i + 1].reshape(1, 3, py, P, px, P).permute(0, 2, 4, 1, 3, 5).reshape(1, py * px, 3 * P * P) + patches_list.append(p) + sizes.append([H, W]) + packed = torch.cat(patches_list, dim=1) + return packed, sizes + + +def process_image_inputs( + tokenizer, processor, image_path: Optional[str], prompt: str, system_prompt: Optional[str] = None +): + """Process image inputs using the model's native processor. + + Uses token directly in the text and the model's own tokenizer/processor. + Each HF-processor tile is pre-patchified into the packed dynamic-resolution + format expected by RADIO when dynamic_resolution=True (temporal_patch_dim=1). + + Returns: + Tuple of (input_ids, packed_pixel_values, num_patches, imgs_sizes). + ``packed_pixel_values`` has shape [1, total_patches, 3*P*P]; ``imgs_sizes`` + has shape [N_tiles, 2] with per-tile (H, W). Both are None for text-only. + """ + if image_path: + image_paths = image_path.split(",") if "," in image_path else [image_path] + images = [load_image(p.strip()) for p in image_paths] + + image_placeholders = "\n".join([""] * len(images)) + text_content = f"{image_placeholders}\n{prompt}" + + messages = [{"role": "user", "content": text_content}] + if system_prompt: + messages.insert(0, {"role": "system", "content": system_prompt}) + + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + inputs = processor(text=[text], images=images, return_tensors="pt") + + pixel_values = inputs.pixel_values # [N_tiles, 3, H, W] + if hasattr(inputs, "num_patches") and inputs.num_patches is not None: + num_patches = inputs.num_patches + else: + num_patches = torch.ones(pixel_values.shape[0], dtype=torch.int) + + # Pre-patchify for the packed dynamic-resolution RADIO path. + packed_pv, sizes = _patchify_pixel_values(pixel_values) # [1, N*py*px, 3*P*P] + imgs_sizes = torch.tensor(sizes, dtype=torch.long) + + print_rank_0( + f"Image: {image_path}, tiles={pixel_values.shape[0]}, " + f"packed_shape={tuple(packed_pv.shape)}, num_patches={num_patches.tolist()}" + ) + return inputs.input_ids, packed_pv, num_patches, imgs_sizes + else: + messages = [{"role": "user", "content": prompt}] + if system_prompt: + messages.insert(0, {"role": "system", "content": system_prompt}) + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + inputs = tokenizer(text, return_tensors="pt") + return inputs.input_ids, None, 0, None + + +def process_video_inputs( + tokenizer, processor, video_path: Optional[str], prompt: str, system_prompt: Optional[str] = None +): + """Process video inputs for the temporal video embedder inference path. + + Mirrors ``NemotronOmniTaskEncoder.encode_sample`` (SFT data pipeline) so the + model sees the same tensor shapes it was trained on: + + - Frames are grouped by ``temporal_patch_size`` (pair of consecutive frames + per ) and the prompt uses the training-time format ("frame i sampled + at t seconds and frame i+1 sampled at t+1 seconds: "). + - One representative frame per group is fed to the HF processor (with + ``max_num_tiles=1``) so that // wrapper tokens render + correctly; ``adjust_image_tokens`` then shrinks each wrapper region down + to exactly one token. + - ALL video frames are patchified to a [1, total_patches, 3*P*P] tensor, + which replaces ``pixel_values`` so RADIO receives pre-patchified + dynamic-resolution input. + - ``imgs_sizes`` is emitted per (pre-grouping) frame, ``num_frames=[N]`` + so RADIO's ``_apply_temporal_grouping`` can fuse frame pairs and route + them through the trained ``video_embedder``. + + Returns: + Tuple of (input_ids, pixel_values, num_patches, imgs_sizes, num_frames). + """ + from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import ( + maybe_path_or_url_to_data_urls, + pil_image_from_base64, + ) + + tps = _VIDEO_TEMPORAL_PATCH_SIZE + + # 1. Extract frames from the video + image_urls, metadata = maybe_path_or_url_to_data_urls( + video_path, + fps=max(0, int(_VIDEO_FPS)), + nframe=max(0, int(_VIDEO_NFRAMES)), + nframe_max=-1, + ) + frames = [pil_image_from_base64(url) for url in image_urls] + print_rank_0(f"Video: extracted {len(frames)} frames, metadata: {metadata}") + + # RADIO's _apply_temporal_grouping fuses `tps` frames per tubelet. To keep + # the math clean and match SFT, round down to a multiple of tps. + usable = (len(frames) // tps) * tps + if usable == 0: + raise ValueError(f"Need at least {tps} frames for temporal video embedder inference; got {len(frames)}") + if usable != len(frames): + print_rank_0(f"Trimming {len(frames) - usable} trailing frames to match tps={tps}") + frames = frames[:usable] + + # 2. Build training-style prompt: one per `tps`-frame group, with + # per-frame timestamps. Also collect a representative frame per group + # for the HF processor to use as placeholder rendering. + fps_for_ts = float(metadata.fps) if (metadata and metadata.fps) else float(_VIDEO_FPS) + paired_images = [] + video_prompt_lines = ["This is a video:"] + for i in range(0, len(frames), tps): + group = frames[i : i + tps] + ts_parts = [f"frame {i + j + 1} sampled at {(i + j) / fps_for_ts:.2f} seconds" for j in range(len(group))] + video_prompt_lines.append(" and ".join(ts_parts) + ": ") + paired_images.append(group[0]) + + content = "\n".join(video_prompt_lines) + "\n" + prompt + messages = [{"role": "user", "content": content}] + if system_prompt: + messages.insert(0, {"role": "system", "content": system_prompt}) + + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + # 3. Process the representative frames to produce input_ids with proper + # ... wrappers around each . Force max_num_tiles=1 to + # match the training encoder (adjust_image_tokens then collapses each + # wrapper region back to a single token). + orig_tiles = getattr(processor.image_processor, "max_num_tiles", None) + if orig_tiles is not None: + processor.image_processor.max_num_tiles = 1 + try: + proc_output = processor(text=[text], images=paired_images, return_tensors="pt") + finally: + if orig_tiles is not None: + processor.image_processor.max_num_tiles = orig_tiles + + input_ids = proc_output.input_ids + num_patches = torch.ones(len(paired_images), dtype=torch.long) + + # 4. Replace the HF processor's pixel_values with an ALL-frames patchified + # tensor of shape [1, total_patches, 3*P*P] (dynamic-resolution input). + all_patches = [_patchify_frame(f, _VIDEO_FRAME_H, _VIDEO_FRAME_W) for f in frames] + packed_pixel_values = torch.cat(all_patches, dim=0).unsqueeze(0) + + imgs_sizes = torch.tensor([[_VIDEO_FRAME_H, _VIDEO_FRAME_W]] * len(frames), dtype=torch.long) + num_frames = torch.tensor([len(frames)], dtype=torch.long) + + return input_ids, packed_pixel_values, num_patches, imgs_sizes, num_frames + + +def process_audio_inputs(tokenizer, processor, audio_path: str, prompt: str, system_prompt: Optional[str] = None): + """Process audio inputs for the Megatron sound encoder. + + Uses the HF processor to expand tokens and extract mel + spectrogram features via ParakeetFeatureExtractor. The Megatron + BridgeSoundEncoder expects pre-processed mel spectrograms, not raw waveforms. + + Returns: + Tuple of (input_ids, sound_clips, sound_length) where sound_clips is a + mel spectrogram tensor and sound_length is the frame count tensor. + """ + from transformers import ParakeetFeatureExtractor + + audio_token = getattr(tokenizer, "audio_token", "") + messages = [{"role": "user", "content": f"{audio_token}\n{prompt}"}] + if system_prompt: + messages.insert(0, {"role": "system", "content": system_prompt}) + + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + # Processor handles audio loading and token expansion + inputs = processor(text=[text], audio=[audio_path], return_tensors="pt") + raw_sound_clips = inputs.pop("sound_clips", None) + + # Extract mel spectrogram features (Megatron BridgeSoundEncoder expects mel, not raw waveform) + feature_extractor = ParakeetFeatureExtractor(sampling_rate=16000, feature_size=128) + audio_features = feature_extractor(raw_sound_clips, sampling_rate=16000, return_tensors="pt") + sound_clips = audio_features.input_features # [batch, frames, mel_bins] + sound_length = torch.tensor([sound_clips.shape[1]], dtype=torch.long) + + # Realign tokens in the prompt to match the encoder's actual output + # length (see _align_sound_tokens docstring for why this can differ from the HF count). + sound_token_id = tokenizer.convert_tokens_to_ids(audio_token) + expected_sound_tokens = _fastconformer_output_length(int(sound_length.item())) + input_ids = _align_sound_tokens(inputs.input_ids, sound_token_id, expected_sound_tokens) + + print_rank_0(f"Audio: {audio_path}") + print_rank_0( + f"Sound clips shape: {sound_clips.shape}, sound_length: {sound_length}, " + f"expected_sound_tokens: {expected_sound_tokens}" + ) + + return input_ids, sound_clips, sound_length + + +def process_video_audio_inputs( + tokenizer, processor, video_path: str, audio_path: str, prompt: str, system_prompt: Optional[str] = None +): + """Process combined video + audio inputs for the temporal video embedder path. + + Same temporal preprocessing as ``process_video_inputs`` (frame pairing, + patchification, imgs_sizes / num_frames), with audio mel features + token + realignment layered on top, matching the SFT training pipeline. + + Returns: + Tuple of (input_ids, pixel_values, num_patches, imgs_sizes, num_frames, + sound_clips, sound_length). + """ + from transformers import ParakeetFeatureExtractor + + from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import ( + maybe_path_or_url_to_data_urls, + pil_image_from_base64, + ) + + tps = _VIDEO_TEMPORAL_PATCH_SIZE + + image_urls, metadata = maybe_path_or_url_to_data_urls( + video_path, + fps=max(0, int(_VIDEO_FPS)), + nframe=max(0, int(_VIDEO_NFRAMES)), + nframe_max=-1, + ) + frames = [pil_image_from_base64(url) for url in image_urls] + print_rank_0(f"Video: extracted {len(frames)} frames, metadata: {metadata}") + + usable = (len(frames) // tps) * tps + if usable == 0: + raise ValueError(f"Need at least {tps} frames for temporal video embedder inference; got {len(frames)}") + if usable != len(frames): + print_rank_0(f"Trimming {len(frames) - usable} trailing frames to match tps={tps}") + frames = frames[:usable] + + fps_for_ts = float(metadata.fps) if (metadata and metadata.fps) else float(_VIDEO_FPS) + paired_images = [] + video_prompt_lines = ["This is a video:"] + for i in range(0, len(frames), tps): + group = frames[i : i + tps] + ts_parts = [f"frame {i + j + 1} sampled at {(i + j) / fps_for_ts:.2f} seconds" for j in range(len(group))] + video_prompt_lines.append(" and ".join(ts_parts) + ": ") + paired_images.append(group[0]) + + audio_token = getattr(tokenizer, "audio_token", "") + content = "\n".join(video_prompt_lines) + f"\nThis is the audio: {audio_token}\n" + prompt + messages = [{"role": "user", "content": content}] + if system_prompt: + messages.insert(0, {"role": "system", "content": system_prompt}) + + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + orig_tiles = getattr(processor.image_processor, "max_num_tiles", None) + if orig_tiles is not None: + processor.image_processor.max_num_tiles = 1 + try: + proc_output = processor(text=[text], images=paired_images, audio=[audio_path], return_tensors="pt") + finally: + if orig_tiles is not None: + processor.image_processor.max_num_tiles = orig_tiles + + # Extract raw sound clips and convert to mel spectrogram for BridgeSoundEncoder + raw_sound_clips = proc_output.pop("sound_clips", None) + feature_extractor = ParakeetFeatureExtractor(sampling_rate=16000, feature_size=128) + audio_features = feature_extractor(raw_sound_clips, sampling_rate=16000, return_tensors="pt") + sound_clips = audio_features.input_features + sound_length = torch.tensor([sound_clips.shape[1]], dtype=torch.long) + + sound_token_id = tokenizer.convert_tokens_to_ids(audio_token) + expected_sound_tokens = _fastconformer_output_length(int(sound_length.item())) + input_ids = _align_sound_tokens(proc_output.input_ids, sound_token_id, expected_sound_tokens) + + print_rank_0( + f"Audio: {audio_path}, sound_clips shape: {sound_clips.shape}, expected_sound_tokens: {expected_sound_tokens}" + ) + + num_patches = torch.ones(len(paired_images), dtype=torch.long) + all_patches = [_patchify_frame(f, _VIDEO_FRAME_H, _VIDEO_FRAME_W) for f in frames] + packed_pixel_values = torch.cat(all_patches, dim=0).unsqueeze(0) + imgs_sizes = torch.tensor([[_VIDEO_FRAME_H, _VIDEO_FRAME_W]] * len(frames), dtype=torch.long) + num_frames = torch.tensor([len(frames)], dtype=torch.long) + + return input_ids, packed_pixel_values, num_patches, imgs_sizes, num_frames, sound_clips, sound_length + + +def main(args) -> None: + """Main function for Nemotron Omni VL generation from HuggingFace models. + + Loads a Nemotron Omni model either from HuggingFace (with optional conversion + to Megatron) or directly from a Megatron checkpoint, then performs greedy + generation using the provided prompt and optional image input. + + Args: + args: Parsed command line arguments containing model paths, prompt, + image path, parallelism settings, and generation parameters + """ + # pylint: disable=C0115,C0116 + tp = args.tp + pp = args.pp + ep = args.ep + etp = args.etp + + # Select vision-backbone config based on input modality. + # + # Image: dynamic_resolution=True with temporal_patch_dim=1 so that RADIO + # runs the packed dynamic-resolution path (is_packed_dynamic_res=True in + # LlavaModel). Each HF-processor tile is pre-patchified into a packed + # [1, N*patches, 3*P*P] tensor and passed with imgs_sizes / + # vision_packed_seq_params. num_image_tiles is recomputed by LlavaModel + # from RADIO output (256 tokens/tile after pixel_shuffle). + # + # Video (and video+audio): dynamic_resolution=True, temporal_patch_dim=2, + # separate_video_embedder=True so RADIO exercises the trained + # `video_embedder`. The matching data pipeline is in `process_video_inputs` + # / `process_video_audio_inputs`. + # + # Audio / text-only: dynamic_resolution=False, temporal_patch_dim=1. + is_video_inference = bool(args.video_path) + is_image_inference = bool(args.image_path) and not is_video_inference + if is_video_inference: + dynamic_resolution = True + temporal_patch_dim = 2 + separate_video_embedder = True + temporal_ckpt_compat = True + elif is_image_inference: + dynamic_resolution = True + temporal_patch_dim = 1 + separate_video_embedder = False + temporal_ckpt_compat = False + else: + dynamic_resolution = False + temporal_patch_dim = 1 + separate_video_embedder = False + temporal_ckpt_compat = False + + # Choose loading method based on arguments + if args.megatron_model_path: + # Load from Megatron checkpoint + print_rank_0(f"Loading Megatron model from: {args.megatron_model_path}") + + # We still need HF config for tokenizer, but we'll load the model from Megatron checkpoint + # Create bridge from HF config only (no weights) + bridge = AutoBridge.from_hf_pretrained(args.hf_model_path, trust_remote_code=True) + + # Initialize model parallel before loading + model_provider = bridge.to_megatron_provider(load_weights=False) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.expert_tensor_parallel_size = etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.dynamic_resolution = dynamic_resolution + model_provider.temporal_patch_dim = temporal_patch_dim + model_provider.separate_video_embedder = separate_video_embedder + model_provider.temporal_ckpt_compat = temporal_ckpt_compat + model_provider.initialize_model_parallel(seed=0) + + # Load the Megatron model directly. The mp_overrides values are applied to + # the loaded model_cfg before the model is built (see model_load_save.py), + # so the temporal/dynamic-resolution overrides must be passed here -- the + # `model_provider` mutations above only affect the throwaway provider used + # for parallel-state init, not the model that load_megatron_model builds. + model = bridge.load_megatron_model( + args.megatron_model_path, + mp_overrides={ + "tensor_model_parallel_size": tp, + "pipeline_model_parallel_size": pp, + "expert_model_parallel_size": ep, + "expert_tensor_parallel_size": etp, + "pipeline_dtype": torch.bfloat16, + "dynamic_resolution": dynamic_resolution, + "temporal_patch_dim": temporal_patch_dim, + "separate_video_embedder": separate_video_embedder, + "temporal_ckpt_compat": temporal_ckpt_compat, + }, + wrap_with_ddp=False, + ) + model[0].module.llava_model.dynamic_resolution = dynamic_resolution + if model[0].module.llava_model.vision_model is not None: + model[0].module.llava_model.vision_model.dynamic_resolution = dynamic_resolution + else: + # Load from HuggingFace and convert to Megatron + print_rank_0(f"Loading HuggingFace model from: {args.hf_model_path}") + bridge = AutoBridge.from_hf_pretrained(args.hf_model_path, trust_remote_code=True) + model_provider = bridge.to_megatron_provider(load_weights=True) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.expert_tensor_parallel_size = etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.dynamic_resolution = dynamic_resolution + model_provider.temporal_patch_dim = temporal_patch_dim + model_provider.separate_video_embedder = separate_video_embedder + model_provider.temporal_ckpt_compat = temporal_ckpt_compat + model_provider.initialize_model_parallel(seed=0) + model_provider.finalize() + model = model_provider.provide_distributed_model(wrap_with_ddp=False) + + model = [m.cuda() for m in model] + for m in model: + m.eval() + + # Clear grad_scale_func for inference. Training checkpoints serialize + # `optimizer.scale_loss` into the model config; reusing it here would + # call the unbound MegatronOptimizer.scale_loss inside the MoE branch of + # forward_step_calc_loss and crash with "missing 1 required positional + # argument: 'loss'". + for m in model: + inner = m.module if hasattr(m, "module") else m + if hasattr(inner, "config"): + inner.config.grad_scale_func = None + if hasattr(inner, "llava_model") and hasattr(inner.llava_model, "config"): + inner.llava_model.config.grad_scale_func = None + + # Initialize tokenizer and processor + tokenizer = AutoTokenizer.from_pretrained(args.hf_model_path, trust_remote_code=True) + processor = AutoProcessor.from_pretrained(args.hf_model_path, trust_remote_code=True) + img_start_token_id = tokenizer.convert_tokens_to_ids("") + img_end_token_id = tokenizer.convert_tokens_to_ids("") + + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + sound_clips = None + sound_length = None + images = None + imgs_sizes = None + num_frames = None + vision_packed_seq_params = None + + if args.video_path and args.audio_path: + ( + input_ids, + pixel_values, + num_patches, + imgs_sizes, + num_frames, + sound_clips, + sound_length, + ) = process_video_audio_inputs( + tokenizer, processor, args.video_path, args.audio_path, args.prompt, args.system_prompt + ) + images = pixel_values.bfloat16() if pixel_values is not None else None + elif args.audio_path: + input_ids, sound_clips, sound_length = process_audio_inputs( + tokenizer, processor, args.audio_path, args.prompt, args.system_prompt + ) + elif args.video_path: + input_ids, pixel_values, num_patches, imgs_sizes, num_frames = process_video_inputs( + tokenizer, processor, args.video_path, args.prompt, args.system_prompt + ) + images = pixel_values.bfloat16() if pixel_values is not None else None + else: + input_ids, pixel_values, num_patches, imgs_sizes = process_image_inputs( + tokenizer, processor, args.image_path, args.prompt, args.system_prompt + ) + images = pixel_values.bfloat16() if pixel_values is not None else None + + if images is not None: + # Adjust image tokens if / wrapper tokens are present. + # The HF processor may expand each into many tokens (one per patch), + # but Megatron LlavaModel expects one token per tile (image path) + # or one token per temporal tubelet (video path). + has_img_wrapper_tokens = ( + img_start_token_id != tokenizer.unk_token_id + and img_end_token_id != tokenizer.unk_token_id + and (input_ids == img_start_token_id).any() + ) + if has_img_wrapper_tokens: + input_ids = adjust_image_tokens(input_ids, num_patches, img_start_token_id, img_end_token_id) + pixel_values = None + + # Move to GPU + input_ids = input_ids.cuda() + if images is not None: + images = images.cuda() + if sound_clips is not None: + sound_clips = sound_clips.bfloat16().cuda() + if sound_length is not None: + sound_length = sound_length.cuda() + if imgs_sizes is not None: + imgs_sizes = imgs_sizes.cuda() + if num_frames is not None: + num_frames = num_frames.cuda() + + position_ids = ( + torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device).unsqueeze(0).expand_as(input_ids) + ) + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + generated_ids = input_ids.clone() + + stop_tokens = [tokenizer.eos_token_id] + + # Greedy generation loop + for step in range(args.max_new_tokens): + with torch.no_grad(): + print_rank_0(f"Generation step {step}") + + # Rebuild RADIO vision packed-seq params each step. RADIO's forward + # mutates `cu_seqlens_q/kv` and `max_seqlen_q/kv` in place to insert + # class tokens (megatron/core/models/vision/radio.py:388-394). Reusing + # the same object across iterations would compound the +class_token_len + # shift on every step and eventually drive cu_seqlens past the end of + # the embedded patch tensor, causing an async CUDA illegal memory + # access in attention. + vision_packed_seq_params = _build_vision_packed_seq_params(imgs_sizes) if imgs_sizes is not None else None + + fwd_bwd_function = get_forward_backward_func() + iterator = SingleBatchIterator( + input_ids, + position_ids, + attention_mask, + pixel_values=pixel_values, + images=images, + sound_clips=sound_clips, + sound_length=sound_length, + imgs_sizes=imgs_sizes, + num_frames=num_frames, + vision_packed_seq_params=vision_packed_seq_params, + ) + + output = fwd_bwd_function( + forward_step_func=vlm_forward_step, + data_iterator=iterator, + model=model, + num_microbatches=1, + forward_only=True, + seq_length=input_ids.size(1), + micro_batch_size=1, + collect_non_loss_data=True, + ) + if isinstance(output, list) and len(output) > 0: + output = output[0] + if isinstance(output, tuple): + output = output[0] + + if parallel_state.is_pipeline_last_stage(): + world_size = parallel_state.get_tensor_model_parallel_world_size() + gathered_tensors = [torch.zeros_like(output) for _ in range(world_size)] + dist.all_gather(gathered_tensors, output, group=parallel_state.get_tensor_model_parallel_group()) + output = torch.cat(gathered_tensors, dim=2) + next_token_ids = torch.argmax(output[:, -1], dim=-1, keepdim=True) + + if step < 5: + print_rank_0(f"Step {step}: output shape={output.shape}, var={output.var():.4f}") + logits = output[0, -1, :] + top5_vals, top5_ids = torch.topk(logits, 5) + top5_tokens = [tokenizer.decode([idx]) for idx in top5_ids] + print_rank_0(f"Top 5: {list(zip(top5_tokens, top5_vals.tolist()))}") + print_rank_0( + f"Selected: '{tokenizer.decode([next_token_ids.item()])}' (id={next_token_ids.item()})" + ) + else: + next_token_ids = torch.ones((1, 1), device=generated_ids.device, dtype=generated_ids.dtype) + + torch.distributed.broadcast(next_token_ids, get_last_rank()) + generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1) + + input_ids = generated_ids + position_ids = ( + torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device) + .unsqueeze(0) + .expand_as(input_ids) + ) + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + + if next_token_ids.item() in stop_tokens: + break + + # Decode the generated sequence + generated_text = tokenizer.decode(list(generated_ids[0])) + print_rank_0("======== GENERATED TEXT OUTPUT ========") + if args.image_path: + print_rank_0(f"Image: {args.image_path}") + if args.video_path: + print_rank_0(f"Video: {args.video_path}") + if args.audio_path: + print_rank_0(f"Audio: {args.audio_path}") + print_rank_0(f"Prompt: {args.prompt}") + print_rank_0(f"Generated: {generated_text}") + print_rank_0("=======================================") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Nemotron Omni VL Generation from HuggingFace Models") + parser.add_argument( + "--hf_model_path", + type=str, + default="nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", + help="Path to the HuggingFace Nemotron Omni VL model.", + ) + parser.add_argument( + "--prompt", + type=str, + default="Describe this image.", + help="Input prompt for vision-language generation.", + ) + parser.add_argument( + "--system_prompt", + type=str, + default="/no_think", + help="System prompt for vision-language generation.", + ) + parser.add_argument( + "--max_new_tokens", + type=int, + default=100, + help="Maximum number of new tokens to generate.", + ) + parser.add_argument("--tp", type=int, default=1, help="Tensor parallelism size") + parser.add_argument("--pp", type=int, default=1, help="Pipeline parallelism size") + parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size") + parser.add_argument("--etp", type=int, default=1, help="Expert tensor parallelism size") + parser.add_argument("--megatron_model_path", type=str, default=None, help="Path to the Megatron model checkpoint") + parser.add_argument( + "--image_path", + type=str, + default=None, + help="Path or URL to the image for vision-language generation (optional). Multiple image paths can be separated" + " with commas.", + ) + parser.add_argument( + "--video_path", + type=str, + default=None, + help="Path or URL to the video for vision-language generation (optional).", + ) + parser.add_argument( + "--audio_path", + type=str, + default=None, + help="Path to an audio file for audio understanding (optional, WAV format).", + ) + args = parser.parse_args() + + main(args) + + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() diff --git a/examples/models/nemotron/nemotron_3_omni/inference.sh b/examples/models/nemotron/nemotron_3_omni/inference.sh new file mode 100755 index 0000000000..cf998fac02 --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/inference.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +# 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. + +# ============================================================================== +# Nemotron-3 Nano Omni - Multimodal generation (image / video / audio / video+audio) +# +# Drives the local hf_to_megatron_generate_nemotron_omni.py helper against the +# imported Megatron checkpoint produced by conversion.sh. Omit --megatron_model_path +# to convert HF -> Megatron on the fly instead. +# ============================================================================== + +set -euo pipefail + +WORKSPACE=${WORKSPACE:-/workspace} +HF_MODEL_ID=${HF_MODEL_ID:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16} +MODEL_NAME=$(basename "$HF_MODEL_ID") +MEGATRON_PATH=${MEGATRON_PATH:-${WORKSPACE}/models/${MODEL_NAME}} + +# Default asset paths — override with your own image / mp4 / wav if desired. +IMAGE_PATH=${IMAGE_PATH:-${WORKSPACE}/assets/table.png} +VIDEO_PATH=${VIDEO_PATH:-${WORKSPACE}/assets/demo.mp4} +AUDIO_PATH=${AUDIO_PATH:-${WORKSPACE}/assets/demo_audio.wav} + +# --------------------------------------------------------------------------- +# Assets: download from the public HF model card if the paths don't exist. +# Audio is extracted from the video with ffmpeg (requires ffmpeg in PATH). +# --------------------------------------------------------------------------- +_HF_ASSETS="https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images" + +mkdir -p "$(dirname "$IMAGE_PATH")" "$(dirname "$VIDEO_PATH")" "$(dirname "$AUDIO_PATH")" + +[[ -f "$IMAGE_PATH" ]] || curl -fL "${_HF_ASSETS}/table.png" -o "$IMAGE_PATH" +[[ -f "$VIDEO_PATH" ]] || curl -fL "${_HF_ASSETS}/demo.mp4" -o "$VIDEO_PATH" +[[ -f "$AUDIO_PATH" ]] || ffmpeg -y -i "$VIDEO_PATH" -vn -acodec pcm_s16le -ar 16000 -ac 1 "$AUDIO_PATH" + +GEN_SCRIPT=examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py + +# 1) Image + Text (1 GPU) +uv run python -m torch.distributed.run --nproc_per_node=1 ${GEN_SCRIPT} \ + --hf_model_path "$HF_MODEL_ID" \ + --megatron_model_path "$MEGATRON_PATH" \ + --image_path "$IMAGE_PATH" \ + --prompt "Describe this image." \ + --max_new_tokens 100 + +# 2) Video + Text (TP=4, EP=4, 8 GPUs) +uv run python -m torch.distributed.run --nproc_per_node=8 ${GEN_SCRIPT} \ + --hf_model_path "$HF_MODEL_ID" \ + --megatron_model_path "$MEGATRON_PATH" \ + --video_path "$VIDEO_PATH" \ + --prompt "Describe what you see." \ + --tp 4 --ep 4 \ + --max_new_tokens 100 + +# 3) Audio + Text (1 GPU) +uv run python -m torch.distributed.run --nproc_per_node=1 ${GEN_SCRIPT} \ + --hf_model_path "$HF_MODEL_ID" \ + --megatron_model_path "$MEGATRON_PATH" \ + --audio_path "$AUDIO_PATH" \ + --prompt "Transcribe the audio." \ + --max_new_tokens 100 + +# 4) Video + Audio + Text (TP=4, EP=2, 8 GPUs) +uv run python -m torch.distributed.run --nproc_per_node=8 ${GEN_SCRIPT} \ + --hf_model_path "$HF_MODEL_ID" \ + --megatron_model_path "$MEGATRON_PATH" \ + --video_path "$VIDEO_PATH" \ + --audio_path "$AUDIO_PATH" \ + --prompt "Describe the video and audio." \ + --tp 4 --ep 2 \ + --max_new_tokens 150 diff --git a/examples/models/nemotron/nemotron_3_omni/slurm_peft_cord_v2.sh b/examples/models/nemotron/nemotron_3_omni/slurm_peft_cord_v2.sh new file mode 100755 index 0000000000..af412f35c0 --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/slurm_peft_cord_v2.sh @@ -0,0 +1,167 @@ +#!/bin/bash +# 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. + +# ============================================================================== +# Nemotron-3 Nano Omni - LoRA PEFT on CORD-V2 (image+text) +# +# Recipe: nemotron_omni_cord_v2_peft_config (HF dataset backend, no shards needed) +# LoRA targets: linear_qkv, linear_proj, in_proj, out_proj +# (LM attention + Mamba projections; vision/sound + projections frozen) +# Default parallelism: TP=2, EP=8, CP=1, MBS=2, GBS=16, packed sequences, +# selective recompute +# Default layout: 1 node / 8 GPUs +# (world_size = PP * max(TP*CP, EP*ETP) = 1 * max(2, 8) = 8) +# +# Override TP/EP/CP/PACKED_SEQ via environment, e.g.: +# TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch slurm_peft_cord_v2.sh +# +# Usage: +# sbatch slurm_peft_cord_v2.sh +# ============================================================================== + +#SBATCH --job-name=nomni-lora-cord-v2 +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=04:00:00 +#SBATCH --partition=gpu +#SBATCH --account=my_account +#SBATCH --output=nomni_lora_cord_v2_%j.out +#SBATCH --error=nomni_lora_cord_v2_%j.err +#SBATCH --exclusive + +set -euo pipefail + +# ============================================================================== +# CONFIGURATION +# ============================================================================== + +WORKSPACE=${WORKSPACE:-/workspace} +HF_MODEL_ID=${HF_MODEL_ID:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16} +MODEL_NAME=$(basename "$HF_MODEL_ID") + +PRETRAINED_CHECKPOINT=${WORKSPACE}/models/${MODEL_NAME} +RECIPE=nemotron_omni_cord_v2_peft_config +DATASET_NAME=cord_v2 + +# Parallelism / batching (override via env: TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch ...) +TP=${TP:-2} +EP=${EP:-8} +CP=${CP:-1} +PACKED_SEQ=${PACKED_SEQ:-true} + +SEQ_LENGTH=4096 +TRAIN_ITERS=4000 +GLOBAL_BATCH_SIZE=16 +MICRO_BATCH_SIZE=2 +EVAL_INTERVAL=50 +EVAL_ITERS=10 +SAVE_INTERVAL=200 +LOG_INTERVAL=1 +WANDB_PROJECT=megatron-bridge-${DATASET_NAME} + +# Container image (required) — use the NeMo 26.04 container or a local .sqsh copy +CONTAINER_IMAGE="" +# CONTAINER_IMAGE="nvcr.io/nvidia/nemo:26.04" +# CONTAINER_IMAGE="/path/to/nemo_26.04.sqsh" + +# Container mounts (optional, space-separated) +CONTAINER_MOUNTS="" +# CONTAINER_MOUNTS="/data:/data /workspace:/workspace" + +# ============================================================================== +# Environment Setup +# ============================================================================== + +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export HTTPX_LOG_LEVEL=WARNING +export PYTHONWARNINGS="ignore::FutureWarning:torch.cuda,ignore::UserWarning:modelopt.torch" +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +# export HF_HOME="/path/to/shared/HF_HOME" +# export HF_TOKEN="hf_your_token_here" +# export WANDB_API_KEY="your_wandb_key_here" +# export WANDB_MODE=disabled + +# ============================================================================== +# Job Execution +# ============================================================================== + +echo "======================================" +echo "Nemotron-3 Nano Omni - LoRA (CORD-V2)" +echo "======================================" +echo "Job ID: ${SLURM_JOB_ID:-N/A}" +echo "Nodes: ${SLURM_JOB_NUM_NODES:-N/A}" +echo "GPUs per node: ${SLURM_GPUS_PER_NODE:-N/A}" +echo "Recipe: $RECIPE" +echo "Checkpoint: $PRETRAINED_CHECKPOINT" +echo "Parallelism: TP=$TP EP=$EP CP=$CP (packed=$PACKED_SEQ, LoRA)" +echo "======================================" + +OUTPUT_DIR=${WORKSPACE}/results/${RECIPE}_lora + +CLI_OVERRIDES="\ + checkpoint.pretrained_checkpoint=$PRETRAINED_CHECKPOINT \ + checkpoint.save=$OUTPUT_DIR/checkpoints \ + checkpoint.save_interval=$SAVE_INTERVAL \ + checkpoint.finetune=True \ + logger.tensorboard_dir=$OUTPUT_DIR/tb_logs \ + model.seq_length=$SEQ_LENGTH \ + model.tensor_model_parallel_size=$TP \ + model.expert_model_parallel_size=$EP \ + model.context_parallel_size=$CP \ + model.sequence_parallel=True \ + model.recompute_granularity=selective \ + model.recompute_modules=[core_attn,mlp,layernorm,moe_act,moe] \ + model.freeze_language_model=False \ + train.train_iters=$TRAIN_ITERS \ + train.global_batch_size=$GLOBAL_BATCH_SIZE \ + train.micro_batch_size=$MICRO_BATCH_SIZE \ + dataset.trust_remote_code=True \ + dataset.pack_sequences_in_batch=$PACKED_SEQ \ + validation.eval_interval=$EVAL_INTERVAL \ + validation.eval_iters=$EVAL_ITERS \ + logger.log_interval=$LOG_INTERVAL \ + logger.wandb_project=$WANDB_PROJECT \ + logger.wandb_exp_name=${RECIPE}_lora" + +CMD="uv run --no-sync python scripts/training/run_recipe.py \ + --recipe $RECIPE \ + --hf_path $HF_MODEL_ID \ + --step_func nemotron_omni_step \ + $CLI_OVERRIDES" + +echo "Executing command..." +echo "======================================" + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: CONTAINER_IMAGE must be set. Please specify a valid container image." + exit 1 +fi + +SRUN_CMD="srun --mpi=pmix --container-image=$CONTAINER_IMAGE" + +if [ -n "$CONTAINER_MOUNTS" ]; then + # pyxis --container-mounts is comma-separated; multiple flags would last-win + MOUNTS_CSV=${CONTAINER_MOUNTS// /,} + SRUN_CMD="$SRUN_CMD --container-mounts=$MOUNTS_CSV" +fi + +$SRUN_CMD bash -c "$CMD" + +echo "======================================" +echo "Job completed" +echo "======================================" diff --git a/examples/models/nemotron/nemotron_3_omni/slurm_peft_valor32k_avqa.sh b/examples/models/nemotron/nemotron_3_omni/slurm_peft_valor32k_avqa.sh new file mode 100755 index 0000000000..f26befe71c --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/slurm_peft_valor32k_avqa.sh @@ -0,0 +1,182 @@ +#!/bin/bash +# 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. + +# ============================================================================== +# Nemotron-3 Nano Omni - LoRA PEFT on VALOR32K-AVQA (audio+video+text) +# +# Recipe: nemotron_omni_valor32k_peft_config (Energon shards) +# LoRA targets: linear_qkv, linear_proj, in_proj, out_proj +# (LM only; vision/sound encoders + projections frozen) +# Temporal video embedder enabled (temporal_patch_dim=2, +# dynamic_resolution=True, separate_video_embedder=True). +# Default parallelism: TP=2, EP=8, CP=1, MBS=2, GBS=64, packed sequences, +# selective recompute, lr=1e-4 +# Default layout: 1 node / 8 GPUs +# (world_size = PP * max(TP*CP, EP*ETP) = 1 * max(2, 8) = 8) +# +# Prerequisite: +# Build the VALOR32K-AVQA Energon shards (one-time): +# python tutorials/data/valor32k-avqa/build_valor32k_avqa_shards.py \ +# --output_dir ${WORKSPACE}/datasets/valor32k_avqa +# +# Override TP/EP/CP/PACKED_SEQ via environment, e.g.: +# TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch slurm_peft_valor32k_avqa.sh +# +# Usage: +# sbatch slurm_peft_valor32k_avqa.sh +# ============================================================================== + +#SBATCH --job-name=nomni-lora-valor32k +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=04:00:00 +#SBATCH --partition=gpu +#SBATCH --account=my_account +#SBATCH --output=nomni_lora_valor32k_%j.out +#SBATCH --error=nomni_lora_valor32k_%j.err +#SBATCH --exclusive + +set -euo pipefail + +# ============================================================================== +# CONFIGURATION +# ============================================================================== + +WORKSPACE=${WORKSPACE:-/workspace} +HF_MODEL_ID=${HF_MODEL_ID:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16} +MODEL_NAME=$(basename "$HF_MODEL_ID") + +PRETRAINED_CHECKPOINT=${WORKSPACE}/models/${MODEL_NAME} +ENERGON_PATH=${ENERGON_PATH:-${WORKSPACE}/datasets/valor32k_avqa/energon} +RECIPE=nemotron_omni_valor32k_peft_config +DATASET_NAME=valor32k_avqa + +# Parallelism / batching (override via env: TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch ...) +# Note: this recipe always uses temporal_patch_dim=2 (set in the recipe itself). +TP=${TP:-2} +EP=${EP:-8} +CP=${CP:-1} +PACKED_SEQ=${PACKED_SEQ:-true} + +SEQ_LENGTH=4096 +TRAIN_ITERS=4000 +GLOBAL_BATCH_SIZE=64 +MICRO_BATCH_SIZE=2 +LR=1e-4 +MIN_LR=1e-5 +EVAL_INTERVAL=50 +EVAL_ITERS=10 +SAVE_INTERVAL=200 +LOG_INTERVAL=1 +WANDB_PROJECT=megatron-bridge-${DATASET_NAME} + +# Container image (required) — use the NeMo 26.04 container or a local .sqsh copy +CONTAINER_IMAGE="" +# CONTAINER_IMAGE="nvcr.io/nvidia/nemo:26.04" +# CONTAINER_IMAGE="/path/to/nemo_26.04.sqsh" + +# Container mounts (optional, space-separated) +CONTAINER_MOUNTS="" +# CONTAINER_MOUNTS="/data:/data /workspace:/workspace" + +# ============================================================================== +# Environment Setup +# ============================================================================== + +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export HTTPX_LOG_LEVEL=WARNING +export PYTHONWARNINGS="ignore::FutureWarning:torch.cuda,ignore::UserWarning:modelopt.torch" +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +# export HF_HOME="/path/to/shared/HF_HOME" +# export HF_TOKEN="hf_your_token_here" +# export WANDB_API_KEY="your_wandb_key_here" +# export WANDB_MODE=disabled + +# ============================================================================== +# Job Execution +# ============================================================================== + +echo "======================================" +echo "Nemotron-3 Nano Omni - LoRA (VALOR32K-AVQA)" +echo "======================================" +echo "Job ID: ${SLURM_JOB_ID:-N/A}" +echo "Nodes: ${SLURM_JOB_NUM_NODES:-N/A}" +echo "GPUs per node: ${SLURM_GPUS_PER_NODE:-N/A}" +echo "Recipe: $RECIPE" +echo "Checkpoint: $PRETRAINED_CHECKPOINT" +echo "Energon shards: $ENERGON_PATH" +echo "Parallelism: TP=$TP EP=$EP CP=$CP (packed=$PACKED_SEQ, temporal video tpd=2, LoRA)" +echo "======================================" + +OUTPUT_DIR=${WORKSPACE}/results/${RECIPE}_lora + +CLI_OVERRIDES="\ + checkpoint.pretrained_checkpoint=$PRETRAINED_CHECKPOINT \ + checkpoint.save=$OUTPUT_DIR/checkpoints \ + checkpoint.save_interval=$SAVE_INTERVAL \ + checkpoint.finetune=True \ + logger.tensorboard_dir=$OUTPUT_DIR/tb_logs \ + dataset.path=$ENERGON_PATH \ + dataset.seq_length=$SEQ_LENGTH \ + dataset.pack_sequences_in_batch=$PACKED_SEQ \ + model.seq_length=$SEQ_LENGTH \ + model.tensor_model_parallel_size=$TP \ + model.expert_model_parallel_size=$EP \ + model.context_parallel_size=$CP \ + model.sequence_parallel=True \ + model.recompute_granularity=selective \ + model.recompute_modules=[core_attn,mlp,layernorm,moe_act,moe] \ + model.freeze_language_model=False \ + train.train_iters=$TRAIN_ITERS \ + train.global_batch_size=$GLOBAL_BATCH_SIZE \ + train.micro_batch_size=$MICRO_BATCH_SIZE \ + optimizer.lr=$LR \ + optimizer.min_lr=$MIN_LR \ + validation.eval_interval=$EVAL_INTERVAL \ + validation.eval_iters=$EVAL_ITERS \ + logger.log_interval=$LOG_INTERVAL \ + logger.wandb_project=$WANDB_PROJECT \ + logger.wandb_exp_name=${RECIPE}_lora" + +CMD="uv run --no-sync python scripts/training/run_recipe.py \ + --recipe $RECIPE \ + --hf_path $HF_MODEL_ID \ + --step_func nemotron_omni_step \ + $CLI_OVERRIDES" + +echo "Executing command..." +echo "======================================" + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: CONTAINER_IMAGE must be set. Please specify a valid container image." + exit 1 +fi + +SRUN_CMD="srun --mpi=pmix --container-image=$CONTAINER_IMAGE" + +if [ -n "$CONTAINER_MOUNTS" ]; then + # pyxis --container-mounts is comma-separated; multiple flags would last-win + MOUNTS_CSV=${CONTAINER_MOUNTS// /,} + SRUN_CMD="$SRUN_CMD --container-mounts=$MOUNTS_CSV" +fi + +$SRUN_CMD bash -c "$CMD" + +echo "======================================" +echo "Job completed" +echo "======================================" diff --git a/examples/models/nemotron/nemotron_3_omni/slurm_sft_cord_v2.sh b/examples/models/nemotron/nemotron_3_omni/slurm_sft_cord_v2.sh new file mode 100755 index 0000000000..f1d89542bf --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/slurm_sft_cord_v2.sh @@ -0,0 +1,168 @@ +#!/bin/bash +# 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. + +# ============================================================================== +# Nemotron-3 Nano Omni - Full SFT on CORD-V2 (image+text) +# +# Recipe: nemotron_omni_cord_v2_sft_config (HF dataset backend, no shards needed) +# Default parallelism: TP=2, EP=8, CP=1, MBS=2, GBS=16, packed sequences, +# selective recompute +# Default layout: 2 nodes / 16 GPUs +# +# Override TP/EP/CP/PACKED_SEQ via environment, e.g.: +# TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch slurm_sft_cord_v2.sh +# +# Usage: +# sbatch slurm_sft_cord_v2.sh +# ============================================================================== + +#SBATCH --job-name=nomni-sft-cord-v2 +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=04:00:00 +#SBATCH --partition=gpu +#SBATCH --account=my_account +#SBATCH --output=nomni_sft_cord_v2_%j.out +#SBATCH --error=nomni_sft_cord_v2_%j.err +#SBATCH --exclusive + +set -euo pipefail + +# ============================================================================== +# CONFIGURATION +# ============================================================================== + +WORKSPACE=${WORKSPACE:-/workspace} +HF_MODEL_ID=${HF_MODEL_ID:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16} +MODEL_NAME=$(basename "$HF_MODEL_ID") + +PRETRAINED_CHECKPOINT=${WORKSPACE}/models/${MODEL_NAME} +RECIPE=nemotron_omni_cord_v2_sft_config +DATASET_NAME=cord_v2 + +# Parallelism / batching (override via env: TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch ...) +TP=${TP:-2} +EP=${EP:-8} +CP=${CP:-1} +PACKED_SEQ=${PACKED_SEQ:-true} + +SEQ_LENGTH=4096 +TRAIN_ITERS=4000 +GLOBAL_BATCH_SIZE=16 +MICRO_BATCH_SIZE=2 +EVAL_INTERVAL=50 +EVAL_ITERS=10 +SAVE_INTERVAL=500 +LOG_INTERVAL=1 +WANDB_PROJECT=megatron-bridge-${DATASET_NAME} + +# Container image (required) — use the NeMo 26.04 container or a local .sqsh copy +CONTAINER_IMAGE="" +# CONTAINER_IMAGE="nvcr.io/nvidia/nemo:26.04" +# CONTAINER_IMAGE="/path/to/nemo_26.04.sqsh" + +# Container mounts (optional, space-separated) +CONTAINER_MOUNTS="" +# CONTAINER_MOUNTS="/data:/data /workspace:/workspace" + +# ============================================================================== +# Environment Setup +# ============================================================================== + +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export HTTPX_LOG_LEVEL=WARNING +export PYTHONWARNINGS="ignore::FutureWarning:torch.cuda,ignore::UserWarning:modelopt.torch" +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +# export HF_HOME="/path/to/shared/HF_HOME" +# export HF_TOKEN="hf_your_token_here" +# export WANDB_API_KEY="your_wandb_key_here" +# export WANDB_MODE=disabled + +# ============================================================================== +# Job Execution +# ============================================================================== + +echo "======================================" +echo "Nemotron-3 Nano Omni - SFT (CORD-V2)" +echo "======================================" +echo "Job ID: ${SLURM_JOB_ID:-N/A}" +echo "Nodes: ${SLURM_JOB_NUM_NODES:-N/A}" +echo "GPUs per node: ${SLURM_GPUS_PER_NODE:-N/A}" +echo "Recipe: $RECIPE" +echo "Checkpoint: $PRETRAINED_CHECKPOINT" +echo "Parallelism: TP=$TP EP=$EP CP=$CP (packed=$PACKED_SEQ)" +echo "======================================" + +OUTPUT_DIR=${WORKSPACE}/results/${RECIPE}_sft + +CLI_OVERRIDES="\ + checkpoint.pretrained_checkpoint=$PRETRAINED_CHECKPOINT \ + checkpoint.save=$OUTPUT_DIR/checkpoints \ + checkpoint.load=$OUTPUT_DIR/checkpoints \ + checkpoint.save_interval=$SAVE_INTERVAL \ + logger.tensorboard_dir=$OUTPUT_DIR/tb_logs \ + model.seq_length=$SEQ_LENGTH \ + model.tensor_model_parallel_size=$TP \ + model.expert_model_parallel_size=$EP \ + model.context_parallel_size=$CP \ + model.sequence_parallel=True \ + model.recompute_granularity=selective \ + model.recompute_modules=[core_attn,mlp,layernorm,moe_act,moe] \ + model.freeze_language_model=False \ + model.freeze_vision_model=False \ + model.freeze_vision_projection=False \ + model.freeze_sound_encoder=False \ + model.freeze_sound_projection=False \ + train.train_iters=$TRAIN_ITERS \ + train.global_batch_size=$GLOBAL_BATCH_SIZE \ + train.micro_batch_size=$MICRO_BATCH_SIZE \ + dataset.trust_remote_code=True \ + dataset.pack_sequences_in_batch=$PACKED_SEQ \ + validation.eval_interval=$EVAL_INTERVAL \ + validation.eval_iters=$EVAL_ITERS \ + logger.log_interval=$LOG_INTERVAL \ + logger.wandb_project=$WANDB_PROJECT \ + logger.wandb_exp_name=${RECIPE}_sft" + +CMD="uv run --no-sync python scripts/training/run_recipe.py \ + --recipe $RECIPE \ + --hf_path $HF_MODEL_ID \ + --step_func nemotron_omni_step \ + $CLI_OVERRIDES" + +echo "Executing command..." +echo "======================================" + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: CONTAINER_IMAGE must be set. Please specify a valid container image." + exit 1 +fi + +SRUN_CMD="srun --mpi=pmix --container-image=$CONTAINER_IMAGE" + +if [ -n "$CONTAINER_MOUNTS" ]; then + # pyxis --container-mounts is comma-separated; multiple flags would last-win + MOUNTS_CSV=${CONTAINER_MOUNTS// /,} + SRUN_CMD="$SRUN_CMD --container-mounts=$MOUNTS_CSV" +fi + +$SRUN_CMD bash -c "$CMD" + +echo "======================================" +echo "Job completed" +echo "======================================" diff --git a/examples/models/nemotron/nemotron_3_omni/slurm_sft_valor32k_avqa.sh b/examples/models/nemotron/nemotron_3_omni/slurm_sft_valor32k_avqa.sh new file mode 100755 index 0000000000..ffb9e598a5 --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/slurm_sft_valor32k_avqa.sh @@ -0,0 +1,180 @@ +#!/bin/bash +# 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. + +# ============================================================================== +# Nemotron-3 Nano Omni - Full SFT on VALOR32K-AVQA (audio+video+text) +# +# Recipe: nemotron_omni_valor32k_sft_config (Energon shards) +# Enables temporal video embedder (temporal_patch_dim=2, +# dynamic_resolution=True, separate_video_embedder=True) + FastConformer +# audio path. +# Default parallelism: TP=2, EP=8, CP=1, MBS=2, GBS=16, packed sequences, +# selective recompute +# Default layout: 2 nodes / 16 GPUs +# +# Prerequisite: +# Build the VALOR32K-AVQA Energon shards (one-time): +# python tutorials/data/valor32k-avqa/build_valor32k_avqa_shards.py \ +# --output_dir ${WORKSPACE}/datasets/valor32k_avqa +# +# Override TP/EP/CP/PACKED_SEQ via environment, e.g.: +# TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch slurm_sft_valor32k_avqa.sh +# +# Usage: +# sbatch slurm_sft_valor32k_avqa.sh +# ============================================================================== + +#SBATCH --job-name=nomni-sft-valor32k +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=04:00:00 +#SBATCH --partition=gpu +#SBATCH --account=my_account +#SBATCH --output=nomni_sft_valor32k_%j.out +#SBATCH --error=nomni_sft_valor32k_%j.err +#SBATCH --exclusive + +set -euo pipefail + +# ============================================================================== +# CONFIGURATION +# ============================================================================== + +WORKSPACE=${WORKSPACE:-/workspace} +HF_MODEL_ID=${HF_MODEL_ID:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16} +MODEL_NAME=$(basename "$HF_MODEL_ID") + +PRETRAINED_CHECKPOINT=${WORKSPACE}/models/${MODEL_NAME} +ENERGON_PATH=${ENERGON_PATH:-${WORKSPACE}/datasets/valor32k_avqa/energon} +RECIPE=nemotron_omni_valor32k_sft_config +DATASET_NAME=valor32k_avqa + +# Parallelism / batching (override via env: TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch ...) +# Note: this recipe always uses temporal_patch_dim=2 (set in the recipe itself). +TP=${TP:-2} +EP=${EP:-8} +CP=${CP:-1} +PACKED_SEQ=${PACKED_SEQ:-true} + +SEQ_LENGTH=4096 +TRAIN_ITERS=4000 +GLOBAL_BATCH_SIZE=16 +MICRO_BATCH_SIZE=2 +EVAL_INTERVAL=50 +EVAL_ITERS=10 +SAVE_INTERVAL=500 +LOG_INTERVAL=1 +WANDB_PROJECT=megatron-bridge-${DATASET_NAME} + +# Container image (required) — use the NeMo 26.04 container or a local .sqsh copy +CONTAINER_IMAGE="" +# CONTAINER_IMAGE="nvcr.io/nvidia/nemo:26.04" +# CONTAINER_IMAGE="/path/to/nemo_26.04.sqsh" + +# Container mounts (optional, space-separated) +CONTAINER_MOUNTS="" +# CONTAINER_MOUNTS="/data:/data /workspace:/workspace" + +# ============================================================================== +# Environment Setup +# ============================================================================== + +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export HTTPX_LOG_LEVEL=WARNING +export PYTHONWARNINGS="ignore::FutureWarning:torch.cuda,ignore::UserWarning:modelopt.torch" +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +# export HF_HOME="/path/to/shared/HF_HOME" +# export HF_TOKEN="hf_your_token_here" +# export WANDB_API_KEY="your_wandb_key_here" +# export WANDB_MODE=disabled + +# ============================================================================== +# Job Execution +# ============================================================================== + +echo "======================================" +echo "Nemotron-3 Nano Omni - SFT (VALOR32K-AVQA)" +echo "======================================" +echo "Job ID: ${SLURM_JOB_ID:-N/A}" +echo "Nodes: ${SLURM_JOB_NUM_NODES:-N/A}" +echo "GPUs per node: ${SLURM_GPUS_PER_NODE:-N/A}" +echo "Recipe: $RECIPE" +echo "Checkpoint: $PRETRAINED_CHECKPOINT" +echo "Energon shards: $ENERGON_PATH" +echo "Parallelism: TP=$TP EP=$EP CP=$CP (packed=$PACKED_SEQ, temporal video tpd=2)" +echo "======================================" + +OUTPUT_DIR=${WORKSPACE}/results/${RECIPE}_sft + +CLI_OVERRIDES="\ + checkpoint.pretrained_checkpoint=$PRETRAINED_CHECKPOINT \ + checkpoint.save=$OUTPUT_DIR/checkpoints \ + checkpoint.load=$OUTPUT_DIR/checkpoints \ + checkpoint.save_interval=$SAVE_INTERVAL \ + logger.tensorboard_dir=$OUTPUT_DIR/tb_logs \ + dataset.path=$ENERGON_PATH \ + dataset.seq_length=$SEQ_LENGTH \ + dataset.pack_sequences_in_batch=$PACKED_SEQ \ + model.seq_length=$SEQ_LENGTH \ + model.tensor_model_parallel_size=$TP \ + model.expert_model_parallel_size=$EP \ + model.context_parallel_size=$CP \ + model.sequence_parallel=True \ + model.recompute_granularity=selective \ + model.recompute_modules=[core_attn,mlp,layernorm,moe_act,moe] \ + model.freeze_language_model=False \ + model.freeze_vision_model=False \ + model.freeze_vision_projection=False \ + model.freeze_sound_encoder=False \ + model.freeze_sound_projection=False \ + train.train_iters=$TRAIN_ITERS \ + train.global_batch_size=$GLOBAL_BATCH_SIZE \ + train.micro_batch_size=$MICRO_BATCH_SIZE \ + validation.eval_interval=$EVAL_INTERVAL \ + validation.eval_iters=$EVAL_ITERS \ + logger.log_interval=$LOG_INTERVAL \ + logger.wandb_project=$WANDB_PROJECT \ + logger.wandb_exp_name=${RECIPE}_sft" + +CMD="uv run --no-sync python scripts/training/run_recipe.py \ + --recipe $RECIPE \ + --hf_path $HF_MODEL_ID \ + --step_func nemotron_omni_step \ + $CLI_OVERRIDES" + +echo "Executing command..." +echo "======================================" + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: CONTAINER_IMAGE must be set. Please specify a valid container image." + exit 1 +fi + +SRUN_CMD="srun --mpi=pmix --container-image=$CONTAINER_IMAGE" + +if [ -n "$CONTAINER_MOUNTS" ]; then + # pyxis --container-mounts is comma-separated; multiple flags would last-win + MOUNTS_CSV=${CONTAINER_MOUNTS// /,} + SRUN_CMD="$SRUN_CMD --container-mounts=$MOUNTS_CSV" +fi + +$SRUN_CMD bash -c "$CMD" + +echo "======================================" +echo "Job completed" +echo "======================================" diff --git a/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py b/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py new file mode 100644 index 0000000000..dc7f5702c5 --- /dev/null +++ b/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py @@ -0,0 +1,618 @@ +# 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 + +""" +VALOR32K-AVQA inference script for Nemotron Omni. + +Runs audio-visual QA inference on VALOR32K-AVQA test samples using a +Megatron checkpoint. Outputs predictions and computes accuracy. + +Vision backbone uses the temporal video embedder path +(dynamic_resolution=True, temporal_patch_dim=2, separate_video_embedder=True), +matching the SFT data pipeline in NemotronOmniTaskEncoder with +use_temporal_video_embedder=True. Frames are pre-patchified into a packed +[1, total_patches, 3*P*P] tensor with imgs_sizes / num_frames so RADIO ViT +exercises the trained `video_embedder`. + +Usage: + torchrun --nproc-per-node=8 examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py \ + --hf_model_path /path/to/nemotron-3-nano-omni-ea1_v1.0 \ + --megatron_model_path /path/to/checkpoint \ + --data_root /path/to/valor32k_avqa \ + --tp 4 --ep 2 \ + --max_samples 10 +""" + +import argparse +import json +import re +from pathlib import Path +from typing import Optional + +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.pipeline_parallel.schedules import get_forward_backward_func +from PIL import Image +from transformers import AutoProcessor, AutoTokenizer, ParakeetFeatureExtractor + +from megatron.bridge import AutoBridge +from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import ( + adjust_image_tokens, + maybe_path_or_url_to_data_urls, + pil_image_from_base64, +) +from megatron.bridge.utils.common_utils import get_last_rank, print_rank_0 + + +# --------------------------------------------------------------------------- +# Temporal video embedder constants (must match nemotron_omni_provider.py and +# the SFT recipe `nemotron_omni_valor32k_sft_config`). +# --------------------------------------------------------------------------- + +_VIDEO_TEMPORAL_PATCH_SIZE = 2 +_VIDEO_FRAME_H = 512 +_VIDEO_FRAME_W = 512 +_VISION_PATCH_DIM = 16 + +# CLIP / RADIO normalization (mirrors NemotronOmniTaskEncoder._patchify_frame) +_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) +_CLIP_STD = (0.26862954, 0.26130258, 0.27577711) + + +def _patchify_frame( + pil_img: Image.Image, target_h: int = _VIDEO_FRAME_H, target_w: int = _VIDEO_FRAME_W +) -> torch.Tensor: + """Resize + normalize a PIL frame and pack into [num_patches, 3*P*P].""" + from torchvision import transforms + + img = pil_img.convert("RGB").resize((target_w, target_h)) + tensor = transforms.ToTensor()(img) + mean = torch.tensor(_CLIP_MEAN).view(3, 1, 1) + std = torch.tensor(_CLIP_STD).view(3, 1, 1) + tensor = (tensor - mean) / std + P = _VISION_PATCH_DIM + py, px = target_h // P, target_w // P + patches = tensor.reshape(3, py, P, px, P).permute(1, 3, 0, 2, 4).reshape(py * px, 3 * P * P) + return patches + + +def _build_vision_packed_seq_params(imgs_sizes: Optional[torch.Tensor]) -> Optional[PackedSeqParams]: + """Vision PackedSeqParams from pre-grouping per-frame (H, W). + + Mirrors `megatron.bridge.training.nemotron_omni_step._build_vision_packed_seq_params`. + """ + if imgs_sizes is None or imgs_sizes.numel() == 0: + return None + sizes = imgs_sizes.tolist() if torch.is_tensor(imgs_sizes) else list(imgs_sizes) + seq_lens = [(int(h) // _VISION_PATCH_DIM) * (int(w) // _VISION_PATCH_DIM) for h, w in sizes] + cu = [0] + for sl in seq_lens: + cu.append(cu[-1] + sl) + device = imgs_sizes.device if torch.is_tensor(imgs_sizes) else torch.device("cpu") + cu_tensor = torch.tensor(cu, dtype=torch.int32, device=device) + max_len = torch.tensor(max(seq_lens) if seq_lens else 0, dtype=torch.int32, device=device) + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_tensor, + cu_seqlens_kv=cu_tensor, + max_seqlen_q=max_len, + max_seqlen_kv=max_len, + ) + + +# --------------------------------------------------------------------------- +# Forward step (same as examples/models/nemotron/nemotron_3_omni/hf_to_megatron_generate_nemotron_omni.py) +# --------------------------------------------------------------------------- + + +class SingleBatchIterator: + """Iterator that yields one prepared inference batch.""" + + def __init__(self, input_ids, position_ids, attention_mask, **kwargs): + self.batch = dict(tokens=input_ids, position_ids=position_ids, attention_mask=attention_mask) + if kwargs.get("images") is not None: + self.batch["images"] = kwargs["images"] + if kwargs.get("sound_clips") is not None: + self.batch["sound_clips"] = kwargs["sound_clips"] + if kwargs.get("sound_length") is not None: + self.batch["sound_length"] = kwargs["sound_length"] + # Temporal video embedder inputs (dynamic-resolution pre-patchified path) + if kwargs.get("imgs_sizes") is not None: + self.batch["imgs_sizes"] = kwargs["imgs_sizes"] + if kwargs.get("num_frames") is not None: + self.batch["num_frames"] = kwargs["num_frames"] + if kwargs.get("vision_packed_seq_params") is not None: + self.batch["vision_packed_seq_params"] = kwargs["vision_packed_seq_params"] + self._yielded = False + + def __iter__(self): + return self + + def __next__(self): + if self._yielded: + raise StopIteration + self._yielded = True + return self.batch + + +def vlm_forward_step(data_iterator, model, **kwargs): + """Run one VLM forward pass for audio-visual generation.""" + + batch = next(data_iterator) + forward_args = { + "input_ids": batch["tokens"], + "position_ids": batch["position_ids"], + "attention_mask": batch.get("attention_mask"), + } + if "images" in batch: + forward_args["images"] = batch["images"] + if "images" not in forward_args: + forward_args["images"] = torch.tensor([], dtype=torch.bfloat16, device=batch["tokens"].device).reshape(0, 0, 0) + if "sound_clips" in batch: + forward_args["sound_clips"] = batch["sound_clips"] + if "sound_length" in batch: + forward_args["sound_length"] = batch["sound_length"] + # Temporal video embedder plumbing (matches nemotron_omni_step.forward_step) + if "imgs_sizes" in batch: + forward_args["imgs_sizes"] = batch["imgs_sizes"] + if "num_frames" in batch: + forward_args["num_frames"] = batch["num_frames"] + if "vision_packed_seq_params" in batch: + forward_args["vision_packed_seq_params"] = batch["vision_packed_seq_params"] + + def loss_func(x, **kw): + return x + + output = model(**forward_args) + if isinstance(output, tuple): + output = output[0] + return output, loss_func + + +# --------------------------------------------------------------------------- +# Data processing +# --------------------------------------------------------------------------- + + +def build_video_id_map(videos_dir: Path) -> dict: + """Map video_id → filename stem (files are {youtube_id}_{start}_{end}.mp4).""" + mapping = {} + for f in videos_dir.iterdir(): + if f.suffix == ".mp4": + parts = f.stem.rsplit("_", 2) + if len(parts) == 3 and parts[0] not in mapping: + mapping[parts[0]] = f.stem + return mapping + + +def process_sample( + qa: dict, + vid_map: dict, + data_root: Path, + tokenizer, + processor, + feature_extractor, + video_fps: float = 1.0, + video_nframes: int = 8, + temporal_patch_size: int = _VIDEO_TEMPORAL_PATCH_SIZE, +): + """Process a single VALOR32K-AVQA sample into model inputs. + + Mirrors the SFT temporal data pipeline (NemotronOmniTaskEncoder with + use_temporal_video_embedder=True): frames are grouped in pairs in the + prompt, all frames are pre-patchified into [1, total_patches, 3*P*P], + and imgs_sizes / num_frames are emitted so RADIO's temporal grouping + + video_embedder run on inference inputs. + """ + import math + + video_id = str(qa["video_id"]) + file_stem = vid_map.get(video_id) + if file_stem is None: + return None + + video_path = data_root / "videos" / f"{file_stem}.mp4" + audio_path = data_root / "audio" / f"{file_stem}.wav" + if not video_path.exists(): + return None + + # Extract video frames + image_urls, metadata = maybe_path_or_url_to_data_urls( + str(video_path), + fps=max(0, int(video_fps)), + nframe=max(0, video_nframes), + nframe_max=-1, + ) + frames = [pil_image_from_base64(url) for url in image_urls] + fps = metadata.fps if metadata and metadata.fps else video_fps + + # RADIO's _apply_temporal_grouping fuses `tps` consecutive frames per + # tubelet, so trim trailing frames to a multiple of tps (matches the SFT + # data layer which also produces tps-aligned frame counts). + tps = temporal_patch_size + usable = (len(frames) // tps) * tps + if usable == 0: + print_rank_0(f"[skip] {video_id}: only {len(frames)} frames, need >= {tps}") + return None + if usable != len(frames): + frames = frames[:usable] + + # Group frames by temporal_patch_size for the prompt: one per pair, + # with the training-time timestamp format. + paired_images = [] + video_prompt_lines = ["This is a video:"] + for i in range(0, len(frames), tps): + group = frames[i : i + tps] + ts_parts = [f"frame {i + j + 1} sampled at {(i + j) / fps:.2f} seconds" for j in range(len(group))] + video_prompt_lines.append(" and ".join(ts_parts) + ": ") + paired_images.append(group[0]) + + # Build question with MCQ options + question = qa["question"] + options = qa.get("options", []) + if options: + option_labels = "ABCD" + option_text = "\n".join(f"{option_labels[i]}. {opt}" for i, opt in enumerate(options)) + question = f"{question}\n{option_text}" + + # Build prompt + content = "\n".join(video_prompt_lines) + "\n" + question + messages = [ + {"role": "system", "content": "/no_think"}, + {"role": "user", "content": content}, + ] + prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + # Run the HF processor on representative frames (one per pair) with + # max_num_tiles=1 (or max_num_patches=1) to emit input_ids with proper + # // wrappers. We only consume `input_ids` here -- + # pixel_values is replaced below with the pre-patchified all-frames tensor. + # Different model versions use different attribute names for the tile limit. + _tile_attr = "max_num_tiles" if hasattr(processor.image_processor, "max_num_tiles") else "max_num_patches" + orig_tiles = getattr(processor.image_processor, _tile_attr) + setattr(processor.image_processor, _tile_attr, 1) + try: + inputs = processor(text=[prompt], images=paired_images, return_tensors="pt") + finally: + setattr(processor.image_processor, _tile_attr, orig_tiles) + + input_ids = inputs.input_ids + # One token per tubelet after adjust_image_tokens (matches the + # number of tubelets that RADIO's temporal grouping will produce). + num_patches = torch.ones(len(paired_images), dtype=torch.long) + + # Adjust image tokens + img_start = tokenizer.convert_tokens_to_ids("") + img_end = tokenizer.convert_tokens_to_ids("") + if (input_ids == img_start).any(): + input_ids = adjust_image_tokens(input_ids, num_patches, img_start, img_end) + + # Pre-patchify ALL frames into [1, total_patches, 3*P*P]: dynamic-resolution + # input that RADIO's _apply_temporal_grouping splits per-frame and fuses + # into tubelets via `video_embedder`. + all_patches = [_patchify_frame(f, _VIDEO_FRAME_H, _VIDEO_FRAME_W) for f in frames] + images = torch.cat(all_patches, dim=0).unsqueeze(0).bfloat16() + imgs_sizes = torch.tensor([[_VIDEO_FRAME_H, _VIDEO_FRAME_W]] * len(frames), dtype=torch.long) + num_frames = torch.tensor([len(frames)], dtype=torch.long) + + # Process audio + sound_clips = None + sound_length = None + if audio_path.exists(): + import soundfile as sf + + waveform, sr = sf.read(str(audio_path), dtype="float32") + if waveform.ndim > 1: + waveform = waveform.mean(axis=1) + if sr != 16000: + import librosa + + waveform = librosa.resample(waveform, orig_sr=sr, target_sr=16000) + waveform = waveform[: int(10.0 * 16000)] # max 10s + + audio_features = feature_extractor([waveform], sampling_rate=16000, return_tensors="pt") + sound_clips = audio_features.input_features.bfloat16() + sound_length = torch.tensor([sound_clips.shape[1]], dtype=torch.long) + + # Compute audio token count and insert into input_ids + mel_len = sound_clips.shape[1] + token_len = float(mel_len) + for _ in range(3): + token_len = math.floor((token_len + 2 - 3) / 2 + 1) + n_sound_tokens = max(1, int(token_len)) + + sound_id = tokenizer.convert_tokens_to_ids("") + so_start_id = tokenizer.convert_tokens_to_ids("") + so_end_id = tokenizer.convert_tokens_to_ids("") + img_end_positions = (input_ids[0] == img_end).nonzero(as_tuple=True)[0] + insert_pos = int(img_end_positions[-1]) + 1 if len(img_end_positions) > 0 else 1 + + sound_block = torch.tensor( + [so_start_id] + [sound_id] * n_sound_tokens + [so_end_id], + dtype=input_ids.dtype, + ).unsqueeze(0) + input_ids = torch.cat([input_ids[:, :insert_pos], sound_block, input_ids[:, insert_pos:]], dim=1) + + # Build answer + correct_idx = qa.get("correct_answer_idx", 0) + answer = options[correct_idx] if options and correct_idx < len(options) else "" + + return { + "input_ids": input_ids, + "images": images, + "imgs_sizes": imgs_sizes, + "num_frames": num_frames, + "sound_clips": sound_clips, + "sound_length": sound_length, + "question": qa["question"], + "options": options, + "correct_answer": answer, + "correct_idx": correct_idx, + "video_id": video_id, + "modality": qa.get("modality", "unknown"), + } + + +# --------------------------------------------------------------------------- +# Generation +# --------------------------------------------------------------------------- + +# Accept answers like "C", "C.", "C)", "C. Yellow" by mapping the leading letter +# to the option text. Reasoning-tuned models often emit just the letter, which +# a pure substring match against the full option text would reject. +_LETTER_RE = re.compile(r"^\s*([A-D])(?:[\.\)\:\-\s]|$)") + + +def grade_prediction(prediction: str, options: list, correct_answer: str) -> bool: + """Grade a model prediction against the VALOR32K-AVQA answer.""" + + m = _LETTER_RE.match(prediction) + if m and options: + idx = ord(m.group(1)) - ord("A") + if idx < len(options) and options[idx].lower() == correct_answer.lower(): + return True + return correct_answer.lower() in prediction.lower() + + +def generate(model, tokenizer, sample, max_new_tokens=50): + """Greedy generation loop.""" + input_ids = sample["input_ids"].cuda() + images = sample["images"].cuda() if sample["images"] is not None else None + sound_clips = sample["sound_clips"].cuda() if sample["sound_clips"] is not None else None + sound_length = sample["sound_length"].cuda() if sample["sound_length"] is not None else None + imgs_sizes = sample["imgs_sizes"].cuda() if sample.get("imgs_sizes") is not None else None + num_frames = sample["num_frames"].cuda() if sample.get("num_frames") is not None else None + + position_ids = torch.arange(input_ids.size(1), device=input_ids.device).unsqueeze(0).expand_as(input_ids) + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + generated_ids = input_ids.clone() + stop_tokens = [tokenizer.eos_token_id] + + for step in range(max_new_tokens): + with torch.no_grad(): + # Rebuild each iteration: RADIO mutates cu_seqlens_q in-place when inserting class tokens, + # so reusing the same object would cause cu_seqlens to grow by class_token_len each step. + vision_packed_seq_params = _build_vision_packed_seq_params(imgs_sizes) + fwd_bwd_function = get_forward_backward_func() + iterator = SingleBatchIterator( + input_ids, + position_ids, + attention_mask, + images=images, + sound_clips=sound_clips, + sound_length=sound_length, + imgs_sizes=imgs_sizes, + num_frames=num_frames, + vision_packed_seq_params=vision_packed_seq_params, + ) + output = fwd_bwd_function( + forward_step_func=vlm_forward_step, + data_iterator=iterator, + model=model, + num_microbatches=1, + forward_only=True, + seq_length=input_ids.size(1), + micro_batch_size=1, + collect_non_loss_data=True, + ) + if isinstance(output, list) and len(output) > 0: + output = output[0] + if isinstance(output, tuple): + output = output[0] + + if parallel_state.is_pipeline_last_stage(): + world_size = parallel_state.get_tensor_model_parallel_world_size() + gathered = [torch.zeros_like(output) for _ in range(world_size)] + dist.all_gather(gathered, output, group=parallel_state.get_tensor_model_parallel_group()) + output = torch.cat(gathered, dim=2) + next_token_ids = torch.argmax(output[:, -1], dim=-1, keepdim=True) + else: + next_token_ids = torch.ones((1, 1), device=generated_ids.device, dtype=generated_ids.dtype) + + dist.broadcast(next_token_ids, get_last_rank()) + generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1) + + input_ids = generated_ids + position_ids = torch.arange(input_ids.size(1), device=input_ids.device).unsqueeze(0).expand_as(input_ids) + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + + if next_token_ids.item() in stop_tokens: + break + + # Decode: return both a cleaned prediction (for grading) and the full raw + # decode (with special tokens and prompt) for inspection. + gen_ids = generated_ids[0, sample["input_ids"].size(1) :] + cleaned = tokenizer.decode(gen_ids.tolist(), skip_special_tokens=True).strip() + full_text = tokenizer.decode(generated_ids[0].tolist(), skip_special_tokens=False) + return cleaned, full_text + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(): + """Run VALOR32K-AVQA inference.""" + + parser = argparse.ArgumentParser(description="VALOR32K-AVQA Inference") + parser.add_argument("--hf_model_path", type=str, required=True) + parser.add_argument( + "--megatron_model_path", + type=str, + default=None, + help="Megatron checkpoint path. If omitted, converts from HF on the fly.", + ) + parser.add_argument("--data_root", type=str, required=True) + parser.add_argument("--split", type=str, default="test") + parser.add_argument("--max_samples", type=int, default=10) + parser.add_argument("--max_new_tokens", type=int, default=50) + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--pp", type=int, default=1) + parser.add_argument("--ep", type=int, default=1) + parser.add_argument("--etp", type=int, default=1) + parser.add_argument("--output", type=str, default=None, help="Save predictions to JSON file") + args = parser.parse_args() + + data_root = Path(args.data_root) + split_name = "val" if args.split == "validation" else args.split + qa_file = data_root / f"combined_dataset_{split_name}_flattened.json" + + # Load model. AVQA is video-only, so always enable the temporal video + # embedder path: dynamic-resolution + tps=2 + separate_video_embedder. + # Inputs are pre-patchified in `process_sample` to match this contract. + bridge = AutoBridge.from_hf_pretrained(args.hf_model_path, trust_remote_code=True) + model_provider = bridge.to_megatron_provider(load_weights=(args.megatron_model_path is None)) + model_provider.tensor_model_parallel_size = args.tp + model_provider.pipeline_model_parallel_size = args.pp + model_provider.expert_model_parallel_size = args.ep + model_provider.expert_tensor_parallel_size = args.etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.dynamic_resolution = True + model_provider.temporal_patch_dim = _VIDEO_TEMPORAL_PATCH_SIZE + model_provider.separate_video_embedder = True + model_provider.temporal_ckpt_compat = True + model_provider.vision_class_token_len = 10 + model_provider.initialize_model_parallel(seed=0) + + if args.megatron_model_path: + print_rank_0(f"Loading Megatron checkpoint from {args.megatron_model_path}") + model = bridge.load_megatron_model( + args.megatron_model_path, + mp_overrides={ + "tensor_model_parallel_size": args.tp, + "pipeline_model_parallel_size": args.pp, + "expert_model_parallel_size": args.ep, + "expert_tensor_parallel_size": args.etp, + "pipeline_dtype": torch.bfloat16, + "dynamic_resolution": True, + "temporal_patch_dim": _VIDEO_TEMPORAL_PATCH_SIZE, + "separate_video_embedder": True, + "temporal_ckpt_compat": True, + "vision_class_token_len": 10, + }, + wrap_with_ddp=False, + ) + model = [m.cuda().eval() for m in model] + # Set grad_scale_func to None for inference (training checkpoints have optimizer config) + for m in model: + inner = m.module if hasattr(m, "module") else m + if hasattr(inner, "config"): + inner.config.grad_scale_func = None + if hasattr(inner, "llava_model") and hasattr(inner.llava_model, "config"): + inner.llava_model.config.grad_scale_func = None + else: + print_rank_0(f"Converting HF model from {args.hf_model_path} on the fly") + model_provider.finalize() + model = model_provider.provide_distributed_model(wrap_with_ddp=False) + model = [m.cuda().bfloat16().eval() for m in model] + + # Load tokenizer and processor + tokenizer = AutoTokenizer.from_pretrained(args.hf_model_path, trust_remote_code=True) + processor = AutoProcessor.from_pretrained(args.hf_model_path, trust_remote_code=True) + feature_extractor = ParakeetFeatureExtractor(sampling_rate=16000, feature_size=128) + + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # Load QA data + with open(qa_file) as f: + qa_pairs = json.load(f) + print_rank_0(f"Loaded {len(qa_pairs)} QA pairs from {qa_file}") + + # Build video ID map + vid_map = build_video_id_map(data_root / "videos") + print_rank_0(f"Video ID map: {len(vid_map)} entries") + + # Run inference + results = [] + correct = 0 + total = 0 + max_samples = min(args.max_samples, len(qa_pairs)) + + for i in range(max_samples): + qa = qa_pairs[i] + sample = process_sample( + qa, + vid_map, + data_root, + tokenizer, + processor, + feature_extractor, + ) + if sample is None: + print_rank_0(f"[{i + 1}/{max_samples}] Skipped: video not found for {qa['video_id']}") + continue + + prediction, prediction_full = generate(model, tokenizer, sample, max_new_tokens=args.max_new_tokens) + + is_correct = grade_prediction(prediction, sample["options"], sample["correct_answer"]) + if is_correct: + correct += 1 + total += 1 + + result = { + "video_id": sample["video_id"], + "question": sample["question"], + "options": sample["options"], + "correct_answer": sample["correct_answer"], + "prediction": prediction, + "prediction_full_decode": prediction_full, + "is_correct": is_correct, + "modality": sample["modality"], + } + results.append(result) + + print_rank_0( + f"[{i + 1}/{max_samples}] Q: {sample['question'][:60]}... " + f"| GT: {sample['correct_answer']} | Pred: {prediction[:60]} " + f"| {'OK' if is_correct else 'WRONG'}" + ) + + # Summary + acc = correct / total * 100 if total > 0 else 0 + print_rank_0(f"\n{'=' * 60}") + print_rank_0(f"Results: {correct}/{total} correct ({acc:.1f}%)") + print_rank_0(f"{'=' * 60}") + + # Save results + if args.output and dist.get_rank() == 0: + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump({"accuracy": acc, "correct": correct, "total": total, "results": results}, f, indent=2) + print_rank_0(f"Results saved to {output_path}") + + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/scripts/training/run_recipe.py b/scripts/training/run_recipe.py index 6059fd1d7b..3473d1d0bc 100755 --- a/scripts/training/run_recipe.py +++ b/scripts/training/run_recipe.py @@ -124,6 +124,7 @@ from megatron.bridge.training.finetune import finetune from megatron.bridge.training.gpt_step import forward_step as gpt_forward_step from megatron.bridge.training.llava_step import forward_step as llava_forward_step +from megatron.bridge.training.nemotron_omni_step import forward_step as nemotron_omni_forward_step from megatron.bridge.training.pretrain import pretrain from megatron.bridge.training.utils.omegaconf_utils import process_config_with_overrides from megatron.bridge.training.vlm_step import forward_step as vlm_forward_step @@ -136,6 +137,7 @@ "qwen3_omni_step": qwen3_omni_forward_step, "qwen3_vl_step": qwen3_vl_forward_step, "llava_step": llava_forward_step, + "nemotron_omni_step": nemotron_omni_forward_step, "flux_step": FluxForwardStep, "wan_step": WanForwardStep, } diff --git a/src/megatron/bridge/data/energon/energon_provider.py b/src/megatron/bridge/data/energon/energon_provider.py index f33ea48dc1..e0d7ea2926 100644 --- a/src/megatron/bridge/data/energon/energon_provider.py +++ b/src/megatron/bridge/data/energon/energon_provider.py @@ -38,6 +38,12 @@ class EnergonProvider(DatasetProvider): def build_datasets(self, context: DatasetBuildContext): assert self.path, "EnergonProvider.path must be set. Use CLI override: dataset.path=" + if ( + self.pack_sequences_in_batch + and self.task_encoder is not None + and hasattr(self.task_encoder, "pack_sequences") + ): + self.task_encoder.pack_sequences = True dataset = EnergonMultiModalDataModule( path=self.path, tokenizer=context.tokenizer if context.tokenizer is not None else self.tokenizer, diff --git a/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py b/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py new file mode 100644 index 0000000000..4ae179633b --- /dev/null +++ b/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py @@ -0,0 +1,669 @@ +# 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. + +"""Nemotron Omni Energon task encoder -- extends HFEncoderVLMTaskEncoder with audio. + +Adds mel spectrogram extraction and ```` token insertion so that +the training step receives ``sound_clips`` / ``sound_length`` alongside the +standard vision + language tensors. +""" + +import dataclasses +import logging +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence, Tuple + +import numpy as np +import torch +from megatron.energon import Batch, DefaultTaskEncoder + +from megatron.bridge.data.energon.task_encoder_utils import ( + IGNORE_INDEX, + ChatMLSample, + _images_to_pil, + _videos_to_pil, + cook_chatml_sample, + find_pattern_indices, + get_ltor_masks_and_position_ids, +) +from megatron.bridge.training.utils.visual_inputs import GenericVisualInputs + + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Sample / batch dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class NemotronOmniTaskSample: + """Encoded sample for Nemotron Omni (vision + audio + language).""" + + __key__: str + __subflavors__: Dict + input_ids: torch.Tensor # [seq_len] + labels: torch.Tensor # [seq_len] + loss_mask: torch.Tensor # [seq_len] + visual_tensors: Dict[str, torch.Tensor] = field(default_factory=dict) + num_patches: Optional[torch.Tensor] = None # [num_images] tile count per image + sound_clips: Optional[torch.Tensor] = None # [frames, mel_bins] + sound_length: Optional[torch.Tensor] = None # scalar + imgs_sizes: Optional[torch.Tensor] = None # [num_frames, 2] per-frame (H, W) + num_frames: Optional[torch.Tensor] = None # [num_media_items] + num_image_tiles: Optional[torch.Tensor] = None # [num_images] LM-side token count per image + + +@dataclass +class NemotronOmniTaskBatch(Batch): + """Batched format for Nemotron Omni.""" + + __keys__: List[str] = field(default_factory=list) + __subflavors__: List[Dict] = field(default_factory=list) + input_ids: torch.Tensor = field(default_factory=lambda: torch.empty(0)) + labels: torch.Tensor = field(default_factory=lambda: torch.empty(0)) + loss_mask: torch.Tensor = field(default_factory=lambda: torch.empty(0)) + attention_mask: Optional[torch.Tensor] = field(default_factory=lambda: torch.empty(0)) + position_ids: torch.Tensor = field(default_factory=lambda: torch.empty(0)) + visual_tensors: Dict[str, Optional[torch.Tensor]] = field(default_factory=dict) + num_patches: Optional[torch.Tensor] = None # tile counts per image + sound_clips: Optional[torch.Tensor] = None # [B, max_frames, mel_bins] + sound_length: Optional[torch.Tensor] = None # [B] + imgs_sizes: Optional[torch.Tensor] = None # [total_frames, 2] + num_frames: Optional[torch.Tensor] = None # [num_media_items] + num_image_tiles: Optional[torch.Tensor] = None # [total_images] LM-side token count per image + # Packed-sequence metadata (only populated when pack_sequences=True). + cu_seqlens: Optional[torch.Tensor] = None + cu_seqlens_unpadded: Optional[torch.Tensor] = None + cu_seqlens_argmin: Optional[torch.Tensor] = None + max_seqlen: Optional[torch.Tensor] = None + + +# --------------------------------------------------------------------------- +# Task encoder +# --------------------------------------------------------------------------- + + +class NemotronOmniTaskEncoder(DefaultTaskEncoder[ChatMLSample, NemotronOmniTaskSample, NemotronOmniTaskBatch, dict]): + """Energon task encoder for Nemotron Omni models. + + Processes ChatML samples that may contain images, videos, AND audio + waveforms (decoded from ``.wav`` fields in the WebDataset shards). + + Audio waveforms are converted to mel spectrograms using + ``compute_mel_features``, and ```` placeholder tokens are + inserted into ``input_ids`` so that ``LLaVAModel.forward()`` can + replace them with the projected sound embeddings. + + Args: + processor: HF ``AutoProcessor`` for the Nemotron Omni model. + seq_length: Maximum sequence length after tokenization. + max_audio_duration: Maximum audio duration in seconds. Longer clips + are truncated. + num_mel_bins: Number of mel frequency bins (must match the sound + encoder config, typically 128 for Parakeet). + visual_keys: Processor output keys to capture as visual tensors. + """ + + def __init__( + self, + processor, + seq_length: int = 4096, + max_audio_duration: float = 30.0, + num_mel_bins: int = 128, + visual_keys: Sequence[str] = ("pixel_values",), + temporal_patch_size: int = 2, + video_fps: float = 1.0, + video_nframes: int = 8, + use_temporal_video_embedder: bool = False, + patch_dim: int = 16, + pack_sequences: bool = False, + ): + super().__init__() + self.processor = processor + self.seq_length = seq_length + self.max_audio_duration = max_audio_duration + self.num_mel_bins = num_mel_bins + self.visual_keys: Tuple[str, ...] = tuple(visual_keys) + self.temporal_patch_size = temporal_patch_size + self.video_fps = video_fps + self.video_nframes = video_nframes + self.use_temporal_video_embedder = use_temporal_video_embedder + self.patch_dim = patch_dim + self.pack_sequences = pack_sequences + + @staticmethod + def _decode_video_bytes(video_bytes: bytes, nframes: int = 8, fps: float = 1.0): + """Decode raw MP4 bytes to a list of PIL frames.""" + import tempfile + + from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import ( + maybe_path_or_url_to_data_urls, + pil_image_from_base64, + ) + + # Write to temp file since video_io needs a file path + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=True) as tmp: + tmp.write(video_bytes) + tmp.flush() + try: + image_urls, _ = maybe_path_or_url_to_data_urls( + tmp.name, + fps=max(0, int(fps)), + nframe=max(0, nframes), + nframe_max=-1, + ) + frames = [pil_image_from_base64(url) for url in image_urls] + return frames if frames else None + except Exception: + logger.warning("Failed to decode video bytes") + return None + + def _patchify_frame(self, pil_img, target_h: int = 512, target_w: int = 512) -> torch.Tensor: + """Convert a PIL image to [num_patches, C*P*P] patches (normalized). + + Matches the HF processor's normalization (CLIP mean/std). + """ + from torchvision import transforms + + img = pil_img.resize((target_w, target_h)) + tensor = transforms.ToTensor()(img) # [3, H, W] + # Normalize with CLIP / RADIO mean/std + mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(3, 1, 1) + std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(3, 1, 1) + tensor = (tensor - mean) / std + P = self.patch_dim + py, px = target_h // P, target_w // P + # [3, (py*P), (px*P)] → [py*px, 3*P*P] + patches = tensor.reshape(3, py, P, px, P).permute(1, 3, 0, 2, 4).reshape(py * px, 3 * P * P) + return patches + + @property + def _tokenizer(self): + return getattr(self.processor, "tokenizer", self.processor) + + @property + def _pad_token_id(self) -> int: + return self._tokenizer.pad_token_id or 0 + + @property + def _eos_token_id(self) -> int: + return self._tokenizer.eos_token_id + + @property + def _sound_token_id(self) -> int: + return self._tokenizer.convert_tokens_to_ids("") + + # ------------------------------------------------------------------ + # encode_sample + # ------------------------------------------------------------------ + + def encode_sample(self, sample: ChatMLSample) -> NemotronOmniTaskSample: + """Encode a single ChatML sample with optional audio into model-ready tensors.""" + import math as _math + + from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import ( + compute_mel_features, + ) + + # 1. Decode video → PIL frames + videos_raw = sample.videos + video_frames = None + if videos_raw is not None: + if isinstance(videos_raw, bytes): + video_frames = self._decode_video_bytes(videos_raw, nframes=self.video_nframes, fps=self.video_fps) + elif isinstance(videos_raw, list) and len(videos_raw) > 0: + if isinstance(videos_raw[0], bytes): + all_f = [] + for vb in videos_raw: + f = self._decode_video_bytes(vb, nframes=self.video_nframes, fps=self.video_fps) + if f: + all_f.extend(f) + video_frames = all_f if all_f else None + else: + video_frames = _videos_to_pil(videos_raw) + + images_pil = _images_to_pil(sample.imgs) if sample.imgs is not None and len(sample.imgs) > 0 else None + + # 2. Process audio → mel spectrogram + compute token count + n_sound_tokens = 0 + sound_clips_t: Optional[torch.Tensor] = None + sound_length_t: Optional[torch.Tensor] = None + + if sample.audio is not None: + waveform = sample.audio + if isinstance(waveform, torch.Tensor): + waveform = waveform.numpy() + waveform = waveform[: int(self.max_audio_duration * 16000)] + + mel = compute_mel_features(waveform, sampling_rate=16000, num_mel_bins=self.num_mel_bins) + sound_clips_t = mel + sound_length_t = torch.tensor(mel.shape[0], dtype=torch.long) + + token_len = float(mel.shape[0]) + for _ in range(3): + token_len = _math.floor((token_len + 2 * 1 - 3) / 2 + 1) + n_sound_tokens = max(1, int(token_len)) + + # 3. Build prompt with temporal frame pairing + audio tokens + # + # Target format (user turn): + # This is a video: + # frame 1 sampled at 0.00 seconds and frame 2 sampled at 1.00 seconds: + # frame 3 sampled at 2.00 seconds and frame 4 sampled at 3.00 seconds: + # [N×] + # {question text} + # + # Each token represents temporal_patch_size consecutive frames combined. + + conversation = cook_chatml_sample(sample.conversation) + tps = self.temporal_patch_size + fps = self.video_fps + + # Group video frames by temporal_patch_size → one per group + paired_images = [] + video_prompt_lines = [] + if video_frames: + video_prompt_lines.append("This is a video:") + for i in range(0, len(video_frames), tps): + group = video_frames[i : i + tps] + ts_parts = [] + for j in range(len(group)): + ts_parts.append(f"frame {i + j + 1} sampled at {(i + j) / fps:.2f} seconds") + video_prompt_lines.append(" and ".join(ts_parts) + ": ") + paired_images.append(group[0]) # representative frame per group + + # Replace conversation content with the structured prompt + all_proc_images = list(images_pil) if images_pil else [] + for turn in conversation: + if not isinstance(turn.get("content"), list): + continue + new_parts = [] + for item in turn["content"]: + if isinstance(item, dict) and item.get("type") == "video": + if video_prompt_lines: + new_parts.append("\n".join(video_prompt_lines)) + all_proc_images.extend(paired_images) + elif isinstance(item, dict) and item.get("type") == "image": + new_parts.append("") + elif isinstance(item, dict) and item.get("type") == "text": + new_parts.append(item.get("text", "")) + elif isinstance(item, str): + new_parts.append(item) + turn["content"] = "\n".join(new_parts) + + # 4. Apply chat template → tokenize + process vision + text_conv = [ + t if isinstance(t["content"], str) else {"role": t["role"], "content": str(t["content"])} + for t in conversation + ] + prompt_text = self._tokenizer.apply_chat_template(text_conv, tokenize=False, add_generation_prompt=False) + + proc_kwargs = {"text": prompt_text, "return_tensors": "pt"} + if all_proc_images: + proc_kwargs["images"] = all_proc_images + + orig_tiles = getattr(self.processor.image_processor, "max_num_tiles", None) + if all_proc_images and orig_tiles is not None: + self.processor.image_processor.max_num_tiles = 1 + proc_output = self.processor(**proc_kwargs) + if orig_tiles is not None: + self.processor.image_processor.max_num_tiles = orig_tiles + + input_ids_np = ( + proc_output["input_ids"][0].numpy() + if proc_output["input_ids"].dim() == 2 + else proc_output["input_ids"].numpy() + ) + + # 5. Insert audio tokens into input_ids after the last + # Format: + N × + + if n_sound_tokens > 0: + sound_id = self._sound_token_id + so_start_id = self._tokenizer.convert_tokens_to_ids("") + so_end_id = self._tokenizer.convert_tokens_to_ids("") + img_end_id = self._tokenizer.convert_tokens_to_ids("") + + img_end_positions = np.where(input_ids_np == img_end_id)[0] + insert_pos = int(img_end_positions[-1]) + 1 if len(img_end_positions) > 0 else 1 + + sound_block = np.array( + [so_start_id] + [sound_id] * n_sound_tokens + [so_end_id], + dtype=input_ids_np.dtype, + ) + input_ids_np = np.concatenate([input_ids_np[:insert_pos], sound_block, input_ids_np[insert_pos:]]) + + # 5b. Build loss mask FIRST on raw input_ids (before adjust_image_tokens), + # then adjust image tokens. The loss mask positions must align with the + # post-adjustment input_ids that the model receives. + pv = proc_output.get("pixel_values") + num_patches = None + if pv is not None: + num_tiles = pv.shape[0] if isinstance(pv, torch.Tensor) else len(pv) + num_patches = torch.ones(num_tiles, dtype=torch.long) + + # 6. Build loss mask — only supervise assistant turns + loss_mask_np = np.zeros(len(input_ids_np), dtype=np.float32) + search_start = 0 + for turn in conversation: + if turn["role"] == "assistant": + answer = turn["content"] if isinstance(turn["content"], str) else turn["content"] + if isinstance(answer, list): + answer = "".join(p.get("text", "") for p in answer if isinstance(p, dict)) + answer_tokens = self._tokenizer.encode(answer, add_special_tokens=False) + ans_start, ans_end = find_pattern_indices(input_ids_np, answer_tokens, search_start) + if ans_start >= 0: + loss_mask_np[ans_start:ans_end] = 1.0 + search_start = ans_end + + # 7. Labels = left-shifted input_ids + labels_np = np.full(len(input_ids_np), IGNORE_INDEX, dtype=np.int64) + labels_np[:-1] = input_ids_np[1:] + shifted_loss = np.zeros_like(loss_mask_np) + shifted_loss[:-1] = loss_mask_np[1:] + labels_np[shifted_loss == 0.0] = IGNORE_INDEX + loss_mask_np = shifted_loss + + # 7b. Adjust image tokens — shrink many tokens per tile to one. + # This must happen AFTER loss mask and labels are built so positions stay aligned. + img_start_id = self._tokenizer.convert_tokens_to_ids("") + img_end_id = self._tokenizer.convert_tokens_to_ids("") + if num_patches is not None and (input_ids_np == img_start_id).any(): + from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import adjust_image_tokens + + proc_num_patches = proc_output.get("num_patches") + if proc_num_patches is not None: + if not isinstance(proc_num_patches, torch.Tensor): + proc_num_patches = torch.tensor(proc_num_patches) + else: + proc_num_patches = num_patches + + # adjust_image_tokens can handle a dict of tensors with matching shapes + adjusted = adjust_image_tokens( + { + "input_ids": torch.from_numpy(input_ids_np).unsqueeze(0), + "labels": torch.from_numpy(labels_np).unsqueeze(0), + "loss_mask": torch.from_numpy(loss_mask_np).unsqueeze(0), + }, + proc_num_patches, + img_start_id, + img_end_id, + ) + input_ids_np = adjusted["input_ids"].squeeze(0).numpy() + labels_np = adjusted["labels"].squeeze(0).numpy() + loss_mask_np = adjusted["loss_mask"].squeeze(0).numpy() + + # 8. Truncate + max_len = self.seq_length + input_ids_np = input_ids_np[:max_len].copy() + labels_np = labels_np[:max_len].copy() + loss_mask_np = loss_mask_np[:max_len].copy() + + # 9. Collect visual tensors (num_patches already computed in step 5b) + visual_tensors: Dict[str, torch.Tensor] = {} + for key in self.visual_keys: + val = proc_output.get(key) + if val is not None: + visual_tensors[key] = val if isinstance(val, torch.Tensor) else torch.tensor(val) + + # 10. Temporal video embedder: patchify ALL video frames and pack + sample_imgs_sizes: Optional[torch.Tensor] = None + sample_num_frames: Optional[torch.Tensor] = None + if self.use_temporal_video_embedder and video_frames: + all_patches = [] + target_h, target_w = 512, 512 + for frame in video_frames: + patches = self._patchify_frame(frame, target_h, target_w) + all_patches.append(patches) + # Pack into [1, total_patches, C*P*P] for dynamic resolution + packed = torch.cat(all_patches, dim=0).unsqueeze(0) # [1, N*num_patches_per_frame, feat] + visual_tensors["pixel_values"] = packed + sample_imgs_sizes = torch.tensor([[target_h, target_w]] * len(video_frames), dtype=torch.long) + sample_num_frames = torch.tensor([len(video_frames)], dtype=torch.long) + # Also include standalone images (each is 1 frame) + if images_pil: + for img in images_pil: + img_patches = self._patchify_frame(img, target_h, target_w) + packed_img = img_patches.unsqueeze(0) + visual_tensors["pixel_values"] = torch.cat([packed_img, visual_tensors["pixel_values"]], dim=1) + sample_imgs_sizes = torch.cat( + [ + torch.tensor([[target_h, target_w]], dtype=torch.long), + sample_imgs_sizes, + ], + dim=0, + ) + sample_num_frames = torch.cat( + [ + torch.tensor([1], dtype=torch.long), + sample_num_frames, + ] + ) + + # Compute per-image num_image_tiles for LM-side image-token expansion + # (new llava_model.py dynamic_resolution path). num_tiles_i = (H/P * W/P) // 4 + # matching the HF collate's shuffled_count computation. + sample_num_image_tiles: Optional[torch.Tensor] = None + if sample_imgs_sizes is not None: + P = self.patch_dim + sample_num_image_tiles = torch.tensor( + [(int(h) // P) * (int(w) // P) // 4 for h, w in sample_imgs_sizes.tolist()], + dtype=torch.int, + ) + + return NemotronOmniTaskSample( + __key__=sample.__key__, + __subflavors__=sample.__subflavors__, + input_ids=torch.from_numpy(input_ids_np), + labels=torch.from_numpy(labels_np), + loss_mask=torch.from_numpy(loss_mask_np), + visual_tensors=visual_tensors, + num_patches=num_patches, + sound_clips=sound_clips_t, + sound_length=sound_length_t, + imgs_sizes=sample_imgs_sizes, + num_frames=sample_num_frames, + num_image_tiles=sample_num_image_tiles, + ) + + # ------------------------------------------------------------------ + # batch + # ------------------------------------------------------------------ + + def batch(self, samples: List[NemotronOmniTaskSample]) -> NemotronOmniTaskBatch: + """Pad-and-collate (default) OR pack samples along the seq dim when + ``pack_sequences=True``. Packing emits ``cu_seqlens`` / ``cu_seqlens_unpadded`` + / ``max_seqlen`` so TE's THD kernels handle cross-sample masking (and CP + partitioning via ``thd_get_partitioned_indices``) without an attention mask. + """ + pad_id = self._pad_token_id + batch_size = len(samples) + + cu_seqlens_t: Optional[torch.Tensor] = None + cu_seqlens_unpadded_t: Optional[torch.Tensor] = None + cu_seqlens_argmin_t: Optional[torch.Tensor] = None + max_seqlen_t: Optional[torch.Tensor] = None + + if self.pack_sequences: + # Concatenate samples along the seq dim into a single [1, total_len] + # microbatch. TE attention kernels use cu_seqlens for per-sample + # masking; no attention_mask needed. + lengths = [int(s.input_ids.size(0)) for s in samples] + cu_seqlens = [0] + for L in lengths: + cu_seqlens.append(cu_seqlens[-1] + L) + + tokens_flat = torch.cat([s.input_ids for s in samples], dim=0) + labels_flat = torch.cat([s.labels for s in samples], dim=0) + loss_mask_flat = torch.cat([s.loss_mask for s in samples], dim=0) + # Per-sample resetting position ids: [0..L1-1, 0..L2-1, ...] + position_ids_flat = torch.cat([torch.arange(L, dtype=torch.long) for L in lengths], dim=0) + + tokens = tokens_flat.unsqueeze(0) + tokens[tokens == pad_id] = 0 + labels = labels_flat.unsqueeze(0) + loss_mask_t = loss_mask_flat.unsqueeze(0) + position_ids = position_ids_flat.unsqueeze(0) + attention_mask = None # TE derives the causal+padding mask from cu_seqlens. + + cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32) + cu_seqlens_unpadded_t = cu_seqlens_t.clone() + # get_packed_seq_params truncates cu_seqlens_padded[: argmin.item()]; the + # trick in the fixed-size-batched case is sentinel=-1 padding with argmin + # pointing at the first sentinel. Here we emit an unpadded cu_seqlens and + # set argmin = len(cu_seqlens) so the slice is a no-op (keeps every entry). + cu_seqlens_argmin_t = torch.tensor(len(cu_seqlens), dtype=torch.int32) + max_seqlen_t = torch.tensor(max(lengths), dtype=torch.int32) + else: + max_seq_len = max(s.input_ids.size(0) for s in samples) + input_ids_mat = np.full((batch_size, max_seq_len), pad_id, dtype=np.int64) + labels_mat = np.full((batch_size, max_seq_len), IGNORE_INDEX, dtype=np.int64) + loss_mask_mat = np.zeros((batch_size, max_seq_len), dtype=np.float32) + + for i, s in enumerate(samples): + seq_len = min(max_seq_len, s.input_ids.size(0)) + input_ids_mat[i, :seq_len] = s.input_ids.numpy()[:seq_len] + labels_mat[i, :seq_len] = s.labels.numpy()[:seq_len] + loss_mask_mat[i, :seq_len] = s.loss_mask.numpy()[:seq_len] + + tokens = torch.from_numpy(input_ids_mat) + tokens[tokens == pad_id] = 0 + labels = torch.from_numpy(labels_mat) + loss_mask_t = torch.from_numpy(loss_mask_mat) + + attention_mask, _, position_ids = get_ltor_masks_and_position_ids( + data=tokens, + eod_token=self._eos_token_id, + eod_mask_loss=False, + reset_attention_mask=False, + reset_position_ids=False, + ) + + # Aggregate visual tensors. + # The temporal video path ships pixel_values as [1, N_i*patches_per_frame, feat] + # per sample. When packing, concat along dim=1 so the whole microbatch becomes a + # single [1, total_patches, feat] packed sequence that matches vision_packed_seq_params + # cu_seqlens. Without packing, preserve the legacy [MBS, patches, feat] stack. + all_visual_keys = set() + for s in samples: + all_visual_keys.update(s.visual_tensors.keys()) + batched_visual: Dict[str, Optional[torch.Tensor]] = {} + for key in all_visual_keys: + tensors = [s.visual_tensors[key] for s in samples if key in s.visual_tensors] + if not tensors: + batched_visual[key] = None + continue + if self.pack_sequences and tensors[0].dim() == 3: + batched_visual[key] = torch.cat(tensors, dim=1) + else: + batched_visual[key] = torch.cat(tensors, dim=0) + + # Aggregate audio: pad mel spectrograms to max length in batch + has_audio = any(s.sound_clips is not None for s in samples) + sound_clips_batch: Optional[torch.Tensor] = None + sound_length_batch: Optional[torch.Tensor] = None + + if has_audio: + mel_list = [] + mel_lengths = [] + for s in samples: + if s.sound_clips is not None: + mel_list.append(s.sound_clips) + mel_lengths.append(s.sound_clips.shape[0]) + else: + mel_list.append(torch.zeros(1, self.num_mel_bins)) + mel_lengths.append(1) + + max_mel_len = max(mel_lengths) + mel_dim = mel_list[0].shape[-1] + sound_clips_batch = torch.zeros(batch_size, max_mel_len, mel_dim) + for i, mel in enumerate(mel_list): + sound_clips_batch[i, : mel.shape[0]] = mel + sound_length_batch = torch.tensor(mel_lengths, dtype=torch.long) + + # Aggregate num_patches + all_patches = [s.num_patches for s in samples if s.num_patches is not None] + num_patches_batch = torch.cat(all_patches, dim=0) if all_patches else None + + # Aggregate imgs_sizes / num_frames (temporal video embedder) + has_temporal = any(s.imgs_sizes is not None for s in samples) + imgs_sizes_batch: Optional[torch.Tensor] = None + num_frames_batch: Optional[torch.Tensor] = None + if has_temporal: + all_imgs_sizes = [s.imgs_sizes for s in samples if s.imgs_sizes is not None] + all_num_frames = [s.num_frames for s in samples if s.num_frames is not None] + imgs_sizes_batch = torch.cat(all_imgs_sizes, dim=0) if all_imgs_sizes else None + num_frames_batch = torch.cat(all_num_frames, dim=0) if all_num_frames else None + all_num_image_tiles = [s.num_image_tiles for s in samples if s.num_image_tiles is not None] + num_image_tiles_batch = torch.cat(all_num_image_tiles, dim=0) if all_num_image_tiles else None + + batch_kwargs: Dict = dict( + __keys__=[s.__key__ for s in samples], + __subflavors__=[s.__subflavors__ for s in samples], + input_ids=tokens, + labels=labels, + loss_mask=loss_mask_t, + attention_mask=attention_mask, + position_ids=position_ids, + visual_tensors=batched_visual, + num_patches=num_patches_batch, + sound_clips=sound_clips_batch, + sound_length=sound_length_batch, + imgs_sizes=imgs_sizes_batch, + num_frames=num_frames_batch, + num_image_tiles=num_image_tiles_batch, + cu_seqlens=cu_seqlens_t, + cu_seqlens_unpadded=cu_seqlens_unpadded_t, + cu_seqlens_argmin=cu_seqlens_argmin_t, + max_seqlen=max_seqlen_t, + ) + _batch_fields = {f.name for f in dataclasses.fields(NemotronOmniTaskBatch)} + if "__key__" in _batch_fields: + batch_kwargs["__key__"] = samples[0].__key__ + if "__restore_key__" in _batch_fields: + batch_kwargs["__restore_key__"] = () + + return NemotronOmniTaskBatch(**batch_kwargs) + + # ------------------------------------------------------------------ + # encode_batch + # ------------------------------------------------------------------ + + def encode_batch(self, batch: NemotronOmniTaskBatch) -> dict: + """Convert batch to dict for the training step.""" + raw = { + "tokens": batch.input_ids, + "labels": batch.labels, + "loss_mask": batch.loss_mask, + "attention_mask": batch.attention_mask, + "position_ids": batch.position_ids, + "num_patches": batch.num_patches, + "sound_clips": batch.sound_clips, + "sound_length": batch.sound_length, + "imgs_sizes": batch.imgs_sizes, + "num_frames": batch.num_frames, + "num_image_tiles": batch.num_image_tiles, + "cu_seqlens": batch.cu_seqlens, + "cu_seqlens_unpadded": batch.cu_seqlens_unpadded, + "cu_seqlens_argmin": batch.cu_seqlens_argmin, + "max_seqlen": batch.max_seqlen, + } + + vt = batch.visual_tensors if batch.visual_tensors else {} + raw["visual_inputs"] = GenericVisualInputs(**{k: v for k, v in vt.items() if v is not None}) + + # Keep sound_clips / sound_length as top-level batch keys + # (nemotron_omni_step picks them up directly) + return raw diff --git a/src/megatron/bridge/data/energon/task_encoder_utils.py b/src/megatron/bridge/data/energon/task_encoder_utils.py index 9cf45baddb..d24dc6df7d 100644 --- a/src/megatron/bridge/data/energon/task_encoder_utils.py +++ b/src/megatron/bridge/data/energon/task_encoder_utils.py @@ -209,11 +209,12 @@ def _videos_to_pil(videos): # --------------------------------------------------------------------------- @dataclass class ChatMLSample(Sample): - """Multi-turn complex samples with images and videos.""" + """Multi-turn complex samples with images, videos, and audio.""" conversation: str # JSON string of GPT-format conversations imgs: Optional[List[torch.Tensor]] = None videos: Optional[List[List[torch.Tensor]]] = None + audio: Optional[torch.Tensor] = None # Raw waveform tensor [num_samples] or pre-computed mel [frames, mel_bins] class videohandler: @@ -238,6 +239,33 @@ def __call__(self, key, data): return data +class audiohandler: + """Webdataset decoder handler for audio fields stored as raw WAV/FLAC bytes.""" + + EXTENSIONS = {"wav", "flac", "mp3", "audio"} + + def __call__(self, key, data): + extension = re.sub(r".*[.]", "", key) + if extension.lower() not in self.EXTENSIONS: + return None + try: + import io + + import soundfile as sf + + waveform, sr = sf.read(io.BytesIO(data), dtype="float32", always_2d=False) + if waveform.ndim > 1: + waveform = waveform.mean(axis=-1) + if sr != 16000: + import librosa + + waveform = librosa.resample(waveform, orig_sr=sr, target_sr=16000) + return torch.from_numpy(waveform.astype(np.float32)) + except Exception: + logging.warning(f"Failed to decode audio for key {key}") + return None + + class ChatMLWebdataset(DefaultDecoderWebdatasetFactory[ChatMLSample]): """Webdataset factory for multi-turn ChatML samples with multimodal support. @@ -262,13 +290,15 @@ class installs its own decoder. def __init__(self, path: EPath, *, auto_decode: bool = True, image_decode_spec: Optional[str] = None, **kwargs): kwargs.pop("decoder", None) - super().__init__(path, auto_decode=auto_decode, **kwargs) + kwargs.pop("auto_decode", None) + super().__init__(path, **kwargs) if auto_decode: spec = image_decode_spec if image_decode_spec is not None else getattr(self, "image_decode", "torchrgb") self._decoder = Decoder( [ imagehandler(spec), videohandler(spec), + audiohandler(), ] ) diff --git a/src/megatron/bridge/data/vlm_datasets/__init__.py b/src/megatron/bridge/data/vlm_datasets/__init__.py index c1ab8abe2a..7cbcef8b08 100644 --- a/src/megatron/bridge/data/vlm_datasets/__init__.py +++ b/src/megatron/bridge/data/vlm_datasets/__init__.py @@ -26,6 +26,7 @@ COLLATE_FNS, default_collate_fn, nemotron_nano_v2_vl_collate_fn, + nemotron_omni_collate_fn, phi4_mm_collate_fn, qwen2_5_collate_fn, qwen2_audio_collate_fn, @@ -66,5 +67,6 @@ "qwen2_5_collate_fn", "phi4_mm_collate_fn", "qwen2_audio_collate_fn", - "qwen2_audio_collate_fn", + "nemotron_nano_v2_vl_collate_fn", + "nemotron_omni_collate_fn", ] diff --git a/src/megatron/bridge/data/vlm_datasets/collate.py b/src/megatron/bridge/data/vlm_datasets/collate.py index f4d1802b9b..5f6baa6527 100644 --- a/src/megatron/bridge/data/vlm_datasets/collate.py +++ b/src/megatron/bridge/data/vlm_datasets/collate.py @@ -377,6 +377,398 @@ def nemotron_nano_v2_vl_collate_fn(examples: list, processor, start_of_response_ return batch +def nemotron_omni_collate_fn( + examples: list, + processor, + start_of_response_token=None, + *, + pack_sequences: bool = False, +) -> dict[str, torch.Tensor]: + """Collate function for Nemotron Omni model (vision + audio + language). + + Extends nemotron_nano_v2_vl_collate_fn with audio support. Each example + may carry an ``audio_path`` field pointing to a 16 kHz mono WAV file. + Audio is converted to mel spectrograms and added to the batch as + ``sound_clips`` / ``sound_length`` tensors consumed by LLaVAModel.forward(). + + When ``pack_sequences=True``, samples in the microbatch are concatenated + along the sequence dim into a single ``[1, sum(L_i)]`` batch, and + ``cu_seqlens`` / ``cu_seqlens_unpadded`` / ``cu_seqlens_argmin`` / + ``max_seqlen`` are emitted so TE's THD attention kernels handle per-sample + masking without an attention mask. Requires ``mbs > 1`` to be meaningful. + """ + from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import ( + compute_mel_features, + load_audio, + ) + from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import adjust_image_tokens + + # Ensure the tokenizer has a pad_token: the processor pads only when one is set, + # and mbs>1 needs padding to collate sequences of different lengths. Safe no-op + # when pad_token is already set. + if processor.tokenizer.pad_token is None and processor.tokenizer.eos_token is not None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + + skipped_tokens = extract_skipped_token_ids(processor) + first_content = examples[0]["conversation"][0]["content"] + is_video = isinstance(first_content, list) and first_content[0].get("type") == "video" + + # --- Vision path --- + # The Nemotron Omni chat template does not expand {"type": "image"} content + # into tokens — it stringifies the list. We must convert conversations + # to use explicit text and pass PIL images via processor(images=...). + if is_video: + from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import ( + maybe_path_or_url_to_data_urls, + pil_image_from_base64, + ) + + assert len(examples) == 1, "Nemotron Omni processor only supports batch size == 1 for video" + frames = [] + video_nframe = 10 + + for example in examples: + video_path = example["conversation"][0]["content"][0]["path"] + image_urls, metadata = maybe_path_or_url_to_data_urls( + video_path, + fps=0, + nframe=max(0, int(video_nframe)), + nframe_max=-1, + ) + frames.append([pil_image_from_base64(image_url) for image_url in image_urls]) + + prompt = processor.apply_chat_template([ex["conversation"] for ex in examples], tokenize=False) + batch = processor(text=prompt, videos=frames, videos_kwargs={"video_metadata": metadata}, return_tensors="pt") + else: + # Convert structured {"type": "image"} content to explicit text + all_images = [] + images_per_ex: list[list] = [] + text_conversations = [] + for example in examples: + images_for_example = [] + text_conv = [] + for turn in example["conversation"]: + if isinstance(turn["content"], list): + text_parts = [] + for item in turn["content"]: + if item["type"] == "image": + text_parts.append("") + images_for_example.append(item["image"]) + elif item["type"] == "text": + text_parts.append(item["text"]) + text_conv.append({"role": turn["role"], "content": "\n".join(text_parts)}) + elif isinstance(turn["content"], str): + text_conv.append(turn) + else: + text_conv.append({"role": turn["role"], "content": str(turn["content"])}) + all_images.extend(images_for_example) + images_per_ex.append(images_for_example) + text_conversations.append(text_conv) + + prompts = [ + processor.tokenizer.apply_chat_template(conv, tokenize=False, add_generation_prompt=False) + for conv in text_conversations + ] + # Normalize audio tokens: replace model-agnostic <|audio_1|> with Nemotron Omni's + audio_token = getattr(processor.tokenizer, "audio_token", "") + prompts = [p.replace("<|audio_1|>", audio_token) for p in prompts] + if all_images: + # Older Nemotron-VL image processors use fixed 512x512 tiles and expose + # `max_num_tiles`; the newer Nemotron-3 Omni Reasoning processor uses + # dynamic-resolution patches (no `max_num_tiles` attr, has + # `max_num_patches` instead). Detect which path we're on. + is_dynamic_res_processor = not hasattr(processor.image_processor, "max_num_tiles") + if is_dynamic_res_processor: + # Variable per-image (H, W) makes ``return_tensors="pt"`` fail to + # stack pixel_values across examples. Process each example + # separately and re-combine: right-pad input_ids across examples, + # keep pixel_values as a flat list of per-image ``[3, H_i, W_i]`` + # tensors (patchified below with per-image (py, px)). + per_ex_batches = [ + processor( + text=[prompt], + images=imgs if imgs else None, + padding=False, + truncation=True, + return_tensors="pt", + ) + for prompt, imgs in zip(prompts, images_per_ex) + ] + pad_id = processor.tokenizer.pad_token_id + if pad_id is None: + pad_id = processor.tokenizer.eos_token_id or 0 + ids_list = [b["input_ids"][0] for b in per_ex_batches] + max_len = max(t.shape[0] for t in ids_list) + padded_ids = torch.full((len(per_ex_batches), max_len), pad_id, dtype=ids_list[0].dtype) + for i, ids in enumerate(ids_list): + padded_ids[i, : ids.shape[0]] = ids + pv_list: list[torch.Tensor] = [] + for b in per_ex_batches: + if "pixel_values" in b and b["pixel_values"] is not None: + pv_b = b["pixel_values"] + if pv_b.dim() == 4: + for img in pv_b: + pv_list.append(img) + elif pv_b.dim() == 3: + pv_list.append(pv_b) + batch = {"input_ids": padded_ids} + if pv_list: + batch["pixel_values"] = pv_list # list[Tensor[3, H_i, W_i]] + else: + # Static-tile path: single-tile per image to match RADIO seq_length. + orig_tiles = processor.image_processor.max_num_tiles + processor.image_processor.max_num_tiles = 1 + batch = processor( + text=prompts, + images=all_images, + padding=processor.tokenizer.pad_token is not None, + truncation=True, + return_tensors="pt", + ) + processor.image_processor.max_num_tiles = orig_tiles + else: + is_dynamic_res_processor = False + batch = processor.tokenizer( + prompts, + padding=processor.tokenizer.pad_token is not None, + truncation=True, + return_tensors="pt", + ) + + # --- Audio path --- + # Support both audio_path (file path) and audio (raw waveform tuple from CV17-style datasets) + has_audio = any(ex.get("audio_path") or ex.get("audio") for ex in examples) + if has_audio: + import numpy as np + + max_dur = examples[0].get("max_audio_duration", 30.0) + max_samples = int(max_dur * 16000) + + mel_list = [] + mel_lengths = [] + n_audio_tokens_list = [] + for ex in examples: + audio_path = ex.get("audio_path") + audio_tuple = ex.get("audio") # (array, sr) from CV17-style datasets + if audio_path: + waveform = load_audio(audio_path, target_sr=16000) + elif audio_tuple is not None: + array, sr = audio_tuple + waveform = np.asarray(array, dtype=np.float32) + if sr != 16000: + import librosa + + waveform = librosa.resample(waveform, orig_sr=sr, target_sr=16000) + else: + mel_list.append(torch.zeros(1, 128)) + mel_lengths.append(1) + n_audio_tokens_list.append(0) + continue + waveform = waveform[:max_samples] + mel = compute_mel_features(waveform, sampling_rate=16000) + mel_list.append(mel) + mel_len = mel.shape[0] + mel_lengths.append(mel_len) + # Compute encoder output length from mel frame count using + # BridgeSoundEncoder._compute_output_lengths formula: + # Conv2D subsampling: floor((L + 2*padding - kernel_size) / stride + 1) + # applied log2(subsampling_factor)=3 times, kernel=3, stride=2, padding=1 + import math as _math + + token_len = float(mel_len) + for _ in range(3): + token_len = _math.floor((token_len + 2 * 1 - 3) / 2 + 1) + n_audio_tokens_list.append(max(1, int(token_len))) + + max_mel_len = max(mel_lengths) + padded_mels = torch.zeros(len(examples), max_mel_len, mel_list[0].shape[-1]) + for i, mel in enumerate(mel_list): + padded_mels[i, : mel.shape[0]] = mel + mel_lengths_t = torch.tensor(mel_lengths, dtype=torch.long) + + sound_token_id = processor.tokenizer.convert_tokens_to_ids("") + + new_input_ids_list = [] + for i, ex in enumerate(examples): + ids = batch["input_ids"][i] + n_tokens = n_audio_tokens_list[i] + if n_tokens > 0: + # Find existing token(s) and replace with correct count + sound_mask = ids == sound_token_id + existing_count = sound_mask.sum().item() + if existing_count > 0: + # Remove existing sound tokens and insert correct count at same position + first_pos = sound_mask.nonzero(as_tuple=True)[0][0].item() + ids_before = ids[:first_pos] + ids_after = ids[first_pos + existing_count :] + sound_tokens = torch.full((n_tokens,), sound_token_id, dtype=ids.dtype) + ids = torch.cat([ids_before, sound_tokens, ids_after]) + else: + # No existing sound token, insert at position 1 + sound_tokens = torch.full((n_tokens,), sound_token_id, dtype=ids.dtype) + ids = torch.cat([ids[:1], sound_tokens, ids[1:]]) + new_input_ids_list.append(ids) + + max_len = max(ids.shape[0] for ids in new_input_ids_list) + pad_id = getattr(processor.tokenizer, "pad_token_id", 0) or 0 + padded_ids = torch.full((len(examples), max_len), pad_id, dtype=new_input_ids_list[0].dtype) + for i, ids in enumerate(new_input_ids_list): + padded_ids[i, : ids.shape[0]] = ids + batch["input_ids"] = padded_ids + batch["sound_clips"] = padded_mels + batch["sound_length"] = mel_lengths_t + + # --- Loss mask (same pattern as nemotron_vl) --- + loss_mask = [ + create_multiturn_loss_mask_by_search(example, input_ids, processor, skipped_tokens) + for example, input_ids in zip(examples, batch["input_ids"]) + ] + + # --- Image token adjustment (only when images are present) --- + img_start_token_id = processor.tokenizer.convert_tokens_to_ids("") + img_end_token_id = processor.tokenizer.convert_tokens_to_ids("") + has_img_tokens = (batch["input_ids"] == img_start_token_id).any() + if has_img_tokens: + # Dynamic-res: one token per image; LM-side expansion is driven + # by per-image ``num_image_tiles`` (set below to shuffled_count_i) with + # ``img_seq_len=1``. Static-tile path keeps the HF processor's num_patches. + if is_dynamic_res_processor: + key_pv = "pixel_values_videos" if is_video else "pixel_values" + pv_ref = batch.get(key_pv) + if pv_ref is None: + n_imgs = 0 + elif isinstance(pv_ref, list): + n_imgs = len(pv_ref) + else: + n_imgs = int(pv_ref.shape[0]) + num_tiles_for_adjust = torch.ones(n_imgs, dtype=torch.long) + else: + num_tiles_for_adjust = batch.get("num_patches", torch.zeros(len(examples), dtype=torch.long)) + adjusted_batch = adjust_image_tokens( + {"input_ids": batch["input_ids"], "loss_mask": torch.tensor(loss_mask)}, + num_tiles_for_adjust, + img_start_token_id, + img_end_token_id, + ) + else: + adjusted_batch = {"input_ids": batch["input_ids"], "loss_mask": torch.tensor(loss_mask)} + + if is_video: + video_token_id = processor.tokenizer.convert_tokens_to_ids("