`bridge.training.utils.flop_utils`#

Module Contents#

Classes#

GlobalFlopsRuntimeStats

Data-parallel-global FLOPS statistics collected during one training step.

Functions#

`_is_gated_delta_net_variant`	Recognize GDN variants across current main and older MCore dev branches.
`_packed_data_exists`	Return True if a packed dataset file exists, for both parquet and npy formats.
`get_model_chunk_vp_stage`	Return the virtual-pipeline stage assigned to a model chunk, if any.
`_accumulator_to_int`	Coerce a FLOPs accumulator (`int` or scalar `Tensor`) to `int`.
`resolve_global_flops_runtime_stats`	Resolve all data-parallel-global FLOPS statistics used by training.
`resolve_global_flops_seqlen_stats`	Resolve sequence statistics and the legacy total vision-patch count.
`_add_flops_accumulator`	Add an int or scalar tensor to a state accumulator.
`_scalar_sum_for_accumulator`	Return a scalar sum without forcing a CUDA host sync inside forward_step.
`_real_subseq_lengths`	Extract sub-sequence lengths from cu_seqlens metadata.
`accumulate_flops_metadata`	Accumulate per-microbatch FLOPS metadata onto `state`.
`_get_vision_config`	Return a direct or thinker-nested vision config when available.
`vit_flops_from_patch_stats`	Calculate ViT FLOPS from additive per-attention-sequence statistics.
`vision_patch_stats_from_grid_thw`	Build additive Qwen ViT patch statistics without a device-to-host sync.
`vit_flops_from_grid_thw`	Calculate exact ViT FLOPS for Qwen-style temporal-height-width grids.
`vit_flops`	Calculate FLOPs for a Vision Transformer (ViT) encoder + patch merger.
`num_floating_point_operations`	Return the number of floating point operations.

Data#

`_lora_seq_stats_cache`
`_mcore_is_gated_delta_net_variant`

API#

bridge.training.utils.flop_utils._lora_seq_stats_cache: dict#: None

bridge.training.utils.flop_utils._mcore_is_gated_delta_net_variant#: ‘cast(…)’

bridge.training.utils.flop_utils._is_gated_delta_net_variant( experimental_attention_variant: str | None, ) → bool#: Recognize GDN variants across current main and older MCore dev branches.

class bridge.training.utils.flop_utils.GlobalFlopsRuntimeStats#

Data-parallel-global FLOPS statistics collected during one training step.

.. attribute:: seqlen_sum

Total padded language tokens, or None when unavailable.

.. attribute:: seqlen_squared_sum

Sum of squared language subsequence lengths.

.. attribute:: num_vision_patches

Legacy aggregate vision-patch count.

.. attribute:: vision_patch_sum

Exact sum of independent vision patch counts.

.. attribute:: vision_patch_squared_sum

Exact sum of squared independent patch counts.

.. attribute:: vision_merged_token_sum

Exact sum of post-merger vision tokens.

.. attribute:: cross_seqlen_sum

Total cross-attention key/value length.

.. attribute:: cross_seqlen_product_sum

Sum of query/key-value length products.

seqlen_sum: int | None#: None

seqlen_squared_sum: int | None#: None

num_vision_patches: int#: 0

vision_patch_sum: int#: 0

vision_patch_squared_sum: int#: 0

vision_merged_token_sum: int#: 0

cross_seqlen_sum: int | None#: None

cross_seqlen_product_sum: int | None#: None

property has_exact_vision_stats: bool#: Return whether exact additive vision-patch statistics were collected.

bridge.training.utils.flop_utils._packed_data_exists(path: str | None) → bool#

Return True if a packed dataset file exists, for both parquet and npy formats.

Parquet specs may be a single file, a glob, or a directory, so they are detected with is_packed_parquet_spec and validated via the packed-parquet resolver rather than a bare Path.exists() check (which would fail for globs/directories and silently disable LoRA-aware FLOP accounting). This mirrors the format detection in calculate_avg_seqlen so a spec that passes this gate also reads correctly.

bridge.training.utils.flop_utils.get_model_chunk_vp_stage(model: torch.nn.Module) → int | None#

Return the virtual-pipeline stage assigned to a model chunk, if any.

Parameters:: model – Model chunk, possibly wrapped by mixed precision or DDP.
Returns:: The integer virtual-pipeline stage, or None for an unchunked model or a model that does not expose the stage.

bridge.training.utils.flop_utils._accumulator_to_int(value) → int#: Coerce a FLOPs accumulator (int or scalar Tensor) to int.

