perf(fix): accumulate per-microbatch FLOPS metadata for accurate…

src/megatron/bridge/training/train.py

@@ -313,7 +313,6 @@ def train(
 
     start_iteration = global_state.train_state.step
     print_rank_0(f"Starting training loop at iteration {start_iteration}")
-    num_floating_point_operations_model = flop_utils.num_floating_point_operations(config, batch_size=1)
     p2p_communicator = P2PCommunicator(pp_group=pg_collection.pp, config=model_config)
     dp_size = pg_collection.dp.size()
     if hasattr(config.model, "dist_train") and getattr(config.model.dist_train, "use_dist_train", False) is True:
@@ -433,6 +432,11 @@ def train(
                 ),
             )
 
+        # Reset per-step FLOPS accumulators (filled by forward_step micro-batches).
+        global_state._flops_seqlen_sum = 0
+        global_state._flops_seqlen_sq_sum = 0
+        global_state._flops_vision_patches = 0
+
         (
             loss_dict,
             skipped_iter,
@@ -537,7 +541,28 @@ def train(
         else:
             assert num_skipped_samples_in_batch == 0
         global_state.train_state.skipped_train_samples += num_skipped_samples_in_batch
-        num_floating_point_operations_in_batch = num_floating_point_operations_model * batch_size
+
+        # Read accumulated FLOPS metadata from forward_step micro-batches.
+        # These are per-DP-rank totals; scale by dp_size for global estimate.
+        local_seqlen_sum = getattr(global_state, '_flops_seqlen_sum', 0)
+        local_seqlen_sq_sum = getattr(global_state, '_flops_seqlen_sq_sum', 0)
+        num_vision_patches = getattr(global_state, '_flops_vision_patches', 0)
+
+        if local_seqlen_sum > 0:
+            seqlen_sum = local_seqlen_sum * dp_size
+            seqlen_squared_sum = local_seqlen_sq_sum * dp_size
+        else:
+            # Fallback for step functions that don't set accumulators
+            seqlen_sum = None
+            seqlen_squared_sum = None
+
+        # Vision patches: local accumulation * dp_size for global
+        num_vision_patches = num_vision_patches * dp_size if num_vision_patches > 0 else 0
+
+        num_floating_point_operations_in_batch = flop_utils.num_floating_point_operations(
+            config, batch_size=batch_size, seqlen_sum=seqlen_sum,
+            seqlen_squared_sum=seqlen_squared_sum, num_vision_patches=num_vision_patches,
+        )
         global_state.train_state.floating_point_operations_so_far += num_floating_point_operations_in_batch
         num_floating_point_operations_so_far = global_state.train_state.floating_point_operations_so_far
         num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch
@@ -580,6 +605,7 @@ def train(
                 model,
                 log_max_attention_logit,
                 loaded_iteration=start_iteration,
+                seq_length=seqlen_sum // batch_size if seqlen_sum else None,
             )
 
         if (

src/megatron/bridge/training/utils/flop_utils.py

@@ -26,8 +26,98 @@
 _lora_seq_stats_cache: dict = {}
 
 
-def num_floating_point_operations(cfg: ConfigContainer, batch_size: int = 1):
-    """Return the number of floating point operations"""
+def vit_flops(
+    batch_size,
+    num_patches,
+    depth,
+    hidden_size,
+    num_heads,
+    intermediate_size,
+    spatial_merge_size=2,
+    out_hidden_size=4096,
+):
+    """Calculate FLOPs for a Vision Transformer (ViT) encoder + patch merger.
+
+    Includes:
+    - ViT transformer layers (bidirectional full attention, not causal)
+    - Patch merger (spatial merge + MLP projection to LLM hidden size)
+
+    Args:
+        batch_size: Batch size.
+        num_patches: Total number of vision patches (before spatial merge).
+        depth: Number of ViT transformer layers.
+        hidden_size: ViT hidden dimension.
+        num_heads: Number of attention heads in ViT.
+        intermediate_size: ViT MLP intermediate dimension.
+        spatial_merge_size: Spatial merge factor (2 means 2x2 merge).
+        out_hidden_size: Output hidden size after patch merger (LLM hidden size).
+
+    Returns:
+        Total training FLOPs (forward * 3 for fwd+bwd).
+    """
+    if num_patches <= 0:
+        return 0
+
+    # ViT Transformer layers (bidirectional attention)
+    per_token_per_layer = (
+        # QKV + O projections: 4 matmuls of h x h => 4 * 2 * h^2 FMA = 8h^2
+        # but standard counting: Q,K,V each h->h (3 * 2h^2) + O h->h (2h^2) = 8h^2
+        8 * hidden_size**2
+        # Attention core (full bidirectional, not causal): QK^T + attn*V
+        # = 2 * 2 * h * num_patches = 4 * h * num_patches
+        + 4 * hidden_size * num_patches
+        # MLP (GELU, 2 matmuls): fc1 h->intermediate + fc2 intermediate->h
+        # = 2 * 2 * h * intermediate = 4 * h * intermediate
+        + 4 * hidden_size * intermediate_size
+    )
+    transformer_flops_val = per_token_per_layer * num_patches * depth
+
+    # Patch Merger: spatial merge (2x2) + MLP projection
+    merge_unit = spatial_merge_size**2
+    merged_hidden = hidden_size * merge_unit  # concatenated hidden dim
+    num_merged_tokens = num_patches // merge_unit if merge_unit > 0 else num_patches
+    merger_flops_val = num_merged_tokens * (
+        2 * merged_hidden * merged_hidden  # fc1: merged_hidden -> merged_hidden
+        + 2 * merged_hidden * out_hidden_size  # fc2: merged_hidden -> out_hidden_size
+    )
+
+    return (transformer_flops_val + merger_flops_val) * batch_size * 3  # 3x for training (fwd + bwd)
+
+
+def num_floating_point_operations(
+    cfg: ConfigContainer,
+    batch_size: int = 1,
+    seqlen_sum: int = None,
+    seqlen_squared_sum: int = None,
+    num_vision_patches: int = 0,
+):
+    """Return the number of floating point operations.
+
+    Args:
+        cfg: Configuration container.
+        batch_size: Batch size.
+        seqlen_sum: Sum of actual sequence lengths across the batch
+            (batch_size * actual_seq_length). When provided, overrides
+            cfg.model.seq_length for more accurate FLOPS estimation with
+            dynamic-length sequences (e.g., VLM with dynamic padding).
+        seqlen_squared_sum: Sum of squared sequence lengths across the batch
+            (batch_size * actual_seq_length^2). Used for attention core FLOPS
+            which scale quadratically with sequence length.
+        num_vision_patches: Total number of vision patches in the batch
+            (before spatial merge). Used to compute ViT encoder FLOPS.
+    """
+    # Compute effective sequence length from actual values or fall back to config.
+    if seqlen_sum is not None and batch_size > 0:
+        effective_seq_length = seqlen_sum / batch_size
+    else:
+        effective_seq_length = cfg.model.seq_length
+        seqlen_sum = batch_size * cfg.model.seq_length
+
+    if seqlen_squared_sum is not None and batch_size > 0:
+        effective_seq_length_sq = seqlen_squared_sum / batch_size
+    else:
+        effective_seq_length_sq = effective_seq_length**2
+
     peft = getattr(cfg, "peft", None)
     is_lora = isinstance(peft, LoRA)
     # If the model provider has a custom TFLOPS calculation method, use it (non-LoRA only).
@@ -387,13 +477,13 @@ def transformer_flops():
                     ## o proj
                     + (cfg.model.num_attention_heads * getattr(cfg.model, "v_head_dim", 64)) * cfg.model.hidden_size
                     ## core attn
-                    + cfg.model.seq_length
+                    + effective_seq_length
                     * (
                         cfg.model.num_attention_heads
                         * (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "qk_pos_emb_head_dim", 0))
                     )
                     / 2
-                    + cfg.model.seq_length * cfg.model.num_attention_heads * getattr(cfg.model, "v_head_dim", 64) / 2
+                    + effective_seq_length * cfg.model.num_attention_heads * getattr(cfg.model, "v_head_dim", 64) / 2
                 )
             )
 
@@ -416,7 +506,7 @@ def transformer_flops():
                     effective_window = window_size[0] + window_size[1] + 1
                 else:
                     effective_window = window_size
-                swa_context = min(effective_window, cfg.model.seq_length)
+                swa_context = min(effective_window, effective_seq_length)
 
                 if window_attn_skip_freq is None:
                     num_swa_layers = num_layers
@@ -433,7 +523,7 @@ def transformer_flops():
                     num_swa_layers = 0
                     num_full_attn_layers = num_layers
 
-                full_core = query_projection_size * cfg.model.seq_length / 2 * 2
+                full_core = query_projection_size * effective_seq_length / 2 * 2
                 swa_core = query_projection_size * swa_context / 2 * 2
 
                 self_attn_term = (
@@ -445,7 +535,7 @@ def transformer_flops():
                     )
                 )
             else:
-                full_core = query_projection_size * cfg.model.seq_length / 2 * 2
+                full_core = query_projection_size * effective_seq_length / 2 * 2
                 self_attn_term = 3 * 2 * num_layers * (proj_per_layer + full_core)
 
         # Handle GDN (Gated DeltaNet) hybrid attention variant.
@@ -525,8 +615,7 @@ def transformer_flops():
             ) + 2 * moe_latent_size
 
         total_floating_point_operations = (
-            batch_size
-            * cfg.model.seq_length
+            seqlen_sum
             * (
                 # MLP
                 3
@@ -556,7 +645,29 @@ def transformer_flops():
                 + 3 * 2 * cfg.model.hidden_size * padded_vocab_size * (mtp_num_layers + 1)
             )
         )
