109109
110110from megatron .bridge import AutoBridge
111111from megatron .bridge .models .hf_pretrained .utils import is_safe_repo
112+ from megatron .bridge .models .nemotron_omni .inference_inputs import (
113+ is_nemotron_omni ,
114+ load_nemotron_omni_video ,
115+ nemotron_omni_reference_metadata ,
116+ prepare_nemotron_omni_inputs ,
117+ )
112118from megatron .bridge .utils .common_utils import disable_mtp_for_inference , get_last_rank , print_rank_0
113119from megatron .bridge .utils .safe_url import is_safe_public_http_url , safe_url_open
114120
@@ -225,6 +231,8 @@ def is_vision_language_model(
225231 ),
226232 ** _hf_revision_kwargs (revision ),
227233 )
234+ if is_nemotron_omni (config ):
235+ return True
228236
229237 # Check for VL model indicators in config
230238 model_type = getattr (config , "model_type" , "" ).lower ()
@@ -280,6 +288,8 @@ def __init__(
280288 image_grid_thw = None ,
281289 inference_context = None ,
282290 mm_token_type_ids = None ,
291+ imgs_sizes = None ,
292+ num_frames = None ,
283293 ):
284294 self .batch = dict (
285295 tokens = input_ids ,
@@ -295,6 +305,10 @@ def __init__(
295305 self .batch ["image_grid_thw" ] = image_grid_thw
296306 if mm_token_type_ids is not None :
297307 self .batch ["mm_token_type_ids" ] = mm_token_type_ids
308+ if imgs_sizes is not None :
309+ self .batch ["imgs_sizes" ] = imgs_sizes
310+ if num_frames is not None :
311+ self .batch ["num_frames" ] = num_frames
298312
299313 self ._yielded = False
300314
@@ -337,6 +351,9 @@ def vlm_forward_step(data_iterator, model, **kwargs) -> torch.Tensor:
337351 forward_args ["image_grid_thw" ] = batch ["image_grid_thw" ]
338352 if "mm_token_type_ids" in batch :
339353 forward_args ["mm_token_type_ids" ] = batch ["mm_token_type_ids" ]
354+ for key in ("imgs_sizes" , "num_frames" ):
355+ if key in batch :
356+ forward_args [key ] = batch [key ]
340357
341358 def loss_func (x , ** kwargs ):
342359 return x
@@ -516,6 +533,14 @@ def _load_hf_model(args, is_vl_model: bool):
516533 ),
517534 ** _hf_revision_kwargs (args .hf_revision ),
518535 }
536+ if getattr (args , "disable_hf_video_pruning" , False ):
537+ config = AutoConfig .from_pretrained (
538+ args .hf_model_path ,
539+ trust_remote_code = load_kwargs ["trust_remote_code" ],
540+ ** _hf_revision_kwargs (args .hf_revision ),
541+ )
542+ config .video_pruning_rate = 0.0
543+ load_kwargs ["config" ] = config
519544 hf_model = model_class .from_pretrained (args .hf_model_path , ** load_kwargs ).to (args .hf_device ).eval ()
520545 print_rank_0 (f"Loaded with { model_class .__name__ } " )
521546
@@ -561,8 +586,13 @@ def _export_and_load_roundtrip_hf_model(args, is_vl_model: bool, megatron_model,
561586 if _is_rank_0 ():
562587 print_rank_0 ("Loading exported HF model for comparison..." )
563588 model_class = get_model_class (args .model_class , is_vl_model )
589+ config_kwargs = {}
590+ if getattr (args , "disable_hf_video_pruning" , False ):
591+ config = AutoConfig .from_pretrained (save_path , trust_remote_code = True )
592+ config .video_pruning_rate = 0.0
593+ config_kwargs ["config" ] = config
564594 hf_model = (
565- model_class .from_pretrained (save_path , torch_dtype = torch .bfloat16 , trust_remote_code = True )
595+ model_class .from_pretrained (save_path , torch_dtype = torch .bfloat16 , trust_remote_code = True , ** config_kwargs )
566596 .to (args .hf_device )
567597 .eval ()
568598 )
@@ -574,10 +604,10 @@ def _export_and_load_roundtrip_hf_model(args, is_vl_model: bool, megatron_model,
574604 return None
575605
576606
577- def _get_hf_forward_model (hf_model , pixel_values ):
607+ def _get_hf_forward_model (hf_model , pixel_values , pixel_values_videos = None ):
578608 """Select a composite model's language backbone for text-only comparison."""
579609 language_model = getattr (hf_model , "language_model" , None )
580- if pixel_values is None and isinstance (language_model , torch .nn .Module ):
610+ if pixel_values is None and pixel_values_videos is None and isinstance (language_model , torch .nn .Module ):
581611 print_rank_0 ("Using the HuggingFace language backbone for a text-only comparison." )
582612 return language_model
583613 return hf_model
@@ -592,6 +622,7 @@ def _run_hf_inference(
592622 * ,
593623 token_type_ids = None ,
594624 mm_token_type_ids = None ,
625+ pixel_values_videos = None ,
595626):
596627 """Run HuggingFace model inference and return results.
597628
@@ -612,7 +643,7 @@ def _run_hf_inference(
612643 if not _is_rank_0 () or hf_model is None :
613644 return None , None , None , None , None
614645
615- hf_forward_model = _get_hf_forward_model (hf_model , pixel_values )
646+ hf_forward_model = _get_hf_forward_model (hf_model , pixel_values , pixel_values_videos )
616647
617648 input_device = input_ids .device
618649 try :
@@ -628,7 +659,13 @@ def _run_hf_inference(
628659 "attention_mask" : torch .ones_like (input_ids , dtype = torch .bool ).to (hf_device ),
629660 }
630661 if pixel_values is not None :
631- hf_inputs ["pixel_values" ] = pixel_values .to (hf_device )
662+ hf_inputs ["pixel_values" ] = (
663+ [value .to (hf_device ) for value in pixel_values ]
664+ if isinstance (pixel_values , (list , tuple ))
665+ else pixel_values .to (hf_device )
666+ )
667+ if pixel_values_videos is not None :
668+ hf_inputs ["pixel_values_videos" ] = pixel_values_videos .to (hf_device )
632669 if image_grid_thw is not None :
633670 hf_inputs ["image_grid_thw" ] = image_grid_thw .to (hf_device )
634671 if token_type_ids is not None :
@@ -671,12 +708,14 @@ def _run_hf_inference(
671708 )
672709
673710
674- def _load_hf_reference_logits (path , input_ids , tokenizer ):
711+ def _load_hf_reference_logits (path , input_ids , tokenizer , * , expected_metadata = None ):
675712 """Load rank-0 HF logits produced by a memory-bounded reference forward."""
676713 if not _is_rank_0 ():
677714 return None , None , None , None , None
678715
679716 reference = torch .load (path , map_location = "cpu" , weights_only = True )
717+ if expected_metadata is not None and reference .get ("metadata" ) != expected_metadata :
718+ raise ValueError ("HF reference does not match the Omni source and native media inputs" )
680719 reference_input_ids = reference .get ("input_ids" )
681720 if reference_input_ids is None or not torch .equal (reference_input_ids .cpu (), input_ids .cpu ()):
682721 raise ValueError ("HF reference logits were produced from different input token IDs" )
@@ -874,6 +913,26 @@ def compare_models_one_step(args) -> None:
874913 # Detect model type
875914 is_vl_model = is_vision_language_model (args .hf_model_path , args .trust_remote_code , args .hf_revision )
876915 print_rank_0 (f"Detected model type: { 'Vision-Language' if is_vl_model else 'Text-only LLM' } " )
916+ config = AutoConfig .from_pretrained (
917+ args .hf_model_path ,
918+ trust_remote_code = is_safe_repo (trust_remote_code = args .trust_remote_code , hf_path = args .hf_model_path ),
919+ ** _hf_revision_kwargs (args .hf_revision ),
920+ )
921+ is_omni = is_nemotron_omni (config )
922+ video_path = getattr (args , "video_path" , None )
923+ disable_video_pruning = getattr (args , "disable_hf_video_pruning" , False )
924+ if video_path and (not is_omni or args .image_path ):
925+ raise ValueError ("Video comparison currently requires Nemotron Omni and no simultaneous image input." )
926+ if disable_video_pruning and not (is_omni and video_path ):
927+ raise ValueError ("--disable-hf-video-pruning applies only to Omni video comparison." )
928+ video_pruning_rate = None
929+ if video_path :
930+ video_pruning_rate = 0.0 if disable_video_pruning else float (getattr (config , "video_pruning_rate" , 0.0 ))
931+ if video_pruning_rate != 0.0 :
932+ raise ValueError (
933+ "Bridge video is unpruned; explicitly use --disable-hf-video-pruning for a matched oracle."
934+ )
935+ print_rank_0 ("Comparing unpruned Omni video (HF video_pruning_rate=0); not original pruned-HF behavior." )
877936
878937 # Validate vision requirements
879938 if args .image_path and not is_vl_model :
@@ -895,11 +954,37 @@ def compare_models_one_step(args) -> None:
895954 # Setup tokenizer and processor
896955 tokenizer , processor = _setup_tokenizer_and_processor (args , is_vl_model )
897956
957+ # Omni uses the native processor once, then adapts only the media layout for Bridge.
958+ omni_inputs = None
959+ reference_metadata = None
960+ imgs_sizes = num_frames = hf_pixel_values_videos = None
898961 # Process inputs
899962 print_rank_0 (f"Processing inputs - Prompt: '{ args .prompt } ', Image: { args .image_path } " )
900- input_ids , pixel_values , image_grid_thw , token_type_ids , mm_token_type_ids = process_inputs (
901- tokenizer , processor , args .image_path , args .prompt , is_vl_model , args .tp
902- )
963+ if is_omni and (args .image_path or video_path ):
964+ images = frames = metadata = None
965+ if video_path :
966+ frames , metadata = load_nemotron_omni_video (video_path , fps = args .video_fps )
967+ else :
968+ images = [load_image (args .image_path ).convert ("RGB" )]
969+ omni_inputs = prepare_nemotron_omni_inputs (
970+ processor , prompt = args .prompt , images = images , video_frames = frames , video_metadata = metadata
971+ )
972+ reference_metadata = nemotron_omni_reference_metadata (
973+ omni_inputs ,
974+ model = args .hf_model_path ,
975+ revision = args .hf_revision ,
976+ video_pruning_rate = video_pruning_rate ,
977+ )
978+ input_ids = pad_input_ids_to_tp_multiple (omni_inputs .hf ["input_ids" ], args .tp , tokenizer .pad_token_id or 0 )
979+ pixel_values = omni_inputs .bridge ["pixel_values" ]
980+ imgs_sizes = omni_inputs .bridge ["imgs_sizes" ]
981+ num_frames = omni_inputs .bridge ["num_frames" ]
982+ hf_pixel_values_videos = omni_inputs .hf .get ("pixel_values_videos" )
983+ image_grid_thw = token_type_ids = mm_token_type_ids = None
984+ else :
985+ input_ids , pixel_values , image_grid_thw , token_type_ids , mm_token_type_ids = process_inputs (
986+ tokenizer , processor , args .image_path , args .prompt , is_vl_model , args .tp
987+ )
903988
904989 # Move to GPU
905990 input_ids = input_ids .cuda ()
@@ -911,24 +996,28 @@ def compare_models_one_step(args) -> None:
911996 token_type_ids = token_type_ids .cuda ()
912997 if mm_token_type_ids is not None :
913998 mm_token_type_ids = mm_token_type_ids .cuda ()
999+ if imgs_sizes is not None :
1000+ imgs_sizes = imgs_sizes .cuda ()
1001+ num_frames = num_frames .cuda ()
9141002
9151003 print_rank_0 (f"Input shape: { input_ids .shape } " )
9161004 print_rank_0 (f"Pixel values shape: { pixel_values .shape if pixel_values is not None else 'None' } " )
9171005
9181006 # Run HF model forward pass
9191007 if args .hf_logits_path :
9201008 hf_logits , hf_next_token , hf_logits_stats , hf_top5_info , logits_shape = _load_hf_reference_logits (
921- args .hf_logits_path , input_ids , tokenizer
1009+ args .hf_logits_path , input_ids , tokenizer , expected_metadata = reference_metadata
9221010 )
9231011 else :
9241012 hf_logits , hf_next_token , hf_logits_stats , hf_top5_info , logits_shape = _run_hf_inference (
9251013 hf_model ,
9261014 input_ids ,
927- pixel_values ,
1015+ omni_inputs . hf . get ( "pixel_values" ) if omni_inputs is not None else pixel_values ,
9281016 image_grid_thw ,
9291017 tokenizer ,
9301018 token_type_ids = token_type_ids ,
9311019 mm_token_type_ids = mm_token_type_ids ,
1020+ pixel_values_videos = hf_pixel_values_videos ,
9321021 )
9331022
9341023 del hf_model
@@ -975,6 +1064,8 @@ def compare_models_one_step(args) -> None:
9751064 pixel_values ,
9761065 image_grid_thw ,
9771066 mm_token_type_ids = mm_token_type_ids ,
1067+ imgs_sizes = imgs_sizes ,
1068+ num_frames = num_frames ,
9781069 )
9791070 megatron_output = fwd_bwd_function (
9801071 forward_step_func = vlm_forward_step ,
@@ -1095,6 +1186,13 @@ def build_parser() -> argparse.ArgumentParser:
10951186 help = "Path or URL to the image for vision-language generation (optional)." ,
10961187 )
10971188 parser .add_argument ("--megatron_model_path" , type = str , default = None , help = "Path to the Megatron model checkpoint" )
1189+ parser .add_argument ("--video_path" , default = None , help = "Video path for Nemotron Omni native-processor comparison." )
1190+ parser .add_argument ("--video-fps" , type = float , default = 2.0 , help = "Video frame sampling rate." )
1191+ parser .add_argument (
1192+ "--disable-hf-video-pruning" ,
1193+ action = "store_true" ,
1194+ help = "Explicitly configure an unpruned Omni HF video oracle matching Bridge; changes the HF reference only." ,
1195+ )
10981196 parser .add_argument (
10991197 "--hf-logits-path" ,
11001198 default = None ,
0 commit comments