perf(recipe): tune Nemotron Omni H100 fine-tuning (#5524)

Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: adityavavreNVDA <avavre@nvidia.com>
Co-authored-by: Kamran Jafari <kjafarisadeg@nvidia.com>
Co-authored-by: adityavavreNVDA <avavre@nvidia.com>

Author: 3 people

examples/model_verification_cards/nemotron-3-nano-omni-30b-a3b-reasoning/card.yaml

@@ -6,15 +6,17 @@ summary: >
   Performance disclaimer: this model has not been performance-tuned; reported
   timing and throughput metrics are sanity checks, not optimized performance
   results. Verification uses the immutable public model and CORD v2 revisions.
-  CPU and distributed GPU import, deterministic Megatron inference, bounded
-  full-model SFT, and LoRA PEFT runs completed. Strict CPU and GPU round trips
-  preserved all 7,349 tensors bitwise, but Transformers 5.8.0 cannot natively
-  reload the local custom-code exports because its dynamic-module cache omits
-  transitive configuration imports. The one-step HF/Megatron comparison
-  predicts the same token but remains below the 0.99 cosine gate. The packed
-  8K long-context recipe completes one optimizer step but does not complete the
-  second because of H100 memory pressure. Unsupported and incomplete workflows
-  remain explicitly identified rather than inferred from focused unit coverage.
+  CPU and distributed GPU import, deterministic Megatron inference, 100-step
+  one-node H100 full-model SFT, in-batch-packed 8K SFT, and LoRA PEFT runs
+  completed. The H100 fine-tuning variants preserve real image-text samples and
+  natural routing while using lower-precision optimizer state; 4K SFT and PEFT
+  use HybridEP, while packed CP2 uses the standard all-to-all dispatcher. 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. Unsupported and incomplete workflows remain explicitly
+  identified rather than inferred from focused unit coverage.
 verification_index:
   model_level:
     verified:
@@ -27,8 +29,8 @@ verification_index:
       - manual_forward_pass
   training:
     H100:
-      verified: [sft, peft]
-      unverified: [sft_export_inference, sft_long_context]
+      verified: [sft, sft_long_context, peft]
+      unverified: [sft_export_inference]
       unsupported: [pretrain, checkpoint_resume]
 model:
   hf_id: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
@@ -180,37 +182,45 @@ items:
     H100:
       status: verified
       precision: bf16
-      enabled_features: {}
+      bridge_commit: 9b69b47a321de14b0453c230dd130e18ae0f00f7  # pragma: allowlist secret
+      enabled_features:
+        moe_dispatcher: hybridep
       command: >
-        ./scripts/training/train.sh --nodes 2 --gpus-per-node 8
-        --recipe nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config
+        ./scripts/training/train.sh --wait --nodes 1 --gpus-per-node 8
+        --recipe nemotron_omni_cord_v2_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 --tensor_model_parallel_size 2
-        --pipeline_model_parallel_size 2 --expert_model_parallel_size 4
-        --expert_tensor_parallel_size 1
+        --max_steps 100
         --save_dir
-        work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/sft-checkpoints-clean
-        --save_interval 10
+        work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/sft-4k-checkpoints
+        --save_interval 50
         '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
+      last_verified: 2026-08-12
       metrics:
-        initial_loss: 1.123339
-        final_loss: 0.4893276
-        last_10_steps_step_time_ms_avg: 24454.11
-        last_10_steps_model_tflops_per_gpu_avg: 60.73
-        last_10_steps_tokens_per_second_per_gpu_avg: 669.990
+        initial_loss: 1.281308
+        final_loss: 0.03330838
+        last_10_steps_step_time_ms_avg: 3603.67
+        last_10_steps_model_tflops_per_gpu_avg: 197.10
+        last_10_steps_tokens_per_second_per_gpu_avg: 9092.952
       expected_result: >
-        The immutable-revision CORD v2 run completes exactly 10 full-SFT
-        optimizer steps on 16 H100 GPUs at TP2/PP2/CP1/EP4/ETP1, GBS/MBS
-        64/1. LM loss is finite from 1.123339 to 0.4893276, all ten recorded
-        steps average 24,454.11 ms and 60.73 TFLOP/s/GPU including first-step
-        compilation, no iteration is skipped or NaN, and a complete
-        iter_0000010 checkpoint is saved.
+        The immutable-revision, real image-text CORD v2 run completes exactly
+        100 full-SFT optimizer steps on eight H100 GPUs at
+        TP2/PP1/CP1/EP8/ETP1 and GBS/MBS 64/4 with natural-routing HybridEP.
+        Execution uses selective MoE/layernorm recompute, fused attention, TE
+        cross entropy, fused weighted squared-ReLU, and router fusion without
+        changing the image-text objective or routing policy. Model parameters,
+        main gradients, and Adam moments use BF16; optimizer main parameters use
+        scaled FP16 without parameter remainders. LM loss is finite from
+        1.281308 to 0.03330838, the final ten steps average 3,603.67 ms and
+        197.10 TFLOP/s/GPU, sampled peak memory is 76,349 MiB, and no iteration
+        is skipped or NaN. Complete optimizer and RNG checkpoints are saved at
+        steps 50 and 100. Direct step-50 resume loads the complete state and
+        reaches step 100 at 197.07 final-ten TFLOP/s/GPU with finite loss and
+        zero skipped or NaN iterations.
 
   sft_export_inference:
     H100:
@@ -224,10 +234,10 @@ items:
           --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
+          work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/sft-4k-checkpoints/iter_0000100
           --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
+          --torch-dtype bfloat16 --tp 2 --pp 1 --ep 8 --etp 1
           --trust-remote-code --distributed-save --not-strict
         - >
           uv run python
@@ -247,72 +257,86 @@ items:
 
   sft_long_context:
     H100:
-      status: unverified
+      status: verified
       precision: bf16
+      bridge_commit: eb7cf72392ac07e7a095b87d9b1d6bf2b38c5b6e  # pragma: allowlist secret
       enabled_features:
         sequence_packing: in_batch
         context_parallel_size: 2
       command: >
-        ./scripts/training/train.sh --nodes 1 --gpus-per-node 8
+        ./scripts/training/train.sh --wait --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
+        --max_steps 100
         --save_dir
-        work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/long-context-checkpoints-clean
-        --save_interval 10
+        work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/sft-8k-packed-checkpoints
+        --save_interval 50
         '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
+      last_verified: 2026-08-11
       metrics:
-        initial_loss: null
-        final_loss: null
-        last_10_steps_step_time_ms_avg: null
-        last_10_steps_model_tflops_per_gpu_avg: null
-        last_10_steps_tokens_per_second_per_gpu_avg: null
+        initial_loss: 1.251364
+        final_loss: 0.01020405
+        last_10_steps_step_time_ms_avg: 15577.9
+        last_10_steps_model_tflops_per_gpu_avg: 93.73
+        last_10_steps_tokens_per_second_per_gpu_avg: 4206.986
       expected_result: >
-        The 8K TP4/PP1/CP2/EP1/ETP4, MBS2 in-batch-packing run uses
-        precision-aware Adam with FP16 main parameters and stored FP32
-        remainders, BF16 gradients, and BF16 moments. Step 1 completes with
-        finite LM loss 1.142561 in 170,179.0 ms at 8.6 TFLOP/s/GPU with no
-        skipped or NaN iteration, but step 2 encounters rank-divergent H100
-        memory exhaustion and does not produce the required 10-step checkpoint.
+        The immutable-revision, real image-text CORD v2 run completes exactly
+        100 full-SFT optimizer steps on eight H100 GPUs at
+        TP4/PP1/CP2/EP8/ETP1 and GBS/MBS 64/2 with 8K in-batch packing and the
+        standard all-to-all dispatcher. Model parameters, main gradients, and
+        Adam moments use BF16; optimizer main parameters use scaled FP16 without
+        parameter remainders. LM loss is finite from 1.251364 to 0.01020405, the
+        final ten steps average 15,577.9 ms and 93.73 TFLOP/s/GPU, sampled peak
+        memory is 75,197 MiB, and no iteration is skipped or NaN. Complete
+        optimizer and RNG checkpoints are saved at steps 50 and 100. Direct
+        step-50 resume loads the complete state and reaches step 100 with finite
+        loss and zero skipped or NaN iterations.
 
   peft:
     H100:
       status: verified
       precision: bf16
-      enabled_features: {}
+      bridge_commit: eb7cf72392ac07e7a095b87d9b1d6bf2b38c5b6e  # pragma: allowlist secret
+      enabled_features:
+        moe_dispatcher: hybridep
       command: >
-        ./scripts/training/train.sh --nodes 1 --gpus-per-node 8
-        --recipe nemotron_omni_cord_v2_peft_4gpu_h100_bf16_config
+        ./scripts/training/train.sh --wait --nodes 1 --gpus-per-node 8
+        --recipe nemotron_omni_cord_v2_peft_8gpu_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
+        --max_steps 100
         --save_dir
-        work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/peft-checkpoints-clean
-        --save_interval 10
+        work/model-verification/nemotron-3-nano-omni-30b-a3b-reasoning/peft-checkpoints
+        --save_interval 50
         '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
+      last_verified: 2026-08-11
       metrics:
-        initial_loss: 1.098394
-        final_loss: 0.3166811
-        last_10_steps_step_time_ms_avg: 41897.41
-        last_10_steps_model_tflops_per_gpu_avg: 22.22
-        last_10_steps_tokens_per_second_per_gpu_avg: 782.101
+        initial_loss: 1.282941
+        final_loss: 0.02634745
+        last_10_steps_step_time_ms_avg: 11621.42
+        last_10_steps_model_tflops_per_gpu_avg: 61.25
+        last_10_steps_tokens_per_second_per_gpu_avg: 2819.621
       expected_result: >
-        The ten-step TP4/PP1/CP1/EP1 LoRA run exits successfully and saves a
-        complete eight-shard iter_0000010 adapter checkpoint. LM loss is finite
-        from 1.098394 to 0.3166811, all ten steps average 41,897.41 ms and
-        22.22 TFLOP/s/GPU including first-step compilation, and no iteration is
-        skipped or NaN.
+        The immutable-revision, real image-text CORD v2 run completes exactly
+        100 LoRA optimizer steps on eight H100 GPUs at
+        TP2/PP1/CP1/EP8/ETP1 and GBS/MBS 64/1 with natural-routing HybridEP.
+        Model parameters, main gradients, and Adam moments use BF16; optimizer
+        main parameters use scaled FP16 without parameter remainders. LM loss is
+        finite from 1.282941 to 0.02634745, the final ten steps average
+        11,621.42 ms and 61.25 TFLOP/s/GPU, sampled peak memory is 44,873 MiB,
+        and no iteration is skipped or NaN. Complete optimizer and RNG
+        checkpoints are saved at steps 50 and 100. Direct step-50 resume loads
+        the complete state and reaches step 100 with finite loss and zero
+        skipped or NaN iterations.
 
   checkpoint_resume:
     all:

src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py

@@ -32,6 +32,7 @@
 
 import torch
 from megatron.core import tensor_parallel
+from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.models.hybrid.hybrid_model import HybridModel
 from megatron.core.models.multimodal.llava_model import pixel_shuffle
 from megatron.core.models.vision.multimodal_projector import MultimodalProjector
@@ -107,6 +108,28 @@ def _pixel_shuffle_dynamic_resolution(
     return shuffled.reshape(batch, (height * width) // 4, hidden * 4)
 
 
+def _project_multimodal_embeddings(
+    projection: torch.nn.Module,
+    embeddings: torch.Tensor,
+) -> torch.Tensor:
+    """Project media rows, padding only the temporary FP8 compute input."""
+    input_shape = embeddings.shape[:-1]
+    flat_embeddings = embeddings.reshape(-1, 1, embeddings.shape[-1])
+    num_embeddings = flat_embeddings.shape[0]
+    projection_config = getattr(projection, "config", None)
+    if getattr(projection_config, "fp8", None):
+        alignment = get_fp8_align_size(projection_config.fp8_recipe)
+        padding = -num_embeddings % alignment
+        if padding:
+            flat_embeddings = torch.cat(
+                (flat_embeddings, flat_embeddings.new_zeros((padding, 1, flat_embeddings.shape[-1]))),
+                dim=0,
+            )
+
+    projected = projection(flat_embeddings)[:num_embeddings]
+    return projected.reshape(*input_shape, projected.shape[-1])
+
+
 class NemotronOmniModel(MegatronModule):
     """Nemotron Omni model whose input sequence is already media-expanded.
 
@@ -389,8 +412,7 @@ def _encode_images(
             encoded = encoded[:, class_tokens:, :]
             encoded = pixel_shuffle(encoded).reshape(-1, encoded.shape[-1] * 4)
 
-        projected = self.vision_projection(encoded.unsqueeze(1))
-        return projected.squeeze(1).contiguous()
+        return _project_multimodal_embeddings(self.vision_projection, encoded).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."""
@@ -419,7 +441,10 @@ def _encode_sound(self, sound_clips: torch.Tensor, sound_length: Optional[torch.
         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 = _project_multimodal_embeddings(
+            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()

src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py

@@ -101,6 +101,9 @@ class NemotronVLModelProvider(HybridModelProvider, ABC):
     def _build_vision_config(self, language_cfg):
         """Build RADIO ViT-H vision encoder config from a language config copy."""
         vision_cfg = copy.deepcopy(language_cfg)
+        if not self.use_vision_backbone_fp8_arch:
+            vision_cfg.fp8 = None
+            vision_cfg.fp8_param = False
         vision_cfg.sequence_parallel = False
         vision_cfg.context_parallel_size = 1
         vision_cfg.tp_comm_overlap = False

src/megatron/bridge/recipes/nemotron_omni/h100/__init__.py

@@ -18,7 +18,9 @@
 __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_peft_8gpu_h100_bf16_config",
     "nemotron_omni_cord_v2_sft_4gpu_h100_bf16_config",
+    "nemotron_omni_cord_v2_sft_8gpu_h100_bf16_config",
     "nemotron_omni_valor32k_peft_4gpu_h100_bf16_config",
     "nemotron_omni_valor32k_sft_4gpu_h100_bf16_config",
 ]