Skip to content

Minor improvements for Dynamic-cp #4226

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuzhongw-nvidia merged 9 commits into from
Jun 8, 2026
Merged

Minor improvements for Dynamic-cp #4226

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
2fd422d
support dcp+gdn and fix for vpp+dcp
xiaoyao0115 Apr 9, 2026
95dc830
fix cu_seqlens broadcast when enabling full-iter cudaGraph
xiaoyao0115 Apr 14, 2026
70fea57
Fix MTP loss logging for variable token counts in Dynamic CP
xiaoyao0115 May 27, 2026
5219f2d
minor fixes according to the comments
xiaoyao0115 May 27, 2026
fea6b8a
fix ut
xiaoyao0115 Jun 3, 2026
e58d248
minor fixes
xiaoyao0115 Jun 3, 2026
7601764
lint and add support for mla
xiaoyao0115 Jun 3, 2026
7b40f70
fix full-iter cudagraph
xiaoyao0115 Jun 4, 2026
3de8fbe
adjust golden value for the calculation of mtp-loss log changed
xiaoyao0115 Jun 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 63 additions & 41 deletions megatron/core/datasets/data_schedule.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
align_sample_id_groups,
broadcast_scalars,
broadcast_tensor,
broadcast_to_pp_group,
build_packed_microbatches,
create_data_iterator,
dcp_get_total_workload,
Expand All @@ -23,6 +22,7 @@
from megatron.core.packed_seq_params import PackedSeqParams
from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler
from megatron.core.process_groups_config import ProcessGroupCollection
from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank


