[training, perf] fix: THD-aware FLOPS via cu_seqlens (Σᵢ sᵢ²)

src/megatron/bridge/models/qwen_omni/qwen3_omni_step.py

@@ -28,6 +28,7 @@
 from megatron.bridge.training.losses import (
     create_masked_next_token_loss_function as _create_loss_function,
 )
+from megatron.bridge.training.utils.flop_utils import accumulate_flops_metadata
 from megatron.bridge.training.utils.padding_utils import (
     pad_or_truncate_2d_to_len,
     pad_or_truncate_attn_to_len,
@@ -248,6 +249,19 @@ def forward_step(
             seq_length=getattr(config, "seq_length", getattr(state.cfg.model, "seq_length", None)),
         )
 
+    # Accumulate FLOPS metadata across micro-batches. Qwen3-Omni does not pack
+    # within a batch, so cu_seqlens is absent and the helper falls back to
+    # BSHD math for the attention term. Vision-patch tracking still applies.
+    # Vision-patch count is model-specific (Qwen reports grid_thw = t*h*w per
+    # image/video); compute it here and pass a scalar to the model-agnostic helper.
+    num_vision_patches = None
+    if isinstance(multimodal_inputs, dict):
+        for grid in (multimodal_inputs.get("image_grid_thw"), multimodal_inputs.get("video_grid_thw")):
+            if grid is not None and grid.numel() > 0:
+                patches = grid.prod(dim=-1).sum()
+                num_vision_patches = patches if num_vision_patches is None else num_vision_patches + patches
+    accumulate_flops_metadata(state, tokens, num_vision_patches=num_vision_patches)
+
     forward_args = {
         "input_ids": tokens,
         "position_ids": position_ids,

src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py

@@ -27,6 +27,7 @@
     create_masked_next_token_loss_function as _create_loss_function,
 )
 from megatron.bridge.training.state import GlobalState
+from megatron.bridge.training.utils.flop_utils import accumulate_flops_metadata
 from megatron.bridge.training.utils.padding_utils import (
     pad_or_truncate_2d_to_len,
     pad_or_truncate_attn_to_len,
@@ -261,6 +262,28 @@ def forward_step(
         force_to_pad_to_seq_len=this_pg_collection.pp.size() > 1 or this_pg_collection.ep.size() > 1,
         seq_length=config.seq_length,
     )
+
+    # Accumulate FLOPS metadata across micro-batches. When in-batch packing is
+    # active, ``packed_seq_params.cu_seqlens_q`` describes the real sub-seq
+    # boundaries used by the THD attention kernel; the helper uses it to
+    # compute the THD-correct Σᵢ sᵢ² instead of pack-length² (BSHD). When not
+    # packed, the helper falls back to BSHD. train.py resets these before each
+    # step and reads accumulated values afterwards.
+    # Vision-patch count is model-specific (Qwen-VL reports grid_thw = t*h*w per
+    # image/video); compute it here and pass a scalar to the model-agnostic helper.
+    num_vision_patches = None
+    if isinstance(multi_modal_inputs, dict):
+        for grid in (multi_modal_inputs.get("image_grid_thw"), multi_modal_inputs.get("video_grid_thw")):
+            if grid is not None and grid.numel() > 0:
+                patches = grid.prod(dim=-1).sum()
+                num_vision_patches = patches if num_vision_patches is None else num_vision_patches + patches
+    accumulate_flops_metadata(
+        state,
+        tokens,
+        cu_seqlens=getattr(packed_seq_params, "cu_seqlens_q", None) if packed_seq_params is not None else None,
+        num_vision_patches=num_vision_patches,
+    )
+
     forward_args = {
         "input_ids": tokens,
         "labels": labels,

src/megatron/bridge/training/gpt_step.py

@@ -42,6 +42,7 @@
 from megatron.bridge.training.losses import masked_next_token_loss
 from megatron.bridge.training.post_training.distillation import loss_func_kd
 from megatron.bridge.training.state import GlobalState
+from megatron.bridge.training.utils.flop_utils import accumulate_flops_metadata
 from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_params
 from megatron.bridge.training.utils.pg_utils import get_pg_collection
 
@@ -353,6 +354,26 @@ def _forward_step_common(
         )
     timers("batch-generator").stop()
 
+    # Accumulate FLOPS metadata across micro-batches. The THD attention term Σᵢ sᵢ² is
+    # derived inline from cu_seqlens (kept on-device, sync-free); see
+    # accumulate_flops_metadata. Falls back to BSHD when cu_seqlens is absent.
+    #
+    # The cu_seqlens-driven THD path is only wired/validated for CP == 1 in this PR.
+    # Under context parallelism the batch (and its cu_seqlens) is CP-partitioned per
+    # rank, so the per-rank Σᵢ sᵢ² accounting here is not yet correct — that is the
+    # follow-up tracked in #4161. Until then, forward cu_seqlens only for CP == 1 so
+    # CP > 1 stays on the BSHD term (the behavior this test passed on before the THD
+    # change), instead of running the not-yet-CP-safe cu_seqlens path.
+    cp_use_thd = pg_collection.cp.size() == 1
+    accumulate_flops_metadata(
+        state,
+        tokens,
+        cu_seqlens=cu_seqlens if cp_use_thd else None,
+        cu_seqlens_argmin=cu_seqlens_argmin if cp_use_thd else None,
+        cu_seqlens_unpadded=cu_seqlens_unpadded if cp_use_thd else None,
+        cu_seqlens_unpadded_argmin=cu_seqlens_unpadded_argmin if cp_use_thd else None,
+    )
+
     forward_args = {
         "input_ids": tokens,
         "position_ids": position_ids,

src/megatron/bridge/training/train.py

@@ -319,6 +319,11 @@ def train(
     print_rank_0(f"Starting training loop at iteration {start_iteration}")
     p2p_communicator = P2PCommunicator(pp_group=pg_collection.pp, config=model_config)
     dp_size = pg_collection.dp.size()
+    # Anchor for interval-average throughput logging: training_log reports the FLOPS
+    # performed over each logging interval as the delta of
+    # floating_point_operations_so_far. Seed it with the current cumulative (0 fresh,
+    # or the checkpoint value on resume) so the first interval's delta is correct.
+    global_state._flops_at_last_log = global_state.train_state.floating_point_operations_so_far
     if hasattr(config.model, "dist_train") and getattr(config.model.dist_train, "use_dist_train", False) is True:
         forward_backward_func = forward_backward_pipelining_without_interleaving
         p2p_communicator = config.model._p2p_communicator
@@ -546,43 +551,18 @@ def train(
             assert num_skipped_samples_in_batch == 0
         global_state.train_state.skipped_train_samples += num_skipped_samples_in_batch
 
-        # 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
-
+        # Resolve this step's data-parallel-global FLOPS sequence stats and fold the
+        # step's FLOPS into the running total. The per-microbatch sub-sequence sums are
+        # accumulated on-device (sync-free) by accumulate_flops_metadata; here we do one
+        # exact SUM all-reduce over the pure DP group (CP ranks share cu_seqlens, so
+        # reducing over DP×CP would double-count) plus one host sync — once per step, at
+        # the existing end-of-step sync boundary alongside the loss all-reduce.
+        seqlen_sum, seqlen_squared_sum, num_vision_patches = flop_utils.resolve_global_flops_seqlen_stats(
+            global_state,
+            data_parallel_size=dp_size,
+            vp_size=config.model.virtual_pipeline_model_parallel_size,
+            dp_group=pg_collection.dp,
+        )
         num_floating_point_operations_in_batch = flop_utils.num_floating_point_operations(
             config,
             batch_size=batch_size,
@@ -632,7 +612,10 @@ def train(
                 model,
                 log_max_attention_logit,
                 loaded_iteration=start_iteration,
-                seq_length=seqlen_sum // batch_size if seqlen_sum else None,
+                # training_log recomputes seqlen/FLOPS from the (log-step) accumulators
+                # itself; pass None so it uses that path (the per-step seqlen_sum is no
+                # longer materialized on the hot path).
+                seq_length=None,
             )
 
         if (