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,51 @@ 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.
+        # In VPP (interleaved pipeline) mode, forward_step_func is called once per
+        # virtual-stage per microbatch (i.e. num_microbatches * vp_size times), but
+        # the FLOPS formula already accounts for ALL layers in the full model.
+        # Therefore we must divide the accumulator by vp_size to avoid over-counting.
+        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)
+        # Coerce to int — getattr on MagicMock test doubles returns a MagicMock
+        # (not the default), which breaks the numeric comparisons below.
+        if not isinstance(local_seqlen_sum, int):
+            local_seqlen_sum = 0
+        if not isinstance(local_seqlen_sq_sum, int):
+            local_seqlen_sq_sum = 0
+        if not isinstance(num_vision_patches, int):
+            num_vision_patches = 0
+
+        # Correct for VPP over-counting: each microbatch's seqlen is accumulated
+        # once per virtual stage, but FLOPS formula already covers all stages.
+        vp_size = config.model.virtual_pipeline_model_parallel_size
+        if isinstance(vp_size, int) and vp_size > 1:
+            local_seqlen_sum = local_seqlen_sum // vp_size
+            local_seqlen_sq_sum = local_seqlen_sq_sum // vp_size
+            num_vision_patches = num_vision_patches // vp_size
+
+        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 +628,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,114 @@
 _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(
+    cfg: ConfigContainer,
+    batch_size: int,
+    num_patches: int,
+):
+    """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:
+        cfg: Configuration container. ViT hyper-parameters are read from
+            ``cfg.model.vision_config`` (``depth``, ``hidden_size``,
+            ``num_heads``, ``intermediate_size``, ``spatial_merge_size``,
+            ``out_hidden_size``). Passing the whole config keeps the public
+            signature stable as the list of required ViT attributes grows.
+        batch_size: Batch size.
+        num_patches: Per-image number of vision patches (before spatial
+            merge). Callers that track the total patch count across the
+            batch should divide by ``batch_size`` before invoking, because
+            ViT attention is per-image (not cross-image) and scales
+            quadratically with the per-image patch count.
+
+    Returns:
+        Total training FLOPs (forward * 3 for fwd+bwd). Returns 0 when
+        no ``vision_config`` is attached or ``num_patches`` is non-positive.
+    """
+    vision_config = getattr(cfg.model, "vision_config", None)
+    if vision_config is None or num_patches <= 0:
+        return 0
+
+    depth = getattr(vision_config, "depth", 0)
+    hidden_size = getattr(vision_config, "hidden_size", 0)
+    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)
+
+    # 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 = None,
+    seqlen_squared_sum: int | None = 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
+            (sum_i actual_seq_length_i^2). Used for attention core FLOPS
+            which scale quadratically with sequence length; when omitted,
+            falls back to ``batch_size * effective_seq_length^2`` so the
+            result matches the legacy constant-length estimate.
+        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
+
+    # Per-layer attention core FLOPS scale as sum_i(s_i^2), while the outer
+    # formula multiplies every per-layer term by ``seqlen_sum``. To account
+    # for the quadratic scaling (and variance) of core attention we replace
+    # the linear ``effective_seq_length`` factor in core-attn expressions
+    # with ``core_attn_seq_factor``. With ``seqlen_sum * core_attn_seq_factor
+    # == sum_i(s_i^2)`` this reproduces the correct quadratic sum; when the
+    # squared sum is unavailable we fall back to ``effective_seq_length`` so
+    # the result matches the legacy constant-length estimate.
+    if seqlen_squared_sum is not None and seqlen_sum > 0:
+        core_attn_seq_factor = seqlen_squared_sum / seqlen_sum
+    else:
+        core_attn_seq_factor = effective_seq_length
+
     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 +493,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
+                    + core_attn_seq_factor
                     * (
                         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
+                    + core_attn_seq_factor * cfg.model.num_attention_heads * getattr(cfg.model, "v_head_dim", 64) / 2
                 )
             )
 
@@ -416,7 +522,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 +539,9 @@ 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 attention is quadratic in seq_len -> use core_attn_seq_factor.
+                # SWA core is bounded by window_size, so keep the averaged bound.
+                full_core = query_projection_size * core_attn_seq_factor / 2 * 2
                 swa_core = query_projection_size * swa_context / 2 * 2
 
                 self_attn_term = (
@@ -445,7 +553,7 @@ def transformer_flops():
                     )
                 )
             else:
-                full_core = query_projection_size * cfg.model.seq_length / 2 * 2
+                full_core = query_projection_size * core_attn_seq_factor / 2 * 2
                 self_attn_term = 3 * 2 * num_layers * (proj_per_layer + full_core)
 
         # Handle GDN (Gated DeltaNet) hybrid attention variant.
@@ -524,39 +632,49 @@ def transformer_flops():
                 / cfg.model.hidden_size
             ) + 2 * moe_latent_size
 
-        total_floating_point_operations = (
-            batch_size
-            * cfg.model.seq_length
+        total_floating_point_operations = seqlen_sum * (
+            # MLP
+            3
+            * 2
+            * cfg.model.hidden_size
             * (
-                # MLP
-                3
-                * 2
-                * cfg.model.hidden_size
-                * (
-                    # dense layers
-                    (cfg.model.ffn_hidden_size * ffn_expansion_factor) * num_dense_layers
-                    # routed experts
-                    + routed_expert_term * num_moe_layers
-                    # Shared Experts.
-                    + (shared_expert_ffn_hidden_size * ffn_expansion_factor) * num_moe_layers
-                )
-                # Self Attention
-                + self_attn_term
-                # MTP norms and proj
-                + 3
-                * 2
-                * mtp_num_layers
-                * (
-                    # MTP eh norm + final norm
-                    3 * cfg.model.hidden_size
-                    # MTP eh proj
-                    + 2 * cfg.model.hidden_size * cfg.model.hidden_size
-                )
-                # Logit.
-                + 3 * 2 * cfg.model.hidden_size * padded_vocab_size * (mtp_num_layers + 1)
+                # dense layers
+                (cfg.model.ffn_hidden_size * ffn_expansion_factor) * num_dense_layers
+                # routed experts
+                + routed_expert_term * num_moe_layers
+                # Shared Experts.
+                + (shared_expert_ffn_hidden_size * ffn_expansion_factor) * num_moe_layers
+            )
+            # Self Attention
+            + self_attn_term
+            # MTP norms and proj
+            + 3
+            * 2
+            * mtp_num_layers
+            * (
+                # MTP eh norm + final norm
+                3 * cfg.model.hidden_size
+                # MTP eh proj
+                + 2 * cfg.model.hidden_size * cfg.model.hidden_size
             )
+            # Logit.
+            + 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 convert to
+        per-image patch count before invoking ``vit_flops`` to get the
+        correct quadratic attention scaling. ``vit_flops`` itself returns
+        0 when ``cfg.model.vision_config`` is absent.
+        """
+        if 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(cfg, batch_size, patches_per_image)
 
     # Main entrypoint for FLOPs calculation.
     if getattr(cfg.model, "is_hybrid_model", False):
@@ -588,9 +706,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 +744,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()