need reset

Signed-off-by: tailaim <tailaim@nvidia.com>

Author: xiaoyao0115

megatron/core/datasets/data_schedule.py

@@ -1641,13 +1641,22 @@ def fill_empty(sample_id_group):
                 sample_id_group = fill_empty(sample_id_group)
             return sample_id_group
 
+        attempts_since_split = 0
         while remainder > 0:
-            assert i >= 0, f'align_sample_id_groups: no tail microbatch has enough ids to split'
+            if i < 0:
+                if attempts_since_split >= len(sample_id_groups):
+                    assert (
+                        False
+                    ), f'align_sample_id_groups: no tail microbatch has enough ids to split'
+                i = len(sample_id_groups) - 1
             group1, group2 = split_group(sample_id_groups[i])
             if group1 is not None and group2 is not None:
                 sample_id_groups[i] = group1
                 sample_id_groups.append(group2)
                 remainder -= 1
+                attempts_since_split = 0
+            else:
+                attempts_since_split += 1
             i -= 1
 
         return sample_id_groups

tests/unit_tests/context_parallel/test_packing_and_hybrid_cp.py

@@ -523,10 +523,7 @@ def set_tp_pp_vpp(tp, pp, cp, vpp=None, destroy_first=True):
     )
     model = model if isinstance(model, list) else [model]
 
-    #debugmtl
-    print("daddy is here!!!",flush=True)
     data_iterator = get_data_iterator(args)
-    print("daddy is out!!!",flush=True)
 
     forward_backward_func = get_forward_backward_func()
     losses_reduced = forward_backward_func(
@@ -976,6 +973,9 @@ def _create_single_sample(seq_len):
         # Call the function under test
         (new_data_iterator, num_micro_batches, num_total_tokens_this_global_batch, sequence_square_sum_this_global_batch) = wrap_dataloader(data_iterator, config, scheduler_type)
 
+        #debugmtl
+        print(f"rank:{torch.distributed.get_rank()}, exit wrap_dataloader",flush=True)
+        
         # check the result
         assert type(num_micro_batches) is int
         assert type(num_total_tokens_this_global_batch) is float
@@ -1025,5 +1025,8 @@ def _check_batch(batch_all, batch_keys):
                 assert new_data_iterator is None
 
     finally:
+        #debugmtl
+        if torch.distributed.get_rank() == 0:
+            print(f"rank:0, exit test_wrap_dataloader successfully with tp:{tp}, pp:{pp}, cp:{cp}, vpp:{vpp}, scheduler_type:{scheduler_type}",flush=True)
         Utils.destroy_model_parallel()
         unset_global_variables()