-        return total_floating_point_operations
+        return total_floating_point_operations + _compute_vit_flops()
+
+    def _compute_vit_flops():
+        """Compute ViT encoder FLOPs if vision config is available.
+
+        Note: num_vision_patches is the *total* patches across the batch.
+        ViT attention is per-image (not cross-image), so we compute FLOPs
+        using per-image patch count to get correct quadratic attention scaling.
+        """
+        vision_config = getattr(cfg.model, "vision_config", None)
+        if vision_config is None or num_vision_patches <= 0:
+            return 0
+        patches_per_image = num_vision_patches / batch_size if batch_size > 0 else num_vision_patches
+        return vit_flops(
+            batch_size=batch_size,
+            num_patches=patches_per_image,
+            depth=getattr(vision_config, "depth", 0),
+            hidden_size=getattr(vision_config, "hidden_size", 0),
+            num_heads=getattr(vision_config, "num_heads", 16),
+            intermediate_size=getattr(vision_config, "intermediate_size", 0),
+            spatial_merge_size=getattr(vision_config, "spatial_merge_size", 2),
+            out_hidden_size=getattr(vision_config, "out_hidden_size", cfg.model.hidden_size),
+        )
 
     # Main entrypoint for FLOPs calculation.
     if getattr(cfg.model, "is_hybrid_model", False):
@@ -588,9 +699,9 @@ def transformer_flops():
         )
 
         # Compute hybrid model FLOPs.
-        return hybrid_flops(
+        llm_flops = hybrid_flops(
             batch_size=batch_size,
-            seq_len=cfg.model.seq_length,
+            seq_len=effective_seq_length,
             hidden_size=cfg.model.hidden_size,
             num_attn_layers=num_attn_layers,
             num_mamba_layers=num_mamba_layers,
@@ -626,6 +737,7 @@ def transformer_flops():
             vocab_size=padded_vocab_size,
             mtp_num_layers=mtp_num_layers,
         )
+        return llm_flops + _compute_vit_flops()
     else:
         # Compute standard Transformer model FLOPs.
         return transformer_flops()

src/megatron/bridge/training/utils/train_utils.py

@@ -400,6 +400,7 @@ def training_log(
     pg_collection: Optional[Any] = None,
     log_max_attention_logit: Optional[float] = None,
     loaded_iteration: int = 0,
+    seq_length: Optional[int] = None,
 ) -> bool:
     """Log training stats (losses, learning rate, timings, etc.).
 
@@ -783,7 +784,19 @@ def training_log(
         # Calculate GPU utilization
         num_flops = None
         if hasattr(config.model, "kv_channels") and hasattr(config.model, "num_attention_heads"):
-            num_flops = num_floating_point_operations(config, batch_size)
+            if seq_length is not None:
+                seqlen_sum = batch_size * seq_length
+                seqlen_squared_sum = batch_size * seq_length ** 2
+                # Use accumulated vision patches from forward_step micro-batches,
+                # scaled by dp_size (consistent with train.py main loop).
+                local_vision_patches = getattr(global_state, '_flops_vision_patches', 0)
+                num_vision_patches = local_vision_patches * config.data_parallel_size if local_vision_patches > 0 else 0
+                num_flops = num_floating_point_operations(
+                    config, batch_size, seqlen_sum=seqlen_sum,
+                    seqlen_squared_sum=seqlen_squared_sum, num_vision_patches=num_vision_patches,
+                )
+            else:
+                num_flops = num_floating_point_operations(config, batch_size)
             per_gpu_tf = num_flops / elapsed_time_per_iteration / get_world_size_safe() / 1e12
             print_rank_0(
                 f"Step Time : {elapsed_time_per_iteration:.2f}s GPU utilization: {per_gpu_tf:.1f}MODEL_TFLOP/s/GPU"

src/megatron/bridge/training/vlm_step.py

@@ -430,6 +430,23 @@ def forward_step(
         ) = get_batch(data_iterator, state.cfg, use_mtp, pg_collection=pg_collection)
     timers("batch-generator").stop()
 
+    # Accumulate FLOPS metadata across micro-batches.
+    # Each micro-batch contributes its actual padded seq_length (not cfg.model.seq_length).
+    # train.py resets these before each step and reads accumulated values afterwards.
+    if tokens is not None:
+        mbs = tokens.shape[0]
+        seq_len = tokens.shape[1]
+        state._flops_seqlen_sum = getattr(state, '_flops_seqlen_sum', 0) + mbs * seq_len
+        state._flops_seqlen_sq_sum = getattr(state, '_flops_seqlen_sq_sum', 0) + mbs * seq_len ** 2
+    if visual_inputs is not None:
+        for attr in ("image_grid_thw", "video_grid_thw"):
+            grid = getattr(visual_inputs, attr, None)
+            if grid is not None and grid.numel() > 0:
+                state._flops_vision_patches = (
+                    getattr(state, '_flops_vision_patches', 0)
+                    + int(grid.prod(dim=-1).sum().item())
+                )
+
     forward_args = {
         "input_ids": tokens,
         "position_ids": position_ids,