bridge.training.utils.flop_utils.resolve_global_flops_runtime_stats( state, *, data_parallel_size: int, vp_size: int | None = None, dp_group=None, include_vision_patch_stats: bool = False, include_cross_attention_stats: bool = False, ) → bridge.training.utils.flop_utils.GlobalFlopsRuntimeStats#

Resolve all data-parallel-global FLOPS statistics used by training.

Reads the accumulators populated by the forward step (_flops_seqlen_sum = Σ padded tokens, _flops_seqlen_sq_sum = Σᵢ sᵢ² over real sub-sequences, _flops_vision_patches, optional exact additive ViT patch statistics, and optional cross-attention key/value and query-key products) and reduces them to global totals across the data-parallel group.

Under variable-length (THD packed) training the per-rank Σᵢ sᵢ² can differ across DP ranks, so a single SUM all-reduce over dp_group is used to get the exact global sum. Dense BSHD training never requests this reduce: every DP rank contributes the same fixed sequence statistics, so extrapolating local * data_parallel_size is exact and avoids an unnecessary collective.

Parameters:

state – Object carrying the _flops_* accumulators (GlobalState).
data_parallel_size – Size of the data-parallel group (used for the extrapolation fallback).
vp_size – Virtual pipeline size. Kept for call-site compatibility; VPP does not rescale these accumulators because they already represent the executed training step consumed by the full-model FLOPS formula.
dp_group – Data-parallel process group to SUM-reduce over. Must be the pure DP group (excluding CP) matching data_parallel_size — CP ranks share the same cu_seqlens and would double-count.
include_vision_patch_stats – Whether the collective includes the three exact additive ViT values. This must be uniform across every rank in dp_group.
include_cross_attention_stats – Whether the collective includes the two optional cross-attention values. This must be uniform across every rank in dp_group; callers should derive it from model capability, not local batch contents.

Returns:

Global statistics with sequence values set to None when no corresponding accumulation happened and vision values set to 0 when no matching metadata was accumulated.

bridge.training.utils.flop_utils.resolve_global_flops_seqlen_stats( state, *, data_parallel_size: int, vp_size: int | None = None, dp_group=None, ) → tuple[int | None, int | None, int]#

Resolve sequence statistics and the legacy total vision-patch count.

This compatibility wrapper preserves the established three-value return contract and its three-integer collective.

bridge.training.utils.flop_utils._add_flops_accumulator(state, name: str, delta) → None#: Add an int or scalar tensor to a state accumulator.

bridge.training.utils.flop_utils._scalar_sum_for_accumulator(value: torch.Tensor) → int | torch.Tensor#: Return a scalar sum without forcing a CUDA host sync inside forward_step.

bridge.training.utils.flop_utils._real_subseq_lengths( cu_seqlens: torch.Tensor | None, cu_seqlens_argmin: torch.Tensor | None = None, cu_seqlens_unpadded: torch.Tensor | None = None, cu_seqlens_unpadded_argmin: torch.Tensor | None = None, ) → torch.Tensor | None#

Extract sub-sequence lengths from cu_seqlens metadata.

Prefers cu_seqlens_unpadded (true sub-sequence boundaries when pad_seq_to_mult > 1) over the padded cu_seqlens. Truncates by the corresponding *_argmin when provided. Returns None when no cu_seqlens info is available.

Runs once per micro-batch, so it must stay free of GPU→CPU syncs: cu_seqlens is a (monotonic non-decreasing) cumulative sum, so the diffs are always >= 0 and we do not filter them — a boolean mask like sub_seq_lens[sub_seq_lens > 0] would force a data-dependent-size device sync every micro-batch (the cause of a ~7% throughput regression). Zero-length entries (padding) contribute 0 to Σᵢ sᵢ² so dropping them is unnecessary; the result is identical.

Accumulate per-microbatch FLOPS metadata onto state.

Under interleaved pipeline parallelism, the forward step runs once per virtual model chunk for the same logical data microbatch. Only virtual stage 0 contributes metadata so model chunking does not multiply the full-model FLOPS estimate. None and 0 both represent the primary/only chunk.

Writes accumulators consumed by train.py at end of step:

_flops_seqlen_sum: mbs * tokens.shape[1] (padded total tokens this microbatch contributes), or mbs * config_seq_len for dense non-packed batches whose tensors were already context-parallel sliced. Drives the linear MLP/proj/logit terms.
_flops_seqlen_sq_sum: the THD attention term Σᵢ sᵢ², computed inline from cu_seqlens (preferring cu_seqlens_unpadded). The per-pack sub-sequence lengths are reduced via :func:_scalar_sum_for_accumulator, which keeps the result on-device (no .item()) — so the per-microbatch path stays sync-free and the single host sync happens once per step in

func:

