[training, perf] fix: only use cu_seqlens THD FLOPS path for CP==1

L0_Launch_training's test_sft_example_runs_with_cp_and_packing (CP=2) hangs
deterministically on this PR (4/4 runs; main at the merged commit is green). The
test exercises the LLM gpt_step under context parallelism + packing, where this PR
newly runs the cu_seqlens-driven THD Σᵢ sᵢ² accounting in the per-microbatch
forward path. That path is only wired/validated for CP==1; under CP the batch and
its cu_seqlens are CP-partitioned per rank, so the per-rank computation is not yet
correct (the follow-up tracked in #4161) and was destabilizing the run.

Forward cu_seqlens to accumulate_flops_metadata only when CP==1; under CP>1 fall
back to the BSHD term — the exact behavior this test passed on before the THD
change. CP==1 (the verified configuration) is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/training/gpt_step.py

@@ -321,13 +321,21 @@ def _forward_step_common(
     # 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,
-        cu_seqlens_argmin=cu_seqlens_argmin,
-        cu_seqlens_unpadded=cu_seqlens_unpadded,
-        cu_seqlens_unpadded_argmin=cu_seqlens_unpadded_argmin,
+        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 = {