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

src/megatron/bridge/data/datasets/sft.py

@@ -1054,6 +1054,24 @@ def collate_fn(self, batch):
                 cu_seqlens_batch["cu_seqlens_unpadded"] = cu_seqlens_unpadded
                 cu_seqlens_batch["cu_seqlens_unpadded_argmin"] = cu_seqlens_unpadded_argmin
 
+            # Precompute the per-pack THD attention term Σᵢ sᵢ² here in the dataloader
+            # worker (free, overlapped with training) so the training loop never runs the
+            # cu_seqlens torch ops per micro-batch — those stalled CUDA run-ahead by ~7%.
+            # This mirrors flop_utils._real_subseq_lengths exactly (prefer unpadded; argmin
+            # truncation; squared-diff sum) so the FLOPS value is identical.
+            sq_cu = cu_seqlens_batch.get("cu_seqlens_unpadded", cu_seqlens_batch["cu_seqlens"])
+            sq_argmin = cu_seqlens_batch.get("cu_seqlens_unpadded_argmin", cu_seqlens_argmin)
+            seqlen_sq = []
+            for row in range(sq_cu.shape[0]):
+                n = int(sq_argmin[row].item())
+                cu_row = sq_cu[row, :n]
+                if cu_row.numel() < 2:
+                    seqlen_sq.append(0)
+                else:
+                    diffs = (cu_row[1:] - cu_row[:-1]).long()
+                    seqlen_sq.append(int((diffs * diffs).sum().item()))
+            cu_seqlens_batch["seqlen_sq"] = torch.LongTensor(seqlen_sq)
+
             processed_batch.update(cu_seqlens_batch)
         else:
             attention_mask = [self._create_attention_mask(max_length) for _ in batch]

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

@@ -27,6 +27,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,
@@ -181,6 +182,19 @@ def forward_step(
     if pg_collection.cp.size() > 1:
         raise NotImplementedError("Qwen3-Omni training supports SP/EP, but CP is not supported yet.")
 
+    # 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

@@ -32,6 +32,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
 
@@ -90,6 +91,7 @@ def _partition_packed_batch_for_cp(batch: dict[str, torch.Tensor], cp_size: int)
         "cu_seqlens_unpadded",
         "cu_seqlens_argmin",
         "cu_seqlens_unpadded_argmin",
+        "seqlen_sq",
         "max_seqlen",
         "token_count",
     }
