diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index b6a6a65dc3c..5e3c1084792 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/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" diff --git a/megatron/training/training.py b/megatron/training/training.py index 21425422ffb..56d6d33b975 100644 --- a/megatron/training/training.py +++ b/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,