Improve default dynamic CP packing scheduler

megatron/core/datasets/data_schedule.py

@@ -11,12 +11,9 @@
     broadcast_tensor,
     build_packed_microbatches,
     create_data_iterator,
-    dcp_get_total_workload,
-    dcp_gpus_needed,
-    dcp_make_buckets_equal,
     get_batch_and_global_seqlens,
     get_cp_slice_for_thd,
-    next_hdp_group,
+    next_hdp_group_packing_aware,
     reroute_samples_to_dcp_ranks,
 )
 from megatron.core.packed_seq_params import (
@@ -424,28 +421,17 @@ def get_groups_and_subsamples(self, sample_id_seqlens):
         """
         mslpr = self.max_seq_len_per_rank
         min_cp = self.min_cp_size
-        workload_fn = lambda seq_len, cp_size=None: dcp_get_total_workload(
-            seq_len, mslpr, cp_size, min_cp
-        )
-        gpus_fn = lambda seq_len: dcp_gpus_needed(seq_len, mslpr, min_cp)
-        buckets_fn = lambda sample_seqlens, compute_est: dcp_make_buckets_equal(
-            sample_seqlens, compute_est, mslpr, min_cp
-        )
 
-        groups = []
         sample_id_groups = []
         sample_id_seqlens = sorted(sample_id_seqlens, key=lambda x: x[1], reverse=True)
+
         while sample_id_seqlens:
-            mb, sample_id_seqlens, exec_times, sample_ids = next_hdp_group(
+            _, sample_id_seqlens, _, sample_ids = next_hdp_group_packing_aware(
                 sample_id_seqlens,
-                workload_fn,
                 self.total_hdp_gpus,
-                gpus_needed_fn=gpus_fn,
-                make_buckets_equal_fn=buckets_fn,
                 max_seq_len_per_rank=mslpr,
-                get_total_workload_fn=workload_fn,
+                min_cp_size=min_cp,
             )
-            groups.append(mb)
             sample_id_groups.append(sample_ids)
 
         if (