resolve_global_flops_seqlen_stats. When cu_seqlens is absent (dense / non-packed) or degenerate, the host-int BSHD fallback mbs * dense_seq_len² is accumulated instead (bit-exact with the pre-fix value). dense_seq_len is config_seq_len when provided, otherwise tokens.shape[1].
_flops_vision_patches: legacy total patch-count approximation.
_flops_vision_patch_sum, _flops_vision_patch_sq_sum, and _flops_vision_merged_token_sum: exact additive ViT statistics that preserve independent media/frame attention boundaries.
_flops_cross_seqlen_sum and _flops_cross_seqlen_product_sum: optional cross-attention Σᵢ kᵢ and Σᵢ qᵢkᵢ terms for model-specific estimators such as WAN.

num_vision_patches remains supported for model callers that only expose a total patch count. vision_patch_stats is the exact (Σp, Σp², Σmerged_tokens) tuple computed by a caller that knows the encoder’s attention boundaries. Each element may be an int or scalar Tensor; device tensors avoid a host sync here. This argument also opts into a conditional DP collective: if one rank supplies it, every rank in that pure DP group must supply a tuple, using (0, 0, 0) on ranks with no media.

For THD packed training (offline packed LLM SFT or VLM in-batch packing), treating the whole pack as one length-seq_len sequence over-counts attention FLOPS by a large factor: actual attention work is Σᵢ sᵢ², not (Σᵢ sᵢ)². Using cu_seqlens here closes that gap.

bridge.training.utils.flop_utils._get_vision_config( cfg: megatron.bridge.training.config.ConfigContainer, )#: Return a direct or thinker-nested vision config when available.

bridge.training.utils.flop_utils.vit_flops_from_patch_stats( cfg: megatron.bridge.training.config.ConfigContainer, *, patch_sum: int | torch.Tensor, patch_squared_sum: int | torch.Tensor, merged_token_sum: int | torch.Tensor, ) → int | torch.Tensor#: Calculate ViT FLOPS from additive per-attention-sequence statistics.

bridge.training.utils.flop_utils.vision_patch_stats_from_grid_thw( grid_thw: torch.Tensor, *, spatial_merge_size: int, ) → tuple[int | torch.Tensor, int | torch.Tensor, int | torch.Tensor]#

Build additive Qwen ViT patch statistics without a device-to-host sync.

Parameters:

grid_thw – Integer tensor shaped [..., 3] with temporal, height, and width patch-grid dimensions.
spatial_merge_size – Vision encoder’s spatial merge size.

Returns:

(patch_sum, patch_squared_sum, merged_token_sum). CUDA inputs produce scalar CUDA tensors; CPU inputs produce Python integers.

Raises:

ValueError – If the grid shape or spatial merge size is invalid.

bridge.training.utils.flop_utils.vit_flops_from_grid_thw( cfg: megatron.bridge.training.config.ConfigContainer, grid_thw: torch.Tensor, ) → int | torch.Tensor#

Calculate exact ViT FLOPS for Qwen-style temporal-height-width grids.

Qwen vision attention treats every temporal frame as an independent THD sequence. For a grid row [t, h, w], linear ViT work therefore sees t*h*w patches while full-attention work sees t*(h*w)^2 rather than (t*h*w)^2. Computing the additive statistics directly avoids expanding frame lengths with repeat_interleave and keeps the microbatch path free of device-to-host synchronization.

Parameters:

cfg – Configuration container with a direct or thinker-nested vision config.
grid_thw – Integer tensor shaped [..., 3]. Each row contains temporal frames, patch-grid height, and patch-grid width.

Returns:

Scalar ViT training FLOPS on the same device as grid_thw. Returns 0 if no vision config or no media rows are present.

Raises:

ValueError – If grid_thw is not shaped [..., 3].

bridge.training.utils.flop_utils.vit_flops( cfg: megatron.bridge.training.config.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)

Parameters:

cfg – Configuration container. ViT hyper-parameters are read from the direct or thinker-nested 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.

bridge.training.utils.flop_utils.num_floating_point_operations( cfg: megatron.bridge.training.config.ConfigContainer, batch_size: int = 1, seqlen_sum: int | None = None, seqlen_squared_sum: int | None = None, num_vision_patches: int = 0, cross_seqlen_sum: int | None = None, cross_seqlen_product_sum: int | None = None, )#

Return the number of floating point operations.

Parameters:

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.
cross_seqlen_sum – Sum of cross-attention key/value sequence lengths.
cross_seqlen_product_sum – Sum of per-sample query and key/value length products.

bridge.training.utils.flop_utils#

Module Contents#

Classes#

Functions#

Data#

API#

`bridge.training.utils.flop_utils`#