Add vision flops calculation for vlm

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

@@ -34,6 +34,101 @@ def num_floating_point_operations(cfg: ConfigContainer, batch_size: int = 1):
     if not is_lora and hasattr(cfg.model, "_get_num_floating_point_operations"):
         return cfg.model._get_num_floating_point_operations(batch_size)
 
+    def vision_floating_point_operations(cfg: ConfigContainer, batch_size: int) -> float:
+        """Estimate FLOPs for the vision tower + projector (patch merger).
+
+        This complements the LM FLOPs estimate with a *vision-side* estimate for VLMs (e.g.,
+        Qwen3-VL). Vision FLOPs depend on runtime inputs (image/video resolution, #frames),
+        so this function prefers runtime statistics populated during `vlm_step.forward_step`.
+        """
+        vision_cfg = getattr(cfg.model, "vision_config", None)
+        if vision_cfg is None:
+            return 0.0
+
+        depth = getattr(vision_cfg, "depth", None)
+        d_model = getattr(vision_cfg, "hidden_size", None)
+        d_ff = getattr(vision_cfg, "intermediate_size", None)
+        n_heads = getattr(vision_cfg, "num_heads", None)
+        spatial_merge_size = getattr(vision_cfg, "spatial_merge_size", None)
+        out_hidden_size = getattr(vision_cfg, "out_hidden_size", None)
+        num_position_embeddings = getattr(vision_cfg, "num_position_embeddings", None)
+        in_channels = getattr(vision_cfg, "in_channels", 3)
+        patch_size = getattr(vision_cfg, "patch_size", None)
+        temporal_patch_size = getattr(vision_cfg, "temporal_patch_size", 1)
+
+        critical = [
+            depth,
+            d_model,
+            d_ff,
+            n_heads,
+            spatial_merge_size,
+            out_hidden_size,
+            num_position_embeddings,
+        ]
+        if any(v is None for v in critical):
+            return 0.0
+
+        merge_sq = int(spatial_merge_size) ** 2
+        n_pre_runtime = getattr(cfg, "_runtime_vision_tokens_pre_per_sample", None)
+        n_post_runtime = getattr(cfg, "_runtime_vision_tokens_post_per_sample", None)
+        if n_pre_runtime is not None:
+            n_pre = float(n_pre_runtime)
+            n_post = float(n_post_runtime) if n_post_runtime is not None else n_pre / float(max(1, merge_sq))
+        else:
+            # Config-only fallback: assume one image per sample at the max grid size.
+            n_pre = float(num_position_embeddings)
+            n_post = n_pre / float(max(1, merge_sq))
+
+        # Vision attention in Qwen3-VL is packed per-frame (each frame is a separate sequence
+        # of length h*w). If runtime statistics are provided, prefer using sum(L^2) instead
+        # of (sum L)^2 for the attention matmul term.
+        sum_seqlen_sq_pre_runtime = getattr(cfg, "_runtime_vision_sum_seqlen_sq_pre_per_sample", None)
+
+        # FLOPs are counted as 2 * MACs for GEMM/conv.
+        # For training, we approximate backward as 2x forward -> total ≈ 3x forward.
+        fwd_to_train_multiplier = 3.0
+
+        # Patch embedding conv3d FLOPs:
+        # FLOPs ≈ 2 * n_pre * d_model * (C * kt * kh * kw)
+        flops_patch_embed_fwd = 0.0
+        if patch_size is not None and temporal_patch_size is not None:
+            kt = int(temporal_patch_size)
+            kh = int(patch_size)
+            kw = int(patch_size)
+            flops_patch_embed_fwd = 2.0 * n_pre * float(d_model) * (int(in_channels) * kt * kh * kw)
+
+        # Vision transformer block FLOPs (ViT-style, rough but standard):
+        # - QKV projections + output projection: 8 * N * D^2
+        # - Attention score + value: 4 * sum(L_i^2) * D (or 4 * N^2 * D if not packed)
+        # - MLP: 4 * N * D * D_ff
+        n = float(n_pre)
+        if sum_seqlen_sq_pre_runtime is not None and float(sum_seqlen_sq_pre_runtime) > 0:
+            attn_quad = 4.0 * float(sum_seqlen_sq_pre_runtime) * float(d_model)
+        else:
+            attn_quad = 4.0 * (n**2) * float(d_model)
+
+        flops_vit_layer_fwd = (8.0 * n * (float(d_model) ** 2)) + attn_quad + (4.0 * n * float(d_model) * float(d_ff))
+        flops_vit_fwd = float(depth) * flops_vit_layer_fwd
+
+        # Patch merger / projector FLOPs:
+        # First projection: (D*merge^2 -> D*merge^2)
+        # Second projection: (D*merge^2 -> out_hidden_size)
+        d_merge = float(d_model) * float(merge_sq)
+        flops_merger_fwd = (2.0 * float(n_post) * d_merge * d_merge) + (
+            2.0 * float(n_post) * d_merge * float(out_hidden_size)
+        )
+
+        # Deepstack mergers: approximate each as a patch merger.
+        deepstack_idxs = getattr(vision_cfg, "deepstack_visual_indexes", None)
+        num_deepstack = len(deepstack_idxs) if isinstance(deepstack_idxs, (list, tuple)) else 0
+        flops_deepstack_mergers_fwd = float(num_deepstack) * flops_merger_fwd
+
+        flops_vision_fwd_per_sample = (
+            flops_patch_embed_fwd + flops_vit_fwd + flops_merger_fwd + flops_deepstack_mergers_fwd
+        )
+        flops_vision_fwd = float(batch_size) * float(flops_vision_fwd_per_sample)
+        return flops_vision_fwd * fwd_to_train_multiplier
+
     def calculate_layer_counts():
         """Calculate the number of attention, Mamba, MLP, and MoE layers."""
         if hasattr(cfg.model, "hybrid_layer_pattern") and cfg.model.hybrid_layer_pattern:
@@ -389,39 +484,71 @@ def transformer_flops():
             logging_enabled=False,
         )
 
-        total_floating_point_operations = (
-            batch_size
-            * cfg.model.seq_length
+        # ---------------------------------------------------------------------
+        # Runtime token statistics (packed / padded)
+        # ---------------------------------------------------------------------
+        tokens_padded_per_sample = getattr(cfg, "_runtime_lm_total_tokens_padded_per_sample", None)
+        sum_seqlen_sq_padded_per_sample = getattr(cfg, "_runtime_lm_sum_seqlen_sq_padded_per_sample", None)
+
+        if tokens_padded_per_sample is None or float(tokens_padded_per_sample) <= 0:
+            tokens_padded_per_sample = float(cfg.model.seq_length)
+        tokens_total = float(batch_size) * float(tokens_padded_per_sample)
+
+        if sum_seqlen_sq_padded_per_sample is None or float(sum_seqlen_sq_padded_per_sample) <= 0:
+            sum_seqlen_sq_total = float(batch_size) * float(cfg.model.seq_length) * float(cfg.model.seq_length)
+        else:
+            sum_seqlen_sq_total = float(batch_size) * float(sum_seqlen_sq_padded_per_sample)
+
+        # MLP (linear in tokens_total).
+        mlp_total = (
+            tokens_total
+            * expansion_factor
+            * float(num_layers)
+            * float(cfg.model.hidden_size)
             * (
-                # MLP
-                expansion_factor
-                * num_layers
-                * cfg.model.hidden_size
-                * (
-                    # dense layer (deepseek v2, v3 style)
-                    (cfg.model.ffn_hidden_size * gated_linear_multiplier) * (num_dense_layers / num_layers)
-                    # routed experts
-                    + (moe_ffn_hidden_size * num_experts_routed_to * gated_linear_multiplier)
-                    * (num_moe_layers / num_layers)
-                    # Shared Experts.
-                    + (shared_expert_ffn_hidden_size * gated_linear_multiplier) * (num_moe_layers / num_layers)
-                )
-                # Self Attention
-                + self_attn_term
-                # MTP norms and proj
-                + 3
-                * 2
-                * mtp_num_layers
-                * (
-                    # MTP eh norm + final nrom
-                    3 * cfg.model.hidden_size
-                    # MTH 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)
+                (float(cfg.model.ffn_hidden_size) * float(gated_linear_multiplier)) * (float(num_dense_layers) / float(num_layers))
+                + (float(moe_ffn_hidden_size) * float(num_experts_routed_to) * float(gated_linear_multiplier))
+                * (float(num_moe_layers) / float(num_layers))
+                + (float(shared_expert_ffn_hidden_size) * float(gated_linear_multiplier)) * (float(num_moe_layers) / float(num_layers))
+            )
+        )
+
+        # MTP norms and proj (linear in tokens_total).
+        mtp_total = (
+            tokens_total
+            * 3
+            * 2
+            * float(mtp_num_layers)
+            * (
+                3 * float(cfg.model.hidden_size)
+                + 2 * float(cfg.model.hidden_size) * float(cfg.model.hidden_size)
             )
         )
+
+        # Logit (linear in tokens_total).
+        logit_total = tokens_total * 3 * 2 * float(cfg.model.hidden_size) * float(padded_vocab_size) * float(mtp_num_layers + 1)
+
+        # Self-attention:
+        # - projection and output terms scale with sum(L_i)
+        # - attention matmul terms scale with sum(L_i^2) (causal uses ~half, reflected by /2)
+        if cfg.model.multi_latent_attention:
+            # Keep the existing MLA approximation (uses cfg.model.seq_length).
+            self_attn_total = float(batch_size) * float(cfg.model.seq_length) * float(self_attn_term)
+        else:
+            attn_linear_factor = 1.0 + (float(num_query_groups) / float(cfg.model.num_attention_heads))
+            attn_base = (
+                expansion_factor
+                * float(num_layers)
+                * float(cfg.model.hidden_size)
+                * float(cfg.model.hidden_size)
+                * float(query_projection_to_hidden_size_ratio)
+            )
+            self_attn_total = attn_base * (
+                attn_linear_factor * tokens_total + (sum_seqlen_sq_total / (2.0 * float(cfg.model.hidden_size)))
+            )
+
+        total_floating_point_operations = mlp_total + self_attn_total + mtp_total + logit_total
+        total_floating_point_operations += vision_floating_point_operations(cfg, batch_size)
         return total_floating_point_operations
 
     # Main entrypoint for FLOPs calculation.
@@ -454,7 +581,7 @@ def transformer_flops():
         )
 
         # Compute hybrid model FLOPs.
-        return hybrid_flops(
+        total_floating_point_operations = hybrid_flops(
             batch_size=batch_size,
             seq_len=cfg.model.seq_length,
             hidden_size=cfg.model.hidden_size,
@@ -486,6 +613,8 @@ def transformer_flops():
             vocab_size=padded_vocab_size,
             mtp_num_layers=mtp_num_layers,
         )
+        total_floating_point_operations += vision_floating_point_operations(cfg, batch_size)
+        return total_floating_point_operations
     else:
         # Compute standard Transformer model FLOPs.
         return transformer_flops()

src/megatron/bridge/training/vlm_step.py

@@ -411,6 +411,81 @@ def forward_step(
         }
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_params)
 
+    # -------------------------------------------------------------------------
+    # Runtime FLOPs support for packed / padded language tokens
+    # -------------------------------------------------------------------------
+    # `flop_utils.num_floating_point_operations(cfg, batch_size)` only receives `(cfg, batch_size)`.
+    # Attention compute depends on the actual per-sample lengths in the microbatch:
+    #   total tokens: sum_i L_i
+    #   attention matmul: sum_i L_i^2 (causal uses ~half, handled in flop_utils)
+    #
+    # We compute per-sample averages here and stash them on the config for more accurate FLOPs.
+    microbatch_size_for_stats = None
+    try:
+        packed = forward_args.get("packed_seq_params", None)
+        cu = None
+        if packed is not None:
+            cu = getattr(packed, "cu_seqlens_q_padded", None)
+            if cu is None:
+                cu = getattr(packed, "cu_seqlens_q", None)
+
+        if cu is not None:
+            lengths = (cu[1:] - cu[:-1]).to(torch.int64)
+            denom = max(1, int(lengths.numel()))
+            total_tokens_padded = int(cu[-1].item())
+            sum_seqlen_sq_padded = int((lengths * lengths).sum().item())
+
+            state.cfg._runtime_lm_total_tokens_padded_per_sample = float(total_tokens_padded) / float(denom)
+            state.cfg._runtime_lm_sum_seqlen_sq_padded_per_sample = float(sum_seqlen_sq_padded) / float(denom)
+            microbatch_size_for_stats = denom
+        else:
+            # Non-packed path: use the padded sequence length directly.
+            seq_len_padded = int(tokens.shape[1])
+            state.cfg._runtime_lm_total_tokens_padded_per_sample = float(seq_len_padded)
+            state.cfg._runtime_lm_sum_seqlen_sq_padded_per_sample = float(seq_len_padded * seq_len_padded)
+            microbatch_size_for_stats = int(tokens.shape[0]) if tokens is not None and tokens.dim() >= 2 else 1
+    except Exception:
+        # FLOPs logging must never break training.
+        microbatch_size_for_stats = int(tokens.shape[0]) if tokens is not None and tokens.dim() >= 2 else 1
+
+    microbatch_size_for_stats = max(1, int(microbatch_size_for_stats or 1))
+
+    # -------------------------------------------------------------------------
+    # Runtime FLOPs support for vision-language models
+    # -------------------------------------------------------------------------
+    # Vision FLOPs depend on the *actual* number of visual tokens (image/video size and frames).
+    # Stash per-sample averages on the config so flop_utils can use them later.
+    vision_cfg = getattr(state.cfg.model, "vision_config", None)
+    if vision_cfg is not None:
+        image_grid_thw = forward_args.get("image_grid_thw", None)
+        video_grid_thw = forward_args.get("video_grid_thw", None)
+
+        vision_tokens_pre_total = 0
+        vision_sum_seqlen_sq_pre_total = 0
+
+        if image_grid_thw is not None:
+            vision_tokens_pre_total += int(image_grid_thw.prod(dim=-1).sum().item())
+            hw = image_grid_thw[:, 1].to(torch.int64) * image_grid_thw[:, 2].to(torch.int64)
+            t = image_grid_thw[:, 0].to(torch.int64)
+            vision_sum_seqlen_sq_pre_total += int((t * hw * hw).sum().item())
+
+        if video_grid_thw is not None:
+            vision_tokens_pre_total += int(video_grid_thw.prod(dim=-1).sum().item())
+            hw = video_grid_thw[:, 1].to(torch.int64) * video_grid_thw[:, 2].to(torch.int64)
+            t = video_grid_thw[:, 0].to(torch.int64)
+            vision_sum_seqlen_sq_pre_total += int((t * hw * hw).sum().item())
+
+        vision_tokens_pre_per_sample = float(vision_tokens_pre_total) / float(microbatch_size_for_stats)
+        vision_sum_seqlen_sq_pre_per_sample = float(vision_sum_seqlen_sq_pre_total) / float(microbatch_size_for_stats)
+
+        spatial_merge_size = getattr(vision_cfg, "spatial_merge_size", None)
+        merge_sq = int(spatial_merge_size) ** 2 if spatial_merge_size is not None else 1
+        vision_tokens_post_per_sample = vision_tokens_pre_per_sample / float(max(1, merge_sq))
+
+        state.cfg._runtime_vision_tokens_pre_per_sample = vision_tokens_pre_per_sample
+        state.cfg._runtime_vision_tokens_post_per_sample = vision_tokens_post_per_sample
+        state.cfg._runtime_vision_sum_seqlen_sq_pre_per_sample = vision_sum_seqlen_sq_pre_per_sample
+
     check_for_nan_in_loss = state.cfg.rerun_state_machine.check_for_nan_in_loss
     check_for_spiky_loss = state.cfg.rerun_state_machine.check_for_spiky_loss
     with straggler_timer: