Improve default dynamic CP packing scheduler - #5154
Conversation
dc78ffd to
b24471e
Compare
|
/ok to test 90e36b3 |
90e36b3 to
aa3c58c
Compare
|
/ok to test aa3c58c |
aa3c58c to
fbdea00
Compare
|
/ok to test fbdea00 |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| for rank in range(group_start_rank, group_end_rank + needed_count + 1): | ||
| new_micro_batches[rank] = micro_batches[group_end_rank] | ||
| new_exec_times[rank] = sum( | ||
| workload(length, next_power) for length in micro_batches[group_end_rank] | ||
| ) | ||
| new_sample_ids_per_gpu[rank] = sample_ids_per_gpu[group_end_rank] |
There was a problem hiding this comment.
[SUGGESTION Correctness] All ranks in the expanded group are assigned the same list objects (micro_batches[group_end_rank] and sample_ids_per_gpu[group_end_rank]). If any downstream code were to mutate one rank's list in-place (e.g. .append()), all sibling ranks would be silently corrupted.
This is not a regression (the old fill_empty_gpus had the same pattern), and current callers only read these lists. But a defensive shallow copy would prevent a subtle future bug:
for rank in range(group_start_rank, group_end_rank + needed_count + 1):
new_micro_batches[rank] = list(micro_batches[group_end_rank])
new_exec_times[rank] = sum(
workload(length, next_power) for length in micro_batches[group_end_rank]
)
new_sample_ids_per_gpu[rank] = list(sample_ids_per_gpu[group_end_rank])
| cp_size = min_needed | ||
| while cp_size <= total_gpus: | ||
| per_gpu_cost = workload(seq_len, cp_size) | ||
|
|
||
| for group_id, size in list(group_size.items()): | ||
| if size != cp_size: | ||
| continue | ||
| if packing_sequence_len.get(group_id, 0) + seq_len / cp_size > max_seq_len_per_rank: | ||
| continue | ||
| members = group_members[group_id] | ||
| member_set = set(members) | ||
| projected_max = max( | ||
| time + per_gpu_cost if rank in member_set else time | ||
| for rank, time in enumerate(exec_times) | ||
| ) | ||
| if projected_max <= cap and (best is None or projected_max < best[0]): | ||
| best = (projected_max, cp_size, "add", group_id, None) | ||
|
|
||
| free_ranks = [rank for rank, group_id in enumerate(gpu_group_id) if group_id is None] | ||
| if len(free_ranks) >= cp_size: | ||
| chosen_members = sorted(free_ranks, key=lambda rank: exec_times[rank])[:cp_size] | ||
| chosen_set = set(chosen_members) | ||
| projected_max = max( | ||
| time + per_gpu_cost if rank in chosen_set else time | ||
| for rank, time in enumerate(exec_times) | ||
| ) | ||
| if projected_max <= cap and (best is None or projected_max < best[0]): | ||
| best = (projected_max, cp_size, "new", None, chosen_members) | ||
|
|
||
| cp_size *= 2 |
There was a problem hiding this comment.
[IMPORTANT Correctness] The search over candidate CP sizes tries both existing groups and new groups at each power-of-2 CP size, but the best-selection only tracks the global minimum projected_max. Consider this scenario:
- Adding to an existing
cp_size=2group givesprojected_max = 95%of cap (barely fits). - Creating a new
cp_size=2group from free ranks givesprojected_max = 60%of cap (much better balance).
Because the inner for group_id, size loop runs first for each cp_size, an "add" candidate at the current cp_size is compared against "new" at the same cp_size, and also against candidates at larger CP sizes. This is correct — the best tuple tracks the global minimum across all candidates.
However, there is a subtle variable-name collision: on line 622, the list comprehension uses group_id as its iteration variable:
free_ranks = [rank for rank, group_id in enumerate(gpu_group_id) if group_id is None]In Python 3 comprehensions have their own scope so this is functionally correct, but group_id is also the loop variable from the for group_id, size in ... on line 608. Using a distinct name (e.g. gid) would prevent confusion during future edits.
| def fill_empty_gpus_once() -> bool: | ||
| nonlocal micro_batches, exec_times, sample_ids_per_gpu | ||
|
|
||
| empty_ranks = [rank for rank, micro_batch in enumerate(micro_batches) if not micro_batch] | ||
| if not empty_ranks: | ||
| return False | ||
|
|
||
| per_gpu_cost = compute_estimator(seq_len) | ||
|
|
||
| packing_sequence_len[best_gid] = packing_sequence_len.get(best_gid, 0) + seq_len / needed | ||
| for r in chosen_members: | ||
| micro_batches[r].append(seq_len) | ||
| exec_times[r] += per_gpu_cost | ||
| sample_ids_per_gpu[r].append(sample_id) | ||
|
|
||
| buckets[bucket_idx].popleft() | ||
| existing_group_sizes = set(group_size.values()) | ||
| if not existing_group_sizes: | ||
| return False | ||
| min_group_size = min(existing_group_sizes) | ||
| next_power = min(min_group_size * 2, total_gpus) | ||
|
|
||
| while buckets and not buckets[0]: | ||
| buckets.pop(0) | ||
| pp_cursor %= max(1, len(buckets)) | ||
| for group_id, size in list(group_size.items()): | ||
| if size != min_group_size: | ||
| continue | ||
|
|
||
| if needed < prev_needed: | ||
| check_balance = True | ||
| members = group_members[group_id] | ||
| needed_count = next_power - min_group_size | ||
| group_start_rank = members[0] | ||
| group_end_rank = members[-1] | ||
| empty_rank = empty_ranks[0] | ||
| if group_end_rank + 1 > empty_rank or group_end_rank + needed_count >= total_gpus: |
There was a problem hiding this comment.
[IMPORTANT Correctness] fill_empty_gpus_once relies on an invariant that all empty ranks are contiguous at the tail of the rank array. This is true because the main scheduling loop always allocates new groups from the lowest-numbered free ranks. But the invariant is not documented or asserted, and a future change to the allocation order could silently break this function.
Specifically, the function only copies data for three rank ranges:
range(0, group_start_rank)— before the grouprange(group_start_rank, group_end_rank + needed_count + 1)— the expanded groupwork_to_pushshifted forward — between the group end andempty_rank
Ranks after empty_rank are left at their default empty values. If the invariant were violated (a non-empty rank existed after the first empty rank), its data would be silently dropped.
Consider adding a guard assertion at the top of the function:
assert all(not micro_batches[r] for r in range(empty_ranks[0], total_gpus)), (
"fill_empty_gpus_once assumes empty ranks are contiguous at the tail"
)
| def workload(seq_len: int, cp_size: int) -> float: | ||
| return (seq_len * seq_len) / cp_size | ||
|
|
||
| if sample_seq_tuple is None: | ||
| break | ||
| sample_seqlens = sorted(sample_seqlens, key=lambda x: x[1], reverse=True) | ||
| local_tall = sample_seqlens[0][1] | ||
| cap = float(local_tall) * float(max_seq_len_per_rank) * (1.0 + _DYNAMIC_CP_WORKLOAD_CAP_DELTA) |
There was a problem hiding this comment.
[SUGGESTION Simplification] The workload function models self-attention cost as seq_len² / cp_size. This is the same model the old dcp_get_total_workload used — good consistency.
One minor nit: the old function had @lru_cache(maxsize=128). The new version is uncached but defined inside next_hdp_group_packing_aware, so it can't be lru_cache'd (unhashable closure). The arithmetic is cheap (O(1)), so the lack of caching has negligible impact. Just confirming this is intentional.
| # DCP: each logical sequence is replicated on local_cp ranks | ||
| # before CP slicing. Scale each rank's local contribution by | ||
| # max_cp / local_cp, then reduce across DPxCP ranks so each | ||
| # sequence contributes exactly max_cp times. | ||
| for batch in batch_all: | ||
| mb_sum = batch['tokens'].long().sum().clone() | ||
| local_cp = batch['local_cp_size'] | ||
| if isinstance(local_cp, torch.Tensor): | ||
| local_cp = local_cp.item() | ||
| mb_cp_group = parallel_state.get_dynamic_data_context_parallel_groups( | ||
| group_size=local_cp | ||
| ) | ||
| torch.distributed.all_reduce( | ||
| mb_sum, op=torch.distributed.ReduceOp.SUM, group=mb_cp_group | ||
| ) | ||
| # all_reduce result = mb_sum * local_cp. | ||
| # Scale to mb_sum * max_cp. | ||
| mb_sum *= max_cp // local_cp | ||
| token_sum_after += mb_sum |
There was a problem hiding this comment.
[IMPORTANT Correctness] Verified: this change is correct and necessary for the packing-aware scheduler.
The old code did a per-microbatch CP all_reduce then scaled by max_cp / local_cp. That worked only because the old scheduler always assigned local_cp == 1 to short sequences — the CP all_reduce within a group of size 1 was a no-op. With the new packing-aware scheduler, short sequences can join larger CP groups (local_cp > 1). The old formula would overcount: each of the local_cp siblings already holds the full (pre-sliced) sequence, so the all_reduce multiplies by local_cp before the dp_cp reduction, resulting in a total of local_cp × sum(S) × max_cp instead of sum(S) × max_cp.
The new code correctly scales each rank's token sum by max_cp / local_cp without the intermediate CP all_reduce. The dp_cp all-reduce then gives: for a sequence S on a CP group of size C, C × sum(S) × max_cp / C = sum(S) × max_cp — matching the "before" side regardless of local_cp.
There was a problem hiding this comment.
Code Review Summary
PR: Improve default dynamic CP packing scheduler
Findings
| Severity | Count |
|---|---|
| CRITICAL | 0 |
| IMPORTANT | 2 |
| SUGGESTION | 2 |
Key Analysis
Algorithm correctness: Traced through the new next_hdp_group_packing_aware algorithm end-to-end. The core scheduling loop is correct: it tries all power-of-2 CP sizes from min_needed up to total_gpus, evaluates both "add to existing group" and "create new group" placements, and picks the one minimizing critical-path workload subject to a cap of local_tall × max_seq_len_per_rank × (1 + δ). The workload model seq_len²/cp_size is preserved from the old code.
fill_empty_gpus_once correctness: The function is correct given the invariant that empty ranks are always contiguous at the tail of the rank array. This invariant holds because the main loop always allocates new groups from the lowest-numbered free ranks (via sorted(free_ranks, key=...)). However, the invariant is undocumented and unasserted — the most impactful finding is to add a guard assertion.
Test token-sum verification change: Verified mathematically that removing the per-microbatch CP all_reduce is both correct and necessary. The old formula was only accidentally correct because the old scheduler never assigned local_cp > 1 to short sequences. With packing-aware CP assignment, the simpler scale-then-reduce formula is the right one.
Backward compatibility: The three removed utility functions (next_hdp_group, dcp_get_total_workload, dcp_make_buckets_equal) are not imported outside data_schedule.py and data_schedule_utils.py. The next_hdp_group method on BalancedCPScheduler in hybrid_cp_schedule.py is a separate, unrelated method. No external breakage.
Unused variable cleanup: The old code built a groups list in get_groups_and_subsamples that was never read after the scheduling loop — correctly removed.
Overall Assessment
Low risk. The PR is a well-scoped refactoring that replaces a bucket-based greedy scheduler with a packing-aware greedy scheduler. The new algorithm is algorithmically sound and the test changes are correct. The two IMPORTANT findings are about defensive hardening (invariant assertion, shared list references), not active bugs.
c27ebc0 to
50ed3c9
Compare
|
/ok to test 50ed3c9 |
50ed3c9 to
3c55ae7
Compare
|
/ok to test 3c55ae7 |
| assert ( | ||
| total_work_after >= total_work_before | ||
| ), f"Samples were removed: {total_work_before} -> {total_work_after}" | ||
| while any(not micro_batch for micro_batch in micro_batches): |
There was a problem hiding this comment.
[IMPORTANT Correctness] next_hdp_group_packing_aware() can return microbatches with empty DPxCP ranks.
The function documents the “no empty DPxCP rank” invariant, but the fill loop silently exits when fill_empty_gpus_once() cannot expand another group:
while any(not micro_batch for micro_batch in micro_batches):
if not fill_empty_gpus_once():
breakDownstream build_packed_microbatches() assumes every rank has at least one sample and indexes sample_ids_this_group[0], so this becomes an IndexError or rank-divergent scheduling failure. A source-equivalent simulation of this algorithm leaves empty ranks for non-power-of-two DPxCP sizes, e.g. total_gpus=14, max_seq_len_per_rank=100, sample_seqlens=[(0, 50), (1, 50)].
Suggestion: Do not return partial groups. Either make full total_gpus a candidate group size for non-power-of-two DPxCP layouts and keep expanding until no rank is empty, or assert/raise here with a clear error before returning. Please also add a unit test with a non-power-of-two DPxCP size.
| # before CP slicing. Scale each rank's local contribution by | ||
| # max_cp / local_cp, then reduce across DPxCP ranks so each | ||
| # sequence contributes exactly max_cp times. | ||
| for batch in batch_all: |
There was a problem hiding this comment.
[IMPORTANT Test Correctness] The DCP token-sum check no longer reconstructs each local CP group before scaling.
get_cp_slice_for_thd() slices tokens per CP rank before this test reads batch['tokens']. Without the local-CP all_reduce, each rank contributes only its CP slice. The final DPxCP all-reduce then computes full_sequence_sum * max_cp / local_cp instead of full_sequence_sum * max_cp for local_cp > 1, which weakens/removes the invariant this test was checking and can hide token-loss bugs in DCP packing.
Suggestion: Restore the per-microbatch local CP all-reduce before scaling:
mb_cp_group = parallel_state.get_dynamic_data_context_parallel_groups(group_size=local_cp)
torch.distributed.all_reduce(mb_sum, op=torch.distributed.ReduceOp.SUM, group=mb_cp_group)
mb_sum *= max_cp // local_cpThen reduce token_sum_after across dp_cp_group as before.
There was a problem hiding this comment.
I think this check is still running before CP slicing.
In this test, batch_all is read directly from the iterator returned by wrap_data_iterator().
Signed-off-by: tailaim <tailaim@nvidia.com>
Signed-off-by: tailaim <tailaim@nvidia.com>
Signed-off-by: tailaim <tailaim@nvidia.com>
|
/ok to test 04e2ccf |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28253750774 |
5 participants
Summary
Testing
git diff --check upstream/devpython -m compileall -q megatron/core/datasets/data_schedule.py megatron/core/datasets/data_schedule_utils.py megatron/core/model_parallel_config.py tests/unit_tests/test_sequence_packing.pyNot run: full pytest, because the local shell Python does not have torch/pytest available.