From 9b0592519846faa429372dbe2c43b5a6773f1973 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 15 Jun 2026 13:07:22 -0700 Subject: [PATCH 1/5] =?UTF-8?q?[training,=20perf]=20fix:=20THD-aware=20FLO?= =?UTF-8?q?PS=20via=20cu=5Fseqlens=20(=CE=A3=E1=B5=A2=20s=E1=B5=A2=C2=B2)?= =?UTF-8?q?=20(re-land=20of=20#3839)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Re-applies #3839 (reverted in #4363 after an accidental merge) with one fix. #3839's last commit added `cfg.checkpoint.load = None` to the CP+packing functional test (test_sft_example_runs_with_cp_and_packing). With pretrained_checkpoint also None, finetune() then fails its precondition (finetune.py:50) with "Finetuning requires a loading from a pretrained checkpoint or resuming from a checkpoint". This drops that line, restoring the pre-#3839 behavior (inherit the recipe's default load) so the test runs. The `use_distributed_optimizer=False` setting added to that test is kept: it works around an NCCL watchdog hang seen only under the distributed optimizer + context parallelism in this test (root-cause tracked separately; the THD-FLOPS code itself is inert under CP>1, taking the BSHD fallback). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Chen Cui --- .../models/qwen_omni/qwen3_omni_step.py | 14 + .../bridge/models/qwen_vl/qwen3_vl_step.py | 23 ++ src/megatron/bridge/training/gpt_step.py | 21 ++ src/megatron/bridge/training/train.py | 59 ++-- .../bridge/training/utils/flop_utils.py | 216 +++++++++++++- .../bridge/training/utils/train_utils.py | 80 ++--- src/megatron/bridge/training/vlm_step.py | 37 ++- .../training/test_seqpacking_cp_example.py | 11 + .../training/utils/test_flop_utils.py | 277 +++++++++++++++++- .../training/utils/test_train_utils.py | 72 +++++ 10 files changed, 702 insertions(+), 108 deletions(-) diff --git a/src/megatron/bridge/models/qwen_omni/qwen3_omni_step.py b/src/megatron/bridge/models/qwen_omni/qwen3_omni_step.py index e46dc2ac60..7a7c719e5c 100644 --- a/src/megatron/bridge/models/qwen_omni/qwen3_omni_step.py +++ b/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, diff --git a/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py b/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py index 869d9e6afe..9665c526e0 100644 --- a/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py +++ b/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, diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index 48aae49433..bbd604ea6a 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/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, diff --git a/src/megatron/bridge/training/train.py b/src/megatron/bridge/training/train.py index 3796a79728..50df9f2eef 100644 --- a/src/megatron/bridge/training/train.py +++ b/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 ( diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index be46651257..71f2c68d7e 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -15,6 +15,8 @@ import importlib from pathlib import Path +import torch + from megatron.bridge.data.datasets.packing_utils import calculate_avg_seqlen from megatron.bridge.peft.lora import LoRA from megatron.bridge.training.config import ConfigContainer @@ -24,6 +26,213 @@ _lora_seq_stats_cache: dict = {} +def _accumulator_to_int(value) -> int: + """Coerce a FLOPs accumulator (``int`` or scalar ``Tensor``) to ``int``.""" + if isinstance(value, bool): + return int(value) + if isinstance(value, int): + return value + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return 0 + return int(value.detach().cpu().item()) + return 0 + + +def 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 data-parallel-global FLOPS sequence stats from per-rank accumulators. + + Reads the three accumulators populated by the forward step + (``_flops_seqlen_sum`` = Σ padded tokens, ``_flops_seqlen_sq_sum`` = Σᵢ sᵢ² + over real sub-sequences, ``_flops_vision_patches``), corrects for VPP + over-counting, and reduces them to global totals across the data-parallel + group. + + Under variable-length (THD packed) training the per-rank ``Σᵢ sᵢ²`` and + vision-patch counts differ across DP ranks, so a single SUM all-reduce over + ``dp_group`` is used to get the exact global sum. When no process group is + available (single process, ``dp_group is None``, or ``torch.distributed`` not + initialized) the values are extrapolated as ``local * data_parallel_size`` — + exact only when every DP rank sees an identical sequence-length distribution. + + Args: + 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; accumulators are divided by it to undo + the per-virtual-stage over-counting. ``None``/``<= 1`` is a no-op. + 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. + + Returns: + ``(seqlen_sum, seqlen_squared_sum, num_vision_patches)``. The first two + are ``None`` when no accumulation happened, signalling the caller to fall + back to a fixed-length estimate. ``num_vision_patches`` is ``0`` when no + vision tokens were seen. + """ + local_seqlen_sum = _accumulator_to_int(getattr(state, "_flops_seqlen_sum", 0)) + local_seqlen_sq_sum = _accumulator_to_int(getattr(state, "_flops_seqlen_sq_sum", 0)) + local_vision_patches = _accumulator_to_int(getattr(state, "_flops_vision_patches", 0)) + + # VPP correction: forward_step_func runs once per virtual stage per microbatch, + # so each accumulator over-counts by vp_size; the FLOPS formula already covers + # all layers. Done per-rank (before the reduce) to recover each rank's true local. + if isinstance(vp_size, int) and vp_size > 1: + local_seqlen_sum //= vp_size + local_seqlen_sq_sum //= vp_size + local_vision_patches //= vp_size + + use_all_reduce = ( + dp_group is not None + and data_parallel_size > 1 + and torch.distributed.is_available() + and torch.distributed.is_initialized() + ) + if use_all_reduce: + device = torch.cuda.current_device() if torch.cuda.is_available() else None + stats = torch.tensor( + [local_seqlen_sum, local_seqlen_sq_sum, local_vision_patches], + dtype=torch.long, + device=device, + ) + torch.distributed.all_reduce(stats, op=torch.distributed.ReduceOp.SUM, group=dp_group) + seqlen_sum, seqlen_squared_sum, num_vision_patches = (int(x) for x in stats.tolist()) + else: + # No process group: extrapolate from the local rank (approximation). + seqlen_sum = local_seqlen_sum * data_parallel_size + seqlen_squared_sum = local_seqlen_sq_sum * data_parallel_size + num_vision_patches = local_vision_patches * data_parallel_size + + if seqlen_sum <= 0: + return None, None, max(num_vision_patches, 0) + return seqlen_sum, seqlen_squared_sum, max(num_vision_patches, 0) + + +def _add_flops_accumulator(state, name: str, delta) -> None: + """Add an int or scalar tensor to a state accumulator.""" + current = getattr(state, name, 0) + if not isinstance(current, (int, torch.Tensor)): + current = 0 + setattr(state, name, current + delta) + + +def _scalar_sum_for_accumulator(value: torch.Tensor) -> int | torch.Tensor: + """Return a scalar sum without forcing a CUDA host sync inside forward_step.""" + total = value.sum() + if total.device.type == "cuda": + return total + return int(total.item()) + + +def _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. + """ + if cu_seqlens_unpadded is not None: + cu = cu_seqlens_unpadded.squeeze() + argmin = cu_seqlens_unpadded_argmin + elif cu_seqlens is not None: + cu = cu_seqlens.squeeze() + argmin = cu_seqlens_argmin + else: + return None + + if argmin is not None: + cu = cu[: int(argmin.item())] + + if cu.numel() < 2: + return cu.new_empty(0, dtype=torch.long) + + # No boolean mask here on purpose (see docstring): keep this sync-free. + return (cu[1:] - cu[:-1]).long() + + +def accumulate_flops_metadata( + state, + tokens: torch.Tensor | None, + *, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_argmin: torch.Tensor | None = None, + cu_seqlens_unpadded: torch.Tensor | None = None, + cu_seqlens_unpadded_argmin: torch.Tensor | None = None, + num_vision_patches: int | torch.Tensor | None = None, +) -> None: + """Accumulate per-microbatch FLOPS metadata onto ``state``. + + Writes three accumulators consumed by ``train.py`` at end of step: + + - ``_flops_seqlen_sum``: ``mbs * tokens.shape[1]`` (padded total tokens + this microbatch contributes). 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 * seq_len²`` is accumulated instead (bit-exact with the pre-fix value). + - ``_flops_vision_patches``: running total of ``num_vision_patches``. + + ``num_vision_patches`` is the precomputed number of vision patches in this + microbatch (drives the ViT term). It is kept model-agnostic on purpose: the + caller — which knows its own encoder's layout — computes the count and passes + a scalar (e.g. Qwen-VL sums ``grid_thw.prod(-1)`` over images and videos). May + be an ``int`` or a scalar ``Tensor`` (a device tensor avoids a host sync here). + + 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. + """ + if tokens is None: + return + + mbs = tokens.shape[0] + seq_len = tokens.shape[1] + _add_flops_accumulator(state, "_flops_seqlen_sum", mbs * seq_len) + + # THD attention term Σᵢ sᵢ², computed inline from cu_seqlens. The squared + # sub-sequence lengths stay on-device (``_scalar_sum_for_accumulator`` returns a + # device tensor, no host sync) so the launch-bound forward path is not stalled; the + # single sync is deferred to the per-step reduce. cu_seqlens is monotonic, so the + # diffs are >= 0 and zero-length padding entries contribute 0 — no boolean mask + # (which would force a data-dependent-size sync) is needed. + sub_seq_lens = _real_subseq_lengths(cu_seqlens, cu_seqlens_argmin, cu_seqlens_unpadded, cu_seqlens_unpadded_argmin) + if sub_seq_lens is not None and sub_seq_lens.numel() > 0: + _add_flops_accumulator(state, "_flops_seqlen_sq_sum", _scalar_sum_for_accumulator(sub_seq_lens.long() ** 2)) + else: + # No cu_seqlens (dense / non-packed) or a degenerate pack with no real + # sub-sequences → BSHD fallback (single pack-length sequence). + _add_flops_accumulator(state, "_flops_seqlen_sq_sum", mbs * seq_len**2) + + if num_vision_patches is not None: + _add_flops_accumulator(state, "_flops_vision_patches", num_vision_patches) + + def vit_flops( cfg: ConfigContainer, batch_size: int, @@ -206,17 +415,19 @@ def attn_layer_flops( num_heads, gqa_groups=8, kv_channels=None, + core_attn_seq_factor=None, ): """Calculate FLOPs for an attention layer.""" p = (kv_channels * num_heads / hidden_size) if kv_channels else 1 g = gqa_groups + core_seq = seq_len if core_attn_seq_factor is None else core_attn_seq_factor return ( 4 * batch_size * seq_len * hidden_size * p - * (hidden_size + (hidden_size * (g / num_heads)) + (seq_len / 2)) + * (hidden_size + (hidden_size * (g / num_heads)) + (core_seq / 2)) ) def mamba_layer_flops( @@ -296,6 +507,7 @@ def hybrid_flops( gdn_conv_kernel_dim=4, vocab_size=256000, mtp_num_layers=0, + core_attn_seq_factor=None, ): """Calculate total FLOPs for the hybrid model.""" flops_fwd = ( @@ -307,6 +519,7 @@ def hybrid_flops( num_attn_heads, gqa_groups, kv_channels, + core_attn_seq_factor, ) + num_mlp_layers * mlp_layer_flops(batch_size, seq_len, hidden_size, mlp_expansion, swiglu) + num_mamba_layers @@ -843,6 +1056,7 @@ def _compute_vit_flops(): gdn_conv_kernel_dim=getattr(cfg.model, "linear_conv_kernel_dim", None) or 4, vocab_size=padded_vocab_size, mtp_num_layers=mtp_num_layers, + core_attn_seq_factor=core_attn_seq_factor, ) return llm_flops + _compute_vit_flops() else: diff --git a/src/megatron/bridge/training/utils/train_utils.py b/src/megatron/bridge/training/utils/train_utils.py index 250cceed71..d34c5939ba 100644 --- a/src/megatron/bridge/training/utils/train_utils.py +++ b/src/megatron/bridge/training/utils/train_utils.py @@ -36,7 +36,6 @@ from megatron.bridge.training.config import ConfigContainer, ProfilingConfig, TrainingConfig from megatron.bridge.training.forward_step_func_types import ForwardStepCallable from megatron.bridge.training.state import GlobalState, TrainState -from megatron.bridge.training.utils.flop_utils import num_floating_point_operations from megatron.bridge.training.utils.mlflow_utils import _sanitize_mlflow_metrics from megatron.bridge.training.utils.pg_utils import get_pg_collection from megatron.bridge.training.utils.theoretical_memory_utils import report_theoretical_memory @@ -1074,62 +1073,33 @@ def training_log( elapsed_time = timers("interval-time").elapsed(barrier=True) elapsed_time_per_iteration = elapsed_time / total_iterations - # Calculate GPU utilization + # Calculate GPU utilization (interval-average achieved throughput). + # + # The main training loop folds each step's DP-exact FLOPS (Σ tokens for the + # linear terms, Σᵢ sᵢ² for THD attention) into + # ``train_state.floating_point_operations_so_far`` every step, so the FLOPS + # performed over this logging interval is exactly the delta since the last + # log. Dividing that interval total by the *interval* elapsed time (not the + # per-iteration average) keeps numerator and denominator over the same window. + # This matters for variable-length THD packing: a per-step numerator (the + # log-boundary step alone) divided by an interval-averaged time would over- or + # under-report whenever that step is heavier/lighter than the interval mean. + # Using the cumulative delta is also the single source of truth — it cannot + # disagree with ``floating_point_operations_so_far`` and needs no extra reduce. num_flops = None if hasattr(config.model, "kv_channels") and hasattr(config.model, "num_attention_heads"): - # Prefer per-microbatch FLOPS accumulators populated by forward_step - # (e.g. vlm_step). They carry the true Σs / Σs² / vision-patches under - # variable-length batches; fall back to the fixed-length assumption - # (batch_size * seq_length) only when no accumulation happened. - # This keeps the per-step TFLOP/s/GPU shown here consistent with the - # `floating_point_operations_so_far` accumulated by the main loop. - # - # VPP correction: forward_step_func is called once per virtual-stage - # per microbatch, so the accumulators over-count by vp_size. Divide - # them back so the FLOPS formula (which already covers all layers) - # receives the correct per-microbatch totals. - # Coerce accumulators to int — getattr on MagicMock test doubles - # returns a MagicMock (not the default), which breaks numeric ops. - local_seqlen_sum = getattr(global_state, "_flops_seqlen_sum", 0) - local_seqlen_sq_sum = getattr(global_state, "_flops_seqlen_sq_sum", 0) - local_vision_patches = getattr(global_state, "_flops_vision_patches", 0) - 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(local_vision_patches, int): - local_vision_patches = 0 - num_vision_patches = local_vision_patches * config.data_parallel_size if local_vision_patches > 0 else 0 - - vp_size = getattr(config.model, "virtual_pipeline_model_parallel_size", None) - 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 * config.data_parallel_size - seqlen_squared_sum = local_seqlen_sq_sum * config.data_parallel_size - num_flops = num_floating_point_operations( - config, - batch_size, - seqlen_sum=seqlen_sum, - seqlen_squared_sum=seqlen_squared_sum, - num_vision_patches=num_vision_patches, - ) - elif seq_length is not None: - seqlen_sum = batch_size * seq_length - seqlen_squared_sum = batch_size * seq_length**2 - num_flops = num_floating_point_operations( - config, - batch_size, - seqlen_sum=seqlen_sum, - seqlen_squared_sum=seqlen_squared_sum, - num_vision_patches=num_vision_patches, - ) - else: - num_flops = num_floating_point_operations(config, batch_size) - per_gpu_tf = num_flops / elapsed_time_per_iteration / get_world_size_safe() / 1e12 + # Coerce to float — floating_point_operations_so_far is a float in + # production, but getattr on a MagicMock test double (and an unset anchor) + # returns a MagicMock; treat non-numerics as 0 so the math stays valid. + flops_so_far = train_state.floating_point_operations_so_far + if not isinstance(flops_so_far, (int, float)): + flops_so_far = 0.0 + prev_flops = getattr(global_state, "_flops_at_last_log", 0.0) + if not isinstance(prev_flops, (int, float)): + prev_flops = flops_so_far + num_flops = max(flops_so_far - prev_flops, 0.0) + global_state._flops_at_last_log = flops_so_far + per_gpu_tf = num_flops / elapsed_time / get_world_size_safe() / 1e12 print_rank_0( f"Step Time : {elapsed_time_per_iteration:.2f}s GPU utilization: {per_gpu_tf:.1f}MODEL_TFLOP/s/GPU" ) diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py index d2bd2d104d..01223ce735 100644 --- a/src/megatron/bridge/training/vlm_step.py +++ b/src/megatron/bridge/training/vlm_step.py @@ -29,6 +29,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.packed_seq_utils import get_packed_seq_params from megatron.bridge.training.utils.padding_utils import ( pad_or_truncate_2d_to_len, @@ -482,21 +483,31 @@ def forward_step( ) = get_batch(data_iterator, state.cfg, use_mtp, pg_collection=pg_collection) timers("batch-generator").stop() - # Accumulate FLOPS metadata across micro-batches. - # Each micro-batch contributes its actual padded seq_length (not cfg.model.seq_length). - # train.py resets these before each step and reads accumulated values afterwards. - if tokens is not None: - mbs = tokens.shape[0] - seq_len = tokens.shape[1] - state._flops_seqlen_sum = getattr(state, "_flops_seqlen_sum", 0) + mbs * seq_len - state._flops_seqlen_sq_sum = getattr(state, "_flops_seqlen_sq_sum", 0) + mbs * seq_len**2 + # Accumulate FLOPS metadata across micro-batches. Passing ``cu_seqlens`` gives + # the THD-correct Σᵢ sᵢ² for the attention term instead of the pack-length² + # BSHD approximation. At CP=1 (and no SP) VLM in-batch packing leaves + # ``cu_seqlens`` equal to the real sub-sequence boundaries, so this counts + # meaningful tokens only. + # NOTE: under CP>1 (or SP), sub-sequences are padded to ``pad_multiple`` (see + # get_batch above), so ``cu_seqlens`` carries that per-sub-seq padding and the + # attention-FLOPS estimate currently includes it (a small over-count). The + # real pre-pad boundaries are not surfaced here yet — tracked as a CP + # follow-up (the linear term also needs a *cp_size correction there, since + # gpt_step CP-shards tokens). train.py resets these before each step and reads + # accumulated values afterwards. + # Vision-patch count is model-specific (Qwen-VL reports it as grid_thw = + # t*h*w per image/video), so compute it here and hand a plain scalar to the + # model-agnostic FLOPS helper. Kept as a device tensor to avoid a host sync. + num_vision_patches = None if visual_inputs is not None: - for attr in ("image_grid_thw", "video_grid_thw"): - grid = getattr(visual_inputs, attr, None) + for grid in ( + getattr(visual_inputs, "image_grid_thw", None), + getattr(visual_inputs, "video_grid_thw", None), + ): if grid is not None and grid.numel() > 0: - state._flops_vision_patches = getattr(state, "_flops_vision_patches", 0) + int( - grid.prod(dim=-1).sum().item() - ) + 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=cu_seqlens, num_vision_patches=num_vision_patches) forward_args = { "input_ids": tokens, diff --git a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py index 5454d492a8..0f6583fb5d 100644 --- a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py +++ b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py @@ -103,12 +103,23 @@ def test_sft_example_runs_with_cp_and_packing(self, tmp_path): cfg.tokenizer.tokenizer_model = "meta-llama/Llama-3.2-1B" cfg.model.calculate_per_token_loss = True cfg.ddp.average_in_collective = False + cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.use_distributed_optimizer = False + cfg.optimizer.use_distributed_optimizer = False # Keep the world-size math simple: tp=1, pp=1, cp=2 -> dp derived from env. cfg.model.tensor_model_parallel_size = 1 cfg.model.pipeline_model_parallel_size = 1 cfg.model.context_parallel_size = 2 + # Keep the L0 test focused on packed SFT + CP plumbing, not full 1B memory coverage. + cfg.model.num_layers = 2 + cfg.model.hidden_size = 128 + cfg.model.ffn_hidden_size = 512 + cfg.model.num_attention_heads = 4 + cfg.model.num_query_groups = 4 + cfg.model.kv_channels = 32 + # Small, fast run cfg.train.train_iters = 2 cfg.train.global_batch_size = 2 diff --git a/tests/unit_tests/training/utils/test_flop_utils.py b/tests/unit_tests/training/utils/test_flop_utils.py index 047043f011..12e644ea9e 100644 --- a/tests/unit_tests/training/utils/test_flop_utils.py +++ b/tests/unit_tests/training/utils/test_flop_utils.py @@ -19,8 +19,14 @@ from unittest.mock import MagicMock, patch import pytest +import torch -from megatron.bridge.training.utils.flop_utils import num_floating_point_operations, vit_flops +from megatron.bridge.training.utils.flop_utils import ( + accumulate_flops_metadata, + num_floating_point_operations, + resolve_global_flops_seqlen_stats, + vit_flops, +) @dataclass @@ -421,6 +427,49 @@ def test_swiglu_scaling_factor(self): ) assert flops_swiglu == expected_swiglu, f"SwiGLU: expected {expected_swiglu:.2e} but got {flops_swiglu:.2e}" + def test_hybrid_attention_term_scales_with_seqlen_squared_sum(self): + """Hybrid attention core FLOPs must use Σ L², not only average sequence length.""" + batch_size = 4 + seq_len = 1024 + hidden_size = 1024 + num_heads = 8 + kv_channels = 128 + cfg = MockConfigContainer( + model=MockModelConfig( + is_hybrid_model=True, + hybrid_layer_pattern="*", + num_layers=1, + hidden_size=hidden_size, + seq_length=seq_len, + ffn_hidden_size=4096, + num_attention_heads=num_heads, + num_query_groups=num_heads, + kv_channels=kv_channels, + vocab_size=32000, + gated_linear_unit=False, + ) + ) + seqlen_sum = batch_size * seq_len + sq_base = batch_size * seq_len**2 + sq_bumped = sq_base + 1_000_000 + + flops_base = num_floating_point_operations( + cfg, + batch_size=batch_size, + seqlen_sum=seqlen_sum, + seqlen_squared_sum=sq_base, + ) + flops_bumped = num_floating_point_operations( + cfg, + batch_size=batch_size, + seqlen_sum=seqlen_sum, + seqlen_squared_sum=sq_bumped, + ) + + query_projection_size = kv_channels * num_heads + expected_delta = 3 * 2 * query_projection_size * (sq_bumped - sq_base) + assert flops_bumped - flops_base == expected_delta + @pytest.mark.unit class TestGDNLayerFlops: @@ -1949,3 +1998,229 @@ def custom(batch_size): assert num_floating_point_operations(cfg, batch_size=4) == sentinel * 4 # Override must have been invoked twice with the right batch_size args. assert captured == [1, 4], f"Override call log mismatch: {captured}" + + +class _State: + """Minimal stand-in for GlobalState — just an attribute bag.""" + + +class TestAccumulateFlopsMetadata: + """Unit tests for ``accumulate_flops_metadata``.""" + + def test_bshd_no_cu_seqlens_uses_pack_length_squared(self): + # Without cu_seqlens, the accumulator falls back to BSHD math — + # mbs * seq_len² — matching the pre-existing behavior on dense + # pretraining / non-packed paths. + state = _State() + tokens = torch.zeros(2, 512) + accumulate_flops_metadata(state, tokens) + assert state._flops_seqlen_sum == 2 * 512 + assert state._flops_seqlen_sq_sum == 2 * 512**2 + + def test_mock_state_accumulators_start_at_zero(self): + state = MagicMock() + tokens = torch.zeros(1, 8) + accumulate_flops_metadata(state, tokens) + assert state._flops_seqlen_sum == 8 + assert state._flops_seqlen_sq_sum == 64 + + def test_thd_cu_seqlens_uses_sum_of_squares(self): + # cu_seqlens = [0, 256, 512, 4096] → sub-seq lengths [256, 256, 3584]. + # THD attention work = 256² + 256² + 3584² = 12,975,488; the BSHD + # approximation (1 × 4096²) would be 16,777,216 — much larger. + state = _State() + tokens = torch.zeros(1, 4096) + cu_seqlens = torch.tensor([0, 256, 512, 4096]) + accumulate_flops_metadata(state, tokens, cu_seqlens=cu_seqlens) + assert state._flops_seqlen_sum == 1 * 4096 + assert state._flops_seqlen_sq_sum == 256**2 + 256**2 + 3584**2 + + def test_thd_padded_cu_seqlens_with_argmin(self): + # Offline packed SFT pads cu_seqlens for CUDA graphs; the real + # entries end at cu_seqlens_argmin. Pad entries past argmin must be + # ignored (here they would otherwise contribute zero-length, but we + # exercise the truncation explicitly). + state = _State() + tokens = torch.zeros(1, 8192) + cu_seqlens = torch.tensor([0, 1024, 4096, 8192, 8192, 8192, 8192]) + argmin = torch.tensor(4) # real entries [0, 1024, 4096, 8192] + accumulate_flops_metadata(state, tokens, cu_seqlens=cu_seqlens, cu_seqlens_argmin=argmin) + assert state._flops_seqlen_sq_sum == 1024**2 + 3072**2 + 4096**2 + + def test_thd_unpadded_takes_precedence_over_padded(self): + # When both cu_seqlens_unpadded and cu_seqlens are present, the + # unpadded variant describes the actual sub-sequence boundaries used + # by the attention kernel (cu_seqlens_q in PackedSeqParams) and must + # be the source of Σᵢ sᵢ². + state = _State() + tokens = torch.zeros(1, 4096) + cu_seqlens_padded = torch.tensor([0, 4096, 4096, 4096]) # 1 pad-aligned sub-seq + cu_seqlens_unpadded = torch.tensor([0, 1000, 3500, 4096]) # 3 real sub-seqs + accumulate_flops_metadata( + state, + tokens, + cu_seqlens=cu_seqlens_padded, + cu_seqlens_unpadded=cu_seqlens_unpadded, + ) + assert state._flops_seqlen_sq_sum == 1000**2 + 2500**2 + 596**2 + + def test_accumulates_additively_across_microbatches(self): + # Each call adds to existing accumulators (microbatch loop semantics). + state = _State() + tokens = torch.zeros(1, 128) + cu_a = torch.tensor([0, 32, 128]) + cu_b = torch.tensor([0, 64, 128]) + accumulate_flops_metadata(state, tokens, cu_seqlens=cu_a) + accumulate_flops_metadata(state, tokens, cu_seqlens=cu_b) + assert state._flops_seqlen_sum == 2 * 128 + assert state._flops_seqlen_sq_sum == (32**2 + 96**2) + (64**2 + 64**2) + + def test_tokens_none_is_noop(self): + state = _State() + accumulate_flops_metadata(state, None) + assert not hasattr(state, "_flops_seqlen_sum") + assert not hasattr(state, "_flops_seqlen_sq_sum") + + def test_num_vision_patches_int(self): + # Vision-patch count is precomputed by the (model-specific) caller and + # passed as a model-agnostic scalar. + state = _State() + tokens = torch.zeros(1, 64) + accumulate_flops_metadata(state, tokens, num_vision_patches=32 + 8) + assert state._flops_vision_patches == 40 + + def test_num_vision_patches_tensor_accumulates_across_microbatches(self): + # A scalar device tensor (no host sync) accumulates correctly. + state = _State() + tokens = torch.zeros(1, 64) + accumulate_flops_metadata(state, tokens, num_vision_patches=torch.tensor(32)) + accumulate_flops_metadata(state, tokens, num_vision_patches=torch.tensor(8)) + assert int(state._flops_vision_patches) == 40 + + def test_num_vision_patches_none_is_noop(self): + state = _State() + tokens = torch.zeros(1, 64) + accumulate_flops_metadata(state, tokens) + assert not hasattr(state, "_flops_vision_patches") + + def test_empty_cu_seqlens_falls_back_to_bshd(self): + # Degenerate cu_seqlens (only one element after argmin truncation) + # yields no sub-seqs, so the helper must fall back to BSHD rather + # than report 0 attention work. + state = _State() + tokens = torch.zeros(1, 256) + cu_seqlens = torch.tensor([0]) + accumulate_flops_metadata(state, tokens, cu_seqlens=cu_seqlens) + assert state._flops_seqlen_sq_sum == 1 * 256**2 + + def test_zero_length_subseq_contributes_zero(self): + # A repeated cu_seqlens value yields a zero-length sub-seq. The Σᵢ sᵢ² + # computation no longer filters with a boolean mask (which would force a + # per-microbatch device sync); zero-length entries must still contribute 0, + # giving the same result as if they were dropped. + state = _State() + tokens = torch.zeros(1, 500) + # sub-seq lengths from [0, 100, 100, 500] -> [100, 0, 400] + cu_seqlens = torch.tensor([0, 100, 100, 500]) + accumulate_flops_metadata(state, tokens, cu_seqlens=cu_seqlens) + assert state._flops_seqlen_sq_sum == 100**2 + 0 + 400**2 + + def test_thd_substantially_smaller_than_bshd_for_short_samples(self): + # Regression check on the headline claim: a pack containing many + # short samples has dramatically less attention work than the BSHD + # approximation would suggest. + state = _State() + tokens = torch.zeros(1, 8192) + # 32 sub-seqs of length 256 → pack length 8192. + cu_seqlens = torch.tensor([i * 256 for i in range(33)]) + accumulate_flops_metadata(state, tokens, cu_seqlens=cu_seqlens) + thd_sq = state._flops_seqlen_sq_sum + bshd_sq = 1 * 8192**2 + # 32 * 256² = 2,097,152 vs 8192² = 67,108,864 → 32× smaller. + assert thd_sq == 32 * 256**2 + assert bshd_sq // thd_sq == 32 + + def test_thd_inline_no_host_sync_on_cpu_returns_int(self): + # On CPU, _scalar_sum_for_accumulator returns a plain int; the inline THD path + # must produce the same Σᵢ sᵢ² as the per-pack analytic value with no buffering. + state = _State() + tokens = torch.zeros(1, 4096) + cu_seqlens = torch.tensor([0, 256, 512, 4096]) + accumulate_flops_metadata(state, tokens, cu_seqlens=cu_seqlens) + assert state._flops_seqlen_sq_sum == 256**2 + 256**2 + 3584**2 + # No deferred cu records: the sum is computed inline, not buffered. + assert getattr(state, "_flops_cu_records", None) is None + + +@pytest.mark.unit +class TestResolveGlobalFlopsSeqlenStats: + """Unit tests for ``resolve_global_flops_seqlen_stats`` (non-distributed paths). + + ``torch.distributed`` is not initialized in unit tests, so the helper takes the + extrapolation fallback (``local * data_parallel_size``). The exact all-reduce + path is covered by functional/distributed tests. + """ + + def test_extrapolates_local_by_dp_size(self): + state = _State() + state._flops_seqlen_sum = 1000 + state._flops_seqlen_sq_sum = 250_000 + state._flops_vision_patches = 64 + seqlen_sum, seqlen_sq_sum, vision = resolve_global_flops_seqlen_stats( + state, data_parallel_size=4, dp_group=None + ) + assert seqlen_sum == 1000 * 4 + assert seqlen_sq_sum == 250_000 * 4 + assert vision == 64 * 4 + + def test_dp_size_one_returns_local(self): + state = _State() + state._flops_seqlen_sum = 512 + state._flops_seqlen_sq_sum = 4096 + seqlen_sum, seqlen_sq_sum, vision = resolve_global_flops_seqlen_stats( + state, data_parallel_size=1, dp_group=None + ) + assert (seqlen_sum, seqlen_sq_sum, vision) == (512, 4096, 0) + + def test_vpp_correction_divides_before_extrapolation(self): + # Accumulators over-count by vp_size; helper divides them back per-rank. + state = _State() + state._flops_seqlen_sum = 1000 * 4 # accumulated once per virtual stage (vp=4) + state._flops_seqlen_sq_sum = 250_000 * 4 + seqlen_sum, seqlen_sq_sum, _ = resolve_global_flops_seqlen_stats( + state, data_parallel_size=2, vp_size=4, dp_group=None + ) + assert seqlen_sum == 1000 * 2 + assert seqlen_sq_sum == 250_000 * 2 + + def test_no_accumulation_returns_none(self): + # Step functions that don't set accumulators → caller falls back to fixed-length. + state = _State() + seqlen_sum, seqlen_sq_sum, vision = resolve_global_flops_seqlen_stats( + state, data_parallel_size=8, dp_group=None + ) + assert seqlen_sum is None + assert seqlen_sq_sum is None + assert vision == 0 + + def test_coerces_scalar_tensor_accumulators(self): + # forward_step may leave accumulators as scalar tensors (deferred host sync). + state = _State() + state._flops_seqlen_sum = torch.tensor(800) + state._flops_seqlen_sq_sum = torch.tensor(160_000) + seqlen_sum, seqlen_sq_sum, _ = resolve_global_flops_seqlen_stats(state, data_parallel_size=2, dp_group=None) + assert seqlen_sum == 1600 + assert seqlen_sq_sum == 320_000 + + def test_dp_group_ignored_when_distributed_not_initialized(self): + # A non-None dp_group must not trigger a collective when torch.distributed + # is not initialized; the helper falls back to extrapolation. + assert not torch.distributed.is_initialized() + state = _State() + state._flops_seqlen_sum = 10 + state._flops_seqlen_sq_sum = 100 + seqlen_sum, seqlen_sq_sum, _ = resolve_global_flops_seqlen_stats( + state, data_parallel_size=4, dp_group=object() + ) + assert seqlen_sum == 40 + assert seqlen_sq_sum == 400 diff --git a/tests/unit_tests/training/utils/test_train_utils.py b/tests/unit_tests/training/utils/test_train_utils.py index 85dac76274..a6e314f82f 100644 --- a/tests/unit_tests/training/utils/test_train_utils.py +++ b/tests/unit_tests/training/utils/test_train_utils.py @@ -419,6 +419,78 @@ def test_tensorboard_logging_interval( mock_global_state.tensorboard_logger.add_scalar.assert_called() mock_global_state.timers.write.assert_called() + @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches") + @mock.patch("megatron.bridge.training.utils.train_utils.reduce_max_stat_across_model_parallel_group") + @mock.patch("megatron.bridge.training.utils.train_utils.get_world_size_safe") + @mock.patch("megatron.bridge.training.utils.train_utils.print_rank_last") + @mock.patch("megatron.bridge.training.utils.train_utils.report_runtime") + @mock.patch("megatron.bridge.training.utils.train_utils.report_throughput") + @mock.patch("megatron.bridge.training.utils.train_utils.report_l2_norm_grad") + def test_throughput_uses_interval_flops_delta( + self, + mock_report_l2_norm_grad, + mock_report_throughput, + mock_report_runtime, + mock_print_rank_last, + mock_get_world_size, + mock_reduce_lr, + mock_get_microbatches, + mock_config, + mock_global_state, + loss_dict, + ): + """Logged TFLOP/s is the interval FLOPs delta ÷ interval time ÷ world size. + + Regression guard for the THD logging fix: ``training_log`` must derive + throughput from the cumulative ``floating_point_operations_so_far`` delta + since the last log (over the full interval elapsed time), not from a single + step's FLOPs over the per-iteration average. It must also advance the + ``_flops_at_last_log`` anchor. + """ + total_loss_dict = self.get_fresh_total_loss_dict() + mock_report_l2_norm_grad.return_value = {} + mock_report_throughput.return_value = {} + mock_report_runtime.return_value = {} + mock_get_microbatches.return_value = 8 + mock_reduce_lr.return_value = 1e-4 + mock_get_world_size.return_value = 8 + + # Log boundary (10 % 5 == 0). interval-time mock returns 0.5s (see fixture). + mock_global_state.train_state.step = 10 + mock_config.logger.log_throughput_to_tensorboard = True + # Interval FLOPs = so_far - anchor = 8e12; per_gpu_tf = 8e12 / 0.5 / 8 / 1e12 = 2.0 + prev_flops = 100.0e12 + mock_global_state._flops_at_last_log = prev_flops + mock_global_state.train_state.floating_point_operations_so_far = prev_flops + 8.0e12 + expected_per_gpu_tf = 8.0e12 / 0.5 / 8 / 1e12 # == 2.0 + + training_log( + loss_dict=loss_dict, + total_loss_dict=total_loss_dict, + learning_rate=1e-4, + decoupled_learning_rate=None, + loss_scale=1024.0, + report_memory_flag=False, + skipped_iter=0, + grad_norm=2.5, + params_norm=15.2, + num_zeros_in_grad=0, + config=mock_config, + global_state=mock_global_state, + history_wct=None, + model=None, + ) + + device_tf_logs = [ + call.args[0]["throughput/tflops/device"] + for call in mock_global_state.wandb_logger.log.call_args_list + if call.args and isinstance(call.args[0], dict) and "throughput/tflops/device" in call.args[0] + ] + assert device_tf_logs, "throughput/tflops/device was not logged" + assert device_tf_logs[-1] == pytest.approx(expected_per_gpu_tf) + # Anchor advanced to the current cumulative for the next interval. + assert mock_global_state._flops_at_last_log == pytest.approx(prev_flops + 8.0e12) + @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches") @mock.patch("megatron.bridge.training.utils.train_utils.reduce_max_stat_across_model_parallel_group") @mock.patch("megatron.bridge.training.utils.train_utils.get_world_size_safe") From 70cbc7dcb7fc8ba9b156e765a5d9c62129c31375 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 15 Jun 2026 16:03:29 -0700 Subject: [PATCH 2/5] test(training): reduce tensor inspect memory footprint Signed-off-by: Chen Cui --- .../training/test_tensor_inspect.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/functional_tests/test_groups/training/test_tensor_inspect.py b/tests/functional_tests/test_groups/training/test_tensor_inspect.py index 949ed231db..667b375d44 100644 --- a/tests/functional_tests/test_groups/training/test_tensor_inspect.py +++ b/tests/functional_tests/test_groups/training/test_tensor_inspect.py @@ -64,11 +64,13 @@ def test_pretrain_with_bf16_tensor_stats(self, tmp_path): torch.distributed.barrier() try: - global_batch_size = 8 + global_batch_size = 2 micro_batch_size = 1 - seq_length = 512 + seq_length = 128 total_iters = 10 + # Keep this L0 focused on tensor inspection. The larger recipe-sized + # shape plus async checkpoint worker can OOM in the shared CI job. model_cfg = GPTModelProvider( normalization="RMSNorm", activation_func=F.silu, @@ -82,7 +84,7 @@ def test_pretrain_with_bf16_tensor_stats(self, tmp_path): persist_layer_norm=True, bias_dropout_fusion=True, apply_rope_fusion=True, - num_query_groups=8, + num_query_groups=4, init_method_std=0.02, layernorm_epsilon=1e-05, rotary_percent=1.0, @@ -90,9 +92,9 @@ def test_pretrain_with_bf16_tensor_stats(self, tmp_path): rope_scaling_factor=32.0, share_embeddings_and_output_weights=True, rotary_base=500_000, - hidden_size=2048, - ffn_hidden_size=8192, - num_attention_heads=32, + hidden_size=512, + ffn_hidden_size=2048, + num_attention_heads=8, tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1, @@ -147,7 +149,7 @@ def test_pretrain_with_bf16_tensor_stats(self, tmp_path): adam_beta1=0.9, adam_beta2=0.95, adam_eps=1e-8, - use_distributed_optimizer=True, + use_distributed_optimizer=False, clip_grad=1.0, lr=3e-3, weight_decay=0.01, @@ -165,11 +167,11 @@ def test_pretrain_with_bf16_tensor_stats(self, tmp_path): ), ddp=DistributedDataParallelConfig( check_for_nan_in_grad=True, - grad_reduce_in_fp32=True, - overlap_grad_reduce=True, - overlap_param_gather=True, + grad_reduce_in_fp32=False, + overlap_grad_reduce=False, + overlap_param_gather=False, average_in_collective=True, - use_distributed_optimizer=True, + use_distributed_optimizer=False, ), dataset=MockGPTDatasetConfig( random_seed=1234, @@ -195,7 +197,7 @@ def test_pretrain_with_bf16_tensor_stats(self, tmp_path): save=checkpoint_dir, ckpt_format="torch_dist", fully_parallel_save=True, - async_save=True, + async_save=False, ), rng=RNGConfig(seed=1234), tensor_inspect=tensor_inspect_cfg, From a386c1f13632915717f0d30f71720061b9b64f36 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 15 Jun 2026 17:36:52 -0700 Subject: [PATCH 3/5] fix(training): limit flops stats all-reduce to THD Signed-off-by: Chen Cui --- src/megatron/bridge/training/train.py | 10 ++-- .../bridge/training/utils/flop_utils.py | 16 ++++--- .../training/utils/test_flop_utils.py | 48 +++++++++++++++++++ 3 files changed, 62 insertions(+), 12 deletions(-) diff --git a/src/megatron/bridge/training/train.py b/src/megatron/bridge/training/train.py index 50df9f2eef..1458639ca1 100644 --- a/src/megatron/bridge/training/train.py +++ b/src/megatron/bridge/training/train.py @@ -445,6 +445,7 @@ def train( global_state._flops_seqlen_sum = 0 global_state._flops_seqlen_sq_sum = 0 global_state._flops_vision_patches = 0 + global_state._flops_requires_global_reduce = False ( loss_dict, @@ -552,11 +553,10 @@ def train( global_state.train_state.skipped_train_samples += num_skipped_samples_in_batch # 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. + # step's FLOPS into the running total. Dense BSHD batches extrapolate exact + # fixed-length stats from the local DP rank; THD batches request one exact SUM + # all-reduce over the pure DP group because packed sub-sequence lengths can + # differ by rank. seqlen_sum, seqlen_squared_sum, num_vision_patches = flop_utils.resolve_global_flops_seqlen_stats( global_state, data_parallel_size=dp_size, diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index 71f2c68d7e..28c7666bb5 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -54,12 +54,12 @@ def resolve_global_flops_seqlen_stats( over-counting, and reduces them to global totals across the data-parallel group. - Under variable-length (THD packed) training the per-rank ``Σᵢ sᵢ²`` and - vision-patch counts differ across DP ranks, so a single SUM all-reduce over - ``dp_group`` is used to get the exact global sum. When no process group is - available (single process, ``dp_group is None``, or ``torch.distributed`` not - initialized) the values are extrapolated as ``local * data_parallel_size`` — - exact only when every DP rank sees an identical sequence-length distribution. + 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. Args: state: Object carrying the ``_flops_*`` accumulators (``GlobalState``). @@ -90,7 +90,8 @@ def resolve_global_flops_seqlen_stats( local_vision_patches //= vp_size use_all_reduce = ( - dp_group is not None + bool(getattr(state, "_flops_requires_global_reduce", False)) + and dp_group is not None and data_parallel_size > 1 and torch.distributed.is_available() and torch.distributed.is_initialized() @@ -223,6 +224,7 @@ def accumulate_flops_metadata( # (which would force a data-dependent-size sync) is needed. sub_seq_lens = _real_subseq_lengths(cu_seqlens, cu_seqlens_argmin, cu_seqlens_unpadded, cu_seqlens_unpadded_argmin) if sub_seq_lens is not None and sub_seq_lens.numel() > 0: + setattr(state, "_flops_requires_global_reduce", True) _add_flops_accumulator(state, "_flops_seqlen_sq_sum", _scalar_sum_for_accumulator(sub_seq_lens.long() ** 2)) else: # No cu_seqlens (dense / non-packed) or a degenerate pack with no real diff --git a/tests/unit_tests/training/utils/test_flop_utils.py b/tests/unit_tests/training/utils/test_flop_utils.py index 12e644ea9e..539705fc53 100644 --- a/tests/unit_tests/training/utils/test_flop_utils.py +++ b/tests/unit_tests/training/utils/test_flop_utils.py @@ -2016,6 +2016,7 @@ def test_bshd_no_cu_seqlens_uses_pack_length_squared(self): accumulate_flops_metadata(state, tokens) assert state._flops_seqlen_sum == 2 * 512 assert state._flops_seqlen_sq_sum == 2 * 512**2 + assert not getattr(state, "_flops_requires_global_reduce", False) def test_mock_state_accumulators_start_at_zero(self): state = MagicMock() @@ -2034,6 +2035,7 @@ def test_thd_cu_seqlens_uses_sum_of_squares(self): accumulate_flops_metadata(state, tokens, cu_seqlens=cu_seqlens) assert state._flops_seqlen_sum == 1 * 4096 assert state._flops_seqlen_sq_sum == 256**2 + 256**2 + 3584**2 + assert state._flops_requires_global_reduce def test_thd_padded_cu_seqlens_with_argmin(self): # Offline packed SFT pads cu_seqlens for CUDA graphs; the real @@ -2088,6 +2090,7 @@ def test_num_vision_patches_int(self): tokens = torch.zeros(1, 64) accumulate_flops_metadata(state, tokens, num_vision_patches=32 + 8) assert state._flops_vision_patches == 40 + assert not getattr(state, "_flops_requires_global_reduce", False) def test_num_vision_patches_tensor_accumulates_across_microbatches(self): # A scalar device tensor (no host sync) accumulates correctly. @@ -2112,6 +2115,7 @@ def test_empty_cu_seqlens_falls_back_to_bshd(self): cu_seqlens = torch.tensor([0]) accumulate_flops_metadata(state, tokens, cu_seqlens=cu_seqlens) assert state._flops_seqlen_sq_sum == 1 * 256**2 + assert not getattr(state, "_flops_requires_global_reduce", False) def test_zero_length_subseq_contributes_zero(self): # A repeated cu_seqlens value yields a zero-length sub-seq. The Σᵢ sᵢ² @@ -2224,3 +2228,47 @@ def test_dp_group_ignored_when_distributed_not_initialized(self): ) assert seqlen_sum == 40 assert seqlen_sq_sum == 400 + + def test_dense_stats_extrapolate_without_all_reduce_when_distributed_initialized(self, monkeypatch): + # Dense BSHD stats are identical across DP ranks, so extrapolation is exact. + # The helper should not create a tiny NCCL collective unless THD metadata + # explicitly requested exact global reduction. + state = _State() + state._flops_seqlen_sum = 10 + state._flops_seqlen_sq_sum = 100 + all_reduce = MagicMock() + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + + seqlen_sum, seqlen_sq_sum, _ = resolve_global_flops_seqlen_stats( + state, data_parallel_size=4, dp_group=object() + ) + + all_reduce.assert_not_called() + assert seqlen_sum == 40 + assert seqlen_sq_sum == 400 + + def test_thd_stats_use_all_reduce_when_distributed_initialized(self, monkeypatch): + # THD packed samples can differ across DP ranks, so a state marked by + # accumulate_flops_metadata must take the exact all-reduce path. + state = _State() + state._flops_seqlen_sum = 10 + state._flops_seqlen_sq_sum = 100 + state._flops_requires_global_reduce = True + + def fake_all_reduce(stats, op=None, group=None): + stats.mul_(4) + + all_reduce = MagicMock(side_effect=fake_all_reduce) + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + seqlen_sum, seqlen_sq_sum, vision = resolve_global_flops_seqlen_stats( + state, data_parallel_size=4, dp_group=object() + ) + + all_reduce.assert_called_once() + assert (seqlen_sum, seqlen_sq_sum, vision) == (40, 400, 0) From 1b9b52521fbe090453996c364ae1254c1b34940a Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 16 Jun 2026 11:52:20 -0700 Subject: [PATCH 4/5] test(training): stabilize CP packed SFT checkpoint test Signed-off-by: Chen Cui --- .../training/test_seqpacking_cp_example.py | 34 +++++++++++++------ 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py index 0f6583fb5d..ed516192bd 100644 --- a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py +++ b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py @@ -17,6 +17,7 @@ import pytest import torch +from datasets import Dataset, DatasetDict from megatron.bridge.data.builders.hf_dataset import HFDatasetConfig from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs @@ -45,16 +46,30 @@ def _make_functional_test_model_small(model: object) -> None: # not the full Llama 3.2 1B model shape. for name, value in { "num_layers": 2, - "hidden_size": 256, - "ffn_hidden_size": 1024, + "hidden_size": 128, + "ffn_hidden_size": 512, "num_attention_heads": 4, "num_query_groups": 4, - "kv_channels": 64, + "kv_channels": 32, "seq_length": 256, }.items(): _set_existing_attr(model, name, value) +def _tiny_squad_dataset() -> DatasetDict: + examples = [ + { + "id": str(i), + "title": "bridge", + "context": "Megatron Bridge converts checkpoints and trains models with packed sequence support.", + "question": "What supports packed sequences?", + "answers": {"text": ["Megatron Bridge"], "answer_start": [0]}, + } + for i in range(16) + ] + return DatasetDict({"train": Dataset.from_list(examples)}) + + class TestPeftSftExample: """Run the PEFT SFT example as a functional test with packed sequences + CP.""" @@ -71,12 +86,14 @@ def test_sft_example_runs_with_cp_and_packing(self, tmp_path): pretrain_tensorboard_dir = os.path.join(shared_dir, "pretrain_tensorboard") sft_checkpoint_dir = os.path.join(shared_dir, "sft_checkpoints") sft_tensorboard_dir = os.path.join(shared_dir, "sft_tensorboard") + dataset_dir = os.path.join(shared_dir, "dataset") if torch.distributed.get_rank() == 0: os.makedirs(pretrain_checkpoint_dir, exist_ok=True) os.makedirs(pretrain_tensorboard_dir, exist_ok=True) os.makedirs(sft_checkpoint_dir, exist_ok=True) os.makedirs(sft_tensorboard_dir, exist_ok=True) + os.makedirs(dataset_dir, exist_ok=True) torch.distributed.barrier() pretrain_cfg = llama32_1b_pretrain_config() @@ -112,14 +129,6 @@ def test_sft_example_runs_with_cp_and_packing(self, tmp_path): cfg.model.pipeline_model_parallel_size = 1 cfg.model.context_parallel_size = 2 - # Keep the L0 test focused on packed SFT + CP plumbing, not full 1B memory coverage. - cfg.model.num_layers = 2 - cfg.model.hidden_size = 128 - cfg.model.ffn_hidden_size = 512 - cfg.model.num_attention_heads = 4 - cfg.model.num_query_groups = 4 - cfg.model.kv_channels = 32 - # Small, fast run cfg.train.train_iters = 2 cfg.train.global_batch_size = 2 @@ -133,6 +142,8 @@ def test_sft_example_runs_with_cp_and_packing(self, tmp_path): # Use a small packed SQuAD dataset to exercise THD/context-parallel slicing cfg.dataset = HFDatasetConfig( dataset_name="rajpurkar/squad", + dataset_dict=_tiny_squad_dataset(), + dataset_root=dataset_dir, process_example_fn=process_squad_example, seq_length=256, dataloader_type="batch", @@ -145,6 +156,7 @@ def test_sft_example_runs_with_cp_and_packing(self, tmp_path): packed_sequence_specs=PackedSequenceSpecs( packed_sequence_size=512, tokenizer_model_name="meta-llama/Llama-3.2-1B", + num_tokenizer_workers=1, pad_seq_to_mult=cfg.model.context_parallel_size * 2, ), rewrite=False, From 91694820d5ca8d699d57b65324e0a3d181806578 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 16 Jun 2026 14:25:39 -0700 Subject: [PATCH 5/5] test(training): probe CP packed SFT larger model Signed-off-by: Chen Cui --- .../test_groups/training/test_seqpacking_cp_example.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py index ed516192bd..dec046a23e 100644 --- a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py +++ b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py @@ -46,11 +46,11 @@ def _make_functional_test_model_small(model: object) -> None: # not the full Llama 3.2 1B model shape. for name, value in { "num_layers": 2, - "hidden_size": 128, - "ffn_hidden_size": 512, + "hidden_size": 256, + "ffn_hidden_size": 1024, "num_attention_heads": 4, "num_query_groups": 4, - "kv_channels": 32, + "kv_channels": 64, "seq_length": 256, }.items(): _set_existing_attr(model, name, value)