diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 2c56125d51..2e2efe6235 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -484,7 +484,7 @@ def process_inputs(tokenizer, processor, image_path: Optional[str], prompt: str, def _load_hf_model(args, is_vl_model: bool): - """Load HuggingFace model on rank 0. + """Load an unsharded HuggingFace model on rank 0. Args: args: Command line arguments. @@ -498,17 +498,15 @@ def _load_hf_model(args, is_vl_model: bool): print_rank_0("Loading HuggingFace model...") model_class = get_model_class(args.model_class, is_vl_model) - hf_model = model_class.from_pretrained( - args.hf_model_path, - torch_dtype=torch.bfloat16, - device_map="cuda", - trust_remote_code=is_safe_repo( + load_kwargs = { + "torch_dtype": torch.bfloat16, + "trust_remote_code": is_safe_repo( trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path, ), **_hf_revision_kwargs(args.hf_revision), - ) - hf_model = hf_model.eval() + } + hf_model = model_class.from_pretrained(args.hf_model_path, **load_kwargs).to(args.hf_device).eval() print_rank_0(f"Loaded with {model_class.__name__}") # Register debug hooks if enabled @@ -553,9 +551,11 @@ def _export_and_load_roundtrip_hf_model(args, is_vl_model: bool, megatron_model, if _is_rank_0(): print_rank_0("Loading exported HF model for comparison...") model_class = get_model_class(args.model_class, is_vl_model) - hf_model = model_class.from_pretrained( - save_path, torch_dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True - ).eval() + hf_model = ( + model_class.from_pretrained(save_path, torch_dtype=torch.bfloat16, trust_remote_code=True) + .to(args.hf_device) + .eval() + ) if args.enable_debug_hooks: print_rank_0("Registering debug hooks for exported HF model...") debugger.register_hooks(hf_model, file_prefix="hf_debug_") @@ -564,6 +564,15 @@ def _export_and_load_roundtrip_hf_model(args, is_vl_model: bool, megatron_model, return None +def _get_hf_forward_model(hf_model, pixel_values): + """Select a composite model's language backbone for text-only comparison.""" + language_model = getattr(hf_model, "language_model", None) + if pixel_values is None and isinstance(language_model, torch.nn.Module): + print_rank_0("Using the HuggingFace language backbone for a text-only comparison.") + return language_model + return hf_model + + def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokenizer, *, token_type_ids=None): """Run HuggingFace model inference and return results. @@ -583,19 +592,29 @@ def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokeniz if not _is_rank_0() or hf_model is None: return None, None, None, None, None + hf_forward_model = _get_hf_forward_model(hf_model, pixel_values) + + input_device = input_ids.device + try: + hf_device = next(hf_forward_model.parameters()).device + except (AttributeError, StopIteration, TypeError): + hf_device = input_device + if not isinstance(hf_device, (torch.device, str, int)): + hf_device = input_device + with torch.no_grad(): hf_inputs = { - "input_ids": input_ids, - "attention_mask": torch.ones_like(input_ids, dtype=torch.bool), + "input_ids": input_ids.to(hf_device), + "attention_mask": torch.ones_like(input_ids, dtype=torch.bool).to(hf_device), } if pixel_values is not None: - hf_inputs["pixel_values"] = pixel_values + hf_inputs["pixel_values"] = pixel_values.to(hf_device) if image_grid_thw is not None: - hf_inputs["image_grid_thw"] = image_grid_thw + hf_inputs["image_grid_thw"] = image_grid_thw.to(hf_device) if token_type_ids is not None: - hf_inputs["token_type_ids"] = token_type_ids + hf_inputs["token_type_ids"] = token_type_ids.to(hf_device) - hf_output = hf_model(**hf_inputs) + hf_output = hf_forward_model(**hf_inputs) # Debug: Check output type print_rank_0(f"HF output type: {type(hf_output)}") @@ -621,7 +640,13 @@ def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokeniz print_rank_0(f"HF next token: {hf_next_token.item()} ('{tokenizer.decode([hf_next_token.item()])}')") print_rank_0(f"HF Top 5: {hf_top5_info}") - return hf_logits, hf_next_token, hf_logits_stats, hf_top5_info, logits_shape + return ( + hf_logits.to(input_device), + hf_next_token.to(input_device), + hf_logits_stats, + hf_top5_info, + logits_shape, + ) def _load_hf_reference_logits(path, input_ids, tokenizer): @@ -1026,6 +1051,11 @@ def build_parser() -> argparse.ArgumentParser: 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( + "--hf-device", + default="cuda", + help="CUDA device used by the rank-0 Hugging Face reference model (for example, cuda:2).", + ) parser.add_argument( "--model_class", type=str, diff --git a/examples/model_verification_cards/nemotron-3-nano-omni-30b-a3b-reasoning/card.yaml b/examples/model_verification_cards/nemotron-3-nano-omni-30b-a3b-reasoning/card.yaml new file mode 100644 index 0000000000..815f651451 --- /dev/null +++ b/examples/model_verification_cards/nemotron-3-nano-omni-30b-a3b-reasoning/card.yaml @@ -0,0 +1,329 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: nemotron_3_nano_omni_30b_a3b_reasoning +summary: > + Performance disclaimer: this model has not been performance-tuned; reported + timing and throughput metrics are sanity checks, not optimized performance + results. Verification uses the immutable public model and CORD v2 revisions. + CPU and distributed GPU import, deterministic Megatron inference, bounded + full-model SFT, and LoRA PEFT runs completed. Strict CPU and GPU round trips + preserved all 7,349 tensors bitwise, but Transformers 5.8.0 cannot natively + reload the local custom-code exports because its dynamic-module cache omits + transitive configuration imports. The one-step HF/Megatron comparison + predicts the same token but remains below the 0.99 cosine gate. The packed + 8K long-context recipe completes one optimizer step but does not complete the + second because of H100 memory pressure. Unsupported and incomplete workflows + remain explicitly identified rather than inferred from focused unit coverage. +verification_index: + model_level: + verified: + - hf_to_megatron_cpu + - hf_to_megatron_gpu + - inference + unverified: + - megatron_to_hf_cpu + - megatron_to_hf_gpu + - manual_forward_pass + training: + H100: + verified: [sft, peft] + unverified: [sft_export_inference, sft_long_context] + unsupported: [pretrain, checkpoint_resume] +model: + hf_id: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + hf_revision: 24e67ea000b7c2837fc8f9488aa2008524fac8ba # pragma: allowlist secret + architecture: NemotronH_Nano_Omni_Reasoning_V3 + min_transformers_version: "5.8.0" +verification_environment: + base_container: nvcr.io/nvidia/nemo:26.06 + bridge_commit: fbbafc7ddfa818ef91d5eb64c8a00dcaa7bee78d # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + --hf-revision 24e67ea000b7c2837fc8f9488aa2008524fac8ba + --megatron-path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/cpu-megatron-clean + --torch-dtype bfloat16 --tp 1 --pp 1 --ep 1 --etp 1 + --trust-remote-code + last_verified: 2026-07-30 + expected_result: > + The offline, immutable-revision CPU import exits successfully after + mapping 7,333 model parameters and creates a reloadable iter_0000000 + torch_dist checkpoint with 33,015,546,816 parameters on the single model + parallel rank. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + --hf-revision 24e67ea000b7c2837fc8f9488aa2008524fac8ba + --megatron-path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/gpu-megatron-clean + --torch-dtype bfloat16 --tp 2 --pp 1 --ep 4 --etp 1 + --trust-remote-code --low-memory-save + last_verified: 2026-07-30 + expected_result: > + The command exits successfully at TP2/PP1/EP4/ETP1 and creates a + reloadable iter_0000000 checkpoint. The paired strict GPU export contains + the source checkpoint's exact 7,349-key set, shapes, dtypes, and values. + + megatron_to_hf_cpu: + status: unverified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + --hf-revision 24e67ea000b7c2837fc8f9488aa2008524fac8ba + --megatron-path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/cpu-megatron-clean/iter_0000000 + --hf-path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/cpu-hf-export-clean + --torch-dtype bfloat16 --trust-remote-code + last_verified: null + expected_result: > + CPU export completes in 14 indexed shards. Its exact comparison contains + all 7,349 source tensors (7,300 BF16, 24 int64, and 25 float32), with + identical keys, shapes, dtypes, and values and maximum difference zero. + The item remains unverified because the Transformers 5.8.0 local + custom-code loader omits transitive configuration modules before + from_pretrained can reload the otherwise bitwise-identical export. + + megatron_to_hf_gpu: + status: unverified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + --hf-revision 24e67ea000b7c2837fc8f9488aa2008524fac8ba + --megatron-path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/gpu-megatron-clean/iter_0000000 + --hf-path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/gpu-hf-export-clean + --torch-dtype bfloat16 --tp 2 --pp 1 --ep 4 --etp 1 + --trust-remote-code --distributed-save --save-every-n-ranks 1 + last_verified: null + expected_result: > + Strict export completed in 17 indexed shards. All 7,349 tensors match the + immutable HF source in keys, shapes, dtypes, and values, with maximum + difference zero. The item remains unverified because Transformers 5.8.0 + local custom-code loading omits the transitive configuration_nemotron_h + and configuration_radio modules from the model cache, preventing a native + from_pretrained reload even though those files exist in the export. + + manual_forward_pass: + status: unverified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + --hf-revision 24e67ea000b7c2837fc8f9488aa2008524fac8ba + --megatron_model_path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/gpu-megatron-clean/iter_0000000 + --tp 1 --pp 4 --ep 2 --etp 1 + --prompt "The capital of France is " --trust-remote-code + last_verified: null + expected_result: > + The pinned one-step run exits successfully and both implementations + predict token ID 6993 (" Paris"), but cosine similarity is 0.969760 and + therefore does not satisfy the required 0.99 verification gate. The + maximum and mean absolute logit differences are 3.527344 and 0.527820, + respectively. + + inference: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + 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 + --hf-revision 24e67ea000b7c2837fc8f9488aa2008524fac8ba + --megatron_model_path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/gpu-megatron-clean/iter_0000000 + --prompt "The capital of France is" --max_new_tokens 4 + --tp 1 --pp 4 --ep 2 --etp 1 + last_verified: 2026-07-30 + expected_result: |- + Deterministic greedy decoding returns the exact 4-token result and + terminates with the model end token. Literal completion: " + Paris." + + pretrain: + all: + status: unsupported + precision: null + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + Megatron Bridge does not publish a pretraining recipe for this + multimodal conditional-generation model; the supported package provides + supervised CORD v2 SFT and LoRA PEFT workflows. + + sft: + H100: + status: verified + precision: bf16 + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config + --mode sft --step_func nemotron_omni_step + --pretrained_checkpoint + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/gpu-megatron-clean/iter_0000000 + --max_steps 10 --tensor_model_parallel_size 2 + --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 + --expert_tensor_parallel_size 1 + --save_dir + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/sft-checkpoints-clean + --save_interval 10 + 'dataset.source.load_kwargs={revision:"7f0115a4b758a71d6473b8d085751692da2fef98"}' + dataset.do_validation=false dataset.do_test=false + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + logger.log_interval=1 logger.log_throughput=true rng.seed=5678 + last_verified: 2026-07-30 + metrics: + initial_loss: 1.123339 + final_loss: 0.4893276 + last_10_steps_step_time_ms_avg: 24454.11 + last_10_steps_model_tflops_per_gpu_avg: 60.73 + expected_result: > + The immutable-revision CORD v2 run completes exactly 10 full-SFT + optimizer steps on 16 H100 GPUs at TP2/PP2/CP1/EP4/ETP1, GBS/MBS + 64/1. LM loss is finite from 1.123339 to 0.4893276, all ten recorded + steps average 24,454.11 ms and 60.73 TFLOP/s/GPU including first-step + compilation, no iteration is skipped or NaN, and a complete + iter_0000010 checkpoint is saved. + + sft_export_inference: + H100: + status: unverified + precision: bf16 + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + --hf-revision 24e67ea000b7c2837fc8f9488aa2008524fac8ba + --megatron-path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/sft-checkpoints-clean/iter_0000010 + --hf-path + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/sft-hf-export-clean + --torch-dtype bfloat16 --tp 2 --pp 1 --ep 4 --etp 1 + --trust-remote-code --distributed-save --not-strict + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/sft-hf-export-clean + --prompt "The capital of France is" --max-new-tokens 32 + --chat-template --disable-thinking --trust-remote-code + --device cuda --dtype bfloat16 + last_verified: null + expected_result: > + Non-strict export succeeds and writes 6,637 indexed tensors in 17 + shards. The CORD v2 recipe intentionally disables sound, so 712 source + audio tensors are absent. Native Transformers generation remains + unverified because local dynamic-code loading omits + configuration_nemotron_h before model weights can reload. + + sft_long_context: + H100: + status: unverified + precision: bf16 + enabled_features: + sequence_packing: in_batch + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 + --recipe nemotron_omni_cord_v2_long_context_sft_8gpu_h100_bf16_config + --mode sft --step_func nemotron_omni_step + --pretrained_checkpoint + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/gpu-megatron-clean/iter_0000000 + --max_steps 10 + --save_dir + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/long-context-checkpoints-clean + --save_interval 10 + 'dataset.source.load_kwargs={revision:"7f0115a4b758a71d6473b8d085751692da2fef98"}' + dataset.do_validation=false dataset.do_test=false + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + logger.log_interval=1 logger.log_throughput=true rng.seed=5678 + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The 8K TP4/PP1/CP2/EP1/ETP4, MBS2 in-batch-packing run uses + precision-aware Adam with FP16 main parameters and stored FP32 + remainders, BF16 gradients, and BF16 moments. Step 1 completes with + finite LM loss 1.142561 in 170,179.0 ms at 8.6 TFLOP/s/GPU with no + skipped or NaN iteration, but step 2 encounters rank-divergent H100 + memory exhaustion and does not produce the required 10-step checkpoint. + + peft: + H100: + status: verified + precision: bf16 + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 + --recipe nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config + --mode lora --step_func nemotron_omni_step + --pretrained_checkpoint + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/gpu-megatron-clean/iter_0000000 + --max_steps 10 + --save_dir + work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/peft-checkpoints-clean + --save_interval 10 + 'dataset.source.load_kwargs={revision:"7f0115a4b758a71d6473b8d085751692da2fef98"}' + dataset.do_validation=false dataset.do_test=false + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + logger.log_interval=1 logger.log_throughput=true rng.seed=5678 + last_verified: 2026-07-30 + metrics: + initial_loss: 1.098394 + final_loss: 0.3166811 + last_10_steps_step_time_ms_avg: 41897.41 + last_10_steps_model_tflops_per_gpu_avg: 22.22 + expected_result: > + The ten-step TP4/PP1/CP1/EP1 LoRA run exits successfully and saves a + complete eight-shard iter_0000010 adapter checkpoint. LM loss is finite + from 1.098394 to 0.3166811, all ten steps average 41,897.41 ms and + 22.22 TFLOP/s/GPU including first-step compilation, and no iteration is + skipped or NaN. + + checkpoint_resume: + all: + status: unsupported + precision: null + depends_on: pretrain + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + resume_comparison: null + expected_result: > + There is no supported pretraining reference workflow for this model, so + no model-wide optimizer and RNG checkpoint-resume contract is available + for sentinel comparison. diff --git a/examples/models/nemotron/nemotron_3_omni/README.md b/examples/models/nemotron/nemotron_3_omni/README.md index 005c33bb70..a66d6b9e6e 100644 --- a/examples/models/nemotron/nemotron_3_omni/README.md +++ b/examples/models/nemotron/nemotron_3_omni/README.md @@ -9,6 +9,17 @@ dynamic-resolution RADIO vision tower and a Parakeet sound encoder. |---|---|---| | Nemotron-3-Nano-Omni-30B-A3B-Reasoning | `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16` | MoE hybrid LM (Mamba+attn) + RADIO vision + Parakeet audio | +AutoBridge, all shipped recipes, Direct-HF collation, Energon collation, and +the inference examples use the canonical processor-expanded +`NemotronOmniModel` path. The collator owns dense padding and complete MCore +THD packing; the model only inserts media embeddings and applies CP/SP +sharding. + +The historical `NemotronOmniLlavaModel`, `NemotronOmniLlavaModelProvider`, +`NemotronOmniLlavaBridge`, and `nemotron_omni_llava_collate_fn` collapse/expand +path remains available only for compatible legacy checkpoints and is +deprecated. Selecting it emits a `FutureWarning`. + > **Verified hardware:** all conversion, inference, and training flows in > this directory have been verified on **NVIDIA H100 80GB** nodes with 8 > GPUs per node. Other GPU SKUs may work but have not been tested. @@ -192,10 +203,19 @@ All training scripts use the Nemotron-3-Nano-Omni-30B-A3B-Reasoning pretrained checkpoint and enable in-batch sequence packing via `dataset.enable_in_batch_packing=True`. Default GPU layout per script: -For the canonical expanded-sequence image path, the collator owns THD packing. -The model receives the final packed tensors and global boundaries, inserts -image embeddings without changing sequence length, and then selects its -rank-local context-parallel shard. +The canonical expanded-sequence collator owns THD packing for text, image, +video, and audio rows. The model receives the final packed tensors and global +boundaries, inserts media without changing sequence length, and applies only +the rank-local CP/SP shard. Alignment gaps are carried as a padding mask for +media validation and consistent CP/SP localization, but the mask is not +forwarded into MCore until +[Megatron-LM #6111](https://github.com/NVIDIA/Megatron-LM/issues/6111) is fixed. +The gaps remain loss-masked but currently count toward MoE routing statistics. + +Compact variable-length packs do not yet restore the original per-row batch +dimension for `seq_aux_loss`; that requires follow-up boundary-aware +unflattening in Megatron-Core. Attention and Mamba sequence boundaries remain +per row in the current implementation. - **Full SFT** — 2 nodes / 16 GPUs (full optimizer state for ~33 B params) - **LoRA PEFT** — 1 node / 8 GPUs diff --git a/examples/models/nemotron/nemotron_3_omni/cord_v2_inference.py b/examples/models/nemotron/nemotron_3_omni/cord_v2_inference.py index 8e641e21eb..681337ea56 100644 --- a/examples/models/nemotron/nemotron_3_omni/cord_v2_inference.py +++ b/examples/models/nemotron/nemotron_3_omni/cord_v2_inference.py @@ -41,7 +41,7 @@ from megatron.bridge import AutoBridge from megatron.bridge.data.sources.hf import HFDatasetSourceConfig, load_and_adapt_hf_dataset from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import ( - inference_merged_sequence_length, + inference_expanded_image_token_counts, inference_num_image_tiles, select_inference_next_token, ) @@ -78,7 +78,7 @@ class SingleBatchIterator: 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_image_tiles", "vision_packed_seq_params"): + for key in ("images", "imgs_sizes", "vision_packed_seq_params"): if kwargs.get(key) is not None: self.batch[key] = kwargs[key] self._yielded = False @@ -106,7 +106,7 @@ def vlm_forward_step(data_iterator, model, **_): 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_image_tiles", "vision_packed_seq_params"): + for key in ("imgs_sizes", "vision_packed_seq_params"): if key in batch: forward_args[key] = batch[key] @@ -131,18 +131,16 @@ def prepare_image_sample(tokenizer, processor, image, prompt, system_prompt=None 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("") pixel_values = inputs.pixel_values num_patches = getattr(inputs, "num_patches", None) if num_patches is None: - num_patches = torch.ones(len(pixel_values), dtype=torch.long) + # This helper processes exactly one source image, so every returned + # processor tile belongs to its single wrapper. + num_patches = torch.tensor([len(pixel_values)], dtype=torch.long) else: num_patches = torch.as_tensor(num_patches, dtype=torch.long).reshape(-1) - 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 every processor tile into RADIO's packed dynamic-resolution path. P = _VISION_PATCH_DIM patches = [] @@ -158,16 +156,20 @@ def prepare_image_sample(tokenizer, processor, image, prompt, system_prompt=None sizes.append([height, width]) pv_patched = torch.cat(patches).unsqueeze(0).contiguous().bfloat16() imgs_sizes = torch.tensor(sizes, dtype=torch.long) - num_image_tiles = inference_num_image_tiles(imgs_sizes, patch_dim=P) + tile_feature_counts = inference_num_image_tiles(imgs_sizes, patch_dim=P) + expanded_counts = inference_expanded_image_token_counts(tile_feature_counts, num_patches) + if img_start_id != tokenizer.unk_token_id and (input_ids == img_start_id).any(): + input_ids = adjust_image_tokens(input_ids, expanded_counts, img_start_id, img_end_id) image_token_id = tokenizer.convert_tokens_to_ids("") num_placeholders = int((input_ids == image_token_id).sum().item()) - if num_image_tiles.numel() != num_placeholders: + expected_placeholders = int(expanded_counts.sum().item()) + if num_placeholders != expected_placeholders: raise ValueError( - "Vision metadata produced " - f"{num_image_tiles.numel()} replacement counts for {num_placeholders} image placeholders." + f"Vision metadata requires {expected_placeholders} expanded image placeholders; " + f"the prompt contains {num_placeholders}." ) - return input_ids, pv_patched, imgs_sizes, num_image_tiles + return input_ids, pv_patched, imgs_sizes @torch.no_grad() @@ -177,7 +179,6 @@ def generate( input_ids, images, imgs_sizes, - num_image_tiles, *, sequence_length, max_new_tokens=200, @@ -188,7 +189,6 @@ def generate( input_ids = input_ids.cuda() images = images.cuda() imgs_sizes = imgs_sizes.cuda() - num_image_tiles = num_image_tiles.cuda() position_ids = ( torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device).unsqueeze(0).expand_as(input_ids) @@ -196,7 +196,6 @@ def generate( attention_mask = torch.ones_like(input_ids, dtype=torch.bool) generated_ids = input_ids.clone() stop_tokens = {tokenizer.eos_token_id} - image_token_id = tokenizer.convert_tokens_to_ids("") fwd_bwd = get_forward_backward_func() for _ in range(max_new_tokens): @@ -209,7 +208,6 @@ def generate( attention_mask, images=images, imgs_sizes=imgs_sizes, - num_image_tiles=num_image_tiles, vision_packed_seq_params=vision_packed_seq_params, ) output = fwd_bwd( @@ -232,12 +230,7 @@ def generate( 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) - merged_sequence_length = inference_merged_sequence_length( - input_ids, - image_token_index=image_token_id, - num_image_tiles=num_image_tiles, - image_seq_len=1, - ) + merged_sequence_length = input_ids.shape[1] next_token_ids = select_inference_next_token(full, merged_sequence_length) else: next_token_ids = torch.ones((1, 1), device=generated_ids.device, dtype=generated_ids.dtype) @@ -303,6 +296,9 @@ def main(): model_provider.separate_video_embedder = True model_provider.temporal_ckpt_compat = True model_provider.vision_class_token_len = 10 + # Canonical multimodal inputs retain their actual expanded length. PP + # stages therefore need shape exchange instead of fixed receive buffers. + model_provider.variable_seq_lengths = args.pp > 1 model_provider.initialize_model_parallel(seed=0) if args.megatron_model_path: @@ -319,6 +315,7 @@ def main(): "separate_video_embedder": True, "temporal_ckpt_compat": True, "vision_class_token_len": 10, + "variable_seq_lengths": args.pp > 1, }, wrap_with_ddp=False, ) @@ -369,9 +366,7 @@ def main(): except Exception as e: print_rank_0(f"WARN: could not save sample image {i}: {e}") - input_ids, pv_patched, imgs_sizes, num_image_tiles = prepare_image_sample( - tokenizer, processor, image, args.prompt - ) + input_ids, pv_patched, imgs_sizes = prepare_image_sample(tokenizer, processor, image, args.prompt) cleaned, prediction_full = generate( model, @@ -379,7 +374,6 @@ def main(): input_ids, pv_patched, imgs_sizes, - num_image_tiles, sequence_length=model_provider.seq_length, max_new_tokens=args.max_new_tokens, ) 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 index 05e18fb5e2..bd12d89175 100644 --- 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 @@ -24,11 +24,10 @@ * Image: 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 + dynamic-resolution path. The ``imgs_sizes`` / ``vision_packed_seq_params`` tensors are built from the - per-tile shapes. ``num_image_tiles`` carries each tile's exact replacement - count to every pipeline stage (256 tokens for a 512x512 tile after - pixel_shuffle). + per-tile shapes, and the prompt is expanded to one image placeholder per + projected RADIO feature. * Audio / text-only: temporal_patch_dim=1 (the vision encoder is unused). * Video (and video+audio): temporal_patch_dim=2, separate_video_embedder=True, temporal_ckpt_compat=True so RADIO ViT @@ -36,7 +35,8 @@ the shared SFT collator (used by `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. + / `num_frames` / `vision_packed_seq_params` are plumbed through to the + canonical ``NemotronOmniModel``. Examples: # Single image: @@ -84,7 +84,7 @@ from megatron.bridge import AutoBridge from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import ( COMPACT_IMAGE_PLACEHOLDER, - inference_merged_sequence_length, + inference_expanded_image_token_counts, inference_num_image_tiles, patchify_temporal_frame, select_inference_next_token, @@ -204,8 +204,6 @@ def __init__(self, input_ids, position_ids, attention_mask, **kwargs): 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("num_image_tiles", None) is not None: - self.batch["num_image_tiles"] = kwargs["num_image_tiles"] if kwargs.get("vision_packed_seq_params", None) is not None: self.batch["vision_packed_seq_params"] = kwargs["vision_packed_seq_params"] @@ -248,7 +246,7 @@ def vlm_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: elif "pixel_values" in batch: forward_args["pixel_values"] = batch["pixel_values"] - # LLaVAModel.forward() requires `images` even for audio-only inference + # Keep the empty-image sentinel used by the training step for audio/text. 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) @@ -263,8 +261,6 @@ def vlm_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: forward_args["imgs_sizes"] = batch["imgs_sizes"] if "num_frames" in batch: forward_args["num_frames"] = batch["num_frames"] - if "num_image_tiles" in batch: - forward_args["num_image_tiles"] = batch["num_image_tiles"] if "vision_packed_seq_params" in batch: forward_args["vision_packed_seq_params"] = batch["vision_packed_seq_params"] @@ -272,7 +268,7 @@ def loss_func(x, **kwargs): return x output = model(**forward_args) - # LlavaModel returns (logits, loss_mask) tuple; pipeline expects a single tensor + # CP training can return (logits, loss_mask); generation needs only logits. if isinstance(output, tuple): output = output[0] return output, loss_func @@ -295,21 +291,47 @@ def load_image(image_path: str) -> Image.Image: 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. +def _patchify_pixel_values( + pv: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...], + patch_dim: int = _VISION_PATCH_DIM, +): + """Pack 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``. + Dynamic-resolution processors return either one stacked ``[N, 3, H, W]`` + tensor when all tiles have the same size, or a list of ``[3, H, W]`` / + ``[N, 3, H, W]`` tensors when tile sizes differ. Normalize both forms and + concatenate their patches so RADIO sees one packed sequence matching + ``imgs_sizes``. """ + if isinstance(pv, torch.Tensor): + pixel_groups = [pv] + elif isinstance(pv, (list, tuple)): + pixel_groups = list(pv) + else: + raise TypeError(f"pixel_values must be a tensor or sequence of tensors, got {type(pv).__name__}") + + tiles = [] + for group in pixel_groups: + if not isinstance(group, torch.Tensor): + raise TypeError(f"Each pixel_values entry must be a tensor, got {type(group).__name__}") + if group.ndim == 3: + tiles.append(group) + elif group.ndim == 4: + tiles.extend(group.unbind(0)) + else: + raise ValueError(f"Each pixel_values tensor must be 3D or 4D, got shape {tuple(group.shape)}") + if not tiles: + raise ValueError("pixel_values must contain at least one image tile") + P = patch_dim patches_list = [] sizes = [] - for i in range(pv.shape[0]): - _, H, W = pv[i].shape + for tile in tiles: + _, H, W = tile.shape + if H % P != 0 or W % P != 0: + raise ValueError(f"Image tile shape {(H, W)} is not divisible by patch_dim={P}") 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) + p = tile.unsqueeze(0).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) @@ -344,18 +366,29 @@ def process_image_inputs( 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] + pixel_values = inputs.pixel_values + # Pre-patchify for the packed dynamic-resolution RADIO path before + # validating ownership because heterogeneous tiles may be a list. + packed_pv, sizes = _patchify_pixel_values(pixel_values) + tile_count = len(sizes) if hasattr(inputs, "num_patches") and inputs.num_patches is not None: - num_patches = inputs.num_patches + num_patches = torch.as_tensor(inputs.num_patches, dtype=torch.long).reshape(-1) + if num_patches.numel() != len(images) or int(num_patches.sum().item()) != tile_count: + raise ValueError( + "num_patches must provide one ownership count per source image and sum to the number of tiles" + ) else: - num_patches = torch.ones(pixel_values.shape[0], dtype=torch.int) + if tile_count != len(images): + raise ValueError( + "The image processor returned multiple tiles per source image without num_patches ownership " + "metadata." + ) + num_patches = torch.ones(len(images), dtype=torch.long) - # 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"Image: {image_path}, tiles={tile_count}, " f"packed_shape={tuple(packed_pv.shape)}, num_patches={num_patches.tolist()}" ) return inputs.input_ids, packed_pv, num_patches, imgs_sizes @@ -583,6 +616,11 @@ def process_video_audio_inputs( return input_ids, packed_pixel_values, num_patches, imgs_sizes, num_frames, sound_clips, sound_length +def _hf_revision_kwargs(revision: str | None) -> dict[str, str]: + """Build keyword arguments for revision-pinned Hugging Face loads.""" + return {"revision": revision} if revision is not None else {} + + def main(args) -> None: """Main function for Nemotron Omni VL generation from HuggingFace models. @@ -605,11 +643,10 @@ def main(args) -> None: # unused for text/audio-only inference. # # Image: 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 + # runs the packed dynamic-resolution path. 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 supplies exact post-shuffle - # replacement counts to PP stages without a vision encoder. + # vision_packed_seq_params. The prompt is expanded before model forward + # so every pipeline stage sees the canonical sequence length. # # Video (and video+audio): temporal_patch_dim=2, # separate_video_embedder=True so RADIO exercises the trained @@ -642,7 +679,11 @@ def main(args) -> None: # 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) + bridge = AutoBridge.from_hf_pretrained( + args.hf_model_path, + trust_remote_code=True, + **_hf_revision_kwargs(args.hf_revision), + ) # Initialize model parallel before loading model_provider = bridge.to_megatron_provider(load_weights=False) @@ -654,6 +695,9 @@ def main(args) -> None: model_provider.temporal_patch_dim = temporal_patch_dim model_provider.separate_video_embedder = separate_video_embedder model_provider.temporal_ckpt_compat = temporal_ckpt_compat + # Canonical multimodal inputs retain their actual expanded length. + # PP stages therefore need shape exchange instead of fixed receive buffers. + model_provider.variable_seq_lengths = pp > 1 model_provider.initialize_model_parallel(seed=0) # Load the Megatron model directly. The mp_overrides values are applied to @@ -672,13 +716,18 @@ def main(args) -> None: "temporal_patch_dim": temporal_patch_dim, "separate_video_embedder": separate_video_embedder, "temporal_ckpt_compat": temporal_ckpt_compat, + "variable_seq_lengths": pp > 1, }, wrap_with_ddp=False, ) 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) + bridge = AutoBridge.from_hf_pretrained( + args.hf_model_path, + trust_remote_code=True, + **_hf_revision_kwargs(args.hf_revision), + ) model_provider = bridge.to_megatron_provider(load_weights=True) model_provider.tensor_model_parallel_size = tp model_provider.pipeline_model_parallel_size = pp @@ -688,6 +737,7 @@ def main(args) -> None: 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.variable_seq_lengths = pp > 1 model_provider.initialize_model_parallel(seed=0) model_provider.finalize() model = model_provider.provide_distributed_model(wrap_with_ddp=False) @@ -709,8 +759,16 @@ def main(args) -> None: 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) + tokenizer = AutoTokenizer.from_pretrained( + args.hf_model_path, + trust_remote_code=True, + **_hf_revision_kwargs(args.hf_revision), + ) + processor = AutoProcessor.from_pretrained( + args.hf_model_path, + trust_remote_code=True, + **_hf_revision_kwargs(args.hf_revision), + ) img_start_token_id = tokenizer.convert_tokens_to_ids("") img_end_token_id = tokenizer.convert_tokens_to_ids("") image_token_id = tokenizer.convert_tokens_to_ids("") @@ -723,7 +781,6 @@ def main(args) -> None: images = None imgs_sizes = None num_frames = None - num_image_tiles = None vision_packed_seq_params = None if args.video_path and args.audio_path: @@ -755,28 +812,32 @@ def main(args) -> None: 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). + tile_feature_counts = inference_num_image_tiles( + imgs_sizes, + patch_dim=_VISION_PATCH_DIM, + num_frames=num_frames, + temporal_patch_size=temporal_patch_dim, + ) + expanded_counts = inference_expanded_image_token_counts( + tile_feature_counts, + num_patches, + feature_multiplier=image_seq_len, + ) + # Normalize processor output to the canonical one-placeholder-per- + # projected-feature contract. 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) - num_image_tiles = inference_num_image_tiles( - imgs_sizes, - patch_dim=_VISION_PATCH_DIM, - num_frames=num_frames, - temporal_patch_size=temporal_patch_dim, - ) + input_ids = adjust_image_tokens(input_ids, expanded_counts, img_start_token_id, img_end_token_id) num_placeholders = int((input_ids == image_token_id).sum().item()) - if num_image_tiles.numel() != num_placeholders: + expected_placeholders = int(expanded_counts.sum().item()) + if num_placeholders != expected_placeholders: raise ValueError( - "Vision metadata produced " - f"{num_image_tiles.numel()} replacement counts for {num_placeholders} image placeholders." + f"Vision metadata requires {expected_placeholders} expanded image placeholders; " + f"the prompt contains {num_placeholders}." ) pixel_values = None @@ -792,8 +853,6 @@ def main(args) -> None: imgs_sizes = imgs_sizes.cuda() if num_frames is not None: num_frames = num_frames.cuda() - if num_image_tiles is not None: - num_image_tiles = num_image_tiles.cuda() position_ids = ( torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device).unsqueeze(0).expand_as(input_ids) @@ -828,7 +887,6 @@ def main(args) -> None: sound_length=sound_length, imgs_sizes=imgs_sizes, num_frames=num_frames, - num_image_tiles=num_image_tiles, vision_packed_seq_params=vision_packed_seq_params, ) @@ -838,8 +896,7 @@ def main(args) -> None: model=model, num_microbatches=1, forward_only=True, - # LLaVA pads PP activations to the configured model width, so - # pipeline receive buffers must use that same fixed length. + # Ignored for PP shape allocation when variable_seq_lengths is enabled. seq_length=model_provider.seq_length, micro_batch_size=1, collect_non_loss_data=True, @@ -854,12 +911,7 @@ def main(args) -> None: 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) - merged_sequence_length = inference_merged_sequence_length( - input_ids, - image_token_index=image_token_id, - num_image_tiles=num_image_tiles, - image_seq_len=image_seq_len, - ) + merged_sequence_length = input_ids.shape[1] next_token_ids = select_inference_next_token(output, merged_sequence_length) if step < 5: @@ -910,6 +962,12 @@ def main(args) -> None: default="nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", help="Path to the HuggingFace Nemotron Omni VL model.", ) + parser.add_argument( + "--hf-revision", + dest="hf_revision", + default=None, + help="Immutable Hugging Face Hub revision used for model, tokenizer, and processor loading.", + ) parser.add_argument( "--prompt", type=str, diff --git a/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py b/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py index 4f39e96907..d65c746c27 100644 --- a/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py +++ b/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py @@ -14,8 +14,8 @@ Vision backbone uses the dynamic-resolution temporal video embedder path (``temporal_patch_dim=2``, ``separate_video_embedder=True``), -matching the shared SFT pipeline in ``nemotron_omni_collate_fn`` with -use_temporal_video_embedder=True. Frames are pre-patchified into a packed +matching the shared SFT pipeline in ``nemotron_omni_expanded_collate_fn`` 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`. @@ -44,13 +44,14 @@ from megatron.bridge import AutoBridge from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import ( COMPACT_IMAGE_PLACEHOLDER, - inference_merged_sequence_length, + inference_expanded_image_token_counts, inference_num_image_tiles, patchify_temporal_frame, select_inference_next_token, temporal_model_frames, ) from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import ( + adjust_image_tokens, maybe_path_or_url_to_data_urls, pil_image_from_base64, ) @@ -113,8 +114,6 @@ def __init__(self, input_ids, position_ids, attention_mask, **kwargs): 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("num_image_tiles") is not None: - self.batch["num_image_tiles"] = kwargs["num_image_tiles"] if kwargs.get("vision_packed_seq_params") is not None: self.batch["vision_packed_seq_params"] = kwargs["vision_packed_seq_params"] self._yielded = False @@ -151,8 +150,6 @@ def vlm_forward_step(data_iterator, model, **kwargs): forward_args["imgs_sizes"] = batch["imgs_sizes"] if "num_frames" in batch: forward_args["num_frames"] = batch["num_frames"] - if "num_image_tiles" in batch: - forward_args["num_image_tiles"] = batch["num_image_tiles"] if "vision_packed_seq_params" in batch: forward_args["vision_packed_seq_params"] = batch["vision_packed_seq_params"] @@ -285,6 +282,18 @@ def process_sample( "Vision metadata produced " f"{num_image_tiles.numel()} replacement counts for {num_placeholders} image placeholders." ) + image_seq_len = (_VIDEO_FRAME_H // _VISION_PATCH_DIM) * (_VIDEO_FRAME_W // _VISION_PATCH_DIM) // 4 + expanded_counts = inference_expanded_image_token_counts( + num_image_tiles, + torch.ones_like(num_image_tiles), + feature_multiplier=image_seq_len, + ) + input_ids = adjust_image_tokens( + input_ids, + expanded_counts, + tokenizer.convert_tokens_to_ids(""), + tokenizer.convert_tokens_to_ids(""), + ) # Process audio sound_clips = None @@ -334,7 +343,6 @@ def process_sample( "images": images, "imgs_sizes": imgs_sizes, "num_frames": num_frames, - "num_image_tiles": num_image_tiles, "sound_clips": sound_clips, "sound_length": sound_length, "question": qa["question"], @@ -375,15 +383,11 @@ def generate(model, tokenizer, sample, *, sequence_length, max_new_tokens=50): 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 - num_image_tiles = sample["num_image_tiles"].cuda() if sample.get("num_image_tiles") 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] - image_token_id = tokenizer.convert_tokens_to_ids("") - image_seq_len = (_VIDEO_FRAME_H // _VISION_PATCH_DIM) * (_VIDEO_FRAME_W // _VISION_PATCH_DIM) // 4 - 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, @@ -399,7 +403,6 @@ def generate(model, tokenizer, sample, *, sequence_length, max_new_tokens=50): sound_length=sound_length, imgs_sizes=imgs_sizes, num_frames=num_frames, - num_image_tiles=num_image_tiles, vision_packed_seq_params=vision_packed_seq_params, ) output = fwd_bwd_function( @@ -408,8 +411,7 @@ def generate(model, tokenizer, sample, *, sequence_length, max_new_tokens=50): model=model, num_microbatches=1, forward_only=True, - # LLaVA pads PP activations to the configured model width, so - # pipeline receive buffers must use that same fixed length. + # Ignored for PP shape allocation when variable_seq_lengths is enabled. seq_length=sequence_length, micro_batch_size=1, collect_non_loss_data=True, @@ -424,12 +426,7 @@ def generate(model, tokenizer, sample, *, sequence_length, max_new_tokens=50): 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) - merged_sequence_length = inference_merged_sequence_length( - input_ids, - image_token_index=image_token_id, - num_image_tiles=num_image_tiles, - image_seq_len=image_seq_len, - ) + merged_sequence_length = input_ids.shape[1] next_token_ids = select_inference_next_token(output, merged_sequence_length) else: next_token_ids = torch.ones((1, 1), device=generated_ids.device, dtype=generated_ids.dtype) @@ -497,6 +494,9 @@ def main(): model_provider.separate_video_embedder = True model_provider.temporal_ckpt_compat = True model_provider.vision_class_token_len = 10 + # Canonical multimodal inputs retain their actual expanded length. PP + # stages therefore need shape exchange instead of fixed receive buffers. + model_provider.variable_seq_lengths = args.pp > 1 model_provider.initialize_model_parallel(seed=0) if args.megatron_model_path: @@ -513,6 +513,7 @@ def main(): "separate_video_embedder": True, "temporal_ckpt_compat": True, "vision_class_token_len": 10, + "variable_seq_lengths": args.pp > 1, }, wrap_with_ddp=False, ) diff --git a/scripts/conversion/utils.py b/scripts/conversion/utils.py index 1d9cdfdcdb..017cc88f2d 100644 --- a/scripts/conversion/utils.py +++ b/scripts/conversion/utils.py @@ -13,6 +13,7 @@ # limitations under the License. """Utilities shared by CPU and distributed GPU conversion backends.""" +import re import shutil from collections.abc import Iterable from pathlib import Path @@ -50,6 +51,8 @@ def resolve_hf_commit_revision(hf_model: str, hf_revision: str | None) -> str | _validate_hf_revision_target(hf_model, hf_revision) if hf_revision is None: return None + if re.fullmatch(r"[0-9a-f]{40}", hf_revision): + return hf_revision from huggingface_hub import HfApi diff --git a/src/megatron/bridge/data/builders/energon.py b/src/megatron/bridge/data/builders/energon.py index f7a3da376d..bace68db0d 100644 --- a/src/megatron/bridge/data/builders/energon.py +++ b/src/megatron/bridge/data/builders/energon.py @@ -102,6 +102,8 @@ class NemotronOmniEnergonTaskEncoderConfig: ``visual_keys`` is retained for configuration compatibility, but Omni owns its visual input contract and supports only ``("pixel_values",)``. + ``collapse_image_tokens=True`` selects the deprecated LLaVA compatibility + path; the default ``False`` selects the canonical expanded-sequence path. """ hf_processor_path: str @@ -113,6 +115,7 @@ class NemotronOmniEnergonTaskEncoderConfig: video_nframes: int use_temporal_video_embedder: bool patch_dim: int + collapse_image_tokens: bool = False trust_remote_code: bool | None = None def validate(self) -> None: @@ -288,6 +291,7 @@ def build_energon_task_encoder(config: EnergonDatasetConfig) -> Any: video_nframes=task_config.video_nframes, use_temporal_video_embedder=task_config.use_temporal_video_embedder, patch_dim=task_config.patch_dim, + collapse_image_tokens=task_config.collapse_image_tokens, pad_to_max_length=config.pad_to_max_length, pad_to_multiple_of=config.pad_to_multiple_of, enable_in_batch_packing=effective_packing, diff --git a/src/megatron/bridge/data/collators/registry.py b/src/megatron/bridge/data/collators/registry.py index 1cefa98c1a..6fd485b555 100644 --- a/src/megatron/bridge/data/collators/registry.py +++ b/src/megatron/bridge/data/collators/registry.py @@ -41,7 +41,7 @@ class _ModelCollateSpec: ), "NemotronH_Nano_Omni_Reasoning_V3Processor": _ModelCollateSpec( "megatron.bridge.models.nemotron_omni.data.collate_fn", - "nemotron_omni_collate_fn", + "nemotron_omni_expanded_collate_fn", required_for_all_examples=True, ), "PixtralProcessor": _ModelCollateSpec( diff --git a/src/megatron/bridge/data/collators/sequence.py b/src/megatron/bridge/data/collators/sequence.py index 6aaf48000d..4995755e5e 100644 --- a/src/megatron/bridge/data/collators/sequence.py +++ b/src/megatron/bridge/data/collators/sequence.py @@ -33,6 +33,7 @@ def prepare_sequence_batch( pad_token_id: int = 0, ignore_index: int = IGNORE_INDEX, sequence_tensor_pad_values: Mapping[str, int | float] | None = None, + emit_packed_padding_mask: bool = False, ) -> None: """Apply the collator's explicit padded or in-batch-packed output policy.""" if enable_in_batch_packing: @@ -43,6 +44,7 @@ def prepare_sequence_batch( ignore_index=ignore_index, pad_to_multiple_of=in_batch_packing_pad_to_multiple_of, sequence_tensor_pad_values=sequence_tensor_pad_values, + emit_padding_mask=emit_packed_padding_mask, ) return pad_or_truncate_sequence_batch( diff --git a/src/megatron/bridge/data/energon/hf_task_encoder.py b/src/megatron/bridge/data/energon/hf_task_encoder.py index 9759a1c30b..d0a94d2ec2 100644 --- a/src/megatron/bridge/data/energon/hf_task_encoder.py +++ b/src/megatron/bridge/data/energon/hf_task_encoder.py @@ -59,6 +59,7 @@ class HFEnergonBatch(Batch): position_ids: torch.Tensor = field(default_factory=lambda: torch.empty(0)) # [B, seq_len] visual_inputs: GenericVisualInputs | None = None attention_mask: torch.Tensor | None = None + padding_mask: torch.Tensor | None = None cu_seqlens_q: torch.Tensor | None = None cu_seqlens_kv: torch.Tensor | None = None cu_seqlens_q_padded: torch.Tensor | None = None @@ -198,6 +199,7 @@ def _collate_batch_kwargs(self, samples: List[HFEnergonSample]) -> tuple[dict[st labels=collated["labels"], loss_mask=collated["loss_mask"], attention_mask=collated.get("attention_mask"), + padding_mask=collated.get("padding_mask"), position_ids=collated["position_ids"], visual_inputs=collated.get("visual_inputs"), cu_seqlens_q=collated.get("cu_seqlens_q"), diff --git a/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py b/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py index 55626ca5aa..6e4aae65a1 100644 --- a/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py +++ b/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py @@ -23,7 +23,8 @@ from megatron.bridge.data.energon.hf_task_encoder import HFEnergonBatch, HFEnergonSample, HFTaskEncoder from megatron.bridge.models.nemotron_omni.data.collate_fn import ( _validate_nemotron_omni_visual_keys, - nemotron_omni_collate_fn, + nemotron_omni_expanded_collate_fn, + nemotron_omni_llava_collate_fn, ) from megatron.bridge.training.utils.visual_inputs import GenericVisualInputs @@ -125,8 +126,9 @@ class NemotronOmniTaskEncoder(HFTaskEncoder): The task encoder owns only source adaptation and configuration. Tokenization, assistant masking, modality-token expansion, padding, and in-batch packing - are performed by :func:`nemotron_omni_collate_fn` for both Direct-HF and - Energon datasets. + are performed by the canonical expanded-sequence collator for both + Direct-HF and Energon datasets. ``collapse_image_tokens=True`` selects the + deprecated LLaVA compatibility contract. """ def __init__( @@ -145,13 +147,15 @@ def __init__( pad_to_multiple_of: int = 128, enable_in_batch_packing: bool = False, in_batch_packing_pad_to_multiple_of: int = 1, + collapse_image_tokens: bool = False, ) -> None: _validate_nemotron_omni_visual_keys(visual_keys) + collate_fn = nemotron_omni_llava_collate_fn if collapse_image_tokens else nemotron_omni_expanded_collate_fn super().__init__( processor=processor, seq_length=seq_length, visual_keys=visual_keys, - collate_fn=nemotron_omni_collate_fn, + collate_fn=collate_fn, pad_to_max_length=pad_to_max_length, pad_to_multiple_of=pad_to_multiple_of, enable_in_batch_packing=enable_in_batch_packing, @@ -164,6 +168,7 @@ def __init__( self.video_nframes = video_nframes self.use_temporal_video_embedder = use_temporal_video_embedder self.patch_dim = patch_dim + self.collapse_image_tokens = collapse_image_tokens def collate_fn(self, examples: list[dict[str, Any]]) -> dict[str, Any]: """Collate normalized Energon examples with the shared Omni path.""" diff --git a/src/megatron/bridge/data/packing/in_batch.py b/src/megatron/bridge/data/packing/in_batch.py index 701b55ec27..24937e7d33 100644 --- a/src/megatron/bridge/data/packing/in_batch.py +++ b/src/megatron/bridge/data/packing/in_batch.py @@ -95,6 +95,7 @@ def build_mcore_thd_sequence_batch_from_rows( ignore_index: int = IGNORE_INDEX, pad_to_multiple_of: int = 1, sequence_tensor_pad_values: Mapping[str, int | float] | None = None, + emit_padding_mask: bool = False, ) -> dict[str, Any]: """Build an MCore THD batch directly from unpadded sequence rows. @@ -107,6 +108,8 @@ def build_mcore_thd_sequence_batch_from_rows( pad_to_multiple_of: Per-sequence alignment multiple for CP/SP. sequence_tensor_pad_values: Additional sequence-aligned tensor keys and the value used for alignment padding. + emit_padding_mask: Whether to emit a boolean mask whose true values + identify physical alignment gaps. Returns: A single-row THD batch with current MCore packed-sequence metadata. @@ -122,7 +125,7 @@ def build_mcore_thd_sequence_batch_from_rows( raise ValueError("sequence_length must be >= 1.") extra_pad_values = dict(sequence_tensor_pad_values or {}) - reserved_keys = {token_key, "position_ids", "labels", "loss_mask", "attention_mask"} + reserved_keys = {token_key, "position_ids", "labels", "loss_mask", "attention_mask", "padding_mask"} if reserved_keys.intersection(extra_pad_values): raise ValueError("Additional sequence tensor keys must not replace standard sequence tensors.") @@ -177,6 +180,10 @@ def build_mcore_thd_sequence_batch_from_rows( ), "attention_mask": None, } + if emit_padding_mask: + # MCore routes the physical THD stream; mask alignment gaps out of MoE + # z/aux losses and expert-bias token counts. + packed["padding_mask"] = torch.ones((1, total_length), dtype=torch.bool, device=first_tokens.device) output_pad_values: dict[str, int | float] = {"labels": ignore_index, "loss_mask": 0, **extra_pad_values} for key, pad_value in output_pad_values.items(): @@ -188,6 +195,8 @@ def build_mcore_thd_sequence_batch_from_rows( for row, length, padded_length in zip(normalized_rows, unpadded_lengths, padded_lengths): packed[token_key][0, offset : offset + length] = row[token_key] packed["position_ids"][0, offset : offset + length] = row["position_ids"] + if emit_padding_mask: + packed["padding_mask"][0, offset : offset + length] = False for key in output_pad_values: if key in packed: packed[key][0, offset : offset + length] = row[key] @@ -227,6 +236,7 @@ def pack_right_padded_sequence_batch_to_mcore_thd( ignore_index: int = IGNORE_INDEX, pad_to_multiple_of: int = 1, sequence_tensor_pad_values: Mapping[str, int | float] | None = None, + emit_padding_mask: bool = False, ) -> None: """Pack a right-padded sequence batch into MCore THD layout. @@ -245,6 +255,8 @@ def pack_right_padded_sequence_batch_to_mcore_thd( pad_to_multiple_of: Optional per-sequence packed length multiple. sequence_tensor_pad_values: Additional sequence-aligned tensor keys and their alignment padding values. + emit_padding_mask: Whether to emit a boolean mask whose true values + identify physical alignment gaps. Raises: ValueError: If required tensors are missing or the batch contains no @@ -304,6 +316,7 @@ def pack_right_padded_sequence_batch_to_mcore_thd( ignore_index=ignore_index, pad_to_multiple_of=pad_to_multiple_of, sequence_tensor_pad_values=sequence_tensor_pad_values, + emit_padding_mask=emit_padding_mask, ) _set_tokens(batch, token_key, packed.pop(token_key)) for key in ( @@ -311,6 +324,7 @@ def pack_right_padded_sequence_batch_to_mcore_thd( "loss_mask", "position_ids", "attention_mask", + "padding_mask", "cu_seqlens_q", "cu_seqlens_kv", "cu_seqlens_q_padded", @@ -322,5 +336,5 @@ def pack_right_padded_sequence_batch_to_mcore_thd( ): if key in packed: batch[key] = packed[key] - elif key in {"cu_seqlens_q_padded", "cu_seqlens_kv_padded"}: + elif key in {"padding_mask", "cu_seqlens_q_padded", "cu_seqlens_kv_padded"}: batch.pop(key, None) diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index 3685b34b6a..a6e8c0588b 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -460,6 +460,7 @@ def from_auto_config(cls, megatron_path: str, hf_model_id: str, trust_remote_cod megatron_hf_cfg_dict = _drop_readonly_config_properties(megatron_hf_cfg_dict, type(hf_cfg)) # 3. Build final bridge from the synthesized config synthesized_config = type(hf_cfg)(**megatron_hf_cfg_dict) + synthesized_config.name_or_path = hf_model_id bridge = cls.from_hf_config(synthesized_config) bridge.hf_model_id = hf_model_id bridge.trust_remote_code = trust_remote_code @@ -1447,30 +1448,34 @@ def import_ckpt( ... low_memory_save=True ... ) """ - # Load the HuggingFace model + # Load the HuggingFace model before creating temporary distributed state. bridge = cls.from_hf_pretrained(hf_model_id, **kwargs) - # Convert to Megatron model - megatron_model = bridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True) + from megatron.bridge.training.model_load_save import temporary_distributed_context + + model_context = nullcontext() if dist.is_initialized() else temporary_distributed_context(backend="gloo") + with model_context: + # Convert to Megatron model + megatron_model = bridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True) - # Save as Megatron checkpoint - hf_tokenizer_kwargs = {} - if hasattr(bridge._model_bridge, "get_hf_tokenizer_kwargs"): - hf_tokenizer_kwargs = bridge._model_bridge.get_hf_tokenizer_kwargs() - if hf_tokenizer_kwargs is None: + # Save as Megatron checkpoint hf_tokenizer_kwargs = {} - if kwargs.get("revision") is not None: - hf_tokenizer_kwargs.setdefault("revision", kwargs["revision"]) - # Forward trust_remote_code to the tokenizer (needed for repos with custom code) - if kwargs.get("trust_remote_code"): - hf_tokenizer_kwargs.setdefault("trust_remote_code", True) - bridge.save_megatron_model( - megatron_model, - megatron_path, - hf_tokenizer_path=hf_model_id, - hf_tokenizer_kwargs=hf_tokenizer_kwargs, - low_memory_save=low_memory_save, - ) + if hasattr(bridge._model_bridge, "get_hf_tokenizer_kwargs"): + hf_tokenizer_kwargs = bridge._model_bridge.get_hf_tokenizer_kwargs() + if hf_tokenizer_kwargs is None: + hf_tokenizer_kwargs = {} + if kwargs.get("revision") is not None: + hf_tokenizer_kwargs.setdefault("revision", kwargs["revision"]) + # Forward trust_remote_code to the tokenizer (needed for repos with custom code) + if kwargs.get("trust_remote_code"): + hf_tokenizer_kwargs.setdefault("trust_remote_code", True) + bridge.save_megatron_model( + megatron_model, + megatron_path, + hf_tokenizer_path=hf_model_id, + hf_tokenizer_kwargs=hf_tokenizer_kwargs, + low_memory_save=low_memory_save, + ) def export_ckpt( self, diff --git a/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py b/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py index 0606cec5bd..53b3ad8119 100644 --- a/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py +++ b/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py @@ -18,6 +18,7 @@ import copy import tempfile +import warnings from collections.abc import Mapping, Sequence from typing import Any @@ -27,6 +28,7 @@ from megatron.bridge.data.collators.sequence import prepare_sequence_batch from megatron.bridge.data.collators.sequence_padding import use_processor_right_padding from megatron.bridge.data.conversation_processing import ( + AssistantMaskBoundaryConfig, assistant_mask_boundary_config_from_markers, build_assistant_loss_mask, chat_template_kwargs_from_example, @@ -49,6 +51,35 @@ _NEMOTRON_OMNI_VISUAL_KEYS = ("pixel_values",) +def _build_padded_assistant_loss_masks( + examples: Sequence[Mapping[str, Any]], + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + processor: Any, + skipped_tokens: torch.Tensor, + *, + boundary_config: AssistantMaskBoundaryConfig, +) -> torch.Tensor: + """Build assistant loss masks without treating batch padding as message boundaries.""" + if input_ids.dim() != 2 or attention_mask.shape != input_ids.shape: + raise ValueError("Nemotron Omni assistant masking expects matching 2D input_ids and attention_mask.") + + loss_masks = [] + for example, token_row, attention_row in zip(examples, input_ids, attention_mask, strict=True): + active_positions = attention_row.to(dtype=torch.bool) + active_mask = build_assistant_loss_mask( + example, + token_row[active_positions], + processor, + skipped_tokens, + boundary_config=boundary_config, + ).to(dtype=torch.int) + padded_mask = torch.zeros_like(token_row, dtype=torch.int) + padded_mask[active_positions] = active_mask + loss_masks.append(padded_mask) + return torch.stack(loss_masks) + + def _validate_nemotron_omni_visual_keys(visual_keys: object = None) -> None: """Validate the model-owned visual input contract retained for API compatibility.""" if visual_keys is None: @@ -823,6 +854,14 @@ def nemotron_omni_collate_fn( Use :func:`nemotron_omni_llava_collate_fn` for the legacy LLaVA collapse/expand contract. """ + if collapse_image_tokens: + warnings.warn( + "The Nemotron Omni LLaVA collapse/expand data contract is deprecated; use " + "nemotron_omni_expanded_collate_fn (the default registry path) with the canonical " + "processor-expanded model.", + FutureWarning, + stacklevel=2, + ) _validate_nemotron_omni_visual_keys(visual_keys) del start_of_response_token, min_pixels, max_pixels if not examples: @@ -879,22 +918,33 @@ def nemotron_omni_collate_fn( assistant_end_fallbacks=("<|im_end|>",), role_start_markers=CHATML_OTHER_ROLE_STARTS, ) - loss_mask = torch.stack( - [ - build_assistant_loss_mask( - example, - input_ids, - processor, - skipped_tokens, - boundary_config=boundary_config, - ).to(dtype=torch.int) - for example, input_ids in zip(mask_examples, batch["input_ids"], strict=True) - ] + loss_mask = _build_padded_assistant_loss_masks( + mask_examples, + batch["input_ids"], + batch["attention_mask"], + processor, + skipped_tokens, + boundary_config=boundary_config, ) if collapse_image_tokens: adjusted, loss_mask = _adjust_image_placeholders(batch, loss_mask, processor, num_tiles) batch["input_ids"] = adjusted["input_ids"] batch["attention_mask"] = adjusted["attention_mask"] + elif use_temporal_video_embedder and num_tiles is not None: + tokens_per_tubelet = _pixel_shuffled_token_count( + height=VISION_FRAME_SIZE, + width=VISION_FRAME_SIZE, + patch_dim=patch_dim, + ) + replacement_counts = torch.full_like(num_tiles, tokens_per_tubelet) + adjusted, loss_mask = _adjust_image_placeholders( + batch, + loss_mask, + processor, + replacement_counts, + ) + batch["input_ids"] = adjusted["input_ids"] + batch["attention_mask"] = adjusted["attention_mask"] if use_per_image_token_counts: _pack_dynamic_images(batch, patch_dim=patch_dim) @@ -952,6 +1002,7 @@ def nemotron_omni_collate_fn( in_batch_packing_pad_to_multiple_of=in_batch_packing_pad_to_multiple_of, pad_token_id=int(pad_token_id), ignore_index=IGNORE_INDEX, + emit_packed_padding_mask=True, ) # Do not synthesize PackedSeqParams.tokens_per_sample: canonical # packing is compact and rows may have different physical lengths. @@ -983,7 +1034,7 @@ def nemotron_omni_collate_fn( def nemotron_omni_llava_collate_fn(*args, **kwargs) -> dict[str, torch.Tensor]: - """Collate inputs for the explicit legacy LLaVA collapse/expand path.""" + """Collate inputs for the deprecated LLaVA collapse/expand path.""" kwargs["collapse_image_tokens"] = True return nemotron_omni_collate_fn(*args, **kwargs) diff --git a/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py b/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py index d86db038cc..a8e47f58ae 100644 --- a/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py +++ b/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py @@ -114,9 +114,8 @@ class NemotronOmniModel(MegatronModule): context parallelism, the model inserts media into the full stream and then selects the rank-local CP shard without changing packed metadata. - Image and text inputs are supported. The sound modules are retained in the - model namespace, but sound insertion remains unsupported until its - one-feature-per-placeholder contract is implemented and tested. + Image, video, sound, and text inputs use the same one-feature-per-placeholder + contract. """ model_owns_packing = False @@ -236,8 +235,8 @@ def __init__( self.vision_model.register_load_state_dict_post_hook(_ignore_transformer_engine_extra_state) self.vision_projection.register_load_state_dict_post_hook(_ignore_transformer_engine_extra_state) - # Preserve the top-level sound-module namespace for checkpoint - # conversion while expanded-sequence sound insertion is unsupported. + # Preserve the top-level sound-module namespace used by checkpoint + # conversion while keeping media insertion local to this model. self.sound_model = sound_model self.sound_projection = sound_projection @@ -339,10 +338,8 @@ def _encode_images( if use_temporal and num_frames is None: raise ValueError( "num_frames is required by the configured RADIO encoder; " - "use one entry with value 1 for each image." + "provide one entry per image or video item." ) - if num_frames is not None and torch.any(num_frames != 1): - raise NotImplementedError("Video insertion is not implemented; num_frames must be 1 for image inputs.") vision_output = self.vision_model( images, @@ -383,6 +380,50 @@ def _encode_images( projected = self.vision_projection(encoded.unsqueeze(1)) return projected.squeeze(1).contiguous() + def _encode_sound(self, sound_clips: torch.Tensor, sound_length: Optional[torch.Tensor]) -> torch.Tensor: + """Encode mel features and return valid projected rows in sample order.""" + + if self.sound_model is None or self.sound_projection is None: + raise RuntimeError("Sound data was provided on a stage without the sound encoder") + if sound_length is None: + raise ValueError("sound_length is required when sound_clips are provided.") + if sound_clips.ndim < 2: + raise ValueError(f"sound_clips must include batch and frame dimensions, got {tuple(sound_clips.shape)}.") + + parameter = next(self.sound_model.parameters()) + sound_clips = sound_clips.to(dtype=parameter.dtype) + sound_embeddings, embedding_lengths = self.sound_model(sound_clips, sound_length) + if sound_embeddings.ndim != 3: + raise ValueError( + "The sound encoder must return [batch, sequence, hidden] embeddings, " + f"got {tuple(sound_embeddings.shape)}." + ) + if embedding_lengths.numel() != sound_embeddings.shape[0]: + raise ValueError( + "The sound encoder must return one valid embedding length per sample; " + f"got {embedding_lengths.numel()} lengths for batch size {sound_embeddings.shape[0]}." + ) + + projection_parameter = next(self.sound_projection.parameters(), None) + if projection_parameter is not None: + sound_embeddings = sound_embeddings.to(dtype=projection_parameter.dtype) + projected = self.sound_projection(sound_embeddings.permute(1, 0, 2).contiguous()).contiguous() + projected_by_sample = projected.permute(1, 0, 2) + if getattr(getattr(self.sound_model, "config", None), "sound_pad_to_clip_duration", False): + return projected_by_sample.reshape(-1, projected.shape[-1]).contiguous() + + valid_embeddings = [] + for sample_embeddings, embedding_length in zip(projected_by_sample, embedding_lengths, strict=True): + valid_length = int(embedding_length.item()) + if valid_length < 0 or valid_length > sample_embeddings.shape[0]: + raise ValueError( + f"Sound embedding length {valid_length} is outside encoded width {sample_embeddings.shape[0]}." + ) + valid_embeddings.append(sample_embeddings[:valid_length]) + if not valid_embeddings: + return projected.new_empty((0, projected.shape[-1])) + return torch.cat(valid_embeddings, dim=0).contiguous() + def _patchify_dynamic_images(self, images: torch.Tensor, imgs_sizes: torch.Tensor) -> torch.Tensor: """Convert padded processor pixels to RADIO's packed patch representation. @@ -502,6 +543,7 @@ def _apply_context_parallel_sharding( attention_mask: Optional[torch.Tensor], labels: Optional[torch.Tensor], loss_mask: Optional[torch.Tensor], + padding_mask: Optional[torch.Tensor], packed_seq_params: Optional[PackedSeqParams], ) -> tuple[ Optional[torch.Tensor], @@ -510,6 +552,7 @@ def _apply_context_parallel_sharding( Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], + Optional[torch.Tensor], bool, ]: """Apply one shared CP index after length-preserving media insertion.""" @@ -520,6 +563,7 @@ def _apply_context_parallel_sharding( (position_ids, 1), (labels, 1), (loss_mask, 1), + (padding_mask, 1), ) full_lengths = {tensor.size(dim) for tensor, dim in sequence_tensors if tensor is not None} if len(full_lengths) > 1: @@ -527,7 +571,7 @@ def _apply_context_parallel_sharding( if not full_lengths: # Intermediate PP stages receive an already CP-local pipeline # tensor and only need the unchanged global THD metadata. - return input_ids, combined_embeddings, position_ids, attention_mask, labels, loss_mask, False + return input_ids, combined_embeddings, position_ids, attention_mask, labels, loss_mask, padding_mask, False total_tokens = full_lengths.pop() if packed_seq_params is not None: @@ -545,13 +589,14 @@ def _apply_context_parallel_sharding( device=device, ) if index is None: - return input_ids, combined_embeddings, position_ids, attention_mask, labels, loss_mask, False + return input_ids, combined_embeddings, position_ids, attention_mask, labels, loss_mask, padding_mask, False input_ids = self._select_sequence(input_ids, index, dim=1) combined_embeddings = self._select_sequence(combined_embeddings, index, dim=0) position_ids = self._select_sequence(position_ids, index, dim=1) labels = self._select_sequence(labels, index, dim=1) loss_mask = self._select_sequence(loss_mask, index, dim=1) + padding_mask = self._select_sequence(padding_mask, index, dim=1) if attention_mask is not None: attention_seq_dim = 1 if attention_mask.dim() == 2 else 2 @@ -564,6 +609,7 @@ def _apply_context_parallel_sharding( attention_mask, labels, loss_mask, + padding_mask, loss_mask is not None, ) @@ -574,6 +620,7 @@ def forward( attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, loss_mask: Optional[torch.Tensor] = None, + padding_mask: Optional[torch.Tensor] = None, inference_context=None, runtime_gather_output: Optional[bool] = None, packed_seq_params: Optional[PackedSeqParams] = None, @@ -595,13 +642,13 @@ def forward( applies a context-parallel shard to the supervision tensors. """ - del kwargs, sound_length + del kwargs if images is None: images = pixel_values - has_sound = sound_clips is not None and sound_clips.numel() > 0 - if has_sound: - raise NotImplementedError("Sound insertion is not implemented; use an image or text input.") + has_sound_inputs = sound_clips is not None and sound_clips.numel() > 0 + if has_sound_inputs and sound_clips.shape == torch.Size([1, 1]): + has_sound_inputs = sound_clips[0, 0].item() != 0 lm_input_ids = input_ids combined_embeddings = None @@ -618,6 +665,11 @@ def forward( else: image_embeddings = None + if has_sound_inputs: + sound_embeddings = self._encode_sound(sound_clips, sound_length) + else: + sound_embeddings = None + # Match LLaVAModel's execution order. Besides keeping the two # implementations directly comparable, this ensures that RADIO's # first distributed forward sees the same runtime/collective state. @@ -632,9 +684,20 @@ def forward( input_ids, image_embeddings, self.image_token_index, - attention_mask, + ~padding_mask if packed_seq_params is not None and padding_mask is not None else attention_mask, ) + if self.sound_token_index > 0: + if sound_embeddings is None: + sound_embeddings = combined_embeddings.new_empty((0, combined_embeddings.shape[-1])) + combined_embeddings = self._merge_projected_media( + combined_embeddings, + input_ids, + sound_embeddings, + self.sound_token_index, + ~padding_mask if packed_seq_params is not None and padding_mask is not None else attention_mask, + ) + if packed_seq_params is not None: # THD tensors and their logical boundaries are final collator # outputs. The model may shard token-aligned tensors for CP, but @@ -648,6 +711,7 @@ def forward( attention_mask, labels, loss_mask, + padding_mask, return_sliced_loss_mask, ) = self._apply_context_parallel_sharding( input_ids=lm_input_ids, @@ -656,6 +720,7 @@ def forward( attention_mask=attention_mask, labels=labels, loss_mask=loss_mask, + padding_mask=padding_mask, packed_seq_params=packed_seq_params, ) @@ -679,6 +744,15 @@ def forward( combined_embeddings, group=self.pg_collection.tp, ).contiguous() + if padding_mask is not None and self.sequence_parallel_lm: + padding_mask = ( + tensor_parallel.scatter_to_sequence_parallel_region( + padding_mask.transpose(0, 1).contiguous(), + group=self.pg_collection.tp, + ) + .transpose(0, 1) + .contiguous() + ) # Match LLaVAModel's external-embedding contract. Once media has been # merged into decoder embeddings, the language model must not receive @@ -690,6 +764,10 @@ def forward( if combined_embeddings is not None and not mtp_enabled: lm_input_ids = None + # TODO(https://github.com/NVIDIA/Megatron-LM/issues/6111): Forward the + # CP/SP-local padding_mask once MCore's expert-bias router supports it. + # Until then, packed alignment gaps remain loss-masked but are counted + # by MoE router auxiliary statistics. output = self.language_model( input_ids=lm_input_ids, position_ids=position_ids, diff --git a/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni_llava.py b/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni_llava.py index 1b4e70eef1..8248ffeeee 100644 --- a/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni_llava.py +++ b/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni_llava.py @@ -12,17 +12,58 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings + +from megatron.core.models.multimodal.llava_model import LLaVAModel + from megatron.bridge.models.nemotron_vl.modeling_nemotron_vl import NemotronVLModel +from megatron.bridge.models.nemotron_vl.nemotron_vl_provider import NemotronVLModelProvider class NemotronOmniLlavaModel(NemotronVLModel): - """Legacy collapse/expand Omni wrapper around MCore ``LLaVAModel``. + """Deprecated collapse/expand Omni wrapper around MCore ``LLaVAModel``. forward() is inherited from NemotronVLModel (which delegates to LLaVAModel), so sound kwargs (sound_clips, sound_length) pass through automatically when the selected LLaVAModel implementation supports them. + + Use :class:`~megatron.bridge.models.nemotron_omni.modeling_nemotron_omni.NemotronOmniModel` + for the canonical processor-expanded sequence and collator-owned packing + contract. """ + def __init__( + self, + config: NemotronVLModelProvider | None = None, + *, + llava_model: LLaVAModel | None = None, + pre_process: bool | None = True, + post_process: bool | None = True, + vp_stage: int | None = None, + ) -> None: + """Construct the deprecated LLaVA compatibility model. + + Args: + config: Provider used to construct the wrapped model. + llava_model: Fully assembled MCore LLaVA model. + pre_process: Whether this pipeline stage owns input processing. + post_process: Whether this pipeline stage owns output processing. + vp_stage: Optional virtual pipeline stage. + """ + warnings.warn( + "NemotronOmniLlavaModel is deprecated; use NemotronOmniModel with the canonical " + "processor-expanded sequence contract.", + FutureWarning, + stacklevel=2, + ) + super().__init__( + config=config, + llava_model=llava_model, + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + ) + def freeze( self, *, diff --git a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py index 427d279618..27ea813930 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py @@ -33,17 +33,21 @@ """ import copy +import warnings +from collections.abc import Iterable from dataclasses import fields +import torch from megatron.core.activations import squared_relu from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry -from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.model_bridge import HFWeightTuple, MegatronModelBridge, WeightConversionTask from megatron.bridge.models.conversion.param_mapping import ( AutoMapping, ReplicatedMapping, ) from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.hf_pretrained.state import SafeTensorsStateSource, StateDict from megatron.bridge.models.nemotron_omni.modeling_nemotron_omni import NemotronOmniModel from megatron.bridge.models.nemotron_omni.nemotron_omni_provider import ( NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT, @@ -82,6 +86,13 @@ def _copy_mapping_with_prefixes(mapping, *, megatron_prefix: str, hf_prefix: str class NemotronOmniBridge(NemotronVLBridge): """Bridge for the canonical expanded-sequence Nemotron-3 Omni model.""" + _HF_PASSTHROUGH_KEYS = ( + "sound_encoder.encoder.feature_extractor.featurizer.fb", + "sound_encoder.encoder.feature_extractor.featurizer.window", + "vision_model.radio_model.input_conditioner.norm_mean", + "vision_model.radio_model.input_conditioner.norm_std", + ) + CONFIG_MAPPING = NemotronVLBridge.CONFIG_MAPPING + [ # HF public Omni config uses layer_norm_epsilon instead of rms_norm_eps. ("layer_norm_epsilon", "layernorm_epsilon"), @@ -117,9 +128,8 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> NemotronOmniMode """Create a NemotronOmniModelProvider from the HF Omni config. Always returns an Omni provider (MoE language model + RADIO ViT - vision + optional Parakeet sound encoder). When ``sound_config`` is - absent on the HF config, ``has_sound=False`` and the sound branch - is skipped at construction time. + vision + optional Parakeet sound encoder). The presence of + ``sound_config`` is the Hugging Face checkpoint's sound capability. """ hf_config = hf_pretrained.config llm_config = hf_config.llm_config @@ -132,10 +142,9 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> NemotronOmniMode if hasattr(hf_config, "projector_hidden_size"): provider_kwargs["vision_proj_ffn_hidden_size"] = hf_config.projector_hidden_size - has_sound = hasattr(hf_config, "sound_config") and hf_config.sound_config is not None - if has_sound: - sc = hf_config.sound_config - provider_kwargs["has_sound"] = True + sc = getattr(hf_config, "sound_config", None) + provider_kwargs["has_sound"] = sc is not None + if sc is not None: provider_kwargs["sound_model_type"] = getattr(sc, "model_type", "parakeet") provider_kwargs["sound_hidden_size"] = sc.hidden_size provider_kwargs["sound_projection_hidden_size"] = sc.projection_hidden_size @@ -173,6 +182,20 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> NemotronOmniMode provider.mtp_hybrid_override_pattern = getattr(llm_config, "mtp_hybrid_override_pattern", None) return provider + @classmethod + def megatron_to_hf_config(cls, provider) -> dict: + """Export sound capability consistently with model construction.""" + hf_config = super().megatron_to_hf_config(provider) + if provider.has_sound: + hf_config["sound_config"] = provider.sound_config + hf_config["sound_context_token_id"] = provider.sound_context_token_id + else: + # Config synthesis fills missing keys from the reference HF config. + # Keep an explicit None so an image-text checkpoint stays sound-free. + hf_config["sound_config"] = None + hf_config["sound_context_token_id"] = None + return hf_config + # ------------------------------------------------------------------ # Parameter mapping # ------------------------------------------------------------------ @@ -217,8 +240,8 @@ def _llava_mapping_registry(self) -> MegatronMappingRegistry: # (conformer layers, subsampling convs, subsampling linear). # Feature extractor buffers (``feature_extractor.featurizer.fb``, # ``.window``) live outside the encoder and are intentionally - # unmapped -- they're skipped on import and regenerated from config - # on export. + # unmapped. They are preserved directly from the source checkpoint + # during export. mapping_list.append( ReplicatedMapping( megatron_param="llava_model.sound_model.encoder.**", @@ -262,11 +285,56 @@ def mapping_registry(self) -> MegatronMappingRegistry: return MegatronMappingRegistry(*mappings) + @torch.no_grad() + def stream_weights_megatron_to_hf( + self, + megatron_model: NemotronOmniModel | list[NemotronOmniModel], + hf_pretrained: PreTrainedCausalLM, + cpu: bool = True, + show_progress: bool = True, + conversion_tasks: list[WeightConversionTask] | None = None, + merge_adapter_weights: bool = True, + weight_dtype: torch.dtype | None = None, + ) -> Iterable[HFWeightTuple]: + """Export model weights and preserve immutable source-only buffers.""" + yield from super().stream_weights_megatron_to_hf( + megatron_model, + hf_pretrained, + cpu=cpu, + show_progress=show_progress, + conversion_tasks=conversion_tasks, + merge_adapter_weights=merge_adapter_weights, + weight_dtype=weight_dtype, + ) + + state = getattr(hf_pretrained, "state", None) + source = getattr(state, "source", None) + if source is None: + source_path = getattr(hf_pretrained, "name_or_path", None) + if not source_path: + return + source = SafeTensorsStateSource(source_path) + state = StateDict(source) + source_keys = set(source.get_all_keys()) + for name in self._HF_PASSTHROUGH_KEYS: + if name in source_keys: + yield from HFWeightTuple(name, state[name]).iter_finalized(cpu=cpu) + class NemotronOmniLlavaBridge(NemotronOmniBridge): - """Explicit fallback bridge for the historical collapse/expand model.""" + """Deprecated fallback bridge for the historical collapse/expand model. + + Use :class:`NemotronOmniBridge`, which is the canonical AutoBridge + registration and consumes processor-expanded media-token sequences. + """ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> NemotronOmniLlavaModelProvider: + warnings.warn( + "NemotronOmniLlavaBridge is deprecated; use NemotronOmniBridge with the canonical " + "processor-expanded sequence contract.", + FutureWarning, + stacklevel=2, + ) provider = super().provider_bridge(hf_pretrained) provider_kwargs = { field.name: getattr(provider, field.name) diff --git a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py index 4097ec98f2..86b80f66ea 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py @@ -13,6 +13,7 @@ # limitations under the License. import copy +import warnings from abc import ABC from dataclasses import dataclass from types import SimpleNamespace @@ -156,6 +157,8 @@ class _NemotronOmniModelProviderBase(NemotronVLModelProvider): # accepts a different 4D tensor contract and is not an Omni configuration. dynamic_resolution: Literal[True] = True + # This is the single source of truth for sound checkpoint capability: + # disabling it omits both the encoder and its dependent projector. has_sound: bool = False sound_model_type: str = "parakeet" sound_hidden_size: int = 1024 @@ -276,6 +279,20 @@ def _build_sound_encoder(self): ) return BridgeSoundEncoder(config) + def _build_sound_modules(self, language_cfg, language_spec, *, add_encoder: bool): + """Build optional sound modules on the encoder pipeline stage.""" + if not (self.has_sound and add_encoder): + return None, None + + sound_model = self._build_sound_encoder() + sound_projection = MultimodalProjector( + config=self._build_sound_projection_config(language_cfg), + submodules=copy.deepcopy(get_language_mlp_submodules(language_spec)), + projector_type="mlp", + input_size=self.sound_hidden_size, + ) + return sound_model, sound_projection + def _provide_llava(self, pre_process=None, post_process=None, vp_stage=None): """Assemble the legacy LLaVA collapse/expand implementation. @@ -300,22 +317,12 @@ def _provide_llava(self, pre_process=None, post_process=None, vp_stage=None): add_encoder_flag = parallel_state.is_pipeline_first_stage() if self.pipeline_model_parallel_size > 1 else True add_decoder_flag = True - # Build sound components (only on PP first stage, only when sound present) - sound_model = None - sound_projection = None sound_token_index = self.sound_context_token_id - - if self.has_sound and add_encoder_flag: - sound_model = self._build_sound_encoder() - - sound_proj_cfg = self._build_sound_projection_config(language_cfg) - sound_proj_spec = copy.deepcopy(get_language_mlp_submodules(language_spec)) - sound_projection = MultimodalProjector( - config=sound_proj_cfg, - submodules=sound_proj_spec, - projector_type="mlp", - input_size=self.sound_hidden_size, - ) + sound_model, sound_projection = self._build_sound_modules( + language_cfg, + language_spec, + add_encoder=add_encoder_flag, + ) llava_model = LLaVAModel( language_transformer_config=language_cfg, @@ -416,16 +423,11 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None): add_encoder = parallel_state.is_pipeline_first_stage() if self.pipeline_model_parallel_size > 1 else True - sound_model = None - sound_projection = None - if self.has_sound and add_encoder: - sound_model = self._build_sound_encoder() - sound_projection = MultimodalProjector( - config=self._build_sound_projection_config(language_cfg), - submodules=copy.deepcopy(get_language_mlp_submodules(language_spec)), - projector_type="mlp", - input_size=self.sound_hidden_size, - ) + sound_model, sound_projection = self._build_sound_modules( + language_cfg, + language_spec, + add_encoder=add_encoder, + ) model = NemotronOmniModel( language_transformer_config=language_cfg, @@ -483,7 +485,11 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None): @dataclass class NemotronOmniLlavaModelProvider(NemotronOmniModelProvider): - """Explicit fallback provider for the historical collapse/expand path.""" + """Deprecated fallback provider for the historical collapse/expand path. + + Use :class:`NemotronOmniModelProvider`, which constructs the canonical + processor-expanded model with collator-owned packing. + """ # Preserve the existing LLaVA provider default for compatibility. radio_interpolate_only_cpe: bool = True @@ -496,5 +502,11 @@ def validate_model_contract(self) -> None: ) def provide(self, pre_process=None, post_process=None, vp_stage=None): + warnings.warn( + "NemotronOmniLlavaModelProvider is deprecated; use NemotronOmniModelProvider with the canonical " + "processor-expanded sequence contract.", + FutureWarning, + stacklevel=2, + ) self.validate_model_contract() return self._provide_llava(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) diff --git a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_utils.py b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_utils.py index 30085c6e9d..0725279ddc 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_utils.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_utils.py @@ -13,6 +13,7 @@ # limitations under the License. import math +import warnings from collections.abc import Sequence from functools import lru_cache from typing import Any, TypeVar, Union @@ -98,11 +99,11 @@ def inference_num_image_tiles( ) -> torch.Tensor: """Build image-placeholder replacement counts for pipeline inference. - The first pipeline stage can derive these counts from vision encoder - outputs, but the last stage needs the same row-major metadata to expand - input positions. Dynamic images contribute their post-pixel-shuffle token - count per compact placeholder. Temporal tubelets contribute one tile each; - ``LLaVAModel.img_seq_len`` supplies their fixed embedding width. + Dynamic images contribute their post-pixel-shuffle feature count per tile. + Temporal tubelets contribute one logical count each; canonical inference + applies the fixed tubelet width through + :func:`inference_expanded_image_token_counts`. The deprecated LLaVA path + instead applies that width inside the model. Args: imgs_sizes: Per-image or per-frame ``(height, width)`` metadata. @@ -140,6 +141,66 @@ def inference_num_image_tiles( return (grid_sizes.prod(dim=1) // (pixel_shuffle_factor**2)).to(dtype=torch.int) +def inference_expanded_image_token_counts( + tile_feature_counts: torch.Tensor, + tiles_per_media: int | Sequence[int] | torch.Tensor, + *, + feature_multiplier: int = 1, +) -> torch.Tensor: + """Aggregate projected feature counts for canonical inference prompts. + + Dynamic image processors can split one source image into multiple RADIO + tiles, while each ``...`` region belongs to the source image. + The canonical model needs one ```` placeholder per projected + feature, so per-tile counts must be summed back to one count per region. + Temporal inference uses one logical tile per tubelet and a fixed feature + multiplier for the post-pixel-shuffle tubelet width. + + Args: + tile_feature_counts: Number of projected feature rows produced by each + RADIO tile or temporal tubelet. + tiles_per_media: Number of entries in ``tile_feature_counts`` owned by + each ``...`` region. + feature_multiplier: Additional projected width per count. Use one for + dynamic images and the tubelet feature width for temporal video. + + Returns: + One expanded placeholder count per ``...`` region. + + Raises: + ValueError: If counts are non-positive or do not account for every + tile/tubelet. + """ + flat_feature_counts = tile_feature_counts.reshape(-1) + if torch.any(flat_feature_counts <= 0): + raise ValueError("tile_feature_counts entries must be greater than 0.") + if feature_multiplier <= 0: + raise ValueError("feature_multiplier must be greater than 0.") + + if isinstance(tiles_per_media, int): + media_tile_counts = [tiles_per_media] + elif isinstance(tiles_per_media, torch.Tensor): + media_tile_counts = [int(count) for count in tiles_per_media.detach().cpu().reshape(-1).tolist()] + else: + media_tile_counts = [int(count) for count in tiles_per_media] + if not media_tile_counts or any(count <= 0 for count in media_tile_counts): + raise ValueError("tiles_per_media entries must be greater than 0.") + if sum(media_tile_counts) != flat_feature_counts.numel(): + raise ValueError( + "tiles_per_media must account for every tile feature count; " + f"got {sum(media_tile_counts)} tiles for {flat_feature_counts.numel()} counts." + ) + + offset = 0 + expanded_counts = [] + for tile_count in media_tile_counts: + expanded_counts.append( + int(flat_feature_counts[offset : offset + tile_count].sum().item()) * feature_multiplier + ) + offset += tile_count + return torch.tensor(expanded_counts, dtype=torch.int, device=tile_feature_counts.device) + + def inference_merged_sequence_length( input_ids: torch.Tensor, *, @@ -147,7 +208,10 @@ def inference_merged_sequence_length( num_image_tiles: torch.Tensor | None, image_seq_len: int, ) -> int: - """Return the unpadded sequence length after vision-token replacement. + """Return the legacy unpadded length after model-owned vision expansion. + + This helper is deprecated because the canonical model consumes an already + expanded sequence; its merged length is simply ``input_ids.shape[1]``. Args: input_ids: One inference prompt row, including generated tokens so far. @@ -158,6 +222,12 @@ def inference_merged_sequence_length( Returns: The real merged sequence length before pipeline padding. """ + warnings.warn( + "inference_merged_sequence_length is deprecated with the Nemotron Omni LLaVA collapse/expand path; " + "canonical expanded-sequence inference uses input_ids.shape[1].", + FutureWarning, + stacklevel=2, + ) if input_ids.ndim != 2 or input_ids.shape[0] != 1: raise ValueError(f"input_ids must have shape [1, S], got {tuple(input_ids.shape)}.") if image_seq_len <= 0: diff --git a/src/megatron/bridge/recipes/nemotron_omni/h100/__init__.py b/src/megatron/bridge/recipes/nemotron_omni/h100/__init__.py index 8fc2540903..0a134caaaf 100644 --- a/src/megatron/bridge/recipes/nemotron_omni/h100/__init__.py +++ b/src/megatron/bridge/recipes/nemotron_omni/h100/__init__.py @@ -16,6 +16,7 @@ __all__ = [ + "nemotron_omni_cord_v2_long_context_sft_8gpu_h100_bf16_config", "nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config", "nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config", "nemotron_omni_valor32k_peft_4gpu_h100_bf16_config", diff --git a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py index 3f308f4b91..b0d00b58a8 100644 --- a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py +++ b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py @@ -31,6 +31,7 @@ from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed _DEFAULT_HF_PATH = "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" @@ -63,15 +64,18 @@ def nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config() -> ConfigContainer: """Return a VL SFT config for Nemotron Omni on CORD v2. Vision-language finetuning on the CORD v2 receipt parsing dataset. + Sound modules are omitted because this dataset contains only image-text samples. Default configuration: 4 GPUs (TP=4). Uses nemotron_omni_step (pass --step_func nemotron_omni_step). """ cfg = _nemotron_omni_base() cfg.model.temporal_patch_dim = 1 + cfg.model.has_sound = False cfg.dataset = DirectHFSFTDatasetConfig( seq_length=4096, preprocessing=ChatSFTPreprocessingConfig(), hf_processor_path=_DEFAULT_HF_PATH, + trust_remote_code=True, source=HFDatasetSourceConfig(dataset_name="cord_v2"), num_workers=2, dataloader_type="cyclic", @@ -88,11 +92,45 @@ def nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config() -> ConfigContainer: return cfg +def nemotron_omni_cord_v2_long_context_sft_8gpu_h100_bf16_config() -> ConfigContainer: + """Return an 8K CORD v2 SFT config with in-batch packing and CP2. + + In-batch packing requires a micro batch greater than one. The TP4/CP2 + topology needs at least eight GPUs and aligns every packed row to the + combined CP/SP multiple. Precision-aware Adam uses FP16 main parameters + with stored FP32 remainders, BF16 gradients, and BF16 moments so first-step + optimizer-state initialization fits within 80 GB H100 memory. + """ + cfg = nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config() + cfg.model.seq_length = 8192 + cfg.model.context_parallel_size = 2 + cfg.model.calculate_per_token_loss = True + cfg.train.micro_batch_size = 2 + cfg.optimizer.use_precision_aware_optimizer = True + cfg.optimizer.main_grads_dtype = torch.bfloat16 + cfg.optimizer.main_params_dtype = torch.float16 + cfg.optimizer.store_param_remainders = True + cfg.optimizer.exp_avg_dtype = torch.bfloat16 + cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16 + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = False + cfg.dataset.seq_length = 8192 + cfg.dataset.enable_in_batch_packing = True + cfg.dataset.in_batch_packing_pad_to_multiple_of = 8 + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + def nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config() -> ConfigContainer: """Return a LoRA PEFT config for Nemotron Omni on CORD v2. LoRA adapters are applied to attention, Mamba, and FC1/FC2 projections. - Vision and sound base modules remain frozen while matching adapters are trainable. + Vision base modules remain frozen and sound modules are omitted. Default configuration: 4 GPUs (TP=4). Uses nemotron_omni_step (pass --step_func nemotron_omni_step). """ @@ -100,6 +138,7 @@ def nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config() -> ConfigContainer: cfg = _nemotron_omni_base() cfg.model.temporal_patch_dim = 1 + cfg.model.has_sound = False cfg.peft = LoRA( target_modules=["linear_qkv", "linear_proj", "in_proj", "out_proj", "linear_fc1", "linear_fc2"], dim=16, @@ -125,6 +164,7 @@ def nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config() -> ConfigContainer: seq_length=4096, preprocessing=ChatSFTPreprocessingConfig(), hf_processor_path=_DEFAULT_HF_PATH, + trust_remote_code=True, source=HFDatasetSourceConfig(dataset_name="cord_v2"), num_workers=2, dataloader_type="cyclic", @@ -278,6 +318,7 @@ def nemotron_omni_valor32k_peft_4gpu_h100_bf16_config() -> ConfigContainer: __all__ = [ + "nemotron_omni_cord_v2_long_context_sft_8gpu_h100_bf16_config", "nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config", "nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config", "nemotron_omni_valor32k_peft_4gpu_h100_bf16_config", diff --git a/src/megatron/bridge/training/nemotron_omni_step.py b/src/megatron/bridge/training/nemotron_omni_step.py index a2dc831295..a215577cec 100644 --- a/src/megatron/bridge/training/nemotron_omni_step.py +++ b/src/megatron/bridge/training/nemotron_omni_step.py @@ -15,7 +15,8 @@ """Nemotron Omni training step -- extends llava_step with sound support. Adds ``sound_clips`` and ``sound_length`` to the model forward kwargs so that -LLaVAModel processes audio embeddings alongside vision embeddings. +the canonical expanded-sequence model and explicit LLaVA compatibility model +can process audio embeddings alongside vision embeddings. """ import logging @@ -82,6 +83,11 @@ def get_batch_from_iterator( if "cu_seqlens_q" in batch: required_device_keys.update(key for key in _PACKED_SEQ_DEVICE_KEYS if key in batch) required_host_keys.update(key for key in _PACKED_SEQ_HOST_KEYS if key in batch) + if batch.get("padding_mask") is not None: + # Preserve the physical THD gap mask through every decoder PP stage + # so CP/SP localization remains ready for routing after MCore #6111. + # The mask is deliberately not forwarded into MCore until then. + required_device_keys.add("padding_mask") if is_first_pp_stage or is_last_pp_stage: input_key = "tokens" if batch.get("tokens") is not None else "input_ids" @@ -176,7 +182,7 @@ def get_batch(data_iterator: Iterable, cfg: ConfigContainer, *, pg_collection) - is_last = is_pp_last_stage(pg_collection.pp) skip_attention_mask = getattr(cfg.dataset, "skip_getting_attention_mask_from_dataset", True) if (not is_first) and (not is_last) and skip_attention_mask and not _uses_packed_sequence_metadata(cfg): - return (None,) * 14 + return (None,) * 15 batch = get_batch_from_iterator( data_iterator, @@ -194,7 +200,7 @@ def get_batch(data_iterator: Iterable, cfg: ConfigContainer, *, pg_collection) - # Leave language tensors in their complete collator-owned layout. The model # inserts media first, then applies one shared CP index to embeddings, - # labels, and the loss mask. + # supervision tensors, and physical alignment padding. if images is not None: batch["images"] = images @@ -221,6 +227,7 @@ def get_batch(data_iterator: Iterable, cfg: ConfigContainer, *, pg_collection) - num_frames, vision_packed_seq_params, num_image_tiles, + batch.get("padding_mask"), ) @@ -257,6 +264,7 @@ def forward_step( num_frames, vision_packed_seq_params, num_image_tiles, + padding_mask, ) = get_batch(data_iterator, state.cfg, pg_collection=pg_collection) timers("batch-generator").stop() @@ -275,6 +283,8 @@ def forward_step( "labels": labels, "loss_mask": loss_mask, } + if padding_mask is not None: + forward_args["padding_mask"] = padding_mask if sound_clips is not None: forward_args["sound_clips"] = sound_clips.to(dtype=torch.bfloat16) diff --git a/tests/unit_tests/conversion/launcher/test_utils.py b/tests/unit_tests/conversion/launcher/test_utils.py index 0c18a8831e..cc946ea3b9 100644 --- a/tests/unit_tests/conversion/launcher/test_utils.py +++ b/tests/unit_tests/conversion/launcher/test_utils.py @@ -63,6 +63,16 @@ def fake_model_info(_self, *, repo_id, revision): assert calls == [{"repo_id": "hf/model", "revision": "release-tag"}] +def test_resolve_hf_commit_revision_preserves_immutable_sha_without_network(monkeypatch): + revision = "0123456789abcdef0123456789abcdef01234567" # pragma: allowlist secret + monkeypatch.setattr( + "huggingface_hub.HfApi.model_info", + lambda *_args, **_kwargs: pytest.fail("an immutable commit SHA must not require Hub metadata"), + ) + + assert resolve_hf_commit_revision("hf/model", revision) == revision + + @pytest.mark.parametrize("resolver", [resolve_hf_commit_revision, resolve_hf_model_revision]) def test_hf_revision_resolvers_reject_local_path(tmp_path, resolver): with pytest.raises(ValueError, match="only to Hugging Face Hub model IDs"): diff --git a/tests/unit_tests/data/builders/test_energon_builder.py b/tests/unit_tests/data/builders/test_energon_builder.py index ff841ce153..40966cc301 100644 --- a/tests/unit_tests/data/builders/test_energon_builder.py +++ b/tests/unit_tests/data/builders/test_energon_builder.py @@ -242,6 +242,7 @@ def test_nemotron_factory_preserves_omni_settings(monkeypatch: pytest.MonkeyPatc assert encoder_cls.call_args.kwargs["processor"] is processor assert encoder_cls.call_args.kwargs["max_audio_duration"] == 10.0 assert encoder_cls.call_args.kwargs["use_temporal_video_embedder"] is True + assert encoder_cls.call_args.kwargs["collapse_image_tokens"] is False assert encoder_cls.call_args.kwargs["enable_in_batch_packing"] is True diff --git a/tests/unit_tests/data/collators/test_model_collators.py b/tests/unit_tests/data/collators/test_model_collators.py index b6163f2bfb..f9c45180ce 100644 --- a/tests/unit_tests/data/collators/test_model_collators.py +++ b/tests/unit_tests/data/collators/test_model_collators.py @@ -38,6 +38,7 @@ kimi_k25_vl_collate_fn=kimi_collate.kimi_k25_vl_collate_fn, ministral3_collate_fn=ministral3_collate.ministral3_collate_fn, nemotron_omni_collate_fn=nemotron_omni_collate.nemotron_omni_collate_fn, + nemotron_omni_expanded_collate_fn=nemotron_omni_collate.nemotron_omni_expanded_collate_fn, nemotron_omni_llava_collate_fn=nemotron_omni_collate.nemotron_omni_llava_collate_fn, qwen2_5_collate_fn=qwen_vl_collate.qwen2_5_collate_fn, qwen2_audio_collate_fn=qwen_audio_collate.qwen2_audio_collate_fn, @@ -55,6 +56,12 @@ def test_only_nemotron_omni_requires_model_collate_for_all_examples(): assert not model_collate_required_for_all_examples("UnknownProcessor") +def test_nemotron_omni_registry_selects_canonical_expanded_contract(): + assert ( + resolve_model_collate("NemotronH_Nano_Omni_Reasoning_V3Processor") is collate.nemotron_omni_expanded_collate_fn + ) + + def test_vlm_collate_keeps_qwen_vl_registration(): assert resolve_model_collate("Qwen2_5_VLProcessor") is collate.qwen2_5_collate_fn @@ -1974,6 +1981,37 @@ def test_nemotron_omni_collate_keeps_chatml_turn_end_token(): assert batch["labels"][0, -5:].tolist() == [21, 22, 102, 103, -100] +def test_nemotron_omni_collate_ignores_end_token_padding_when_building_loss_masks(): + proc = _NemotronOmniProcessor( + tokenized_rows=[ + [199, 10, 102, 103, 101, 21, 22, 102, 103], + [199, 11, 102, 103, 101, 31, 32, 33, 34, 35, 102, 103], + ] + ) + proc.tokenizer.pad_token_id = 102 + examples = [ + { + "conversation": [ + {"role": "user", "content": "short question"}, + {"role": "assistant", "content": "short answer"}, + ] + }, + { + "conversation": [ + {"role": "user", "content": "long question"}, + {"role": "assistant", "content": "longer answer"}, + ] + }, + ] + + batch = collate.nemotron_omni_collate_fn(examples, proc, pad_to_multiple_of=1) + + assert batch["attention_mask"][0].tolist() == [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0] + assert torch.all(batch["loss_mask"][0, 9:] == 0) + assert batch["loss_mask"][0].sum() > 0 + assert batch["loss_mask"][1].sum() > 0 + + def test_nemotron_omni_collate_rejects_unsupported_visual_keys(): proc = _NemotronOmniProcessor(tokenized_rows=[[5, 6]]) @@ -1989,13 +2027,14 @@ def test_nemotron_omni_llava_collate_packs_heterogeneous_image_rows_at_post_merg processor = _DynamicNemotronOmniProcessor() monkeypatch.setattr(nemotron_omni_collate, "build_assistant_loss_mask", _sentinel_assistant_loss_mask) - batch = collate.nemotron_omni_llava_collate_fn( - _heterogeneous_nemotron_examples(), - processor, - enable_in_batch_packing=True, - sequence_length=24, - in_batch_packing_pad_to_multiple_of=4, - ) + with pytest.warns(FutureWarning, match="collapse/expand data contract is deprecated"): + batch = collate.nemotron_omni_llava_collate_fn( + _heterogeneous_nemotron_examples(), + processor, + enable_in_batch_packing=True, + sequence_length=24, + in_batch_packing_pad_to_multiple_of=4, + ) assert batch["input_ids"].tolist() == [ [ @@ -2538,3 +2577,31 @@ def test_nemotron_omni_llava_collate_checks_temporal_model_expansion_before_trun use_temporal_video_embedder=True, patch_dim=16, ) + + +def test_nemotron_omni_expanded_collate_emits_one_placeholder_per_temporal_feature(monkeypatch): + processor = _NemotronOmniProcessor() + input_ids = torch.tensor([[10, NEMO_IMG_START_TOKEN_ID, NEMO_IMAGE_TOKEN_ID, NEMO_IMG_END_TOKEN_ID, 11]]) + prepared = { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + "visual_inputs": GenericVisualInputs(pixel_values=torch.ones(1, 1, 768)), + } + examples = [{"conversation": [{"role": "user", "content": "one tubelet"}]}] + monkeypatch.setattr( + nemotron_omni_collate, + "_prepare_temporal_rows", + lambda *args, **kwargs: (prepared, examples, torch.ones(1, dtype=torch.long)), + ) + monkeypatch.setattr(nemotron_omni_collate, "build_assistant_loss_mask", _zero_assistant_loss_mask) + + batch = collate.nemotron_omni_expanded_collate_fn( + examples, + processor, + use_temporal_video_embedder=True, + patch_dim=16, + pad_to_multiple_of=1, + ) + + assert int((batch["input_ids"] == NEMO_IMAGE_TOKEN_ID).sum().item()) == 256 + assert int(batch["attention_mask"].sum().item()) == 260 diff --git a/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py b/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py index 5fd3b2e7a9..76dc961de6 100644 --- a/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py +++ b/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py @@ -292,6 +292,7 @@ def test_energon_temporal_video_is_processed_in_shared_collator(monkeypatch): use_temporal_video_embedder=True, patch_dim=16, pad_to_multiple_of=1, + collapse_image_tokens=True, ) frames = [Image.new("RGB", (16, 16), color=value) for value in (0, 64, 128)] encoded = encoder.encode_sample( @@ -342,6 +343,7 @@ def test_energon_single_frame_video_uses_temporal_embedder_contract(monkeypatch) use_temporal_video_embedder=True, patch_dim=16, pad_to_multiple_of=1, + collapse_image_tokens=True, ) encoded = encoder.encode_sample( _sample( @@ -374,6 +376,38 @@ def test_energon_raw_video_bytes_remain_one_owned_video_per_sample(): assert encoded.example["videos"] == [raw_video] +def test_energon_temporal_video_defaults_to_expanded_contract(monkeypatch): + from PIL import Image + + monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens) + monkeypatch.setattr( + omni_collate, + "_patchify_frame", + lambda frame, *, height, width, patch_dim: torch.ones(2, 3), + ) + processor = _Processor([[1, IMG_START_ID, IMAGE_TOKEN_ID, IMG_END_ID, 21, PAD_AND_END_ID]]) + encoder = NemotronOmniTaskEncoder( + processor=processor, + seq_length=512, + temporal_patch_size=2, + use_temporal_video_embedder=True, + patch_dim=16, + pad_to_multiple_of=1, + ) + encoded = encoder.encode_sample( + _sample( + [{"role": "user", "content": [{"type": "video"}]}], + videos=[[Image.new("RGB", (16, 16)), Image.new("RGB", (16, 16))]], + ) + ) + + batch = encoder.batch([encoded]) + + assert int((batch.input_ids == IMAGE_TOKEN_ID).sum().item()) == 256 + assert batch.attention_mask.sum().item() == 261 + assert batch.num_frames.tolist() == [2] + + def test_energon_multiple_raw_video_bytes_keep_placeholder_order(): raw_videos = [b"first-mp4", b"second-mp4"] encoder = NemotronOmniTaskEncoder(processor=_Processor([[1]]), pad_to_multiple_of=1) @@ -410,6 +444,213 @@ def _fake_decode(path, *, video_fps, video_nframes): assert sampled_fps == 2.5 +def test_energon_llava_multimodal_packing_uses_post_merge_boundaries(monkeypatch): + from PIL import Image + + monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens) + monkeypatch.setattr( + omni_collate, + "_patchify_frame", + lambda frame, *, height, width, patch_dim: torch.ones(2, 3), + ) + encoder = NemotronOmniTaskEncoder( + processor=_Processor( + [ + [1, IMG_START_ID, IMAGE_TOKEN_ID, IMG_END_ID, 21, PAD_AND_END_ID], + [2, IMG_START_ID, IMAGE_TOKEN_ID, IMG_END_ID, 31, 32, PAD_AND_END_ID], + ] + ), + seq_length=768, + enable_in_batch_packing=True, + use_temporal_video_embedder=True, + in_batch_packing_pad_to_multiple_of=8, + pad_to_multiple_of=1, + collapse_image_tokens=True, + ) + samples = [ + encoder.encode_sample( + _sample( + [{"role": "user", "content": [{"type": "video"}]}], + key=f"row-{row_index}", + videos=[[Image.new("RGB", (16, 16)), Image.new("RGB", (16, 16))]], + ) + ) + for row_index in range(2) + ] + + batch = encoder.batch(samples) + + assert batch.input_ids.tolist() == [ + [ + 1, + IMG_START_ID, + IMAGE_TOKEN_ID, + IMG_END_ID, + 21, + PAD_AND_END_ID, + PAD_AND_END_ID, + PAD_AND_END_ID, + PAD_AND_END_ID, + 2, + IMG_START_ID, + IMAGE_TOKEN_ID, + IMG_END_ID, + 31, + 32, + PAD_AND_END_ID, + PAD_AND_END_ID, + PAD_AND_END_ID, + ] + ] + assert batch.attention_mask is None + assert batch.cu_seqlens_q.tolist() == [0, 261, 523] + assert batch.cu_seqlens_q_padded.tolist() == [0, 264, 528] + assert batch.max_seqlen_q.item() == 264 + assert batch.total_tokens == 528 + assert batch.num_image_tiles.tolist() == [1, 1] + packed_seq_params = get_packed_seq_params(encoder.encode_batch(batch)) + assert packed_seq_params.seq_idx.shape == (1, 528) + assert packed_seq_params.seq_idx[0, :264].unique().tolist() == [0] + assert packed_seq_params.seq_idx[0, 264:].unique().tolist() == [1] + + +@pytest.mark.parametrize("collapse_image_tokens", [False, True], ids=["canonical", "llava"]) +def test_hf_and_energon_packing_are_identical_for_image_video_audio(monkeypatch, collapse_image_tokens): + from PIL import Image + + monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens) + monkeypatch.setattr( + omni_collate, + "_patchify_frame", + lambda frame, *, height, width, patch_dim: torch.ones(2, 3), + ) + monkeypatch.setattr( + "megatron.bridge.models.nemotron_omni.nemotron_omni_utils.compute_mel_features", + lambda waveform, sampling_rate=16000, num_mel_bins=4: torch.ones(9, num_mel_bins), + ) + rows = [ + [1, IMG_START_ID, IMAGE_TOKEN_ID, IMG_END_ID, 21, PAD_AND_END_ID], + [ + 2, + IMG_START_ID, + IMAGE_TOKEN_ID, + IMG_END_ID, + IMG_START_ID, + IMAGE_TOKEN_ID, + IMG_END_ID, + 31, + 32, + PAD_AND_END_ID, + ], + ] + image = Image.new("RGB", (16, 16), color=32) + frames = [Image.new("RGB", (16, 16), color=value) for value in (0, 64, 128)] + source_samples = [ + _sample( + [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "image"}]}], + key="image-row", + imgs=[image], + audio=torch.tensor([0.0, 0.1, -0.1]), + ), + _sample( + [{"role": "user", "content": [{"type": "video"}, {"type": "text", "text": "video"}]}], + key="video-row", + videos=[frames], + audio=torch.tensor([0.2, 0.0, -0.2]), + ), + ] + normalizer = NemotronOmniTaskEncoder(processor=_Processor(rows), pad_to_multiple_of=1) + normalized_samples = [normalizer.encode_sample(sample) for sample in source_samples] + examples = [sample.example for sample in normalized_samples] + collate_kwargs = { + "sequence_length": 1024, + "enable_in_batch_packing": True, + "in_batch_packing_pad_to_multiple_of": 8, + "use_temporal_video_embedder": True, + "temporal_patch_size": 2, + "num_mel_bins": 4, + "pad_to_multiple_of": 1, + } + + hf_collate_fn = ( + omni_collate.nemotron_omni_llava_collate_fn + if collapse_image_tokens + else omni_collate.nemotron_omni_expanded_collate_fn + ) + hf_batch = hf_collate_fn(examples, _Processor(rows), **collate_kwargs) + energon_encoder = NemotronOmniTaskEncoder( + processor=_Processor(rows), + seq_length=1024, + enable_in_batch_packing=True, + in_batch_packing_pad_to_multiple_of=8, + use_temporal_video_embedder=True, + temporal_patch_size=2, + num_mel_bins=4, + pad_to_multiple_of=1, + collapse_image_tokens=collapse_image_tokens, + ) + energon_batch = energon_encoder.encode_batch(energon_encoder.batch(normalized_samples)) + + tensor_keys = ( + "input_ids", + "labels", + "loss_mask", + "position_ids", + "sound_clips", + "sound_length", + "imgs_sizes", + "num_frames", + "num_image_tiles", + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + "max_seqlen_q", + "max_seqlen_kv", + ) + for key in tensor_keys: + assert torch.equal(hf_batch[key], energon_batch[key]), key + assert hf_batch["attention_mask"] is energon_batch["attention_mask"] is None + assert hf_batch["total_tokens"] == energon_batch["total_tokens"] == 800 + assert hf_batch["cu_seqlens_q"].tolist() == [0, 265, 789] + assert hf_batch["cu_seqlens_q_padded"].tolist() == [0, 272, 800] + if not collapse_image_tokens: + assert hf_batch["input_ids"].shape == (1, 800) + assert hf_batch["padding_mask"].sum().item() == 11 + assert torch.equal(hf_batch["padding_mask"], energon_batch["padding_mask"]) + assert torch.equal( + hf_batch["visual_inputs"].pixel_values, + energon_batch["visual_inputs"].pixel_values, + ) + + +def test_energon_llava_temporal_video_refuses_unsafe_sequence_truncation(monkeypatch): + from PIL import Image + + monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens) + monkeypatch.setattr( + omni_collate, + "_patchify_frame", + lambda frame, *, height, width, patch_dim: torch.ones(2, 3), + ) + encoder = NemotronOmniTaskEncoder( + processor=_Processor([[1, IMG_START_ID, 97, IMG_END_ID, IMG_START_ID, 97, IMG_END_ID, 21, PAD_AND_END_ID]]), + seq_length=6, + use_temporal_video_embedder=True, + pad_to_multiple_of=1, + collapse_image_tokens=True, + ) + encoded = encoder.encode_sample( + _sample( + [{"role": "user", "content": [{"type": "video"}]}], + videos=[[Image.new("RGB", (16, 16)), Image.new("RGB", (16, 16)), Image.new("RGB", (16, 16))]], + ) + ) + + with pytest.raises(ValueError, match="cannot fit the rectangular multimodal batch"): + encoder.batch([encoded]) + + def test_energon_canonical_collator_owns_complete_thd_packing(monkeypatch): monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens) encoder = NemotronOmniTaskEncoder( @@ -429,10 +670,10 @@ def test_energon_canonical_collator_owns_complete_thd_packing(monkeypatch): assert batch.input_ids.tolist() == [[1, 2, 3, PAD_AND_END_ID, 4, 5, PAD_AND_END_ID, PAD_AND_END_ID]] assert batch.attention_mask is None - assert getattr(batch, "padding_mask", None) is None + assert batch.padding_mask.tolist() == [[False, False, False, True, False, False, True, True]] assert batch.cu_seqlens_q.tolist() == [0, 3, 5] assert batch.cu_seqlens_q_padded.tolist() == [0, 4, 8] assert batch.total_tokens == 8 assert encoded["tokens"] is batch.input_ids - assert encoded.get("padding_mask") is None + assert encoded["padding_mask"] is batch.padding_mask assert get_packed_seq_params(encoded).tokens_per_sample is None diff --git a/tests/unit_tests/data/packing/test_in_batch.py b/tests/unit_tests/data/packing/test_in_batch.py index 4c357c83f6..015731b65a 100644 --- a/tests/unit_tests/data/packing/test_in_batch.py +++ b/tests/unit_tests/data/packing/test_in_batch.py @@ -162,6 +162,41 @@ def test_packing_with_pad_to_multiple_of(self): # max_seqlen should be 6 (longest padded sequence) assert max_seqlen.item() == 6 + def test_packing_marks_only_physical_alignment_gaps_as_padding(self): + """MoE routing can exclude collator-inserted THD alignment gaps.""" + batch = { + "tokens": torch.tensor([[1, 2, 3, 0, 0], [4, 5, 0, 0, 0]]), + "labels": torch.tensor([[2, 3, -100, -100, -100], [5, -100, -100, -100, -100]]), + "loss_mask": torch.tensor([[1.0, 1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0, 0.0]]), + "attention_mask": torch.tensor([[1, 1, 1, 0, 0], [1, 1, 0, 0, 0]]), + "position_ids": torch.arange(5).unsqueeze(0).expand(2, -1), + } + + pack_right_padded_sequence_batch_to_mcore_thd( + batch, + pad_token_id=0, + pad_to_multiple_of=4, + emit_padding_mask=True, + ) + + assert batch["padding_mask"].dtype == torch.bool + assert batch["padding_mask"].tolist() == [[False, False, False, True, False, False, True, True]] + assert batch["cu_seqlens_q"].tolist() == [0, 3, 5] + assert batch["cu_seqlens_q_padded"].tolist() == [0, 4, 8] + + def test_packing_removes_stale_padding_mask_when_not_emitted(self): + """Packing without mask emission removes an incompatible input mask.""" + batch = { + "tokens": torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]]), + "position_ids": torch.arange(4).unsqueeze(0).expand(2, -1), + "padding_mask": torch.tensor([[False, False, False, True], [False, False, True, True]]), + } + + pack_right_padded_sequence_batch_to_mcore_thd(batch, pad_token_id=0) + + assert batch["tokens"].shape == (1, 5) + assert "padding_mask" not in batch + def test_packing_with_larger_multiple(self): """Test packing with larger pad_to_multiple_of (e.g., for CP=4).""" tokens = torch.tensor( diff --git a/tests/unit_tests/examples/test_nemotron_omni_inference.py b/tests/unit_tests/examples/test_nemotron_omni_inference.py index a098ad23cc..ee478ec5d3 100644 --- a/tests/unit_tests/examples/test_nemotron_omni_inference.py +++ b/tests/unit_tests/examples/test_nemotron_omni_inference.py @@ -22,6 +22,15 @@ _EXAMPLE_ROOT = Path(__file__).parents[3] / "examples" / "models" / "nemotron" / "nemotron_3_omni" +@pytest.mark.unit +def test_hf_revision_kwargs(): + script_globals = runpy.run_path(_EXAMPLE_ROOT / "hf_to_megatron_generate_nemotron_omni.py") + revision_kwargs = script_globals["_hf_revision_kwargs"] + + assert revision_kwargs(None) == {} + assert revision_kwargs("immutable-revision") == {"revision": "immutable-revision"} + + @pytest.mark.unit @pytest.mark.parametrize( "script_name", @@ -31,7 +40,7 @@ "valor32k_avqa_inference.py", ], ) -def test_inference_forward_step_forwards_num_image_tiles_to_pipeline_stages(script_name): +def test_inference_forward_step_uses_canonical_expanded_sequence_contract(script_name): script_globals = runpy.run_path(_EXAMPLE_ROOT / script_name) iterator_cls = script_globals["SingleBatchIterator"] forward_step = script_globals["vlm_forward_step"] @@ -55,4 +64,43 @@ def __call__(self, **kwargs): output, _ = forward_step(iterator, _Model()) assert output.shape == (1, 3, 8) - assert seen["num_image_tiles"] is num_image_tiles + assert "num_image_tiles" not in seen + + +@pytest.mark.unit +def test_generic_inference_processes_heterogeneous_source_images(monkeypatch): + script_globals = runpy.run_path(_EXAMPLE_ROOT / "hf_to_megatron_generate_nemotron_omni.py") + process_inputs = script_globals["process_image_inputs"] + pixel_values = [ + torch.arange(3 * 32 * 16, dtype=torch.float32).reshape(3, 32, 16), + torch.arange(3 * 16 * 32, dtype=torch.float32).reshape(3, 16, 32), + ] + + class _Tokenizer: + def apply_chat_template(self, messages, **kwargs): + assert messages[-1]["content"].count("") == 2 + return "rendered prompt" + + class _Inputs: + input_ids = torch.tensor([[1, 2, 3]]) + num_patches = torch.tensor([1, 1]) + + def __init__(self): + self.pixel_values = pixel_values + + class _Processor: + def __call__(self, *, text, images, return_tensors): + assert text == ["rendered prompt"] + assert images == ["first.png", "second.png"] + assert return_tensors == "pt" + return _Inputs() + + monkeypatch.setitem(process_inputs.__globals__, "load_image", lambda path: path) + input_ids, packed, num_patches, imgs_sizes = process_inputs( + _Tokenizer(), _Processor(), "first.png,second.png", "describe" + ) + + assert torch.equal(input_ids, torch.tensor([[1, 2, 3]])) + assert packed.shape == (1, 4, 3 * 16 * 16) + assert torch.equal(num_patches, torch.tensor([1, 1])) + assert torch.equal(imgs_sizes, torch.tensor([[32, 16], [16, 32]])) diff --git a/tests/unit_tests/models/nemotron_omni/test_collator_owned_packing_distributed.py b/tests/unit_tests/models/nemotron_omni/test_collator_owned_packing_distributed.py index 1efa07faa5..3ed2b0b487 100644 --- a/tests/unit_tests/models/nemotron_omni/test_collator_owned_packing_distributed.py +++ b/tests/unit_tests/models/nemotron_omni/test_collator_owned_packing_distributed.py @@ -70,6 +70,10 @@ def test_collator_owned_thd_tensors_use_one_real_cp_partition_index() -> None: position_ids = torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3]], device="cuda") labels = input_ids.clone() loss_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0]], device="cuda") + padding_mask = torch.tensor( + [[False, False, False, True, False, False, False, True]], + device="cuda", + ) cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32, device="cuda") cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32, device="cuda") packed_seq_params = PackedSeqParams( @@ -90,6 +94,7 @@ def test_collator_owned_thd_tensors_use_one_real_cp_partition_index() -> None: local_attention_mask, local_labels, local_loss_mask, + local_padding_mask, loss_mask_was_sliced, ) = model._apply_context_parallel_sharding( input_ids=input_ids, @@ -98,6 +103,7 @@ def test_collator_owned_thd_tensors_use_one_real_cp_partition_index() -> None: attention_mask=None, labels=labels, loss_mask=loss_mask, + padding_mask=padding_mask, packed_seq_params=packed_seq_params, ) @@ -111,6 +117,7 @@ def test_collator_owned_thd_tensors_use_one_real_cp_partition_index() -> None: assert torch.equal(local_position_ids, position_ids.index_select(1, expected_index)) assert torch.equal(local_labels, labels.index_select(1, expected_index)) assert torch.equal(local_loss_mask, loss_mask.index_select(1, expected_index)) + assert torch.equal(local_padding_mask, padding_mask.index_select(1, expected_index)) assert local_attention_mask is None assert loss_mask_was_sliced is True assert packed_seq_params.cu_seqlens_q is cu_seqlens diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_conversion.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_conversion.py index f79bb5c3a1..9340261414 100644 --- a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_conversion.py +++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_conversion.py @@ -13,16 +13,18 @@ # limitations under the License. from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import MagicMock, Mock, patch import pytest import torch from megatron.core.activations import squared_relu +from safetensors.torch import save_file from torch import nn +from transformers import PretrainedConfig from megatron.bridge.models.conversion.auto_bridge import AutoBridge from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry -from megatron.bridge.models.conversion.model_bridge import get_model_bridge +from megatron.bridge.models.conversion.model_bridge import HFWeightTuple, get_model_bridge from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM from megatron.bridge.models.nemotron_omni import nemotron_omni_provider as provider_module from megatron.bridge.models.nemotron_omni.modeling_nemotron_omni import NemotronOmniModel @@ -38,6 +40,7 @@ NemotronOmniModelProvider, ) from megatron.bridge.models.nemotron_vl.modeling_nemotron_vl import NemotronVLModel +from megatron.bridge.models.nemotron_vl.nemotron_vl_bridge import NemotronVLBridge from megatron.bridge.training.config import ConfigContainer @@ -162,6 +165,47 @@ def test_nemotron_omni_provider_bridge_maps_public_config_fields(): assert provider.temporal_ckpt_compat is True serialized = ConfigContainer._convert_value_to_dict(provider) assert serialized["nemotron_omni_contract"] == NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT + assert serialized["has_sound"] is True + assert "add_sound_encoder" not in serialized + + +def test_nemotron_omni_provider_bridge_omits_sound_when_config_is_absent(): + hf_config = _mock_omni_hf_config() + del hf_config.sound_config + hf_pretrained = Mock(spec=PreTrainedCausalLM) + hf_pretrained.config = hf_config + + provider = NemotronOmniBridge().provider_bridge(hf_pretrained) + + assert provider.has_sound is False + assert provider.sound_config is None + assert provider.sound_context_token_id == 0 + + +def test_nemotron_omni_hf_config_export_preserves_sound_capability(): + provider = NemotronOmniModelProvider( + has_sound=True, + sound_context_token_id=27, + sound_config={"hidden_size": 128}, + ) + + hf_config = NemotronOmniBridge.megatron_to_hf_config(provider) + + assert hf_config["sound_config"] == {"hidden_size": 128} + assert hf_config["sound_context_token_id"] == 27 + + +def test_nemotron_omni_hf_config_export_omits_disabled_sound_capability(): + provider = NemotronOmniModelProvider( + has_sound=False, + sound_context_token_id=27, + sound_config={"hidden_size": 128}, + ) + + hf_config = NemotronOmniBridge.megatron_to_hf_config(provider) + + assert hf_config["sound_config"] is None + assert hf_config["sound_context_token_id"] is None def test_nemotron_omni_provider_rejects_static_resolution(): @@ -181,7 +225,12 @@ def test_nemotron_omni_provider_rejects_nonpositive_image_token_index(image_toke def test_nemotron_omni_provider_rejects_nonpositive_sound_token_index(): - provider = NemotronOmniModelProvider(image_token_index=18, has_sound=True, sound_context_token_id=0) + provider = NemotronOmniModelProvider( + image_token_index=18, + has_sound=True, + sound_context_token_id=0, + sound_config={}, + ) with pytest.raises(ValueError, match="requires a positive sound_context_token_id"): provider.finalize() @@ -213,6 +262,31 @@ def test_canonical_provider_builds_dedicated_model(monkeypatch): llava_factory.assert_not_called() +def test_nemotron_omni_provider_can_omit_sound_modules(): + provider = NemotronOmniModelProvider(has_sound=False) + + sound_model, sound_projection = provider._build_sound_modules(None, None, add_encoder=True) + + assert provider.has_sound is False + assert sound_model is None + assert sound_projection is None + + +def test_nemotron_omni_provider_builds_sound_modules_when_enabled(monkeypatch): + provider = NemotronOmniModelProvider(has_sound=True) + expected_sound_model = object() + expected_sound_projection = object() + monkeypatch.setattr(provider, "_build_sound_encoder", lambda: expected_sound_model) + monkeypatch.setattr(provider, "_build_sound_projection_config", lambda _: object()) + monkeypatch.setattr(provider_module, "get_language_mlp_submodules", lambda _: object()) + monkeypatch.setattr(provider_module, "MultimodalProjector", lambda **_: expected_sound_projection) + + sound_model, sound_projection = provider._build_sound_modules(None, None, add_encoder=True) + + assert sound_model is expected_sound_model + assert sound_projection is expected_sound_projection + + def test_nemotron_omni_vision_projection_uses_squared_relu(): provider = NemotronOmniModelProvider() @@ -239,6 +313,49 @@ def test_nemotron_omni_mapping_registry_includes_sound_mappings(): assert all(not name.startswith("llava_model.") for name in names) +def test_nemotron_omni_export_preserves_source_only_buffers(): + bridge = NemotronOmniBridge() + hf_pretrained = Mock(spec=PreTrainedCausalLM) + source_tensors = { + name: torch.full((2,), index, dtype=torch.float32) for index, name in enumerate(bridge._HF_PASSTHROUGH_KEYS) + } + hf_pretrained.state = MagicMock() + hf_pretrained.state.source.get_all_keys.return_value = [ + "language_model.weight", + *source_tensors, + ] + hf_pretrained.state.__getitem__ = Mock(side_effect=source_tensors.__getitem__) + converted = HFWeightTuple("language_model.weight", torch.ones(1)) + + with patch.object(NemotronVLBridge, "stream_weights_megatron_to_hf", return_value=iter([converted])): + exported = list(bridge.stream_weights_megatron_to_hf([], hf_pretrained)) + + assert exported[0] == converted + exported_buffers = {item.param_name: item.weight for item in exported[1:]} + assert exported_buffers.keys() == source_tensors.keys() + for name, source_tensor in source_tensors.items(): + assert torch.equal(exported_buffers[name], source_tensor) + + +def test_nemotron_omni_config_only_export_preserves_source_only_buffers(tmp_path): + bridge = NemotronOmniBridge() + source_tensors = { + name: torch.full((2,), index, dtype=torch.float32) for index, name in enumerate(bridge._HF_PASSTHROUGH_KEYS) + } + save_file(source_tensors, tmp_path / "model.safetensors") + hf_config = PretrainedConfig() + hf_config.name_or_path = str(tmp_path) + converted = HFWeightTuple("language_model.weight", torch.ones(1)) + + with patch.object(NemotronVLBridge, "stream_weights_megatron_to_hf", return_value=iter([converted])): + exported = list(bridge.stream_weights_megatron_to_hf([], hf_config)) + + exported_buffers = {item.param_name: item.weight for item in exported[1:]} + assert exported_buffers.keys() == source_tensors.keys() + for name, source_tensor in source_tensors.items(): + assert torch.equal(exported_buffers[name], source_tensor) + + def test_canonical_bridge_maps_super_mtp_config(): hf_config = _mock_omni_hf_config() hf_config.architectures = ["NemotronH_Super_Omni_Reasoning_V3"] @@ -292,7 +409,8 @@ def test_llava_bridge_retains_legacy_wrapper_namespace(): hf_pretrained = Mock(spec=PreTrainedCausalLM) hf_pretrained.config = _mock_omni_hf_config() - provider = NemotronOmniLlavaBridge().provider_bridge(hf_pretrained) + with pytest.warns(FutureWarning, match="NemotronOmniLlavaBridge is deprecated"): + provider = NemotronOmniLlavaBridge().provider_bridge(hf_pretrained) registry = NemotronOmniLlavaBridge().mapping_registry() assert isinstance(provider, NemotronOmniLlavaModelProvider) diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py index 444bd2eacc..726c02765d 100644 --- a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py +++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py @@ -32,12 +32,14 @@ NemotronOmniModel, _pixel_shuffle_dynamic_resolution, ) +from megatron.bridge.models.nemotron_omni.modeling_nemotron_omni_llava import NemotronOmniLlavaModel from megatron.bridge.models.nemotron_omni.nemotron_omni_provider import ( NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT, NEMOTRON_OMNI_LLAVA_CONTRACT, NemotronOmniLlavaModelProvider, NemotronOmniModelProvider, ) +from megatron.bridge.models.nemotron_vl.modeling_nemotron_vl import NemotronVLModel class _FakeLanguageModel(nn.Module): @@ -58,20 +60,47 @@ def forward(self, *, decoder_input, **kwargs): class _BoundaryModel(NemotronOmniModel): """CPU-only shell that exercises the real expanded-sequence forward.""" - def __init__(self, image_features): + def __init__(self, image_features, sound_features=None): nn.Module.__init__(self) self.pre_process = True self.image_token_index = 18 + self.sound_token_index = 19 self.context_parallel_lm = 1 self.sequence_parallel_lm = False self.config = SimpleNamespace(mtp_num_layers=None) self.language_model = _FakeLanguageModel() self.image_features = image_features + self.sound_features = torch.empty(0, 3) if sound_features is None else sound_features def _encode_images(self, images, imgs_sizes, vision_packed_seq_params, num_frames): del images, imgs_sizes, vision_packed_seq_params, num_frames return self.image_features + def _encode_sound(self, sound_clips, sound_length): + del sound_clips, sound_length + return self.sound_features + + +class _FakeSoundModel(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + self.config = SimpleNamespace(sound_pad_to_clip_duration=False) + + def forward(self, sound_clips, sound_length): + del sound_clips, sound_length + embeddings = torch.arange(12, dtype=torch.float32).reshape(2, 3, 2) + return embeddings, torch.tensor([2, 1]) + + +class _SoundEncoderBoundaryModel(NemotronOmniModel): + def __init__(self): + nn.Module.__init__(self) + self.sound_model = _FakeSoundModel() + self.sound_projection = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16) + with torch.no_grad(): + self.sound_projection.weight.copy_(torch.eye(2, dtype=torch.bfloat16)) + @dataclass class _TinyOmniProvider(NemotronOmniModelProvider): @@ -234,6 +263,22 @@ def test_llava_provider_preserves_existing_radio_cpe_default(): assert provider.radio_interpolate_only_cpe is True +def test_llava_model_emits_deprecation_notice(monkeypatch): + monkeypatch.setattr(NemotronVLModel, "__init__", lambda *_args, **_kwargs: None) + + with pytest.warns(FutureWarning, match="NemotronOmniLlavaModel is deprecated"): + NemotronOmniLlavaModel() + + +def test_llava_provider_emits_deprecation_notice(monkeypatch): + provider = NemotronOmniLlavaModelProvider(nemotron_omni_contract=NEMOTRON_OMNI_LLAVA_CONTRACT) + legacy_model = object() + monkeypatch.setattr(provider, "_provide_llava", lambda **_: legacy_model) + + with pytest.warns(FutureWarning, match="NemotronOmniLlavaModelProvider is deprecated"): + assert provider.provide() is legacy_model + + def test_dynamic_resolution_pixel_shuffle_groups_spatial_2x2_blocks(): features = torch.arange(2 * 4 * 2, dtype=torch.float32).reshape(1, 8, 2) @@ -271,6 +316,71 @@ def test_image_forward_replaces_expanded_placeholders_without_changing_length(): assert torch.equal(output[3, 0], torch.tensor([9.0, 9.0, 9.0])) +def test_audio_forward_replaces_expanded_placeholders_without_changing_length(): + sound_features = torch.tensor([[101.0, 102.0, 103.0], [201.0, 202.0, 203.0]]) + model = _BoundaryModel(torch.empty(0, 3), sound_features) + input_ids = torch.tensor([[7, 19, 19, 9]]) + + output = model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids, dtype=torch.bool), + sound_clips=torch.ones(1, 8, 2), + sound_length=torch.tensor([8]), + ) + + assert output.shape == (4, 1, 3) + assert torch.equal(output[0, 0], torch.tensor([7.0, 7.0, 7.0])) + assert torch.equal(output[1, 0], sound_features[0]) + assert torch.equal(output[2, 0], sound_features[1]) + assert torch.equal(output[3, 0], torch.tensor([9.0, 9.0, 9.0])) + + +def test_sound_encoder_drops_padded_rows_and_preserves_sample_order(): + model = _SoundEncoderBoundaryModel() + + encoded = model._encode_sound( + torch.ones(2, 8, 2), + torch.tensor([8, 4]), + ) + assert torch.equal( + encoded, + torch.tensor( + [ + [0.0, 1.0], + [2.0, 3.0], + [6.0, 7.0], + ], + dtype=torch.bfloat16, + ), + ) + + +def test_real_parakeet_sound_encoder_matches_subsampled_placeholder_count(): + from megatron.bridge.models.nemotron_omni.nemotron_omni_sound import BridgeSoundEncoder + + config = SimpleNamespace( + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + intermediate_size=64, + num_mel_bins=8, + subsampling_factor=8, + conv_kernel_size=9, + use_bias=False, + sound_pad_to_clip_duration=False, + ) + model = _SoundEncoderBoundaryModel() + model.sound_model = BridgeSoundEncoder(config) + model.sound_projection = nn.Linear(config.hidden_size, 3, bias=False) + sound_length = torch.tensor([64, 40]) + + encoded = model._encode_sound(torch.randn(2, 64, config.num_mel_bins), sound_length) + + expected_lengths = model.sound_model.encoder._get_subsampling_output_length(sound_length) + assert encoded.shape == (int(expected_lengths.sum().item()), 3) + assert torch.isfinite(encoded).all() + + def test_text_only_control_preserves_language_embeddings(): model = _BoundaryModel(torch.empty(0, 3)) input_ids = torch.tensor([[7, 8, 9]]) @@ -370,6 +480,7 @@ def fake_get_packed_seq_cp_partition_indices(packed_seq_params, **kwargs): position_ids = torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3]]) labels = input_ids.clone() loss_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0]]) + padding_mask = torch.tensor([[False, False, False, True, False, False, False, True]]) cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32) cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32) packed_seq_params = PackedSeqParams( @@ -388,6 +499,7 @@ def fake_get_packed_seq_cp_partition_indices(packed_seq_params, **kwargs): position_ids=position_ids, labels=labels, loss_mask=loss_mask, + padding_mask=padding_mask, packed_seq_params=packed_seq_params, images=torch.ones(1), ) @@ -404,6 +516,7 @@ def fake_get_packed_seq_cp_partition_indices(packed_seq_params, **kwargs): assert torch.equal(local_loss_mask, loss_mask.index_select(1, cp_index)) assert model.language_model.last_kwargs["packed_seq_params"] is packed_seq_params assert torch.equal(model.language_model.last_kwargs["labels"], labels.index_select(1, cp_index)) + assert "padding_mask" not in model.language_model.last_kwargs assert model.language_model.last_kwargs["attention_mask"] is None @@ -465,6 +578,7 @@ def test_real_radio_image_forward_with_collator_owned_cp1_packing( with torch.no_grad(): output = model( input_ids=input_ids, + padding_mask=torch.zeros_like(input_ids, dtype=torch.bool), packed_seq_params=caller_packed_seq_params, pixel_values=torch.randn(1, 3, 32, 32, device="cuda"), imgs_sizes=torch.tensor([[32, 32]], dtype=torch.int32, device="cuda"), @@ -486,6 +600,7 @@ def test_real_packed_multimodal_optimizer_step(single_rank_model_parallel): input_ids = torch.tensor([[7, 18, 9, 0, 11, 18, 12, 0]], device="cuda") labels = torch.tensor([[18, 9, -100, -100, 18, 12, -100, -100]], device="cuda") loss_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0]], device="cuda") + padding_mask = torch.tensor([[False, False, False, True, False, False, False, True]], device="cuda") cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32, device="cuda") cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32, device="cuda") packed_seq_params = PackedSeqParams( @@ -504,6 +619,7 @@ def test_real_packed_multimodal_optimizer_step(single_rank_model_parallel): input_ids=input_ids, labels=labels, loss_mask=loss_mask, + padding_mask=padding_mask, packed_seq_params=packed_seq_params, pixel_values=torch.randn(2, 3, 32, 32, device="cuda"), imgs_sizes=torch.tensor([[32, 32], [32, 32]], dtype=torch.int32, device="cuda"), @@ -523,6 +639,27 @@ def test_real_packed_multimodal_optimizer_step(single_rank_model_parallel): assert not torch.equal(updated_parameter, parameter_before_step) +@pytest.mark.run_only_on("GPU") +def test_real_radio_multiframe_video_forward(single_rank_model_parallel): + del single_rank_model_parallel + provider = _TinyOmniProvider() + provider.finalize() + model = provider.provide().cuda().eval() + input_ids = torch.tensor([[7, 18, 9, 10]], device="cuda") + + with torch.no_grad(): + output = model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids, dtype=torch.bool), + pixel_values=torch.randn(2, 3, 32, 32, device="cuda"), + imgs_sizes=torch.tensor([[32, 32], [32, 32]], dtype=torch.int32, device="cuda"), + num_frames=torch.tensor([2], dtype=torch.int32, device="cuda"), + ) + + assert output.shape == (1, 4, 128) + assert torch.isfinite(output).all() + + @pytest.mark.run_only_on("GPU") def test_packed_mamba_resets_state_between_samples(single_rank_model_parallel): del single_rank_model_parallel @@ -530,7 +667,7 @@ def test_packed_mamba_resets_state_between_samples(single_rank_model_parallel): provider.finalize() model = provider.provide().cuda().eval() - def forward(input_ids, cu_seqlens, cu_seqlens_padded): + def forward(input_ids, padding_mask, cu_seqlens, cu_seqlens_padded): caller_packed_seq_params = PackedSeqParams( qkv_format="thd", cu_seqlens_q=cu_seqlens, @@ -543,22 +680,26 @@ def forward(input_ids, cu_seqlens, cu_seqlens_padded): ) return model( input_ids=input_ids, + padding_mask=padding_mask, packed_seq_params=caller_packed_seq_params, ) input_ids = torch.tensor([[7, 8, 9, 0, 11, 12, 0, 0]], device="cuda") + padding_mask = torch.tensor([[False, False, False, True, False, False, True, True]], device="cuda") cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32, device="cuda") cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32, device="cuda") with torch.no_grad(): - packed_output = forward(input_ids, cu_seqlens, cu_seqlens_padded) + packed_output = forward(input_ids, padding_mask, cu_seqlens, cu_seqlens_padded) first_output = forward( input_ids[:, :4], + padding_mask[:, :4], torch.tensor([0, 3], dtype=torch.int32, device="cuda"), torch.tensor([0, 4], dtype=torch.int32, device="cuda"), ) second_output = forward( input_ids[:, 4:], + padding_mask[:, 4:], torch.tensor([0, 2], dtype=torch.int32, device="cuda"), torch.tensor([0, 4], dtype=torch.int32, device="cuda"), ) diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_utils.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_utils.py index 371633d057..0c9713518e 100644 --- a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_utils.py +++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_utils.py @@ -16,11 +16,13 @@ import torch from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import ( + inference_expanded_image_token_counts, inference_merged_sequence_length, inference_num_image_tiles, select_inference_next_token, temporal_model_frames, ) +from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import adjust_image_tokens def test_temporal_model_frames_duplicates_single_frame_for_temporal_embedder(): @@ -83,22 +85,126 @@ def test_inference_num_image_tiles_rejects_unshufflable_image_grid(): inference_num_image_tiles(torch.tensor([[528, 512]]), patch_dim=16) -def test_inference_merged_sequence_length_uses_exact_image_replacements(): - input_ids = torch.tensor([[10, -200, 11, -200, 12]]) +def test_inference_expanded_image_token_counts_aggregates_dynamic_tiles_by_media(): + counts = inference_expanded_image_token_counts( + torch.tensor([256, 128, 64]), + torch.tensor([2, 1]), + ) + + assert counts.tolist() == [384, 64] + + +def test_inference_expanded_image_token_counts_applies_temporal_feature_width(): + counts = inference_expanded_image_token_counts( + torch.ones(3, dtype=torch.int), + torch.ones(3, dtype=torch.int), + feature_multiplier=256, + ) + + assert counts.tolist() == [256, 256, 256] + + +def test_inference_expanded_image_token_counts_rejects_incomplete_tile_ownership(): + with pytest.raises(ValueError, match="account for every tile"): + inference_expanded_image_token_counts(torch.tensor([256, 128]), torch.tensor([1])) + - dynamic_length = inference_merged_sequence_length( - input_ids, - image_token_index=-200, - num_image_tiles=torch.tensor([3, 2]), - image_seq_len=1, +def test_canonical_dynamic_pre_expansion_preserves_legacy_merged_length(): + image_token_id = -200 + img_start_id = -201 + img_end_id = -202 + processor_input_ids = torch.tensor([[10, img_start_id, image_token_id, img_end_id, 11]]) + tile_feature_counts = inference_num_image_tiles( + torch.tensor([[512, 512], [512, 256]]), + patch_dim=16, + ) + + legacy_compact_ids = adjust_image_tokens( + processor_input_ids, + torch.tensor([2]), + img_start_id, + img_end_id, ) - temporal_length = inference_merged_sequence_length( - input_ids, - image_token_index=-200, - num_image_tiles=torch.tensor([1, 1]), - image_seq_len=256, + with pytest.warns(FutureWarning, match="deprecated"): + legacy_merged_length = inference_merged_sequence_length( + legacy_compact_ids, + image_token_index=image_token_id, + num_image_tiles=tile_feature_counts, + image_seq_len=1, + ) + + expanded_counts = inference_expanded_image_token_counts(tile_feature_counts, torch.tensor([2])) + canonical_input_ids = adjust_image_tokens( + processor_input_ids, + expanded_counts, + img_start_id, + img_end_id, ) + assert tile_feature_counts.tolist() == [256, 128] + assert expanded_counts.tolist() == [384] + assert canonical_input_ids.shape[1] == legacy_merged_length + assert int((canonical_input_ids == image_token_id).sum()) == 384 + + +def test_canonical_temporal_pre_expansion_preserves_legacy_merged_length(): + image_token_id = -200 + img_start_id = -201 + img_end_id = -202 + processor_input_ids = torch.tensor( + [[10, img_start_id, image_token_id, img_end_id, 11, img_start_id, image_token_id, img_end_id, 12]] + ) + tubelet_counts = inference_num_image_tiles( + torch.tensor([[512, 512]] * 4), + patch_dim=16, + num_frames=torch.tensor([4]), + temporal_patch_size=2, + ) + + with pytest.warns(FutureWarning, match="deprecated"): + legacy_merged_length = inference_merged_sequence_length( + processor_input_ids, + image_token_index=image_token_id, + num_image_tiles=tubelet_counts, + image_seq_len=256, + ) + + expanded_counts = inference_expanded_image_token_counts( + tubelet_counts, + torch.ones_like(tubelet_counts), + feature_multiplier=256, + ) + canonical_input_ids = adjust_image_tokens( + processor_input_ids, + expanded_counts, + img_start_id, + img_end_id, + ) + + assert tubelet_counts.tolist() == [1, 1] + assert expanded_counts.tolist() == [256, 256] + assert canonical_input_ids.shape[1] == legacy_merged_length + assert int((canonical_input_ids == image_token_id).sum()) == 512 + + +def test_inference_merged_sequence_length_uses_exact_image_replacements(): + input_ids = torch.tensor([[10, -200, 11, -200, 12]]) + + with pytest.warns(FutureWarning, match="deprecated"): + dynamic_length = inference_merged_sequence_length( + input_ids, + image_token_index=-200, + num_image_tiles=torch.tensor([3, 2]), + image_seq_len=1, + ) + with pytest.warns(FutureWarning, match="deprecated"): + temporal_length = inference_merged_sequence_length( + input_ids, + image_token_index=-200, + num_image_tiles=torch.tensor([1, 1]), + image_seq_len=256, + ) + assert dynamic_length == 8 assert temporal_length == 515 @@ -114,10 +220,11 @@ def test_select_inference_next_token_ignores_pipeline_padding_logits(): def test_inference_merged_sequence_length_rejects_misaligned_image_metadata(): - with pytest.raises(ValueError, match="Expected 2 num_image_tiles entries"): - inference_merged_sequence_length( - torch.tensor([[10, -200, 11, -200, 12]]), - image_token_index=-200, - num_image_tiles=torch.tensor([3]), - image_seq_len=1, - ) + with pytest.warns(FutureWarning, match="deprecated"): + with pytest.raises(ValueError, match="Expected 2 num_image_tiles entries"): + inference_merged_sequence_length( + torch.tensor([[10, -200, 11, -200, 12]]), + image_token_index=-200, + num_image_tiles=torch.tensor([3]), + image_seq_len=1, + ) diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index f8a8ee225f..b71ca936d9 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -17,6 +17,7 @@ """ import json +from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, PropertyMock, patch @@ -1072,7 +1073,9 @@ def test_from_auto_config_happy_path(self, tmp_path): "megatron.bridge.models.conversion.utils.conform_config_to_reference", return_value={"vocab_size": 64000}, ) as mock_conform: - with patch.object(AutoBridge, "from_hf_config", side_effect=[first_bridge, second_bridge]): + with patch.object( + AutoBridge, "from_hf_config", side_effect=[first_bridge, second_bridge] + ) as mock_from_config: bridge = AutoBridge.from_auto_config(str(ckpt_dir), hf_model_id) assert bridge is second_bridge @@ -1080,6 +1083,7 @@ def test_from_auto_config_happy_path(self, tmp_path): mock_auto_cfg.assert_called_once_with(hf_model_id, trust_remote_code=False) mock_load_cfg.assert_called_once_with(str(ckpt_dir)) mock_conform.assert_called_once_with({"vocab_size": 64000}, {"vocab_size": 32000}) + assert mock_from_config.call_args_list[1].args[0].name_or_path == hf_model_id def test_from_auto_config_uses_latest_iter_run_config(self, tmp_path): """from_auto_config falls back to latest iter_* directory for run_config.yaml.""" @@ -1793,7 +1797,11 @@ def test_import_ckpt_basic(self, mock_from_hf_pretrained, mock_to_megatron_model mock_bridge.save_megatron_model = Mock() # Test import_ckpt - AutoBridge.import_ckpt("meta-llama/Meta-Llama-3-8B", "./megatron_checkpoint") + with patch( + "megatron.bridge.training.model_load_save.temporary_distributed_context", + return_value=nullcontext(), + ): + AutoBridge.import_ckpt("meta-llama/Meta-Llama-3-8B", "./megatron_checkpoint") # Assertions mock_from_hf_pretrained.assert_called_once_with("meta-llama/Meta-Llama-3-8B") @@ -1821,13 +1829,17 @@ def test_import_ckpt_with_kwargs(self, mock_from_hf_pretrained, mock_to_megatron mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {} # Test import_ckpt with kwargs - AutoBridge.import_ckpt( - "./local_model", - "./megatron_checkpoint", - torch_dtype=torch.float16, - device_map="auto", - revision="0123456789abcdef", # pragma: allowlist secret - ) + with patch( + "megatron.bridge.training.model_load_save.temporary_distributed_context", + return_value=nullcontext(), + ): + AutoBridge.import_ckpt( + "./local_model", + "./megatron_checkpoint", + torch_dtype=torch.float16, + device_map="auto", + revision="0123456789abcdef", # pragma: allowlist secret + ) # Assertions mock_from_hf_pretrained.assert_called_once_with( @@ -1859,12 +1871,16 @@ def test_import_ckpt_with_low_memory_save( mock_bridge.save_megatron_model = Mock() mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {} - AutoBridge.import_ckpt( - "meta-llama/Meta-Llama-3-8B", - "./megatron_checkpoint", - low_memory_save=True, - torch_dtype=torch.bfloat16, - ) + with patch( + "megatron.bridge.training.model_load_save.temporary_distributed_context", + return_value=nullcontext(), + ): + AutoBridge.import_ckpt( + "meta-llama/Meta-Llama-3-8B", + "./megatron_checkpoint", + low_memory_save=True, + torch_dtype=torch.bfloat16, + ) mock_from_hf_pretrained.assert_called_once_with( "meta-llama/Meta-Llama-3-8B", @@ -1878,6 +1894,50 @@ def test_import_ckpt_with_low_memory_save( low_memory_save=True, ) + @patch("megatron.bridge.training.model_load_save.temporary_distributed_context") + @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=False) + @patch.object(AutoBridge, "from_hf_pretrained") + def test_import_ckpt_scopes_standalone_cpu_state_to_gloo_context( + self, + mock_from_hf_pretrained, + mock_dist_is_initialized, + mock_temporary_distributed_context, + ): + """Standalone CPU import uses the shared temporary Gloo lifecycle.""" + mock_bridge = Mock(spec=AutoBridge) + mock_bridge.to_megatron_model.return_value = [Mock()] + mock_bridge.save_megatron_model = Mock() + mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {} + mock_from_hf_pretrained.return_value = mock_bridge + + AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") + + mock_dist_is_initialized.assert_called_once_with() + mock_temporary_distributed_context.assert_called_once_with(backend="gloo") + mock_temporary_distributed_context.return_value.__enter__.assert_called_once_with() + mock_temporary_distributed_context.return_value.__exit__.assert_called_once() + + @patch("megatron.bridge.training.model_load_save.temporary_distributed_context") + @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=True) + @patch.object(AutoBridge, "from_hf_pretrained") + def test_import_ckpt_preserves_existing_distributed_context( + self, + mock_from_hf_pretrained, + mock_dist_is_initialized, + mock_temporary_distributed_context, + ): + """Import reuses distributed state owned by its caller.""" + mock_bridge = Mock(spec=AutoBridge) + mock_bridge.to_megatron_model.return_value = [Mock()] + mock_bridge.save_megatron_model = Mock() + mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {} + mock_from_hf_pretrained.return_value = mock_bridge + + AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") + + mock_dist_is_initialized.assert_called_once_with() + mock_temporary_distributed_context.assert_not_called() + def test_export_ckpt_basic(self): """Test basic export_ckpt functionality.""" # Setup mocks diff --git a/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py b/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py index f5112d4e32..1e3645c6b2 100644 --- a/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py +++ b/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py @@ -25,7 +25,7 @@ NemotronOmniEnergonTaskEncoderConfig, ) from megatron.bridge.data.collators.registry import resolve_model_collate -from megatron.bridge.models.nemotron_omni.data.collate_fn import nemotron_omni_collate_fn +from megatron.bridge.models.nemotron_omni.data.collate_fn import nemotron_omni_expanded_collate_fn from megatron.bridge.training.config import ConfigContainer from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global @@ -46,6 +46,7 @@ class _FakeModelCfg: dynamic_resolution = True + has_sound = True def finalize(self): return None @@ -145,20 +146,49 @@ def test_cord_v2_sft_recipe_uses_hf_dataset_config(fake_processor): _assert_common_config(cfg) assert isinstance(cfg.dataset, DirectHFSFTDatasetConfig) assert cfg.dataset.hf_processor_path == _TEST_HF_ID + assert cfg.dataset.trust_remote_code is True assert cfg.dataset.source.dataset_name == "cord_v2" - assert resolve_model_collate("NemotronH_Nano_Omni_Reasoning_V3Processor") is nemotron_omni_collate_fn + assert resolve_model_collate("NemotronH_Nano_Omni_Reasoning_V3Processor") is nemotron_omni_expanded_collate_fn assert cfg.dataset.enable_in_batch_packing is False assert cfg.dataset.dataloader_type == "cyclic" assert cfg.model.temporal_patch_dim == 1 + assert cfg.model.has_sound is False assert cfg.model.freeze_sound_projection is False assert cfg.peft is None +def test_cord_v2_long_context_sft_recipe_enables_packing_and_cp(fake_processor): + cfg = _build_config( + _h100_recipe_module.nemotron_omni_cord_v2_long_context_sft_8gpu_h100_bf16_config, + fake_processor, + ) + + assert isinstance(cfg.dataset, DirectHFSFTDatasetConfig) + assert cfg.dataset.trust_remote_code is True + assert cfg.model.seq_length == 8192 + assert cfg.model.context_parallel_size == 2 + assert cfg.model.calculate_per_token_loss is True + assert cfg.train.global_batch_size == 64 + assert cfg.train.micro_batch_size == 2 + assert cfg.optimizer.use_precision_aware_optimizer is True + assert cfg.optimizer.main_grads_dtype == torch.bfloat16 + assert cfg.optimizer.main_params_dtype == torch.float16 + assert cfg.optimizer.store_param_remainders is True + assert cfg.optimizer.exp_avg_dtype == torch.bfloat16 + assert cfg.optimizer.exp_avg_sq_dtype == torch.bfloat16 + assert cfg.mixed_precision.grad_reduce_in_fp32 is False + assert cfg.ddp.grad_reduce_in_fp32 is False + assert cfg.dataset.seq_length == 8192 + assert cfg.dataset.enable_in_batch_packing is True + assert cfg.dataset.in_batch_packing_pad_to_multiple_of == 8 + + def test_cord_v2_peft_recipe_configures_lora_and_freezing(fake_processor): cfg = _build_config(_recipe_module.nemotron_omni_cord_v2_peft_config, fake_processor) _assert_common_config(cfg) assert isinstance(cfg.dataset, DirectHFSFTDatasetConfig) + assert cfg.dataset.trust_remote_code is True assert cfg.dataset.dataloader_type == "cyclic" assert cfg.peft is not None assert cfg.peft.target_modules == [ @@ -173,6 +203,7 @@ def test_cord_v2_peft_recipe_configures_lora_and_freezing(fake_processor): assert cfg.peft.alpha == 32 assert cfg.checkpoint.load is None assert cfg.model.freeze_vision_projection is True + assert cfg.model.has_sound is False assert cfg.model.freeze_sound_projection is True @@ -189,9 +220,11 @@ def test_valor32k_sft_recipe_uses_temporal_omni_task_encoder_config(fake_process assert cfg.dataset.task_encoder.num_mel_bins == 128 assert cfg.dataset.task_encoder.use_temporal_video_embedder is True assert cfg.dataset.task_encoder.patch_dim == 16 + assert cfg.dataset.task_encoder.collapse_image_tokens is False assert cfg.model.temporal_patch_dim == 2 assert cfg.model.separate_video_embedder is True assert cfg.model.temporal_ckpt_compat is True + assert cfg.model.has_sound is True assert cfg.model.freeze_sound_projection is False assert cfg.peft is None @@ -216,4 +249,5 @@ def test_valor32k_peft_recipe_configures_lora_and_freezing(fake_processor): assert cfg.peft.alpha == 32 assert cfg.checkpoint.load is None assert cfg.model.freeze_vision_projection is True + assert cfg.model.has_sound is True assert cfg.model.freeze_sound_projection is True diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index ba08125a19..7a34025aaa 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -269,6 +269,16 @@ def test_hf_path_receives_ones_like_attention_mask(self): assert call_kwargs["attention_mask"].shape == input_ids.shape assert torch.equal(call_kwargs["attention_mask"], expected_mask) + def test_hf_text_only_path_selects_composite_language_backbone(self): + """Text-only comparisons bypass a composite model's media-required forward.""" + composite_model = MagicMock() + language_model = torch.nn.Linear(3, 3) + composite_model.language_model = language_model + + with patch.object(compare, "print_rank_0"): + assert compare._get_hf_forward_model(composite_model, pixel_values=None) is language_model + assert compare._get_hf_forward_model(composite_model, pixel_values=torch.ones(1)) is composite_model + def test_hf_path_receives_multimodal_token_type_ids(self): """Gemma 3 token types reach HF so its image attention mask matches Megatron.""" mock_hf_model = MagicMock() @@ -367,3 +377,36 @@ def test_hf_revision_is_parsed_and_forwarded(self): assert args.hf_revision == revision assert compare._hf_revision_kwargs(args.hf_revision) == {"revision": revision} assert compare._hf_revision_kwargs(None) == {} + + def test_hf_loader_uses_one_device_without_hf_tensor_parallelism(self): + """Load the HF reference on one device without a Transformers TP plan.""" + args = compare.build_parser().parse_args( + [ + "--hf_model_path", + "org/model", + "--prompt", + "Hello", + ] + ) + loaded_model = MagicMock() + model_class = MagicMock() + model_class.__name__ = "MockModel" + model_class.from_pretrained.return_value = loaded_model + loaded_model.to.return_value = loaded_model + loaded_model.eval.return_value = loaded_model + + with ( + patch.object(compare, "_is_rank_0", return_value=True), + patch.object(compare, "get_model_class", return_value=model_class), + patch.object(compare, "is_safe_repo", return_value=True), + patch.object(compare, "print_rank_0"), + ): + result = compare._load_hf_model(args, is_vl_model=False) + + assert result is loaded_model + model_class.from_pretrained.assert_called_once() + load_kwargs = model_class.from_pretrained.call_args.kwargs + assert "device_map" not in load_kwargs + assert "tp_plan" not in load_kwargs + assert "tp_size" not in load_kwargs + loaded_model.to.assert_called_once_with("cuda") diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index ee93303cb5..833ef4cef8 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/tests/unit_tests/training/test_model_load_save.py @@ -195,7 +195,11 @@ def test_temporary_distributed_context_gloo(self, mock_os, mock_socket, mock_par mock_socket_instance.getsockname.return_value = ("localhost", 12345) mock_socket.socket.return_value.__enter__.return_value = mock_socket_instance - with temporary_distributed_context(backend="gloo"): + with ( + patch("megatron.bridge.training.model_load_save.torch.cuda.is_available", return_value=False), + patch("megatron.core.tensor_parallel.model_parallel_cuda_manual_seed") as mock_seed, + temporary_distributed_context(backend="gloo"), + ): pass mock_dist.init_process_group.assert_called_once_with( @@ -204,6 +208,7 @@ def test_temporary_distributed_context_gloo(self, mock_os, mock_socket, mock_par mock_parallel_state.initialize_model_parallel.assert_called_once() mock_parallel_state.destroy_model_parallel.assert_called_once() mock_dist.destroy_process_group.assert_called_once() + mock_seed.assert_not_called() @patch("megatron.bridge.training.model_load_save.dist") @patch("megatron.bridge.training.model_load_save.parallel_state") diff --git a/tests/unit_tests/training/test_nemotron_omni_step.py b/tests/unit_tests/training/test_nemotron_omni_step.py index 195ff8c596..ae08084009 100644 --- a/tests/unit_tests/training/test_nemotron_omni_step.py +++ b/tests/unit_tests/training/test_nemotron_omni_step.py @@ -85,6 +85,7 @@ def _packed_pipeline_batch(): "input_ids": tokens, "labels": tokens.clone(), "loss_mask": torch.ones_like(tokens, dtype=torch.float32), + "padding_mask": torch.zeros_like(tokens, dtype=torch.bool), "position_ids": torch.arange(4).unsqueeze(0), "attention_mask": None, "visual_inputs": SimpleNamespace(pixel_values=torch.ones(1, 4, 8)), @@ -135,6 +136,7 @@ def test_middle_pipeline_stage_preserves_only_packed_attention_metadata(monkeypa assert result[0] is None assert result[7]["cu_seqlens_q"].tolist() == [0, 2, 4] assert result[7]["total_tokens"] == 4 + assert result[14].tolist() == [[False, False, False, False]] def test_middle_unpacked_pipeline_stage_does_not_consume_iterator(monkeypatch): @@ -144,7 +146,7 @@ def test_middle_unpacked_pipeline_stage_does_not_consume_iterator(monkeypatch): result = get_batch(data_iterator, _pipeline_cfg(packed=False), pg_collection=SimpleNamespace(pp=object())) - assert result == (None,) * 14 + assert result == (None,) * 15 assert next(data_iterator)["input_ids"].tolist() == [[18, 1, 18, 2]] @@ -166,6 +168,7 @@ def test_last_pipeline_stage_keeps_label_expansion_inputs_without_media(monkeypa assert moved["sound_clips"] is None assert moved["imgs_sizes"] is None assert moved["cu_seqlens_q"] is batch["cu_seqlens_q"] + assert moved["padding_mask"] is batch["padding_mask"] def test_packed_middle_pipeline_forward_uses_boundaries_without_input_tensors(monkeypatch): @@ -220,6 +223,7 @@ def __call__(self, **kwargs): assert model.kwargs["images"] is None assert model.kwargs["input_ids"] is None assert model.kwargs["packed_seq_params"].cu_seqlens_q.tolist() == [0, 2, 4] + assert model.kwargs["padding_mask"].tolist() == [[False, False, False, False]] def test_forward_unwraps_model_output_and_uses_expanded_loss_mask(monkeypatch): @@ -266,6 +270,7 @@ def __call__(self, **kwargs): None, None, None, + None, ), ) monkeypatch.setattr(