From 0542a771df31a49570d39bf4ccbe66f6276c487c Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 21 Jul 2026 20:57:34 -0700 Subject: [PATCH 01/30] feat(model): enable Nemotron Omni audio and video Signed-off-by: Chen Cui --- .../valor32k_avqa_inference.py | 23 +- src/megatron/bridge/data/builders/energon.py | 2 + .../bridge/data/collators/registry.py | 2 +- .../energon/nemotron_omni_task_encoder.py | 13 +- .../models/nemotron_omni/data/collate_fn.py | 15 ++ .../nemotron_omni/modeling_nemotron_omni.py | 79 +++++- .../data/collators/test_model_collators.py | 33 +++ .../test_nemotron_omni_task_encoder.py | 231 ++++++++++++++++++ .../nemotron_omni/test_nemotron_omni_model.py | 115 ++++++++- .../test_nemotron_omni_recipes.py | 4 +- 10 files changed, 486 insertions(+), 31 deletions(-) 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..17d1954bd6 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,13 @@ from megatron.bridge import AutoBridge from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import ( COMPACT_IMAGE_PLACEHOLDER, - inference_merged_sequence_length, 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, ) @@ -285,6 +285,13 @@ 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 + input_ids = adjust_image_tokens( + input_ids, + torch.full_like(num_image_tiles, image_seq_len), + tokenizer.convert_tokens_to_ids(""), + tokenizer.convert_tokens_to_ids(""), + ) # Process audio sound_clips = None @@ -381,9 +388,6 @@ def generate(model, tokenizer, sample, *, sequence_length, max_new_tokens=50): 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, @@ -424,12 +428,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) diff --git a/src/megatron/bridge/data/builders/energon.py b/src/megatron/bridge/data/builders/energon.py index 708488c4f6..4c38d5f432 100644 --- a/src/megatron/bridge/data/builders/energon.py +++ b/src/megatron/bridge/data/builders/energon.py @@ -100,6 +100,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: @@ -272,6 +273,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/energon/nemotron_omni_task_encoder.py b/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py index 55626ca5aa..9ba2ea5e4f 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 + explicit legacy 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/models/nemotron_omni/data/collate_fn.py b/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py index 0606cec5bd..5b24a8aa07 100644 --- a/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py +++ b/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py @@ -895,6 +895,21 @@ def nemotron_omni_collate_fn( 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) 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..e15b64f232 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 before model-owned packing. """ 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. @@ -595,13 +636,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.") + if has_sound and sound_clips.shape == torch.Size([1, 1]): + has_sound = sound_clips[0, 0].item() != 0 lm_input_ids = input_ids combined_embeddings = None @@ -618,6 +659,11 @@ def forward( else: image_embeddings = None + if has_sound: + 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. @@ -635,6 +681,17 @@ def forward( 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, + 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 diff --git a/tests/unit_tests/data/collators/test_model_collators.py b/tests/unit_tests/data/collators/test_model_collators.py index 791e11994f..315163151e 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 @@ -2471,3 +2478,29 @@ 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..9843e05aa5 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) @@ -436,3 +470,200 @@ def test_energon_canonical_collator_owns_complete_thd_packing(monkeypatch): assert encoded["tokens"] is batch.input_ids assert encoded.get("padding_mask") is None assert get_packed_seq_params(encoded).tokens_per_sample is None + + +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] + + +def test_hf_and_energon_llava_packing_are_identical_for_image_video_audio(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), + ) + 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_batch = omni_collate.nemotron_omni_llava_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=True, + ) + 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] + 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]) 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..9fb753af3e 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 @@ -58,20 +58,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): @@ -271,6 +298,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]]) @@ -523,6 +615,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 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..ea15aad206 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 @@ -146,7 +146,7 @@ def test_cord_v2_sft_recipe_uses_hf_dataset_config(fake_processor): assert isinstance(cfg.dataset, DirectHFSFTDatasetConfig) assert cfg.dataset.hf_processor_path == _TEST_HF_ID 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 From 94cfad7a9b8235f5901b932d0d46d77b0ab4796b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 27 Jul 2026 16:56:14 -0700 Subject: [PATCH 02/30] refactor(model): move Nemotron Omni packing to collator Signed-off-by: Chen Cui --- .../models/nemotron/nemotron_3_omni/README.md | 14 ++-- .../bridge/data/collators/sequence.py | 2 + .../bridge/data/energon/hf_task_encoder.py | 2 + src/megatron/bridge/data/packing/in_batch.py | 18 ++++- .../models/nemotron_omni/data/collate_fn.py | 1 + .../nemotron_omni/modeling_nemotron_omni.py | 28 ++++++-- .../bridge/training/nemotron_omni_step.py | 15 +++- .../test_nemotron_omni_task_encoder.py | 70 +++++++++++-------- .../unit_tests/data/packing/test_in_batch.py | 22 ++++++ ...test_collator_owned_packing_distributed.py | 7 ++ .../nemotron_omni/test_nemotron_omni_model.py | 15 +++- .../training/test_nemotron_omni_step.py | 7 +- 12 files changed, 154 insertions(+), 47 deletions(-) diff --git a/examples/models/nemotron/nemotron_3_omni/README.md b/examples/models/nemotron/nemotron_3_omni/README.md index 005c33bb70..62e80da132 100644 --- a/examples/models/nemotron/nemotron_3_omni/README.md +++ b/examples/models/nemotron/nemotron_3_omni/README.md @@ -192,10 +192,16 @@ 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 so +they do not affect 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/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 232b3e5360..c52f776039 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 @@ -197,6 +198,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/packing/in_batch.py b/src/megatron/bridge/data/packing/in_batch.py index 701b55ec27..000b0acf5a 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,12 @@ def build_mcore_thd_sequence_batch_from_rows( ), "attention_mask": None, } + if emit_padding_mask: + # MoE routing operates on the physical THD stream, including alignment + # gaps between logical rows. Keep those gaps explicit so callers can + # exclude them from router statistics without reconstructing packing + # state inside the model. + 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 +197,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 +238,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 +257,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 +318,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 +326,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", 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 5b24a8aa07..e023c0f283 100644 --- a/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py +++ b/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py @@ -967,6 +967,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. 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 e15b64f232..c95ae7d765 100644 --- a/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py +++ b/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py @@ -115,7 +115,7 @@ class NemotronOmniModel(MegatronModule): selects the rank-local CP shard without changing packed metadata. Image, video, sound, and text inputs use the same one-feature-per-placeholder - contract before model-owned packing. + contract. """ model_owns_packing = False @@ -543,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], @@ -551,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.""" @@ -561,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: @@ -568,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: @@ -586,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 @@ -605,6 +609,7 @@ def _apply_context_parallel_sharding( attention_mask, labels, loss_mask, + padding_mask, loss_mask is not None, ) @@ -615,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, @@ -678,7 +684,7 @@ 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: @@ -689,7 +695,7 @@ def forward( input_ids, sound_embeddings, self.sound_token_index, - attention_mask, + ~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: @@ -705,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, @@ -713,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, ) @@ -736,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 @@ -758,6 +775,7 @@ def forward( inference_params=inference_params, runtime_gather_output=runtime_gather_output, packed_seq_params=language_packed_seq_params, + padding_mask=padding_mask, ) if return_sliced_loss_mask: return output, loss_mask diff --git a/src/megatron/bridge/training/nemotron_omni_step.py b/src/megatron/bridge/training/nemotron_omni_step.py index a2dc831295..bb7d23c505 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,10 @@ 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: + # Every decoder PP stage needs the physical THD gap mask for MoE + # routing. The model applies CP/SP sharding after media insertion. + 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 +181,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 +199,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 +226,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 +263,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 +282,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/data/energon/test_nemotron_omni_task_encoder.py b/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py index 9843e05aa5..9d71598bf6 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 @@ -444,34 +444,6 @@ def _fake_decode(path, *, video_fps, video_nframes): assert sampled_fps == 2.5 -def test_energon_canonical_collator_owns_complete_thd_packing(monkeypatch): - monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens) - encoder = NemotronOmniTaskEncoder( - processor=_Processor([[1, 2, 3], [4, 5]]), - seq_length=8, - enable_in_batch_packing=True, - in_batch_packing_pad_to_multiple_of=4, - pad_to_multiple_of=1, - ) - samples = [ - encoder.encode_sample(_sample([{"role": "user", "content": "text"}], key=f"row-{row_index}")) - for row_index in range(2) - ] - - batch = encoder.batch(samples) - encoded = encoder.encode_batch(batch) - - 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.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 get_packed_seq_params(encoded).tokens_per_sample is None - - def test_energon_llava_multimodal_packing_uses_post_merge_boundaries(monkeypatch): from PIL import Image @@ -542,7 +514,8 @@ def test_energon_llava_multimodal_packing_uses_post_merge_boundaries(monkeypatch assert packed_seq_params.seq_idx[0, 264:].unique().tolist() == [1] -def test_hf_and_energon_llava_packing_are_identical_for_image_video_audio(monkeypatch): +@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) @@ -599,7 +572,12 @@ def test_hf_and_energon_llava_packing_are_identical_for_image_video_audio(monkey "pad_to_multiple_of": 1, } - hf_batch = omni_collate.nemotron_omni_llava_collate_fn(examples, _Processor(rows), **collate_kwargs) + 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, @@ -609,7 +587,7 @@ def test_hf_and_energon_llava_packing_are_identical_for_image_video_audio(monkey temporal_patch_size=2, num_mel_bins=4, pad_to_multiple_of=1, - collapse_image_tokens=True, + collapse_image_tokens=collapse_image_tokens, ) energon_batch = energon_encoder.encode_batch(energon_encoder.batch(normalized_samples)) @@ -636,6 +614,10 @@ def test_hf_and_energon_llava_packing_are_identical_for_image_video_audio(monkey 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, @@ -667,3 +649,29 @@ def test_energon_llava_temporal_video_refuses_unsafe_sequence_truncation(monkeyp 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( + processor=_Processor([[1, 2, 3], [4, 5]]), + seq_length=8, + enable_in_batch_packing=True, + in_batch_packing_pad_to_multiple_of=4, + pad_to_multiple_of=1, + ) + samples = [ + encoder.encode_sample(_sample([{"role": "user", "content": "text"}], key=f"row-{row_index}")) + for row_index in range(2) + ] + + batch = encoder.batch(samples) + encoded = encoder.encode_batch(batch) + + 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 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["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..f777208a34 100644 --- a/tests/unit_tests/data/packing/test_in_batch.py +++ b/tests/unit_tests/data/packing/test_in_batch.py @@ -162,6 +162,28 @@ 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_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/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_model.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py index 9fb753af3e..243f65b28d 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 @@ -462,6 +462,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( @@ -480,6 +481,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), ) @@ -496,6 +498,10 @@ 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 torch.equal( + model.language_model.last_kwargs["padding_mask"], + padding_mask.index_select(1, cp_index), + ) assert model.language_model.last_kwargs["attention_mask"] is None @@ -557,6 +563,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(attention_mask), 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"), @@ -643,7 +650,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, @@ -656,22 +663,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/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( From 183cf51d149c6487e88feaaeb1d307038b24f3aa Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 27 Jul 2026 17:42:53 -0700 Subject: [PATCH 03/30] fix(model): make canonical Nemotron Omni path the default Signed-off-by: Chen Cui --- .../models/nemotron/nemotron_3_omni/README.md | 11 ++ .../nemotron_3_omni/cord_v2_inference.py | 48 +++--- .../hf_to_megatron_generate_nemotron_omni.py | 145 +++++++++++------- .../valor32k_avqa_inference.py | 22 +-- src/megatron/bridge/data/builders/energon.py | 2 + .../energon/nemotron_omni_task_encoder.py | 2 +- .../models/nemotron_omni/data/collate_fn.py | 11 +- .../modeling_nemotron_omni_llava.py | 6 +- .../nemotron_omni/nemotron_omni_bridge.py | 13 +- .../nemotron_omni/nemotron_omni_provider.py | 13 +- .../nemotron_omni/nemotron_omni_utils.py | 82 +++++++++- .../data/builders/test_energon_builder.py | 1 + .../data/collators/test_model_collators.py | 15 +- .../examples/test_nemotron_omni_inference.py | 43 +++++- .../test_nemotron_omni_conversion.py | 3 +- .../nemotron_omni/test_nemotron_omni_model.py | 11 ++ .../nemotron_omni/test_nemotron_omni_utils.py | 66 +++++--- .../test_nemotron_omni_recipes.py | 1 + 18 files changed, 361 insertions(+), 134 deletions(-) diff --git a/examples/models/nemotron/nemotron_3_omni/README.md b/examples/models/nemotron/nemotron_3_omni/README.md index 62e80da132..8a49315f36 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. 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..a2cb26590a 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 @@ -605,11 +638,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 @@ -654,6 +686,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,6 +707,7 @@ 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, ) @@ -688,6 +724,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) @@ -723,7 +760,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 +791,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 +832,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 +866,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 +875,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 +890,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: 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 17d1954bd6..d65c746c27 100644 --- a/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py +++ b/examples/models/nemotron/nemotron_3_omni/valor32k_avqa_inference.py @@ -44,6 +44,7 @@ from megatron.bridge import AutoBridge from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import ( COMPACT_IMAGE_PLACEHOLDER, + inference_expanded_image_token_counts, inference_num_image_tiles, patchify_temporal_frame, select_inference_next_token, @@ -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"] @@ -286,9 +283,14 @@ def process_sample( 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, - torch.full_like(num_image_tiles, image_seq_len), + expanded_counts, tokenizer.convert_tokens_to_ids(""), tokenizer.convert_tokens_to_ids(""), ) @@ -341,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"], @@ -382,7 +383,6 @@ 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) @@ -403,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( @@ -412,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, @@ -496,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: @@ -512,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/src/megatron/bridge/data/builders/energon.py b/src/megatron/bridge/data/builders/energon.py index 4c38d5f432..a5c355e356 100644 --- a/src/megatron/bridge/data/builders/energon.py +++ b/src/megatron/bridge/data/builders/energon.py @@ -89,6 +89,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 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 9ba2ea5e4f..6e4aae65a1 100644 --- a/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py +++ b/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py @@ -128,7 +128,7 @@ class NemotronOmniTaskEncoder(HFTaskEncoder): assistant masking, modality-token expansion, padding, and in-batch packing are performed by the canonical expanded-sequence collator for both Direct-HF and Energon datasets. ``collapse_image_tokens=True`` selects the - explicit legacy LLaVA compatibility contract. + deprecated LLaVA compatibility contract. """ def __init__( 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 e023c0f283..f6c80559af 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 @@ -823,6 +824,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: @@ -999,7 +1008,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_llava.py b/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni_llava.py index 1b4e70eef1..cae398f665 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 @@ -16,11 +16,15 @@ 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 freeze( 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..68ffa89aa2 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py @@ -33,6 +33,7 @@ """ import copy +import warnings from dataclasses import fields from megatron.core.activations import squared_relu @@ -264,9 +265,19 @@ def mapping_registry(self) -> MegatronMappingRegistry: 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..47719827fd 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 @@ -483,7 +484,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 +501,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/tests/unit_tests/data/builders/test_energon_builder.py b/tests/unit_tests/data/builders/test_energon_builder.py index 128c2f911a..0f860ead6d 100644 --- a/tests/unit_tests/data/builders/test_energon_builder.py +++ b/tests/unit_tests/data/builders/test_energon_builder.py @@ -198,6 +198,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 315163151e..7f42eaca24 100644 --- a/tests/unit_tests/data/collators/test_model_collators.py +++ b/tests/unit_tests/data/collators/test_model_collators.py @@ -1929,13 +1929,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() == [ [ diff --git a/tests/unit_tests/examples/test_nemotron_omni_inference.py b/tests/unit_tests/examples/test_nemotron_omni_inference.py index a098ad23cc..d142fc90eb 100644 --- a/tests/unit_tests/examples/test_nemotron_omni_inference.py +++ b/tests/unit_tests/examples/test_nemotron_omni_inference.py @@ -31,7 +31,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 +55,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_nemotron_omni_conversion.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_conversion.py index f79bb5c3a1..77561ac0f1 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 @@ -292,7 +292,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 243f65b28d..ddae06b82e 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 @@ -261,6 +261,15 @@ def test_llava_provider_preserves_existing_radio_cpe_default(): assert provider.radio_interpolate_only_cpe is True +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) @@ -585,6 +594,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( @@ -603,6 +613,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"), 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..213e7653fc 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,6 +16,7 @@ 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, @@ -83,21 +84,47 @@ def test_inference_num_image_tiles_rejects_unshufflable_image_grid(): inference_num_image_tiles(torch.tensor([[528, 512]]), patch_dim=16) +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])) + + def test_inference_merged_sequence_length_uses_exact_image_replacements(): input_ids = torch.tensor([[10, -200, 11, -200, 12]]) - dynamic_length = inference_merged_sequence_length( - input_ids, - image_token_index=-200, - num_image_tiles=torch.tensor([3, 2]), - image_seq_len=1, - ) - 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"): + 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 +141,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/recipes/nemotron_omni/test_nemotron_omni_recipes.py b/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py index ea15aad206..eda5dd0f13 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 @@ -189,6 +189,7 @@ 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 From 42f6d9efdd09c63be1714cf583bef226e06db5b1 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 13:48:55 -0700 Subject: [PATCH 04/30] fix(model): deprecate Nemotron Omni LLaVA model Signed-off-by: Chen Cui --- .../modeling_nemotron_omni_llava.py | 37 +++++++++++++++++++ .../nemotron_omni/test_nemotron_omni_model.py | 9 +++++ 2 files changed, 46 insertions(+) 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 cae398f665..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,7 +12,12 @@ # 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): @@ -27,6 +32,38 @@ class NemotronOmniLlavaModel(NemotronVLModel): 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/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py index ddae06b82e..6d6320b0ca 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): @@ -261,6 +263,13 @@ 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() From 071e98d3fccd1a25acb1c4a2147176f3b6b0ac54 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 13:53:53 -0700 Subject: [PATCH 05/30] test(model): preserve Nemotron Omni inference expansion Signed-off-by: Chen Cui --- .../nemotron_omni/test_nemotron_omni_utils.py | 79 +++++++++++++++++++ 1 file changed, 79 insertions(+) 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 213e7653fc..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 @@ -22,6 +22,7 @@ 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(): @@ -108,6 +109,84 @@ def test_inference_expanded_image_token_counts_rejects_incomplete_tile_ownership inference_expanded_image_token_counts(torch.tensor([256, 128]), torch.tensor([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, + ) + 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]]) From a5502ad7ae32d9e67e232198101538cc5f7eff54 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 14:03:36 -0700 Subject: [PATCH 06/30] docs(model): explain packed padding mask Signed-off-by: Chen Cui --- src/megatron/bridge/data/packing/in_batch.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/data/packing/in_batch.py b/src/megatron/bridge/data/packing/in_batch.py index 000b0acf5a..2828fe74d0 100644 --- a/src/megatron/bridge/data/packing/in_batch.py +++ b/src/megatron/bridge/data/packing/in_batch.py @@ -181,10 +181,8 @@ def build_mcore_thd_sequence_batch_from_rows( "attention_mask": None, } if emit_padding_mask: - # MoE routing operates on the physical THD stream, including alignment - # gaps between logical rows. Keep those gaps explicit so callers can - # exclude them from router statistics without reconstructing packing - # state inside the model. + # 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} From a2b03d2fcbca47f96b7c54de50cd687800897d81 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 14:23:23 -0700 Subject: [PATCH 07/30] feat(model): allow omitting Nemotron Omni sound encoder Signed-off-by: Chen Cui --- .../nemotron_omni/nemotron_omni_provider.py | 52 ++++++++++--------- .../nemotron_omni/h100/nemotron_omni.py | 5 +- .../test_nemotron_omni_conversion.py | 12 +++++ .../test_nemotron_omni_recipes.py | 5 ++ 4 files changed, 48 insertions(+), 26 deletions(-) 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 47719827fd..cbefa858ce 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py @@ -158,6 +158,9 @@ class _NemotronOmniModelProviderBase(NemotronVLModelProvider): dynamic_resolution: Literal[True] = True has_sound: bool = False + # Keep checkpoint capability separate from per-run construction. Disabling + # the encoder also omits its dependent projector. + add_sound_encoder: bool = True sound_model_type: str = "parakeet" sound_hidden_size: int = 1024 sound_projection_hidden_size: int = 4096 @@ -277,6 +280,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 self.add_sound_encoder 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. @@ -301,22 +318,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, @@ -417,16 +424,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, 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..7d26b41731 100644 --- a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py +++ b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py @@ -63,11 +63,13 @@ 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.add_sound_encoder = False cfg.dataset = DirectHFSFTDatasetConfig( seq_length=4096, preprocessing=ChatSFTPreprocessingConfig(), @@ -92,7 +94,7 @@ 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 +102,7 @@ def nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config() -> ConfigContainer: cfg = _nemotron_omni_base() cfg.model.temporal_patch_dim = 1 + cfg.model.add_sound_encoder = False cfg.peft = LoRA( target_modules=["linear_qkv", "linear_proj", "in_proj", "out_proj", "linear_fc1", "linear_fc2"], dim=16, 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 77561ac0f1..28e0290b08 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 @@ -138,6 +138,7 @@ def test_nemotron_omni_provider_bridge_maps_public_config_fields(): assert isinstance(provider, NemotronOmniModelProvider) assert provider.nemotron_omni_contract == NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT assert provider.has_sound is True + assert provider.add_sound_encoder is True assert provider.language_model_type == "nemotron6-moe" assert provider.hidden_size == 256 assert provider.ffn_hidden_size == 512 @@ -162,6 +163,7 @@ 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["add_sound_encoder"] is True def test_nemotron_omni_provider_rejects_static_resolution(): @@ -213,6 +215,16 @@ 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=True, add_sound_encoder=False) + + sound_model, sound_projection = provider._build_sound_modules(None, None, add_encoder=True) + + assert provider.has_sound is True + assert sound_model is None + assert sound_projection is None + + def test_nemotron_omni_vision_projection_uses_squared_relu(): provider = NemotronOmniModelProvider() 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 eda5dd0f13..7f3058e639 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 @@ -46,6 +46,7 @@ class _FakeModelCfg: dynamic_resolution = True + add_sound_encoder = True def finalize(self): return None @@ -150,6 +151,7 @@ def test_cord_v2_sft_recipe_uses_hf_dataset_config(fake_processor): 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.add_sound_encoder is False assert cfg.model.freeze_sound_projection is False assert cfg.peft is None @@ -173,6 +175,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.add_sound_encoder is False assert cfg.model.freeze_sound_projection is True @@ -193,6 +196,7 @@ def test_valor32k_sft_recipe_uses_temporal_omni_task_encoder_config(fake_process 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.add_sound_encoder is True assert cfg.model.freeze_sound_projection is False assert cfg.peft is None @@ -217,4 +221,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.add_sound_encoder is True assert cfg.model.freeze_sound_projection is True From d64f5e7307636804a1bd702ccee258fc7d0f8cba Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 14:33:35 -0700 Subject: [PATCH 08/30] refactor(model): remove Nemotron Omni has_sound flag Signed-off-by: Chen Cui --- .../nemotron_omni/modeling_nemotron_omni.py | 8 +++---- .../nemotron_omni/nemotron_omni_bridge.py | 21 ++++++++---------- .../nemotron_omni/nemotron_omni_provider.py | 14 +++++------- .../test_nemotron_omni_recipes_finetune.py | 5 +++-- .../test_nemotron_omni_conversion.py | 22 ++++++++++++++----- .../nemotron_omni/test_nemotron_omni_model.py | 6 +++-- 6 files changed, 42 insertions(+), 34 deletions(-) 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 c95ae7d765..0db845d4e4 100644 --- a/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py +++ b/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py @@ -646,9 +646,9 @@ def forward( if images is None: images = pixel_values - has_sound = sound_clips is not None and sound_clips.numel() > 0 - if has_sound and sound_clips.shape == torch.Size([1, 1]): - has_sound = sound_clips[0, 0].item() != 0 + 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 @@ -665,7 +665,7 @@ def forward( else: image_embeddings = None - if has_sound: + if has_sound_inputs: sound_embeddings = self._encode_sound(sound_clips, sound_length) else: sound_embeddings = None 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 68ffa89aa2..7dad255849 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py @@ -118,9 +118,7 @@ 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 + Parakeet sound encoder). """ hf_config = hf_pretrained.config llm_config = hf_config.llm_config @@ -133,15 +131,14 @@ 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 - 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 - provider_kwargs["sound_context_token_id"] = hf_config.sound_context_token_id - provider_kwargs["sound_config"] = sc.to_dict() if hasattr(sc, "to_dict") else dict(sc) + sc = getattr(hf_config, "sound_config", None) + if sc is None: + raise ValueError("Nemotron Omni requires sound_config in the Hugging Face checkpoint configuration.") + 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 + provider_kwargs["sound_context_token_id"] = hf_config.sound_context_token_id + provider_kwargs["sound_config"] = sc.to_dict() if hasattr(sc, "to_dict") else dict(sc) provider_kwargs["language_model_type"] = "nemotron6-moe" provider_kwargs["image_token_index"] = getattr(hf_config, "img_context_token_id", 18) 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 cbefa858ce..51e9f65785 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py @@ -157,9 +157,7 @@ class _NemotronOmniModelProviderBase(NemotronVLModelProvider): # accepts a different 4D tensor contract and is not an Omni configuration. dynamic_resolution: Literal[True] = True - has_sound: bool = False - # Keep checkpoint capability separate from per-run construction. Disabling - # the encoder also omits its dependent projector. + # Disabling the sound encoder also omits its dependent projector. add_sound_encoder: bool = True sound_model_type: str = "parakeet" sound_hidden_size: int = 1024 @@ -187,13 +185,13 @@ def _validate_omni_config(self) -> None: "Nemotron Omni requires a positive image_token_index from the checkpoint configuration; " f"got {self.image_token_index}. Construct the provider through AutoBridge or set it explicitly." ) - if self.has_sound and self.sound_context_token_id <= 0: + if self.sound_context_token_id <= 0: raise ValueError( - "Sound-enabled Nemotron Omni requires a positive sound_context_token_id from the checkpoint " + "Nemotron Omni requires a positive sound_context_token_id from the checkpoint " f"configuration; got {self.sound_context_token_id}." ) - if self.has_sound and self.sound_config is None: - raise ValueError("Sound-enabled Nemotron Omni requires sound_config from the checkpoint configuration.") + if self.sound_config is None: + raise ValueError("Nemotron Omni requires sound_config from the checkpoint configuration.") def finalize(self) -> None: """Finalize a dynamic-resolution Nemotron Omni provider.""" @@ -282,7 +280,7 @@ def _build_sound_encoder(self): 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 self.add_sound_encoder and add_encoder): + if not (self.add_sound_encoder and add_encoder): return None, None sound_model = self._build_sound_encoder() diff --git a/tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py b/tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py index e83a02d4d1..e821c3ea30 100644 --- a/tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py +++ b/tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py @@ -16,7 +16,7 @@ import copy import os -from dataclasses import dataclass +from dataclasses import dataclass, field import pytest @@ -36,7 +36,8 @@ class _TinyNemotronOmniModelProvider(NemotronOmniModelProvider): """Small Omni provider used only for functional recipe smoke tests.""" - has_sound: bool = False + add_sound_encoder: bool = False + sound_config: dict = field(default_factory=dict) language_model_type: str = "nemotron6-moe" hidden_size: int = 128 ffn_hidden_size: int = 256 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 28e0290b08..24f148da93 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 @@ -137,7 +137,7 @@ def test_nemotron_omni_provider_bridge_maps_public_config_fields(): assert isinstance(provider, NemotronOmniModelProvider) assert provider.nemotron_omni_contract == NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT - assert provider.has_sound is True + assert not hasattr(provider, "has_sound") assert provider.add_sound_encoder is True assert provider.language_model_type == "nemotron6-moe" assert provider.hidden_size == 256 @@ -163,9 +163,20 @@ 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 "has_sound" not in serialized assert serialized["add_sound_encoder"] is True +def test_nemotron_omni_provider_bridge_requires_sound_config(): + hf_config = _mock_omni_hf_config() + del hf_config.sound_config + hf_pretrained = Mock(spec=PreTrainedCausalLM) + hf_pretrained.config = hf_config + + with pytest.raises(ValueError, match="requires sound_config"): + NemotronOmniBridge().provider_bridge(hf_pretrained) + + def test_nemotron_omni_provider_rejects_static_resolution(): provider = NemotronOmniModelProvider() provider.dynamic_resolution = False @@ -183,14 +194,14 @@ 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, sound_context_token_id=0, sound_config={}) with pytest.raises(ValueError, match="requires a positive sound_context_token_id"): provider.finalize() -def test_nemotron_omni_provider_requires_sound_config_when_enabled(): - provider = NemotronOmniModelProvider(image_token_index=18, has_sound=True, sound_context_token_id=27) +def test_nemotron_omni_provider_requires_sound_config(): + provider = NemotronOmniModelProvider(image_token_index=18, sound_context_token_id=27) with pytest.raises(ValueError, match="requires sound_config"): provider.finalize() @@ -216,11 +227,10 @@ def test_canonical_provider_builds_dedicated_model(monkeypatch): def test_nemotron_omni_provider_can_omit_sound_modules(): - provider = NemotronOmniModelProvider(has_sound=True, add_sound_encoder=False) + provider = NemotronOmniModelProvider(add_sound_encoder=False) sound_model, sound_projection = provider._build_sound_modules(None, None, add_encoder=True) - assert provider.has_sound is True assert sound_model is None assert sound_projection is None 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 6d6320b0ca..1448b60b1d 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 @@ -16,7 +16,7 @@ import datetime import os import socket -from dataclasses import dataclass +from dataclasses import dataclass, field from types import SimpleNamespace import pytest @@ -106,7 +106,9 @@ def __init__(self): class _TinyOmniProvider(NemotronOmniModelProvider): """One-layer image model for the real RADIO/NemotronH Stage 1 smoke.""" - has_sound: bool = False + add_sound_encoder: bool = False + sound_context_token_id: int = 19 + sound_config: dict = field(default_factory=dict) language_model_type: str = "nemotron6-moe" hidden_size: int = 128 ffn_hidden_size: int = 256 From 26515b409f4984b87e737ad678cd8062f7eeedd1 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 14:51:22 -0700 Subject: [PATCH 09/30] fix(model): unify Nemotron Omni sound capability Signed-off-by: Chen Cui --- .../nemotron_omni/nemotron_omni_bridge.py | 31 +++++--- .../nemotron_omni/nemotron_omni_provider.py | 15 ++-- .../nemotron_omni/h100/nemotron_omni.py | 4 +- .../test_nemotron_omni_recipes_finetune.py | 5 +- .../test_nemotron_omni_conversion.py | 71 ++++++++++++++++--- .../nemotron_omni/test_nemotron_omni_model.py | 6 +- .../test_nemotron_omni_recipes.py | 10 +-- 7 files changed, 102 insertions(+), 40 deletions(-) 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 7dad255849..662e8fffea 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py @@ -118,7 +118,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 + Parakeet sound encoder). + 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,13 +133,13 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> NemotronOmniMode provider_kwargs["vision_proj_ffn_hidden_size"] = hf_config.projector_hidden_size sc = getattr(hf_config, "sound_config", None) - if sc is None: - raise ValueError("Nemotron Omni requires sound_config in the Hugging Face checkpoint configuration.") - 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 - provider_kwargs["sound_context_token_id"] = hf_config.sound_context_token_id - provider_kwargs["sound_config"] = sc.to_dict() if hasattr(sc, "to_dict") else dict(sc) + 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 + provider_kwargs["sound_context_token_id"] = hf_config.sound_context_token_id + provider_kwargs["sound_config"] = sc.to_dict() if hasattr(sc, "to_dict") else dict(sc) provider_kwargs["language_model_type"] = "nemotron6-moe" provider_kwargs["image_token_index"] = getattr(hf_config, "img_context_token_id", 18) @@ -171,6 +172,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 # ------------------------------------------------------------------ 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 51e9f65785..86b80f66ea 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py @@ -157,8 +157,9 @@ class _NemotronOmniModelProviderBase(NemotronVLModelProvider): # accepts a different 4D tensor contract and is not an Omni configuration. dynamic_resolution: Literal[True] = True - # Disabling the sound encoder also omits its dependent projector. - add_sound_encoder: bool = 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 sound_projection_hidden_size: int = 4096 @@ -185,13 +186,13 @@ def _validate_omni_config(self) -> None: "Nemotron Omni requires a positive image_token_index from the checkpoint configuration; " f"got {self.image_token_index}. Construct the provider through AutoBridge or set it explicitly." ) - if self.sound_context_token_id <= 0: + if self.has_sound and self.sound_context_token_id <= 0: raise ValueError( - "Nemotron Omni requires a positive sound_context_token_id from the checkpoint " + "Sound-enabled Nemotron Omni requires a positive sound_context_token_id from the checkpoint " f"configuration; got {self.sound_context_token_id}." ) - if self.sound_config is None: - raise ValueError("Nemotron Omni requires sound_config from the checkpoint configuration.") + if self.has_sound and self.sound_config is None: + raise ValueError("Sound-enabled Nemotron Omni requires sound_config from the checkpoint configuration.") def finalize(self) -> None: """Finalize a dynamic-resolution Nemotron Omni provider.""" @@ -280,7 +281,7 @@ def _build_sound_encoder(self): def _build_sound_modules(self, language_cfg, language_spec, *, add_encoder: bool): """Build optional sound modules on the encoder pipeline stage.""" - if not (self.add_sound_encoder and add_encoder): + if not (self.has_sound and add_encoder): return None, None sound_model = self._build_sound_encoder() 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 7d26b41731..346631d5fb 100644 --- a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py +++ b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py @@ -69,7 +69,7 @@ def nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config() -> ConfigContainer: """ cfg = _nemotron_omni_base() cfg.model.temporal_patch_dim = 1 - cfg.model.add_sound_encoder = False + cfg.model.has_sound = False cfg.dataset = DirectHFSFTDatasetConfig( seq_length=4096, preprocessing=ChatSFTPreprocessingConfig(), @@ -102,7 +102,7 @@ def nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config() -> ConfigContainer: cfg = _nemotron_omni_base() cfg.model.temporal_patch_dim = 1 - cfg.model.add_sound_encoder = False + cfg.model.has_sound = False cfg.peft = LoRA( target_modules=["linear_qkv", "linear_proj", "in_proj", "out_proj", "linear_fc1", "linear_fc2"], dim=16, diff --git a/tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py b/tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py index e821c3ea30..e83a02d4d1 100644 --- a/tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py +++ b/tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py @@ -16,7 +16,7 @@ import copy import os -from dataclasses import dataclass, field +from dataclasses import dataclass import pytest @@ -36,8 +36,7 @@ class _TinyNemotronOmniModelProvider(NemotronOmniModelProvider): """Small Omni provider used only for functional recipe smoke tests.""" - add_sound_encoder: bool = False - sound_config: dict = field(default_factory=dict) + has_sound: bool = False language_model_type: str = "nemotron6-moe" hidden_size: int = 128 ffn_hidden_size: int = 256 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 24f148da93..0ea107c06e 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 @@ -137,8 +137,7 @@ def test_nemotron_omni_provider_bridge_maps_public_config_fields(): assert isinstance(provider, NemotronOmniModelProvider) assert provider.nemotron_omni_contract == NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT - assert not hasattr(provider, "has_sound") - assert provider.add_sound_encoder is True + assert provider.has_sound is True assert provider.language_model_type == "nemotron6-moe" assert provider.hidden_size == 256 assert provider.ffn_hidden_size == 512 @@ -163,18 +162,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 "has_sound" not in serialized - assert serialized["add_sound_encoder"] is True + assert serialized["has_sound"] is True + assert "add_sound_encoder" not in serialized -def test_nemotron_omni_provider_bridge_requires_sound_config(): +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 - with pytest.raises(ValueError, match="requires sound_config"): - NemotronOmniBridge().provider_bridge(hf_pretrained) + 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(): @@ -194,14 +222,19 @@ 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, sound_context_token_id=0, sound_config={}) + 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() -def test_nemotron_omni_provider_requires_sound_config(): - provider = NemotronOmniModelProvider(image_token_index=18, sound_context_token_id=27) +def test_nemotron_omni_provider_requires_sound_config_when_enabled(): + provider = NemotronOmniModelProvider(image_token_index=18, has_sound=True, sound_context_token_id=27) with pytest.raises(ValueError, match="requires sound_config"): provider.finalize() @@ -227,14 +260,30 @@ def test_canonical_provider_builds_dedicated_model(monkeypatch): def test_nemotron_omni_provider_can_omit_sound_modules(): - provider = NemotronOmniModelProvider(add_sound_encoder=False) + 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() 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 1448b60b1d..6d6320b0ca 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 @@ -16,7 +16,7 @@ import datetime import os import socket -from dataclasses import dataclass, field +from dataclasses import dataclass from types import SimpleNamespace import pytest @@ -106,9 +106,7 @@ def __init__(self): class _TinyOmniProvider(NemotronOmniModelProvider): """One-layer image model for the real RADIO/NemotronH Stage 1 smoke.""" - add_sound_encoder: bool = False - sound_context_token_id: int = 19 - sound_config: dict = field(default_factory=dict) + has_sound: bool = False language_model_type: str = "nemotron6-moe" hidden_size: int = 128 ffn_hidden_size: int = 256 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 7f3058e639..4ae0e70352 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 @@ -46,7 +46,7 @@ class _FakeModelCfg: dynamic_resolution = True - add_sound_encoder = True + has_sound = True def finalize(self): return None @@ -151,7 +151,7 @@ def test_cord_v2_sft_recipe_uses_hf_dataset_config(fake_processor): 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.add_sound_encoder is False + assert cfg.model.has_sound is False assert cfg.model.freeze_sound_projection is False assert cfg.peft is None @@ -175,7 +175,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.add_sound_encoder is False + assert cfg.model.has_sound is False assert cfg.model.freeze_sound_projection is True @@ -196,7 +196,7 @@ def test_valor32k_sft_recipe_uses_temporal_omni_task_encoder_config(fake_process 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.add_sound_encoder is True + assert cfg.model.has_sound is True assert cfg.model.freeze_sound_projection is False assert cfg.peft is None @@ -221,5 +221,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.add_sound_encoder is True + assert cfg.model.has_sound is True assert cfg.model.freeze_sound_projection is True From 8f3942b82359803f7c00ea638a659f77d2c8f8b3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 18:51:58 -0700 Subject: [PATCH 10/30] fix(model): defer Nemotron Omni padding mask routing Signed-off-by: Chen Cui --- .../bridge/models/nemotron_omni/modeling_nemotron_omni.py | 5 ++++- .../models/nemotron_omni/test_nemotron_omni_model.py | 5 +---- 2 files changed, 5 insertions(+), 5 deletions(-) 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 0db845d4e4..a8e47f58ae 100644 --- a/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py +++ b/src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py @@ -764,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, @@ -775,7 +779,6 @@ def forward( inference_params=inference_params, runtime_gather_output=runtime_gather_output, packed_seq_params=language_packed_seq_params, - padding_mask=padding_mask, ) if return_sliced_loss_mask: return output, loss_mask 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 6d6320b0ca..78c7aa081b 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 @@ -516,10 +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 torch.equal( - model.language_model.last_kwargs["padding_mask"], - padding_mask.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 From 91c02ea162ef9ee47778753ff1d7eb7f5630a82b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 28 Jul 2026 19:04:00 -0700 Subject: [PATCH 11/30] docs(model): clarify Nemotron Omni padding routing Signed-off-by: Chen Cui --- examples/models/nemotron/nemotron_3_omni/README.md | 7 +++++-- src/megatron/bridge/training/nemotron_omni_step.py | 5 +++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/examples/models/nemotron/nemotron_3_omni/README.md b/examples/models/nemotron/nemotron_3_omni/README.md index 8a49315f36..a66d6b9e6e 100644 --- a/examples/models/nemotron/nemotron_3_omni/README.md +++ b/examples/models/nemotron/nemotron_3_omni/README.md @@ -206,8 +206,11 @@ pretrained checkpoint and enable in-batch sequence packing via 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 so -they do not affect MoE routing statistics. +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 diff --git a/src/megatron/bridge/training/nemotron_omni_step.py b/src/megatron/bridge/training/nemotron_omni_step.py index bb7d23c505..a215577cec 100644 --- a/src/megatron/bridge/training/nemotron_omni_step.py +++ b/src/megatron/bridge/training/nemotron_omni_step.py @@ -84,8 +84,9 @@ def get_batch_from_iterator( 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: - # Every decoder PP stage needs the physical THD gap mask for MoE - # routing. The model applies CP/SP sharding after media insertion. + # 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: From a3c8d076f3fd6992715c2c9e30f21e33ef986bba Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Wed, 29 Jul 2026 14:49:15 -0700 Subject: [PATCH 12/30] test(model): repair Nemotron Omni packed image smoke Signed-off-by: Chen Cui --- tests/unit_tests/data/collators/test_model_collators.py | 2 ++ .../unit_tests/data/energon/test_nemotron_omni_task_encoder.py | 2 ++ .../unit_tests/models/nemotron_omni/test_nemotron_omni_model.py | 2 +- 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/data/collators/test_model_collators.py b/tests/unit_tests/data/collators/test_model_collators.py index 7f42eaca24..8418e61a59 100644 --- a/tests/unit_tests/data/collators/test_model_collators.py +++ b/tests/unit_tests/data/collators/test_model_collators.py @@ -2479,6 +2479,8 @@ 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]]) 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 9d71598bf6..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 @@ -649,6 +649,8 @@ def test_energon_llava_temporal_video_refuses_unsafe_sequence_truncation(monkeyp 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( 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 78c7aa081b..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 @@ -578,7 +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(attention_mask), + 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"), From 7888d7aba2d04220d6bfc35fad7db4209e61d6ac Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Wed, 29 Jul 2026 14:50:06 -0700 Subject: [PATCH 13/30] docs(model): add Nemotron Omni verification card Signed-off-by: Chen Cui --- .../card.yaml | 186 ++++++++++++++++++ 1 file changed, 186 insertions(+) create mode 100644 examples/model_verification_cards/nemotron-3-nano-omni-30b-a3b-reasoning/card.yaml 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..3184655b6d --- /dev/null +++ b/examples/model_verification_cards/nemotron-3-nano-omni-30b-a3b-reasoning/card.yaml @@ -0,0 +1,186 @@ +# 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. On the recorded Bridge commit, 144 focused tests passed, including + canonical collator-owned packing, image/video/audio insertion, and a real + tiny-model forward, backward, and optimizer update. Two-rank + context-parallel sharding passed on both ranks. The official ten-frame + Valor32K audio/video sample produced 1,280 video and 126 audio placeholders + with finite loss 11.72488499; a three-step legacy/canonical comparison had + maximum loss delta 0.0. These are focused implementation checks rather than + completed end-to-end model verification items, so the required inventory + remains unverified. +verification_index: + model_level: + unverified: + - hf_to_megatron_cpu + - hf_to_megatron_gpu + - megatron_to_hf_cpu + - megatron_to_hf_gpu + - manual_forward_pass + - inference + training: + H100: + unverified: [pretrain, sft, sft_export_inference, sft_long_context, peft, 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: 7dff2df3e773adf346cd6018596909760679efcd # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + No full CPU import is part of this focused stack-refresh validation. + + hf_to_megatron_gpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + No full distributed GPU import is part of this focused stack-refresh + validation. + + megatron_to_hf_cpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + No full CPU export is part of this focused stack-refresh validation. + + megatron_to_hf_gpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + No full distributed GPU export is part of this focused stack-refresh + validation. + + manual_forward_pass: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + No full Hugging Face and Megatron logit-correlation run is part of this + focused stack-refresh validation. + + inference: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + No full-checkpoint deterministic generation run is part of this focused + stack-refresh validation. + + pretrain: + H100: + status: unverified + precision: bf16 + 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: > + No bounded full-model pretraining run is part of this focused + stack-refresh validation. + + sft: + H100: + status: unverified + precision: bf16 + 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: > + Focused tiny-model optimizer coverage does not satisfy the full-model + SFT verification gate. + + sft_export_inference: + H100: + status: unverified + precision: bf16 + depends_on: sft + commands: null + last_verified: null + expected_result: > + No post-SFT export and deterministic Hugging Face generation run is + part of this focused stack-refresh validation. + + sft_long_context: + H100: + status: unverified + precision: bf16 + 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: > + Focused packed-sequence and context-parallel tests do not satisfy the + full-model long-context SFT verification gate. + + peft: + H100: + status: unverified + precision: bf16 + 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: > + No full-model parameter-efficient finetuning run is part of this + focused stack-refresh validation. + + checkpoint_resume: + H100: + status: unverified + precision: bf16 + 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: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: false + expected_result: > + No full-state checkpoint continuation is part of this focused + stack-refresh validation. From 4e1618f7fe0595e5ba6fe4dad6a1f242bc4e8534 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 17:25:55 -0700 Subject: [PATCH 14/30] fix(examples): support large HF parity checks Signed-off-by: Chen Cui --- .../compare_hf_and_megatron/compare.py | 72 ++++++++++++++++--- .../unit_tests/test_compare_mask_handling.py | 36 ++++++++++ 2 files changed, 97 insertions(+), 11 deletions(-) diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 610e14b5b4..5bd06d4b54 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -492,6 +492,22 @@ def process_inputs(tokenizer, processor, image_path: Optional[str], prompt: str, return input_ids, None, None, None +def _is_transformers_none_tp_plan_error(error: TypeError) -> bool: + """Return whether Transformers failed allocator warmup on an unset TP plan.""" + if str(error) != "object of type 'NoneType' has no len()": + return False + + traceback = error.__traceback__ + while traceback is not None: + frame = traceback.tb_frame + if frame.f_code.co_name == "get_total_byte_count" and frame.f_globals.get("__name__") == ( + "transformers.modeling_utils" + ): + return True + traceback = traceback.tb_next + return False + + def _load_hf_model(args, is_vl_model: bool): """Load HuggingFace model on rank 0. @@ -507,16 +523,31 @@ 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), - ) + } + try: + hf_model = model_class.from_pretrained( + args.hf_model_path, + device_map=args.hf_device, + **load_kwargs, + ) + except TypeError as error: + if not _is_transformers_none_tp_plan_error(error): + raise + + print_rank_0( + "HuggingFace model has an unset tensor-parallel plan; retrying through CPU before moving it to CUDA." + ) + gc.collect() + torch.cuda.empty_cache() + hf_model = model_class.from_pretrained(args.hf_model_path, **load_kwargs).to(args.hf_device) + hf_model = hf_model.eval() print_rank_0(f"Loaded with {model_class.__name__}") @@ -591,15 +622,23 @@ 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 + input_device = input_ids.device + try: + hf_device = next(hf_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) hf_output = hf_model(**hf_inputs) @@ -627,7 +666,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_megatron_model(args): @@ -990,6 +1035,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/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index 8d1f80dd93..613f49adda 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -262,3 +262,39 @@ 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_retries_unset_transformers_tp_plan_through_cpu(self): + """Retry the Transformers allocator bug without weakening other TypeErrors.""" + args = compare.build_parser().parse_args( + [ + "--hf_model_path", + "org/model", + "--prompt", + "Hello", + ] + ) + loaded_model = MagicMock() + tp_plan_error = TypeError("object of type 'NoneType' has no len()") + + model_class = MagicMock() + model_class.__name__ = "MockModel" + model_class.from_pretrained.side_effect = [tp_plan_error, 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_transformers_none_tp_plan_error", return_value=True), + patch.object(compare, "is_safe_repo", return_value=True), + patch.object(compare, "print_rank_0"), + patch.object(compare.gc, "collect"), + patch.object(compare.torch.cuda, "empty_cache"), + ): + result = compare._load_hf_model(args, is_vl_model=False) + + assert result is loaded_model + assert model_class.from_pretrained.call_count == 2 + assert model_class.from_pretrained.call_args_list[0].kwargs["device_map"] == "cuda" + assert "device_map" not in model_class.from_pretrained.call_args_list[1].kwargs + loaded_model.to.assert_called_once_with("cuda") From d10a1d8e7b63cc4c7ca1668ac9dec1aca1ee3284 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 17:25:55 -0700 Subject: [PATCH 15/30] fix(conversion): preserve Nemotron Omni HF buffers Signed-off-by: Chen Cui --- .../nemotron_omni/nemotron_omni_bridge.py | 46 +++++++++++++++++-- .../test_nemotron_omni_conversion.py | 29 +++++++++++- 2 files changed, 70 insertions(+), 5 deletions(-) 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 662e8fffea..a7b070e5b5 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py @@ -34,12 +34,14 @@ 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, @@ -83,6 +85,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"), @@ -230,8 +239,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.**", @@ -275,6 +284,37 @@ 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: + return + 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): """Deprecated fallback bridge for the historical collapse/expand model. 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 0ea107c06e..d8159928e5 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,7 +13,7 @@ # limitations under the License. from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import MagicMock, Mock, patch import pytest import torch @@ -22,7 +22,7 @@ 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 +38,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 @@ -310,6 +311,30 @@ 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_canonical_bridge_maps_super_mtp_config(): hf_config = _mock_omni_hf_config() hf_config.architectures = ["NemotronH_Super_Omni_Reasoning_V3"] From 0f7bb5e91fd13cffd2814575959acc1eaf0762ea Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 17:25:55 -0700 Subject: [PATCH 16/30] feat(recipes): add Nemotron Omni long-context SFT Signed-off-by: Chen Cui --- .../recipes/nemotron_omni/h100/__init__.py | 1 + .../nemotron_omni/h100/nemotron_omni.py | 19 +++++++++++++++++++ .../test_nemotron_omni_recipes.py | 17 +++++++++++++++++ 3 files changed, 37 insertions(+) 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 346631d5fb..665dd2c25c 100644 --- a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py +++ b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py @@ -90,6 +90,24 @@ 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. + """ + 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.dataset.seq_length = 8192 + cfg.dataset.enable_in_batch_packing = True + cfg.dataset.in_batch_packing_pad_to_multiple_of = 8 + return cfg + + def nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config() -> ConfigContainer: """Return a LoRA PEFT config for Nemotron Omni on CORD v2. @@ -281,6 +299,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/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py b/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py index 4ae0e70352..e1137a87b1 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 @@ -156,6 +156,23 @@ def test_cord_v2_sft_recipe_uses_hf_dataset_config(fake_processor): 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.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.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) From 6adfd7b0627a84e355efc42db47d9e84d223fc38 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 17:31:10 -0700 Subject: [PATCH 17/30] fix(examples): compare composite text backbones Signed-off-by: Chen Cui --- .../conversion/compare_hf_and_megatron/compare.py | 15 +++++++++++++-- tests/unit_tests/test_compare_mask_handling.py | 10 ++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 5bd06d4b54..e49a7f2e04 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -604,6 +604,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): """Run HuggingFace model inference and return results. @@ -622,9 +631,11 @@ 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_model.parameters()).device + 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)): @@ -640,7 +651,7 @@ def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokeniz if image_grid_thw is not None: hf_inputs["image_grid_thw"] = image_grid_thw.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)}") diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index 613f49adda..398fbadde2 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -210,6 +210,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_broadcast_uses_model_output_vocab_size(self): """Test that non-rank-0 buffers use the HF logits size instead of tokenizer vocab size.""" broadcast_shapes = [] From ffb09059595fb43e0e9d6884f876da6c445591dc Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 17:37:44 -0700 Subject: [PATCH 18/30] fix(examples): pin Nemotron Omni inference revision Signed-off-by: Chen Cui --- .../hf_to_megatron_generate_nemotron_omni.py | 35 ++++++++++++++++--- .../examples/test_nemotron_omni_inference.py | 9 +++++ 2 files changed, 40 insertions(+), 4 deletions(-) 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 a2cb26590a..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 @@ -616,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. @@ -674,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) @@ -714,7 +723,11 @@ def main(args) -> None: 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 @@ -746,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("") @@ -941,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/tests/unit_tests/examples/test_nemotron_omni_inference.py b/tests/unit_tests/examples/test_nemotron_omni_inference.py index d142fc90eb..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", From 960e41687fde2a8bcd7e1e297383c8783defc420 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 17:54:27 -0700 Subject: [PATCH 19/30] fix(recipes): trust Nemotron Omni processor code Signed-off-by: Chen Cui --- .../bridge/recipes/nemotron_omni/h100/nemotron_omni.py | 2 ++ .../recipes/nemotron_omni/test_nemotron_omni_recipes.py | 3 +++ 2 files changed, 5 insertions(+) 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 665dd2c25c..d3a31e1059 100644 --- a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py +++ b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py @@ -74,6 +74,7 @@ def nemotron_omni_cord_v2_sft_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", @@ -146,6 +147,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", 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 e1137a87b1..40166140f0 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 @@ -146,6 +146,7 @@ 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_expanded_collate_fn assert cfg.dataset.enable_in_batch_packing is False @@ -163,6 +164,7 @@ def test_cord_v2_long_context_sft_recipe_enables_packing_and_cp(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 @@ -178,6 +180,7 @@ def test_cord_v2_peft_recipe_configures_lora_and_freezing(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 == [ From c9853c04b94d5ceb18f69ad25e05800aaca019e8 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 18:05:23 -0700 Subject: [PATCH 20/30] fix(conversion): initialize CPU models with Gloo Signed-off-by: Chen Cui --- src/megatron/bridge/models/model_provider.py | 6 +++-- .../models/test_model_provider_mixin.py | 26 +++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/model_provider.py b/src/megatron/bridge/models/model_provider.py index 1de2dd49f5..2edf28f745 100644 --- a/src/megatron/bridge/models/model_provider.py +++ b/src/megatron/bridge/models/model_provider.py @@ -252,8 +252,10 @@ def provide_distributed_model( os.environ["WORLD_SIZE"] = os.environ.get("WORLD_SIZE", "1") os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "localhost") os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "12355") - torch.cuda.set_device(get_local_rank_preinit()) - torch.distributed.init_process_group("nccl") + backend = "gloo" if use_cpu_initialization else "nccl" + if backend == "nccl": + torch.cuda.set_device(get_local_rank_preinit()) + torch.distributed.init_process_group(backend) # If pg_collection is provided (e.g., from use_decentralized_pg=True), # use it directly. Otherwise, initialize model parallel state and get pg_collection from MPU. diff --git a/tests/unit_tests/models/test_model_provider_mixin.py b/tests/unit_tests/models/test_model_provider_mixin.py index 88a90c6d78..bfbe3a6c52 100644 --- a/tests/unit_tests/models/test_model_provider_mixin.py +++ b/tests/unit_tests/models/test_model_provider_mixin.py @@ -228,6 +228,32 @@ def __init__(self): assert provider._pg_collection is pg_instance +@patch("megatron.bridge.models.model_provider.ProcessGroupCollection.use_mpu_process_groups") +@patch("megatron.bridge.models.model_provider.get_model") +@patch("megatron.bridge.models.model_provider.torch.cuda.set_device") +@patch("megatron.bridge.models.model_provider.torch.distributed") +@patch("megatron.bridge.models.model_provider.parallel_state.is_initialized", return_value=True) +def test_cpu_initialization_starts_gloo_without_selecting_cuda( + mock_ps_init, + mock_dist, + mock_set_device, + mock_get_model, + mock_use_pg, + provider, +): + """Standalone CPU model construction must not require a CUDA driver.""" + mock_dist.is_initialized.return_value = False + mock_model = [MockMegatronModule()] + mock_get_model.return_value = mock_model + mock_use_pg.return_value = Mock() + + result = provider.provide_distributed_model(wrap_with_ddp=False, use_cpu_initialization=True) + + assert result is mock_model + mock_dist.init_process_group.assert_called_once_with("gloo") + mock_set_device.assert_not_called() + + def test_hook_registration_and_composition(provider): """Test hook registration order and composition.""" # Initially, no hooks are registered From 23cc39f8634c26bde8a54453517b0337ceea25e6 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 18:24:53 -0700 Subject: [PATCH 21/30] fix(data): ignore padding in Omni loss masks Signed-off-by: Chen Cui --- .../models/nemotron_omni/data/collate_fn.py | 48 ++++++++++++++----- .../data/collators/test_model_collators.py | 31 ++++++++++++ 2 files changed, 68 insertions(+), 11 deletions(-) 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 f6c80559af..53b3ad8119 100644 --- a/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py +++ b/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py @@ -28,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, @@ -50,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: @@ -888,17 +918,13 @@ 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) diff --git a/tests/unit_tests/data/collators/test_model_collators.py b/tests/unit_tests/data/collators/test_model_collators.py index 8418e61a59..425c127cab 100644 --- a/tests/unit_tests/data/collators/test_model_collators.py +++ b/tests/unit_tests/data/collators/test_model_collators.py @@ -1914,6 +1914,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]]) From fbbafc7ddfa818ef91d5eb64c8a00dcaa7bee78d Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 19:58:50 -0700 Subject: [PATCH 22/30] fix(nemotron-omni): complete model verification Signed-off-by: Chen Cui --- scripts/conversion/utils.py | 3 +++ .../bridge/models/conversion/auto_bridge.py | 1 + .../nemotron_omni/nemotron_omni_bridge.py | 7 ++++++- .../nemotron_omni/h100/nemotron_omni.py | 14 ++++++++++++- .../conversion/launcher/test_utils.py | 10 +++++++++ .../test_nemotron_omni_conversion.py | 21 +++++++++++++++++++ tests/unit_tests/models/test_auto_bridge.py | 5 ++++- .../test_nemotron_omni_recipes.py | 8 +++++++ 8 files changed, 66 insertions(+), 3 deletions(-) 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/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index 7f69c61a1e..819d24913e 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -404,6 +404,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 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 a7b070e5b5..27ea813930 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py @@ -47,6 +47,7 @@ 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, @@ -309,7 +310,11 @@ def stream_weights_megatron_to_hf( state = getattr(hf_pretrained, "state", None) source = getattr(state, "source", None) if source is None: - return + 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: 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 d3a31e1059..308dbd322d 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" @@ -96,13 +97,24 @@ def nemotron_omni_cord_v2_long_context_sft_8gpu_h100_bf16_config() -> ConfigCont 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. + 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 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/models/nemotron_omni/test_nemotron_omni_conversion.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_conversion.py index d8159928e5..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 @@ -18,7 +18,9 @@ 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 @@ -335,6 +337,25 @@ def test_nemotron_omni_export_preserves_source_only_buffers(): 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"] diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index 90807e9a78..a0b64f0c0e 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -810,7 +810,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 @@ -818,6 +820,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.""" 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 40166140f0..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 @@ -170,6 +170,14 @@ def test_cord_v2_long_context_sft_recipe_enables_packing_and_cp(fake_processor): 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 From 5aa403472cc7b9e0b369e23db928fc3ecbf8206b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 30 Jul 2026 22:14:54 -0700 Subject: [PATCH 23/30] docs(model): verify Nemotron Omni workflows Signed-off-by: Chen Cui --- .../card.yaml | 297 +++++++++++++----- 1 file changed, 220 insertions(+), 77 deletions(-) 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 index 3184655b6d..815f651451 100644 --- 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 @@ -5,27 +5,31 @@ 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. On the recorded Bridge commit, 144 focused tests passed, including - canonical collator-owned packing, image/video/audio insertion, and a real - tiny-model forward, backward, and optimizer update. Two-rank - context-parallel sharding passed on both ranks. The official ten-frame - Valor32K audio/video sample produced 1,280 video and 126 audio placeholders - with finite loss 11.72488499; a three-step legacy/canonical comparison had - maximum loss delta 0.0. These are focused implementation checks rather than - completed end-to-end model verification items, so the required inventory - remains unverified. + 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: - unverified: + verified: - hf_to_megatron_cpu - hf_to_megatron_gpu + - inference + unverified: - megatron_to_hf_cpu - megatron_to_hf_gpu - manual_forward_pass - - inference training: H100: - unverified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] + 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 @@ -33,65 +37,131 @@ model: min_transformers_version: "5.8.0" verification_environment: base_container: nvcr.io/nvidia/nemo:26.06 - bridge_commit: 7dff2df3e773adf346cd6018596909760679efcd # pragma: allowlist secret + bridge_commit: fbbafc7ddfa818ef91d5eb64c8a00dcaa7bee78d # pragma: allowlist secret items: hf_to_megatron_cpu: - status: unverified + status: verified precision: bf16 - command: null - last_verified: null + 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: > - No full CPU import is part of this focused stack-refresh validation. + 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: unverified + status: verified precision: bf16 - command: null - last_verified: null + 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: > - No full distributed GPU import is part of this focused stack-refresh - validation. + 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: null + 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: > - No full CPU export is part of this focused stack-refresh validation. + 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: null + 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: > - No full distributed GPU export is part of this focused stack-refresh - validation. + 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: null + 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: > - No full Hugging Face and Megatron logit-correlation run is part of this - focused stack-refresh validation. + 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: unverified + status: verified precision: bf16 - command: null - last_verified: null - expected_result: > - No full-checkpoint deterministic generation run is part of this focused - stack-refresh validation. + 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: - H100: - status: unverified - precision: bf16 + all: + status: unsupported + precision: null enabled_features: {} command: null last_verified: null @@ -101,42 +171,99 @@ items: last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - No bounded full-model pretraining run is part of this focused - stack-refresh validation. + 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: unverified + status: verified precision: bf16 enabled_features: {} - command: null - last_verified: null + 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: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null + 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: > - Focused tiny-model optimizer coverage does not satisfy the full-model - SFT verification gate. + 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: null + 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: > - No post-SFT export and deterministic Hugging Face generation run is - part of this focused stack-refresh validation. + 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: {} - command: null + 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 @@ -144,29 +271,49 @@ items: last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - Focused packed-sequence and context-parallel tests do not satisfy the - full-model long-context SFT verification gate. + 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: unverified + status: verified precision: bf16 enabled_features: {} - command: null - last_verified: null + 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: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null + 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: > - No full-model parameter-efficient finetuning run is part of this - focused stack-refresh validation. + 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: - H100: - status: unverified - precision: bf16 + all: + status: unsupported + precision: null depends_on: pretrain command: null last_verified: null @@ -175,12 +322,8 @@ items: final_loss: null last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null - resume_comparison: - reference_item: pretrain - sentinel_steps: [51, 100] - loss_relative_tolerance: 1.0e-2 - loss_absolute_tolerance: 1.0e-6 - sentinels_match: false + resume_comparison: null expected_result: > - No full-state checkpoint continuation is part of this focused - stack-refresh validation. + 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. From 0607c1f8e18f972683752b0a3270eca9b1bbeb5a Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Sun, 2 Aug 2026 12:47:51 -0700 Subject: [PATCH 24/30] fix(recipe): declare long-context Omni environment Signed-off-by: Chen Cui --- .../bridge/recipes/nemotron_omni/h100/nemotron_omni.py | 5 +++++ 1 file changed, 5 insertions(+) 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 308dbd322d..b0d00b58a8 100644 --- a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py +++ b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py @@ -118,6 +118,11 @@ def nemotron_omni_cord_v2_long_context_sft_8gpu_h100_bf16_config() -> ConfigCont 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 From dc2eb3de1f4cebb970914c183d32bbb27d48a792 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 3 Aug 2026 15:31:44 -0700 Subject: [PATCH 25/30] fix(conversion): release temporary import process groups Signed-off-by: Chen Cui --- .../bridge/models/conversion/auto_bridge.py | 66 ++++--- tests/unit_tests/models/test_auto_bridge.py | 167 ++++++++++++++++++ 2 files changed, 212 insertions(+), 21 deletions(-) diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index 819d24913e..9d7ec282f9 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -30,6 +30,7 @@ if TYPE_CHECKING: from megatron.bridge.peft.base import PEFT +from megatron.core import parallel_state from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig from safetensors.torch import save_file @@ -1389,30 +1390,53 @@ def import_ckpt( ... low_memory_save=True ... ) """ - # Load the HuggingFace model - bridge = cls.from_hf_pretrained(hf_model_id, **kwargs) + owns_distributed_state = not dist.is_initialized() + owns_model_parallel_state = not parallel_state.is_initialized() + body_failed = False + try: + # Load the HuggingFace model + 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) + # 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, + ) + except BaseException: + body_failed = True + raise + finally: + try: + if owns_model_parallel_state: + parallel_state.destroy_model_parallel() + except Exception: + if not body_failed: + raise + logger.exception("Failed to release model-parallel state after checkpoint import failure") + finally: + try: + if owns_distributed_state and dist.is_initialized(): + dist.destroy_process_group() + except Exception: + if not body_failed: + raise + logger.exception("Failed to release the process group after checkpoint import failure") def export_ckpt( self, diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index a0b64f0c0e..f5a960ded1 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -1619,6 +1619,173 @@ def test_import_ckpt_with_low_memory_save( low_memory_save=True, ) + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel") + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) + @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") + @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", side_effect=[False, True]) + @patch.object(AutoBridge, "from_hf_pretrained") + def test_import_ckpt_cleans_up_distributed_state_it_creates( + self, + mock_from_hf_pretrained, + mock_dist_is_initialized, + mock_destroy_process_group, + mock_parallel_state_is_initialized, + mock_destroy_model_parallel, + ): + """Import releases a temporary Gloo group before later GPU work starts.""" + 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") + + assert mock_dist_is_initialized.call_count == 2 + mock_parallel_state_is_initialized.assert_called_once_with() + mock_destroy_model_parallel.assert_called_once_with() + mock_destroy_process_group.assert_called_once_with() + + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel") + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) + @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") + @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=True) + @patch.object(AutoBridge, "from_hf_pretrained") + def test_import_ckpt_cleans_up_model_parallel_state_it_creates( + self, + mock_from_hf_pretrained, + mock_dist_is_initialized, + mock_destroy_process_group, + mock_parallel_state_is_initialized, + mock_destroy_model_parallel, + ): + """Import preserves a caller's process group but releases model-parallel groups it creates.""" + 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_parallel_state_is_initialized.assert_called_once_with() + mock_destroy_model_parallel.assert_called_once_with() + mock_destroy_process_group.assert_not_called() + + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel") + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) + @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") + @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=True) + @patch.object(AutoBridge, "from_hf_pretrained") + def test_import_ckpt_cleans_up_partial_model_parallel_state_on_failure( + self, + mock_from_hf_pretrained, + mock_dist_is_initialized, + mock_destroy_process_group, + mock_parallel_state_is_initialized, + mock_destroy_model_parallel, + ): + """Import releases partial model-parallel state even when initialization fails.""" + mock_bridge = Mock(spec=AutoBridge) + mock_bridge.to_megatron_model.side_effect = RuntimeError("model-parallel initialization failed") + mock_from_hf_pretrained.return_value = mock_bridge + + with pytest.raises(RuntimeError, match="model-parallel initialization failed"): + AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") + + mock_dist_is_initialized.assert_called_once_with() + mock_parallel_state_is_initialized.assert_called_once_with() + mock_destroy_model_parallel.assert_called_once_with() + mock_destroy_process_group.assert_not_called() + + @patch( + "megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel", + side_effect=RuntimeError("model-parallel cleanup failed"), + ) + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) + @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") + @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", side_effect=[False, True]) + @patch.object(AutoBridge, "from_hf_pretrained", side_effect=RuntimeError("checkpoint import failed")) + def test_import_ckpt_preserves_failure_and_finishes_process_group_cleanup( + self, + mock_from_hf_pretrained, + mock_dist_is_initialized, + mock_destroy_process_group, + mock_parallel_state_is_initialized, + mock_destroy_model_parallel, + ): + """Cleanup failures do not mask an import failure or skip default-group teardown.""" + with pytest.raises(RuntimeError, match="checkpoint import failed"): + AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") + + mock_from_hf_pretrained.assert_called_once_with("./local_model") + assert mock_dist_is_initialized.call_count == 2 + mock_parallel_state_is_initialized.assert_called_once_with() + mock_destroy_model_parallel.assert_called_once_with() + mock_destroy_process_group.assert_called_once_with() + + @patch( + "megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel", + side_effect=RuntimeError("model-parallel cleanup failed"), + ) + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) + @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") + @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", side_effect=[False, True]) + @patch.object(AutoBridge, "from_hf_pretrained") + def test_import_ckpt_propagates_cleanup_failure_inside_outer_exception_handler( + self, + mock_from_hf_pretrained, + mock_dist_is_initialized, + mock_destroy_process_group, + mock_parallel_state_is_initialized, + mock_destroy_model_parallel, + ): + """An outer exception context does not make a successful import suppress cleanup failures.""" + 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 + + try: + raise ValueError("outer caller error") + except ValueError: + with pytest.raises(RuntimeError, match="model-parallel cleanup failed"): + AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") + + assert mock_dist_is_initialized.call_count == 2 + mock_parallel_state_is_initialized.assert_called_once_with() + mock_destroy_model_parallel.assert_called_once_with() + mock_destroy_process_group.assert_called_once_with() + + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel") + @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=True) + @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") + @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_state( + self, + mock_from_hf_pretrained, + mock_dist_is_initialized, + mock_destroy_process_group, + mock_parallel_state_is_initialized, + mock_destroy_model_parallel, + ): + """Import does not tear down 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_parallel_state_is_initialized.assert_called_once_with() + mock_destroy_model_parallel.assert_not_called() + mock_destroy_process_group.assert_not_called() + def test_export_ckpt_basic(self): """Test basic export_ckpt functionality.""" # Setup mocks From ef8f8bc35739af0fd47e04f9640ad78ab8adb2e9 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 15:41:24 -0700 Subject: [PATCH 26/30] fix(examples): avoid Hugging Face tensor parallel loading Signed-off-by: Chen Cui --- .../compare_hf_and_megatron/compare.py | 45 +++---------------- .../unit_tests/test_compare_mask_handling.py | 19 ++++---- 2 files changed, 15 insertions(+), 49 deletions(-) diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 2084ba65b3..e968afee66 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -483,24 +483,8 @@ def process_inputs(tokenizer, processor, image_path: Optional[str], prompt: str, return input_ids, None, None, None -def _is_transformers_none_tp_plan_error(error: TypeError) -> bool: - """Return whether Transformers failed allocator warmup on an unset TP plan.""" - if str(error) != "object of type 'NoneType' has no len()": - return False - - traceback = error.__traceback__ - while traceback is not None: - frame = traceback.tb_frame - if frame.f_code.co_name == "get_total_byte_count" and frame.f_globals.get("__name__") == ( - "transformers.modeling_utils" - ): - return True - traceback = traceback.tb_next - return False - - 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. @@ -522,24 +506,7 @@ def _load_hf_model(args, is_vl_model: bool): ), **_hf_revision_kwargs(args.hf_revision), } - try: - hf_model = model_class.from_pretrained( - args.hf_model_path, - device_map=args.hf_device, - **load_kwargs, - ) - except TypeError as error: - if not _is_transformers_none_tp_plan_error(error): - raise - - print_rank_0( - "HuggingFace model has an unset tensor-parallel plan; retrying through CPU before moving it to CUDA." - ) - gc.collect() - torch.cuda.empty_cache() - hf_model = model_class.from_pretrained(args.hf_model_path, **load_kwargs).to(args.hf_device) - - 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 @@ -584,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_") diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index ea195db738..7f7444393d 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -357,8 +357,8 @@ def test_hf_revision_is_parsed_and_forwarded(self): assert compare._hf_revision_kwargs(args.hf_revision) == {"revision": revision} assert compare._hf_revision_kwargs(None) == {} - def test_hf_loader_retries_unset_transformers_tp_plan_through_cpu(self): - """Retry the Transformers allocator bug without weakening other TypeErrors.""" + 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", @@ -368,27 +368,24 @@ def test_hf_loader_retries_unset_transformers_tp_plan_through_cpu(self): ] ) loaded_model = MagicMock() - tp_plan_error = TypeError("object of type 'NoneType' has no len()") - model_class = MagicMock() model_class.__name__ = "MockModel" - model_class.from_pretrained.side_effect = [tp_plan_error, loaded_model] + 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_transformers_none_tp_plan_error", return_value=True), patch.object(compare, "is_safe_repo", return_value=True), patch.object(compare, "print_rank_0"), - patch.object(compare.gc, "collect"), - patch.object(compare.torch.cuda, "empty_cache"), ): result = compare._load_hf_model(args, is_vl_model=False) assert result is loaded_model - assert model_class.from_pretrained.call_count == 2 - assert model_class.from_pretrained.call_args_list[0].kwargs["device_map"] == "cuda" - assert "device_map" not in model_class.from_pretrained.call_args_list[1].kwargs + 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") From 878464457120ad39579ec75ce33ff462d930814b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 16:38:42 -0700 Subject: [PATCH 27/30] fix(conversion): scope CPU import distributed state Signed-off-by: Chen Cui --- .../bridge/models/conversion/auto_bridge.py | 32 +-- src/megatron/bridge/models/model_provider.py | 6 +- tests/unit_tests/models/test_auto_bridge.py | 190 ++++-------------- .../models/test_model_provider_mixin.py | 26 --- .../training/test_model_load_save.py | 7 +- 5 files changed, 54 insertions(+), 207 deletions(-) diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index 9d7ec282f9..bd1ccba0e0 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -30,7 +30,6 @@ if TYPE_CHECKING: from megatron.bridge.peft.base import PEFT -from megatron.core import parallel_state from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig from safetensors.torch import save_file @@ -1390,13 +1389,13 @@ def import_ckpt( ... low_memory_save=True ... ) """ - owns_distributed_state = not dist.is_initialized() - owns_model_parallel_state = not parallel_state.is_initialized() - body_failed = False - try: - # Load the HuggingFace model - bridge = cls.from_hf_pretrained(hf_model_id, **kwargs) + # Load the HuggingFace model before creating temporary distributed state. + bridge = cls.from_hf_pretrained(hf_model_id, **kwargs) + + 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) @@ -1418,25 +1417,6 @@ def import_ckpt( hf_tokenizer_kwargs=hf_tokenizer_kwargs, low_memory_save=low_memory_save, ) - except BaseException: - body_failed = True - raise - finally: - try: - if owns_model_parallel_state: - parallel_state.destroy_model_parallel() - except Exception: - if not body_failed: - raise - logger.exception("Failed to release model-parallel state after checkpoint import failure") - finally: - try: - if owns_distributed_state and dist.is_initialized(): - dist.destroy_process_group() - except Exception: - if not body_failed: - raise - logger.exception("Failed to release the process group after checkpoint import failure") def export_ckpt( self, diff --git a/src/megatron/bridge/models/model_provider.py b/src/megatron/bridge/models/model_provider.py index 2edf28f745..1de2dd49f5 100644 --- a/src/megatron/bridge/models/model_provider.py +++ b/src/megatron/bridge/models/model_provider.py @@ -252,10 +252,8 @@ def provide_distributed_model( os.environ["WORLD_SIZE"] = os.environ.get("WORLD_SIZE", "1") os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "localhost") os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "12355") - backend = "gloo" if use_cpu_initialization else "nccl" - if backend == "nccl": - torch.cuda.set_device(get_local_rank_preinit()) - torch.distributed.init_process_group(backend) + torch.cuda.set_device(get_local_rank_preinit()) + torch.distributed.init_process_group("nccl") # If pg_collection is provided (e.g., from use_decentralized_pg=True), # use it directly. Otherwise, initialize model parallel state and get pg_collection from MPU. diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index f5a960ded1..239816f399 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 @@ -1534,7 +1535,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") @@ -1562,13 +1567,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( @@ -1600,12 +1609,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", @@ -1619,47 +1632,16 @@ def test_import_ckpt_with_low_memory_save( low_memory_save=True, ) - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel") - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) - @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") - @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", side_effect=[False, True]) - @patch.object(AutoBridge, "from_hf_pretrained") - def test_import_ckpt_cleans_up_distributed_state_it_creates( - self, - mock_from_hf_pretrained, - mock_dist_is_initialized, - mock_destroy_process_group, - mock_parallel_state_is_initialized, - mock_destroy_model_parallel, - ): - """Import releases a temporary Gloo group before later GPU work starts.""" - 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") - - assert mock_dist_is_initialized.call_count == 2 - mock_parallel_state_is_initialized.assert_called_once_with() - mock_destroy_model_parallel.assert_called_once_with() - mock_destroy_process_group.assert_called_once_with() - - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel") - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) - @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") - @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=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_cleans_up_model_parallel_state_it_creates( + def test_import_ckpt_scopes_standalone_cpu_state_to_gloo_context( self, mock_from_hf_pretrained, mock_dist_is_initialized, - mock_destroy_process_group, - mock_parallel_state_is_initialized, - mock_destroy_model_parallel, + mock_temporary_distributed_context, ): - """Import preserves a caller's process group but releases model-parallel groups it creates.""" + """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() @@ -1669,110 +1651,20 @@ def test_import_ckpt_cleans_up_model_parallel_state_it_creates( AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") mock_dist_is_initialized.assert_called_once_with() - mock_parallel_state_is_initialized.assert_called_once_with() - mock_destroy_model_parallel.assert_called_once_with() - mock_destroy_process_group.assert_not_called() - - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel") - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) - @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") - @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=True) - @patch.object(AutoBridge, "from_hf_pretrained") - def test_import_ckpt_cleans_up_partial_model_parallel_state_on_failure( - self, - mock_from_hf_pretrained, - mock_dist_is_initialized, - mock_destroy_process_group, - mock_parallel_state_is_initialized, - mock_destroy_model_parallel, - ): - """Import releases partial model-parallel state even when initialization fails.""" - mock_bridge = Mock(spec=AutoBridge) - mock_bridge.to_megatron_model.side_effect = RuntimeError("model-parallel initialization failed") - mock_from_hf_pretrained.return_value = mock_bridge - - with pytest.raises(RuntimeError, match="model-parallel initialization failed"): - AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") - - mock_dist_is_initialized.assert_called_once_with() - mock_parallel_state_is_initialized.assert_called_once_with() - mock_destroy_model_parallel.assert_called_once_with() - mock_destroy_process_group.assert_not_called() - - @patch( - "megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel", - side_effect=RuntimeError("model-parallel cleanup failed"), - ) - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) - @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") - @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", side_effect=[False, True]) - @patch.object(AutoBridge, "from_hf_pretrained", side_effect=RuntimeError("checkpoint import failed")) - def test_import_ckpt_preserves_failure_and_finishes_process_group_cleanup( - self, - mock_from_hf_pretrained, - mock_dist_is_initialized, - mock_destroy_process_group, - mock_parallel_state_is_initialized, - mock_destroy_model_parallel, - ): - """Cleanup failures do not mask an import failure or skip default-group teardown.""" - with pytest.raises(RuntimeError, match="checkpoint import failed"): - AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") - - mock_from_hf_pretrained.assert_called_once_with("./local_model") - assert mock_dist_is_initialized.call_count == 2 - mock_parallel_state_is_initialized.assert_called_once_with() - mock_destroy_model_parallel.assert_called_once_with() - mock_destroy_process_group.assert_called_once_with() - - @patch( - "megatron.bridge.models.conversion.auto_bridge.parallel_state.destroy_model_parallel", - side_effect=RuntimeError("model-parallel cleanup failed"), - ) - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=False) - @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") - @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", side_effect=[False, True]) - @patch.object(AutoBridge, "from_hf_pretrained") - def test_import_ckpt_propagates_cleanup_failure_inside_outer_exception_handler( - self, - mock_from_hf_pretrained, - mock_dist_is_initialized, - mock_destroy_process_group, - mock_parallel_state_is_initialized, - mock_destroy_model_parallel, - ): - """An outer exception context does not make a successful import suppress cleanup failures.""" - 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 - - try: - raise ValueError("outer caller error") - except ValueError: - with pytest.raises(RuntimeError, match="model-parallel cleanup failed"): - AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") - - assert mock_dist_is_initialized.call_count == 2 - mock_parallel_state_is_initialized.assert_called_once_with() - mock_destroy_model_parallel.assert_called_once_with() - mock_destroy_process_group.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.models.conversion.auto_bridge.parallel_state.destroy_model_parallel") - @patch("megatron.bridge.models.conversion.auto_bridge.parallel_state.is_initialized", return_value=True) - @patch("megatron.bridge.models.conversion.auto_bridge.dist.destroy_process_group") + @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_state( + def test_import_ckpt_preserves_existing_distributed_context( self, mock_from_hf_pretrained, mock_dist_is_initialized, - mock_destroy_process_group, - mock_parallel_state_is_initialized, - mock_destroy_model_parallel, + mock_temporary_distributed_context, ): - """Import does not tear down distributed state owned by its caller.""" + """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() @@ -1782,9 +1674,7 @@ def test_import_ckpt_preserves_existing_distributed_state( AutoBridge.import_ckpt("./local_model", "./megatron_checkpoint") mock_dist_is_initialized.assert_called_once_with() - mock_parallel_state_is_initialized.assert_called_once_with() - mock_destroy_model_parallel.assert_not_called() - mock_destroy_process_group.assert_not_called() + mock_temporary_distributed_context.assert_not_called() def test_export_ckpt_basic(self): """Test basic export_ckpt functionality.""" diff --git a/tests/unit_tests/models/test_model_provider_mixin.py b/tests/unit_tests/models/test_model_provider_mixin.py index bfbe3a6c52..88a90c6d78 100644 --- a/tests/unit_tests/models/test_model_provider_mixin.py +++ b/tests/unit_tests/models/test_model_provider_mixin.py @@ -228,32 +228,6 @@ def __init__(self): assert provider._pg_collection is pg_instance -@patch("megatron.bridge.models.model_provider.ProcessGroupCollection.use_mpu_process_groups") -@patch("megatron.bridge.models.model_provider.get_model") -@patch("megatron.bridge.models.model_provider.torch.cuda.set_device") -@patch("megatron.bridge.models.model_provider.torch.distributed") -@patch("megatron.bridge.models.model_provider.parallel_state.is_initialized", return_value=True) -def test_cpu_initialization_starts_gloo_without_selecting_cuda( - mock_ps_init, - mock_dist, - mock_set_device, - mock_get_model, - mock_use_pg, - provider, -): - """Standalone CPU model construction must not require a CUDA driver.""" - mock_dist.is_initialized.return_value = False - mock_model = [MockMegatronModule()] - mock_get_model.return_value = mock_model - mock_use_pg.return_value = Mock() - - result = provider.provide_distributed_model(wrap_with_ddp=False, use_cpu_initialization=True) - - assert result is mock_model - mock_dist.init_process_group.assert_called_once_with("gloo") - mock_set_device.assert_not_called() - - def test_hook_registration_and_composition(provider): """Test hook registration order and composition.""" # Initially, no hooks are registered 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") From 64a20525a375a498795eca99c182a11b50c24872 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 16:38:47 -0700 Subject: [PATCH 28/30] fix(data): clear stale packed padding masks Signed-off-by: Chen Cui --- src/megatron/bridge/data/packing/in_batch.py | 2 +- tests/unit_tests/data/packing/test_in_batch.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/data/packing/in_batch.py b/src/megatron/bridge/data/packing/in_batch.py index 2828fe74d0..24937e7d33 100644 --- a/src/megatron/bridge/data/packing/in_batch.py +++ b/src/megatron/bridge/data/packing/in_batch.py @@ -336,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/tests/unit_tests/data/packing/test_in_batch.py b/tests/unit_tests/data/packing/test_in_batch.py index f777208a34..015731b65a 100644 --- a/tests/unit_tests/data/packing/test_in_batch.py +++ b/tests/unit_tests/data/packing/test_in_batch.py @@ -184,6 +184,19 @@ def test_packing_marks_only_physical_alignment_gaps_as_padding(self): 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( From 558e6b085f8cd5270c158909096b34b6b9510797 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Wed, 5 Aug 2026 08:52:50 -0700 Subject: [PATCH 29/30] fix(distributed): avoid temporary rendezvous port collisions Signed-off-by: Chen Cui --- .../bridge/training/model_load_save.py | 15 ++---- .../training/test_model_load_save.py | 49 +++++++------------ 2 files changed, 20 insertions(+), 44 deletions(-) diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index a077b8a520..714eb2384b 100644 --- a/src/megatron/bridge/training/model_load_save.py +++ b/src/megatron/bridge/training/model_load_save.py @@ -14,8 +14,6 @@ import argparse import logging -import os -import socket from contextlib import contextmanager from pathlib import Path from typing import Any, Generator, Literal, Optional, Union @@ -128,16 +126,9 @@ def temporary_distributed_context(backend: str = "gloo") -> Generator[None, None Yields: None. """ - if "MASTER_ADDR" in os.environ and "MASTER_PORT" in os.environ: - init_method = None - else: - # Find an available port dynamically - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("localhost", 0)) - addr, port = s.getsockname() - init_method = f"tcp://{addr}:{port}" - - dist.init_process_group(backend=backend, init_method=init_method, world_size=1, rank=0) + # This context is single-process, so keep rendezvous in-process and avoid TCP port collisions. + store = dist.HashStore() + dist.init_process_group(backend=backend, store=store, world_size=1, rank=0) parallel_state.initialize_model_parallel() # Initialize RNG tracker for model initialization diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index 833ef4cef8..ea4258c3a2 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/tests/unit_tests/training/test_model_load_save.py @@ -183,17 +183,9 @@ class TestTemporaryDistributedContext: @patch("megatron.bridge.training.model_load_save.dist") @patch("megatron.bridge.training.model_load_save.parallel_state") - @patch("megatron.bridge.training.model_load_save.socket") - @patch("megatron.bridge.training.model_load_save.os") - def test_temporary_distributed_context_gloo(self, mock_os, mock_socket, mock_parallel_state, mock_dist): + def test_temporary_distributed_context_gloo(self, mock_parallel_state, mock_dist): """Test temporary distributed context with gloo backend.""" - # Mock environment to not have MASTER_ADDR and MASTER_PORT - mock_os.environ = {} - - # Mock socket for port selection - mock_socket_instance = Mock() - mock_socket_instance.getsockname.return_value = ("localhost", 12345) - mock_socket.socket.return_value.__enter__.return_value = mock_socket_instance + mock_store = mock_dist.HashStore.return_value with ( patch("megatron.bridge.training.model_load_save.torch.cuda.is_available", return_value=False), @@ -202,9 +194,8 @@ def test_temporary_distributed_context_gloo(self, mock_os, mock_socket, mock_par ): pass - mock_dist.init_process_group.assert_called_once_with( - backend="gloo", init_method="tcp://localhost:12345", world_size=1, rank=0 - ) + mock_dist.HashStore.assert_called_once_with() + mock_dist.init_process_group.assert_called_once_with(backend="gloo", store=mock_store, world_size=1, rank=0) 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() @@ -212,37 +203,31 @@ def test_temporary_distributed_context_gloo(self, mock_os, mock_socket, mock_par @patch("megatron.bridge.training.model_load_save.dist") @patch("megatron.bridge.training.model_load_save.parallel_state") - @patch("megatron.bridge.training.model_load_save.os") - def test_temporary_distributed_context_with_env_vars(self, mock_os, mock_parallel_state, mock_dist): - """Test temporary distributed context when env vars are already set.""" - mock_os.environ = {"MASTER_ADDR": "localhost", "MASTER_PORT": "12345"} + def test_temporary_distributed_context_ignores_env_rendezvous(self, mock_parallel_state, mock_dist): + """Test temporary distributed context does not use the process rendezvous.""" + mock_store = mock_dist.HashStore.return_value - with temporary_distributed_context(backend="gloo"): + with ( + patch.dict(os.environ, {"MASTER_ADDR": "localhost", "MASTER_PORT": "12345"}), + temporary_distributed_context(backend="gloo"), + ): pass - mock_dist.init_process_group.assert_called_once_with(backend="gloo", init_method=None, world_size=1, rank=0) + mock_dist.HashStore.assert_called_once_with() + mock_dist.init_process_group.assert_called_once_with(backend="gloo", store=mock_store, world_size=1, rank=0) @patch("megatron.bridge.training.model_load_save.dist") @patch("megatron.bridge.training.model_load_save.parallel_state") - @patch("megatron.bridge.training.model_load_save.socket") - @patch("megatron.bridge.training.model_load_save.os") @patch("megatron.core.tensor_parallel.model_parallel_cuda_manual_seed") - def test_temporary_distributed_context_nccl(self, mock_seed, mock_os, mock_socket, mock_parallel_state, mock_dist): + def test_temporary_distributed_context_nccl(self, mock_seed, mock_parallel_state, mock_dist): """Test temporary distributed context with nccl backend.""" - # Mock environment to not have MASTER_ADDR and MASTER_PORT - mock_os.environ = {} - - # Mock socket for port selection - mock_socket_instance = Mock() - mock_socket_instance.getsockname.return_value = ("localhost", 12345) - mock_socket.socket.return_value.__enter__.return_value = mock_socket_instance + mock_store = mock_dist.HashStore.return_value with temporary_distributed_context(backend="nccl"): pass - mock_dist.init_process_group.assert_called_once_with( - backend="nccl", init_method="tcp://localhost:12345", world_size=1, rank=0 - ) + mock_dist.HashStore.assert_called_once_with() + mock_dist.init_process_group.assert_called_once_with(backend="nccl", store=mock_store, world_size=1, rank=0) mock_seed.assert_called_once_with(0) mock_parallel_state.initialize_model_parallel.assert_called_once() mock_parallel_state.destroy_model_parallel.assert_called_once() From a554f71cb533ba81cf8b764663ad3068166d844a Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Wed, 5 Aug 2026 09:07:16 -0700 Subject: [PATCH 30/30] revert: remove temporary rendezvous hardening This reverts commit 558e6b085f8cd5270c158909096b34b6b9510797. Signed-off-by: Chen Cui --- .../bridge/training/model_load_save.py | 15 ++++-- .../training/test_model_load_save.py | 49 ++++++++++++------- 2 files changed, 44 insertions(+), 20 deletions(-) diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index 714eb2384b..a077b8a520 100644 --- a/src/megatron/bridge/training/model_load_save.py +++ b/src/megatron/bridge/training/model_load_save.py @@ -14,6 +14,8 @@ import argparse import logging +import os +import socket from contextlib import contextmanager from pathlib import Path from typing import Any, Generator, Literal, Optional, Union @@ -126,9 +128,16 @@ def temporary_distributed_context(backend: str = "gloo") -> Generator[None, None Yields: None. """ - # This context is single-process, so keep rendezvous in-process and avoid TCP port collisions. - store = dist.HashStore() - dist.init_process_group(backend=backend, store=store, world_size=1, rank=0) + if "MASTER_ADDR" in os.environ and "MASTER_PORT" in os.environ: + init_method = None + else: + # Find an available port dynamically + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("localhost", 0)) + addr, port = s.getsockname() + init_method = f"tcp://{addr}:{port}" + + dist.init_process_group(backend=backend, init_method=init_method, world_size=1, rank=0) parallel_state.initialize_model_parallel() # Initialize RNG tracker for model initialization diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index ea4258c3a2..833ef4cef8 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/tests/unit_tests/training/test_model_load_save.py @@ -183,9 +183,17 @@ class TestTemporaryDistributedContext: @patch("megatron.bridge.training.model_load_save.dist") @patch("megatron.bridge.training.model_load_save.parallel_state") - def test_temporary_distributed_context_gloo(self, mock_parallel_state, mock_dist): + @patch("megatron.bridge.training.model_load_save.socket") + @patch("megatron.bridge.training.model_load_save.os") + def test_temporary_distributed_context_gloo(self, mock_os, mock_socket, mock_parallel_state, mock_dist): """Test temporary distributed context with gloo backend.""" - mock_store = mock_dist.HashStore.return_value + # Mock environment to not have MASTER_ADDR and MASTER_PORT + mock_os.environ = {} + + # Mock socket for port selection + mock_socket_instance = Mock() + mock_socket_instance.getsockname.return_value = ("localhost", 12345) + mock_socket.socket.return_value.__enter__.return_value = mock_socket_instance with ( patch("megatron.bridge.training.model_load_save.torch.cuda.is_available", return_value=False), @@ -194,8 +202,9 @@ def test_temporary_distributed_context_gloo(self, mock_parallel_state, mock_dist ): pass - mock_dist.HashStore.assert_called_once_with() - mock_dist.init_process_group.assert_called_once_with(backend="gloo", store=mock_store, world_size=1, rank=0) + mock_dist.init_process_group.assert_called_once_with( + backend="gloo", init_method="tcp://localhost:12345", world_size=1, rank=0 + ) 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() @@ -203,31 +212,37 @@ def test_temporary_distributed_context_gloo(self, mock_parallel_state, mock_dist @patch("megatron.bridge.training.model_load_save.dist") @patch("megatron.bridge.training.model_load_save.parallel_state") - def test_temporary_distributed_context_ignores_env_rendezvous(self, mock_parallel_state, mock_dist): - """Test temporary distributed context does not use the process rendezvous.""" - mock_store = mock_dist.HashStore.return_value + @patch("megatron.bridge.training.model_load_save.os") + def test_temporary_distributed_context_with_env_vars(self, mock_os, mock_parallel_state, mock_dist): + """Test temporary distributed context when env vars are already set.""" + mock_os.environ = {"MASTER_ADDR": "localhost", "MASTER_PORT": "12345"} - with ( - patch.dict(os.environ, {"MASTER_ADDR": "localhost", "MASTER_PORT": "12345"}), - temporary_distributed_context(backend="gloo"), - ): + with temporary_distributed_context(backend="gloo"): pass - mock_dist.HashStore.assert_called_once_with() - mock_dist.init_process_group.assert_called_once_with(backend="gloo", store=mock_store, world_size=1, rank=0) + mock_dist.init_process_group.assert_called_once_with(backend="gloo", init_method=None, world_size=1, rank=0) @patch("megatron.bridge.training.model_load_save.dist") @patch("megatron.bridge.training.model_load_save.parallel_state") + @patch("megatron.bridge.training.model_load_save.socket") + @patch("megatron.bridge.training.model_load_save.os") @patch("megatron.core.tensor_parallel.model_parallel_cuda_manual_seed") - def test_temporary_distributed_context_nccl(self, mock_seed, mock_parallel_state, mock_dist): + def test_temporary_distributed_context_nccl(self, mock_seed, mock_os, mock_socket, mock_parallel_state, mock_dist): """Test temporary distributed context with nccl backend.""" - mock_store = mock_dist.HashStore.return_value + # Mock environment to not have MASTER_ADDR and MASTER_PORT + mock_os.environ = {} + + # Mock socket for port selection + mock_socket_instance = Mock() + mock_socket_instance.getsockname.return_value = ("localhost", 12345) + mock_socket.socket.return_value.__enter__.return_value = mock_socket_instance with temporary_distributed_context(backend="nccl"): pass - mock_dist.HashStore.assert_called_once_with() - mock_dist.init_process_group.assert_called_once_with(backend="nccl", store=mock_store, world_size=1, rank=0) + mock_dist.init_process_group.assert_called_once_with( + backend="nccl", init_method="tcp://localhost:12345", world_size=1, rank=0 + ) mock_seed.assert_called_once_with(0) mock_parallel_state.initialize_model_parallel.assert_called_once() mock_parallel_state.destroy_model_parallel.assert_called_once()