class BasePackingScheduler:
Expand Down Expand Up @@ -178,14 +178,17 @@ def run(
Steps:
1. Fetch batches and gather global sequence lengths
2. Check required sample keys
3. Schedule samples into groups
4. Reroute samples to DCP ranks
5. Build packed microbatches
6. Calculate FLOPs info
7. Broadcast to PP group (for middle PP stages)
8. Broadcast to TP group (for non-TP-0 ranks)
3. Strip data fields not needed by this PP stage
4. Schedule samples into groups
5. Reroute samples to DCP ranks
6. Build packed microbatches
7. Calculate FLOPs info
8. Broadcast scalars to TP group (for non-TP-0 ranks)
9. Handle VPP if enabled

Note: There is no PP-group broadcast. In packed-sequence mode
is_dataset_built_on_rank returns True for every PP stage on TP rank 0

Args:
data_iterator: The data iterator.
num_microbatches: The number of microbatches to fetch.
Expand All @@ -204,27 +207,40 @@ def run(
"""

total_dcp_gpus = dp_cp_group.size()
is_first_pp = pp_group.rank() == 0
is_last_pp = pp_group.rank() == pp_group.size() - 1

mtp_on_this_pp = mtp_on_this_rank(config, ignore_virtual=True)
vpp_size = config.virtual_pipeline_model_parallel_size or 1

# Handle VPP: extract the correct data_iterator for this PP stage.
# When VPP is enabled, data_iterator is a list with one entry per VPP stage.
# We only need one data_iterator to run the schedule (all VPP stages on the
# same PP rank share the same underlying dataset), so pick the first non-None.
# Record which VPP stages had data so create_data_iterator knows which ones
# need full samples vs metadata only.
vpp_has_data = None
if (
config.virtual_pipeline_model_parallel_size is not None
and config.virtual_pipeline_model_parallel_size > 1
):
assert len(data_iterator) == config.virtual_pipeline_model_parallel_size
vpp_has_data = [di is not None for di in data_iterator]
# Determine which VPP stages need full data based on pipeline position and MTP.
vpp_needs_data = None
if vpp_size > 1:
assert len(data_iterator) == vpp_size
extracted = None
for di in data_iterator:
if di is not None:
extracted = di
break
data_iterator = extracted

# Only first VPP on first PP and last VPP on last PP need full data.
# MTP VPP stages also need full data (both tokens and labels).
# Middle VPP stages only need metadata (cu_seqlens, max_seqlen, etc.).
vpp_needs_data = [False] * vpp_size
if is_first_pp:
vpp_needs_data[0] = True
if is_last_pp:
vpp_needs_data[-1] = True
if mtp_on_this_pp:
for vp_i in range(vpp_size):
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).
if data_iterator is not None:
Expand All @@ -241,7 +257,24 @@ def run(
key in batch[0]
), f"Batch missing required key {key}, provided keys: {batch[0].keys()}"

# Step 3: Schedule samples into groups
# Step 3: Strip data fields not needed by this PP stage to avoid
# unnecessary all-to-all communication. First PP needs tokens/position_ids,
# last PP needs labels/loss_mask. MTP stages need all four.
# NOTE: this assumes _unpack_batch produces only the six keys below
# (tokens, position_ids, labels, loss_mask, original_seq_len,
# padded_seq_len). Any custom dataset metadata key outside this set
# would be silently dropped here; extend keys_to_keep if needed.
keys_to_keep = {'original_seq_len', 'padded_seq_len'}
if is_first_pp or mtp_on_this_pp:
keys_to_keep.update(['tokens', 'position_ids'])
if is_last_pp or mtp_on_this_pp:
keys_to_keep.update(['labels', 'loss_mask'])
for sample in batch:
for key in list(sample.keys()):
if key not in keys_to_keep:
del sample[key]
Comment on lines +267 to +275

@claude claude Bot Apr 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION Simplification] The stripping logic is correct and well-structured. One minor robustness note: this assumes the only keys in the raw (unpacked) samples are {tokens, labels, loss_mask, position_ids, original_seq_len, padded_seq_len}. If a custom dataset adds extra metadata keys beyond these (e.g., a sample_id or language_tag), they would be silently dropped.

This is unlikely to be a problem in practice since _unpack_batch only produces those six keys, but a brief inline comment noting the assumption (that only these keys exist post-unpack) would help future maintainers.


# Step 4: Schedule samples into groups
sample_id_groups = self.get_groups_and_subsamples(global_id_seqlens)

# Validate scheduling result
Expand All @@ -254,7 +287,7 @@ def run(
f"global_id_seqlens length: {len(global_id_seqlens)}"
)

# Step 4: Reroute samples to DCP ranks
# Step 5: Reroute samples to DCP ranks
samples_this_rank_with_id = reroute_samples_to_dcp_ranks(
batch,
global_ids_this_rank,
Expand All @@ -270,12 +303,12 @@ def run(
dcp_rank = dp_cp_group.rank()
num_micro_batches = len(sample_id_groups)

# Step 5: Build packed microbatches
# Step 6: Build packed microbatches
new_samples = build_packed_microbatches(
samples_this_rank_with_id, sample_id_groups, dcp_rank, dev, self.is_dynamic_cp
)

# Step 6: Calculate FLOPs info
# Step 7: Calculate FLOPs info
seqlen_sum_this_global_batch = float(sum(seqlens_gathered))
seqlen_squared_sum_this_global_batch = float(
sum(seqlen**2 for seqlen in seqlens_gathered)
Expand All @@ -288,24 +321,7 @@ def run(
seqlen_squared_sum_this_global_batch,
) = (None, None, None, None)

# Step 7: Broadcast to PP group (for middle PP stages)
if tp_group.rank() == 0:
(
new_samples,
num_micro_batches,
seqlen_sum_this_global_batch,
seqlen_squared_sum_this_global_batch,
) = broadcast_to_pp_group(
new_samples,
num_micro_batches,
seqlen_sum_this_global_batch,
seqlen_squared_sum_this_global_batch,
pp_group,
dev,
is_dynamic_cp=self.is_dynamic_cp,
)
Comment on lines -291 to -306

@yuzhongw-nvidia yuzhongw-nvidia Jun 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why remove PP broadcast here? For middle PP stages, their new_samples are None, and they do not have chances to get the metadata for packed seq params.

Please correct me if I misunderstand it.

@xiaoyao0115 xiaoyao0115 Jun 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

now, all pp stages would create dataloader. The docstring expired and need to be modified. I will make the changes.


# Step 8: Broadcast to TP group (for non-TP-0 ranks)
# Broadcast to TP group (for non-TP-0 ranks)
(num_micro_batches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch) = (
broadcast_scalars(
[
Expand All @@ -319,9 +335,9 @@ def run(
)
num_micro_batches = int(num_micro_batches)

# Step 9: create data_iterator and handle VPP if enabled
# Step 8: Broadcast to TP group and create data_iterator
new_data_iterator = create_data_iterator(
new_samples, tp_group, config, vpp_has_data, self.is_dynamic_cp
new_samples, tp_group, config, vpp_needs_data, self.is_dynamic_cp
)

return (
Expand Down Expand Up @@ -551,7 +567,13 @@ def get_batch_on_this_rank_for_sequence_packing(

if is_first_or_last_stage or mtp_on_this_rank:
if is_tp_rank_0:
total_tokens = torch.tensor(batch['tokens'].size(0), dtype=torch.int32, device=dev)
# Use whichever data field is available (first stage has tokens, last has labels).
# Avoid `tokens or labels`: PyTorch tensors raise on truthiness when they have
# more than one element ("Boolean value of Tensor ... is ambiguous").
_data_field = batch.get('tokens')
if _data_field is None:
_data_field = batch.get('labels')
total_tokens = torch.tensor(_data_field.size(0), dtype=torch.int32, device=dev)
else:
total_tokens = torch.empty(1, dtype=torch.int32, device=dev)
broadcast_tensor(total_tokens, tp_src_rank, tp_group)
Expand Down
Loading
Loading