@@ -141,6 +143,10 @@ def get_batch_from_iterator(
         required_host_keys.add("max_seqlen")
         if "cu_seqlens_unpadded_argmin" in batch:
             required_host_keys.add("cu_seqlens_unpadded_argmin")
+        # Precomputed per-pack Σᵢ sᵢ² for FLOPS — keep on host (it's a tiny int tensor),
+        # so the training loop adds an int instead of running cu_seqlens torch ops.
+        if "seqlen_sq" in batch:
+            required_host_keys.add("seqlen_sq")
 
     if not include_full_batch_fields:
         if is_first_pp_stage or use_mtp:
@@ -173,6 +179,7 @@ def get_batch(
     torch.Tensor,
     torch.Tensor | None,
     torch.Tensor | None,
+    torch.Tensor | None,
 ]:
     """Generate a batch.
 
@@ -183,16 +190,16 @@ def get_batch(
 
     Returns:
         tuple of tensors containing tokens, labels, loss_mask, attention_mask, position_ids,
-        cu_seqlens, cu_seqlens_argmin, max_seqlen, cu_seqlens_unpadded, and
-        cu_seqlens_unpadded_argmin
+        cu_seqlens, cu_seqlens_argmin, max_seqlen, cu_seqlens_unpadded,
+        cu_seqlens_unpadded_argmin, and seqlen_sq (precomputed per-pack Σᵢ sᵢ² for FLOPS)
     """
     # Determine pipeline stage role via process group collection
     is_first = is_pp_first_stage(pg_collection.pp)
     is_last = is_pp_last_stage(pg_collection.pp)
     is_middle = (not is_first) and (not is_last)
     include_full_batch_fields = is_middle and _middle_pp_stage_needs_batch(cfg)
     if is_middle and not include_full_batch_fields:
-        return None, None, None, None, None, None, None, None, None, None
+        return None, None, None, None, None, None, None, None, None, None, None
 
     batch = get_batch_from_iterator(
         data_iterator,
@@ -224,6 +231,7 @@ def get_batch(
         batch.get("max_seqlen"),
         batch.get("cu_seqlens_unpadded"),
         batch.get("cu_seqlens_unpadded_argmin"),
+        batch.get("seqlen_sq"),
     )
 
 
@@ -261,9 +269,25 @@ def _forward_step_common(
             max_seqlen,
             cu_seqlens_unpadded,
             cu_seqlens_unpadded_argmin,
+            seqlen_sq,
         ) = get_batch(data_iterator, state.cfg, use_mtp, pg_collection=pg_collection)
     timers("batch-generator").stop()
 
+    # Accumulate FLOPS metadata across micro-batches. ``seqlen_sq`` is the per-pack
+    # THD Σᵢ sᵢ² precomputed in the dataloader worker (see GPTSFTPackedDataset.collate_fn),
+    # so the training loop just sums an int — no per-micro-batch cu_seqlens torch ops on
+    # the launch-bound forward path. When it's absent (datasets that don't precompute),
+    # the helper falls back to deriving Σᵢ sᵢ² from cu_seqlens.
+    accumulate_flops_metadata(
+        state,
+        tokens,
+        seqlen_sq=seqlen_sq,
+        cu_seqlens=cu_seqlens,
+        cu_seqlens_argmin=cu_seqlens_argmin,
+        cu_seqlens_unpadded=cu_seqlens_unpadded,
+        cu_seqlens_unpadded_argmin=cu_seqlens_unpadded_argmin,
+    )
+
     forward_args = {
         "input_ids": tokens,
         "position_ids": position_ids,

src/megatron/bridge/training/train.py

@@ -316,6 +316,13 @@ 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()
+    # FLOPS accounting: per-step stats are buffered (sync-free) and drained at log /
+    # checkpoint boundaries by flush_flops_buffer, so the hot path performs no per-step
+    # device→host sync or DP collective. DP context is stashed on state so saves can flush.
+    global_state._flops_buffer = []
+    global_state._flops_since_last_log = 0.0
+    global_state._flops_dp_group = pg_collection.dp
+    global_state._flops_dp_size = dp_size
     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
@@ -437,6 +444,9 @@ def train(
         global_state._flops_seqlen_sum = 0
         global_state._flops_seqlen_sq_sum = 0
         global_state._flops_vision_patches = 0
+        # Per-step cu_seqlens references; the THD Σᵢ sᵢ² is computed from these once
+        # per step in buffer_flops_metadata (off the per-microbatch forward path).
+        global_state._flops_cu_records = []
 
         (
             loss_dict,
@@ -543,53 +553,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
-
-        num_floating_point_operations_in_batch = flop_utils.num_floating_point_operations(
-            config,
+        # Buffer this step's per-microbatch FLOPS stats (sync-free: no .item(), no
+        # collective). The DP reduction and the FLOPS formula are deferred to log /
+        # checkpoint boundaries (flush_flops_buffer), keeping the hot path free of the
+        # per-step device→host sync that stalls CUDA run-ahead. Deferral is exact: each
+        # buffered step is reduced and evaluated as its own row, so ViT / SSM /
+        # sliding-window (per-step-average) terms are unchanged.
+        flop_utils.buffer_flops_metadata(
+            global_state,
             batch_size=batch_size,
-            seqlen_sum=seqlen_sum,
-            seqlen_squared_sum=seqlen_squared_sum,
-            num_vision_patches=num_vision_patches,
+            vp_size=config.model.virtual_pipeline_model_parallel_size,
         )
-        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
 
         # Logging.
         if not config.logger.skip_train_metrics_log:
@@ -629,7 +604,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 (
@@ -687,13 +665,20 @@ def train(
         # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC).
         # Some of these only happen at specific iterations.
         maybe_synchronize_training_step(config.train.train_sync_interval, global_state.train_state.step)
+        # Drain the FLOPS buffer at log boundaries so the cumulative + since-last-log
+        # counters (used for straggler throughput reporting) are current. One reduce +
+        # one host sync per log interval — not per step.
+        if config.logger.log_interval and global_state.train_state.step % config.logger.log_interval == 0:
+            flop_utils.flush_flops_buffer(global_state, config, data_parallel_size=dp_size, dp_group=pg_collection.dp)
+            num_floating_point_operations_so_far = global_state.train_state.floating_point_operations_so_far
         num_floating_point_operations_since_last_log_event = maybe_report_stragglers(
             config.logger.log_interval,
             bool(getattr(config.straggler, "log_straggler", False)),
             straggler_timer,
             global_state.train_state.step,
-            num_floating_point_operations_since_last_log_event,
+            global_state._flops_since_last_log,
         )
+        global_state._flops_since_last_log = num_floating_point_operations_since_last_log_event
         maybe_check_weight_hash_across_dp_replicas(
             model,
             config.train.check_weight_hash_across_dp_replicas_interval,
@@ -750,6 +735,10 @@ def train(
                 callback_manager=callback_manager,
             )
 
+    # Drain any FLOPS buffered after the last flush so the final cumulative is complete
+    # (idempotent; the final save above already flushes when it runs).
+    flop_utils.flush_flops_buffer(global_state, config, data_parallel_size=dp_size, dp_group=pg_collection.dp)
+
     _delete_cuda_graphs(cuda_graph_helper)
 
     # Flush TensorBoard, WandB writers and one-logger.
@@ -1273,6 +1262,18 @@ def save_checkpoint_and_time(
     timers = state.timers
     energy_monitor = state.energy_monitor
 
+    # Drain any buffered per-step FLOPS so the cumulative saved in the checkpoint is
+    # current. flush is idempotent (empty buffer → no-op) and runs only on saves, so it
+    # keeps the per-step hot path sync-free while guaranteeing correctness at save time.
+    flushed = flop_utils.flush_flops_buffer(
+        state,
+        state.cfg,
+        data_parallel_size=getattr(state, "_flops_dp_size", 1),
+        dp_group=getattr(state, "_flops_dp_group", None),
+    )
+    if flushed:
+        num_floating_point_operations_so_far = state.train_state.floating_point_operations_so_far
+
     # Stop timer to get accurate train interval time and exclude checkpointing duration
     timers("interval-time").stop()