From 7bbe447e79a71cf4f5a5d08f752dc3f06f8fc0c7 Mon Sep 17 00:00:00 2001 From: ilml Date: Wed, 10 Jun 2026 16:55:55 -0700 Subject: [PATCH] Sync Dynamic-CP feature from dev to main Port the Dynamic CP (dynamic context parallelism) feature from the dev branch, superseding the Hybrid Context Parallel implementation (#2282) with the evolved design from dev PRs #2000, #3405, #4226 and follow-up fixes. - Copy the packing schedulers (BasePackingScheduler, DpBalancedScheduler, DefaultDynamicCPScheduler), wrap_data_iterator and get_batch_on_this_rank_for_sequence_packing into megatron/core/datasets/data_schedule.py, and add megatron/core/datasets/data_schedule_utils.py (both byte-identical to dev), keeping HybridCPDataLoaderWrapper for compatibility. - Rename --hybrid-context-parallel to --dynamic-context-parallel; keep hybrid_context_parallel as a deprecated alias on ModelParallelConfig. - Add min_dynamic_context_parallel_size and sequence_packing_scheduler knobs with validation. - Add get_thd_partitioned_indices to transformer_engine.py and thread packed-seq/cu_seqlens plumbing through attention, MLA, MTP, mamba/GDN and the GPT model, matching dev. - Replace the HybridCPDataLoaderWrapper wiring in training.py and pretrain_gpt.py with dev's wrap_data_iterator pattern, and update parallel_state dynamic DPxCP group helpers (byte-identical to dev). - Update main-only callers (pretrain_hybrid.py, tools/prepare_cache.py, megatron/elastification/pretrain_hybrid_flex.py) and existing unit tests for the rename. Co-Authored-By: Claude Fable 5 --- megatron/core/datasets/data_schedule.py | 677 ++++++++++++- megatron/core/datasets/data_schedule_utils.py | 952 ++++++++++++++++++ megatron/core/datasets/gpt_dataset.py | 4 +- megatron/core/datasets/readme.md | 60 ++ .../core/extensions/transformer_engine.py | 52 +- megatron/core/model_parallel_config.py | 51 +- megatron/core/models/gpt/gpt_model.py | 5 +- megatron/core/packed_seq_params.py | 14 + megatron/core/parallel_state.py | 71 +- megatron/core/pipeline_parallel/schedules.py | 19 - megatron/core/ssm/gated_delta_net.py | 102 +- megatron/core/ssm/mamba_context_parallel.py | 8 + megatron/core/ssm/mamba_mixer.py | 9 +- megatron/core/transformer/attention.py | 9 + .../absorbed_mla.py | 14 +- .../transformer/multi_latent_attention.py | 14 +- .../transformer/multi_token_prediction.py | 249 +++-- .../core/transformer/transformer_config.py | 36 +- megatron/core/utils.py | 49 +- .../elastification/pretrain_hybrid_flex.py | 6 +- megatron/training/arguments.py | 59 +- megatron/training/datasets/data_samplers.py | 77 +- megatron/training/datasets/sft_dataset.py | 35 +- megatron/training/initialize.py | 3 +- megatron/training/training.py | 62 +- pretrain_gpt.py | 81 +- pretrain_hybrid.py | 10 +- tests/unit_tests/data/test_get_batch.py | 4 +- .../models/test_hybrid_moe_model.py | 3 + tests/unit_tests/test_parallel_state.py | 10 +- tools/prepare_cache.py | 1 - 31 files changed, 2418 insertions(+), 328 deletions(-) create mode 100644 megatron/core/datasets/data_schedule_utils.py diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 0f016473b6a..b6a6a65dc3c 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -1,12 +1,687 @@ # Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved. -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional, Type import torch from megatron.core import parallel_state +from megatron.core.datasets.data_schedule_utils import ( + align_sample_id_groups, + broadcast_scalars, + 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, + reroute_samples_to_dcp_ranks, +) +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: + """Base class for sequence packing schedulers.""" + + def __init__( + self, + max_seqlen_per_dp_cp_rank: int, + cp_size: int, + dp_size: int, + microbatch_group_size_per_vp_stage: Optional[int], + ): + """ + Args: + max_seqlen_per_dp_cp_rank: The maximum sequence length per DPxCP rank. + cp_size: The context parallel size. + dp_size: The data parallel size. + microbatch_group_size_per_vp_stage: The microbatch group size per virtual + pipeline stage, only used when enabling VPP, otherwise None. + """ + self.max_seqlen_per_dp_cp_rank = max_seqlen_per_dp_cp_rank + self.cp_size = cp_size + self.dp_size = dp_size + self.microbatch_group_size_per_vp_stage = microbatch_group_size_per_vp_stage + + def get_required_sample_keys(self): + """Return the required key of each batch.""" + raise NotImplementedError + + def get_groups_and_subsamples(self, sample_id_seqlens): + """schedule the samples into groups""" + raise NotImplementedError + + def run( + self, + data_iterator, + num_microbatches, + dp_group, + tp_group, + pp_group, + dp_cp_group, + dev, + config, + ): + """ + Run the scheduler and return the new data_iterator. + + Args: + data_iterator: The data iterator. + num_microbatches: The number of microbatches to fetch. + dp_group: Data parallel process group. + tp_group: Tensor parallel process group. + pp_group: Pipeline parallel process group. + dp_cp_group: Data parallel + context parallel process group. + dev: CUDA device. + config: Model parallel config. + + Returns: + new_data_iterator: The new data iterator (or list for VPP). + num_micro_batches: Number of micro batches after scheduling. + seqlen_sum_this_global_batch: Total tokens for FLOPs calculation. + seqlen_squared_sum_this_global_batch: Sum of squared seqlens for FLOPs. + """ + raise NotImplementedError + + +class DpBalancedScheduler(BasePackingScheduler): + """Packs sequences in their original order until reaching the max limit of sequence length.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.max_seq_len_all_ranks = self.max_seqlen_per_dp_cp_rank * self.cp_size + self.is_dynamic_cp = False + + def get_required_sample_keys(self): + """Return the required key of each batch.""" + return [ + "tokens", + "labels", + "loss_mask", + "position_ids", + "original_seq_len", # Length of the original sequence length, should be a gpu tensor. + "padded_seq_len", # Length of the padded sequence length, should be a gpu tensor. + ] + + def get_groups_and_subsamples(self, sample_id_seqlens): + """ + Packs sequences in their original order until reaching the max limit of sequence length. + """ + sample_id_groups = [] + packed_id_groups = [] + sum_seqlen = 0 + single_microbatch = [] + + for i in range(len(sample_id_seqlens)): + if sum_seqlen + sample_id_seqlens[i][1] <= self.max_seq_len_all_ranks: + single_microbatch.append(i) + sum_seqlen += sample_id_seqlens[i][1] + else: + packed_id_groups.append(single_microbatch) + single_microbatch = [i] + sum_seqlen = sample_id_seqlens[i][1] + if len(single_microbatch) > 0: + packed_id_groups.append(single_microbatch) + + # we want the number of packed sequences to be multiple of dp_size + # so we move few samples from previous microbatch + # to the end of the microbatches if needed + num_packed_sequence = len(packed_id_groups) + + # when enabling vpp, we want the number of packed sequences to be + # multiple of dp_size * microbatch_group_size_per_vp_stage + multiple = self.dp_size * ( + self.microbatch_group_size_per_vp_stage + if self.microbatch_group_size_per_vp_stage is not None + else 1 + ) + if num_packed_sequence % multiple != 0: + remainder = num_packed_sequence % multiple + num_to_move = multiple - remainder + i = num_packed_sequence - 1 + while num_to_move > 0: + assert i >= 0, "Not enough samples to move" + if len(packed_id_groups[i]) > 1: + seq_id = packed_id_groups[i].pop() + packed_id_groups.append([seq_id]) + num_to_move -= 1 + else: + i -= 1 + + num_micro_batches = int(len(packed_id_groups) / self.dp_size) + for i in range(num_micro_batches): + sample_id_groups.append([]) + for j in range(self.cp_size * self.dp_size): + seq_id = int(i * self.dp_size + j / self.cp_size) + sample_id_groups[i].append(packed_id_groups[seq_id]) + return sample_id_groups + + def run( + self, + data_iterator, + num_microbatches: int, + dp_group, + tp_group, + pp_group, + dp_cp_group, + dev: torch.device, + config, + ): + """ + Run the complete scheduling pipeline. + + Steps: + 1. Fetch batches and gather global sequence lengths + 2. Check required sample keys + 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. + dp_group: Data parallel process group. + tp_group: Tensor parallel process group. + pp_group: Pipeline parallel process group. + dp_cp_group: Data parallel + context parallel process group. + dev: CUDA device. + config: Model parallel config. + + Returns: + new_data_iterator: The new data iterator (or list for VPP). + num_micro_batches: Number of micro batches after scheduling. + seqlen_sum_this_global_batch: Total tokens for FLOPs calculation. + seqlen_squared_sum_this_global_batch: Sum of squared seqlens for FLOPs. + """ + + 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. + # 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: + assert tp_group.rank() == 0, "Only TP rank 0 should have data_iterator" + + # Step 1: Fetch batches and gather global sequence lengths + batch, global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered = ( + get_batch_and_global_seqlens(data_iterator, num_microbatches, dp_group) + ) + + # Step 2: Check required sample keys + for key in self.get_required_sample_keys(): + assert ( + key in batch[0] + ), f"Batch missing required key {key}, provided keys: {batch[0].keys()}" + + # 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] + + # Step 4: Schedule samples into groups + sample_id_groups = self.get_groups_and_subsamples(global_id_seqlens) + + # Validate scheduling result + set_gbs = set() + for group in sample_id_groups: + for sub in group: + set_gbs.update(sub) + assert len(set_gbs) == len(global_id_seqlens), ( + f"set_gbs length: {len(set_gbs)} != " + f"global_id_seqlens length: {len(global_id_seqlens)}" + ) + + # Step 5: Reroute samples to DCP ranks + samples_this_rank_with_id = reroute_samples_to_dcp_ranks( + batch, + global_ids_this_rank, + global_id_seqlens, + sample_id_groups, + offsets, + dp_group, + tp_group, + dp_cp_group, + total_dcp_gpus, + ) + + dcp_rank = dp_cp_group.rank() + num_micro_batches = len(sample_id_groups) + + # 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 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) + ) + else: + ( + new_samples, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) = (None, None, None, None) + + # 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( + [ + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ], + tp_group, + dev, + ) + ) + num_micro_batches = int(num_micro_batches) + + # Step 8: Broadcast to TP group and create data_iterator + new_data_iterator = create_data_iterator( + new_samples, tp_group, config, vpp_needs_data, self.is_dynamic_cp + ) + + return ( + new_data_iterator, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) + + +class DefaultDynamicCPScheduler(DpBalancedScheduler): + """ + Dynamic CP scheduler that balances workload across variable CP sizes. + """ + + def __init__(self, *args, min_cp_size=1, **kwargs): + super().__init__(*args, **kwargs) + self.is_dynamic_cp = True + self.max_seq_len_per_rank = self.max_seqlen_per_dp_cp_rank + self.total_hdp_gpus = self.dp_size * self.cp_size + self.min_cp_size = min_cp_size + + def get_groups_and_subsamples(self, sample_id_seqlens): + """ + This function recursively forms groups of sub-samples such that all DPxCP ranks + have a roughly balanced workload in the group. + """ + 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, + 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, + ) + groups.append(mb) + sample_id_groups.append(sample_ids) + + if ( + self.microbatch_group_size_per_vp_stage is not None + and self.microbatch_group_size_per_vp_stage > 1 + ): + sample_id_groups = align_sample_id_groups( + sample_id_groups, self.microbatch_group_size_per_vp_stage + ) + + return sample_id_groups + + +scheduler_map: Dict[str, Type[BasePackingScheduler]] = { + "dp_balanced": DpBalancedScheduler, + "default_dynamic_cp": DefaultDynamicCPScheduler, +} + + +def wrap_data_iterator( + data_iterator, config, num_microbatches, pg_collection: Optional[ProcessGroupCollection] = None +): + """ + A wrapper function that wraps around an existing data_iterator + and return the num_micro_batches for sequence packing. + + Args: + data_iterator: The original data_iterator to wrap around + config: The config object containing the max_seqlen_per_dp_cp_rank + dp_cp_group: Data parallel context parallel group. + pg_collection: The process group collection. + """ + + if pg_collection is None: + dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True) + dp_group = parallel_state.get_data_parallel_group() + tp_group = parallel_state.get_tensor_model_parallel_group() + pp_group = parallel_state.get_pipeline_model_parallel_group() + else: + dp_cp_group = pg_collection.dp_cp + dp_group = pg_collection.dp + tp_group = pg_collection.tp + pp_group = pg_collection.pp + assert ( + dp_cp_group is not None + and dp_group is not None + and tp_group is not None + and pp_group is not None + ), "dp_cp_group, dp_group, tp_group must not be None when using sequence packing" + + dev = torch.cuda.current_device() + dp_size = dp_group.size() + cp_size = dp_cp_group.size() // dp_size + + # Look up the scheduler class by name + scheduler_type = config.sequence_packing_scheduler + + scheduler_kwargs = {} + if scheduler_type == 'default_dynamic_cp': + scheduler_kwargs['min_cp_size'] = config.min_dynamic_context_parallel_size + + scheduler = scheduler_map[scheduler_type]( + config.max_seqlen_per_dp_cp_rank, + cp_size, + dp_size, + ( + None + if config.virtual_pipeline_model_parallel_size is None + else config.microbatch_group_size_per_vp_stage + ), + **scheduler_kwargs, + ) + + ( + new_data_iterator, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) = scheduler.run( + data_iterator, num_microbatches, dp_group, tp_group, pp_group, dp_cp_group, dev, config + ) + + return ( + new_data_iterator, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) + + +def get_batch_on_this_rank_for_sequence_packing( + data_iterator, + vpp_size: Optional[int] = None, + mtp_on_this_rank: bool = False, + vp_stage: Optional[int] = None, + dynamic_cp: bool = False, + pg_collection: Optional[ProcessGroupCollection] = None, +): + """ + Get a batch of data for sequence packing. + Args: + data_iterator (Iterator): The data iterator to get the batch from. + mtp_on_this_rank (bool): Whether to use multi-token prediction. + vp_stage (Optional[int]): The stage of the pipeline. + Returns: + tuple of (tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params) + """ + + if pg_collection is None: + tp_group = parallel_state.get_tensor_model_parallel_group() + pp_group = parallel_state.get_pipeline_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + else: + tp_group = pg_collection.tp + pp_group = pg_collection.pp + cp_group = pg_collection.cp + + tp_src_rank = torch.distributed.get_process_group_ranks(tp_group)[0] + + is_tp_rank_0 = tp_group.rank() == 0 + is_first_stage = pp_group.rank() == 0 and (vp_stage is None or vp_stage == 0) + is_last_stage = pp_group.rank() == pp_group.size() - 1 and ( + vp_stage is None or vp_stage == vpp_size - 1 + ) + + is_first_or_last_stage = is_first_stage or is_last_stage + dev = torch.cuda.current_device() + + # data_iterator should return a batch including the following keys. + batch_keys = ['cu_seqlens', 'cu_seqlens_padded', 'max_seqlen'] + if dynamic_cp: + batch_keys.append('local_cp_size') + if is_first_stage or mtp_on_this_rank: + batch_keys.append('tokens') + batch_keys.append('position_ids') + if is_last_stage or mtp_on_this_rank: + batch_keys.append('labels') + batch_keys.append('loss_mask') + + # Get a batch from data_iterator or create an emtpy batch. + if is_tp_rank_0: + assert data_iterator is not None + batch = next(data_iterator) + for key in batch_keys: + assert key in batch, f"{key} is missing in current batch." + else: + assert data_iterator is None, "Non TP 0 rank should not have data_iterator" + batch = {} + + # For dynamic CP, determine the correct cp_group from batch on TP rank 0. + if dynamic_cp and is_tp_rank_0: + local_cp_size_val = batch['local_cp_size'] + if isinstance(local_cp_size_val, torch.Tensor): + local_cp_size_val = local_cp_size_val.item() + cp_group = parallel_state.get_dynamic_data_context_parallel_groups( + group_size=local_cp_size_val + ) + + # Partition tokens, position_ids, labels, loss_mask for context parallel. + # Only TP rank 0 on stages that have data (first/last PP stage or MTP stage) needs this. + if is_tp_rank_0 and (is_first_or_last_stage or mtp_on_this_rank): + get_cp_slice_for_thd(batch, cp_group) + + # Broadcast cu_seqlens_size because we need it to create placeholder for cu_seqlens and + # cu_seqlens_padded for non TP 0 ranks. + if is_tp_rank_0: + cu_seqlen_size = torch.tensor(batch['cu_seqlens'].size(0), dtype=torch.int32, device=dev) + else: + cu_seqlen_size = torch.empty(1, dtype=torch.int32, device=dev) + broadcast_tensor(cu_seqlen_size, tp_src_rank, tp_group) + cu_seqlen_size = cu_seqlen_size.item() + + # Broadcast total_tokens because we need it to create placeholder for tokens, position_ids, + # labels, loss_mask for non TP 0 ranks. Only first stage, last stage, + # and stage with mtp need this. + + if is_first_or_last_stage or mtp_on_this_rank: + if is_tp_rank_0: + # 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) + total_tokens = total_tokens.item() + + # Step1: Prepare "tokens", "position_ids" for first stage and stage with mtp on all TP ranks. + if is_first_stage or mtp_on_this_rank: + if is_tp_rank_0: + assert batch['tokens'].dtype == torch.int64 + assert batch['position_ids'].dtype == torch.int64 + batch['tokens'] = batch['tokens'].view(1, total_tokens) + batch['position_ids'] = batch['position_ids'].view(1, total_tokens) + else: + batch['tokens'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + batch['position_ids'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + else: + # Non first stage rank doesn't need tokens and position_ids. + batch['tokens'] = None + batch['position_ids'] = None + + # Step2: Prepare "labels", "loss_mask" for last stage and stage with mtp on all TP ranks. + if is_last_stage or mtp_on_this_rank: + if is_tp_rank_0: + assert batch['labels'].dtype == torch.int64 + assert batch['loss_mask'].dtype == torch.float32 + batch['labels'] = batch['labels'].view(1, total_tokens) + batch['loss_mask'] = batch['loss_mask'].view(1, total_tokens) + else: + batch['labels'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + batch['loss_mask'] = torch.empty([1, total_tokens], dtype=torch.float32, device=dev) + else: + # Non last stage rank doesn't need labels and loss_mask. + batch['labels'] = None + batch['loss_mask'] = None + + # Step3: Prepare "cu_seqlens", "cu_seqlens_padded", "max_seqlen" on all ranks. + if is_tp_rank_0: + assert batch['cu_seqlens'].dtype == torch.int32 + assert batch['cu_seqlens_padded'].dtype == torch.int32 + assert batch['cu_seqlens'].dim() == 1 + assert batch['cu_seqlens_padded'].dim() == 1 + if type(batch['max_seqlen']) == int: + batch['max_seqlen'] = torch.tensor(batch['max_seqlen'], dtype=torch.int32, device=dev) + else: + assert batch['max_seqlen'].dtype == torch.int32 + assert batch['max_seqlen'].numel() == 1 + else: + batch['cu_seqlens'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev) + batch['cu_seqlens_padded'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev) + batch['max_seqlen'] = torch.empty(1, dtype=torch.int32, device=dev) + + # Step4: Prepare "local_cp_size" if dynamic context parallel is enabled. + if dynamic_cp: + if is_tp_rank_0: + if type(batch['local_cp_size']) == int: + batch['local_cp_size'] = torch.tensor( + batch['local_cp_size'], dtype=torch.int32, device=dev + ) + else: + assert batch['local_cp_size'].dtype == torch.int32 + assert batch['local_cp_size'].numel() == 1 + else: + batch['local_cp_size'] = torch.empty(1, dtype=torch.int32, device=dev) + else: + batch['local_cp_size'] = None + + # Broadcast batch inside TP group. + broadcast_tensor(batch['tokens'], tp_src_rank, tp_group) + broadcast_tensor(batch['position_ids'], tp_src_rank, tp_group) + broadcast_tensor(batch['labels'], tp_src_rank, tp_group) + broadcast_tensor(batch['loss_mask'], tp_src_rank, tp_group) + broadcast_tensor(batch['cu_seqlens'], tp_src_rank, tp_group) + broadcast_tensor(batch['cu_seqlens_padded'], tp_src_rank, tp_group) + broadcast_tensor(batch['max_seqlen'], tp_src_rank, tp_group) + broadcast_tensor(batch['local_cp_size'], tp_src_rank, tp_group) + + # Extract the data from batch after broadcasting. + tokens = batch['tokens'] + position_ids = batch['position_ids'] + labels = batch['labels'] + loss_mask = batch['loss_mask'] + cu_seqlens = batch['cu_seqlens'] + cu_seqlens_padded = batch['cu_seqlens_padded'] + max_seqlen = batch['max_seqlen'].item() + local_cp_size = batch['local_cp_size'].item() if dynamic_cp else None + cp_group = ( + parallel_state.get_dynamic_data_context_parallel_groups(group_size=local_cp_size) + if dynamic_cp + else None + ) + + # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as cu_seqlens to + # get the correct result. + # TODO: Revert this workaround once TE fixes the issue. + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_padded, + cu_seqlens_kv=cu_seqlens_padded, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + local_cp_size=local_cp_size, + cp_group=cp_group, + ) + + # "attention_mask" is not valid for sequence packing, so set it to None. + return tokens, labels, loss_mask, None, position_ids, packed_seq_params class HybridCPDataLoaderWrapper: diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py new file mode 100644 index 00000000000..51be6282ffe --- /dev/null +++ b/megatron/core/datasets/data_schedule_utils.py @@ -0,0 +1,952 @@ +# Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved. + +from collections import deque +from functools import lru_cache +from math import ceil, log2 +from typing import Callable, Dict, List, Optional, Tuple + +import torch + +from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices +from megatron.core.rerun_state_machine import RerunDataIterator + + +def get_cp_slice_for_thd(batch, cp_group): + """Partition sequence data for context parallelism in THD format. + + Uses TE's THD partitioned indices to split the packed sequence across CP ranks. + Only keys present in the batch are sliced. + + Args: + batch: Dict with packed sequence data. + cp_group: Context parallel process group. + """ + cp_size = cp_group.size() + if cp_size <= 1: + return + cp_rank = cp_group.rank() + # Use whichever data field is available to determine total_tokens + for _key in ['tokens', 'labels', 'loss_mask', 'position_ids']: + if _key in batch and batch[_key] is not None: + total_tokens = batch[_key].size(0) + break + else: + raise ValueError("Cannot determine total_tokens: no data field found in batch") + # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as + # cu_seqlens to get the correct result. + # TODO: Revert this workaround once TE fixes the issue. + cu_seqlens = batch["cu_seqlens_padded"] + index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) + for key in ['tokens', 'position_ids', 'labels', 'loss_mask']: + if key in batch: + batch[key] = batch[key].index_select(0, index) + + +def _unpack_batch(batch: List[Dict[str, torch.Tensor]]) -> List[Dict[str, torch.Tensor]]: + """ + Unpacks the packed samples into a list of sub-samples. + Since each sub-sample may be routed to different DPxCP ranks, + we unpack the sample here to avoid unnecessarily transferring + the entire packed sample. + + Two input shapes are accepted: + + * **Pre-packed** (e.g. :class:`SFTDataset`): each sample carries a + ``cu_seqlens`` tensor and the tokens of multiple sub-samples + concatenated together. We slice them apart and synthesize + ``original_seq_len`` / ``padded_seq_len`` from the cu_seqlens deltas. + + * **Already unpacked** (e.g. :class:`VarlenDataset`): each sample is a + single sub-sample that already carries ``padded_seq_len`` (and + usually ``original_seq_len``). We just normalize the leading batch + dimension introduced by the default collate_fn and return as-is. + """ + # Short-circuit for datasets that already emit one sub-sample per index. + if batch and "padded_seq_len" in batch[0]: + for sample in batch: + for key in sample.keys(): + if sample[key].ndim == 2 and sample[key].shape[0] == 1: + # Drop the redundant batch dim added by collate_fn. + sample[key] = sample[key].squeeze(0) + if "original_seq_len" not in sample: + sample["original_seq_len"] = sample["padded_seq_len"].clone() + return batch + + batch_unpacked = [] + dev = batch[0]["cu_seqlens"].device + original_seq_lens = [] + padded_seq_lens = [] + # Determine which data fields exist in the batch + data_keys = [k for k in ["tokens", "labels", "loss_mask", "position_ids"] if k in batch[0]] + for sample in batch: + for key in sample.keys(): + if len(sample[key].shape) == 2: + # squeeze the redundant batch dimension added by + # default collate_fn in pytorch dataloader + # we need a custom collate_fn for THD to avoid this + # current THD does not support micro_batch_size > 1 due to sft_dataset.py and + # data_loader in data_samples.py + sample[key] = sample[key].squeeze(0) + for sub_sample in range(sample["cu_seqlens"].shape[0] - 1): + sub_sample_dict = {} + start_idx = sample["cu_seqlens"][sub_sample] + end_idx = sample["cu_seqlens"][sub_sample + 1] + if end_idx - start_idx == 0: + continue + for key in data_keys: + sub_sample_dict[key] = sample[key][start_idx:end_idx] + # Since sft_dataset.py does not provide cu_seqlens_original, + # we assume original_seq_len equals padded_seq_len here. + # Ideally the dataset should define the pre-padding seq_len. + seq_len = (end_idx - start_idx).item() + original_seq_lens.append(seq_len) + padded_seq_lens.append(seq_len) + batch_unpacked.append(sub_sample_dict) + + # Single H2D transfer for all seq lens + original_seq_lens_cuda = torch.tensor(original_seq_lens, device=dev) + padded_seq_lens_cuda = torch.tensor(padded_seq_lens, device=dev) + for i, sub_sample_dict in enumerate(batch_unpacked): + sub_sample_dict["original_seq_len"] = original_seq_lens_cuda[i : i + 1] + sub_sample_dict["padded_seq_len"] = padded_seq_lens_cuda[i : i + 1] + + return batch_unpacked + + +def _get_global_seqlens_and_ids(subsample_seqlens: torch.Tensor, dp_group): + """ + Gathers the sequence lengths of all subsamples from all DP ranks and calculates global IDs. + """ + # Collect the number of subsamples from all ranks + num_local_subsamples = subsample_seqlens.shape[0] + local_len = torch.tensor([num_local_subsamples], dtype=torch.int32).cuda() + dp_subsample_count = [torch.zeros_like(local_len) for _ in range(dp_group.size())] + torch.distributed.all_gather(dp_subsample_count, local_len, group=dp_group) + + # Find the max number of subsamples across all ranks and pad subsample_seqlens to max length + dp_subsample_counts = torch.stack(dp_subsample_count, dim=0).cpu().view(-1) + max_sub_samples = int(dp_subsample_counts.max().item()) + + if num_local_subsamples < max_sub_samples: + subsample_seqlens_padded = torch.cat( + [ + subsample_seqlens, + torch.zeros(max_sub_samples - num_local_subsamples, dtype=torch.int32).cuda(), + ], + dim=0, + ) + else: + subsample_seqlens_padded = subsample_seqlens + + # Gather the subsample_seqlens from all ranks + seqlens_gathered = [torch.empty_like(subsample_seqlens_padded) for _ in range(dp_group.size())] + torch.distributed.all_gather(seqlens_gathered, subsample_seqlens_padded, group=dp_group) + + # Trim each seqlens_gathered to the length of the correct sample + for dp_rank, seqlen in enumerate(seqlens_gathered): + seqlens_gathered[dp_rank] = seqlen[: dp_subsample_counts[dp_rank]] + + seqlens_gathered = torch.cat(seqlens_gathered, dim=0) + seqlens_gathered = seqlens_gathered.cpu().tolist() + + # Calculate the offsets to assign unique global ID to each subsample. + csum = torch.cumsum(dp_subsample_counts, dim=0, dtype=torch.int32) + offsets = torch.cat([torch.zeros(1, dtype=torch.int32), csum], dim=0) + + # Calculate global ID for each subsample + dp_rank = dp_group.rank() + global_ids = torch.arange(len(seqlens_gathered), dtype=torch.int32).cuda() + + # Create a list of (global_id, seqlen) tuples for scheduling + global_id_seqlens = [(i, seqlens_gathered[i]) for i in range(len(global_ids))] + + # Get the global IDs locally present on this rank + start_idx = offsets[dp_rank] + end_idx = offsets[dp_rank + 1] + + global_ids_this_rank = global_ids[start_idx:end_idx] + + return global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered + + +def _pack_sequences( + samples: List, + padded_lengths: torch.Tensor, + original_lengths: torch.Tensor, + local_cp_size: Optional[torch.Tensor], + dev: torch.device, +) -> Dict[str, torch.Tensor]: + """Pack multiple samples into a single packed sample.""" + + def _pack_tensors(tensors): + return torch.cat([t.reshape(-1) for t in tensors], dim=0) + + new_sample = {} + for key in ['tokens', 'labels', 'loss_mask', 'position_ids']: + if key in samples[0]: + new_sample[key] = _pack_tensors([sample[key] for sample in samples]) + + padded_lengths = padded_lengths.to(device=dev, dtype=torch.int32, non_blocking=True).reshape(-1) + cu_seqlens_padded = torch.empty(padded_lengths.numel() + 1, device=dev, dtype=torch.int32) + cu_seqlens_padded[0] = 0 + cu_seqlens_padded[1:] = torch.cumsum(padded_lengths, dim=0) + max_seqlen = torch.max(padded_lengths).to(dtype=torch.int32) + + new_sample["cu_seqlens_padded"] = cu_seqlens_padded + new_sample["max_seqlen"] = max_seqlen + + original_lengths = original_lengths.to( + device=dev, dtype=torch.int32, non_blocking=True + ).reshape(-1) + cu_seqlens = torch.empty(original_lengths.numel() + 1, device=dev, dtype=torch.int32) + cu_seqlens[0] = 0 + cu_seqlens[1:] = torch.cumsum(original_lengths, dim=0).reshape(-1) + new_sample["cu_seqlens"] = cu_seqlens + + if local_cp_size is not None: + new_sample["local_cp_size"] = local_cp_size + + return new_sample + + +def broadcast_tensor(item, src_rank, group) -> None: + """Broadcast a tensor from src_rank to all ranks in the group.""" + if item is not None: + torch.distributed.broadcast(item, src_rank, group=group) + + +def broadcast_scalars(values: List, group, dev, dtype=torch.float32) -> List: + """ + Broadcast scalar values from rank 0 to all ranks in the group. + + Args: + values: List of scalar values to broadcast (only used on rank 0). + group: The process group to broadcast within. + dev: The device to use for the tensor. + dtype: The data type for the tensor. + + Returns: + List of broadcasted values. + """ + if group.size() <= 1: + return values + + src_rank = torch.distributed.get_process_group_ranks(group)[0] + num_values = len(values) + + if group.rank() == 0: + info_to_broadcast = torch.tensor(values, dtype=dtype, device=dev) + else: + info_to_broadcast = torch.zeros(num_values, dtype=dtype, device=dev) + + broadcast_tensor(info_to_broadcast, src_rank, group) + + if group.rank() != 0: + values = info_to_broadcast.cpu().tolist() + + return values + + +def create_data_iterator( + new_samples, tp_group, config, vpp_needs_data=None, is_dynamic_cp: bool = False +): + """Handle virtual pipeline parallelism. + + For VPP, each PP rank needs a list of data iterators (one per VPP stage). + VPP stages that need full data (first/last pipeline stage, or MTP) get + full samples; others get metadata only (cu_seqlens, cu_seqlens_padded, + max_seqlen). + + Args: + new_samples: The packed samples after scheduling. + tp_group: Tensor parallel process group. + config: Model parallel config. + vpp_needs_data: A list of booleans (one per VPP stage) indicating which + VPP stages need full samples (data fields). None if VPP is disabled. + """ + if ( + config.virtual_pipeline_model_parallel_size is not None + and config.virtual_pipeline_model_parallel_size > 1 + ): + vpp_size = config.virtual_pipeline_model_parallel_size + if tp_group.rank() == 0: + metadata_keys = ["max_seqlen", "cu_seqlens", "cu_seqlens_padded"] + if is_dynamic_cp: + metadata_keys.append("local_cp_size") + new_data_iterator = [] + for i in range(vpp_size): + if vpp_needs_data is not None and vpp_needs_data[i]: + # Give each data-carrying VPP stage its own shallow-copied + # sample dicts. + samples_copy = [dict(sample) for sample in new_samples] + new_data_iterator.append(RerunDataIterator(iter(samples_copy))) + else: + # Create independent metadata dicts to avoid shared-reference mutation + metadata = [ + {k: sample[k] for k in metadata_keys if k in sample} + for sample in new_samples + ] + new_data_iterator.append(RerunDataIterator(iter(metadata))) + else: + new_data_iterator = [None for _ in range(vpp_size)] + else: + new_data_iterator = RerunDataIterator(iter(new_samples)) if tp_group.rank() == 0 else None + + return new_data_iterator + + +def reroute_samples_to_dcp_ranks( + batch, + global_ids_this_rank, + global_id_seqlens, + sample_id_groups, + offsets, + dp_group, + tp_group, + dp_cp_group, + total_dcp_gpus, +): + """ + Reroutes the sub-samples to the correct rank after scheduling. + + For each key in the batch dict, we perform an all-to-all communication + to transfer the data to the correct ranks. + """ + + def _gid_to_src_rank(gid: int) -> int: + dp_src_rank = torch.bucketize(gid, offsets[1:] - 1) + dcp_rank = ( + torch.distributed.get_process_group_ranks(dp_group)[dp_src_rank] // tp_group.size() + ) % dp_cp_group.size() + return dcp_rank + + gid2local_id = {int(gid): i for i, gid in enumerate(global_ids_this_rank)} + dcp_rank = dp_cp_group.rank() + dp_ranks = torch.distributed.get_process_group_ranks(dp_group) + dp_ranks = [(r // tp_group.size()) % dp_cp_group.size() for r in dp_ranks] + + data_keys = batch[0].keys() + + # Create the send plan + combined_sample_id_groups: List[List[int]] = [[] for _ in range(total_dcp_gpus)] + for d in range(total_dcp_gpus): + for sample_id_group in sample_id_groups: + combined_sample_id_groups[d].extend(sample_id_group[d]) + for dest_rank in range(total_dcp_gpus): + combined_sample_id_groups[dest_rank].sort() + + send_ids_sorted = [ + gid for d in dp_ranks for gid in combined_sample_id_groups[d] if gid in global_ids_this_rank + ] + + send_num_split = [0] * total_dcp_gpus + send_lens_split = [0] * total_dcp_gpus + for dest_rank in range(total_dcp_gpus): + if dest_rank in dp_ranks: + send_seq_lens = [ + global_id_seqlens[gid][1] + for gid in combined_sample_id_groups[dest_rank] + if gid in global_ids_this_rank + ] + send_num_split[dest_rank] = len(send_seq_lens) + send_lens_split[dest_rank] = sum(send_seq_lens) + else: + send_lens_split[dest_rank] = 0 + + # Create the recv plan + recv_sample_id_groups = [[] for _ in range(total_dcp_gpus)] + for gid in combined_sample_id_groups[dcp_rank]: + src_rank = _gid_to_src_rank(gid) + recv_sample_id_groups[src_rank].append(gid) + + recv_lens_split = [0] * total_dcp_gpus + for src_rank in range(total_dcp_gpus): + recv_lens_split[src_rank] = sum( + [global_id_seqlens[gid][1] for gid in recv_sample_id_groups[src_rank]] + ) + + recv_ids_sorted = [gid for d in range(total_dcp_gpus) for gid in recv_sample_id_groups[d]] + recv_counts = [len(recv_sample_id_groups[d]) for d in range(total_dcp_gpus)] + + recv_samples = [{k: None for k in data_keys} for _ in range(sum(recv_counts))] + + def _pack_sample_by_key(key: str) -> torch.Tensor: + flattened_tensors = [] + for gid in send_ids_sorted: + t = batch[gid2local_id[gid]][key].to(torch.cuda.current_device(), non_blocking=True) + flattened_tensors.append(t.reshape(-1)) + return ( + torch.cat(flattened_tensors, dim=0) + if flattened_tensors + else torch.empty(0, device=torch.cuda.current_device(), dtype=batch[0][key].dtype) + ) + + def _unpack_sample_by_key(key: str, recv_tensor: torch.Tensor): + cursor = 0 + for i, gid in enumerate(recv_ids_sorted): + sample_len = ( + 1 if key in ["original_seq_len", "padded_seq_len"] else global_id_seqlens[gid][1] + ) + recv_samples[i][key] = recv_tensor[cursor : cursor + sample_len] + cursor += sample_len + + for key in data_keys: + output_split_sizes, input_split_sizes = ( + (recv_counts, send_num_split) + if key in ["original_seq_len", "padded_seq_len"] + else (recv_lens_split, send_lens_split) + ) + send_tensor = _pack_sample_by_key(key) + recv_tensor_size = sum(output_split_sizes) + recv_tensor = torch.empty( + recv_tensor_size, device=torch.cuda.current_device(), dtype=send_tensor.dtype + ) + torch.distributed.all_to_all_single( + output=recv_tensor, + input=send_tensor, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=dp_cp_group, + ) + _unpack_sample_by_key(key, recv_tensor) + + recv_sample_with_id = {recv_id: recv_samples[i] for i, recv_id in enumerate(recv_ids_sorted)} + return recv_sample_with_id + + +def build_packed_microbatches( + samples_this_rank_with_id: Dict[int, Dict[str, torch.Tensor]], + sample_id_groups: List[List[List[int]]], + dcp_rank: int, + dev: torch.device, + is_dynamic_cp: bool = False, +) -> List[Dict[str, torch.Tensor]]: + """Build packed samples for each microbatch. + + Args: + samples_this_rank_with_id: Mapping from global sample ID to sample dict, + as returned by reroute_samples_to_dcp_ranks. + sample_id_groups: Per-microbatch, per-rank lists of global sample IDs. + dcp_rank: This rank's index within the DP×CP group. + dev: Target device. + is_dynamic_cp: Whether dynamic context parallel is enabled. + """ + num_micro_batches = len(sample_id_groups) + seg_starts: List[int] = [0] + original_lens_tensors = [] + padded_lens_tensors = [] + + grouped_samples = [ + [ + samples_this_rank_with_id[sub_sample_id] + for sub_sample_id in sample_id_groups[i][dcp_rank] + ] + for i in range(num_micro_batches) + ] + + local_cp_sizes_gpu = None + if is_dynamic_cp: + local_cp_sizes_cpu: List[int] = [] + for i in range(num_micro_batches): + sample_ids_this_group = sample_id_groups[i][dcp_rank] + local_cp_sizes_cpu.append( + len( + [ + 1 + for sample_ids in sample_id_groups[i] + if sample_ids_this_group[0] in sample_ids + ] + ) + ) + local_cp_sizes_gpu = torch.tensor(local_cp_sizes_cpu, dtype=torch.int32, device=dev) + + for i in range(num_micro_batches): + samples = grouped_samples[i] + seg_starts.append(seg_starts[-1] + len(samples)) + original_lens_tensors.extend([s["original_seq_len"].reshape(-1) for s in samples]) + padded_lens_tensors.extend([s["padded_seq_len"].reshape(-1) for s in samples]) + + padded_lens_all_gpu = torch.cat(padded_lens_tensors, dim=0).to(dtype=torch.int32) + original_lens_all_gpu = torch.cat(original_lens_tensors, dim=0).to(dtype=torch.int32) + + new_samples: List[Dict[str, torch.Tensor]] = [] + for i in range(num_micro_batches): + samples = grouped_samples[i] + lens_padded = padded_lens_all_gpu[seg_starts[i] : seg_starts[i + 1]] + lens_original = original_lens_all_gpu[seg_starts[i] : seg_starts[i + 1]] + local_cp_size = local_cp_sizes_gpu[i] if is_dynamic_cp else None + new_sample = _pack_sequences(samples, lens_padded, lens_original, local_cp_size, dev) + new_samples.append(new_sample) + + return new_samples + + +def get_batch_and_global_seqlens(data_iterator, num_microbatches, dp_group): + """ + Get the batch and global sequence lengths. + Each DP rank loads the same number of sequences, so we need to gather the sequence + lengths from all ranks then we can schedule the sequences into groups. + Args: + data_iterator: The data iterator. + num_microbatches: The number of microbatches. + dp_group: The data parallel group. + + Returns: + batch: The batch. + global_id_seqlens: The global sequence lengths. + global_ids_this_rank: The global IDs locally present on this rank. + """ + + batch_list = [next(data_iterator) for _ in range(num_microbatches)] + + batch = [] + for item in batch_list: + if isinstance(item, dict): + batch.append(item) + elif isinstance(item, list): + batch.extend(item) + else: + raise ValueError(f"Invalid item type: {type(item)}") + + # in sft_dataset.py, sequences are already packed before rescheduling, + # so we need to unpack them here and repack after rescheduling. + # This is only to adapt to the current megatron-lm sft_dataset. + # If you implement your own dataset, just have __getitem__ return List[Dict] + # and this step can be skipped. + batch = _unpack_batch(batch) + + subsample_seqlens = torch.cat([sample["padded_seq_len"] for sample in batch]).to( + dtype=torch.int32, device=torch.cuda.current_device() + ) + + global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered = ( + _get_global_seqlens_and_ids(subsample_seqlens, dp_group) + ) + + return batch, global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered + + +# ============================================================================= +# Dynamic CP scheduling algorithms (used by DefaultDynamicCPScheduler) +# ============================================================================= + + +def next_hdp_group( + sample_seqlens: List[Tuple[int, int]], + compute_estimator: Callable[[int], float], + total_gpus: int, + gpus_needed_fn: Callable[[int], int], + make_buckets_equal_fn: Callable, + max_seq_len_per_rank: float, + get_total_workload_fn: Callable, + delta: float = 0.05, + strategy: str = "dp", + eps_bucket: float = 0.10, +) -> Tuple[List[List[int]], List[Tuple[int, int]], List[float], List[List[int]]]: + """Form one balanced micro-batch group across DPxCP ranks. + + This is a standalone version of the scheduling algorithm extracted from + DefaultDynamicCPScheduler so it can live in a utils module. + + Extra args compared to the method version: + gpus_needed_fn: callable(seq_len) -> int + make_buckets_equal_fn: callable(sample_seqlens, compute_estimator) -> list[deque] + max_seq_len_per_rank: max tokens per rank for packing + get_total_workload_fn: callable(seq_len, cp_size) -> float + """ + if not sample_seqlens: + return ( + [[] for _ in range(total_gpus)], + [], + [0.0 for _ in range(total_gpus)], + [[] for _ in range(total_gpus)], + ) + + buckets = make_buckets_equal_fn(sample_seqlens, compute_estimator) + + micro_batches = [[] for _ in range(total_gpus)] + exec_times = [0.0 for _ in range(total_gpus)] + sample_ids_per_gpu = [[] for _ in range(total_gpus)] + packing_sequence_len = {} + + gpu_group_id = [None] * total_gpus + group_members = {} + group_size = {} + next_gid = 0 + + pp_cursor = 0 + prev_needed = None + check_balance = False + + while buckets: + sample_seq_tuple = bucket_idx = None + needed = None + + scan_order = ( + range(len(buckets)) + if strategy == "dp" + else [(pp_cursor + i) % len(buckets) for i in range(len(buckets))] + ) + + for idx in scan_order: + if not buckets[idx]: + continue + cand_tuple = buckets[idx][0] + cand_seq_len = cand_tuple[1] + needed = gpus_needed_fn(cand_seq_len) + + candidate_gids = [gid for gid, sz in group_size.items() if sz == needed] + free_ranks = [r for r, gid in enumerate(gpu_group_id) if gid is None] + if candidate_gids or len(free_ranks) >= needed: + sample_seq_tuple, bucket_idx = cand_tuple, idx + break + + if sample_seq_tuple is None: + break + + if strategy == "pp": + pp_cursor = (bucket_idx + 1) % len(buckets) + + sample_id, seq_len = sample_seq_tuple + needed = gpus_needed_fn(seq_len) + if prev_needed is None: + prev_needed = needed + + candidate_gids = [ + gid + for gid, sz in group_size.items() + if sz == needed and packing_sequence_len[gid] + seq_len / needed <= max_seq_len_per_rank + ] + if candidate_gids: + best_gid, best_load = min( + ((gid, max(exec_times[r] for r in group_members[gid])) for gid in candidate_gids), + key=lambda t: t[1], + ) + else: + best_gid, best_load = None, float("inf") + + free_ranks = [r for r, gid in enumerate(gpu_group_id) if gid is None] + if len(free_ranks) >= needed: + free_sorted = sorted(free_ranks, key=lambda r: exec_times[r]) + new_members = free_sorted[:needed] + new_load = exec_times[new_members[-1]] + + if new_load < best_load: + best_gid = None + chosen_members = new_members + else: + chosen_members = group_members[best_gid] + else: + if best_gid is None: + break + chosen_members = group_members[best_gid] + + if best_gid is None: + best_gid = next_gid + next_gid += 1 + group_members[best_gid] = chosen_members + group_size[best_gid] = needed + for r in chosen_members: + gpu_group_id[r] = best_gid + + 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() + + while buckets and not buckets[0]: + buckets.pop(0) + pp_cursor %= max(1, len(buckets)) + + if needed < prev_needed: + check_balance = True + + if ( + check_balance + and buckets + and max(exec_times) - min(exec_times) <= delta * max(exec_times) + ): + break + + leftovers = [] + for b in buckets: + for sample_seq_tuple in b: + leftovers.append(sample_seq_tuple) + + def trim_overload(): + while True: + cur_max = max(exec_times) + cur_min = min(exec_times) + cur_slack = cur_max - cur_min + if cur_slack <= delta * cur_max: + break + if cur_min == 0: + break + + max_r = exec_times.index(cur_max) + gid = gpu_group_id[max_r] + members = group_members[gid] + + if not micro_batches[max_r] or len(micro_batches[max_r]) <= 1: + break + + seq = micro_batches[max_r][-1] + per_gpu_cost = compute_estimator(seq) + + proj_times = exec_times[:] + for r in members: + proj_times[r] -= per_gpu_cost + + proj_slack = max(proj_times) - min(proj_times) + + if proj_slack < cur_slack: + sample_id_to_remove = sample_ids_per_gpu[max_r][-1] + for r in members: + micro_batches[r].pop() + exec_times[r] -= per_gpu_cost + sample_ids_per_gpu[r].pop() + leftovers.append((sample_id_to_remove, seq)) + else: + break + + # TODO(tailaim): uncomment this to support different ranks have different num_microbatches + # trim_overload() + + total_work_before = sum(len(mb) for mb in micro_batches) + + def fill_empty_gpus(micro_batches, exec_times, sample_ids_per_gpu, group_members, group_size): + empty_gpus = [i for i in range(total_gpus) if not micro_batches[i]] + if not empty_gpus: + return (micro_batches, exec_times, sample_ids_per_gpu, group_members, group_size) + + existing_group_sizes = set(group_size.values()) + assert ( + existing_group_sizes + ), "There should be at least one group existing, cannot redistribute, " + "try to increase 'max-seqlen-per-dp-cp-rank'." + + min_group_size = min(existing_group_sizes) + next_power = min(min_group_size * 2, total_gpus) + + for gid, size in group_size.items(): + if size == min_group_size: + members = group_members[gid] + needed_count = next_power - min_group_size + group_start_gpu = members[0] + group_end_gpu = members[-1] + empty_gpu = [idx for idx, work in enumerate(micro_batches) if not work][0] + assert not all( + work for work in micro_batches[empty_gpu : empty_gpu + needed_count] + ), "Empty GPUs were detected but not enough to expand." + work_to_push = micro_batches[group_end_gpu + 1 : empty_gpu] + exec_times_to_push = exec_times[group_end_gpu + 1 : empty_gpu] + sample_ids_to_push = sample_ids_per_gpu[group_end_gpu + 1 : empty_gpu] + + new_micro_batches = [[]] * len(micro_batches) + new_exec_times = [0.0] * len(exec_times) + new_sample_ids_per_gpu = [[]] * len(sample_ids_per_gpu) + + for i in range(group_start_gpu): + new_micro_batches[i] = micro_batches[i] + new_exec_times[i] = exec_times[i] + new_sample_ids_per_gpu[i] = sample_ids_per_gpu[i] + + for i in range(group_start_gpu, group_end_gpu + needed_count + 1): + new_micro_batches[i] = micro_batches[group_end_gpu] + new_exec_times[i] = get_total_workload_fn( + micro_batches[group_end_gpu][0], next_power + ) + new_sample_ids_per_gpu[i] = sample_ids_per_gpu[group_end_gpu] + + for i, work in enumerate(work_to_push): + new_micro_batches[group_end_gpu + needed_count + 1 + i] = work + new_exec_times[group_end_gpu + needed_count + 1 + i] = exec_times_to_push[i] + new_sample_ids_per_gpu[group_end_gpu + needed_count + 1 + i] = ( + sample_ids_to_push[i] + ) + + group_size[gid] = next_power + group_members[gid] = list(range(members[0], members[-1] + needed_count + 1)) + for pushed_gid in group_size.keys(): + if pushed_gid > gid: + group_members[pushed_gid] = [ + x + needed_count for x in group_members[pushed_gid] + ] + + return ( + new_micro_batches, + new_exec_times, + new_sample_ids_per_gpu, + group_members, + group_size, + ) + + empty_gpus = any([not micro_batches[i] for i in range(total_gpus)]) + while empty_gpus: + micro_batches, exec_times, sample_ids_per_gpu, group_members, group_size = fill_empty_gpus( + micro_batches, exec_times, sample_ids_per_gpu, group_members, group_size + ) + empty_gpus = any([not micro_batches[i] for i in range(total_gpus)]) + + total_work_after = sum(len(mb) for mb in micro_batches) + assert ( + total_work_after >= total_work_before + ), f"Samples were removed: {total_work_before} -> {total_work_after}" + + return micro_batches, leftovers, exec_times, sample_ids_per_gpu + + +def align_sample_id_groups(sample_id_groups: List, microbatch_group_size_per_vp_stage: int) -> List: + """Align len(sample_id_groups) to microbatch_group_size_per_vp_stage when VPP is enabled. + + Standalone version extracted from DefaultDynamicCPScheduler. + """ + multiple = int(microbatch_group_size_per_vp_stage) + remainder = (-len(sample_id_groups)) % multiple + i = len(sample_id_groups) - 1 + + def split_group(sample_id_group): + total_hdp_ranks = len(sample_id_group) + cu_ranks = [0] + prev_cp_size = 0 + + while cu_ranks[-1] != total_hdp_ranks: + start_rank = cu_ranks[-1] + sid0 = sample_id_group[start_rank][0] + cp_size = 0 + for r in range(start_rank, total_hdp_ranks): + if sid0 in sample_id_group[r]: + cp_size += 1 + else: + break + assert ( + prev_cp_size == 0 or cp_size <= prev_cp_size + ), f"split_group: CP size is not decreasing: prev={prev_cp_size}, cur={cp_size}" + cu_ranks.append(start_rank + cp_size) + prev_cp_size = cp_size + if len(cu_ranks) == 2: + return None, None + + k = 0 + while cu_ranks[k] < total_hdp_ranks // 2: + k += 1 + + old_mb = sample_id_group[: cu_ranks[k]] + [[] for _ in range(total_hdp_ranks - cu_ranks[k])] + new_mb = sample_id_group[cu_ranks[k] :] + [[] for _ in range(cu_ranks[k])] + old_mb = fill_empty_by_expanding_cp(old_mb) + new_mb = fill_empty_by_expanding_cp(new_mb) + return new_mb, old_mb + + def fill_empty_by_expanding_cp(sample_id_group): + def fill_empty(sample_id_group): + empty_size = sum(1 for x in sample_id_group if len(x) == 0) + i = len(sample_id_group) - 1 - empty_size + prev_cp_size = 0 + while i >= 0: + sid0 = sample_id_group[i][0] + cp_size = 0 + while sid0 in sample_id_group[i] and i >= 0: + cp_size += 1 + i -= 1 + if cp_size > prev_cp_size and prev_cp_size != 0: + start_idx = i + 1 + cp_size + end_idx = -empty_size + prev_cp_size if -empty_size + prev_cp_size < 0 else None + sample_id_group[start_idx + 2 * prev_cp_size : end_idx] = sample_id_group[ + start_idx + prev_cp_size : -empty_size + ] + sample_id_group[start_idx + prev_cp_size : start_idx + 2 * prev_cp_size] = ( + sample_id_group[start_idx : start_idx + prev_cp_size] + ) + break + elif cp_size <= empty_size and i == -1: + end_idx = -empty_size + cp_size if -empty_size + cp_size < 0 else None + sample_id_group[2 * cp_size : end_idx] = sample_id_group[cp_size:-empty_size] + sample_id_group[cp_size : 2 * cp_size] = sample_id_group[0:cp_size] + break + prev_cp_size = cp_size + return sample_id_group + + while len(sample_id_group[-1]) == 0: + sample_id_group = fill_empty(sample_id_group) + return sample_id_group + + attempts_since_split = 0 + while remainder > 0: + if i < 0: + if attempts_since_split >= len(sample_id_groups): + assert False, '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 + + +# ============================================================================= +# Workload estimation helpers for dynamic CP scheduling +# ============================================================================= + + +@lru_cache(maxsize=128) +def dcp_gpus_needed(seq_len: int, max_seq_len_per_rank: int, min_cp_size: int = 1) -> int: + """Number of GPUs needed, rounded up to the next power of 2, lower-bounded by min_cp_size.""" + raw = max(1, 2 ** ceil(log2(seq_len / max_seq_len_per_rank))) + return max(min_cp_size, raw) + + +@lru_cache(maxsize=128) +def dcp_get_total_workload( + seq_length: int, max_seq_len_per_rank: int, cp_size: Optional[int] = None, min_cp_size: int = 1 +) -> float: + """Estimate workload of a sub-sample for scheduling balance.""" + if cp_size is None: + cp_size = dcp_gpus_needed(seq_length, max_seq_len_per_rank, min_cp_size) + return (seq_length * seq_length) / cp_size + + +def dcp_make_buckets_equal( + sample_seqlens: List[Tuple[int, int]], + compute_estimator: Callable, + max_seq_len_per_rank: int, + min_cp_size: int = 1, +) -> List[deque]: + """Split samples into buckets of roughly equal work, one per unique CP size.""" + seqlens = [seq_len for _, seq_len in sample_seqlens] + k = len({dcp_gpus_needed(L, max_seq_len_per_rank, min_cp_size) for L in seqlens}) + + work = [] + for _, s in sample_seqlens: + cp_size = dcp_gpus_needed(s, max_seq_len_per_rank, min_cp_size) + work.append(compute_estimator(s, cp_size)) + total_work = sum(work) + target = total_work / k + buckets, cur, cur_work = [], [], 0.0 + remaining_k = k + + for i, (sample_id, seq_len) in enumerate(sample_seqlens): + w = compute_estimator(seq_len) + projected = cur_work + w + if cur and ( + projected > target * 1.1 or len(sample_seqlens) - i <= remaining_k - len(buckets) + ): + buckets.append(deque(cur)) + cur, cur_work = [], 0.0 + remaining_k -= 1 + cur.append((sample_id, seq_len)) + cur_work += w + + if cur: + buckets.append(deque(cur)) + return buckets diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 42146d1acd2..82ef605aa6d 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/megatron/core/datasets/gpt_dataset.py @@ -58,8 +58,8 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig): Set to 0 if sequence parallel is not enabled regardless of TP size. """ - hybrid_context_parallel: bool = False - """Option to enable hybrid context parallelism. When setting this to True, + dynamic_context_parallel: bool = False + """Option to enable dynamic context parallelism. When setting this to True, each sample should be divisible by the data parallel size * context parallel size * 2. If sequence parallel is enabled, it should be divisible by the data parallel size * context parallel size * sequence parallel size * 2. diff --git a/megatron/core/datasets/readme.md b/megatron/core/datasets/readme.md index 58721b7471b..af3bccdebf4 100644 --- a/megatron/core/datasets/readme.md +++ b/megatron/core/datasets/readme.md @@ -192,6 +192,66 @@ To query the `BlendedDataset` for the _k_-th sample we do the following To save time during initialization, each index is built/cached sequentially on one process rank and subsequently loaded in parallel on other process ranks. The cached indices are unique to a hash generated in the `BlendedDataset.__init__` function. +## Packing Scheduler + +The packing scheduler re-schedules variable-length sequences across DP×CP ranks to improve GPU utilization. It is built around two modules: `data_schedule.py` (high-level logic and entry points) and `data_schedule_utils.py` (utility functions). + +### Call Hierarchy + +The scheduling pipeline has two phases connected by the data iterator: `wrap_data_iterator` consumes the **original** data iterator, performs global-batch scheduling, and produces a **wrapped** (packed) data iterator; `get_batch_on_this_rank_for_sequence_packing` then consumes this **wrapped** data iterator to fetch individual packed microbatches during training. + +``` + original wrapped (packed) + data_iterator data_iterator + │ │ + ▼ ▼ + ┌────────────────────────┐ ┌────────────────────────────────────┐ + │ wrap_data_iterator() │ │ get_batch_on_this_rank_for_ │ +Phase 1 │ (once per global │ ────────► │ sequence_packing() │ Phase 2 +(scheduling) │ batch) │ returns │ (once per microbatch, │ (fetching) + │ │ wrapped │ called by training loop) │ + └───────────┬────────────┘ iterator └──────────────┬─────────────────────┘ + │ │ + ▼ ▼ + DpBalancedScheduler.run() next(wrapped_data_iterator) + │ ├─ get_thd_partitioned_indices() [TE] + ├─ get_batch_and_global_seqlens() [utils] ├─ broadcast_tensor() [utils] + ├─ get_groups_and_subsamples() └─ PackedSeqParams(...) + ├─ reroute_samples_to_dcp_ranks() [utils] + ├─ build_packed_microbatches() [utils] + ├─ broadcast_scalars() [utils] + └─ create_data_iterator() [utils] +``` + +### `data_schedule.py` + +#### Entry Points + +- **`wrap_data_iterator(original_data_iterator) → wrapped_data_iterator`** — Top-level entry point called once per global batch. Takes the **original** data iterator as input, resolves the scheduler class from `scheduler_map`, instantiates it, and delegates to `scheduler.run()` which consumes all microbatches from the original iterator, re-schedules them, and produces a **wrapped** (packed) data iterator along with the updated `num_microbatches` and FLOPs statistics. + +- **`get_batch_on_this_rank_for_sequence_packing(wrapped_data_iterator)`** — Per-microbatch entry point called by the training loop. Takes the **wrapped** data iterator returned by `wrap_data_iterator` as input. Fetches one packed microbatch via `next(wrapped_data_iterator)`, broadcasts batch fields across TP ranks, optionally partitions sequences across CP ranks using Transformer Engine's `thd_get_partitioned_indices`, and constructs `PackedSeqParams` (with `cu_seqlens`, `max_seqlen`, `qkv_format=thd`). + +#### Scheduler Classes + +- **`BasePackingScheduler`** — Abstract base class. Defines the interface: + - `get_groups_and_subsamples()` — pure scheduling algorithm (must be overridden). + - `run()` — full pipeline: fetch → schedule → reroute → pack → broadcast → VPP handling. + +- **`DpBalancedScheduler(BasePackingScheduler)`** — Concrete scheduler that packs sequences in their original order until reaching `max_seqlen_per_dp_cp_rank × cp_size`. Aligns the number of microbatches to `dp_size` (and VPP stage multiples when applicable). + +### `data_schedule_utils.py` + +Utility functions consumed by the schedulers above: + +| Function | Role | +|---|---| +| `get_batch_and_global_seqlens()` | Fetch `num_microbatches` batches from the data iterator and all-gather sequence lengths across DP ranks. | +| `reroute_samples_to_dcp_ranks()` | All-to-all communication to transfer sub-samples to their scheduled DP×CP rank. | +| `build_packed_microbatches()` | Concatenate sub-samples within each microbatch group and produce `cu_seqlens`. | +| `broadcast_scalars()` | Broadcast scalar values (e.g. `num_microbatches`, FLOPs stats) across a process group. | +| `broadcast_tensor()` | Broadcast a single tensor within a process group. | +| `create_data_iterator()` | Wrap packed sample lists into a data iterator; handles VPP stage splitting. | + ## Offline cache preparation For GPT-style training, the dataset caches described above can be prepared ahead of time with `tools/prepare_cache.py` instead of waiting for rank 0 to build them during training startup. diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index b84565dd1f3..7381cb8883c 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1791,21 +1791,22 @@ def forward( """Forward.""" if packed_seq_params is not None: # If Dynamic CP group is provided, update TE DPA CP group - if packed_seq_params.cp_group is not None: - self.cp_group = packed_seq_params.cp_group - super().set_context_parallel_group( - self.cp_group, - torch.distributed.get_process_group_ranks(self.cp_group), - TEDotProductAttention.cp_stream, - self.cp_comm_type, - ) - # If cp_group is None but local_cp_size is provided, - # Indicates to turn off CP dynamically - elif packed_seq_params.local_cp_size is not None: - assert ( - packed_seq_params.local_cp_size == 1 - ), "local_cp_size must be == 1 if provided without cp_group" - super().set_context_parallel_group(None, None, None, self.cp_comm_type) + if packed_seq_params.local_cp_size is not None: + if packed_seq_params.local_cp_size == 1: + super().set_context_parallel_group(None, None, None, self.cp_comm_type) + else: + assert ( + packed_seq_params.cp_group is not None + ), "cp_group is not set in packed_seq_params for dynamic CP" + self.cp_group = packed_seq_params.cp_group + if TEDotProductAttention.cp_stream is None: + TEDotProductAttention.cp_stream = torch.cuda.Stream() + super().set_context_parallel_group( + self.cp_group, + torch.distributed.get_process_group_ranks(self.cp_group), + TEDotProductAttention.cp_stream, + self.cp_comm_type, + ) self.kept_packed_seq_params.discard("cp_group") self.kept_packed_seq_params.discard("local_cp_size") @@ -3383,3 +3384,24 @@ def set_save_original_input(module): from transformer_engine.pytorch.float8_tensor import Float8Tensor except ImportError: Float8Tensor = None + + +def get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank): + """Get partitioned indices for THD format data in context parallel. + + Args: + cu_seqlens: Cumulative sequence lengths tensor. + total_tokens: Total number of tokens. + cp_size: Context parallel world size. + cp_rank: Context parallel rank. + + Returns: + Partitioned indices tensor. + """ + assert is_te_min_version("1.10.0"), ( + "Please update Transformer Engine to >= 1.10 to use " + "Context Parallel with THD format data" + ) + import transformer_engine_torch as tex + + return tex.thd_get_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 5170c692ac5..18706ffcb03 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -59,14 +59,29 @@ class ModelParallelConfig: can handle without overflowing the memory. Typically, a good starting point is to set this to maximum sequence length / context parallel size. This is used to calculate the number and length of sub-samples assigned to - each rank when using hybrid_context_parallel. + each rank when sequence_packing_scheduler is not None. """ - hybrid_context_parallel: bool = False + dynamic_context_parallel: bool = False """ - If true, enables hybrid context parallel. This is used to balance the workload of + If true, enables dynamic context parallel. This is used to balance the workload of each CP rank when we use packed samples with variable sequence lengths. - Please set max_seqlen_per_dp_cp_rank when using hybrid_context_parallel. + Dynamic CP forms variable-sized CP groups from the DPxCP ranks dynamically. + Please set max_seqlen_per_dp_cp_rank. + """ + + min_dynamic_context_parallel_size: int = 1 + """Minimum CP group size for dynamic context parallel. Default 1 (no CP). + The maximum is dp_size * context_parallel_size (the full DPxCP group).""" + + hybrid_context_parallel: bool = False + """Deprecated. Use ``dynamic_context_parallel`` instead.""" + + sequence_packing_scheduler: Optional[Literal['dp_balanced', 'default_dynamic_cp']] = None + """ + Scheduler for sequence packing and dynamic context parallel. + dp_balanced: DP-balanced scheduler for sequence packing. + default_dynamic_cp: Dynamic-CP scheduler for packed sequence balancing. """ expert_model_parallel_size: int = 1 @@ -418,6 +433,34 @@ def __post_init__(self): See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more details. """ + if self.hybrid_context_parallel: + warnings.warn( + "hybrid_context_parallel is deprecated and will be removed in a future release. " + "Use dynamic_context_parallel instead.", + DeprecationWarning, + ) + if self.dynamic_context_parallel: + raise ValueError( + "Cannot set both hybrid_context_parallel and dynamic_context_parallel. " + "Please use dynamic_context_parallel only." + ) + self.dynamic_context_parallel = True + + if self.dynamic_context_parallel: + if self.sequence_packing_scheduler is None: + self.sequence_packing_scheduler = 'default_dynamic_cp' + if self.sequence_packing_scheduler != 'default_dynamic_cp': + raise ValueError( + 'Dynamic context parallelism requires ' + 'sequence_packing_scheduler=default_dynamic_cp' + ) + + if self.min_dynamic_context_parallel_size < 1: + raise ValueError( + f"min_dynamic_context_parallel_size must be >= 1, " + f"got {self.min_dynamic_context_parallel_size}" + ) + if self.sequence_parallel: if self.tensor_model_parallel_size <= 1: raise ValueError("Cannot use sequence parallelism without tensor parallelism") diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index ac2e3f8bab1..080beb60ebe 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -20,7 +20,7 @@ RotaryEmbedding, ) from megatron.core.models.common.language_module.language_module import LanguageModule -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( FineGrainedActivationOffloadingInterface as off_interface, ) @@ -671,6 +671,7 @@ def _postprocess( self._decoder_hidden_states_cache = hidden_states else: # In training/eval, use the utility function for processing MTP loss/scaling. + mtp_cp_group = resolve_cp_group(self.pg_collection.cp, packed_seq_params) hidden_states = process_mtp_loss( hidden_states=hidden_states, labels=labels, @@ -681,7 +682,7 @@ def _postprocess( is_training=self.training, compute_language_model_loss=self.compute_language_model_loss, config=self.config, - cp_group=self.pg_collection.cp, + cp_group=mtp_cp_group, tp_group=self.tp_group, packed_seq_params=packed_seq_params, scale_logits_fn=self._scale_logits if self.config.use_mup else None, diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 322f12a4122..b1b4275fee1 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -64,3 +64,17 @@ def __post_init__(self): .to(torch.int32) .unsqueeze(0) # Add a batch dimension ) + + +def resolve_cp_group( + static_cp_group: dist.ProcessGroup, packed_seq_params: PackedSeqParams = None +) -> dist.ProcessGroup: + """Return the dynamic CP group from packed_seq_params when available, else the static one. + + Dynamic CP assigns a per-microbatch CP group that may differ from the + process-group stored at model construction time. This helper centralises + the resolution logic used by GPTModel, GatedDeltaNet, and MTP layers. + """ + if packed_seq_params is not None and packed_seq_params.cp_group is not None: + return packed_seq_params.cp_group + return static_cp_group diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py index 337485b4d12..863b5d55d9d 100644 --- a/megatron/core/parallel_state.py +++ b/megatron/core/parallel_state.py @@ -115,8 +115,8 @@ _CONTEXT_PARALLEL_GLOBAL_RANKS = None # Hierarchical context parallel groups _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS = None -# Hybrid context parallel groups -_HYBRID_DP_CP_GROUPS = {} +# Dynamic context parallel groups +_DYNAMIC_DP_CP_GROUPS = {} # Data parallel group information with context parallel combined. _DATA_PARALLEL_GROUP_WITH_CP = None @@ -418,29 +418,27 @@ def create_hierarchical_groups( return hierarchical_groups, hierarchical_groups_gloo -def create_hybrid_dp_cp_groups(rank, ranks, pg_options): +def create_dynamic_dp_cp_groups(rank, ranks, pg_options, min_cp_size=1): """ - Creates groups required for hybrid DPxCP. - Creates a new group for every power of 2 up to the number of DPxCP ranks. + Creates groups required for dynamic DPxCP. + Creates a new group for every power of 2 from min_cp_size up to len(ranks). Returns a dictionary indexed by group size. """ - hybrid_dp_cp_groups = {} - # Generate group for every power of 2 up to the number of CP ranks - # We limit the allowed group sizes in order to avoid excessive overhead. - group_sizes = [2**i for i in range(int(log2(len(ranks))))][1:] + dynamic_dp_cp_groups = {} + group_sizes = [2**i for i in range(int(log2(len(ranks)))) if 2**i >= min_cp_size] for group_size in group_sizes: for i in range(0, len(ranks), group_size): group = create_group( ranks[i : i + group_size], pg_options=pg_options, - group_desc=f"HYBRID_DP_CP_GROUP_{group_size}", + group_desc=f"DYNAMIC_DP_CP_GROUP_{group_size}", ) if rank in ranks[i : i + group_size]: assert ( - group_size not in hybrid_dp_cp_groups - ), f"Rank {rank} appears in multiple Hybrid DP CP groups of size {group_size}" - hybrid_dp_cp_groups[group_size] = group - return hybrid_dp_cp_groups + group_size not in dynamic_dp_cp_groups + ), f"Rank {rank} appears in multiple Dynamic DP CP groups of size {group_size}" + dynamic_dp_cp_groups[group_size] = group + return dynamic_dp_cp_groups class RankGenerator(object): @@ -552,7 +550,8 @@ def initialize_model_parallel( use_sharp: bool = False, context_parallel_size: int = 1, hierarchical_context_parallel_sizes: Optional[List[int]] = None, - hybrid_context_parallel: bool = False, + dynamic_context_parallel: bool = False, + min_dynamic_context_parallel_size: int = 1, expert_model_parallel_size: int = 1, num_distributed_optimizer_instances: int = 1, expert_tensor_parallel_size: Optional[int] = None, @@ -919,18 +918,34 @@ def initialize_model_parallel( if "NCCL_COLLNET_ENABLE" in os.environ: del os.environ["NCCL_COLLNET_ENABLE"] - if hybrid_context_parallel: - global _HYBRID_DP_CP_GROUPS + if dynamic_context_parallel: + # TODO: Are gloo groups needed for Dynamic CP? + global _DYNAMIC_DP_CP_GROUPS for ranks_with_cp in decoder_rank_generator.get_ranks('dp-cp'): assert ( len(ranks_with_cp) % 2 == 0 - ), "Hybrid context parallel requires an even number of ranks" - _HYBRID_DP_CP_GROUPS.update( - create_hybrid_dp_cp_groups( - rank, ranks_with_cp, get_nccl_options("dp_cp", nccl_comm_cfgs) + ), "Dynamic context parallel requires an even number of ranks" + _DYNAMIC_DP_CP_GROUPS.update( + create_dynamic_dp_cp_groups( + rank, + ranks_with_cp, + get_nccl_options("dp_cp", nccl_comm_cfgs), + min_cp_size=min_dynamic_context_parallel_size, ) ) - # TODO: Are gloo groups needed for hybrid cp? + + data_parallel_size_with_cp = data_parallel_size * context_parallel_size + group_sizes = [ + 2**i + for i in range(int(log2(data_parallel_size_with_cp))) + if 2**i >= min_dynamic_context_parallel_size + ] + if data_parallel_size_with_cp not in group_sizes: + group_sizes.append(data_parallel_size_with_cp) + for group_size in group_sizes: + group = get_dynamic_data_context_parallel_groups(group_size=group_size) + torch.distributed.barrier(group=group, device_ids=[torch.cuda.current_device()]) + torch.cuda.synchronize() for ranks in decoder_rank_generator.get_ranks('dp'): group = create_group( @@ -1523,16 +1538,15 @@ def get_hierarchical_context_parallel_groups(check_initialized=True): return _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS -def get_hybrid_data_context_parallel_groups(check_initialized=True, group_size=None): - """Get the hybrid context parallel groups the caller rank belongs to.""" - # If the group size is the same as the entire DPxCP group, return the original group +def get_dynamic_data_context_parallel_groups(check_initialized=True, group_size=None): + """Get the dynamic context parallel groups the caller rank belongs to.""" if get_data_parallel_world_size(with_context_parallel=True) == group_size: if check_initialized: assert _DATA_PARALLEL_GROUP_WITH_CP is not None return _DATA_PARALLEL_GROUP_WITH_CP if check_initialized: - assert _HYBRID_DP_CP_GROUPS is not None - return _HYBRID_DP_CP_GROUPS[group_size] + assert _DYNAMIC_DP_CP_GROUPS is not None + return _DYNAMIC_DP_CP_GROUPS[group_size] def get_embedding_group(check_initialized=True): @@ -2113,6 +2127,9 @@ def destroy_model_parallel(): global _CONTEXT_PARALLEL_GLOBAL_RANKS _CONTEXT_PARALLEL_GLOBAL_RANKS = None + global _DYNAMIC_DP_CP_GROUPS + _DYNAMIC_DP_CP_GROUPS = {} + global _EMBEDDING_GROUP _EMBEDDING_GROUP = None diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index 36ccb2df98b..c5a582a2403 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -39,7 +39,6 @@ combined_1f1b_schedule_for_interleaved_pipelining, combined_1f1b_schedule_for_no_pipelining, ) -from .hybrid_cp_schedule import hybrid_context_parallel_forward_backward # Types Shape = Union[List[int], torch.Size] @@ -719,24 +718,6 @@ def forward_backward_no_pipelining( total_num_tokens, partial(check_first_val_step, first_val_step, forward_only), ) - elif config.hybrid_context_parallel: - forward_data_store, total_num_tokens = hybrid_context_parallel_forward_backward( - forward_step_func, - data_iterator, - model, - num_microbatches, - input_tensor, - output_tensor_grad, - forward_data_store, - config, - collect_non_loss_data, - first_val_step, - forward_only, - no_sync_func, - total_num_tokens, - check_first_val_step, - model_type, - ) else: with no_sync_func(): for i in range(num_microbatches - 1): diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 2521145c467..9ad4533e7dd 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -21,7 +21,7 @@ from megatron.core.fp8_utils import get_fp8_align_size from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.mamba_context_parallel import ( _all_to_all_cp2hp, @@ -57,6 +57,12 @@ logger = logging.getLogger(__name__) +# Triton's autotune key for causal_conv1d includes cdiv(total_tokens, 1024). +# Dynamic CP causes total_tokens to vary per microbatch, triggering repeated +# autotuning. Aligning to this boundary collapses most variations into a +# small number of buckets. +_CONV_PAD_ALIGNMENT = 4096 + @dataclass class GatedDeltaNetSubmodules: @@ -140,6 +146,22 @@ def __init__( self.qk_dim_local_tp = self.qk_dim // self.tp_size self.v_dim_local_tp = self.v_dim // self.tp_size + # GDN uses head-parallel CP: each CP rank handles a slice of heads. + # The static cp_size (== max dynamic cp_size) must evenly divide the + # per-TP head counts so that every possible runtime cp_size also divides. + num_key_heads_per_tp = self.num_key_heads // self.tp_size + num_value_heads_per_tp = self.num_value_heads // self.tp_size + assert num_key_heads_per_tp % self.cp_size == 0, ( + f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_key_heads per TP rank ({num_key_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + assert num_value_heads_per_tp % self.cp_size == 0, ( + f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_value_heads per TP rank ({num_value_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + # Input projection (hidden_states -> q, k, v, gate, beta, alpha) # TODO: for now, output gate is forced for GDN. # We may remove this restriction in the future. @@ -297,8 +319,11 @@ def forward( inference_context = deprecate_inference_params(inference_context, inference_params) + cp_group = resolve_cp_group(self.pg_collection.cp, packed_seq_params) + cp_size = cp_group.size() + seq_len, batch, _ = hidden_states.shape - seq_len = seq_len * self.sp_size * self.cp_size + seq_len = seq_len * self.sp_size * cp_size if inference_context is not None: assert ( @@ -348,7 +373,7 @@ def forward( nvtx_range_pop(suffix="in_proj") # CP All to All: CP to HP - if self.cp_size > 1: + if cp_size > 1: # # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a. head_perm = _build_head_perm_for_split_sections( ( @@ -359,7 +384,7 @@ def forward( self.num_value_heads // self.tp_size, self.num_value_heads // self.tp_size, ), - self.pg_collection.cp.size(), + cp_size, torch.cuda.current_device(), ) qkvzba = qkvzba.index_select(-1, head_perm) @@ -368,21 +393,19 @@ def forward( qkvzba, seq_dim=0, head_dim=-1, - cp_group=self.pg_collection.cp, + cp_group=cp_group, undo_attention_load_balancing=False, ) - if self.cp_size > 1: + if cp_size > 1: # Permute at the seq dim so that a single unsectioned a2a # is equivalent to per-sequence a2a. # This also folds the ``_undo_attention_load_balancing`` step. thd_cp_a2a_idx, thd_cp_a2a_inv = _build_thd_cp_a2a_perm( - cu_seqlens_q, self.cp_size, seq_len + cu_seqlens_q, cp_size, seq_len ) qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx) else: - qkvzba = tensor_a2a_cp2hp( - qkvzba, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) + qkvzba = tensor_a2a_cp2hp(qkvzba, seq_dim=0, head_dim=-1, cp_group=cp_group) # Transpose: s b x --> b s x # From sbhd to bshd format @@ -392,10 +415,10 @@ def forward( qkv, gate, beta, alpha = torch.split( qkvzba, [ - (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size, - self.v_dim_local_tp // self.cp_size, - self.num_value_heads // self.tp_size // self.cp_size, - self.num_value_heads // self.tp_size // self.cp_size, + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size, + self.v_dim_local_tp // cp_size, + self.num_value_heads // self.tp_size // cp_size, + self.num_value_heads // self.tp_size // cp_size, ], dim=-1, ) @@ -412,16 +435,13 @@ def forward( self.v_dim_local_tp, ] conv1d_weight = get_parameter_local_cp( - self.conv1d.weight, - dim=0, - cp_group=self.pg_collection.cp, - split_sections=qkv_channels_split_sections, + self.conv1d.weight, dim=0, cp_group=cp_group, split_sections=qkv_channels_split_sections ) conv1d_bias = ( get_parameter_local_cp( self.conv1d.bias, dim=0, - cp_group=self.pg_collection.cp, + cp_group=cp_group, split_sections=qkv_channels_split_sections, ) if self.conv_bias @@ -436,36 +456,47 @@ def forward( stride=self.conv1d.stride, padding=self.conv1d.padding, dilation=self.conv1d.dilation, - groups=self.conv_dim_local_tp // self.cp_size, + groups=self.conv_dim_local_tp // cp_size, ) qkv = self.act_fn(conv_out[..., :seq_len]) qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d else: assert self.activation in ["silu", "swish"] + _orig_seq = qkv.shape[1] + _pad_n = -_orig_seq % _CONV_PAD_ALIGNMENT + _conv_input = qkv + _conv_cu_seqlens = cu_seqlens_q + if _pad_n > 0: + _conv_input = torch.nn.functional.pad(qkv, (0, 0, 0, _pad_n)) + # cu_seqlens_q is None in non-packed-sequence mode; only the + # last-segment offset needs to grow to cover the padding tail. + if cu_seqlens_q is not None: + _conv_cu_seqlens = cu_seqlens_q.clone() + _conv_cu_seqlens[-1] += _pad_n qkv, _ = causal_conv1d( - x=qkv, # FLA conv1d accepts [b, s, d] format input + x=_conv_input, # FLA conv1d accepts [b, s, d] format input weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w bias=conv1d_bias, activation=self.activation, initial_state=None, output_final_state=False, - cu_seqlens=cu_seqlens_q, + cu_seqlens=_conv_cu_seqlens, ) + if _pad_n > 0: + qkv = qkv[:, :_orig_seq, :] nvtx_range_pop(suffix="conv1d") # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous) nvtx_range_push(suffix="prepare_qkv_for_gated_delta_rule") query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( - qkv, gate, beta, alpha, batch, seq_len + qkv, gate, beta, alpha, batch, seq_len, cp_size ) nvtx_range_pop(suffix="prepare_qkv_for_gated_delta_rule") # Calculate g and beta nvtx_range_push(suffix="g_and_beta") - A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp) - dt_bias_local_cp = get_parameter_local_cp( - self.dt_bias, dim=0, cp_group=self.pg_collection.cp - ) + A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group) + dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group) g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") @@ -496,19 +527,17 @@ def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor): # CP all to all: HP to CP if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - if self.cp_size > 1: + if cp_size > 1: norm_out_hp = norm_out_hp.index_select(0, thd_cp_a2a_inv) norm_out = tensor_a2a_hp2cp( norm_out_hp, seq_dim=0, head_dim=-1, - cp_group=self.pg_collection.cp, + cp_group=cp_group, redo_attention_load_balancing=False, ) else: - norm_out = tensor_a2a_hp2cp( - norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) + norm_out = tensor_a2a_hp2cp(norm_out_hp, seq_dim=0, head_dim=-1, cp_group=cp_group) return norm_out @@ -541,16 +570,14 @@ def _apply_gated_norm(self, x, gate): return y @jit_fuser - def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len): + def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len, cp_size): """ Prepare query, key, value, gate, beta, alpha tensors for gated delta rule. Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations. """ # Split qkv into query_key and value query_key, value = torch.split( - qkv, - [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], - dim=-1, + qkv, [2 * self.qk_dim_local_tp // cp_size, self.v_dim_local_tp // cp_size], dim=-1 ) # Reshape query_key and value @@ -562,8 +589,7 @@ def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_l query_key = l2norm(query_key.contiguous()) # Split query and key - split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size - query, key = torch.split(query_key, [split_size, split_size], dim=2) + query, key = query_key.chunk(2, dim=2) # Expand query and key if needed (grouped query attention) if self.num_value_heads // self.num_key_heads > 1: diff --git a/megatron/core/ssm/mamba_context_parallel.py b/megatron/core/ssm/mamba_context_parallel.py index 5c040716069..0968571c38d 100644 --- a/megatron/core/ssm/mamba_context_parallel.py +++ b/megatron/core/ssm/mamba_context_parallel.py @@ -92,6 +92,9 @@ def __init__( self.D_cp1 = D_cp1 self.D_has_hdim = D_has_hdim + self._set_cp_params() + + def _set_cp_params(self) -> None: self.cp_size = self.cp_group.size() if self.cp_size == 1: @@ -137,6 +140,11 @@ def __init__( # and also `nheads_local_tpcp = nheads_local_tp // cp_size` whilst ngroups_local_tpcp is # either 1 or `ngroups_local_tp // cp_size` + def set_context_parallel_group(self, cp_group: torch.distributed.ProcessGroup): + """Set the context parallel group.""" + self.cp_group = cp_group + self._set_cp_params() + def pre_conv_ssm( self, input_: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None ) -> torch.Tensor: diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index 820e5a5a452..9006f4bdda7 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -24,7 +24,7 @@ tensor_merge, ) from megatron.core.inference.utils import InferenceMode -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update from megatron.core.ssm.ops.mamba_ssm import selective_state_update @@ -470,6 +470,11 @@ def forward( out, out_bias = self._decode(hidden_states, conv_state, ssm_state) return out, out_bias + _orig_cp_group = self.cp.cp_group + _resolved_cp_group = resolve_cp_group(_orig_cp_group, packed_seq_params) + if _resolved_cp_group is not _orig_cp_group: + self.cp.set_context_parallel_group(_resolved_cp_group) + zxBCdt, _ = self.in_proj(hidden_states) zxBCdt = self.cp.pre_conv_ssm(zxBCdt, packed_seq_params) @@ -487,6 +492,8 @@ def forward( out, out_bias = self.out_proj(y) + if _resolved_cp_group is not _orig_cp_group: + self.cp.set_context_parallel_group(_orig_cp_group) return out, out_bias def _dynamic_inference(self, hidden_states: torch.Tensor, context: DynamicInferenceContext): diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index b27f90c53d0..384548f0ab3 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1045,6 +1045,13 @@ def forward( (Tuple[Tensor, Tensor]) Attention output and bias. """ + + # here we need to set the right cp group for dynamic-cp + _orig_cp_group = self.pg_collection.cp + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" + self.pg_collection.cp = packed_seq_params.cp_group + # Check if we need to skip RoPE # no_rope is 0-indexed array and self.layer_number is 1-indexed no_rope = ( @@ -1173,6 +1180,7 @@ def forward( out = output.transpose(0, 1).contiguous() context_layer = out.view(out.size(0), out.size(1), -1) output, bias = apply_module(self.linear_proj)(context_layer) + self.pg_collection.cp = _orig_cp_group return output, bias if ( @@ -1346,6 +1354,7 @@ def forward( ) nvtx_range_pop(suffix="linear_proj") + self.pg_collection.cp = _orig_cp_group return output, bias @jit_fuser diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index 48e6c76ea2f..1cfc129b7c7 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -363,11 +363,6 @@ def get_query_key_value_tensors( assert ( hidden_states.ndim == 3 ), f"hidden_states should be 3D, [s, b, h], got {hidden_states.ndim}D" - if packed_seq_params is not None: - assert ( - packed_seq_params.local_cp_size is None - ), "dynamic context parallel is not supported with MLA yet and is planned for future. \ - Please disable dynamic context parallel." inference_context = deprecate_inference_params(inference_context, inference_params) @@ -732,6 +727,14 @@ def forward( inference_context is None and inference_params is None ), "Inference is not supported for AbsorbedMLA" + # Set the right cp group for dynamic-cp. Mirrors Attention.forward: + # downstream RoPE uses self.pg_collection.cp, which must point at this + # microbatch's dynamic CP group. Restored before every return. + _orig_cp_group = self.pg_collection.cp + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" + self.pg_collection.cp = packed_seq_params.cp_group + # ===================== # Query, Key, and Value # ===================== @@ -813,6 +816,7 @@ def forward( # ================= output, bias = self.linear_proj(core_attn_out) + self.pg_collection.cp = _orig_cp_group return output, bias def backward_dw(self) -> NoReturn: diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 8023f53056e..2a41a35325b 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -334,6 +334,14 @@ def forward( if self.config.cache_mla_latents: self.prepare_for_absorption() + # Set the right cp group for dynamic-cp. Mirrors Attention.forward: + # downstream RoPE uses self.pg_collection.cp, which must point at this + # microbatch's dynamic CP group. Restored before every return. + _orig_cp_group = self.pg_collection.cp + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" + self.pg_collection.cp = packed_seq_params.cp_group + # ===================== # Query, Key, and Value # ===================== @@ -465,6 +473,7 @@ def forward( output, name="attn_proj", forced_released_tensors=[core_attn_out] ) + self.pg_collection.cp = _orig_cp_group return output, bias @@ -671,11 +680,6 @@ def get_query_key_value_tensors( assert ( hidden_states.ndim == 3 ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" - if packed_seq_params is not None: - assert ( - packed_seq_params.local_cp_size is None - ), "hybrid_context_parallel is not supported with MLA yet and is planned for future. \ - Please disable hybrid_context_parallel." inference_context = deprecate_inference_params(inference_context, inference_params) diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 200620b7483..16d96de57c8 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -17,7 +17,7 @@ from megatron.core.fp8_utils import get_fp8_context from megatron.core.inference.utils import InferenceMode from megatron.core.models.backends import BackendSpecProvider, LocalSpecProvider -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group from megatron.core.pipeline_parallel.utils import is_vp_last_stage from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel import ( @@ -246,7 +246,16 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No dims == -1 or dims == tensor.dim() - 1 ), "Packed sequence roll only supports the last dimension." assert shifts == -1, "Packed sequence roll only supports a single-token left shift." - cu_seqlens = packed_seq_params.cu_seqlens_q + # Prefer the padded cumulative seqlens because, with CP, the local THD layout is + # produced by `tex.thd_get_partitioned_indices(cu_seqlens_padded, ...)` and requires + # each per-sequence padded length to be divisible by 2*cp_size. Indexing with the + # unpadded cu_seqlens then produces wrong local boundaries when seqlens are not + # already multiples of 2*cp_size (e.g. odd seqlens). + cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if getattr(packed_seq_params, 'cu_seqlens_q_padded', None) is not None + else packed_seq_params.cu_seqlens_q + ) assert cu_seqlens is not None, "Packed sequence parameters must provide cu_seqlens_q." rolled_tensor = tensor.clone() @@ -351,21 +360,15 @@ def save_metrics_to_tracker( total: torch.Tensor, layer_number: int, num_layers: int, - reduce_group: torch.distributed.ProcessGroup = None, - avg_group: torch.distributed.ProcessGroup = None, + reduce_group: Optional[torch.distributed.ProcessGroup] = None, + avg_group: Optional[torch.distributed.ProcessGroup] = None, ): - """Save the mtp metrics (loss, correct, total) for logging. + """Save normalized MTP loss and acceptance counts for logging. - Args: - loss (torch.Tensor): The normalized loss value for this MTP layer. - correct (torch.Tensor): Number of correct predictions. - total (torch.Tensor): Total number of predictions. - layer_number (int): Layer index of the loss. - num_layers (int): The number of total layers. - reduce_group (torch.distributed.ProcessGroup): The group for reducing the loss. - avg_group (torch.distributed.ProcessGroup): The group for averaging the loss. + This compatibility path is used by tests and callers that already + computed a normalized per-layer loss. Dynamic-CP code should use + ``save_loss_to_tracker`` so loss is weighted by token counts. """ - # Skip mtp loss logging if layer_number is None. if layer_number is None: return @@ -383,6 +386,56 @@ def save_metrics_to_tracker( tracker["reduce_group"] = reduce_group tracker["avg_group"] = avg_group + @staticmethod + def save_loss_to_tracker( + loss_sum: torch.Tensor, + num_tokens: torch.Tensor, + layer_number: int, + num_layers: int, + correct: Optional[torch.Tensor] = None, + total: Optional[torch.Tensor] = None, + reduce_group: Optional[torch.distributed.ProcessGroup] = None, + avg_group: Optional[torch.distributed.ProcessGroup] = None, + ): + """Save the mtp loss sum and token count for logging. + + Stores raw sums so that the global per-token loss can be computed + correctly after all-reduce, even when token counts differ across + ranks (e.g. Dynamic CP) or microbatches. + + Args: + loss_sum (torch.Tensor): Sum of per-element losses on this rank. + num_tokens (torch.Tensor): Number of valid tokens on this rank. + layer_number (int): Layer index of the loss. + num_layers (int): The number of total layers. + correct (Optional[torch.Tensor]): Number of correct MTP predictions. + total (Optional[torch.Tensor]): Total number of MTP predictions. + reduce_group (torch.distributed.ProcessGroup): The group for sum-reducing losses. + avg_group (torch.distributed.ProcessGroup): The group for sum-reducing before averaging. + """ + if layer_number is None: + return + + tracker = MTPLossLoggingHelper.tracker + if "loss_sums" not in tracker: + tracker["loss_sums"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + tracker["num_tokens"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + tracker["loss_sums"][layer_number] += loss_sum.detach() + tracker["num_tokens"][layer_number] += num_tokens.detach() + if correct is not None and total is not None: + if "correct_values" not in tracker: + tracker["correct_values"] = torch.zeros( + num_layers, device=torch.cuda.current_device() + ) + if "total_values" not in tracker: + tracker["total_values"] = torch.zeros( + num_layers, device=torch.cuda.current_device() + ) + tracker["correct_values"][layer_number] += correct.detach() + tracker["total_values"][layer_number] += total.detach() + tracker["reduce_group"] = reduce_group + tracker["avg_group"] = avg_group + @staticmethod def clean_metrics_in_tracker(): """Clear the mtp metrics.""" @@ -400,16 +453,15 @@ def clean_metrics_in_tracker(): def reduce_metrics_in_tracker(): """Collect and reduce the mtp metrics across ranks.""" tracker = MTPLossLoggingHelper.tracker - if "loss_values" not in tracker: - return - loss_values = tracker["loss_values"] - if tracker.get('reduce_group') is not None: - torch.distributed.all_reduce(loss_values, group=tracker.get('reduce_group')) - if tracker.get('avg_group') is not None: - torch.distributed.all_reduce( - loss_values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG - ) + if "loss_values" in tracker: + loss_values = tracker["loss_values"] + if tracker.get('reduce_group') is not None: + torch.distributed.all_reduce(loss_values, group=tracker.get('reduce_group')) + if tracker.get('avg_group') is not None: + torch.distributed.all_reduce( + loss_values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG + ) for key in ["correct_values", "total_values"]: if key not in tracker: @@ -422,47 +474,83 @@ def reduce_metrics_in_tracker(): values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.SUM ) + @staticmethod + def clean_loss_in_tracker(): + """Clear the mtp losses.""" + tracker = MTPLossLoggingHelper.tracker + if "loss_sums" in tracker: + tracker["loss_sums"].zero_() + tracker["num_tokens"].zero_() + if "values" in tracker: + tracker["values"].zero_() + if "correct_values" in tracker: + tracker["correct_values"].zero_() + if "total_values" in tracker: + tracker["total_values"].zero_() + tracker["reduce_group"] = None + tracker["avg_group"] = None + + @staticmethod + def reduce_loss_in_tracker(): + """Collect and reduce the mtp losses across ranks. + + Packs loss sums and token counts into a single tensor for one + all-reduce, then computes per-token loss. This produces correct + weighted-average results even when ranks hold different numbers + of tokens (e.g. Dynamic CP with variable CP sizes). + """ + tracker = MTPLossLoggingHelper.tracker + if "loss_sums" not in tracker: + return + packed = torch.cat([tracker["loss_sums"], tracker["num_tokens"]]) + for group_key in ('reduce_group', 'avg_group'): + group = tracker.get(group_key) + if group is not None: + torch.distributed.all_reduce(packed, group=group) + loss_sums, num_tokens = packed.chunk(2) + tracker["values"] = loss_sums / num_tokens.clamp(min=1) + @staticmethod def track_mtp_metrics(loss_scale, iteration, writer, wandb_writer=None, total_loss_dict=None): """Track the Multi-Token Prediction (MTP) metrics for logging.""" + MTPLossLoggingHelper.reduce_loss_in_tracker() MTPLossLoggingHelper.reduce_metrics_in_tracker() tracker = MTPLossLoggingHelper.tracker - if "loss_values" not in tracker: + if "loss_sums" in tracker and "values" in tracker: + mtp_losses = tracker["values"] * loss_scale + elif "loss_values" in tracker: + mtp_losses = tracker["loss_values"] * loss_scale + else: return - mtp_losses = tracker["loss_values"] * loss_scale - mtp_corrects = tracker.get("correct_values", torch.zeros_like(mtp_losses)) - mtp_totals = tracker.get("total_values", torch.ones_like(mtp_losses)) - - # Process-local logging state; cumulative rates intentionally reset after restart/resume. - if ( - "cumulative_correct_values" not in tracker - or tracker["cumulative_correct_values"].shape != mtp_corrects.shape - ): - tracker["cumulative_correct_values"] = torch.zeros_like(mtp_corrects) - if ( - "cumulative_total_values" not in tracker - or tracker["cumulative_total_values"].shape != mtp_totals.shape - ): - tracker["cumulative_total_values"] = torch.zeros_like(mtp_totals) - - tracker["cumulative_correct_values"] += mtp_corrects - tracker["cumulative_total_values"] += mtp_totals - mtp_cumulative_corrects = tracker["cumulative_correct_values"] - mtp_cumulative_totals = tracker["cumulative_total_values"] + has_acceptance = "correct_values" in tracker and "total_values" in tracker + if has_acceptance: + mtp_corrects = tracker["correct_values"] + mtp_totals = tracker["total_values"] + + # Process-local logging state; cumulative rates intentionally + # reset after restart/resume. + if ( + "cumulative_correct_values" not in tracker + or tracker["cumulative_correct_values"].shape != mtp_corrects.shape + ): + tracker["cumulative_correct_values"] = torch.zeros_like(mtp_corrects) + if ( + "cumulative_total_values" not in tracker + or tracker["cumulative_total_values"].shape != mtp_totals.shape + ): + tracker["cumulative_total_values"] = torch.zeros_like(mtp_totals) + + tracker["cumulative_correct_values"] += mtp_corrects + tracker["cumulative_total_values"] += mtp_totals + mtp_cumulative_corrects = tracker["cumulative_correct_values"] + mtp_cumulative_totals = tracker["cumulative_total_values"] mtp_num_layers = mtp_losses.shape[0] for i in range(mtp_num_layers): loss_name = f"mtp_{i+1} loss" - step_acc_name = f"mtp_{i+1}_acceptance_rate" - cum_acc_name = f"mtp_{i+1}_cumulative_acceptance_rate" loss = mtp_losses[i] - # Empty masks can leave no valid MTP positions, so clamp denominators to avoid NaNs. - step_rate = (mtp_corrects[i] / torch.clamp(mtp_totals[i], min=1)) * 100.0 - cum_rate = ( - mtp_cumulative_corrects[i] / torch.clamp(mtp_cumulative_totals[i], min=1) - ) * 100.0 if total_loss_dict is not None: total_loss_dict[loss_name] = ( @@ -471,13 +559,26 @@ def track_mtp_metrics(loss_scale, iteration, writer, wandb_writer=None, total_lo if writer is not None: writer.add_scalar(loss_name, loss, iteration) - writer.add_scalar(step_acc_name, step_rate, iteration) - writer.add_scalar(cum_acc_name, cum_rate, iteration) if wandb_writer is not None: wandb_writer.log({f"{loss_name}": loss}, iteration) - wandb_writer.log({f"{step_acc_name}": step_rate}, iteration) - wandb_writer.log({f"{cum_acc_name}": cum_rate}, iteration) + if has_acceptance: + step_acc_name = f"mtp_{i+1}_acceptance_rate" + cum_acc_name = f"mtp_{i+1}_cumulative_acceptance_rate" + # Empty masks can leave no valid MTP positions, so clamp denominators to avoid NaNs. + step_rate = (mtp_corrects[i] / torch.clamp(mtp_totals[i], min=1)) * 100.0 + cum_rate = ( + mtp_cumulative_corrects[i] / torch.clamp(mtp_cumulative_totals[i], min=1) + ) * 100.0 + + if writer is not None: + writer.add_scalar(step_acc_name, step_rate, iteration) + writer.add_scalar(cum_acc_name, cum_rate, iteration) + if wandb_writer is not None: + wandb_writer.log({f"{step_acc_name}": step_rate}, iteration) + wandb_writer.log({f"{cum_acc_name}": cum_rate}, iteration) + + MTPLossLoggingHelper.clean_loss_in_tracker() MTPLossLoggingHelper.clean_metrics_in_tracker() @@ -631,6 +732,12 @@ def mtp_on_this_rank( - If no custom layout is provided, assumes all MTP layers (if any) are placed on the last pipeline stage. The function returns True only on the last pipeline stage. """ + if layout is not None and hasattr(layout, "pipeline_model_parallel_layout"): + # Backward-compat: some callers pass a TransformerConfig as the first + # positional argument instead of (layout, mtp_num_layers). Unpack it. + _config = layout + layout = _config.pipeline_model_parallel_layout + mtp_num_layers = _config.mtp_num_layers mtp_on_this_rank = False pp_rank = parallel_state.get_pipeline_model_parallel_rank() if layout is not None: @@ -857,19 +964,17 @@ def process_mtp_loss( mtp_loss = loss_mask * mtp_loss if is_training: - mtp_loss_for_log = ( - torch.sum(mtp_loss) * (num_tokens > 0).to(mtp_loss.dtype) - ) / num_tokens.clamp(min=1) correct, total = _compute_mtp_acceptance_counts( mtp_logits, mtp_labels, loss_mask, output_layer, runtime_gather_output, tp_group ) - MTPLossLoggingHelper.save_metrics_to_tracker( - mtp_loss_for_log, - correct, - total, + MTPLossLoggingHelper.save_loss_to_tracker( + torch.sum(mtp_loss), + num_tokens, mtp_layer_number, config.mtp_num_layers, + correct=correct, + total=total, avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True), ) mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers @@ -1077,27 +1182,20 @@ def _get_embeddings( sequence length, b is the batch size, and h is the hidden size. packed_seq_params (PackedSeqParams): Parameters for packed sequence processing. """ - # Calc logits for the current Multi-Token Prediction (MTP) layers. + cp_group = resolve_cp_group(self.cp_group, packed_seq_params) + input_ids, _ = roll_tensor( - input_ids, - shifts=-1, - dims=-1, - cp_group=self.cp_group, - packed_seq_params=packed_seq_params, + input_ids, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params ) position_ids, _ = roll_tensor( - position_ids, - shifts=-1, - dims=-1, - cp_group=self.cp_group, - packed_seq_params=packed_seq_params, + position_ids, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params ) if padding_mask is not None: padding_mask, _ = roll_tensor( padding_mask, shifts=-1, dims=-1, - cp_group=self.cp_group, + cp_group=cp_group, packed_seq_params=packed_seq_params, ) # embedding @@ -1479,6 +1577,8 @@ def forward( [s, b, h], and optionally the updated context tensor if cross-attention is used. """ assert context is None, "multi token prediction + cross attention is not yet supported." + _orig_cp_group = self.cp_group + self.cp_group = resolve_cp_group(self.cp_group, packed_seq_params) input_ids, position_ids, padding_mask, decoder_input, hidden_states = self._get_embeddings( input_ids=input_ids, position_ids=position_ids, @@ -1521,6 +1621,7 @@ def forward( sequence_len_offset=sequence_len_offset, ) + self.cp_group = _orig_cp_group return hidden_states, input_ids, position_ids, padding_mask def sharded_state_dict( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index ab84abd3a17..7020e84b8e3 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -2388,7 +2388,7 @@ def _scope_to_str(s): "full-iteration CUDA graphs" ) - if self.moe_token_dispatcher_type in ["allgather"]: + if self.num_moe_experts is not None and self.moe_token_dispatcher_type in ["allgather"]: if self.variable_seq_lengths is True: raise ValueError( f"Token dispatcher type: {self.moe_token_dispatcher_type} does not support " @@ -2582,6 +2582,40 @@ def _scope_to_str(s): self.attention_backend == AttnBackend.flash ), "Batch invariant mode only supports FlashAttention" + if self.sequence_packing_scheduler is not None: + # Check TE version. + if not HAVE_PACKAGING: + raise ImportError( + "packaging is not installed. Please install it with `pip install packaging`." + ) + # TODO: remove this after we fix the convergence issue with TE < 2.9. + if not ( + is_te_min_version("2.9.0") or get_te_version() == PkgVersion("2.9.0.dev0+5b3092a") + ): + raise ValueError( + "SFT sequence packing requires Transformer Engine >= 2.9.0 " + f"but got {get_te_version()} (TE < 2.9.0 may have convergence issues)." + ) + + # Needed for passing variable sequences between pp stages. + self.variable_seq_lengths = True + + if self.num_moe_experts is not None: + assert self.moe_token_dispatcher_type in ("alltoall", "flex"), ( + f"sequence_packing only supports moe_token_dispatcher_type in " + f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'" + ) + + supported_schedulers = ['dp_balanced', 'default_dynamic_cp'] + if ( + self.sequence_packing_scheduler is not None + and self.sequence_packing_scheduler not in supported_schedulers + ): + raise ValueError( + f"Unsupported scheduler: {self.sequence_packing_scheduler}. " + f"Available schedulers: {supported_schedulers}" + ) + @dataclass class MLATransformerConfig(TransformerConfig): diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 2e916482433..eec06854e6a 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -43,6 +43,7 @@ from megatron.core import parallel_state from megatron.core.dist_checkpointing.mapping import ShardedTensor +from megatron.core.packed_seq_params import PackedSeqParams try: from packaging.version import Version as PkgVersion @@ -2368,7 +2369,7 @@ def get_pretrain_batch_on_this_cp_rank( def get_batch_on_this_cp_rank( batch: Dict[str, Any], - is_hybrid_cp: bool, + is_hybrid_cp: bool = False, cp_group: Optional[torch.distributed.ProcessGroup] = None, hybrid_cp_group_func: Optional[Callable[[int], torch.distributed.ProcessGroup]] = None, ): @@ -2401,6 +2402,12 @@ def get_batch_on_this_cp_rank( to this CP rank. """ + if cp_group is None: + # Backward-compatible fallback for callers that pass only ``batch`` + # (pretrain entrypoints historically read the global CP group + # internally): use the current context-parallel group. + cp_group = parallel_state.get_context_parallel_group() + if batch.get("cu_seqlens") is not None: # NOTE(asolergi-nv): SFT & HybridCP case if is_hybrid_cp: assert ( @@ -2417,6 +2424,46 @@ def get_batch_on_this_cp_rank( return batch +def get_thd_batch_on_this_cp_rank( + batch: Dict[str, Any], + cu_seqlens: torch.Tensor, + cu_seqlens_padded: torch.Tensor, + max_seqlen: torch.Tensor, + cp_size: Optional[int] = None, + cp_rank: Optional[int] = None, +): + """Slice each sub-sample in a packed sample batch input along + sequence dimension into multiple chunks, which are parallelized + across GPUs in a context parallel group. + """ + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=int(max_seqlen[0].item()), + max_seqlen_kv=int(max_seqlen[0].item()), + ) + + cp_size = parallel_state.get_context_parallel_world_size() if cp_size is None else cp_size + cp_rank = parallel_state.get_context_parallel_rank() if cp_rank is None else cp_rank + if cp_size > 1: # slice batch along sequence dimension for context parallelism + assert tex is not None and is_te_min_version("1.10.0"), ( + "Please update Transformer Engine to >= 1.10 to use " + "Context Parallel with THD format data" + ) + index = tex.thd_get_partitioned_indices( + cu_seqlens_padded, batch['tokens'].size(1), cp_size, cp_rank + ) + for key, data in batch.items(): + if key in {'attention_mask', 'cu_seqlens', 'cu_seqlens_padded', 'max_seqlen'}: + continue + batch[key] = data.index_select(1, index) + + return batch, packed_seq_params + + ###################### ### NVTX profiling ### ###################### diff --git a/megatron/elastification/pretrain_hybrid_flex.py b/megatron/elastification/pretrain_hybrid_flex.py index 846284167d7..10b3d414438 100644 --- a/megatron/elastification/pretrain_hybrid_flex.py +++ b/megatron/elastification/pretrain_hybrid_flex.py @@ -21,7 +21,7 @@ get_context_parallel_group, get_data_parallel_rank, get_data_parallel_world_size, - get_hybrid_data_context_parallel_groups, + get_dynamic_data_context_parallel_groups, get_pipeline_model_parallel_rank, get_pipeline_model_parallel_world_size, get_tensor_model_parallel_group, @@ -177,7 +177,7 @@ def get_batch(data_iterator, vp_stage=None): cp_size = args.context_parallel_size tp_rank = mpu.get_tensor_model_parallel_rank() is_sft = args.sft - is_hybrid_cp = args.hybrid_context_parallel + is_hybrid_cp = args.dynamic_context_parallel mtp_on_this_rank = mtp_on_this_rank_func( layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, @@ -225,7 +225,7 @@ def get_batch(data_iterator, vp_stage=None): batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), - hybrid_cp_group_func=get_hybrid_data_context_parallel_groups, + hybrid_cp_group_func=get_dynamic_data_context_parallel_groups, ) # cu_seqlens / max_seqlen arrive with the dataloader's batch dim (shape (1, n) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e7415dc3019..30c5df5c8e8 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1386,12 +1386,59 @@ def validate_args(args, defaults={}): if args.tp_comm_overlap: assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' - if args.hybrid_context_parallel: - assert not args.pipeline_model_parallel_size > 1, 'Hybrid context parallelism not supported with pipeline parallelism' - assert not args.enable_cuda_graph, 'Hybrid context parallelism not supported with CUDA Graph' - assert not args.use_megatron_fsdp, 'Hybrid context parallelism not supported with Megatron FSDP' - assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type' - assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss' + if args.dynamic_context_parallel: + assert ( + not args.enable_cuda_graph + ), 'Dynamic context parallelism not supported with CUDA Graph' + assert ( + not args.use_megatron_fsdp + ), 'Dynamic context parallelism not supported with Megatron FSDP' + assert ( + args.dataloader_type == 'single' + ), 'Dynamic context parallelism only supported with single dataloader type' + assert ( + args.calculate_per_token_loss + ), 'Dynamic context parallelism must be used with --calculate-per-token-loss' + if args.sequence_packing_scheduler is None: + args.sequence_packing_scheduler = 'default_dynamic_cp' + if args.sequence_packing_scheduler != 'default_dynamic_cp': + raise ValueError( + 'Dynamic context parallelism requires ' + 'sequence_packing_scheduler=default_dynamic_cp' + ) + + dp_cp_size = args.data_parallel_size * args.context_parallel_size + assert args.min_dynamic_context_parallel_size <= dp_cp_size, ( + f'min_dynamic_context_parallel_size ({args.min_dynamic_context_parallel_size}) ' + f'must be <= dp_size * cp_size ({dp_cp_size})' + ) + + import warnings + + warnings.warn( + f"Dynamic CP enabled: dp_size * context_parallel_size=" + f"{args.data_parallel_size * args.context_parallel_size} " + f"will be used as the maximum dynamic CP group size. " + f"Dynamic CP groups will range from " + f"min_dynamic_context_parallel_size={args.min_dynamic_context_parallel_size} " + f"to {args.data_parallel_size * args.context_parallel_size}." + ) + + if args.sequence_packing_scheduler is not None: + if args.sequence_packing_scheduler == 'dp_balanced': + total_cp_ranks = args.context_parallel_size + else: + total_cp_ranks = args.data_parallel_size * args.context_parallel_size + if args.max_seqlen_per_dp_cp_rank is None: + args.max_seqlen_per_dp_cp_rank = getattr(args, "max_seqlen_per_cp_rank", None) + if args.max_seqlen_per_dp_cp_rank is None: + args.max_seqlen_per_dp_cp_rank = ( + args.seq_length + total_cp_ranks - 1 + ) // total_cp_ranks + assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, ( + f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) ' + f'must be >= single sequence max length ({args.seq_length})' + ) # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index 296acc97941..1143704e470 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -38,7 +38,6 @@ def build_pretraining_data_loader(dataset, consumed_samples): # Use eval-specific batch sizes for validation/test splits is_eval = split in (Split.valid, Split.test) micro_batch_size = getattr(args, 'eval_micro_batch_size', args.micro_batch_size) if is_eval else args.micro_batch_size - global_batch_size = getattr(args, 'eval_global_batch_size', args.global_batch_size) if is_eval else args.global_batch_size if split == Split.valid and args.full_validation: batch_sampler = MegatronFullValidationSampler( @@ -46,22 +45,15 @@ def build_pretraining_data_loader(dataset, consumed_samples): data_parallel_rank=mpu.get_data_parallel_rank(), data_parallel_size=mpu.get_data_parallel_world_size()) elif args.dataloader_type == 'single': - if args.hybrid_context_parallel: - batch_sampler = HybridCPMegatronPretrainingSampler( - total_samples=len(dataset), - consumed_samples=consumed_samples, - micro_batch_size=micro_batch_size, - global_batch_size=global_batch_size, - data_parallel_rank=mpu.get_data_parallel_rank(), - data_parallel_size=mpu.get_data_parallel_world_size()) - else: - # Megatron sampler - batch_sampler = MegatronPretrainingSampler( - total_samples=len(dataset), - consumed_samples=consumed_samples, - micro_batch_size=micro_batch_size, - data_parallel_rank=mpu.get_data_parallel_rank(), - data_parallel_size=mpu.get_data_parallel_world_size()) + # Packing schedulers consume one microbatch at a time and form + # global/DCP batches themselves. + batch_sampler = MegatronPretrainingSampler( + total_samples=len(dataset), + consumed_samples=consumed_samples, + micro_batch_size=micro_batch_size, + data_parallel_rank=mpu.get_data_parallel_rank(), + data_parallel_size=mpu.get_data_parallel_world_size(), + ) elif args.dataloader_type == 'cyclic': batch_sampler = MegatronPretrainingRandomSampler( dataset, @@ -97,9 +89,10 @@ def close_nvidia_fds(): maybe_worker_init_fn = ( worker_init_fn if args.num_workers > 0 else None ) - # Torch dataloader. - if args.hybrid_context_parallel: - extra_kwargs = {"collate_fn": lambda x: x,} + # Identity collate for packing-scheduler paths; they emit one + # variable-length dict per sample, not stack-able by the default collate. + if args.dynamic_context_parallel or args.sequence_packing_scheduler is not None: + extra_kwargs = {"collate_fn": lambda x: x} else: extra_kwargs = {} return torch.utils.data.DataLoader( @@ -181,50 +174,6 @@ def __iter__(self): start_idx, end_idx = self.get_start_end_idx() yield batch[start_idx:end_idx] -class HybridCPMegatronPretrainingSampler(MegatronPretrainingSampler): - """ - Data sampler for hybrid context parallel (Hybrid CP) format. - This data sampler pulls in the entire global batch at once across all data parallel ranks. - This helps provide the Hybrid CP Dataloader Wrapper to schedule and load balance sub-samples - of the entire global batch. - """ - - def __init__(self, total_samples, consumed_samples, micro_batch_size, global_batch_size, - data_parallel_rank, data_parallel_size, drop_last=True): - super().__init__(total_samples, consumed_samples, micro_batch_size, data_parallel_rank, data_parallel_size, drop_last) - self.global_batch_size = global_batch_size - self.data_parallel_size = data_parallel_size - self.num_micro_batches = self.global_batch_size // self.micro_batch_times_data_parallel_size - - def __len__(self): - return self.total_samples - - def get_start_end_idx_global_batch(self): - start_idx = [self.data_parallel_rank * self.micro_batch_size + i * self.micro_batch_size * self.data_parallel_size for i in range(self.num_micro_batches)] - end_idx = [start_idx[i] + self.micro_batch_size for i in range(self.num_micro_batches)] - return start_idx, end_idx - - def __iter__(self): - batch = [] - # Last batch will be dropped if drop_last is not set False - for idx in range(self.consumed_samples, self.total_samples): - batch.append(idx) - if len(batch) == self.micro_batch_times_data_parallel_size * self.num_micro_batches: - start_idx, end_idx = self.get_start_end_idx_global_batch() - global_batch_idx = [] - for i in range(self.num_micro_batches): - global_batch_idx.extend(batch[start_idx[i]:end_idx[i]]) - yield global_batch_idx - batch = [] - - # Check the last partial batch and see drop_last is set - if len(batch) > 0 and not self.drop_last: - start_idx, end_idx = self.get_start_end_idx_global_batch() - global_batch_idx = [] - for i in range(self.num_micro_batches): - global_batch_idx.extend(batch[start_idx[i]:end_idx[i]]) - yield global_batch_idx - class MegatronFullValidationSampler: """Sampler for full validation that handles small datasets gracefully. diff --git a/megatron/training/datasets/sft_dataset.py b/megatron/training/datasets/sft_dataset.py index 9de5d2a52fe..bbb43484e02 100644 --- a/megatron/training/datasets/sft_dataset.py +++ b/megatron/training/datasets/sft_dataset.py @@ -88,6 +88,30 @@ def _split_conversations(self, merged_conversations): split_conversations.append(current) return split_conversations + def _calculate_padding_divisor(self) -> int: + """ + Calculate the divisor used for sequence padding. + tp_pad = tp_size * 2 if tp_size > 1 else 1 + cp_pad = cp_size * 2 if cp_size > 1 else 1 + cp_pad = cp_pad * dp_size if dynamic_cp else cp_pad + divisor = cp_pad * tp_pad + """ + if self.config.dynamic_context_parallel: + # Dynamic CP: consider both CP and DP + cp_pad = self.config.data_parallel_size * self.config.context_parallel_size * 2 + else: + # Standard CP: only consider CP + cp_pad = ( + self.config.context_parallel_size * 2 + if self.config.context_parallel_size > 1 + else 1 + ) + tp_pad = self.config.sequence_parallel_size if self.config.sequence_parallel_size > 0 else 1 + divisor = cp_pad * tp_pad + # TODO(tailaim): do we need to pad for FP8 execution? + # divisor = ((divisor + 15) // 16) * 16 + return divisor + def __getitem__(self, idx: int) -> Dict[str, Any]: tokenizer = self.config.tokenizer @@ -124,12 +148,11 @@ def extend_with_padding(tokens, targets, positions, pad_len): assert not self.config.reset_position_ids pack_positions.extend(range(len(tokens_list))) - if self.config.context_parallel_size > 1: - pad_granularity = self.config.context_parallel_size * 2 - mod_token_count = len(pack_tokens) % pad_granularity - if mod_token_count != 0: - pad_len = pad_granularity - mod_token_count - extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) + pad_granularity = self._calculate_padding_divisor() + mod_token_count = len(pack_tokens) % pad_granularity + if mod_token_count != 0: + pad_len = pad_granularity - mod_token_count + extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) # TODO(duncan): Consider also padding to multiple of number of tokens here. This might # be needed for efficiency (and potentially set via command-line argument). diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index ff655502019..debc50ab098 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -346,7 +346,8 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s use_sharp=args.use_sharp, context_parallel_size=args.context_parallel_size, hierarchical_context_parallel_sizes=args.hierarchical_context_parallel_sizes, - hybrid_context_parallel=args.hybrid_context_parallel, + dynamic_context_parallel=args.dynamic_context_parallel, + min_dynamic_context_parallel_size=args.min_dynamic_context_parallel_size, expert_model_parallel_size=args.expert_model_parallel_size, num_distributed_optimizer_instances=args.num_distributed_optimizer_instances, expert_tensor_parallel_size=args.expert_tensor_parallel_size, diff --git a/megatron/training/training.py b/megatron/training/training.py index 3930cc46a21..1671eaddf2d 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -190,7 +190,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): except ImportError: HAVE_FSDP2 = False -from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper +from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper, wrap_data_iterator from megatron.core.distributed import finalize_model_grads from megatron.core.enums import ModelType from megatron.core.inference.symmetric_memory import SymmetricMemoryManager @@ -202,7 +202,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): destroy_global_memory_buffer, destroy_model_parallel, get_context_parallel_group, - get_hybrid_data_context_parallel_groups, + get_dynamic_data_context_parallel_groups, update_pg_timeout, ) from megatron.core.rerun_state_machine import ( @@ -2117,7 +2117,7 @@ def dummy_train_step(data_iterator): args = get_args() tp_rank = mpu.get_tensor_model_parallel_rank() is_sft = getattr(args, 'sft', False) - is_hybrid_cp = args.hybrid_context_parallel + is_hybrid_cp = args.dynamic_context_parallel BATCH_KEYS = [ "tokens", "labels", "loss_mask", "position_ids", "attention_mask", @@ -2155,7 +2155,7 @@ def dummy_train_step(data_iterator): batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), - hybrid_cp_group_func=get_hybrid_data_context_parallel_groups, + hybrid_cp_group_func=get_dynamic_data_context_parallel_groups, ) @@ -2217,6 +2217,27 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch if isinstance(optim_instance, DistributedOptimizer): optim_instance._copy_main_params_to_param_buffer() + if config.sequence_packing_scheduler is not None: + # This wrapper is designed to support DP-balanced THD and dynamic-CP. + # Before wrapping, the data_iterator returns either a single sequence per get_item call, or a list where each element is a sequence. + # The wrapper is responsible for: + # 1. scheduling the sequences across ranks + # 2. packing them into THD format + # 3. broadcast flops parametes and num_microbatches to TP ranks to support unfixed num_microbatches + # 4. broadcast metadata(cu_seqlens, cu_seqlens_padded, max_seqlen, etc.) to PP ranks to + # 5. returning the packed data iterator and the FLOPs parameters + ( + data_iterator, + num_microbatches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) = wrap_data_iterator(data_iterator, config, get_num_microbatches()) + else: + # data_iterator unchanged + num_microbatches = get_num_microbatches() + seqlen_sum_this_global_batch = args.seq_length * args.global_batch_size + seqlen_squared_sum_this_global_batch = args.seq_length**2 * args.global_batch_size + # Forward pass. if save_activations_in_this_iteration: enable_activation_logging(model, args.save) @@ -2228,7 +2249,7 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch forward_step_func=forward_step_func, data_iterator=data_iterator, model=model, - num_microbatches=get_num_microbatches(), + num_microbatches=num_microbatches, seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, decoder_seq_length=args.decoder_seq_length, @@ -2577,7 +2598,10 @@ def training_log( # Log MTP metrics. if args.mtp_num_layers is not None: - mtp_loss_scale = 1 / get_num_microbatches() + # MTP tracker stores raw loss sums and token counts, so after reduction + # tracker["values"] already equals the per-token loss (loss_sum / num_tokens) + # aggregated across all ranks and microbatches. No further scaling needed. + mtp_loss_scale = 1.0 MTPLossLoggingHelper.track_mtp_metrics( mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict ) @@ -3172,9 +3196,6 @@ def train( energy_monitor = get_energy_monitor() one_logger = get_one_logger() - if args.hybrid_context_parallel: - train_data_iterator = iter(HybridCPDataLoaderWrapper(train_data_iterator, config)) - if args.run_workload_inspector_server: try: import threading @@ -3898,11 +3919,30 @@ def evaluate( # Don't care about timing during evaluation config.timers = None ft_integration.on_eval_step_start() + if config.sequence_packing_scheduler is not None: + # This wrapper is designed to support DP-balanced THD and dynamic-CP. + # Before wrapping, the data_iterator returns either a single sequence per get_item call, or a list where each element is a sequence. + # The wrapper is responsible for: + # 1. scheduling the sequences across ranks + # 2. packing them into THD format + # 3. broadcast flops parametes and num_microbatches to TP ranks to support unfixed num_microbatches + # 4. broadcast metadata(cu_seqlens, cu_seqlens_padded, max_seqlen, etc.) to PP ranks to + # 5. returning the packed data iterator and the FLOPs parameters + try: + (packed_data_iterator, scheduled_eval_num_microbatches, _, _) = ( + wrap_data_iterator(data_iterator, config, eval_num_microbatches) + ) + except StopIteration: + # Validation data iterator exhausted, stop evaluation early. + break + else: + packed_data_iterator = data_iterator + scheduled_eval_num_microbatches = eval_num_microbatches loss_dicts = forward_backward_func( forward_step_func=forward_step_func, - data_iterator=data_iterator, + data_iterator=packed_data_iterator, model=model, - num_microbatches=eval_num_microbatches, + num_microbatches=scheduled_eval_num_microbatches, seq_length=args.seq_length, micro_batch_size=eval_micro_batch_size, decoder_seq_length=args.decoder_seq_length, diff --git a/pretrain_gpt.py b/pretrain_gpt.py index e475129847b..1da3ae84334 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -26,13 +26,14 @@ from gpt_builders import gpt_builder from megatron.core import mpu from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder +from megatron.core.datasets.data_schedule import get_batch_on_this_rank_for_sequence_packing from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset from megatron.core.enums import ModelType from megatron.core.models.gpt import GPTModel from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( get_context_parallel_group, - get_hybrid_data_context_parallel_groups, + get_dynamic_data_context_parallel_groups, ) from megatron.core.rerun_state_machine import get_rerun_state_machine from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer @@ -54,7 +55,7 @@ print_rank_0, set_startup_timestamps, ) -from megatron.training.argument_utils import pretrain_cfg_container_from_args, gpt_config_from_args +from megatron.training.argument_utils import gpt_config_from_args, pretrain_cfg_container_from_args from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig from megatron.training.datasets.sft_dataset import SFTDataset @@ -86,7 +87,21 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): is_sft = args.sft create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader mtp_on_this_rank = mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) - is_hybrid_cp = args.hybrid_context_parallel + is_hybrid_cp = args.dynamic_context_parallel + + if args.sequence_packing_scheduler is not None: + # Sequence-packing (THD) scheduler path: the wrapped data iterator + # emits packed micro-batches; this helper broadcasts them across TP, + # slices for (dynamic) CP, and returns + # (tokens, labels, loss_mask, attention_mask, position_ids, + # packed_seq_params) with a ready-made PackedSeqParams. + return get_batch_on_this_rank_for_sequence_packing( + data_iterator, + vpp_size=config.virtual_pipeline_model_parallel_size, + mtp_on_this_rank=mtp_on_this_rank, + vp_stage=vp_stage, + dynamic_cp=args.dynamic_context_parallel, + ) if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank and not is_sft: return [None for _ in BATCH_KEYS] @@ -103,7 +118,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): assert is_sft return None, batch['cu_seqlens'], batch['cu_seqlens_padded'], None, None, None, None, batch['max_seqlen'], None, None - batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_hybrid_data_context_parallel_groups) + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_dynamic_data_context_parallel_groups) # Return values in BATCH_KEYS order so callers can unpack into the fixed # names regardless of any provenance fields wrappers like BlendedDataset @@ -196,6 +211,14 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa global stimer with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") + batch_values = get_batch(data_iterator, vp_stage) + + if args.sequence_packing_scheduler is not None: + # Sequence-packing scheduler path: get_batch already returns a + # ready-made PackedSeqParams via + # get_batch_on_this_rank_for_sequence_packing. + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = batch_values + else: ( attention_mask, cu_seqlens, @@ -207,30 +230,30 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa max_seqlen, position_ids, tokens, - ) = get_batch(data_iterator, vp_stage) - - packed_seq_params = None - if cu_seqlens is not None: - # cu_seqlens / cu_seqlens_padded carry the dataloader's batch dim (1, n). - # PackedSeqParams (and TE attention) expect 1-D, so squeeze before use. - cu_seqlens = cu_seqlens[0] - if cu_seqlens_padded is not None: - cu_seqlens_padded = cu_seqlens_padded[0] - # Use real (unpadded) cu_seqlens to feed the FLOPs accounting: varlen - # attention only computes work for real tokens within each chunk. - update_seqlen_stats_from_cu_seqlens(cu_seqlens) - cu_seqlens_for_params = cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens # TODO(asolergi-nv): Currently there is a bug forcing cu_seqlens to be cu_seqlens_padded - packed_seq_params = PackedSeqParams( - qkv_format="thd", - cu_seqlens_q=cu_seqlens_for_params, - cu_seqlens_kv=cu_seqlens_for_params, - cu_seqlens_q_padded=cu_seqlens_padded, - cu_seqlens_kv_padded=cu_seqlens_padded, - max_seqlen_q=int(max_seqlen.item()), - max_seqlen_kv=int(max_seqlen.item()), - local_cp_size=int(local_cp_size.item()) if local_cp_size is not None else None, - cp_group=hybrid_cp_group, - ) + ) = batch_values + + packed_seq_params = None + if cu_seqlens is not None: + # cu_seqlens / cu_seqlens_padded carry the dataloader's batch dim (1, n). + # PackedSeqParams (and TE attention) expect 1-D, so squeeze before use. + cu_seqlens = cu_seqlens[0] + if cu_seqlens_padded is not None: + cu_seqlens_padded = cu_seqlens_padded[0] + # Use real (unpadded) cu_seqlens to feed the FLOPs accounting: varlen + # attention only computes work for real tokens within each chunk. + update_seqlen_stats_from_cu_seqlens(cu_seqlens) + cu_seqlens_for_params = cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens # TODO(asolergi-nv): Currently there is a bug forcing cu_seqlens to be cu_seqlens_padded + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_for_params, + cu_seqlens_kv=cu_seqlens_for_params, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=int(max_seqlen.item()), + max_seqlen_kv=int(max_seqlen.item()), + local_cp_size=int(local_cp_size.item()) if local_cp_size is not None else None, + cp_group=hybrid_cp_group, + ) timers('batch-generator').stop() @@ -302,7 +325,7 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: "context_parallel_size": args.context_parallel_size, "data_parallel_size": args.data_parallel_size, "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, - "hybrid_context_parallel": args.hybrid_context_parallel, + "dynamic_context_parallel": args.dynamic_context_parallel, } # add FIM args to the config diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 110ffabc11e..1586fd9161c 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -31,7 +31,7 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( get_context_parallel_group, - get_hybrid_data_context_parallel_groups, + get_dynamic_data_context_parallel_groups, ) from megatron.core.rerun_state_machine import get_rerun_state_machine from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer @@ -85,7 +85,7 @@ def get_batch(data_iterator, vp_stage=None): is_sft = args.sft create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader mtp_on_this_rank = mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) - is_hybrid_cp = args.hybrid_context_parallel + is_dynamic_cp = args.dynamic_context_parallel if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank and not is_sft: return [None for _ in BATCH_KEYS] @@ -96,13 +96,13 @@ def get_batch(data_iterator, vp_stage=None): for key in BATCH_KEYS: batch[key] = batch[key].cuda(non_blocking=True) if key in batch and batch[key] is not None else None - batch = get_batch_on_this_tp_rank(batch, broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), broadcast_group=mpu.get_tensor_model_parallel_group(), is_sft=is_sft, is_hybrid_cp=is_hybrid_cp, create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, cp_size=cp_size, tp_rank=tp_rank, micro_batch_size=args.micro_batch_size, seq_length=args.seq_length, mtp_on_this_rank=mtp_on_this_rank, pipeline_model_parallel_size=args.pipeline_model_parallel_size, is_pipeline_first_stage=mpu.is_pipeline_first_stage(), is_pipeline_last_stage=mpu.is_pipeline_last_stage()) + batch = get_batch_on_this_tp_rank(batch, broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), broadcast_group=mpu.get_tensor_model_parallel_group(), is_sft=is_sft, is_hybrid_cp=is_dynamic_cp, create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, cp_size=cp_size, tp_rank=tp_rank, micro_batch_size=args.micro_batch_size, seq_length=args.seq_length, mtp_on_this_rank=mtp_on_this_rank, pipeline_model_parallel_size=args.pipeline_model_parallel_size, is_pipeline_first_stage=mpu.is_pipeline_first_stage(), is_pipeline_last_stage=mpu.is_pipeline_last_stage()) if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank: assert is_sft return None, batch['cu_seqlens'], batch['cu_seqlens_padded'], None, None, None, None, batch['max_seqlen'], None, None - batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_hybrid_data_context_parallel_groups) + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_dynamic_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_dynamic_data_context_parallel_groups) # Return values in BATCH_KEYS order so callers can unpack into the fixed # names regardless of any provenance fields wrappers like BlendedDataset @@ -294,7 +294,7 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: context_parallel_size=args.context_parallel_size, data_parallel_size=args.data_parallel_size, sequence_parallel_size=args.tensor_model_parallel_size * args.sequence_parallel, - hybrid_context_parallel=args.hybrid_context_parallel, + dynamic_context_parallel=args.dynamic_context_parallel, ) diff --git a/tests/unit_tests/data/test_get_batch.py b/tests/unit_tests/data/test_get_batch.py index c136dcfa807..daed71faa81 100644 --- a/tests/unit_tests/data/test_get_batch.py +++ b/tests/unit_tests/data/test_get_batch.py @@ -36,7 +36,7 @@ def initialize_test_environment( args.sequence_parallel = True if tp_size > 1 else False args.pipeline_model_parallel_size = pp_size args.context_parallel_size = cp_size - args.hybrid_context_parallel = hybrid_context_parallel + args.dynamic_context_parallel = hybrid_context_parallel args.max_seqlen_per_cp_rank = max_seqlen_per_cp_rank args.sft = sft args.micro_batch_size = micro_batch_size @@ -60,7 +60,7 @@ def initialize_test_environment( tensor_model_parallel_size=tp_size, pipeline_model_parallel_size=pp_size, context_parallel_size=cp_size, - hybrid_context_parallel=hybrid_context_parallel, + dynamic_context_parallel=hybrid_context_parallel, ) return args diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index c781dd11dd8..ce5fbed4233 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -300,8 +300,11 @@ "min_offloaded_tensor_size": 1024 * 1024, "offload_modules": [], "fine_grained_offloading_max_inflight_offloads": None, + "dynamic_context_parallel": False, + "min_dynamic_context_parallel_size": 1, "hybrid_context_parallel": False, "max_seqlen_per_dp_cp_rank": None, + "sequence_packing_scheduler": None, "inference_cuda_graph_scope": { "__objclass__": "megatron.core.transformer.enums.InferenceCudaGraphScope", "_name_": "none", diff --git a/tests/unit_tests/test_parallel_state.py b/tests/unit_tests/test_parallel_state.py index 65b0d5ca91a..64c3a5b36a6 100644 --- a/tests/unit_tests/test_parallel_state.py +++ b/tests/unit_tests/test_parallel_state.py @@ -508,9 +508,9 @@ def golden_rank_result_from_past_code( "world_size, tp_size, cp_size, dp_size", [(8, 1, 2, 4), (8, 1, 1, 8)], # 8 GPUs, 1 TP, 2 CP, 4 DP # 8 GPUs, 1 TP, 1 CP, 8 DP ) -def test_hybrid_dp_cp_groups(world_size, tp_size, cp_size, dp_size): +def test_dynamic_dp_cp_groups(world_size, tp_size, cp_size, dp_size): """ - Test that hybrid DPxCP groups are created correctly. + Test that dynamic DPxCP groups are created correctly. """ Utils.destroy_model_parallel() @@ -521,13 +521,13 @@ def test_hybrid_dp_cp_groups(world_size, tp_size, cp_size, dp_size): Utils.initialize_model_parallel( tensor_model_parallel_size=tp_size, context_parallel_size=cp_size, - hybrid_context_parallel=True, + dynamic_context_parallel=True, ) dp_cp_size = ps.get_data_parallel_world_size(with_context_parallel=True) - group_sizes = [2**i for i in range(int(log2(dp_cp_size)))][1:] + group_sizes = [2**i for i in range(int(log2(dp_cp_size)))] for group_size in group_sizes: - group = ps.get_hybrid_data_context_parallel_groups(group_size=group_size) + group = ps.get_dynamic_data_context_parallel_groups(group_size=group_size) assert group.size() == group_size Utils.destroy_model_parallel() diff --git a/tools/prepare_cache.py b/tools/prepare_cache.py index 98d57db498a..8eebf1500a9 100644 --- a/tools/prepare_cache.py +++ b/tools/prepare_cache.py @@ -156,7 +156,6 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: context_parallel_size=args.context_parallel_size, data_parallel_size=args.data_parallel_size, sequence_parallel_size=args.tensor_model_parallel_size * args.sequence_parallel, - hybrid_context_parallel=args.hybrid_context_parallel, )