fix tflops calculation when sequence_packing_scheduler is not none

megatron/core/datasets/data_schedule.py

@@ -241,8 +241,12 @@ def run(
                     if mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_i):
                         vpp_needs_data[vp_i] = True
 
-        # data_iterator is not None on TP rank 0 for PP stages that need data
-        # (first stage, last stage, or any stage with MTP).
+        # In packed-sequence mode is_dataset_built_on_rank returns True for every
+        # PP stage on TP rank 0, so data_iterator is not None on TP rank 0 of
+        # every PP stage (and every stage independently fetches data and computes
+        # the global seqlen stats). vpp_needs_data / keys_to_keep below only decide
+        # which data fields are kept per stage; they do not affect whether
+        # data_iterator is None.
         if data_iterator is not None:
             assert tp_group.rank() == 0, "Only TP rank 0 should have data_iterator"
 

megatron/training/training.py

@@ -2411,7 +2411,18 @@ def _save_state_dict(attr_name, label):
 
     should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit()
     if should_exit:
-        return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0
+        return (
+            {},
+            True,
+            should_checkpoint,
+            should_exit,
+            exit_code,
+            None,
+            None,
+            0,
+            seqlen_sum_this_global_batch,
+            seqlen_squared_sum_this_global_batch,
+        )
 
     # Empty unused memory.
     if args.empty_unused_memory_level >= 1:
@@ -2494,6 +2505,8 @@ def _save_state_dict(attr_name, label):
             grad_norm,
             num_zeros_in_grad,
             log_max_attention_logit,
+            seqlen_sum_this_global_batch,
+            seqlen_squared_sum_this_global_batch,
         )
     return (
         {},
@@ -2504,6 +2517,8 @@ def _save_state_dict(attr_name, label):
         grad_norm,
         num_zeros_in_grad,
         log_max_attention_logit,
+        seqlen_sum_this_global_batch,
+        seqlen_squared_sum_this_global_batch,
     )
 
 
@@ -3695,6 +3710,8 @@ def trace_handler(p):
             grad_norm = 0.0
             num_zeros_in_grad = 0
             max_attention_logit = None
+            seqlen_sum_this_global_batch = None
+            seqlen_squared_sum_this_global_batch = None
         else:
             ft_integration.on_training_step_start()
             (
@@ -3706,6 +3723,8 @@ def trace_handler(p):
                 grad_norm,
                 num_zeros_in_grad,
                 max_attention_logit,
+                seqlen_sum_this_global_batch,
+                seqlen_squared_sum_this_global_batch,
             ) = train_step(
                 forward_step_func,
                 train_data_iterator,
@@ -3811,14 +3830,23 @@ def trace_handler(p):
         else:
             assert num_skipped_samples_in_batch == 0
         args.skipped_train_samples += num_skipped_samples_in_batch
-        # Drain the per-iteration packed-sequence stats so the FLOPs computation
-        # reflects THD per-chunk causal attention AND excludes padding tokens
-        # from token-linear work. Returns ``(None, None)`` for unpacked BSHD
-        # runs (no collective issued), letting ``num_floating_point_operations``
-        # fall back to its closed-form defaults.
-        total_real_tokens_in_batch, seqlen_squared_sum_in_batch = (
-            consume_seqlen_stats_in_iteration()
-        )
+        if config.sequence_packing_scheduler is not None and not args.skip_train:
+            # Scheduler-based packing does not feed the packed-sequence accumulator.
+            # The scheduler already computed these from the real per-sample lengths
+            # before CP padding/rerouting, so use them directly here.
+            assert seqlen_sum_this_global_batch is not None
+            assert seqlen_squared_sum_this_global_batch is not None
+            total_real_tokens_in_batch = seqlen_sum_this_global_batch
+            seqlen_squared_sum_in_batch = seqlen_squared_sum_this_global_batch
+        else:
+            # Drain the per-iteration packed-sequence stats so the FLOPs computation
+            # reflects THD per-chunk causal attention AND excludes padding tokens
+            # from token-linear work. Returns ``(None, None)`` for unpacked BSHD
+            # runs (no collective issued), letting ``num_floating_point_operations``
+            # fall back to its closed-form defaults.
+            total_real_tokens_in_batch, seqlen_squared_sum_in_batch = (
+                consume_seqlen_stats_in_iteration()
+            )
         num_floating_point_operations_in_batch = num_floating_point_operations(
             args,
             batch_size,