[training, perf] fix: THD-aware FLOPS via cu_seqlens (Σᵢ sᵢ²) (re-land of #3839)

[cuichenx]

Re-lands #3839, which was reverted in #4363 after an accidental merge while CI was red.

This branch is no longer a byte-for-byte reland of #3839. It keeps the THD-aware FLOPs accounting and includes the follow-up fixes needed to make the reland correct and CI-stable.

What changed

1. THD-aware FLOPs accounting via cu_seqlens.

For packed/THD batches, attention FLOPs are computed from the real sub-sequence lengths (Σᵢ sᵢ²) instead of treating the whole pack as one dense sequence (seq_len²). The metadata is accumulated in each forward step and folded into train_state.floating_point_operations_so_far once per completed training step.

Throughput logging now uses the cumulative FLOPs delta since the previous log event, so console, WandB, TensorBoard, MLflow, and Comet throughput use the same interval-level FLOPs source as the cumulative counter. This avoids the old current-step-only over/under-reporting behavior for variable-length THD logging intervals.

2. The DP FLOPs all-reduce is conditional.

The earlier implementation reduced FLOPs metadata whenever data_parallel_size > 1, including dense BSHD training where every DP rank has identical fixed-length sequence stats. The final version adds _flops_requires_global_reduce, which is set only when real THD cu_seqlens metadata is seen.

• Dense / BSHD path: extrapolates local * dp_size for fixed-shape dense batches and avoids the extra collective.
• THD packed path: performs one SUM all-reduce over the pure DP group, because per-rank Σᵢ sᵢ² can differ.

3. CP+packing functional test is stabilized.

The test now runs a proper pretrain -> SFT flow: it pretrains a tiny CP=2 checkpoint, then finetunes from that checkpoint, satisfying finetune()'s checkpoint precondition. The pretrain and SFT model shapes are kept identical so checkpoint load validation is meaningful. The test also uses a tiny in-memory SQuAD-shaped dataset to avoid depending on external rajpurkar/squad cache state in CI.

use_distributed_optimizer=False remains in this CP+packing test as a test-local stability/memory workaround.

4. test_tensor_inspect memory footprint is reduced.

The tensor-inspect functional test uses smaller model and batch shapes and disables overlap/distributed-optimizer settings that are unnecessary for the behavior under test. This keeps the shared L0 functional-test job below CI runner memory limits.

CI failure context

The original CI symptom looked like a hang/NCCL watchdog timeout, but the evidence pointed to CUDA OOM in the shared single-process L0 functional-test session. Once one rank OOMs, the remaining ranks can block in the next collective and eventually hit the watchdog.

The final code avoids the unnecessary dense-path FLOPs all-reduce and trims the two memory-sensitive functional tests. The CP+packing test does not exercise the THD exact-reduce path in this configuration: it uses BSHD fallback FLOPs metadata and does not issue the FLOPs DP all-reduce.

Validation

• Focused cw-dfw repro reproduced the checkpoint shape mismatch before the final test fix.
• Focused cw-dfw validation passed after the final test fix.
• GitHub CI is green, including L0_Launch_training and Nemo_CICD_Test.

Reverts the revert (#4363).

[claude]

Nit: "the behavior this test passed on before the THD change" references a specific test from this PR context — it'll go stale quickly. Consider rephrasing to something like "so CP > 1 stays on the BSHD fallback (which is correct for the non-THD path)".

[claude]

Light Code Review

Clean re-land. The refactor centralises FLOPS accumulation into accumulate_flops_metadata / resolve_global_flops_seqlen_stats, replaces the per-step log-time FLOPS recomputation with a cumulative delta, and adds an all-reduce for exact DP-global stats under variable-length packing. No critical bugs found.

Minor observations

• Comment references a test in production code (gpt_step.py:367-368): the phrase will go stale. Left an inline suggestion.
• _real_subseq_lengths argmin sync: int(argmin.item()) on line 165 of flop_utils.py does force a GPU-to-CPU sync when the argmin tensor is on CUDA. The docstring says the function must stay sync-free. In practice this is a single scalar on the offline-packed LLM path only, so the cost is negligible, but the docstring claim is slightly overstated.
• train.py:612 comment says training_log recomputes from accumulators, but training_log now uses the cumulative floating_point_operations_so_far delta. The comment could be clearer.

Test coverage

Good coverage. Unit tests for accumulate_flops_metadata (14 cases) and resolve_global_flops_seqlen_stats (6 cases) cover BSHD fallback, THD cu_seqlens, argmin truncation, unpadded precedence, VPP correction, tensor coercion, and the MagicMock edge case. The test_throughput_uses_interval_flops_delta regression test validates the new delta-based logging path. The functional test (test_seqpacking_cp_example) is appropriately slimmed down.

Suggested test cases: No perf tests impacted.

[cuichenx]

/ok to test 1b9b525

[zhongbozhu]

what does this condition mean and why does it indicate THD layout?

[cuichenx]

this mean we only do the DP all-reduce when trining with THD. For fixed seqlen training, we don't want to run allreduce

[zhongbozhu]

LGTM