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

[cuichenx]

Summary

The FLOPS calculator treated every batch as BSHD with the attention term scaling as pack_length² (seq_length² after #3529). For THD packed training — both offline-packed LLM SFT (via PackedSequenceSpecs) and VLM in-batch packing — the actual attention work is Σᵢ sᵢ² over the sub-sequence lengths within each pack, not pack_length². This PR threads cu_seqlens from the dataloader into the FLOPS accumulator so the attention term reports the true work, and fixes the per-interval throughput logging to use the same source.

The fix

accumulate_flops_metadata() in src/megatron/bridge/training/utils/flop_utils.py, called once per micro-batch from the step functions:

• Accumulates seqlen_sum (Σ padded tokens — drives the linear MLP/proj/logit terms) and the THD attention term seqlen_sq_sum = Σᵢ sᵢ², computed inline from cu_seqlens (preferring cu_seqlens_unpadded when pad_seq_to_mult > 1; the matching *_argmin truncates the padded tail). The squared sub-sequence lengths stay on-device (_scalar_sum_for_accumulator, no .item()), so the per-micro-batch forward path takes no host sync.
• Falls back to the BSHD value mbs * seq_len² when cu_seqlens is absent or degenerate (dense pretraining / non-packed) — bit-exact with pre-fix behavior.
• Also accumulates num_vision_patches (model-agnostic: each step computes the count from its own image/video grid) for the ViT term.

Per step, the main loop resolves the accumulators to a data-parallel-global total via resolve_global_flops_seqlen_stats — an exact SUM all-reduce over the pure DP group (CP ranks share the same cu_seqlens, so reducing over DP×CP would double-count) — feeds num_floating_point_operations, and folds the result into floating_point_operations_so_far. One tiny (3-element) reduce + one host sync per step, at the existing end-of-step sync boundary.

Wired into the step functions:

• gpt_step.py — LLM, offline THD via PackedSequenceSpecs
• vlm_step.py — legacy VLM in-batch packing
• qwen_vl/qwen3_vl_step.py — Qwen3-VL family in-batch packing
• qwen_omni/qwen3_omni_step.py — Qwen3-Omni (vision-patch tracking + BSHD seqlen²)

Throughput logging fix

training_log now derives the per-interval throughput/tflops (WandB / TensorBoard / MLflow / Comet / console) from the cumulative FLOPs delta since the last log ÷ the interval elapsed time — numerator and denominator over the same window. Previously it used the log-boundary step's FLOPS ÷ the interval-average iteration time, which over- or under-reports under variable-length THD (the per-step FLOPS vary). The cumulative floating_point_operations_so_far is now the single source of truth, so the logged throughput cannot disagree with it, and training_log's own resolve_… call (a second per-log DP all-reduce) is removed.

Why

For a packed batch of length L_pack holding N real sub-sequences of lengths s₁ … s_N:

• BSHD approximation (old): attention FLOPS ∝ L_pack²
• THD truth (new): attention FLOPS ∝ Σᵢ sᵢ²

Equality only when N == 1 (a single sub-seq fills the pack). For typical packed SFT or in-batch VLM packing, Σᵢ sᵢ² ≪ L_pack² because cross-sub-seq attention does not happen.

This is a metric correction, not a throughput regression

Replacing the inflated BSHD pack_length² with the accurate (smaller) THD Σᵢ sᵢ² lowers the reported TFLOP/s by ~7% on packed SFT — at unchanged wall-clock. Confirmed on cw-dfw (qwen3_8b_sft, DP=8):

Variant	s/step (wall-clock)	reported TFLOP/s/GPU
baseline / BSHD over-count	1.3666	~578
this PR (THD truth)	1.3623–1.3671	~543

Same clock, lower number — the old number was an over-count. The reported-TFLOP/s drop is the fix working, not a slowdown. (An earlier iteration of this PR added precompute/defer/buffer machinery to "recover" the 7%; that was chasing a measurement artifact in the reported metric, so it has been removed in favor of the straightforward inline implementation above. CI perf goldens for the affected recipes need re-baselining to the corrected values — tracked separately.)

Verification (FLOPS-value correctness)

Two identical workspaces on cw-dfw — only the source files differ. Same model, dataset, seed, parallelism. Identical loss and iteration times across paired runs confirm only the reported FLOPS changes.

Recipe	seq	Packing	Baseline (BSHD)	Fix (THD)	Baseline over-count
qwen3_8b_sft_config	2048	offline THD (SQuAD)	162.6 TFLOP/s/GPU	155.8 TFLOP/s/GPU	+4.4%
qwen3_8b_sft_config	4096	offline THD (SQuAD)	339.9 TFLOP/s/GPU	156.7 TFLOP/s/GPU	+117% (2.17×)
qwen35_vl_9b_sft_config	4096	in-batch (CORD-v2)	261.6 TFLOP/s/GPU	88.9 TFLOP/s/GPU	+194% (~3×)

The seq=2048 vs 4096 pair on the same LLM recipe is the cleanest demonstration: the fix reports near-identical TFLOP/s/GPU (155.8 vs 156.7) because actual attention work is set by per-sample lengths, not pack length. The baseline doubles (162.6 → 339.9) because its pack_length² formula scales with pack size while the real work doesn't — a clear, growing over-count.

Tests

tests/unit_tests/training/utils/test_flop_utils.py:

• BSHD fallback (no cu_seqlens) → mbs * seq_len²
• THD with cu_seqlens → Σᵢ sᵢ²; padded cu_seqlens + cu_seqlens_argmin truncation; cu_seqlens_unpadded precedence; additive accumulation across micro-batches; zero-length sub-seq contributes 0; 32×256 pack ≈ 32× smaller than BSHD; tokens-None no-op; inline-path computes without buffering
• resolve_global_flops_seqlen_stats: extrapolation fallback, VPP correction, empty → None, scalar-tensor coercion, dp_group-ignored-when-dist-not-initialized

tests/unit_tests/training/test_gpt_step.py: get_batch 10-tuple (packed-seq metadata) unchanged.

tests/unit_tests/training/utils/test_train_utils.py: training_log throughput = interval-FLOPs-delta ÷ interval-time ÷ world-size, and the _flops_at_last_log anchor advances.

Out of scope

• llava_step.py and audio_lm_step.py still fall back to cfg.model.seq_length (helper not wired). Easy follow-up.
• Hybrid model path (hybrid_flops) still passes scalar seq_len, dropping quadratic accuracy for hybrid attention — independent of this PR.
• CP > 1: the linear term uses per-CP-rank tokens while Σᵢ sᵢ² comes from the full (un-sharded) cu_seqlens — a cp_size scale mismatch. Not introduced here and not exercised by the CP=1 verification runs. Fixed in follow-up [training, perf] fix: CP-correct token count in FLOPS (cp_size) #4161.

Labels

bug · area:perf · area:training · needs-review

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 2edde36

[yaoyu-33]

/ok to test a78f863

[cuichenx]

error not reproducible with release container. pulling gha container to debug
edit: also not reproducible with gha container. test is flaky

[cuichenx]

/ok to test 5d61c29

[cuichenx]

/ok to test 9d3f709

[cuichenx]

/ok to test 8f9d4f5

[cuichenx]

/ok to test 004b734

[cuichenx]

/ok to test 2b8925c