diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 0f016473b6a..ec2a0820e64 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -1,14 +1,72 @@ # Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved. -from typing import Any, List, Optional +import enum +from typing import Any, Dict, List, Optional, Type import torch from megatron.core import parallel_state +from megatron.core.datasets.data_schedule_utils import ( + broadcast_scalars, + broadcast_tensor, + broadcast_to_pp_group, + build_packed_microbatches, + create_data_iterator, + get_batch_and_global_seqlens, + reroute_samples_to_dcp_ranks, +) +from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices +from megatron.core.packed_seq_params import ( + PackedSeqParams, + get_thd_padding_kwargs, + pad_sequence_for_thd, +) from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler from megatron.core.process_groups_config import ProcessGroupCollection +def _build_thd_padding_mask( + cu_seqlens: torch.Tensor, cu_seqlens_padded: torch.Tensor +) -> torch.Tensor: + """Build a 1D THD padding mask from scheduler sequence metadata.""" + assert cu_seqlens.dim() == 1 + assert cu_seqlens_padded.dim() == 1 + assert cu_seqlens.numel() == cu_seqlens_padded.numel() + + total_tokens = int(cu_seqlens_padded[-1].item()) + if total_tokens == 0: + return torch.empty((0,), dtype=torch.bool, device=cu_seqlens.device) + + num_sequences = cu_seqlens.numel() - 1 + if num_sequences <= 0: + return torch.ones((total_tokens,), dtype=torch.bool, device=cu_seqlens.device) + + positions = torch.arange( + total_tokens, dtype=cu_seqlens_padded.dtype, device=cu_seqlens_padded.device + ) + seq_indices = torch.searchsorted(cu_seqlens_padded[1:].contiguous(), positions, right=True) + + valid_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).clamp(min=0) + valid_ends = cu_seqlens_padded[:-1] + valid_lengths + return positions >= valid_ends[seq_indices] + + +def _sanitize_thd_padding_values(batch: Dict[str, Any], padding_mask: torch.Tensor) -> None: + """Replace padded token-like slots with safe neutral values in-place.""" + assert padding_mask.dim() == 1 + pad_values = {'tokens': 0, 'labels': 0, 'loss_mask': 0.0, 'position_ids': 0} + for key, pad_value in pad_values.items(): + tensor = batch.get(key) + if tensor is None: + continue + assert tensor.dim() == 1, f"{key} must be 1D before CP slicing, got {tensor.dim()}D" + assert tensor.numel() == padding_mask.numel(), ( + f"{key} length ({tensor.numel()}) must match padding_mask length " + f"({padding_mask.numel()}) before CP slicing." + ) + batch[key] = tensor.masked_fill(padding_mask, pad_value) + + class HybridCPDataLoaderWrapper: """ A wrapper class that wraps around an existing data_iterator. @@ -57,7 +115,6 @@ def get_global_seqlens(self, subsample_seqlens: torch.Tensor) -> List[int]: Gathers the sequence lengths of all subsamples from all DP ranks. Each DP rank loads the same number of microbatches but each microbatch may have a different number of subsamples. - We find the number of subsamples each rank holds and then gather the sequence lengths of all subsamples from all ranks. """ @@ -299,3 +356,661 @@ def __next__(self) -> Any: batch, global_ids_this_rank, global_id_seqlens, sample_id_groups, offsets ) return samples_this_rank_with_id, sample_id_groups + + +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], + max_num_seqs: Optional[int] = None, + ): + """ + 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. + max_num_seqs: Optional cap on the number of real packed sequences + per microbatch. This excludes any dummy sequence later appended for + THD padding. + """ + 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 + self.max_num_seqs = max_num_seqs + + 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 + + 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 and ( + self.max_num_seqs is None or len(single_microbatch) < self.max_num_seqs + ): + 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. Schedule samples into groups + 4. Reroute samples to DCP ranks + 5. Build packed microbatches + 6. Calculate FLOPs info + 7. Broadcast to PP group (for middle PP stages) + 8. Broadcast to TP group (for non-TP-0 ranks) + 9. Handle VPP if enabled + + 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() + + # Handle VPP: extract the correct data_iterator for this PP stage + if ( + config.virtual_pipeline_model_parallel_size is not None + and config.virtual_pipeline_model_parallel_size > 1 + ): + # if enable VPP, data_iterator is a list of data_iterators for each VPP stage, + # and only the first and last stage rank will have data_iterator, + # other stages will have None. + assert len(data_iterator) == config.virtual_pipeline_model_parallel_size + if pp_group.rank() == 0: + # the first stage + data_iterator = data_iterator[0] + elif pp_group.rank() == pp_group.size() - 1: + # the last stage + data_iterator = data_iterator[-1] + else: + data_iterator = None + + # data_iterator is not None when TP rank 0, with PP stage 0 or -1. + if data_iterator is not None: + assert tp_group.rank() == 0 and ( + pp_group.rank() == 0 or pp_group.rank() == pp_group.size() - 1 + ), f"Only TP rank 0 and PP stage 0 or -1 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: 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 4: 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) + + 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) + ] + + # Step 5: Build packed microbatches + new_samples = build_packed_microbatches(grouped_samples, dev) + + # Step 6: 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) + + # Step 7: Broadcast to PP group (for middle PP stages) + if tp_group.rank() == 0: + ( + new_samples, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) = broadcast_to_pp_group( + new_samples, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + pp_group, + dev, + ) + + # Step 8: 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 9: create data_iterator and handle VPP if enabled + new_data_iterator = create_data_iterator(new_samples, pp_group, tp_group, config) + + return ( + new_data_iterator, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) + + +class PackingSchedulerEnum(enum.Enum): + """Enum for supported sequence packing algorithms.""" + + DP_BALANCED = "dp_balanced" + + +scheduler_map: Dict[PackingSchedulerEnum, Type[BasePackingScheduler]] = { + PackingSchedulerEnum.DP_BALANCED: DpBalancedScheduler +} + + +def _get_scheduler_max_real_num_seqs(config) -> Optional[int]: + """Return the scheduler cap for real THD sequences. + + ``thd_max_packed_sequences`` is the final static THD capacity, including the + optional dummy sequence appended for a padding tail. The dp_balanced + scheduler only packs real sequences, so reserve one slot when dummy-tail + padding is enabled. + """ + max_num_seqs = getattr(config, 'thd_max_packed_sequences', None) + if max_num_seqs is None: + return None + + max_num_seqs = int(max_num_seqs) + if max_num_seqs < 1: + raise ValueError(f"thd_max_packed_sequences must be >= 1, got {max_num_seqs}.") + + if getattr(config, 'pad_packed_seq_alignment', None) is not None and getattr( + config, 'pad_packed_seq_by_appending_dummy_seq', True + ): + if max_num_seqs < 2: + raise ValueError( + "thd_max_packed_sequences must be >= 2 when THD padding appends a dummy " + "sequence, because thd_max_packed_sequences includes that dummy sequence." + ) + return max_num_seqs - 1 + + return max_num_seqs + + +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 + + # Convert string to enum + scheduler_type = config.sequence_packing_scheduler + scheduler_type = PackingSchedulerEnum[scheduler_type.upper()] + + scheduler_max_num_seqs = ( + _get_scheduler_max_real_num_seqs(config) + if scheduler_type == PackingSchedulerEnum.DP_BALANCED + else getattr(config, 'thd_max_packed_sequences', None) + ) + + scheduler = scheduler_map[scheduler_type]( + config.max_seqlen_per_dp_cp_rank, + cp_size, + dp_size, + # When VPP is enabled, align num_micro_batches to this multiple. + ( + None + if config.virtual_pipeline_model_parallel_size is None + else config.microbatch_group_size_per_vp_stage + ), + max_num_seqs=scheduler_max_num_seqs, + ) + + ( + 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, + pg_collection: Optional[ProcessGroupCollection] = None, + config=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. + config: Model parallel config used for optional THD packed-sequence padding. + When None or config.pad_packed_seq_alignment is None, no padding is applied. + Returns: + tuple of (tokens, labels, loss_mask, attention_mask, position_ids, + packed_seq_params, padding_mask) + """ + + 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 is_first_stage: + batch_keys.append('tokens') + batch_keys.append('position_ids') + if is_last_stage: + 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 = {} + + # Build padding_mask before CP slicing while tensors still have the full + # packed length represented by cu_seqlens_padded[-1]. + if is_tp_rank_0: + batch['padding_mask'] = _build_thd_padding_mask( + batch['cu_seqlens'], batch['cu_seqlens_padded'] + ) + _sanitize_thd_padding_values(batch, batch['padding_mask']) + + # Partition padding_mask for context parallel on every PP stage. Partition + # token-like tensors only on stages that own them. + if is_tp_rank_0: + cp_size = cp_group.size() + cp_rank = cp_group.rank() + # If cp_size == 1, no need to do further processing. + if cp_size > 1: + # 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"] + total_tokens = int(cu_seqlens[-1].item()) + index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) + cp_slice_keys = ['padding_mask'] + if is_first_or_last_stage or mtp_on_this_rank: + cp_slice_keys.extend(['tokens', 'position_ids', 'labels', 'loss_mask']) + for key in cp_slice_keys: + if key in batch and batch[key] is not None: + batch[key] = batch[key].index_select(0, index) + + # 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 padding_mask is prepared on every PP stage. + # Tokens/labels/loss_mask/position_ids use the same length on stages that own them. + if is_tp_rank_0: + total_tokens = torch.tensor(batch['padding_mask'].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" on all 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" on all ranks. + if is_last_stage: + 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 "padding_mask" on all TP ranks. + if is_tp_rank_0: + assert batch['padding_mask'].dtype == torch.bool + batch['padding_mask'] = batch['padding_mask'].view(1, total_tokens) + else: + batch['padding_mask'] = torch.empty([1, total_tokens], dtype=torch.bool, device=dev) + + # Step4: 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) + + # 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['padding_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) + + # Extract the data from batch after broadcasting. + tokens = batch['tokens'] + position_ids = batch['position_ids'] + labels = batch['labels'] + loss_mask = batch['loss_mask'] + padding_mask = batch['padding_mask'] + cu_seqlens = batch['cu_seqlens'] + cu_seqlens_padded = batch['cu_seqlens_padded'] + max_seqlen = batch['max_seqlen'].item() + + # Use padded cumulative lengths for THD partitioning so token slices follow + # the padded sequence boundaries consumed by attention kernels. + 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=None, + cp_group=None, + pad_between_seqs=False, + ) + + # Pad the already-packed THD tensors at the end when requested. CUDA Graph + # additionally pads cu_seqlens tensors to thd_max_packed_sequences + 1 entries. + pad_alignment = ( + getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None + ) + if pad_alignment is not None and packed_seq_params is not None: + alignment, target_len, max_num_seqs = get_thd_padding_kwargs( + pad_alignment, + getattr(config, 'max_seqlen_per_dp_cp_rank', None), + getattr(config, 'thd_max_packed_sequences', None), + getattr(config, 'cuda_graph_impl', 'none') != 'none', + ) + tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = ( + pad_sequence_for_thd( + tokens, + labels, + loss_mask, + position_ids, + packed_seq_params, + alignment=alignment, + target_len=target_len, + max_num_seqs=max_num_seqs, + pad_by_appending_dummy_seq=getattr( + config, 'pad_packed_seq_by_appending_dummy_seq', True + ), + padding_mask=padding_mask, + 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, padding_mask diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py new file mode 100644 index 00000000000..86eb7d2e1fc --- /dev/null +++ b/megatron/core/datasets/data_schedule_utils.py @@ -0,0 +1,540 @@ +# Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved. + +from typing import Dict, List + +import torch + +from megatron.core.rerun_state_machine import RerunDataIterator + + +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]["tokens"].device + original_seq_lens = [] + padded_seq_lens = [] + 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 ["tokens", "labels", "loss_mask", "position_ids"]: + 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, 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) + + tokens = _pack_tensors([sample["tokens"] for sample in samples]) + labels = _pack_tensors([sample["labels"] for sample in samples]) + loss_mask = _pack_tensors([sample["loss_mask"] for sample in samples]) + position_ids = _pack_tensors([sample["position_ids"] for sample in samples]) + + new_sample = {} + new_sample["tokens"] = tokens + new_sample["labels"] = labels + new_sample["loss_mask"] = loss_mask + new_sample["position_ids"] = position_ids + + 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 + + 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_to_pp_group( + new_samples, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + pp_group, + dev, +): + """ + Broadcast num_micro_batches, seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch and metadata to middle PP stages. + Before this broadcast, the new_samples on middle PP stages are None, + after this broadcast, the new_samples on middle PP stages contain the metadata but + without tokens, labels, loss_mask, position_ids. + """ + + pp_src_rank = torch.distributed.get_process_group_ranks(pp_group)[0] + + if pp_group.size() > 2: + if pp_group.rank() == 0: + cu_seqlens_lengths = torch.tensor( + [sample["cu_seqlens"].numel() for sample in new_samples], + dtype=torch.float32, + device=dev, + ) + cu_seqlens_padded_lengths = torch.tensor( + [sample["cu_seqlens_padded"].numel() for sample in new_samples], + dtype=torch.float32, + device=dev, + ) + tensor_list = [ + torch.tensor( + [ + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ], + dtype=torch.float32, + device=dev, + ) + ] + for sample in new_samples: + tensor_list.append(sample["max_seqlen"].reshape(1)) + tensor_list.append(cu_seqlens_lengths) + tensor_list.append(cu_seqlens_padded_lengths) + for sample in new_samples: + tensor_list.append(sample["cu_seqlens"]) + tensor_list.append(sample["cu_seqlens_padded"]) + info_to_broadcast = torch.cat(tensor_list, dim=0).to(device=dev, dtype=torch.float32) + info_length_tensor = torch.tensor( + info_to_broadcast.shape[0], dtype=torch.int32, device=dev + ) + broadcast_tensor(info_length_tensor, pp_src_rank, pp_group) + broadcast_tensor(info_to_broadcast, pp_src_rank, pp_group) + else: + info_length_tensor = torch.tensor(0, dtype=torch.int32, device=dev) + broadcast_tensor(info_length_tensor, pp_src_rank, pp_group) + info_to_broadcast = torch.empty( + info_length_tensor.item(), dtype=torch.float32, device=dev + ) + broadcast_tensor(info_to_broadcast, pp_src_rank, pp_group) + if pp_group.rank() != pp_group.size() - 1: + # Middle PP stages receive the broadcasted info and unpack it. + # Cu-seqlens lengths are encoded explicitly so zero values inside + # the payload cannot be mistaken for tensor boundaries. + num_micro_batches = int(info_to_broadcast[0].item()) + seqlen_sum_this_global_batch = info_to_broadcast[1].item() + seqlen_squared_sum_this_global_batch = info_to_broadcast[2].item() + + cursor = 3 + max_seqlens = info_to_broadcast[cursor : cursor + num_micro_batches] + cursor += num_micro_batches + cu_seqlens_lengths = info_to_broadcast[ + cursor : cursor + num_micro_batches + ].to(torch.int64) + cursor += num_micro_batches + cu_seqlens_padded_lengths = info_to_broadcast[ + cursor : cursor + num_micro_batches + ].to(torch.int64) + cursor += num_micro_batches + + new_samples = [] + for i in range(num_micro_batches): + cu_seqlens_len = int(cu_seqlens_lengths[i].item()) + cu_seqlens_padded_len = int(cu_seqlens_padded_lengths[i].item()) + new_sample = {} + new_sample["max_seqlen"] = max_seqlens[i].to(torch.int32) + new_sample["cu_seqlens"] = info_to_broadcast[ + cursor : cursor + cu_seqlens_len + ].to(torch.int32) + cursor += cu_seqlens_len + new_sample["cu_seqlens_padded"] = info_to_broadcast[ + cursor : cursor + cu_seqlens_padded_len + ].to(torch.int32) + cursor += cu_seqlens_padded_len + new_samples.append(new_sample) + + return ( + new_samples, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) + + +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, pp_group, tp_group, config): + """Handle virtual pipeline parallelism.""" + 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: + if pp_group.rank() == 0 or pp_group.rank() == pp_group.size() - 1: + metadata = [ + {k: sample[k] for k in ["max_seqlen", "cu_seqlens", "cu_seqlens_padded"]} + for sample in new_samples + ] + if pp_group.rank() == 0: + new_data_iterator = [RerunDataIterator(iter(new_samples))] + [ + RerunDataIterator(iter(metadata)) for _ in range(vpp_size - 1) + ] + else: + new_data_iterator = [ + RerunDataIterator(iter(metadata)) for _ in range(vpp_size - 1) + ] + [RerunDataIterator(iter(new_samples))] + else: + # on middle PP stages, the new_samples are the metadata + metadata = new_samples + new_data_iterator = [RerunDataIterator(iter(metadata)) for _ in range(vpp_size)] + 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(1, 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( + grouped_samples: List[List[Dict[str, torch.Tensor]]], dev: torch.device +) -> List[Dict[str, torch.Tensor]]: + """Build packed samples for each microbatch.""" + num_micro_batches = len(grouped_samples) + seg_starts: List[int] = [0] + original_lens_tensors = [] + padded_lens_tensors = [] + + 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]] + new_sample = _pack_sequences(samples, lens_padded, lens_original, 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 diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 166a47c084a..f707321a6ee 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/megatron/core/datasets/gpt_dataset.py @@ -81,6 +81,26 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig): """When True, return cu_seqlens marking document boundaries within each sample so that attention is restricted to individual documents.""" + sft_mock_dataset_config_json: Optional[str] = None + """This config provides the necessary information for the mock dataset.""" + + sequence_packing_scheduler: Optional[str] = None + """Scheduler for sequence packing and hybrid context parallel. + dp_balanced: DP-balanced scheduler for sequence packing. + """ + + varlen_mock_dataset_config_json: Optional[str] = None + """Mock-dataset config (same JSON schema as ``sft_mock_dataset_config_json``) + used by the ``--use-varlen-dataset`` path; kept separate so the varlen path + does not implicitly inherit SFT-specific knobs.""" + + varlen_sbhd_validation: bool = False + """When True, :class:`VarlenDataset.__getitem__` emits SBHD samples padded + to ``sequence_length`` (no ``cu_seqlens`` / ``original_seq_len`` / + ``padded_seq_len``), bypassing the packed-sequence path. Used to obtain a + SBHD reference run that mirrors the THD path's tokenization but skips all + packing — useful for THD numerical-correctness validation.""" + def __post_init__(self) -> None: """Do asserts and set fields post init""" super().__post_init__() @@ -91,6 +111,12 @@ def __post_init__(self) -> None: assert self.reset_attention_mask is not None assert self.eod_mask_loss is not None + if self.varlen_sbhd_validation: + assert not self.hybrid_context_parallel, ( + "--varlen-sbhd-validation is incompatible with " + "--hybrid-context-parallel (SBHD mode is not packed)." + ) + self.token_dtype_code = ( None if self.tokenizer.vocab_size is None diff --git a/megatron/core/datasets/readme.md b/megatron/core/datasets/readme.md index 58721b7471b..40019642469 100644 --- a/megatron/core/datasets/readme.md +++ b/megatron/core/datasets/readme.md @@ -204,6 +204,28 @@ If the later training job does not specify `--global-batch-size` (which is neede `tools/prepare_cache.py` does not support `--mock-data`, `--sft`, `--fim-data`, or `--step-batch-size-schedule`. +## Packing Scheduler + +The packing scheduler re-schedules variable-length sequences across DP×CP ranks to improve GPU utilization. It is built around the following modules: + +### `data_schedule` + +This module contains the high-level scheduling logic and entry points: + +- **`HybridCPDataLoaderWrapper`**: A wrapper class for hybrid context parallel (CP) scheduling. For every `__next__` call, it: (1) pulls a batch of packed samples from each DP rank, (2) gathers sequence lengths across the DP group, (3) schedules sub-samples using the `BalancedCPScheduler`, (4) reroutes sub-samples to the correct DPxCP ranks via all-to-all communication. + +- **`BasePackingScheduler`**: Abstract base class for packing schedulers. Defines the interface for `get_groups_and_subsamples()` (scheduling algorithm) and `run()` (full scheduling pipeline including fetch, schedule, reroute, pack, broadcast, and VPP handling). + +- **`DpBalancedScheduler`**: A concrete scheduler that packs sequences in their original order until reaching the max sequence length limit per DPxCP rank. Supports aligning the number of microbatches to DP size and VPP stage multiples. + +- **`wrap_data_iterator()`**: Top-level entry point that wraps an existing `data_iterator`. It creates the appropriate scheduler, runs the scheduling pipeline, broadcast metadata and new num_microbatches, returns a new data iterator along with the updated number of microbatches and FLOPs statistics. + +- **`get_batch_on_this_rank_for_sequence_packing()`**: Fetches and broadcasts a single packed microbatch for the current rank. Handles TP/PP broadcasting, constructs `PackedSeqParams` (with `cu_seqlens`, `max_seqlen`, `qkv_format=thd`), and optionally partitions sequences across CP ranks using Transformer Engine's `thd_get_partitioned_indices`. + +### `data_schedule_utils.py` + +This module contains the utility functions used by the schedulers. + ## Fast DataLoader initialization Especially for large-scale runs, DataLoader initialization can take several minutes, since it involves opening and memory-mapping multiple files and can significantly stress the filesystem. To speed up this process, we have developed the following three optimizations, controlled by configuration flags: diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index ac7d5c1da9b..e742d88daac 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -3393,3 +3393,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/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 19ff501af91..4bc29b4982e 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1771,6 +1771,7 @@ def apply_rotary_emb_query( cp_group=cp_group, mscale=mscale, mla_rotary_interleaved=config.multi_latent_attention, + max_seqlen=query_emb.size(0), ) return query diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 88bb070e105..0812b281a40 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -2,11 +2,26 @@ import warnings from dataclasses import dataclass, field -from typing import Callable, ContextManager, Literal, Optional +from typing import Callable, ContextManager, Literal, Optional, Union import torch +def _parse_pad_packed_seq_alignment(value): + """Parse THD packed-sequence padding alignment. + + Accepts ``"max"`` or a positive integer alignment. + """ + if value == "max": + return value + try: + return int(value) + except (TypeError, ValueError) as exc: + raise ValueError( + "pad_packed_seq_alignment must be 'max' or a positive integer alignment." + ) from exc + + @dataclass class ModelParallelConfig: """Base configuration for Megatron Core @@ -59,7 +74,7 @@ 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 @@ -69,6 +84,37 @@ class ModelParallelConfig: Please set max_seqlen_per_dp_cp_rank when using hybrid_context_parallel. """ + sequence_packing_scheduler: Optional[Literal['dp_balanced']] = None + """ + Scheduler for sequence packing and hybrid context parallel. + dp_balanced: DP-balanced scheduler for sequence packing. + """ + + pad_packed_seq_alignment: Optional[Union[int, Literal["max"]]] = field( + default=None, + metadata={ + "argparse_meta": { + "arg_names": ["--pad-packed-seq-alignment"], + "type": _parse_pad_packed_seq_alignment, + } + }, + ) + """Pad THD packed sequence tensors after packing. + + If set to ``max``, token-like tensors are padded to + max_seqlen_per_dp_cp_rank. If set to a positive integer N, token-like + tensors are padded to a multiple of N. + """ + + pad_packed_seq_by_appending_dummy_seq: bool = True + """Represent a THD packed-sequence padding tail by appending a dummy sequence. + + When disabled, token-like tensors are still padded according to + pad_packed_seq_alignment, but cu_seqlens sequence boundaries are not extended + for the padding tail. CUDA Graph static-input padding may still pad the + cu_seqlens tensors to thd_max_packed_sequences + 1 entries. + """ + expert_model_parallel_size: int = 1 """Distributes Moe Experts across sub data parallel dimension.""" @@ -423,6 +469,28 @@ def __post_init__(self): See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more details. """ + if self.pad_packed_seq_alignment is not None: + self.pad_packed_seq_alignment = _parse_pad_packed_seq_alignment( + self.pad_packed_seq_alignment + ) + if self.max_seqlen_per_dp_cp_rank is None: + raise ValueError( + "max_seqlen_per_dp_cp_rank must be set when pad_packed_seq_alignment " + "is enabled." + ) + if self.pad_packed_seq_alignment != "max": + if self.pad_packed_seq_alignment <= 0: + raise ValueError( + "pad_packed_seq_alignment must be 'max' or a positive integer " "alignment." + ) + if self.pad_packed_seq_alignment > self.max_seqlen_per_dp_cp_rank: + raise ValueError( + "pad_packed_seq_alignment must not exceed " + "max_seqlen_per_dp_cp_rank " + f"({self.max_seqlen_per_dp_cp_rank}), got " + f"{self.pad_packed_seq_alignment}." + ) + 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/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index 9fab25a3fae..b77fbd4443c 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -195,18 +195,22 @@ def _apply_rotary_pos_emb_thd( mscale: float = 1.0, cp_group: torch.distributed.ProcessGroup = None, multi_latent_attention: Optional[bool] = None, + max_seqlen: Optional[int] = None, ) -> Tensor: - """A baseline implementation of applying RoPE for `thd` format. + """Apply RoPE for `thd` format using pure CUDA ops (CUDA Graph compatible). + + Replaces the original Python-loop + .tolist() implementation with vectorized + CUDA operations. No GPU->CPU syncs, compatible with CUDA Graph capture. Args: - t (Tensor): Input tensor T is of shape [t, h, d] - cu_seqlens(Tensor): Cumulative sum of sequence lengths in a batch for `t`, - with shape [b + 1] and dtype torch.int32. - freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d] - cp_group (torch.distributed.ProcessGroup): The context parallel group + t (Tensor): Input tensor of shape [total_tokens, h, d] + cu_seqlens (Tensor): Cumulative sequence lengths, shape [num_seqs + 1], int32. + freqs (Tensor): RoPE frequencies, shape [max_s, 1, 1, d] or [total_tokens, 1, 1, d] + cp_group: Context parallel group + max_seqlen: Global max sequence length for this packed batch when known. Returns: - Tensor: Shape [t, h, d]. The input tensor after applying RoPE. + Tensor: Shape [total_tokens, h, d]. Input with RoPE applied. """ if multi_latent_attention is not None: warnings.warn( @@ -219,53 +223,68 @@ def _apply_rotary_pos_emb_thd( raise ValueError("cp_group must be provided for THD format RoPE") cp_size = cp_group.size() cp_rank = cp_group.rank() - seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist() - sequence_splits = torch.split(t, seqlens) - total_seqlen = int(cu_seqlens[-1].item()) - has_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == total_seqlen - - # Handle two different frequency tensor formats: - # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains positions for the whole packed - # batch. Each sequence must therefore use its cu_seqlens offset when selecting the local CP - # front/back slices. For example, with cu_seqlens=[0, 4, 8], cp_size=2, rank 0 should use - # positions [0, 3, 4, 7], not [0, 3, 0, 3]. - # 2. Otherwise: freqs contains only max sequence length positions. Each packed sequence should - # reuse positions starting from 0, preserving the legacy THD behavior. - if has_packed_freqs: - # CASE 1: Exact mapping with offsets - local_freqs = [] - for i, x in enumerate(sequence_splits): - # cu_seqlens[i] is the starting offset of this sequence in the original batch - seq_start_offset = cu_seqlens[i].item() - local_freqs.append( - _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) - ) - freqs = torch.cat(local_freqs, dim=0) - return _apply_rotary_pos_emb_bshd( - t.unsqueeze(1), - freqs, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - ).squeeze(1) - - # CASE 2: Traditional mapping without offsets. Apply RoPE one sequence at a time so the second - # and later packed sequences do not look like continuations of the first sequence. - output = torch.empty_like(t) - output_offset = 0 - for x in sequence_splits: - freq_slice = _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs) - output_slice = _apply_rotary_pos_emb_bshd( - x.unsqueeze(1), - freq_slice, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - ).squeeze(1) - output.narrow(0, output_offset, x.size(0)).copy_(output_slice) - output_offset += x.size(0) - return output + total_tokens = t.shape[0] + device = t.device + + token_pos = torch.arange(total_tokens, device=device, dtype=torch.int64) + + # `cu_seqlens` describes the global packed sequence. With CP, `t` is already + # CP-partitioned, so build a local cumulative-length view before assigning + # local tokens to packed sequences. + cu_seqlens_i64 = cu_seqlens.to(torch.int64) + global_seq_lens = cu_seqlens_i64[1:] - cu_seqlens_i64[:-1] + local_seq_lens = global_seq_lens // cp_size if cp_size > 1 else global_seq_lens + local_cu_seqlens = torch.zeros_like(cu_seqlens_i64) + local_cu_seqlens[1:] = torch.cumsum(local_seq_lens, dim=0) + + # `searchsorted(..., right=True) - 1` returns the local sequence index. The + # clamp guards padded tokens that sit beyond the final real local token; they + # get a harmless frequency and are later masked out. + seq_idx = torch.searchsorted(local_cu_seqlens, token_pos, right=True) - 1 + seq_idx = seq_idx.clamp(min=0, max=cu_seqlens.shape[0] - 2) + + local_seq_start = local_cu_seqlens[seq_idx] + local_pos = token_pos - local_seq_start + local_seq_len = local_seq_lens[seq_idx] + global_seq_start = cu_seqlens_i64[seq_idx] + + if cp_size > 1: + cp_seg = local_seq_len // 2 + full_seqlen = local_seq_len * cp_size + is_first_half = local_pos < cp_seg + freq_pos = torch.where( + is_first_half, + cp_rank * cp_seg + local_pos, + full_seqlen - (cp_rank + 1) * cp_seg + (local_pos - cp_seg), + ) + else: + freq_pos = local_pos.to(torch.int64) + + assert max_seqlen is not None, ( + "max_seqlen must be provided for THD RoPE so packed-frequency offset " + "detection does not silently depend on tensor shape heuristics." + ) + exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) > max_seqlen + if exact_packed_freqs: + # `freqs` covers all positions across all sequences (used for non-1D + # RoPE / VLMs); shift by the per-sequence start offset so each token + # samples its absolute position. When `freqs` only spans one max-len + # sequence, no shift is needed. + freq_pos = freq_pos + global_seq_start + + # Padded positions can sit outside the frequency table. Clamp them into + # range; downstream padding masks exclude those positions from the result. + freq_pos = freq_pos.clamp(min=0, max=freqs.shape[0] - 1) + freqs_packed = freqs[freq_pos] + + return _apply_rotary_pos_emb_bshd( + t.unsqueeze(1), + freqs_packed, + rotary_interleaved=rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + ).squeeze(1) def apply_rotary_pos_emb( @@ -276,6 +295,7 @@ def apply_rotary_pos_emb( mscale: float = 1.0, cp_group: torch.distributed.ProcessGroup = None, mla_rotary_interleaved: bool = False, + max_seqlen: Optional[int] = None, ): """ Reroute to the appropriate apply_rotary_pos_emb function depending on @@ -343,6 +363,7 @@ def apply_rotary_pos_emb( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, cp_group=cp_group, + max_seqlen=max_seqlen, ) diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index bd598bb557a..c2398165d1b 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -1,8 +1,10 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. from dataclasses import dataclass +from typing import Literal, Optional, Tuple, Union import torch import torch.distributed as dist +import torch.nn.functional as F from torch import Tensor @@ -25,6 +27,7 @@ class PackedSeqParams: total_tokens: int = None seq_idx: Tensor = None tokens_per_sample: int = None + pad_between_seqs: Optional[bool] = None def __post_init__(self): """Pre-compute seq_idx for Mamba mixer CUDA graph compatibility. @@ -65,3 +68,411 @@ def __post_init__(self): .to(torch.int32) .unsqueeze(0) # Add a batch dimension ) + + +def _pad_seq_tensor(t: Optional[Tensor], target_len: int) -> Optional[Tensor]: + """Pad a [..., seq] tensor to ``target_len`` along the last dim with zeros. + + Asserts the actual length does not exceed ``target_len``: an oversize input + would silently desync the captured graph from replay shapes. + """ + if t is None: + return None + actual_len = t.shape[-1] + assert actual_len <= target_len, ( + f"Sequence-length tensor (last dim = {actual_len}) exceeds target " + f"({target_len}); refusing to silently truncate. Increase " + f"--max-seqlen-per-dp-cp-rank or filter overlong samples upstream." + ) + if actual_len == target_len: + return t + return F.pad(t, (0, target_len - actual_len), value=0) + + +def _pad_padding_mask(mask: Tensor, target_len: int) -> Tensor: + """Pad a [..., seq] bool padding mask to ``target_len`` with True.""" + actual_len = mask.shape[-1] + assert actual_len <= target_len, ( + f"Padding mask length ({actual_len}) exceeds target ({target_len}); " + "refusing to silently truncate." + ) + if actual_len == target_len: + return mask + + pad_shape = list(mask.shape) + pad_shape[-1] = target_len - actual_len + tail = torch.ones(pad_shape, dtype=mask.dtype, device=mask.device) + return torch.cat((mask, tail), dim=-1) + + +def _pad_cu_seqlens(cu_seqlens: Optional[Tensor], target_entries: int) -> Optional[Tensor]: + """Pad a cu_seqlens tensor to exactly ``target_entries`` entries. + + Asserts the actual entry count does not exceed ``target_entries``. An + oversized pack cannot be represented by the configured static cu_seqlens + buffer and would not match captured CUDA Graph replay shapes. + """ + if cu_seqlens is None: + return None + actual_entries = cu_seqlens.shape[0] + assert actual_entries <= target_entries, ( + f"Actual num_seqs ({actual_entries - 1}) exceeds thd_max_packed_sequences " + f"({target_entries - 1}). Increase --thd-max-packed-sequences, decrease " + f"--max-seqlen-per-dp-cp-rank, or filter shorter samples upstream so " + f"the packing scheduler stops earlier." + ) + if actual_entries == target_entries: + return cu_seqlens + pad_value = cu_seqlens[-1].item() + padded = torch.full( + (target_entries,), pad_value, dtype=cu_seqlens.dtype, device=cu_seqlens.device + ) + padded[:actual_entries] = cu_seqlens + return padded + + +def _append_dummy_seq(cu_seqlens: Optional[Tensor], dummy_end: int) -> Optional[Tensor]: + """Append a dummy sequence boundary to a cu_seqlens tensor. + + ``dummy_end`` is the padded target length. Appending it to both + ``cu_seqlens_*`` and ``cu_seqlens_*_padded`` represents the post-pack + alignment tail as an ordinary dummy sequence. That keeps every token row + covered by THD metadata without enabling TE's pad-between-sequences mode. + """ + if cu_seqlens is None: + return None + + dummy = torch.full((1,), int(dummy_end), dtype=cu_seqlens.dtype, device=cu_seqlens.device) + return torch.cat((cu_seqlens, dummy), dim=0) + + +def _round_up_to_alignment(value: int, alignment: int) -> int: + assert alignment > 0, f"Packed sequence padding alignment must be > 0, got {alignment}." + return ((value + alignment - 1) // alignment) * alignment + + +def get_thd_padding_kwargs( + pad_packed_seq_alignment: Union[int, Literal["max"]], + max_seqlen_per_dp_cp_rank: Optional[int], + thd_max_packed_sequences: Optional[int], + cuda_graph_static: bool, +) -> Tuple[Optional[int], Optional[int], Optional[int]]: + """Resolve ``pad_sequence_for_thd`` kwargs from the training config. + + ``--pad-packed-seq-alignment`` has two forms: + + - ``max`` pads token-like tensors to ``max_seqlen_per_dp_cp_rank``; + - a positive value pads token-like tensors to a multiple of that value. + + Padding cu_seqlens to ``thd_max_packed_sequences + 1`` is a CUDA Graph static-input + requirement. Eager pad-to-max should preserve sequence metadata so kernels + continue to see the real packed sequence boundaries. + """ + if cuda_graph_static: + return None, int(max_seqlen_per_dp_cp_rank), thd_max_packed_sequences + + if pad_packed_seq_alignment == "max": + return None, int(max_seqlen_per_dp_cp_rank), None + + return int(pad_packed_seq_alignment), None, None + + +def _resolve_thd_padding_lengths( + tokens: Optional[Tensor], + labels: Optional[Tensor], + loss_mask: Optional[Tensor], + position_ids: Optional[Tensor], + packed_seq_params: PackedSeqParams, + target_len: Optional[int], + alignment: Optional[int], + cp_group: Optional[dist.ProcessGroup] = None, + cp_size: Optional[int] = None, + cp_rank: Optional[int] = None, +) -> Tuple[int, int, int, int, torch.device]: + """Resolve local/global THD padding lengths without changing tensors. + + Returns: + local_actual_T: Current rank's token-like tensor length. + global_actual_T: Global packed length represented by THD metadata. + local_target_len: Current rank's padded token-like tensor length. + global_target_len: Global padded endpoint represented by THD metadata. + mask_device: Device used to build the returned padding mask. + """ + + cp_size, cp_rank = _resolve_thd_cp_geometry( + packed_seq_params, cp_group=cp_group, cp_size=cp_size, cp_rank=cp_rank + ) + + # Find the first token-like tensor that carries this rank's local length. + local_tensor_T = None + mask_device = None + for candidate in (tokens, labels, loss_mask, position_ids): + if candidate is not None: + local_tensor_T = int(candidate.shape[-1]) + mask_device = candidate.device + break + + # Prefer THD metadata for the global packed length when it is available. + has_local_tensor = local_tensor_T is not None + if packed_seq_params.cu_seqlens_q is not None: + global_actual_T = int(packed_seq_params.cu_seqlens_q[-1].item()) + if mask_device is None: + mask_device = packed_seq_params.cu_seqlens_q.device + else: + assert has_local_tensor, ( + "packed_seq_params.cu_seqlens_q must be available to derive padding_mask " + "when tokens/labels/loss_mask/position_ids are all None." + ) + global_actual_T = local_tensor_T * cp_size + + # Tensor path: use the already-sliced local shape and scale to the global endpoint. + if has_local_tensor: + local_actual_T = local_tensor_T + local_target_len = ( + int(target_len) + if target_len is not None + else _round_up_to_alignment(local_actual_T, alignment) + ) + global_target_len = local_target_len * cp_size + return local_actual_T, global_actual_T, local_target_len, global_target_len, mask_device + + # Metadata-only path: resolve the global padded endpoint first. + global_target_len = ( + int(target_len) * cp_size + if target_len is not None + else _round_up_to_alignment(global_actual_T, alignment) + ) + + # Under CP, ask TE which packed rows this rank would receive. + if cp_size > 1: + from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices + + partition_cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + # The number of selected rows is this rank's local actual length. + local_actual_T = int( + get_thd_partitioned_indices( + partition_cu_seqlens, global_actual_T, cp_size, cp_rank + ).numel() + ) + # Do the same for the padded endpoint; THD CP is not simple equal split. + local_target_len = int( + get_thd_partitioned_indices( + partition_cu_seqlens, global_target_len, cp_size, cp_rank + ).numel() + ) + else: + # Without CP, local and global metadata lengths are identical. + local_actual_T = global_actual_T + local_target_len = global_target_len + + return local_actual_T, global_actual_T, local_target_len, global_target_len, mask_device + + +def _resolve_thd_cp_geometry( + packed_seq_params: PackedSeqParams, + cp_group: Optional[dist.ProcessGroup] = None, + cp_size: Optional[int] = None, + cp_rank: Optional[int] = None, +) -> Tuple[int, int]: + """Resolve CP geometry for THD padding. + + Callers with a known CP group or explicit size/rank should pass it here. + Falling back to ``parallel_state`` preserves legacy call sites. + """ + if cp_group is not None: + return int(dist.get_world_size(group=cp_group)), int(dist.get_rank(group=cp_group)) + + if cp_size is not None: + cp_size = int(cp_size) + if cp_rank is not None: + return cp_size, int(cp_rank) + if cp_size == 1: + return cp_size, 0 + + if packed_seq_params.cp_group is not None: + cp_group = packed_seq_params.cp_group + return int(dist.get_world_size(group=cp_group)), int(dist.get_rank(group=cp_group)) + + if cp_size is None and packed_seq_params.local_cp_size is not None: + cp_size = int(packed_seq_params.local_cp_size) + if cp_size == 1: + return cp_size, 0 + + # Last resort for compatibility with older callers that do not thread CP + # geometry through PackedSeqParams. + from megatron.core import parallel_state + + if cp_size is None: + cp_size = int(parallel_state.get_context_parallel_world_size()) + if cp_rank is None: + cp_rank = int(parallel_state.get_context_parallel_rank()) if cp_size > 1 else 0 + return int(cp_size), int(cp_rank) + + +def pad_sequence_for_thd( + tokens: Optional[Tensor], + labels: Optional[Tensor], + loss_mask: Optional[Tensor], + position_ids: Optional[Tensor], + packed_seq_params: PackedSeqParams, + alignment: Optional[int] = None, + target_len: Optional[int] = None, + max_num_seqs: Optional[int] = None, + pad_by_appending_dummy_seq: bool = True, + padding_mask: Optional[Tensor] = None, + cp_group: Optional[dist.ProcessGroup] = None, + cp_size: Optional[int] = None, + cp_rank: Optional[int] = None, +) -> Tuple[ + Optional[Tensor], + Optional[Tensor], + Optional[Tensor], + Optional[Tensor], + PackedSeqParams, + Optional[Tensor], +]: + """Pad packed THD tensors after packing. + + This appends padding tokens to token-like tensors and returns a padding mask + for MoE auxiliary-loss/routing paths. + + Args: + tokens: Packed token tensor with sequence length on the last dimension, + or None on pipeline stages that do not own tokens. + labels: Packed label tensor with sequence length on the last dimension, + or None. + loss_mask: Packed loss mask tensor with sequence length on the last + dimension, or None. + position_ids: Packed position id tensor with sequence length on the + last dimension, or None. + packed_seq_params: THD metadata for the packed batch. + alignment: If set, round each CP-local token-like tensor length up to + this multiple. Exactly one of ``alignment`` and ``target_len`` must + be provided. + target_len: If set, pad token-like tensors to this CP-local length. + Exactly one of ``alignment`` and ``target_len`` must be provided. + max_num_seqs: If set, pad cu_seqlens tensors to + ``max_num_seqs + 1`` entries for static CUDA Graph inputs. + pad_by_appending_dummy_seq: If true, represent the post-pack padding + tail as an extra dummy sequence in cu_seqlens metadata. + padding_mask: Existing bool padding mask for already-packed tokens, + with True marking padding positions. + cp_group: Context-parallel process group for resolving local/global + THD padding lengths. If omitted, ``packed_seq_params.cp_group`` is + used when available. + cp_size: Explicit context-parallel world size used when no CP group is + available. + cp_rank: Explicit context-parallel rank used with ``cp_size``. + + Notes: + - THD CP slicing is defined by Transformer Engine. On metadata-only + stages, Megatron asks TE which packed rows this CP rank would receive + and uses that row count as the local length instead of assuming equal + division by CP size. + - When ``pad_by_appending_dummy_seq`` is true, the padding tail is also + represented as an ordinary dummy sequence in cu_seqlens metadata. + - ``max_num_seqs`` pads all four cu_seqlens tensors; this is required + by CUDA Graph replay because those tensors are graph inputs. + + Returns: + Padded (tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask) + padding_mask: [1, target] bool tensor, True at padding positions. + """ + assert (alignment is None) != ( + target_len is None + ), "Exactly one of alignment or target_len must be provided for THD padding." + + local_actual_T, global_actual_T, local_target_len, global_target_len, mask_device = ( + _resolve_thd_padding_lengths( + tokens, + labels, + loss_mask, + position_ids, + packed_seq_params, + target_len=target_len, + alignment=alignment, + cp_group=cp_group, + cp_size=cp_size, + cp_rank=cp_rank, + ) + ) + + # Reject individual packed sequences that cannot fit the resolved target. + if packed_seq_params.cu_seqlens_q is not None: + _cu = packed_seq_params.cu_seqlens_q + _individual_lens = _cu[1:] - _cu[:-1] + _max_individual = int(_individual_lens.max().item()) if _individual_lens.numel() > 0 else 0 + assert _max_individual <= global_target_len, ( + f"Individual request length ({_max_individual}) exceeds the global max sequence length " + f"({global_target_len}). Increase --max-seqlen-per-dp-cp-rank / alignment, " + f"or filter out overlong requests." + ) + + # Pad token-like tensors to the CP-local target length. + tokens = _pad_seq_tensor(tokens, local_target_len) + labels = _pad_seq_tensor(labels, local_target_len) + loss_mask = _pad_seq_tensor(loss_mask, local_target_len) + position_ids = _pad_seq_tensor(position_ids, local_target_len) + + # Copy THD metadata before optionally appending/padding sequence boundaries. + cu_seqlens_q = packed_seq_params.cu_seqlens_q + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + cu_seqlens_q_padded = packed_seq_params.cu_seqlens_q_padded + cu_seqlens_kv_padded = packed_seq_params.cu_seqlens_kv_padded + + # Represent post-pack padding as a dummy sequence when requested. + target_cu_entries = None if max_num_seqs is None else max_num_seqs + 1 + has_dummy_padding_seq = pad_by_appending_dummy_seq and global_target_len > global_actual_T + dummy_seq_len = global_target_len - global_actual_T if has_dummy_padding_seq else 0 + + if has_dummy_padding_seq: + cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, global_target_len) + cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, global_target_len) + cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len) + cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len) + + # Pad cu_seqlens entry counts for static CUDA Graph inputs. + if target_cu_entries is not None: + cu_seqlens_q = _pad_cu_seqlens(cu_seqlens_q, target_cu_entries) + cu_seqlens_kv = _pad_cu_seqlens(cu_seqlens_kv, target_cu_entries) + cu_seqlens_q_padded = _pad_cu_seqlens(cu_seqlens_q_padded, target_cu_entries) + cu_seqlens_kv_padded = _pad_cu_seqlens(cu_seqlens_kv_padded, target_cu_entries) + + # Rebuild PackedSeqParams with the padded tensor and metadata shapes. + padded_params = PackedSeqParams( + qkv_format=packed_seq_params.qkv_format, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + max_seqlen_q=( + global_target_len + if target_cu_entries is not None + else max(packed_seq_params.max_seqlen_q, dummy_seq_len) + ), + max_seqlen_kv=( + global_target_len + if target_cu_entries is not None + else max(packed_seq_params.max_seqlen_kv, dummy_seq_len) + ), + local_cp_size=packed_seq_params.local_cp_size, + cp_group=packed_seq_params.cp_group, + total_tokens=local_target_len if target_cu_entries is None else None, + tokens_per_sample=packed_seq_params.tokens_per_sample, + pad_between_seqs=False if has_dummy_padding_seq else packed_seq_params.pad_between_seqs, + ) + + # True marks padded local token slots for routing/loss paths. + tail_padding_mask = ( + torch.arange(local_target_len, device=mask_device).unsqueeze(0) >= local_actual_T + ) + if padding_mask is None: + padding_mask = tail_padding_mask + else: + padding_mask = _pad_padding_mask(padding_mask, local_target_len) | tail_padding_mask + + return tokens, labels, loss_mask, position_ids, padded_params, padding_mask diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index e67c498e2cc..1bfa5cbc57e 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -706,6 +706,9 @@ def forward_backward_no_pipelining( pg_collection.tp_dp_cp = parallel_state.get_tensor_and_data_parallel_group( with_context_parallel=True ) + pg_collection.dp = parallel_state.get_data_parallel_group( + with_context_parallel=False, partial_data_parallel=False + ) elif pg_collection is not None: assert hasattr(pg_collection, 'tp'), "pg_collection must have tp" @@ -1036,6 +1039,9 @@ def forward_backward_pipelining_with_interleaving( pg_collection.tp_dp_cp = parallel_state.get_tensor_and_data_parallel_group( with_context_parallel=True ) + pg_collection.dp = parallel_state.get_data_parallel_group( + with_context_parallel=False, partial_data_parallel=False + ) elif p2p_communicator is not None and pg_collection is not None: model_type = get_model_type(model[0]) @@ -2195,6 +2201,9 @@ def forward_backward_pipelining_without_interleaving( pg_collection.tp_dp_cp = parallel_state.get_tensor_and_data_parallel_group( with_context_parallel=True ) + pg_collection.dp = parallel_state.get_data_parallel_group( + with_context_parallel=False, partial_data_parallel=False + ) elif p2p_communicator is not None and pg_collection is not None: assert hasattr(p2p_communicator, 'config'), "p2p_communicator must have a config" diff --git a/megatron/core/tokenizers/text/libraries/null_tokenizer.py b/megatron/core/tokenizers/text/libraries/null_tokenizer.py index 96a0d3afd57..160aaa8bcb3 100644 --- a/megatron/core/tokenizers/text/libraries/null_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/null_tokenizer.py @@ -93,7 +93,7 @@ def eod(self): @property def pad_id(self): - """Returns pad token.""" + """Returns id of padding token.""" return self._pad_id @property diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 92e2bccb8cf..755400dc609 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1451,8 +1451,11 @@ def forward( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv else: cu_seqlens_q = cu_seqlens_kv = None + rope_max_seqlen_q = rope_max_seqlen_kv = None if split_qkv: if q_pos_emb is not None: @@ -1465,6 +1468,7 @@ def forward( cu_seqlens=cu_seqlens_q, mscale=self._yarn_concentration_factor, cp_group=self.pg_collection.cp, + max_seqlen=rope_max_seqlen_q, ) else: query = inference_context.apply_rotary_emb_query( @@ -1483,6 +1487,7 @@ def forward( cu_seqlens=cu_seqlens_kv, mscale=self._yarn_concentration_factor, cp_group=self.pg_collection.cp, + max_seqlen=rope_max_seqlen_kv, ) else: query, key, value = apply_fused_qkv_rotary_pos_emb( diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 210f39fa217..fb5b9cee934 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -20,6 +20,7 @@ import torch from torch.utils._pytree import tree_map as tree_map_pyt +from megatron.core import parallel_state from megatron.core.num_microbatches_calculator import get_num_microbatches from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import ( @@ -1678,7 +1679,14 @@ class TECudaGraphHelper: """ def __init__( - self, model, config, seq_length, micro_batch_size, optimizers=[], pg_collection=None + self, + model, + config, + seq_length, + micro_batch_size, + optimizers=[], + pg_collection=None, + thd_sequence_length_upper_bound=None, ): assert HAVE_TE_GRAPHS, "CUDA Graphs are not supported without TE." assert ( @@ -1694,12 +1702,14 @@ def __init__( self.model = model self.config = config self.seq_length = seq_length + self.thd_sequence_length_upper_bound = thd_sequence_length_upper_bound self.micro_batch_size = micro_batch_size self.optimizers = optimizers self.pg_collection = pg_collection if self.pg_collection is None: self.pg_collection = ProcessGroupCollection.use_mpu_process_groups() self.tp_group = self.pg_collection.tp + self.dp_group = self.pg_collection.dp self.dp_cp_group = self.pg_collection.dp_cp self.pp_group = self.pg_collection.pp from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator @@ -1904,6 +1914,12 @@ def get_rotary_pos_emb(transformer_module, transformer_input): static_inputs = layer.get_layer_static_inputs(self.seq_length, self.micro_batch_size) + if self._needs_full_local_padding_mask(layer, chunk_of_the_layer, static_inputs): + local_slen = self.config.max_seqlen_per_dp_cp_rank + static_inputs["padding_mask"] = torch.zeros( + 1, local_slen, dtype=torch.bool, device=torch.cuda.current_device() + ) + from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.transformer_layer import TransformerLayer @@ -2075,6 +2091,160 @@ def _get_amax_reduction_group(self, with_context_parallel=False, tp_only_amax_re assert self.pg_collection.tp is not None return self.pg_collection.tp + def _should_use_dynamic_microbatch_slots(self) -> bool: + """Whether to capture a bounded number of graph slots and reuse them by modulo.""" + return bool(getattr(self.config, "cuda_graph_dynamic_microbatches", False)) + + def _needs_full_local_padding_mask(self, layer, chunk, static_inputs) -> bool: + """Whether this layer's static padding_mask needs full max_seqlen_per_dp_cp_rank. + + For the post_process chunk (last PP/VPP chunk that holds labels), + padding_mask arrives at the full CP-local length because: + 1. labels are present -> actual_T_is_local=True -> no CP re-partition; + 2. pre_process=False -> _preprocess does not scatter under SP. + Other non-pre_process chunks (intermediate VPP) have no data, so their + captured padding_mask stays at the default scattered size from + `get_layer_static_inputs` (~max_seqlen/CP/TP). + + Returns True only for that post_process-with-data case under THD CUDA + Graph + SP + PP>1. + """ + return ( + hasattr(layer, "_is_thd_cuda_graph") + and layer._is_thd_cuda_graph() + and self.config.sequence_parallel + and self.config.pipeline_model_parallel_size > 1 + and not getattr(chunk, "pre_process", True) + and getattr(chunk, "post_process", False) + and "padding_mask" in static_inputs + ) + + @staticmethod + def _get_required_num_microbatch_slots_from_order(order, num_model_chunks): + """Infer the minimum safe slot count from a PP/VPP order. + + The slot count is defined as the maximum number of real microbatches whose forward has + happened but whose corresponding backward for the same chunk has not completed yet. + This is the exact liveness condition for whether a static buffer/graph slot can be reused. + """ + outstanding = [0] * num_model_chunks + max_outstanding = [0] * num_model_chunks + + for c_id in order: + if ceil(c_id) != c_id: + continue + model_chunk_idx = abs(int(ceil(c_id))) - 1 + if c_id > 0: + outstanding[model_chunk_idx] += 1 + max_outstanding[model_chunk_idx] = max( + max_outstanding[model_chunk_idx], outstanding[model_chunk_idx] + ) + else: + outstanding[model_chunk_idx] -= 1 + assert outstanding[model_chunk_idx] >= 0, ( + "Invalid PP/VPP schedule: negative outstanding microbatches while " + f"inferring CUDA graph slots for chunk {model_chunk_idx}." + ) + + assert all(count == 0 for count in outstanding), ( + "Invalid PP/VPP schedule: outstanding microbatches did not drain to zero when " + f"inferring CUDA graph slots. outstanding={outstanding}" + ) + return max(1, max(max_outstanding, default=1)) + + def _get_probe_num_microbatches_for_dynamic_slots(self): + """Return a topology-only probe microbatch count for slot inference.""" + pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size() + if pipeline_parallel_size == 1 and not self.config.overlap_moe_expert_parallel_comm: + return 1 + + group_size = self.config.microbatch_group_size_per_vp_stage + if group_size is None: + group_size = pipeline_parallel_size + + return max( + pipeline_parallel_size * max(1, self.num_model_chunks) * 4, + group_size * max(1, self.num_model_chunks) * 2, + 1, + ) + + @staticmethod + def _get_dp_balanced_thd_max_num_microbatches( + global_batch_size, + dp_size, + cp_size, + max_seqlen_per_dp_cp_rank, + max_sequence_length, + microbatch_group_size_per_vp_stage=None, + max_num_seqs=None, + ): + """Return the packed-microbatch upper bound for dp_balanced THD packing.""" + assert global_batch_size >= 1 + assert dp_size >= 1 + assert cp_size >= 1 + assert max_seqlen_per_dp_cp_rank >= 1 + assert max_sequence_length >= 1 + + max_seq_len_all_ranks = max_seqlen_per_dp_cp_rank * cp_size + seqs_per_pack = max(1, max_seq_len_all_ranks // max_sequence_length) + if max_num_seqs is not None: + seqs_per_pack = min(seqs_per_pack, max(1, int(max_num_seqs))) + + num_packed_sequences = math.ceil(global_batch_size / seqs_per_pack) + multiple = dp_size * ( + microbatch_group_size_per_vp_stage + if microbatch_group_size_per_vp_stage is not None + else 1 + ) + num_packed_sequences = math.ceil(num_packed_sequences / multiple) * multiple + return max(1, num_packed_sequences // dp_size) + + def _get_thd_varlen_max_num_microbatches( + self, runtime_num_microbatches, microbatch_group_size_per_vp_stage + ): + """Return the THD packing upper bound used for dynamic CUDA graph capture.""" + if self.config.sequence_packing_scheduler != 'dp_balanced': + return runtime_num_microbatches, "runtime" + if self.config.max_seqlen_per_dp_cp_rank is None: + return runtime_num_microbatches, "runtime" + + dp_size = self.dp_group.size() + cp_size = self.dp_cp_group.size() // dp_size + global_batch_size = runtime_num_microbatches * self.micro_batch_size * dp_size + # Use the dataset-produced padded sequence length upper bound when available. + # Do not use max_seqlen_per_dp_cp_rank here: under CP it is only the per-rank + # token budget, not the max length of one input sample before packing. + max_sequence_length = ( + self.thd_sequence_length_upper_bound + if self.thd_sequence_length_upper_bound is not None + else self.seq_length + ) + + max_num_seqs = getattr(self.config, 'thd_max_packed_sequences', None) + if max_num_seqs is not None: + max_num_seqs = int(max_num_seqs) + if getattr(self.config, 'pad_packed_seq_alignment', None) is not None and getattr( + self.config, 'pad_packed_seq_by_appending_dummy_seq', True + ): + max_num_seqs -= 1 + + return ( + self._get_dp_balanced_thd_max_num_microbatches( + global_batch_size, + dp_size, + cp_size, + int(self.config.max_seqlen_per_dp_cp_rank), + int(max_sequence_length), + microbatch_group_size_per_vp_stage=( + None + if self.config.virtual_pipeline_model_parallel_size is None + else microbatch_group_size_per_vp_stage + ), + max_num_seqs=max_num_seqs, + ), + "thd_varlen_upper_bound", + ) + def _get_cuda_graph_input_data(self): """ Create the CUDA Graph capturing input data. @@ -2087,26 +2257,93 @@ def _get_cuda_graph_input_data(self): get_schedule_table, ) + microbatch_group_size_per_vp_stage = self.config.microbatch_group_size_per_vp_stage + if microbatch_group_size_per_vp_stage is None: + microbatch_group_size_per_vp_stage = ( + parallel_state.get_pipeline_model_parallel_world_size() + ) + # If PP is not enabled, we only need to capture one microbatch. if self.pp_group.size() == 1 and not self.config.overlap_moe_expert_parallel_comm: assert ( self.num_model_chunks == 1 ), "If PP is not enabled, there should be only one model chunk." self.num_microbatches = 1 + elif self._should_use_dynamic_microbatch_slots(): + probe_num_microbatches = self._get_probe_num_microbatches_for_dynamic_slots() + from megatron.core.pipeline_parallel.schedules import ( + get_pp_rank_microbatches as _probe_get_pp, + ) + from megatron.core.pipeline_parallel.schedules import ( + get_schedule_table as _probe_get_st, + ) + + _, _, _probe_warmup, _ = _probe_get_pp( + probe_num_microbatches, + self.num_model_chunks, + microbatch_group_size_per_vp_stage, + False, + overlap_moe_expert_parallel_comm=self.config.overlap_moe_expert_parallel_comm, + ) + _probe_st = _probe_get_st( + probe_num_microbatches, self.num_model_chunks, microbatch_group_size_per_vp_stage + ) + _probe_order = convert_schedule_table_to_order( + _probe_warmup, self.num_model_chunks, _probe_st + ) + auto_num_slots = self._get_required_num_microbatch_slots_from_order( + _probe_order, self.num_model_chunks + ) + pp_group = parallel_state.get_pipeline_model_parallel_group() + if pp_group is not None and pp_group.size() > 1: + auto_num_slots_tensor = torch.tensor( + [auto_num_slots], dtype=torch.int32, device=torch.cuda.current_device() + ) + torch.distributed.all_reduce( + auto_num_slots_tensor, op=torch.distributed.ReduceOp.MAX, group=pp_group + ) + auto_num_slots = int(auto_num_slots_tensor.item()) + runtime_num_microbatches = get_num_microbatches() + max_num_microbatches, capture_mode = self._get_thd_varlen_max_num_microbatches( + runtime_num_microbatches, microbatch_group_size_per_vp_stage + ) + if self.config.overlap_moe_expert_parallel_comm or self.config.delay_wgrad_compute: + self.num_microbatches = runtime_num_microbatches + capture_mode = "runtime" + fallback_reason = "overlap_moe_expert_parallel_comm/delay_wgrad_compute" + else: + # auto_num_slots is a topology-only theoretical lower bound for PP/VPP graph + # slot liveness. THD varlen packing can produce different real microbatch + # counts across iterations, so capture uses the THD/GBS-derived upper + # bound instead of the reduced slot count for safety. Currently TE cuda + # graph backend may crash if use the auto_num_slots. + self.num_microbatches = max(runtime_num_microbatches, max_num_microbatches) + fallback_reason = None + log_on_each_pipeline_stage( + logger=logger, + tp_group=None, + dp_cp_group=None, + level=logging.INFO, + msg=f'Rank {torch.distributed.get_rank()}: dynamic CUDA graph slots ' + f'enabled. runtime_num_microbatches={runtime_num_microbatches}, ' + f'auto_num_slots={auto_num_slots}, ' + f'max_num_microbatches={max_num_microbatches}, ' + f'capture_num_microbatches={self.num_microbatches}, ' + f'capture_mode={capture_mode}' + + (f', fallback_reason={fallback_reason}' if fallback_reason else ''), + ) else: self.num_microbatches = get_num_microbatches() _, _, num_warmup_microbatches, _ = get_pp_rank_microbatches( self.num_microbatches, self.num_model_chunks, - self.config.microbatch_group_size_per_vp_stage, + microbatch_group_size_per_vp_stage, forward_only=False, p2p_communicator=self.p2p_communicator, ) schedule_table = get_schedule_table( - self.num_microbatches, - self.num_model_chunks, - self.config.microbatch_group_size_per_vp_stage, + self.num_microbatches, self.num_model_chunks, microbatch_group_size_per_vp_stage ) order = convert_schedule_table_to_order( num_warmup_microbatches, self.num_model_chunks, schedule_table diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index fccf674d785..2472d0ec2fb 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -442,8 +442,11 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv else: cu_seqlens_q = cu_seqlens_kv = None + rope_max_seqlen_q = rope_max_seqlen_kv = None # ========================================= # Q down projection @@ -631,6 +634,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_max_seqlen_q, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -641,6 +645,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_max_seqlen_kv, ) # query: [num_tokens, n, (kv_lora_rank + qk_pos_emb_head_dim)] diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index 6a55bfbc348..90250c8d58f 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -224,6 +224,13 @@ def _te_cuda_graph_backward_dw_graph(self, microbatch_idx): return self.cuda_graphs[cg_index].backward_dw() + def _is_thd_cuda_graph(self): + """Check if THD format with CUDA Graph is being used.""" + return ( + getattr(self.config, 'sequence_packing_scheduler', None) is not None + and self.config.cuda_graph_impl != "none" + ) + def get_layer_static_inputs(self, seq_length, micro_batch_size): """ Get the static inputs for the layer. @@ -231,26 +238,47 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): from the seq_length, micro_batch_size, and parallel config. Override this method if the module has other inputs. + For THD + CUDA Graph, hidden_states uses the padded max sequence length with + micro_batch_size=1 (packed sequence format). + Returns: Dict[str, torch.Tensor]: A dictionary containing the static inputs for the layer. """ # Calculate data shape related values. context_parallel_size = self.config.context_parallel_size - slen_per_cp = seq_length // context_parallel_size sequence_parallel = self.config.sequence_parallel tensor_model_parallel_size = self.config.tensor_model_parallel_size - slen_per_cptp = ( - slen_per_cp // tensor_model_parallel_size if sequence_parallel else slen_per_cp - ) - static_inputs = {} - static_inputs["hidden_states"] = torch.ones( - (slen_per_cptp, micro_batch_size, self.config.hidden_size), - dtype=torch.bfloat16, - requires_grad=True, - device=torch.cuda.current_device(), - ) - return static_inputs + if self._is_thd_cuda_graph(): + # THD + CUDA Graph: pre-padded packed-sequence buffer, batch dim = 1. + assert ( + self.config.max_seqlen_per_dp_cp_rank is not None + ), "max_seqlen_per_dp_cp_rank must be set when using THD format with CUDA Graph." + slen_full = self.config.max_seqlen_per_dp_cp_rank + batch = 1 + else: + # SBHD path: per-rank seq is split by CP and (optionally) by TP under SP. + slen_full = seq_length // context_parallel_size + batch = micro_batch_size + slen_per_cptp = slen_full // tensor_model_parallel_size if sequence_parallel else slen_full + + # Static input dtype must match the runtime activation dtype that flows + # through the captured graph. + if self.config.bf16: + dtype = torch.bfloat16 + elif self.config.fp16: + dtype = torch.float16 + else: + dtype = torch.float32 + + return { + "hidden_states": torch.ones( + (slen_per_cptp, batch, self.config.hidden_size), + dtype=dtype, + requires_grad=True, + device=torch.cuda.current_device(), + ) + } def setup_manual_hooks(self, make_hook_func): """ diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index deebd3472ea..50d12b0ea62 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -11,6 +11,7 @@ from megatron.core import tensor_parallel, utils from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.inference.utils import InferenceMode +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.moe.moe_utils import ( @@ -435,13 +436,20 @@ def setup_delayed_wgrad_for_dispatch_backward_overlap(self): self._delayed_wgrad_stream = torch.cuda.Stream(device="cuda") @maybe_skip_or_early_return_by_cudagraph("route") - def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def route( + self, + hidden_states: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """Compute token routing for preprocessing. This method uses the router to determine which experts to send each token to, producing routing probabilities and a mapping. """ - probs, routing_map = apply_module(self.router)(hidden_states, padding_mask) + probs, routing_map = apply_module(self.router)( + hidden_states, padding_mask, packed_seq_params + ) return probs, routing_map @maybe_skip_or_early_return_by_cudagraph("preprocess") @@ -600,6 +608,7 @@ def forward( hidden_states: torch.Tensor, intermediate_tensors=None, padding_mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """Forward pass for the MoE layer. @@ -611,9 +620,9 @@ def forward( Args: hidden_states (torch.Tensor): The input tensor shape [seq_length, bsz, hidden_size]. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. Returns: A tuple containing the output tensor and the MLP bias, if any. """ @@ -634,6 +643,29 @@ def forward( else: self.token_dispatcher = self._training_token_dispatcher self.shared_expert_overlap = self.config.moe_shared_expert_overlap + + # Align padding_mask to hidden_states sequence dimension before transpose. + # padding_mask arrives as [bsz, seq_length] but may need SP scatter when + # hidden_states is already TP-scattered (seq_length / TP). + if padding_mask is not None and padding_mask.shape[1] != hidden_states.shape[0]: + if ( + self.config.sequence_parallel + and padding_mask.shape[1] % self.config.tensor_model_parallel_size == 0 + and padding_mask.shape[1] // self.config.tensor_model_parallel_size + == hidden_states.shape[0] + ): + padding_mask = ( + tensor_parallel.scatter_to_sequence_parallel_region( + padding_mask.transpose(0, 1).contiguous() + ) + .transpose(0, 1) + .contiguous() + ) + else: + raise AssertionError( + f"padding_mask shape {padding_mask.shape} cannot be aligned to " + f"hidden_states sequence length {hidden_states.shape[0]}" + ) # Transpose from [bsz, seq_length] to [seq_length, bsz] to align with hidden_states if padding_mask is not None: padding_mask = padding_mask.transpose(0, 1).bool() @@ -643,7 +675,9 @@ def custom_forward(hidden_states, intermediate_tensors=None, padding_mask=None): try: if "route" in self.fwd_execution_map: shared_expert_output = self.shared_experts_compute(hidden_states) - probs, routing_map = self.route(hidden_states, padding_mask) + probs, routing_map = self.route( + hidden_states, padding_mask, packed_seq_params + ) hidden_states, probs = self.preprocess(hidden_states, probs, routing_map) if intermediate_tensors is not None: diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 79b45156ed8..f6b0c2d886d 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -3,7 +3,7 @@ import functools import math from dataclasses import dataclass -from typing import List, Optional, Tuple, Union +from typing import List, Optional, Sequence, Tuple, Union import torch @@ -56,7 +56,7 @@ def switch_load_balancing_loss_func( probs: torch.Tensor, tokens_per_expert: torch.Tensor, - total_num_tokens: int, + total_num_tokens: Union[int, torch.Tensor], topk: int, num_experts: int, moe_aux_loss_coeff: float, @@ -106,7 +106,7 @@ def switch_load_balancing_loss_func( Shape in [num_tokens, num_experts]. tokens_per_expert (torch.Tensor): Number of tokens assigned to each expert in the batch. Shape in [num_experts] - total_num_tokens (int): Total number of tokens in the batch. + total_num_tokens (int or torch.Tensor): Total number of tokens in the batch. topk (int): The number of experts selected for each token. num_experts (int): The number of experts. moe_aux_loss_coeff (float): The coefficient for the auxiliary loss. @@ -224,22 +224,31 @@ def get_capacity( def get_tokens_per_expert_and_token_count( routing_map: torch.Tensor, reduce_group: torch.distributed.ProcessGroup, + reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, topk: int = None, with_padding_mask: bool = False, -) -> torch.Tensor: +) -> Tuple[torch.Tensor, Union[int, torch.Tensor], Union[int, torch.Tensor]]: """ Compute global_tokens_per_expert, local_num_tokens and total_num_tokens with padding mask. """ local_tokens_per_expert = routing_map.sum(dim=0) - global_tokens_per_expert = reduce_from_tensor_model_parallel_region( - local_tokens_per_expert, reduce_group - ) + if reduce_groups is None: + reduce_groups = (reduce_group,) + + global_tokens_per_expert = local_tokens_per_expert + reduce_world_size = 1 + for group in reduce_groups: + global_tokens_per_expert = reduce_from_tensor_model_parallel_region( + global_tokens_per_expert, group + ) + reduce_world_size *= group.size() + if with_padding_mask: - local_num_tokens = local_tokens_per_expert.sum() / topk - total_num_tokens = global_tokens_per_expert.sum() / topk + local_num_tokens = local_tokens_per_expert.sum() // topk + total_num_tokens = global_tokens_per_expert.sum() // topk else: local_num_tokens = routing_map.shape[0] - total_num_tokens = local_num_tokens * reduce_group.size() + total_num_tokens = local_num_tokens * reduce_world_size return global_tokens_per_expert, local_num_tokens, total_num_tokens diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 7414c8a7ab0..0b44c84f623 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -1,12 +1,14 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from abc import ABC, abstractmethod -from typing import Optional, Union +from dataclasses import dataclass +from typing import Optional, Sequence, Union import torch from megatron.core.inference.utils import InferenceMode from megatron.core.jit import jit_fuser +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker from megatron.core.transformer.moe.moe_utils import ( @@ -21,12 +23,26 @@ sinkhorn, switch_load_balancing_loss_func, topk_routing_with_score_function, - z_loss_func, ) from megatron.core.transformer.moe.router_replay import RouterReplay from megatron.core.transformer.transformer_config import TransformerConfig +@dataclass(frozen=True) +class _AuxLossGroupConfig: + """Process groups for local aux/seq-aux loss and its metric logging.""" + + loss_reduce_groups: Sequence[torch.distributed.ProcessGroup] + metric_reduce_group: Optional[torch.distributed.ProcessGroup] + metric_avg_group: Optional[torch.distributed.ProcessGroup] + metric_needs_dp_avg: bool + + @property + def metric_pre_reduce_groups(self) -> Optional[Sequence[torch.distributed.ProcessGroup]]: + """Groups to reduce eagerly before recording metrics, if tracker reduction is unsafe.""" + return self.loss_reduce_groups if self.metric_avg_group is not None else None + + class Router(ABC, MegatronModule): """Base Router class""" @@ -292,16 +308,19 @@ def _apply_aux_loss( scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor, with_padding_mask: bool = False, + packed_seq_params: Optional[PackedSeqParams] = None, ): """Apply the auxiliary loss for the given scores and routing map.""" aux_loss_coeff = self.get_aux_loss_coeff("aux_loss") if aux_loss_coeff == 0: return probs + aux_loss_groups = self._get_aux_loss_groups(packed_seq_params) global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, - reduce_group=self.tp_cp_group, + reduce_group=aux_loss_groups.loss_reduce_groups[0], + reduce_groups=aux_loss_groups.loss_reduce_groups, topk=self.topk, with_padding_mask=with_padding_mask, ) @@ -321,8 +340,13 @@ def _apply_aux_loss( aux_loss_coeff, aux_loss, "load_balancing_loss", - self.tp_cp_group, + aux_loss_groups.metric_reduce_group, + avg_group=aux_loss_groups.metric_avg_group, + needs_dp_avg=aux_loss_groups.metric_needs_dp_avg, valid_token_count=local_num_tokens, + aux_loss_logging_reduce_groups=aux_loss_groups.metric_pre_reduce_groups, + aux_loss_scale_reduce_groups=aux_loss_groups.loss_reduce_groups, + aux_loss_scale_num_tokens=total_num_tokens, ) return probs @@ -334,6 +358,7 @@ def _apply_seq_aux_loss( seq_length: int, bsz: int, with_padding_mask: bool = False, + packed_seq_params: Optional[PackedSeqParams] = None, ): """Apply the sequence-level auxiliary loss for the given scores and routing map. @@ -349,10 +374,12 @@ def _apply_seq_aux_loss( scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1) routing_map = routing_map.reshape(seq_length, -1) + aux_loss_groups = self._get_aux_loss_groups(packed_seq_params) global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, - reduce_group=self.tp_cp_group, + reduce_group=aux_loss_groups.loss_reduce_groups[0], + reduce_groups=aux_loss_groups.loss_reduce_groups, with_padding_mask=with_padding_mask, topk=self.topk * bsz, ) @@ -376,8 +403,13 @@ def _apply_seq_aux_loss( seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", - self.tp_cp_group, + aux_loss_groups.metric_reduce_group, + avg_group=aux_loss_groups.metric_avg_group, + needs_dp_avg=aux_loss_groups.metric_needs_dp_avg, valid_token_count=local_num_tokens, + aux_loss_logging_reduce_groups=aux_loss_groups.metric_pre_reduce_groups, + aux_loss_scale_reduce_groups=aux_loss_groups.loss_reduce_groups, + aux_loss_scale_num_tokens=total_num_tokens, ) return probs @@ -393,7 +425,8 @@ def _apply_global_aux_loss( if global_aux_loss_coeff == 0: return probs - # Use unified function to compute tokens_per_expert and num_tokens + # Global aux loss intentionally uses the full static TP x DP x CP domain. + # Dynamic CP subgroups only affect local aux/seq-aux domains. global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, @@ -424,18 +457,51 @@ def _apply_global_aux_loss( self.tp_dp_cp_group, needs_dp_avg=False, valid_token_count=local_num_tokens, + # The global aux-loss statistics/logging domain is TP x DP x CP, but + # per-token-loss gradient normalization already reduces the denominator + # across DP x CP in finalize_model_grads. Scale the aux-loss numerator + # over TP x CP only, matching the original static behavior while still + # using an exact valid-token count when padding is present. + aux_loss_scale_reduce_groups=(self.tp_cp_group,), ) return probs + def _get_aux_loss_groups( + self, packed_seq_params: Optional[PackedSeqParams] = None + ) -> _AuxLossGroupConfig: + """Return process groups for MoE aux-loss statistics and logging.""" + if ( + packed_seq_params is not None + and packed_seq_params.local_cp_size is not None + and packed_seq_params.cp_group is not None + ): + return _AuxLossGroupConfig( + loss_reduce_groups=(packed_seq_params.cp_group, self.tp_group), + metric_reduce_group=None, + metric_avg_group=self.tp_dp_cp_group, + metric_needs_dp_avg=False, + ) + + return _AuxLossGroupConfig( + loss_reduce_groups=(self.tp_cp_group,), + metric_reduce_group=self.tp_cp_group, + metric_avg_group=None, + metric_needs_dp_avg=True, + ) + def attach_and_log_load_balancing_loss( self, activation: torch.Tensor, aux_loss_coeff: float, aux_loss: torch.Tensor, aux_loss_name: str, - reduce_group: torch.distributed.ProcessGroup, + reduce_group: Optional[torch.distributed.ProcessGroup], + avg_group: Optional[torch.distributed.ProcessGroup] = None, needs_dp_avg: bool = True, valid_token_count: Optional[Union[int, torch.Tensor]] = None, + aux_loss_logging_reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, + aux_loss_scale_reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, + aux_loss_scale_num_tokens: Optional[Union[int, torch.Tensor]] = None, ): """Attach aux loss function to activation and add to logging. @@ -444,7 +510,10 @@ def attach_and_log_load_balancing_loss( aux_loss_coeff (float): Coefficient for the aux loss. aux_loss (torch.Tensor): Computed aux loss. aux_loss_name (str): Name of the aux loss for logging. - reduce_group (torch.distributed.ProcessGroup): Process group for reduction. + reduce_group (torch.distributed.ProcessGroup, optional): Process group for deferred + logging reduction. + avg_group (torch.distributed.ProcessGroup, optional): Process group for deferred + logging average. needs_dp_avg (bool): Whether to average this metric across DP ranks after reduce_group. valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor). @@ -473,44 +542,47 @@ def attach_and_log_load_balancing_loss( else: layer_number = self.layer_number + metric_value = aux_loss / aux_loss_coeff + if aux_loss_logging_reduce_groups is not None: + metric_value = metric_value.detach().clone() + for group in aux_loss_logging_reduce_groups: + torch.distributed.all_reduce(metric_value, group=group) + get_moe_metrics_tracker().record( aux_loss_name, - aux_loss / aux_loss_coeff, + metric_value, layer_number, num_layers, reduce_group=reduce_group, + avg_group=avg_group, needs_dp_avg=needs_dp_avg, ) if self.calculate_per_token_loss: - # Target final scaling on aux_loss gradients: 1 / (num_micro_batches * dp_size), - # matching the !calculate_per_token_loss path. - # - # --calculate-per-token-loss already divides every parameter gradient by - # total_global_tokens (the global non-padded token count summed in - # finalize_model_grads). The router's `num_local_tokens` (= activation.shape[0]) - # is sequence-parallel sharded — the router weight is marked - # `sequence_parallel=True` in Router.reset_parameters (see - # `setattr(self.weight, 'sequence_parallel', ...)` above), so each TP rank - # computes a partial gradient on the router weight from its local sequence - # shard, and `_allreduce_non_tensor_model_parallel_grads` SUMS those partial - # gradients across the TP group. Re-expressing total_global_tokens in terms of the - # router's `num_local_tokens`: - # total_global_tokens - # = num_micro_batches * dp_cp_size * loss_func_local_tokens - # = num_micro_batches * dp_cp_size * tp_size * num_local_tokens - # = num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()) - # (using loss_func_local_tokens = tp_size * num_local_tokens, then regrouping - # dp_cp_size * tp_size as dp_size * tp_cp_group.size()). - # - # So pre-multiplying aux_loss by num_local_tokens * tp_cp_group.size() cancels - # that same factor in total_global_tokens above, leaving 1 / (num_micro_batches * - # dp_size) as the effective scaling on the aux_loss gradient — the target. - # Use valid_token_count (excluding padding) if provided, otherwise use total tokens. - num_local_tokens = ( - valid_token_count if valid_token_count is not None else activation.shape[0] - ) + # --calculate-per-token-loss divides all parameter gradients by the global + # non-padded token count in finalize_model_grads. Pre-multiplying by the + # valid-token count from this aux-loss domain makes the final objective a + # token-weighted average of per-domain aux losses. Use the reduced count + # directly: with THD padding or dynamic CP, valid token counts can differ + # by rank/group, so local_num_tokens * group_size is not generally correct. + if aux_loss_scale_num_tokens is None: + num_local_tokens = ( + valid_token_count if valid_token_count is not None else activation.shape[0] + ) + if torch.is_tensor(num_local_tokens): + aux_loss_scale_num_tokens = num_local_tokens.clone().to( + device=activation.device + ) + else: + aux_loss_scale_num_tokens = torch.tensor( + num_local_tokens, device=activation.device + ) + if aux_loss_scale_reduce_groups is None: + assert reduce_group is not None, "reduce_group is required for aux-loss scaling" + aux_loss_scale_reduce_groups = (reduce_group,) + for group in aux_loss_scale_reduce_groups: + torch.distributed.all_reduce(aux_loss_scale_num_tokens, group=group) activation = MoEAuxLossAutoScaler.apply( - activation, aux_loss * num_local_tokens * self.tp_cp_group.size() + activation, aux_loss * aux_loss_scale_num_tokens ) else: activation = MoEAuxLossAutoScaler.apply(activation, aux_loss) @@ -522,48 +594,45 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None): Args: logits (torch.Tensor): The logits of the router. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape in [num_tokens]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [num_tokens]. True = padding, + False = valid. Defaults to None. Returns: torch.Tensor: The logits after applying the z-loss. """ if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled(): # Skip Z loss calculations when using torch.no_grad() or checkpointing. - moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() - z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask) - if self.calculate_per_token_loss: - # Same derivation as in attach_and_log_load_balancing_loss: - # - Target final scaling on z_loss gradients: 1 / (num_micro_batches * dp_size). - # - In terms of the router's `num_local_tokens`, the total_global_tokens - # divisor that finalize_model_grads applies factors as - # num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()). - # - Pre-multiplying z_loss by num_local_tokens * tp_cp_group.size() cancels - # that same factor in total_global_tokens, leaving - # 1 / (num_micro_batches * dp_size) as the effective scaling — the target. - # The /tp_cp_group.size() on moe_z_loss_coeff above is a separate forward-side - # correction: z_loss is computed independently on each TP+CP rank's local - # logits and must be averaged across TP+CP rather than summed. - # Count valid tokens: sum of inverted mask (False -> True = valid) - num_local_tokens = ( - (~padding_mask).sum() if padding_mask is not None else logits.shape[0] - ) - logits = MoEAuxLossAutoScaler.apply( - logits, z_loss * num_local_tokens * self.tp_cp_group.size() - ) + logsum = torch.logsumexp(logits, dim=-1) + z_loss_values = torch.square(logsum) + if padding_mask is not None: + valid_mask = ~padding_mask + z_loss_values = z_loss_values * valid_mask + num_local_tokens = valid_mask.sum() else: - logits = MoEAuxLossAutoScaler.apply(logits, z_loss) + num_local_tokens = torch.tensor(logits.shape[0], device=logits.device) + + z_loss_sum = z_loss_values.sum() + z_loss_mean = z_loss_sum / torch.clamp(num_local_tokens, min=1) - # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times. - # To avoid accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers - # so the total loss is correct. + mtp_loss_scale = 1 if ( self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None ): - z_loss = z_loss / self.config.mtp_num_layers + mtp_loss_scale = self.config.mtp_num_layers + + if self.calculate_per_token_loss: + # --calculate-per-token-loss divides gradients by the global non-padded + # token count. Attach the local z-loss numerator directly so the final + # objective is a token-weighted z-loss over valid tokens. + z_loss = z_loss_sum * self.config.moe_z_loss_coeff / mtp_loss_scale + logits = MoEAuxLossAutoScaler.apply(logits, z_loss) + else: + moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() + z_loss = z_loss_mean * moe_z_loss_coeff / mtp_loss_scale + logits = MoEAuxLossAutoScaler.apply(logits, z_loss) num_layers = self.config.num_layers if self.config.mtp_num_layers is not None: @@ -575,7 +644,7 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None): layer_number = self.layer_number get_moe_metrics_tracker().record( - "z_loss", z_loss / moe_z_loss_coeff, layer_number, num_layers + "z_loss", z_loss_mean / mtp_loss_scale, layer_number, num_layers ) return logits @@ -614,14 +683,19 @@ def _apply_expert_bias( routing_map = routing_map & (~padding_mask) self.local_tokens_per_expert += routing_map.sum(dim=0) - def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def routing( + self, + logits: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """Top-k routing function Args: logits (torch.Tensor): Logits tensor after gating. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. Returns: probs (torch.Tensor): The probabilities of token to experts assignment. @@ -681,6 +755,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N scores_for_aux_loss, routing_map_for_aux_loss, with_padding_mask=padding_mask is not None, + packed_seq_params=packed_seq_params, ) probs = self._apply_seq_aux_loss( probs, @@ -689,6 +764,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N seq_length, bsz, with_padding_mask=padding_mask is not None, + packed_seq_params=packed_seq_params, ) probs = self._apply_global_aux_loss( probs, @@ -708,15 +784,20 @@ def reset_global_aux_loss_tracker(self): self.global_tokens_per_expert.zero_() self.ga_steps.zero_() - def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """ Forward pass of the router. Args: input (torch.Tensor): Input tensor. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. """ self._maintain_float32_expert_bias() @@ -734,7 +815,9 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No logits, self.config.moe_router_force_biased, self.layer_number ) - probs, routing_map = self.routing(logits, padding_mask=padding_mask) + probs, routing_map = self.routing( + logits, padding_mask=padding_mask, packed_seq_params=packed_seq_params + ) return probs, routing_map @@ -832,7 +915,12 @@ def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = N ) return probs.squeeze(1), top_indices.squeeze(1) - def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """Simplified forward pass for inference - returns dense tensors only. Args: @@ -846,6 +934,6 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No """ if not InferenceMode.is_active(): - return super().forward(input, padding_mask) + return super().forward(input, padding_mask, packed_seq_params) return self._forward(input, padding_mask) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 202034986db..2e19fefada6 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -716,8 +716,11 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv else: cu_seqlens_q = cu_seqlens_kv = None + rope_max_seqlen_q = rope_max_seqlen_kv = None # ========================================= # QKV down projection and layernorm @@ -929,6 +932,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_max_seqlen_q, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -939,6 +943,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_max_seqlen_kv, ) # query: [num_tokens, n, (qk_head_dim + v_head_dim)] diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index f8779d674e2..6c1ec006ab2 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1064,6 +1064,26 @@ class TransformerConfig(ModelParallelConfig): CudaGraphScope instances deserialized from pre-refactor checkpoints are converted to their string names before normalization so existing CUDA_GRAPH_MODULES_DEPRECATIONS handles them.""" + thd_max_packed_sequences: int = field( + default=32, metadata={"argparse_meta": {"arg_names": ["--thd-max-packed-sequences"]}} + ) + """Maximum number of THD packed sequences per microbatch, including any dummy + sequence appended for a padding tail. The dp_balanced packing scheduler reserves + that dummy slot when THD padding appends one. When CUDA Graph is enabled, cu_seqlens + tensors are padded to this size + 1. + + Sizing guidance: choose a value that comfortably covers the worst-case packing, + roughly ceil(max_seqlen_per_dp_cp_rank * cp_size / min_seq_len_after_filter). + Setting it too small results in more microbatches with smaller packs (wasted + token budget); setting it too large just allocates a slightly larger cu_seqlens + buffer.""" + + cuda_graph_dynamic_microbatches: bool = False + """Allow CUDA graph replay when runtime microbatch count varies across iterations. + This option is only meaningful for cuda_graph_impl=transformer_engine. For THD sequence + packing, capture uses a conservative upper bound on the packed microbatch count so graph + replay can cover iterations whose real packed microbatch count changes.""" + #################### # miscellaneous #################### @@ -1320,6 +1340,12 @@ def __post_init__(self): self.linear_attention_freq is not None ), f"linear_attention_freq must be set for linear gated_delta_net." + if self.pad_packed_seq_alignment is not None: + assert self.pad_packed_seq_by_appending_dummy_seq, ( + "gated_delta_net with pad_packed_seq_alignment requires " + "pad_packed_seq_by_appending_dummy_seq." + ) + # Check required parameters assert ( self.linear_conv_kernel_dim is not None @@ -2445,6 +2471,12 @@ def _scope_to_str(s): if CudaGraphModule.moe_preprocess not in self.cuda_graph_modules: self.cuda_graph_modules.append(CudaGraphModule.moe_preprocess) + if self.cuda_graph_impl != "transformer_engine": + assert not self.cuda_graph_dynamic_microbatches, ( + "cuda_graph_dynamic_microbatches is only supported with " + "cuda_graph_impl=transformer_engine." + ) + assert ( CudaGraphModule.moe not in self.cuda_graph_modules or CudaGraphModule.moe_router not in self.cuda_graph_modules @@ -2812,6 +2844,55 @@ def _scope_to_str(s): self.attention_backend == AttnBackend.flash ), "Batch invariant mode only supports FlashAttention" + if self.cuda_graph_impl != "none" and ( + self.sequence_packing_scheduler is not None or self.hybrid_context_parallel + ): + assert ( + self.pad_packed_seq_alignment is not None + ), "THD CUDA Graph requires --pad-packed-seq-alignment to be set." + assert ( + self.pad_packed_seq_alignment == "max" + or self.pad_packed_seq_alignment == self.max_seqlen_per_dp_cp_rank + ), ( + "THD CUDA Graph requires --pad-packed-seq-alignment='max' " + "or --pad-packed-seq-alignment equal to max_seqlen_per_dp_cp_rank " + f"({self.max_seqlen_per_dp_cp_rank}), got {self.pad_packed_seq_alignment}." + ) + + 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'] + 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/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index f6ea382077e..2dfbf8cad37 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -865,6 +865,10 @@ def _forward_mlp( InferenceMode.is_active() and self.config.inference_fuse_tp_communication ) + moe_kwargs = {} + if self.is_moe_layer and packed_seq_params is not None: + moe_kwargs["packed_seq_params"] = packed_seq_params + if self.recompute_mlp: if self.config.fp8 or self.config.fp4: # import here to avoid circular import @@ -877,10 +881,13 @@ def _forward_mlp( self.pg_collection.tp, pre_mlp_layernorm_output, padding_mask=padding_mask, + **moe_kwargs, ) else: mlp_output_with_bias = tensor_parallel.checkpoint( - functools.partial(apply_module(self.mlp), padding_mask=padding_mask), + functools.partial( + apply_module(self.mlp), padding_mask=padding_mask, **moe_kwargs + ), False, pre_mlp_layernorm_output, ) @@ -1089,19 +1096,50 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): Get the static inputs for the transformer layer. Besides the hidden_states that is generated in GraphableMegatronModule, we also add the attention_mask. + For THD + CUDA Graph: generates cu_seqlens and padding_mask static tensors + instead of attention_mask. + Returns: Dict[str, torch.Tensor]: A dictionary containing the static inputs for the layer. """ static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) + device = torch.cuda.current_device() - if not isinstance(self.self_attention, IdentityOp) and ( + # Captured forward needs attention-side static input only when this + # layer's attention is inside the captured scope. + attn_in_graph = not isinstance(self.self_attention, IdentityOp) and ( not self.config.cuda_graph_modules or CudaGraphModule.attn in self.config.cuda_graph_modules - ): + ) + + if self._is_thd_cuda_graph(): + if attn_in_graph: + # Static cu_seqlens shaped [thd_max_packed_sequences + 1]. We seed it as + # one full-length sequence (covers the worst case at capture): + # cu_seqlens = [0, max_T, max_T, ..., max_T] + # which represents a single packed sequence followed by zero-length + # entries. cu_seqlens_q / kv / *_padded all share this layout. + max_T = self.config.max_seqlen_per_dp_cp_rank * self.config.context_parallel_size + max_num_seqs = self.config.thd_max_packed_sequences + cu_seqlens = torch.zeros(max_num_seqs + 1, dtype=torch.int32, device=device) + cu_seqlens[1:] = max_T + + static_inputs["cu_seqlens_q"] = cu_seqlens + static_inputs["cu_seqlens_kv"] = cu_seqlens.clone() + static_inputs["cu_seqlens_q_padded"] = cu_seqlens.clone() + static_inputs["cu_seqlens_kv_padded"] = cu_seqlens.clone() + + slen_for_mask = self.config.max_seqlen_per_dp_cp_rank + if self.config.sequence_parallel: + slen_for_mask //= self.config.tensor_model_parallel_size + static_inputs["padding_mask"] = torch.zeros( + 1, slen_for_mask, dtype=torch.bool, device=device + ) + elif attn_in_graph: slen_per_cp = seq_length // self.config.context_parallel_size static_inputs["attention_mask"] = ( ~(torch.tril(torch.ones((slen_per_cp, seq_length))).bool()) - .to(torch.cuda.current_device()) + .to(device) .reshape(1, 1, slen_per_cp, seq_length) .tile(micro_batch_size, 1, 1, 1) ) @@ -1135,6 +1173,47 @@ def _get_submodules_under_cudagraphs(self): submodules += [self.mlp.shared_experts] return submodules + @staticmethod + def _decompose_packed_seq_params_to_kwargs(kwargs): + """Decompose PackedSeqParams into individual tensor kwargs for CUDA graph. + + CUDA graph requires all inputs to be tensors. This extracts the cu_seqlens + tensor fields from PackedSeqParams into individual kwargs. max_seqlen_q/kv + are omitted because they are static and reading them from a CUDA tensor is + forbidden during graph capture. They are restored from config in + _reconstruct_packed_seq_params_from_kwargs. + """ + packed_seq_params = kwargs.pop('packed_seq_params', None) + if packed_seq_params is None: + return + kwargs['cu_seqlens_q'] = packed_seq_params.cu_seqlens_q + kwargs['cu_seqlens_kv'] = packed_seq_params.cu_seqlens_kv + kwargs['cu_seqlens_q_padded'] = packed_seq_params.cu_seqlens_q_padded + kwargs['cu_seqlens_kv_padded'] = packed_seq_params.cu_seqlens_kv_padded + + def _reconstruct_packed_seq_params_from_kwargs(self, kwargs): + """Reconstruct PackedSeqParams from individual tensor kwargs (CUDA graph path). + + During CUDA graph capture/replay, PackedSeqParams fields are decomposed into + individual cu_seqlens tensor kwargs. This method reassembles them into a + PackedSeqParams. max_seqlen_q/kv are taken from config since they are always + the padded static value and cannot be read from CUDA tensors during graph capture. + """ + if 'cu_seqlens_q' not in kwargs: + return + max_seqlen = self.config.max_seqlen_per_dp_cp_rank * self.config.context_parallel_size + packed_seq_params = PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=kwargs.pop('cu_seqlens_q'), + cu_seqlens_kv=kwargs.pop('cu_seqlens_kv'), + cu_seqlens_q_padded=kwargs.pop('cu_seqlens_q_padded'), + cu_seqlens_kv_padded=kwargs.pop('cu_seqlens_kv_padded'), + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + pad_between_seqs=False, + ) + kwargs['packed_seq_params'] = packed_seq_params + def _te_cuda_graph_capture(self, *args, **kwargs): """ CUDA Graph capture for this layer using TE interface. @@ -1142,7 +1221,10 @@ def _te_cuda_graph_capture(self, *args, **kwargs): 1. In some conditions CUDA graph cannot cover the entire layer. The `cuda_graph_modules` attribute can be set to control the scope of the CUDA graph. 2. If context is None, it cannot be returned as output. + For THD format, PackedSeqParams is reconstructed from tensor kwargs. """ + self._reconstruct_packed_seq_params_from_kwargs(kwargs) + # Record the backward event on cuda graph stream in backward pass. # This is to ensure the main stream waits for computing on cuda graph stream to complete, # and overlaps with the H2D transfer on reload stream. @@ -1178,7 +1260,11 @@ def _te_cuda_graph_capture(self, *args, **kwargs): ) ) ): - hidden_states = self._forward_mlp(hidden_states) + hidden_states = self._forward_mlp( + hidden_states, + padding_mask=kwargs.get("padding_mask", None), + packed_seq_params=kwargs.get("packed_seq_params", None), + ) if not isinstance(hidden_states, list) and not isinstance(hidden_states, tuple): cuda_graph_outputs = [hidden_states] else: @@ -1198,8 +1284,10 @@ def _te_cuda_graph_replay(self, *args, **kwargs): interface. TransformerEngine versions>=1.10 allow keyword arguments with CUDA graph. However, CUDA graph accepts only Tensor inputs. Hence, `inference_context` and `packed_seq_params` are excluded from input list. + For THD format, PackedSeqParams is decomposed into individual tensor kwargs. """ context = None + padding_mask = kwargs.get("padding_mask", None) if ( self.config.cuda_graph_modules and CudaGraphModule.attn not in self.config.cuda_graph_modules @@ -1207,6 +1295,10 @@ def _te_cuda_graph_replay(self, *args, **kwargs): hidden_states, context = self._forward_attention(*args, **kwargs) args = (hidden_states,) kwargs = {} + if padding_mask is not None: + kwargs["padding_mask"] = padding_mask + else: + self._decompose_packed_seq_params_to_kwargs(kwargs) assert (kwargs.get('inference_context') is None) and ( kwargs.get('packed_seq_params') is None @@ -1334,7 +1426,18 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): return residual, hidden_states, probs, shared_expert_output # CUDA Graph does not capture the MLP/MoE part at all. - output = self._forward_mlp(*cuda_graph_output) + # The first CUDA Graph output is hidden_states for the uncaptured + # MLP path. Pass padding_mask as a keyword so it is not consumed as + # a positional output. + assert ( + len(cuda_graph_output) >= 1 + ), "expected at least hidden_states in cuda_graph_output" + hidden_states = cuda_graph_output[0] + output = self._forward_mlp( + hidden_states, + padding_mask=kwargs.get("padding_mask", None), + packed_seq_params=kwargs.get("packed_seq_params", None), + ) return output, context def _get_te_cuda_graph_replay_args(self, *args, **kwargs): @@ -1619,7 +1722,7 @@ def _restore_token_dispatcher_attrs(self): obj, name = self._resolve_token_dispatcher_attr(attr_name) setattr(obj, name, attr) - def _forward_mlp_router(self, hidden_states, padding_mask=None): + def _forward_mlp_router(self, hidden_states, padding_mask=None, packed_seq_params=None): """ Executes the router phase of the MoE block. @@ -1644,7 +1747,10 @@ def _forward_mlp_router(self, hidden_states, padding_mask=None): residual = residual.float() router_outputs = apply_module(self.mlp)( - pre_mlp_layernorm_output, intermediate_tensors=(), padding_mask=padding_mask + pre_mlp_layernorm_output, + intermediate_tensors=(), + padding_mask=padding_mask, + packed_seq_params=packed_seq_params, ) if is_graph_capturing() and not is_graph_warmup(): @@ -1717,10 +1823,10 @@ def _forward_mlp( ) def _forward_mlp_partial_cudagraphs( - hidden_states, inference_context=None, padding_mask=None + hidden_states, inference_context=None, padding_mask=None, packed_seq_params=None ): residual, hidden_states, probs, shared_expert_output = self._forward_mlp_router( - hidden_states, padding_mask=padding_mask + hidden_states, padding_mask=padding_mask, packed_seq_params=packed_seq_params ) # After the router graph replays, the captured .copy_() operations that update @@ -1751,17 +1857,24 @@ def _forward_mlp_partial_cudagraphs( parallel_state.get_tensor_model_parallel_group(), hidden_states, padding_mask=padding_mask, + packed_seq_params=packed_seq_params, ) else: result = tensor_parallel.checkpoint( functools.partial( - _forward_mlp_partial_cudagraphs, padding_mask=padding_mask + _forward_mlp_partial_cudagraphs, + padding_mask=padding_mask, + packed_seq_params=packed_seq_params, ), False, hidden_states, ) else: - result = _forward_mlp_partial_cudagraphs(hidden_states, padding_mask=padding_mask) + result = _forward_mlp_partial_cudagraphs( + hidden_states, + padding_mask=padding_mask, + packed_seq_params=packed_seq_params, + ) result = self._maybe_reflatten_from_moe(result, packed_seq_params, moe_unflatten_mbs) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index ae8496845c8..34954f19b9f 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -6,15 +6,16 @@ import dataclasses import json import os -from pathlib import Path import re import types +from pathlib import Path import torch +from megatron.core.model_parallel_config import _parse_pad_packed_seq_alignment +from megatron.core.msc_utils import MultiStorageClientFeature from megatron.core.rerun_state_machine import RerunStateMachine from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.transformer.cuda_graph_config import ( ALLOWED_INFERENCE_SCOPES, get_deprecated_cuda_graph_modules_migration, @@ -23,23 +24,24 @@ validate_deprecated_cuda_graph_modules_migration_inputs, ) from megatron.core.transformer.enums import AttnBackend, CudaGraphModule, InferenceCudaGraphScope +from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.utils import ( get_torch_version, is_flashinfer_min_version, is_te_min_version, is_torch_min_version, ) +from megatron.training.argument_utils import ( # noqa: F401 # pylint: disable=unused-import + ArgumentGroupFactory, + core_transformer_config_from_args, +) from megatron.training.global_vars import set_global_variables from megatron.training.utils import ( get_device_arch_version, - update_use_dist_ckpt, print_rank_0, + update_use_dist_ckpt, warn_rank_0, ) -from megatron.core.msc_utils import MultiStorageClientFeature - -from megatron.training.argument_utils import ArgumentGroupFactory, core_transformer_config_from_args # noqa: F401 # pylint: disable=unused-import - def add_megatron_arguments(parser: argparse.ArgumentParser): @@ -79,6 +81,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): parser = _add_msc_args(parser) parser = _add_kitchen_quantization_arguments(parser) parser = _add_sft_args(parser) + parser = _add_varlen_dataset_args(parser) parser = _add_fault_injector_args(parser) @@ -398,8 +401,9 @@ def validate_args(args, defaults={}): 'Currently only global and local checkpoints are supported' if args.non_persistent_ckpt_type == 'local': try: - from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import \ - LocalCheckpointManager + from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import ( + LocalCheckpointManager, + ) except ModuleNotFoundError as e: raise RuntimeError('nvidia_resiliency_ext is required for local checkpointing') from e @@ -719,8 +723,10 @@ def validate_args(args, defaults={}): ) from megatron.core.models.hybrid.hybrid_layer_allocation import ( - Symbols, parse_hybrid_pattern, get_hybrid_total_layer_count, + Symbols, + get_hybrid_total_layer_count, get_hybrid_total_pipeline_segment_count, + parse_hybrid_pattern, ) sep = Symbols.MTP_SEPARATOR @@ -1182,13 +1188,6 @@ def validate_args(args, defaults={}): if args.rl_use_sequence_packing: args.consumed_train_bins = 0 - # Support for variable sequence lengths across batches/microbatches. - # set it if the dataloader supports generation of variable sequence lengths - # across batches/microbatches. Due to additional communication overhead - # during pipeline parallelism, it should not be set if sequence length - # is constant during training. - args.variable_seq_lengths = False - # Iteration-based training. # Skip these checks when skip_train is set: LR config is irrelevant. if args.train_iters and not args.skip_train: @@ -1366,6 +1365,62 @@ def validate_args(args, defaults={}): 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' + # Support for variable sequence lengths across batches/microbatches. + # set it if the dataloader supports generation of variable sequence lengths + # across batches/microbatches. Due to additional communication overhead + # during pipeline parallelism, it should not be set if sequence length + # is constant during training. + args.variable_seq_lengths = False + if args.mock_data and args.sft and args.sft_mock_dataset_config_json is None: + args.sft_mock_dataset_config_json = json.dumps( + { + "mode": "distribution", + "type": "lognormal", + "min_seq_len": args.seq_length // 2, + "max_seq_len": args.seq_length, + "mean_seq_len": args.seq_length // 4 * 3, + "lognormal_sigma": 1.1, + } + ) + + if getattr(args, 'pad_packed_seq_alignment', None) is not None: + args.pad_packed_seq_alignment = _parse_pad_packed_seq_alignment( + args.pad_packed_seq_alignment + ) + if args.max_seqlen_per_dp_cp_rank is None: + raise ValueError( + '--max-seqlen-per-dp-cp-rank must be set when ' + '--pad-packed-seq-alignment is enabled.' + ) + if args.pad_packed_seq_alignment != 'max': + if args.pad_packed_seq_alignment <= 0: + raise ValueError( + "--pad-packed-seq-alignment must be 'max' or a positive integer " + "alignment." + ) + if args.pad_packed_seq_alignment > args.max_seqlen_per_dp_cp_rank: + raise ValueError( + '--pad-packed-seq-alignment must not exceed ' + f'--max-seqlen-per-dp-cp-rank ({args.max_seqlen_per_dp_cp_rank}), ' + f'got {args.pad_packed_seq_alignment}.' + ) + + if args.cuda_graph_impl != "none" and ( + args.sequence_packing_scheduler is not None or args.hybrid_context_parallel + ): + if getattr(args, 'pad_packed_seq_alignment', None) is None: + raise ValueError('THD CUDA Graph requires --pad-packed-seq-alignment to be set.') + if ( + args.pad_packed_seq_alignment != 'max' + and args.pad_packed_seq_alignment != args.max_seqlen_per_dp_cp_rank + ): + raise ValueError( + "THD CUDA Graph requires --pad-packed-seq-alignment='max' " + 'or --pad-packed-seq-alignment equal to ' + f'--max-seqlen-per-dp-cp-rank ({args.max_seqlen_per_dp_cp_rank}), ' + f'got {args.pad_packed_seq_alignment}.' + ) + # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \ @@ -1479,6 +1534,47 @@ def validate_args(args, defaults={}): if args.ckpt_format == "fsdp_dtensor": assert args.use_megatron_fsdp, "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." + # --use-varlen-dataset: independent of --sft. Cannot be combined with --sft + # because they are mutually-exclusive top-level dataset selectors that both + # drive the packed-sequence (THD) path. + if args.use_varlen_dataset: + assert not args.sft, ( + "--use-varlen-dataset and --sft are mutually exclusive; both " + "select the packed-sequence dataset family. Pick one." + ) + if args.varlen_sbhd_validation: + assert args.sequence_packing_scheduler is None, ( + "--varlen-sbhd-validation does not use a sequence packing " + "scheduler; drop --sequence-packing-scheduler." + ) + # SBHD validation is a real-data numerical-reference path only; + # MockVarlenDataset does not implement it. + assert not args.mock_data, ( + "--varlen-sbhd-validation is not supported with --mock-data; " + "SBHD validation requires a real dataset." + ) + else: + # VarlenDataset emits one unpacked sample per __getitem__; it + # relies on an upstream packing scheduler to group variable-length + # samples into THD batches. Auto-pick a default scheduler when + # the user did not request one explicitly: + # Otherwise fall back to ``dp_balanced`` (static packing). + if args.sequence_packing_scheduler is None: + args.sequence_packing_scheduler = 'dp_balanced' + + # Packed-sequence buffer-size check. Placed after varlen scheduler + # auto-select so it validates the final resolved scheduler. + if args.sequence_packing_scheduler is not None: + args.variable_seq_lengths = True + assert args.max_seqlen_per_dp_cp_rank is not None, ( + "--max-seqlen-per-dp-cp-rank must be set when using sequence packing" + ) + total_cp_ranks = args.context_parallel_size + 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})' + ) + # Data blend checks assert args.mock_data + \ bool(args.data_path) + \ @@ -2124,6 +2220,9 @@ def _add_network_size_args(parser): "bias_dropout_fusion", "apply_rope_fusion", "mamba_training_ssm_states_dtype", + "max_seqlen_per_dp_cp_rank", + "hybrid_context_parallel", + "sequence_packing_scheduler", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") @@ -2555,8 +2654,7 @@ def _add_rl_args(parser): return parser def _add_training_args(parser): - from megatron.training.config import TrainingConfig - from megatron.training.config import ProfilingConfig + from megatron.training.config import ProfilingConfig, TrainingConfig prof_factory = ArgumentGroupFactory(ProfilingConfig) prof_group = prof_factory.build_group(parser, "profiling") @@ -2895,6 +2993,14 @@ def _add_distributed_args(parser): 'all layers will share the same communication type. Users can also ' 'specify separated types for each layer like ' '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p') + group.add_argument('--max-seqlen-per-dp-cp-rank', type=int, default=None, + help='Maximum sequence length per CP rank. This is used to calculate the ' + 'number of sub-samples assigned to each CP rank when using heterogeneous context parallel.') + group.add_argument('--hybrid-context-parallel', action='store_true', default=False, + help='Enables hybrid context parallel. This is used to balance the workload ' + 'of each CP rank when we use packed samples with variable sequence lengths. ' + 'Requires --max-seqlen-per-dp-cp-rank to be set.') + group.add_argument('--sequence-packing-scheduler', type=str, default=None, choices=['dp_balanced']) group.add_argument('--fake-process-group', action='store_true', default=False, help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ This is quite useful for profiling memory usage of distributed training with just one GPU. \ @@ -3415,7 +3521,7 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): If kitchen isn't available, nothing to do here, return unchanged parser """ try: - from megatron.core.extensions.kitchen import KitchenSpecProvider, HAVE_KITCHEN + from megatron.core.extensions.kitchen import HAVE_KITCHEN, KitchenSpecProvider except (ImportError, ModuleNotFoundError): HAVE_KITCHEN = False @@ -3442,8 +3548,69 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): def _add_sft_args(parser): group = parser.add_argument_group(title='sft') group.add_argument('--sft', action="store_true", help='Megatron SFT training') - group.add_argument('--sft-tokenizer-prompt-format', type=str, default="nemotron-h-aligned", - help='SFT prompt format.') + group.add_argument( + '--sft-tokenizer-prompt-format', + type=str, + default="nemotron-h-aligned", + help='SFT prompt format.', + ) + group.add_argument( + '--sft-mock-dataset-config-json', + type=str, + default=None, + help='This config provides the necessary information for the mock dataset. ' + 'Accepts either an inline JSON literal or a path to a JSON file containing ' + 'the same schema. You can either specify a CSV file that contains sequence lengths, ' + 'where each line stores the length of a sequence, for example: ' + '{"mode":"file","path":"/path/to/file"}. Alternatively, you can specify a distribution ' + '(currently only supporting lognormal distribution) along with the required parameters, ' + 'for example, {"mode":"distribution","type":"lognormal","min_seq_len":1024,' + '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, where sigma controls ' + 'the variability of the lognormal distribution. ' + 'If not specified and --mock-data is set, defaults to a lognormal distribution with ' + 'min_seq_len=seq_length//2, max_seq_len=seq_length, mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.', + ) + return parser + + +def _add_varlen_dataset_args(parser): + group = parser.add_argument_group(title='varlen dataset') + group.add_argument( + '--use-varlen-dataset', + action="store_true", + help='Train with VarlenDataset, a variable-length packed (THD) dataset ' + 'that consumes instruction-tuning data from a HuggingFace Hub repo id, ' + 'a local parquet file, or a local jsonl file. Schema (alpaca / sharegpt ' + '/ openai-messages) is auto-detected from the dataset columns. ' + 'Mutually exclusive with --sft. Auto-picks a sequence packing ' + 'scheduler when none is given: ``dp_balanced``. ' + 'Combine with --mock-data for a synthetic lognormal sequence-length ' + 'distribution; see --varlen-mock-dataset-config-json.', + ) + group.add_argument( + '--varlen-sbhd-validation', + action="store_true", + help='Reference SBHD mode for THD numerical verification. When set, ' + 'VarlenDataset emits SBHD-style samples right-padded to ' + '--seq-length (no cu_seqlens, no packing scheduler), so the run can ' + 'be compared against the THD path to validate correctness. ' + 'Incompatible with --sequence-packing-scheduler.', + ) + group.add_argument( + '--varlen-mock-dataset-config-json', + type=str, + default=None, + help='Mock-dataset config for --use-varlen-dataset --mock-data. ' + 'Accepts either an inline JSON literal or a path to a JSON file containing ' + 'the same schema as --sft-mock-dataset-config-json: either ' + '{"mode":"file","path":"/path/to/lengths.csv"}, ' + '{"mode":"distribution","type":"lognormal","min_seq_len":1024,' + '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, or ' + '{"mode":"verification","data_path":"/prefix/of/IndexedDataset"}. ' + 'If not specified, defaults to a lognormal distribution with ' + 'min_seq_len=seq_length//2, max_seq_len=seq_length, ' + 'mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.', + ) return parser def _add_logits_distillation_args(parser): diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index 296acc97941..27c27393f82 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -39,14 +39,16 @@ def build_pretraining_data_loader(dataset, consumed_samples): 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( total_samples=len(dataset), 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: + if ( + getattr(args, "hybrid_context_parallel", False) + and getattr(args, "sequence_packing_scheduler", None) is None + ): batch_sampler = HybridCPMegatronPretrainingSampler( total_samples=len(dataset), consumed_samples=consumed_samples, @@ -55,7 +57,8 @@ 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()) else: - # Megatron sampler + # Megatron sampler. Packing schedulers consume one microbatch at a + # time and form packed global batches themselves. batch_sampler = MegatronPretrainingSampler( total_samples=len(dataset), consumed_samples=consumed_samples, @@ -97,9 +100,21 @@ 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 VarlenDataset and packing-scheduler paths; + # they emit one variable-length dict per sample, not stack-able by + # the default collate. --varlen-sbhd-validation is excluded: it bypasses + # packing and emits fixed-length [seq_length] samples that the default + # collate stacks normally. + if ( + ( + getattr(args, "use_varlen_dataset", False) + and not getattr(args, "varlen_sbhd_validation", False) + ) + or getattr(args, "hybrid_context_parallel", False) + or getattr(args, "sequence_packing_scheduler", None) is not None + or getattr(args, "use_vanilla_collate_fn", False) + ): + extra_kwargs = {"collate_fn": lambda x: x} else: extra_kwargs = {} return torch.utils.data.DataLoader( @@ -225,7 +240,6 @@ def __iter__(self): 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 3f93927387d..80c625d665c 100644 --- a/megatron/training/datasets/sft_dataset.py +++ b/megatron/training/datasets/sft_dataset.py @@ -1,15 +1,18 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -import atexit, json +import atexit from collections import Counter -from typing import Any, Dict, Optional +import math +from typing import Any, Dict, List, Optional, Union import numpy as np +import pandas as pd import torch from megatron.core.datasets.gpt_dataset import GPTDatasetConfig from megatron.core.datasets.megatron_dataset import LowLevelDataset, MegatronDataset from megatron.core.datasets.utils import Split +from megatron.training.datasets.utils import load_json_arg IGNORE_INDEX = -100 @@ -61,6 +64,8 @@ def __init__( config: GPTDatasetConfig, ) -> None: super().__init__(dataset, dataset_path, indices, num_samples, index_split, config) + # Pre-calculate padding divisor to avoid redundant computation in get_item + self.padding_divisor = self._calculate_padding_divisor() @staticmethod def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int: @@ -88,6 +93,26 @@ 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 hybrid_cp else cp_pad + divisor = cp_pad * tp_pad + """ + if self.config.hybrid_context_parallel: + # Hybrid 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 +149,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.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). @@ -199,3 +223,171 @@ def extend_with_padding(tokens, targets, positions, pad_len): 'cu_seqlens': padded_cu_seqlens, 'max_seqlen': max_seqlen, } + + +class MockSFTLowLevelDataset: + """The low-level mock dataset for SFT + + Args: + mode (str): Either 'file' or 'distribution'. + **kwargs: Additional arguments depending on mode. + For mode='file': path (str) - path to a CSV file with sequence lengths. + For mode='distribution': type (str), min_seq_len (int), max_seq_len (int), + mean_seq_len (int), and distribution-specific params (e.g. lognormal_sigma). + """ + + seed: int = 0 + """The hard-coded random seed to use to set the NumPy RNG""" + + size: int = 1000000 + """The hard-coded number of sequence to generate""" + + def __init__(self, mode: str, **kwargs) -> None: + np.random.seed(self.seed) + + if mode == "file": + self.sequence_lengths = np.array(pd.read_csv(kwargs["path"])).flatten() + self.size = len(self.sequence_lengths) + elif mode == "distribution": + min_seq_len = kwargs["min_seq_len"] + max_seq_len = kwargs["max_seq_len"] + mean_seq_len = kwargs["mean_seq_len"] + if kwargs["type"] == "lognormal": + lognormal_sigma = kwargs["lognormal_sigma"] + self.sequence_lengths = self.generate_lognormal_samples( + self.size, mean_seq_len, lognormal_sigma, min_seq_len, max_seq_len + ) + else: + raise ValueError(f"Unsupported distribution type {kwargs['type']}") + else: + raise ValueError(f"Unsupported mode '{mode}', must be 'file' or 'distribution'") + + def generate_lognormal_samples(self, size, mean, sigma, min_seq_len, max_seq_len): + mu = np.log(mean) - sigma**2 / 2 + samples = np.random.lognormal(mu, sigma, size) + samples = np.clip(samples, min_seq_len, max_seq_len) + return samples.astype(int) + + def __len__(self) -> int: + return self.size + + def __getitem__(self, idx: int) -> List[np.ndarray]: + # the length of sample is 'length', but only length-1 elements are generated here, + # because an eod token will be appended at the end later in SFTDataset + + length = self.sequence_lengths[idx % self.size] + sample = np.arange(1, length, dtype=np.int64) + return sample + + +class MockSFTDataset(SFTDataset): + """The mock dataset used during SFT""" + + def __init__( + self, + dataset: LowLevelDataset, + dataset_path: Optional[str], + indices: np.ndarray, + num_samples: Optional[int], + index_split: Split, + config: GPTDatasetConfig, + ) -> None: + super().__init__(dataset, dataset_path, indices, num_samples, index_split, config) + + @staticmethod + def build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> LowLevelDataset: + if config.sft_mock_dataset_config_json is None: + mock_config = { + "mode": "distribution", + "type": "lognormal", + "min_seq_len": config.sequence_length // 2, + "max_seq_len": config.sequence_length, + "mean_seq_len": config.sequence_length // 4 * 3, + "lognormal_sigma": 1.1, + } + else: + mock_config = load_json_arg(config.sft_mock_dataset_config_json) + return MockSFTLowLevelDataset(**mock_config) + + def __len__(self) -> int: + return self.num_samples + + def __getitem__(self, idx: int) -> Dict[str, Any]: + + tokenizer = self.config.tokenizer + pack_length = self.config.sequence_length + eod = tokenizer.eod + pad = tokenizer.pad + + tokens = self.dataset[int(self.indices[idx % len(self.indices)])] + + def extend_with_padding(tokens, targets, positions, pad_len): + tokens.extend([pad] * pad_len) + targets.extend([pad] * pad_len) + positions.extend(range(positions[-1] + 1, positions[-1] + 1 + pad_len)) + + # Convert tokens to list and add EOD + tokens_list = tokens.tolist() + if tokens_list[-1] != eod: + tokens_list.append(eod) + targets_list = list(tokens_list) + + pack_tokens = list(tokens_list) + pack_targets = list(targets_list) + pack_positions = list(range(len(tokens_list))) + cu_seqlens = [0] + + # Pad to padding_divisor alignment + if self.padding_divisor > 1: + mod_token_count = len(pack_tokens) % self.padding_divisor + if mod_token_count != 0: + pad_len = self.padding_divisor - mod_token_count + extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) + + # Record padded boundary after padding + cu_seqlens.append(len(pack_tokens)) + + # Handle any necessary truncation + if len(pack_tokens) >= pack_length + 1: # +1 here to account for later alignment + max_body = pack_length - 1 + pack_tokens = pack_tokens[:max_body] + pack_targets = pack_targets[:max_body] + pack_tokens.extend([eod, pad]) + pack_targets.extend([eod, pad]) + pack_positions = pack_positions[:pack_length + 1] + cu_seqlens[-1] = len(pack_tokens) - 1 + + # Handle any necessary padding + if len(pack_tokens) < pack_length + 1: # +1 here to account for later alignment + pad_len = pack_length + 1 - len(pack_tokens) + extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) + cu_seqlens[-1] = len(pack_tokens) - 1 + + assert len(pack_tokens) == pack_length + 1 + assert len(pack_targets) == pack_length + 1 + assert len(pack_positions) == pack_length + 1 + + # Align and convert to tensors + input_ids = torch.tensor(pack_tokens[:-1], dtype=torch.int64) + labels = torch.tensor(pack_targets[1:], dtype=torch.int64) + position_ids = torch.tensor(pack_positions[:-1], dtype=torch.int64) + + # Loss mask + loss_mask = torch.ones(pack_length, dtype=torch.float32) + loss_mask[labels == pad] = 0.0 + loss_mask[labels == IGNORE_INDEX] = 0.0 + + assert len(cu_seqlens) >= 2 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32) + # Calculating max_seqlen here because of possible effects of truncation and padding + adjacent_diffs = cu_seqlens[1:] - cu_seqlens[:-1] + max_seqlen = adjacent_diffs.max() # max_seqlen is a 0-D tensor + + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': position_ids, + 'cu_seqlens': cu_seqlens, + 'max_seqlen': max_seqlen, + } diff --git a/megatron/training/datasets/utils.py b/megatron/training/datasets/utils.py new file mode 100644 index 00000000000..1fe6d7ef83e --- /dev/null +++ b/megatron/training/datasets/utils.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Shared utilities for training-side dataset helpers.""" + +import json +import os +from typing import Any, Optional + + +def load_json_arg(spec: Optional[str]) -> Optional[Any]: + """Parse a CLI JSON argument that may be either a JSON literal or a path + to a JSON file. + + The argument is interpreted as a file path when ``spec`` points to an + existing regular file on the local filesystem; otherwise it is parsed as + a JSON literal string. Returns ``None`` when ``spec`` itself is ``None``, + so callers can use it transparently for optional CLI flags. + + Used by the ``--sft-mock-dataset-config-json`` and + ``--varlen-mock-dataset-config-json`` flags, which both accept either an + inline JSON snippet or the path to a file containing the same JSON + document. + """ + if spec is None: + return None + if os.path.isfile(spec): + with open(spec, "r") as f: + return json.load(f) + return json.loads(spec) diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py new file mode 100644 index 00000000000..c2533f795bb --- /dev/null +++ b/megatron/training/datasets/varlen_dataset.py @@ -0,0 +1,548 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Variable-length packed (THD) dataset for SFT-style instruction data. + +This dataset is the entry point for the ``--use-varlen-dataset`` flag. It is +independent of the ``--sft`` flag (no implicit coupling) but shares the same +THD packing / cu_seqlens / dynamic-CP padding logic by extending the existing +:class:`SFTDataset` family. The variable-length aspect is what matters here: +samples have wildly different lengths and are packed into THD format for +training throughput. + +Compared to :class:`SFTDataset`, this dataset adds: + + * **Multi-source loading** — accepts HuggingFace Hub repo ids + (``owner/repo``), local ``.parquet`` files, and local ``.jsonl/.json`` + files; the latter are read via pandas to sidestep pyarrow's per-chunk + JSON schema inference which fails when sample fields vary across rows. + + * **Auto schema detection** — four input layouts are auto-detected by column + name. The three instruction-tuning layouts are normalized to the messages + list format expected by the parent ``SFTDataset.__getitem__``; the + ``pretrain-text`` fallback instead returns a raw string handled separately + in :meth:`VarlenDataset.__getitem__`: + + * **openai-messages** — column ``messages`` (Llama post-training, + HuggingFaceH4/no_robots, ...) + * **sharegpt** — column ``conversations`` (OpenOrca, Vicuna, ...) + * **alpaca / dolly** — at least one of + ``instruction|prompt|query|question`` + one of + ``output|response|completion|answer``, plus optional context field + ``input|context``. + * **pretrain-text** — column ``text``; returns the raw string (no + messages list, no role masking), tokenized as plain pretraining text. + + * **Mock variant** — :class:`MockVarlenDataset` mirrors + :class:`MockSFTDataset` end-to-end (synthetic lognormal sequence-length + distribution / fixed-length file / verification mode from an + ``IndexedDataset``), configured via + ``--varlen-mock-dataset-config-json``. + +Limitations (raise a clear ``ValueError`` instead of silently mishandling): + + * Sample content/value must be a plain string — multi-modal content lists + (image+text parts) are not supported. + * Tree-structured (OpenAssistant oasst1) and preference (chosen/rejected) + datasets are out of scope. + * For HF Hub repos, only ``split="train"`` is loaded. +""" + +import os +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple + +import numpy as np +import torch + +from megatron.core.datasets.gpt_dataset import GPTDatasetConfig +from megatron.core.datasets.megatron_dataset import LowLevelDataset +from megatron.core.datasets.utils import Split +from megatron.training.datasets.sft_dataset import ( + IGNORE_INDEX, + MockSFTDataset, + MockSFTLowLevelDataset, + SFTDataset, + SFTLowLevelDataset, +) +from megatron.training.datasets.utils import load_json_arg + +# Field-name synonyms (probed in order; first non-empty wins). +_INSTRUCTION_FIELDS: Tuple[str, ...] = ( + "instruction", "prompt", "query", "question", +) +_OUTPUT_FIELDS: Tuple[str, ...] = ( + "output", "response", "completion", "answer", +) +# Supplementary user-turn context: Stanford Alpaca's "input", Dolly's "context". +_EXTRA_INPUT_FIELDS: Tuple[str, ...] = ("input", "context") + +# ShareGPT "from" value -> chat-template "role". Unknown values fall back to +# "user" so downstream tokenization does not crash on unfamiliar speakers. +_SHAREGPT_ROLE_MAP: Dict[str, str] = { + "human": "user", + "user": "user", + "gpt": "assistant", + "assistant": "assistant", + "model": "assistant", + "chatgpt": "assistant", + "bing": "assistant", + "bard": "assistant", + "system": "system", + "tool": "tool", + "function": "tool", + "observation": "tool", +} + + +def _looks_like_hf_id(path: str) -> bool: + """Heuristic: does ``path`` look like an ``owner/repo`` HF dataset id? + + True iff ``path`` contains ``/``, is not an absolute/relative file path, + and does not exist on the local filesystem. + """ + if not path: + return False + if os.path.exists(path): + return False + if path.startswith(("/", "./", "../")): + return False + return "/" in path + + +def _first_present( + sample: Dict[str, Any], fields: Iterable[str] +) -> Optional[str]: + """Return the first non-empty string value among the given fields, or None.""" + for f in fields: + v = sample.get(f) + if v in (None, ""): + continue + if not isinstance(v, str): + raise ValueError( + f"VarlenDataset: field '{f}' must be a string, " + f"got {type(v).__name__}." + ) + return v + return None + + +def _ensure_str_content(content: Any, where: str) -> str: + """Validate that a turn's content is a plain string (reject multi-modal lists).""" + if content is None: + return "" + if not isinstance(content, str): + raise ValueError( + f"VarlenDataset: {where} content must be a string, " + f"got {type(content).__name__}. Multi-modal datasets (e.g. " + "content as a list of image/text parts) are not supported." + ) + return content + + +def _alpaca_to_messages(sample: Dict[str, Any]) -> List[Dict[str, str]]: + """Convert an Alpaca/Dolly-style sample to a 3-turn messages list.""" + instruction = _first_present(sample, _INSTRUCTION_FIELDS) or "" + extra_input = _first_present(sample, _EXTRA_INPUT_FIELDS) or "" + output = _first_present(sample, _OUTPUT_FIELDS) or "" + user_content = ( + f"{instruction}\n\n{extra_input}" if extra_input else instruction + ) + return [ + {"role": "system", "content": ""}, + {"role": "user", "content": user_content}, + {"role": "assistant", "content": output}, + ] + + +def _sharegpt_to_messages(sample: Dict[str, Any]) -> List[Dict[str, str]]: + """Convert a ShareGPT ``conversations`` sample to a messages list. + + Prepends an empty ``system`` turn unless the conversation already starts + with one, so ``SFTDataset._split_conversations`` treats the sample as a + single conversation. + """ + conv = sample.get("conversations") or [] + out: List[Dict[str, str]] = [] + first_speaker = (conv[0].get("from") or "").lower() if conv else "" + if first_speaker != "system": + out.append({"role": "system", "content": ""}) + for turn in conv: + speaker = (turn.get("from") or "").lower() + role = _SHAREGPT_ROLE_MAP.get(speaker, "user") + content = _ensure_str_content(turn.get("value"), f"sharegpt turn role={role}") + out.append({"role": role, "content": content}) + return out + + +def _messages_passthrough(sample: Dict[str, Any]) -> List[Dict[str, str]]: + """Pass through an OpenAI ``messages`` sample, ensuring a leading system turn. + + Strips any keys other than ``role``/``content`` (e.g. ``name``, + ``tool_calls``) since they are not part of the chat-template input + expected by SFTTokenizer. + """ + raw = list(sample.get("messages") or []) + if raw and raw[0].get("role") != "system": + raw = [{"role": "system", "content": ""}] + raw + out: List[Dict[str, str]] = [] + for m in raw: + role = m.get("role") or "user" + content = _ensure_str_content(m.get("content"), f"messages turn role={role}") + out.append({"role": role, "content": content}) + return out + + +def _raw_text_loader(sample: Dict[str, Any]) -> str: + """Return the ``text`` column unchanged for pretrain-style packed runs. + + Unlike the SFT schemas this returns a plain string (no messages list). + :class:`VarlenDataset.__getitem__` dispatches on the return type to pick + a tokenization path that skips chat templating and prompt masking. + """ + text = sample.get("text") or "" + if not isinstance(text, str): + raise ValueError( + f"VarlenDataset (pretrain-text schema): 'text' must be a string, " + f"got {type(text).__name__}." + ) + return text + + +def _select_converter( + column_names: List[str], +) -> Tuple[Callable[[Dict[str, Any]], Any], str]: + """Pick a sample converter based on dataset column names. + + Priority (most explicit first): openai-messages > sharegpt > alpaca/dolly + > pretrain-text. ``pretrain-text`` is the fallback for datasets that + only carry a single ``text`` column (e.g. Dolma / OLMo midtraining + corpora) — long-context pretraining packed through the same THD path + as SFT. + """ + cols = set(column_names) + if "messages" in cols: + return _messages_passthrough, "openai-messages" + if "conversations" in cols: + return _sharegpt_to_messages, "sharegpt" + has_instr = any(f in cols for f in _INSTRUCTION_FIELDS) + has_out = any(f in cols for f in _OUTPUT_FIELDS) + if has_instr and has_out: + return _alpaca_to_messages, "alpaca" + if "text" in cols: + return _raw_text_loader, "pretrain-text" + raise ValueError( + "VarlenDataset cannot infer schema from columns " + f"{sorted(cols)}. Supported schemas: " + f"alpaca/dolly ({'|'.join(_INSTRUCTION_FIELDS)} + " + f"{'|'.join(_OUTPUT_FIELDS)} [+ optional {'|'.join(_EXTRA_INPUT_FIELDS)}]), " + "sharegpt (conversations), openai-messages (messages), " + "pretrain-text (text)." + ) + + +class VarlenLowLevelDataset(SFTLowLevelDataset): + """Low-level loader: HF Hub repo / local parquet / local jsonl, normalized. + + Dataset path interpretation: + + * HF Hub repo id (e.g. ``Yukang/LongAlpaca-12k``) — contains ``/`` and + does not exist on the local filesystem; loaded via + ``datasets.load_dataset(path, split="train")``. + * Local ``.parquet`` — loaded via + ``datasets.load_dataset("parquet", data_files=path, split="all")``; + parquet's footer schema makes chunked loading safe. + * Otherwise local jsonl/json — loaded via pandas + ``read_json(lines=True)`` and wrapped in ``Dataset.from_pandas``. + We avoid ``datasets.load_dataset("json", ...)`` for local files + because its pyarrow-based JSON reader infers schema per parallel + chunk and fails with ``CastError`` when the union of fields varies + between rows (e.g. LongAlpaca-12k). + + A per-sample converter is selected once at construction time based on + column names and applied at access time. The instruction-tuning schemas + convert to a messages list; the ``pretrain-text`` fallback returns the raw + string instead. + """ + + def __init__(self, dataset_path: str) -> None: + try: + from datasets import Dataset, load_dataset + except ImportError as exc: + raise ImportError( + "VarlenDataset requires the `datasets` library " + "(pip install datasets)." + ) from exc + + if _looks_like_hf_id(dataset_path): + self.dataset = load_dataset(dataset_path, split="train") + elif dataset_path.endswith(".parquet"): + self.dataset = load_dataset( + "parquet", data_files=dataset_path, split="all" + ) + else: + try: + import pandas as pd + except ImportError as exc: + raise ImportError( + "VarlenDataset requires `pandas` to load local jsonl " + "files (pip install pandas)." + ) from exc + df = pd.read_json(dataset_path, lines=True) + self.dataset = Dataset.from_pandas(df, preserve_index=False) + + self._converter, self._schema_name = _select_converter( + list(self.dataset.column_names) + ) + + @property + def schema_name(self) -> str: + """Detected schema name: ``alpaca`` / ``sharegpt`` / ``openai-messages`` / + ``pretrain-text`` (the raw ``text``-column fallback).""" + return self._schema_name + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> List[Dict[str, str]]: + return self._converter(self.dataset[idx]) + + +class VarlenDataset(SFTDataset): + """Variable-length single-sample SFT dataset for the packed-sequence path. + + Each ``__getitem__`` returns **one tokenized conversation** in unpacked + form: ``tokens``/``labels``/``loss_mask``/``position_ids`` whose length + equals the sample's actual token count (padded to ``pad_granularity``, + NOT to ``sequence_length``), plus ``original_seq_len``/``padded_seq_len`` + tensors that the upstream packing scheduler consumes directly via + :func:`get_batch_and_global_seqlens`. + + This is the schema described in :class:`BasePackingScheduler.get_required_sample_keys`. + It deliberately skips the multi-conversation pre-packing that + :class:`SFTDataset.__getitem__` does, letting the upstream scheduler + pack variable-length samples across the DP×CP grid with no per-sample + padding waste. + + Truncation: samples longer than ``config.sequence_length`` are truncated + on the right; an EOD token is appended if the truncation removed it. + """ + + def __init__( + self, + dataset: LowLevelDataset, + dataset_path: Optional[str], + indices: np.ndarray, + num_samples: Optional[int], + index_split: Split, + config: GPTDatasetConfig, + ) -> None: + super().__init__(dataset, dataset_path, indices, num_samples, index_split, config) + + @staticmethod + def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int: + return len(low_level_dataset) + + @staticmethod + def build_low_level_dataset( + dataset_path: str, config: GPTDatasetConfig + ) -> LowLevelDataset: + return VarlenLowLevelDataset(dataset_path) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + tokenizer = self.config.tokenizer + max_len = self.config.sequence_length + # HuggingFaceTokenizer returns None for ``pad`` when the underlying + # tokenizer has no explicit pad token (common for raw pretraining + # tokenizers like Qwen3). Fall back to eod for padding — irrelevant + # for loss because loss_mask zeros pad positions out. + eod = tokenizer.eod + pad = tokenizer.pad if tokenizer.pad is not None else eod + assert eod is not None, ( + "VarlenDataset requires the tokenizer to expose an EOD/EOS token id." + ) + + # 1. Pull a single item from the low-level dataset. For SFT schemas + # (alpaca / sharegpt / openai-messages) this is a messages list; + # for the pretrain-text schema it is a raw string. + item = self.dataset[int(self.indices[idx % len(self.indices)])] + + assert not self.config.reset_position_ids + assert not self.config.create_attention_mask and not self.config.reset_attention_mask + + # 2. Tokenize. SFT schemas go through tokenize_conversation (chat + # template + role-aware target masking); pretrain-text bypasses + # chat templating and uses the plain ``tokenize`` interface, + # treating every token as a target (no prompt masking). + if isinstance(item, str): + ids = list(tokenizer.tokenize(item)) + tokens_list = ids + targets_list = list(ids) + else: + tokens, targets = tokenizer.tokenize_conversation( + item, return_target=True, add_generation_prompt=False + ) + tokens_list = tokens.tolist() + targets_list = targets.tolist() + + # 2b. Guard against an empty tokenization (e.g. a blank ``pretrain-text`` + # row where ``tokenizer.tokenize("")`` returns no ids). Represent it + # as a single end-of-document token so the next-token shift still + # yields a valid 1-token sample instead of raising on + # ``tokens_list[-1]`` below or producing a zero-length sequence. + if len(tokens_list) == 0: + tokens_list = [eod, eod] + targets_list = [eod, eod] + + # 3. Right-truncate to ``sequence_length + 1`` (we drop the last token + # after the input/label shift below). Keep an EOD at the end so a + # truncated assistant turn still has a valid stop token. + if len(tokens_list) > max_len + 1: + tokens_list = tokens_list[: max_len + 1] + targets_list = targets_list[: max_len + 1] + if tokens_list[-1] != eod: + tokens_list[-1] = eod + targets_list[-1] = eod + + # 4. Ensure EOD is the last token (unconditional for short samples). + if tokens_list[-1] != eod: + tokens_list.append(eod) + targets_list.append(eod) + + valid_len = len(tokens_list) - 1 + + # 5a. SBHD validation mode: right-pad to sequence_length + 1, drop + # packing metadata, return shape [sequence_length]. Useful as a + # numerical reference for THD path verification (no scheduler). + if self.config.varlen_sbhd_validation: + pad_len = max_len + 1 - len(tokens_list) + if pad_len > 0: + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + assert len(tokens_list) == max_len + 1 + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + loss_mask = torch.ones(max_len, dtype=torch.float32) + loss_mask[valid_len:] = 0.0 # mask the right-padded tail by position + loss_mask[labels == IGNORE_INDEX] = 0.0 + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': torch.arange(max_len, dtype=torch.int64), + } + + original_seq_len = len(tokens_list) - 1 # length after the shift below + + # 5b. THD path: pad to pad_granularity (dp_size * cp_size * 2 * sp), + # the minimum alignment required by CP slicing. We deliberately + # do NOT pad to sequence_length — the upstream packing scheduler + # will combine variable-length samples up to + # max_seqlen_per_dp_cp_rank. + pad_granularity = self._calculate_padding_divisor() + mod = original_seq_len % pad_granularity + if mod != 0: + pad_len = pad_granularity - mod + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + padded_seq_len = len(tokens_list) - 1 + + # 6. Apply the next-token shift. + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + position_ids = torch.arange(padded_seq_len, dtype=torch.int64) + loss_mask = torch.ones(padded_seq_len, dtype=torch.float32) + loss_mask[valid_len:] = 0.0 # mask the right-padded tail by position + loss_mask[labels == IGNORE_INDEX] = 0.0 + + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': position_ids, + # The packing scheduler consumes these directly; cu_seqlens / + # max_seqlen are produced downstream in _pack_sequences. + 'original_seq_len': torch.tensor([original_seq_len], dtype=torch.int32), + 'padded_seq_len': torch.tensor([padded_seq_len], dtype=torch.int32), + } + + +class MockVarlenDataset(MockSFTDataset): + """Mock variable-length dataset for benchmarking the varlen path. + + Uses :class:`MockSFTLowLevelDataset` for sequence-length sampling (lognormal + distribution / per-line CSV / IndexedDataset verification mode — same JSON + schema as ``--sft-mock-dataset-config-json``, just consumed via + ``--varlen-mock-dataset-config-json``). + + Output shape mirrors :class:`VarlenDataset.__getitem__` (not the inherited + :meth:`MockSFTDataset.__getitem__`) so the mock and real-data paths + exercise exactly the same downstream pipeline: + + * THD mode: emits **one unpacked sample** padded to ``pad_granularity`` + with ``original_seq_len`` / ``padded_seq_len`` tensors. The upstream + scheduler packs across the DP×CP grid. + + ``--varlen-sbhd-validation`` is intentionally not implemented for mock + data; it is guarded against in argument validation. + """ + + @staticmethod + def build_low_level_dataset( + dataset_path: str, config: GPTDatasetConfig + ) -> LowLevelDataset: + if config.varlen_mock_dataset_config_json is None: + mock_config = { + "mode": "distribution", + "type": "lognormal", + "min_seq_len": config.sequence_length // 2, + "max_seq_len": config.sequence_length, + "mean_seq_len": config.sequence_length // 4 * 3, + "lognormal_sigma": 1.1, + } + else: + mock_config = load_json_arg(config.varlen_mock_dataset_config_json) + return MockSFTLowLevelDataset(**mock_config) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + tokenizer = self.config.tokenizer + max_len = self.config.sequence_length + eod = tokenizer.eod + pad = tokenizer.pad if tokenizer.pad is not None else eod + + # MockSFTLowLevelDataset returns ``length - 1`` token ids; append EOD + # to make the conversation end on a stop token, mirroring the real + # VarlenDataset path. + raw = self.dataset[int(self.indices[idx % len(self.indices)])] + tokens_list = raw.tolist() + tokens_list.append(eod) + # Mock data uses ``tokens == targets`` (no role masking). + targets_list = list(tokens_list) + + # MockVarlenDataset only implements the THD (packed) path; SBHD + # validation is a real-data numerical-reference mode (guarded against + # --mock-data in validate_args). + # THD mode: unpacked single sample, pad to pad_granularity only. + if len(tokens_list) > max_len + 1: + tokens_list = tokens_list[: max_len - 1] + [eod] + targets_list = targets_list[: max_len - 1] + [eod] + original_seq_len = len(tokens_list) - 1 + + pad_granularity = self._calculate_padding_divisor() + mod = original_seq_len % pad_granularity + if mod != 0: + pad_len = pad_granularity - mod + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + padded_seq_len = len(tokens_list) - 1 + + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + loss_mask = torch.ones(padded_seq_len, dtype=torch.float32) + loss_mask[original_seq_len:] = 0.0 # mask the right-padded tail by position + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': torch.arange(padded_seq_len, dtype=torch.int64), + 'original_seq_len': torch.tensor([original_seq_len], dtype=torch.int32), + 'padded_seq_len': torch.tensor([padded_seq_len], dtype=torch.int32), + } diff --git a/megatron/training/training.py b/megatron/training/training.py index dde1585c602..326f789c39f 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -42,7 +42,7 @@ # First-party. from megatron.core import mpu, nccl_allocator, tensor_parallel -from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper +from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper, wrap_data_iterator from megatron.core.distributed import DistributedDataParallel as DDP from megatron.core.distributed import ( DistributedDataParallelConfig, @@ -252,6 +252,7 @@ # never call ``update_*`` so the flag stays ``False`` and no collective fires. _seqlen_stats_in_iteration: Optional[torch.Tensor] = None _seqlen_stats_active: bool = False +_seqlen_stats_are_global: bool = False # Only report memory for first 3 checkpoint saves. num_checkpoints_memory_reported = 0 @@ -313,7 +314,7 @@ def update_seqlen_stats_from_cu_seqlens(cu_seqlens): the all-reduce; BSHD callers that never invoke this function leave the flag at ``False`` and pay zero collective cost. """ - global _seqlen_stats_in_iteration, _seqlen_stats_active + global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global if cu_seqlens is None or cu_seqlens.numel() < 2: return # Pin the accumulator to the current CUDA device when available so the @@ -332,6 +333,21 @@ def update_seqlen_stats_from_cu_seqlens(cu_seqlens): _seqlen_stats_in_iteration[0] += seqlens.sum() _seqlen_stats_in_iteration[1] += (seqlens * seqlens).sum() _seqlen_stats_active = True + _seqlen_stats_are_global = False + + +def set_seqlen_stats_in_iteration(total_real_tokens, seqlen_squared_sum): + """Seed per-iteration THD FLOPs stats that were already computed globally.""" + global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global + if total_real_tokens is None or seqlen_squared_sum is None: + return + if _seqlen_stats_in_iteration is None: + device = torch.device(f'cuda:{torch.cuda.current_device()}') if torch.cuda.is_available() else 'cpu' + _seqlen_stats_in_iteration = torch.zeros(2, dtype=torch.float64, device=device) + _seqlen_stats_in_iteration[0] = float(total_real_tokens) + _seqlen_stats_in_iteration[1] = float(seqlen_squared_sum) + _seqlen_stats_active = True + _seqlen_stats_are_global = True def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float]]: @@ -359,13 +375,15 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float replicated across TP/CP/PP); the world all-reduce therefore overcounts by a factor of ``TP * CP * PP``, which we divide out. """ - global _seqlen_stats_in_iteration, _seqlen_stats_active + global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global if not _seqlen_stats_active: # BSHD path: never allocated the tensor; tell the caller to use the # closed-form defaults. return None, None t = _seqlen_stats_in_iteration - if torch.distributed.is_initialized() and mpu.model_parallel_is_initialized(): + if _seqlen_stats_are_global: + dedup = 1 + elif torch.distributed.is_initialized() and mpu.model_parallel_is_initialized(): torch.distributed.all_reduce(t) tp_size = max(mpu.get_tensor_model_parallel_world_size(), 1) cp_size = max(mpu.get_context_parallel_world_size(), 1) @@ -381,6 +399,7 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float # iterations reuse it without reallocating. t.zero_() _seqlen_stats_active = False + _seqlen_stats_are_global = False return total_real_tokens / dedup, seqlen_squared_sum / dedup @@ -2341,6 +2360,20 @@ 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 getattr(config, "sequence_packing_scheduler", None) is not None: + ( + data_iterator, + scheduled_num_microbatches, + total_real_tokens_in_batch, + seqlen_squared_sum_in_batch, + ) = wrap_data_iterator(data_iterator, config, get_num_microbatches()) + set_seqlen_stats_in_iteration( + total_real_tokens_in_batch, + seqlen_squared_sum_in_batch, + ) + else: + scheduled_num_microbatches = get_num_microbatches() + # Forward pass. if save_activations_in_this_iteration: enable_activation_logging(model, args.save) @@ -2352,7 +2385,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=scheduled_num_microbatches, seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, decoder_seq_length=args.decoder_seq_length, @@ -2401,7 +2434,17 @@ def _save_state_dict(attr_name, label): should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit() if should_exit: - return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0 + return ( + {}, + True, + should_checkpoint, + should_exit, + exit_code, + None, + None, + 0, + scheduled_num_microbatches, + ) # Empty unused memory. if args.empty_unused_memory_level >= 1: @@ -2495,8 +2538,19 @@ def _save_state_dict(attr_name, label): grad_norm, num_zeros_in_grad, log_max_attention_logit, + scheduled_num_microbatches, ) - return {}, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad, log_max_attention_logit + return ( + {}, + skipped_iter, + should_checkpoint, + should_exit, + exit_code, + grad_norm, + num_zeros_in_grad, + log_max_attention_logit, + scheduled_num_microbatches, + ) def training_log( @@ -2515,6 +2569,7 @@ def training_log( is_first_iteration=False, seqlen_squared_sum_in_batch: float | None = None, total_real_tokens_in_batch: float | None = None, + num_microbatches: int | None = None, ): """Log training information such as losses, timing, ....""" args = get_args() @@ -2684,7 +2739,7 @@ def training_log( # Log MoE metrics. moe_log_string = "" if args.num_experts is not None: - moe_loss_scale = 1 / get_num_microbatches() + moe_loss_scale = 1 / (num_microbatches or get_num_microbatches()) track_names = [] if "aux_loss" in args.moe_router_load_balancing_type: track_names.append("load_balancing_loss") @@ -3621,6 +3676,7 @@ def trace_handler(p): seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, optimizers=[optimizer], + thd_sequence_length_upper_bound=_get_thd_sequence_length_upper_bound(args), ) # Run training iterations till done. @@ -3659,9 +3715,9 @@ def trace_handler(p): # Standard microbatch update (sequence packing overrides this in rl_utils.py) update_num_microbatches(args.consumed_train_samples, consistency_check=False, verbose=True) # Skip automatic checkpoint on microbatch changes when sequence packing is active - # as it intentionally reconfigures microbatches + # as it intentionally reconfigures microbatches. if get_num_microbatches() != num_microbatches and iteration != 0: - if args.rl_use_sequence_packing: + if args.rl_use_sequence_packing or args.sequence_packing_scheduler is not None: print_rank_0( f"[Sequence Packing] Skipping automatic checkpoint at iteration {iteration} " f"(microbatch change: {num_microbatches} -> {get_num_microbatches()})" @@ -3699,6 +3755,9 @@ def trace_handler(p): # Completely skip iteration if needed. if (iteration + 1) in args.iterations_to_skip: + assert ( + getattr(config, "sequence_packing_scheduler", None) is None + ), "Sequence packing scheduler is not supported in skip iteration mode" # Dummy train_step to fast forward train_data_iterator. dummy_train_step(train_data_iterator) if iteration == start_iteration: @@ -3745,6 +3804,7 @@ def trace_handler(p): grad_norm = 0.0 num_zeros_in_grad = 0 max_attention_logit = None + num_microbatches = get_num_microbatches() else: ft_integration.on_training_step_start() ( @@ -3756,6 +3816,7 @@ def trace_handler(p): grad_norm, num_zeros_in_grad, max_attention_logit, + num_microbatches, ) = train_step( forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration, pg_collection=pg_collection, @@ -3901,6 +3962,7 @@ def trace_handler(p): is_first_iteration=is_first_iteration, seqlen_squared_sum_in_batch=seqlen_squared_sum_in_batch, total_real_tokens_in_batch=total_real_tokens_in_batch, + num_microbatches=num_microbatches, ) is_first_iteration = False @@ -4131,11 +4193,21 @@ def evaluate( # Don't care about timing during evaluation config.timers = None ft_integration.on_eval_step_start() + if getattr(config, "sequence_packing_scheduler", None) is not None: + try: + (packed_data_iterator, scheduled_eval_num_microbatches, _, _) = ( + wrap_data_iterator(data_iterator, config, eval_num_microbatches) + ) + except StopIteration: + 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, @@ -4588,3 +4660,40 @@ def should_disable_forward_pre_hook(args): ) and args.overlap_param_gather ) + + +def _get_thd_sequence_length_upper_bound(args): + """Return the padded per-sample THD length upper bound used for graph sizing.""" + max_sequence_length = getattr(args, "seq_length", None) + mock_config_spec = None + if getattr(args, "use_varlen_dataset", False): + mock_config_spec = getattr(args, "varlen_mock_dataset_config_json", None) + elif getattr(args, "sft", False): + mock_config_spec = getattr(args, "sft_mock_dataset_config_json", None) + + if mock_config_spec is not None: + from megatron.training.datasets.utils import load_json_arg + + mock_config = load_json_arg(mock_config_spec) + if isinstance(mock_config, dict) and mock_config.get("max_seq_len") is not None: + max_sequence_length = int(mock_config["max_seq_len"]) + + if max_sequence_length is None: + return None + + if getattr(args, "seq_length", None) is not None: + max_sequence_length = min(int(max_sequence_length), int(args.seq_length)) + + cp_size = int(getattr(args, "context_parallel_size", 1) or 1) + if getattr(args, "hybrid_context_parallel", False): + cp_pad = int(getattr(args, "data_parallel_size", 1) or 1) * cp_size * 2 + else: + cp_pad = cp_size * 2 if cp_size > 1 else 1 + + sp_pad = ( + int(getattr(args, "tensor_model_parallel_size", 1) or 1) + if getattr(args, "sequence_parallel", False) + else 1 + ) + pad_granularity = cp_pad * sp_pad + return int(math.ceil(max_sequence_length / pad_granularity) * pad_granularity) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 06bae9965f1..ede1a47f7aa 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -26,10 +26,15 @@ 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.packed_seq_params import ( + PackedSeqParams, + get_thd_padding_kwargs, + pad_sequence_for_thd, +) from megatron.core.parallel_state import ( get_context_parallel_group, get_hybrid_data_context_parallel_groups, @@ -58,7 +63,8 @@ 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 +from megatron.training.datasets.sft_dataset import MockSFTDataset, SFTDataset +from megatron.training.datasets.varlen_dataset import MockVarlenDataset, VarlenDataset from megatron.training.training import update_seqlen_stats_from_cu_seqlens from megatron.training.utils import get_blend_and_blend_per_split, is_first_or_last_pipeline_stage from model_provider import model_provider @@ -95,6 +101,22 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): args = get_args() config = core_transformer_config_from_args(args) + if args.sequence_packing_scheduler is not None: + # `get_batch_on_this_rank_for_sequence_packing` owns scheduler THD metadata + # and returns a 7-tuple including `padding_mask`. + 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_func( + layout=config.pipeline_model_parallel_layout, + mtp_num_layers=config.mtp_num_layers, + ignore_virtual=False, + vp_stage=vp_stage, + ), + vp_stage=vp_stage, + config=config, + ) + cp_size = args.context_parallel_size tp_rank = mpu.get_tensor_model_parallel_rank() is_sft = args.sft @@ -289,44 +311,91 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa timers('batch-generator', log_level=2).start() with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") - ( - attention_mask, - cu_seqlens, - cu_seqlens_padded, - hybrid_cp_group, - labels, - local_cp_size, - loss_mask, - max_seqlen, - position_ids, - tokens, - ) = get_batch(data_iterator, vp_stage) - - packed_seq_params = None - if cu_seqlens is not None: - # Squeeze the batch dim: the batch dict keeps cu_seqlens as (1, N) - # for consistency, but PackedSeqParams and TE expect 1-D. - cu_seqlens = cu_seqlens.squeeze(0) - if cu_seqlens_padded is not None: - cu_seqlens_padded = cu_seqlens_padded.squeeze(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, - tokens_per_sample=args.seq_length, - ) + batch = get_batch(data_iterator, vp_stage) + + if len(batch) == 7: + ( + tokens, + labels, + loss_mask, + attention_mask, + position_ids, + packed_seq_params, + padding_mask, + ) = batch + elif len(batch) == 6: + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = batch + padding_mask = None + else: + ( + attention_mask, + cu_seqlens, + cu_seqlens_padded, + hybrid_cp_group, + labels, + local_cp_size, + loss_mask, + max_seqlen, + position_ids, + tokens, + ) = batch + + padding_mask = None + packed_seq_params = None + if cu_seqlens is not None: + # Squeeze the batch dim: the batch dict keeps cu_seqlens as (1, N) + # for consistency, but PackedSeqParams and TE expect 1-D. + cu_seqlens = cu_seqlens.squeeze(0) + if cu_seqlens_padded is not None: + cu_seqlens_padded = cu_seqlens_padded.squeeze(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, + tokens_per_sample=args.seq_length, + ) + + # Pad the already-packed THD tensors at the end when requested. + # CUDA Graph additionally pads cu_seqlens tensors to + # thd_max_packed_sequences + 1 entries. + config = core_transformer_config_from_args(args) + if config.pad_packed_seq_alignment is not None: + alignment, target_len, max_num_seqs = get_thd_padding_kwargs( + config.pad_packed_seq_alignment, + config.max_seqlen_per_dp_cp_rank, + config.thd_max_packed_sequences, + config.cuda_graph_impl != "none", + ) + ( + tokens, + labels, + loss_mask, + position_ids, + packed_seq_params, + padding_mask, + ) = pad_sequence_for_thd( + tokens, + labels, + loss_mask, + position_ids, + packed_seq_params, + alignment=alignment, + target_len=target_len, + max_num_seqs=max_num_seqs, + pad_by_appending_dummy_seq=config.pad_packed_seq_by_appending_dummy_seq, + ) timers('batch-generator').stop() @@ -336,7 +405,13 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa args.overlap_moe_expert_parallel_comm ), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" schedule_plan = model.build_schedule_plan( - tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask + tokens, + position_ids, + attention_mask, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, ) return schedule_plan, partial(loss_func, loss_mask, model=model) else: @@ -347,6 +422,7 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa labels=labels, loss_mask=loss_mask, packed_seq_params=packed_seq_params, + padding_mask=padding_mask, ) # [ModelOpt]: model is needed to access ModelOpt distillation losses @@ -410,6 +486,10 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, "hybrid_context_parallel": args.hybrid_context_parallel, "inter_document_masking": args.dataloader_inter_document_masking, + "sft_mock_dataset_config_json": args.sft_mock_dataset_config_json, + "sequence_packing_scheduler": args.sequence_packing_scheduler, + "varlen_mock_dataset_config_json": args.varlen_mock_dataset_config_json, + "varlen_sbhd_validation": args.varlen_sbhd_validation, } # add FIM args to the config @@ -448,8 +528,22 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None is_packed_sequence = False if args.sft: - dataset_type = SFTDataset + if args.mock_data: + dataset_type = MockSFTDataset + else: + dataset_type = SFTDataset is_packed_sequence = True # SFT always uses packed sequence + elif args.use_varlen_dataset: + # Variable-length packed (THD) dataset, independent of --sft. + # Reuses SFTDataset's THD packing internally but is gated + # by its own top-level flag. + if args.mock_data: + dataset_type = MockVarlenDataset + else: + dataset_type = VarlenDataset + # SBHD validation mode runs the non-packed pipeline; THD mode + # is the packed-sequence path. + is_packed_sequence = not args.varlen_sbhd_validation else: if args.mock_data: dataset_type = MockGPTDataset diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py new file mode 100644 index 00000000000..4dfc3db7061 --- /dev/null +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -0,0 +1,796 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for :mod:`megatron.training.datasets.varlen_dataset`. + +These tests cover the schema-detection and message-normalization helpers and +the :class:`VarlenLowLevelDataset` loader. The end-to-end SFTDataset packing +behavior is exercised by the existing SFT test suite; here we focus on the +varlen-specific contracts (auto-detect schema, normalize to messages, +ValueError on unsupported shapes). +""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +# Import via the public module path so this test gets discovered through the +# regular pytest entry point. The functions under test are pure Python and do +# not require torch.distributed. +from megatron.training.datasets.sft_dataset import IGNORE_INDEX +from megatron.training.datasets.varlen_dataset import ( + MockVarlenDataset, + VarlenDataset, + VarlenLowLevelDataset, + _alpaca_to_messages, + _looks_like_hf_id, + _messages_passthrough, + _raw_text_loader, + _select_converter, + _sharegpt_to_messages, +) + +# ---------------------------------------------------------------------------- +# _looks_like_hf_id heuristic +# ---------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "path,expected", + [ + ("Yukang/LongAlpaca-12k", True), + ("HuggingFaceH4/no_robots", True), + ("databricks/databricks-dolly-15k", True), + ("/tmp/foo.jsonl", False), + ("./local.jsonl", False), + ("../up.jsonl", False), + ("singlename", False), + ("", False), + (None, False), + ], +) +def test_looks_like_hf_id(path, expected): + assert _looks_like_hf_id(path) is expected + + +# ---------------------------------------------------------------------------- +# Schema converters +# ---------------------------------------------------------------------------- + + +def test_alpaca_canonical_with_input(): + out = _alpaca_to_messages( + {"instruction": "Summarize.", "input": "Long passage", "output": "It says X."} + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + assert out[1]["content"] == "Summarize.\n\nLong passage" + assert out[2]["content"] == "It says X." + + +def test_alpaca_without_input(): + out = _alpaca_to_messages({"instruction": "Hi.", "output": "Hello."}) + assert out[0] == {"role": "system", "content": ""} + assert out[1]["content"] == "Hi." + assert out[2]["content"] == "Hello." + + +@pytest.mark.parametrize( + "instr_key,out_key", + [ + ("prompt", "response"), + ("query", "answer"), + ("question", "completion"), + ("instruction", "answer"), + ], +) +def test_alpaca_field_synonyms(instr_key, out_key): + out = _alpaca_to_messages({instr_key: "Q?", out_key: "A."}) + assert out[1]["content"] == "Q?" + assert out[2]["content"] == "A." + + +def test_dolly_instruction_context_response(): + """Dolly-15k: instruction + context + response, all via synonyms.""" + out = _alpaca_to_messages( + { + "instruction": "Who wrote 1984?", + "context": "1984 was written in 1948.", + "response": "George Orwell.", + } + ) + assert out[1]["content"] == "Who wrote 1984?\n\n1984 was written in 1948." + assert out[2]["content"] == "George Orwell." + + +def test_sharegpt_human_gpt(): + out = _sharegpt_to_messages( + {"conversations": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello"}]} + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + assert out[1]["content"] == "hi" + + +def test_sharegpt_preserves_existing_system_turn(): + out = _sharegpt_to_messages( + { + "conversations": [ + {"from": "system", "value": "be terse"}, + {"from": "human", "value": "hi"}, + {"from": "gpt", "value": "hello"}, + ] + } + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + assert out[0]["content"] == "be terse" + + +@pytest.mark.parametrize( + "speaker,expected_role", + [ + ("human", "user"), + ("user", "user"), + ("gpt", "assistant"), + ("assistant", "assistant"), + ("model", "assistant"), + ("chatgpt", "assistant"), + ("tool", "tool"), + ("function", "tool"), + ("alien", "user"), # unknown speakers fall back to user + ], +) +def test_sharegpt_role_map(speaker, expected_role): + out = _sharegpt_to_messages({"conversations": [{"from": speaker, "value": "x"}]}) + # First entry is the prepended system turn; second is the actual content. + assert out[1]["role"] == expected_role + + +def test_messages_passthrough_prepends_system_when_missing(): + out = _messages_passthrough( + {"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]} + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + + +def test_messages_passthrough_keeps_existing_system(): + out = _messages_passthrough( + { + "messages": [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok"}, + ] + } + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + assert out[0]["content"] == "be terse" + + +def test_messages_passthrough_strips_extra_keys(): + """OpenAI-style messages may carry ``name`` / ``tool_calls`` etc.; + chat-template input only wants ``role`` and ``content``.""" + out = _messages_passthrough( + { + "messages": [ + {"role": "user", "content": "hi", "name": "alice"}, + {"role": "assistant", "content": "hi alice", "tool_calls": [{"function": "foo"}]}, + ] + } + ) + for m in out: + assert set(m.keys()) == {"role", "content"} + + +# ---------------------------------------------------------------------------- +# Shape validation: reject multi-modal / non-string content +# ---------------------------------------------------------------------------- + + +def test_messages_rejects_list_content(): + with pytest.raises(ValueError, match="must be a string"): + _messages_passthrough( + {"messages": [{"role": "user", "content": [{"type": "image", "url": "x.png"}]}]} + ) + + +def test_alpaca_rejects_non_string_field(): + with pytest.raises(ValueError, match="must be a string"): + _alpaca_to_messages({"instruction": ["a", "b"], "output": "x"}) + + +def test_sharegpt_rejects_list_value(): + with pytest.raises(ValueError, match="must be a string"): + _sharegpt_to_messages({"conversations": [{"from": "human", "value": [1, 2, 3]}]}) + + +# ---------------------------------------------------------------------------- +# Pretrain-text schema +# ---------------------------------------------------------------------------- + + +def test_raw_text_loader_returns_string(): + """``text``-column samples are returned as plain strings (not messages).""" + out = _raw_text_loader({"text": "Once upon a time...", "id": "doc-1"}) + assert isinstance(out, str) + assert out == "Once upon a time..." + + +def test_raw_text_loader_handles_empty(): + assert _raw_text_loader({"text": None}) == "" + assert _raw_text_loader({}) == "" + + +def test_raw_text_rejects_non_string(): + with pytest.raises(ValueError, match="must be a string"): + _raw_text_loader({"text": [1, 2, 3]}) + + +# ---------------------------------------------------------------------------- +# Schema selector priority +# ---------------------------------------------------------------------------- + + +def test_select_converter_alpaca(): + fn, name = _select_converter(["instruction", "output", "file"]) + assert name == "alpaca" + assert fn is _alpaca_to_messages + + +def test_select_converter_alpaca_via_synonyms(): + fn, name = _select_converter(["prompt", "response"]) + assert name == "alpaca" + + +def test_select_converter_dolly_columns(): + fn, name = _select_converter(["instruction", "context", "response", "category"]) + assert name == "alpaca" + + +def test_select_converter_sharegpt(): + fn, name = _select_converter(["conversations", "id"]) + assert name == "sharegpt" + assert fn is _sharegpt_to_messages + + +def test_select_converter_messages(): + fn, name = _select_converter(["messages"]) + assert name == "openai-messages" + assert fn is _messages_passthrough + + +def test_select_converter_priority_messages_over_alpaca(): + # When both ``messages`` and alpaca-style columns are present, the more + # explicit ``messages`` schema wins. + fn, name = _select_converter(["messages", "instruction", "output"]) + assert name == "openai-messages" + + +def test_select_converter_unrecognized_columns(): + with pytest.raises(ValueError, match="cannot infer schema"): + _select_converter(["foo", "bar"]) + + +def test_select_converter_alpaca_missing_output(): + """Having an instruction column but no output column is not a match.""" + with pytest.raises(ValueError, match="cannot infer schema"): + _select_converter(["instruction", "category"]) + + +def test_select_converter_pretrain_text(): + fn, name = _select_converter(["text", "id"]) + assert name == "pretrain-text" + assert fn is _raw_text_loader + + +def test_select_converter_pretrain_text_with_metadata(): + """Real corpora (e.g. Dolma) have ``text`` + ``url`` + ``metadata``.""" + fn, name = _select_converter(["text", "url", "metadata", "id"]) + assert name == "pretrain-text" + + +def test_select_converter_alpaca_beats_pretrain_text(): + """When both ``instruction``/``output`` and ``text`` are present (rare), + the alpaca schema is more specific and should win.""" + fn, name = _select_converter(["text", "instruction", "output"]) + assert name == "alpaca" + + +def test_select_converter_messages_beats_pretrain_text(): + fn, name = _select_converter(["text", "messages"]) + assert name == "openai-messages" + + +# ---------------------------------------------------------------------------- +# VarlenLowLevelDataset on local jsonl (no HF Hub network needed) +# ---------------------------------------------------------------------------- + + +def _write_jsonl(tmp_path: Path, rows): + p = tmp_path / "data.jsonl" + with p.open("w") as f: + for row in rows: + f.write(json.dumps(row) + "\n") + return str(p) + + +def test_low_level_loads_jsonl_alpaca(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + {"instruction": "i1", "output": "o1"}, + {"instruction": "i2", "output": "o2", "file": "extra"}, + ], + ) + ll = VarlenLowLevelDataset(path) + assert len(ll) == 2 + assert ll.schema_name == "alpaca" + sample = ll[0] + assert [m["role"] for m in sample] == ["system", "user", "assistant"] + assert sample[1]["content"] == "i1" + assert sample[2]["content"] == "o1" + + +def test_low_level_loads_jsonl_sharegpt(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + {"conversations": [{"from": "human", "value": "q1"}, {"from": "gpt", "value": "a1"}]}, + {"conversations": [{"from": "human", "value": "q2"}, {"from": "gpt", "value": "a2"}]}, + ], + ) + ll = VarlenLowLevelDataset(path) + assert len(ll) == 2 + assert ll.schema_name == "sharegpt" + sample = ll[1] + # system prepended + 2 turns from the conversation + assert [m["role"] for m in sample] == ["system", "user", "assistant"] + assert sample[1]["content"] == "q2" + + +def test_low_level_loads_jsonl_messages(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + } + ], + ) + ll = VarlenLowLevelDataset(path) + assert ll.schema_name == "openai-messages" + sample = ll[0] + assert [m["role"] for m in sample] == ["system", "user", "assistant"] + + +def test_low_level_jsonl_heterogeneous_columns(tmp_path): + """Real datasets often mix rows that have / lack an optional field. Our + pandas-based loader must accept the union schema without ``CastError``.""" + pytest.importorskip("datasets") + pytest.importorskip("pandas") + rows = [{"instruction": "a", "output": "x"}] * 100 + [ + {"instruction": "b", "output": "y", "file": "extra"} + ] * 100 + path = _write_jsonl(tmp_path, rows) + ll = VarlenLowLevelDataset(path) + assert len(ll) == 200 + # Both halves should normalize to the same messages structure. + assert [m["role"] for m in ll[0]] == ["system", "user", "assistant"] + assert [m["role"] for m in ll[150]] == ["system", "user", "assistant"] + + +def test_low_level_rejects_unknown_schema(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl(tmp_path, [{"foo": "bar"}]) + with pytest.raises(ValueError, match="cannot infer schema"): + VarlenLowLevelDataset(path) + + +def test_low_level_loads_jsonl_pretrain_text(tmp_path): + """Pretrain-text corpora (Dolma / OLMo midtraining) typically have + ``text`` + extra fields like ``id`` / ``url`` / ``metadata``.""" + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + {"text": "Doc one body...", "id": "1", "url": "https://x/1"}, + {"text": "Doc two body...", "id": "2", "url": "https://x/2"}, + ], + ) + ll = VarlenLowLevelDataset(path) + assert ll.schema_name == "pretrain-text" + assert len(ll) == 2 + # Each item is a raw string, NOT a messages list. + assert ll[0] == "Doc one body..." + assert ll[1] == "Doc two body..." + + +# ---------------------------------------------------------------------------- +# VarlenDataset / MockVarlenDataset __getitem__ (fake tokenizer, no GPU) +# +# These bypass the heavy SFTDataset.__init__ and inject the minimal attributes +# __getitem__ reads, so the EOD handling / position-based loss masking / +# pad-to-divisor / packing-metadata contracts can be unit tested without a +# real tokenizer or torch.distributed. +# ---------------------------------------------------------------------------- + + +class _FakeTokenizer: + """Minimal tokenizer for exercising VarlenDataset.__getitem__. + + ``tokenize`` maps each character to a non-zero id (so plain text never + collides with ``eod``/``pad``); ``tokenize("")`` returns ``[]`` to exercise + the empty-row guard. ``tokenize_conversation`` masks non-assistant turns + with ``IGNORE_INDEX`` in the targets. + """ + + def __init__(self, eod: int = 0, pad=None): + self._eod = eod + self._pad = pad + + @property + def eod(self): + return self._eod + + @property + def pad(self): + return self._pad + + def tokenize(self, text): + return [ord(c) % 100 + 1 for c in text] # always >= 1, never eod (0) + + def tokenize_conversation(self, messages, return_target=True, add_generation_prompt=False): + tokens, targets = [], [] + for m in messages: + ids = self.tokenize(m["content"]) + tokens.extend(ids) + # Only assistant turns contribute to the loss; prompt is masked. + targets.extend(ids if m["role"] == "assistant" else [IGNORE_INDEX] * len(ids)) + return (torch.tensor(tokens, dtype=torch.int64), torch.tensor(targets, dtype=torch.int64)) + + +def _make_config(tokenizer, seq_length=64, *, cp=1, dp=1, sp=1, sbhd=False): + return SimpleNamespace( + tokenizer=tokenizer, + sequence_length=seq_length, + reset_position_ids=False, + create_attention_mask=False, + reset_attention_mask=False, + varlen_sbhd_validation=sbhd, + data_parallel_size=dp, + context_parallel_size=cp, + sequence_parallel_size=sp, + ) + + +def _make_varlen(items, config): + ds = VarlenDataset.__new__(VarlenDataset) + ds.config = config + ds.dataset = items + ds.indices = np.arange(len(items)) + return ds + + +def _make_mock_varlen(token_arrays, config): + ds = MockVarlenDataset.__new__(MockVarlenDataset) + ds.config = config + ds.dataset = token_arrays # each item exposes .tolist() + ds.indices = np.arange(len(token_arrays)) + return ds + + +def test_getitem_thd_pretrain_text_keys_and_shapes(): + tok = _FakeTokenizer(eod=0, pad=7) + ds = _make_varlen(["hello world"], _make_config(tok, seq_length=64)) + out = ds[0] + assert set(out) == { + "tokens", + "labels", + "loss_mask", + "position_ids", + "original_seq_len", + "padded_seq_len", + } + n = out["tokens"].numel() + assert out["labels"].numel() == n + assert out["loss_mask"].numel() == n + assert out["position_ids"].numel() == n + assert int(out["padded_seq_len"].item()) == n + + +def test_getitem_thd_sft_prompt_is_masked(): + tok = _FakeTokenizer(eod=0, pad=7) + messages = [ + {"role": "system", "content": ""}, + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + ds = _make_varlen([messages], _make_config(tok, seq_length=64)) + out = ds[0] + # Prompt (user) tokens are IGNORE_INDEX in labels and must be masked out; + # assistant tokens must contribute to the loss. + labels = out["labels"] + loss_mask = out["loss_mask"] + assert torch.all(loss_mask[labels == IGNORE_INDEX] == 0.0) + assert loss_mask.sum() > 0 # assistant span still contributes + + +def test_getitem_thd_pad_masked_by_position_keeps_real_eod(): + """Regression: with pad falling back to eod, the real end-of-document EOD + target must stay in the loss (masked by position, not by value).""" + tok = _FakeTokenizer(eod=0, pad=None) # pad falls back to eod + # cp=2 -> pad divisor = cp*2 = 4, so a 3-token doc gets a padding tail. + ds = _make_varlen(["abc"], _make_config(tok, seq_length=64, cp=2)) + out = ds[0] + loss_mask = out["loss_mask"].tolist() + labels = out["labels"].tolist() + # tokens=[a,b,c,eod] padded to 4 -> labels=[b,c,eod,eod(pad)] + assert len(loss_mask) == 4 + # index 2 is the real end-of-document EOD target -> kept (would be wrongly + # dropped by value-based ``labels == pad`` masking). + assert labels[2] == tok.eod and loss_mask[2] == 1.0 + # index 3 is the appended pad -> masked. + assert loss_mask[3] == 0.0 + + +def test_getitem_thd_padded_to_divisor(): + tok = _FakeTokenizer(eod=0, pad=7) + ds = _make_varlen(["abcde"], _make_config(tok, seq_length=64, cp=2)) # divisor 4 + out = ds[0] + assert int(out["padded_seq_len"].item()) % 4 == 0 + + +def test_getitem_thd_empty_text_does_not_crash(): + """A blank pretrain-text row tokenizes to [] -> must not crash and must + yield a valid (non-zero-length) sample.""" + tok = _FakeTokenizer(eod=0, pad=7) + ds = _make_varlen([""], _make_config(tok, seq_length=64)) + out = ds[0] + assert out["tokens"].numel() >= 1 + assert out["labels"].numel() == out["tokens"].numel() + assert out["loss_mask"].numel() == out["tokens"].numel() + + +def test_getitem_sbhd_pads_to_seq_length_and_masks_tail(): + tok = _FakeTokenizer(eod=0, pad=None) + ds = _make_varlen(["abc"], _make_config(tok, seq_length=8, sbhd=True)) + out = ds[0] + # SBHD emits fixed [seq_length] samples with no packing metadata. + assert set(out) == {"tokens", "labels", "loss_mask", "position_ids"} + assert out["tokens"].numel() == 8 + loss_mask = out["loss_mask"].tolist() + # tokens=[a,b,c,eod]: valid_len=3 -> first 3 kept (incl. real eod), rest masked. + assert loss_mask[0:3] == [1.0, 1.0, 1.0] + assert all(v == 0.0 for v in loss_mask[3:]) + + +def test_mock_getitem_thd_keys_and_pad_fallback(): + tok = _FakeTokenizer(eod=0, pad=None) # exercise the eod fallback (no crash) + ds = _make_mock_varlen([np.array([1, 2, 3, 4], dtype=np.int64)], _make_config(tok, cp=2)) + out = ds[0] + assert set(out) == { + "tokens", + "labels", + "loss_mask", + "position_ids", + "original_seq_len", + "padded_seq_len", + } + n = out["tokens"].numel() + assert out["labels"].numel() == n and out["loss_mask"].numel() == n + assert int(out["padded_seq_len"].item()) % 4 == 0 + + +# ---------------------------------------------------------------------------- +# THD handoff: _unpack_batch contract for VarlenDataset-style samples +# +# VarlenDataset already emits one unpacked sub-sample carrying ``padded_seq_len``, +# so _unpack_batch must short-circuit (no cu_seqlens slicing) and only normalize +# the collate batch dim. SFTDataset-style pre-packed samples (cu_seqlens, no +# padded_seq_len) still take the slicing path. +# ---------------------------------------------------------------------------- + + +def test_unpack_batch_short_circuits_for_varlen_samples(): + from megatron.core.datasets.data_schedule_utils import _unpack_batch + + # Two VarlenDataset-style samples, each already a single sub-sample with a + # leading batch dim (as added by the default collate_fn) and padded_seq_len. + batch = [ + { + "tokens": torch.arange(4, dtype=torch.int64).view(1, 4), + "labels": torch.arange(4, dtype=torch.int64).view(1, 4), + "loss_mask": torch.ones(1, 4), + "position_ids": torch.arange(4, dtype=torch.int64).view(1, 4), + "padded_seq_len": torch.tensor([4], dtype=torch.int32), + }, + { + "tokens": torch.arange(8, dtype=torch.int64).view(1, 8), + "labels": torch.arange(8, dtype=torch.int64).view(1, 8), + "loss_mask": torch.ones(1, 8), + "position_ids": torch.arange(8, dtype=torch.int64).view(1, 8), + "padded_seq_len": torch.tensor([8], dtype=torch.int32), + "original_seq_len": torch.tensor([8], dtype=torch.int32), + }, + ] + out = _unpack_batch(batch) + # Short-circuit: same number of samples (no slicing into sub-samples). + assert len(out) == 2 + # Leading collate batch dim dropped. + assert out[0]["tokens"].shape == (4,) + assert out[1]["tokens"].shape == (8,) + # Missing original_seq_len synthesized from padded_seq_len. + assert "original_seq_len" in out[0] + assert int(out[0]["original_seq_len"].item()) == 4 + # Existing original_seq_len preserved. + assert int(out[1]["original_seq_len"].item()) == 8 + + +def test_unpack_batch_slices_prepacked_cu_seqlens_samples(): + from megatron.core.datasets.data_schedule_utils import _unpack_batch + + # SFTDataset-style pre-packed sample: two sub-sequences [0:3) and [3:5), + # described by cu_seqlens, NO padded_seq_len -> takes the slicing path. + batch = [ + { + "tokens": torch.arange(5, dtype=torch.int64), + "labels": torch.arange(5, dtype=torch.int64), + "loss_mask": torch.ones(5), + "position_ids": torch.arange(5, dtype=torch.int64), + "cu_seqlens": torch.tensor([0, 3, 5], dtype=torch.int32), + } + ] + out = _unpack_batch(batch) + # One packed sample with two sub-sequences -> two unpacked samples. + assert len(out) == 2 + assert out[0]["tokens"].numel() == 3 + assert out[1]["tokens"].numel() == 2 + assert int(out[0]["padded_seq_len"].item()) == 3 + assert int(out[1]["padded_seq_len"].item()) == 2 + + +# ---------------------------------------------------------------------------- +# DataLoader collate selection (distributed; run under torch.distributed.run). +# +# Validates the build_pretraining_data_loader contract for the varlen paths: +# * --varlen-sbhd-validation emits fixed-length [seq_length] samples that the +# DEFAULT collate stacks into a [mbs, seq_length] batch. +# * The THD path (--use-varlen-dataset without SBHD) uses the identity collate +# (variable-length dicts are returned as a list, not stacked). +# ---------------------------------------------------------------------------- + + +def _build_varlen_for_loader(items, config, num_samples): + from megatron.core.datasets.utils import Split + + ds = VarlenDataset.__new__(VarlenDataset) + ds.config = config + ds.dataset = items + ds.indices = np.arange(len(items)) + ds.num_samples = num_samples + ds.index_split = Split.train + return ds + + +def _loader_args(*, use_varlen, sbhd, scheduler, mbs, gbs=None): + return SimpleNamespace( + dataloader_type='single', + micro_batch_size=mbs, + global_batch_size=mbs if gbs is None else gbs, + full_validation=False, + num_workers=0, + use_varlen_dataset=use_varlen, + varlen_sbhd_validation=sbhd, + sequence_packing_scheduler=scheduler, + ) + + +def test_sbhd_validation_dataloader_uses_default_collate(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tok = _FakeTokenizer(eod=0, pad=7) + seq_len, mbs = 16, 2 + # One global batch needs micro_batch_size * data_parallel_size samples; + # size the dataset off the runtime DP world size so this passes under + # any --nproc-per-node (the CI default is 8 ranks -> dp=8). + dp = parallel_state.get_data_parallel_world_size() + n = mbs * dp * 4 + cfg = _make_config(tok, seq_length=seq_len, sbhd=True) + ds = _build_varlen_for_loader(["hello world"] * n, cfg, num_samples=n) + set_args(_loader_args(use_varlen=True, sbhd=True, scheduler=None, mbs=mbs)) + loader = build_pretraining_data_loader(ds, consumed_samples=0) + batch = next(iter(loader)) + # Default collate stacks fixed-length SBHD samples into a tensor batch. + assert isinstance(batch, dict) + assert batch["tokens"].shape == (mbs, seq_len) + assert batch["labels"].shape == (mbs, seq_len) + assert batch["loss_mask"].shape == (mbs, seq_len) + finally: + destroy_global_vars() + Utils.destroy_model_parallel() + + +def test_thd_dataloader_uses_identity_collate(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tok = _FakeTokenizer(eod=0, pad=7) + mbs = 2 + dp = parallel_state.get_data_parallel_world_size() + n = mbs * dp * 4 + cfg = _make_config(tok, seq_length=64, sbhd=False) + # Variable-length samples so identity collate is required. + variable = ["a", "abcdef", "xy", "qwerty"] + items = [variable[i % len(variable)] for i in range(n)] + ds = _build_varlen_for_loader(items, cfg, num_samples=n) + set_args(_loader_args(use_varlen=True, sbhd=False, scheduler="dp_balanced", mbs=mbs)) + loader = build_pretraining_data_loader(ds, consumed_samples=0) + batch = next(iter(loader)) + # Identity collate returns the raw list of per-sample dicts (unstacked). + assert isinstance(batch, list) + assert len(batch) == mbs + assert "padded_seq_len" in batch[0] + finally: + destroy_global_vars() + Utils.destroy_model_parallel() + + +def test_packing_scheduler_dataloader_yields_microbatches(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tok = _FakeTokenizer(eod=0, pad=7) + mbs = 2 + num_microbatches = 3 + dp = parallel_state.get_data_parallel_world_size() + gbs = mbs * dp * num_microbatches + n = gbs * 2 + cfg = _make_config(tok, seq_length=64, dp=dp, cp=1) + variable = ["a", "abcdef", "xy", "qwerty"] + items = [variable[i % len(variable)] for i in range(n)] + ds = _build_varlen_for_loader(items, cfg, num_samples=n) + set_args( + _loader_args( + use_varlen=True, + sbhd=False, + scheduler="dp_balanced", + mbs=mbs, + gbs=gbs, + ) + ) + loader = build_pretraining_data_loader(ds, consumed_samples=0) + batch = next(iter(loader)) + # The packing scheduler calls next(data_iterator) num_microbatches times; + # each loader step must therefore be one local microbatch, not all + # local samples from the global batch. + assert isinstance(batch, list) + assert len(batch) == mbs + assert "padded_seq_len" in batch[0] + finally: + destroy_global_vars() + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 810f48092ee..e10c370332e 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -54,7 +54,9 @@ def test_packed_freqs_returns_offset_mapped_output_for_context_parallel(self): t = torch.randn(4, 2, 8) freqs = torch.randn(8, 1, 1, 8) - out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) expected_freqs = torch.cat([freqs[0:1], freqs[3:4], freqs[4:5], freqs[7:8]], dim=0) expected = rope_utils_module._apply_rotary_pos_emb_bshd( @@ -69,7 +71,9 @@ def test_max_seqlen_freqs_returns_sequence_mapped_output_for_context_parallel(se t = torch.randn(4, 2, 8) freqs = torch.randn(4, 1, 1, 8) - out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) expected_freqs = torch.cat([freqs[1:2], freqs[2:3]], dim=0) expected_slices = [] @@ -97,6 +101,7 @@ def _test_fused_apply_mla_rope_for_q(input_format): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -142,6 +147,7 @@ def _test_fused_apply_mla_rope_for_q(input_format): freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, @@ -183,6 +189,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -241,6 +248,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 194cb2a285b..019cad9622b 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -72,6 +72,7 @@ "cuda_graph_retain_backward_graph": False, "cuda_graph_modules": [], "cuda_graph_use_single_mempool": True, + "cuda_graph_dynamic_microbatches": False, "cuda_graph_scope": None, "cuda_graph_warmup_steps": 3, "deallocate_pipeline_outputs": True, @@ -281,6 +282,7 @@ "symmetric_ar_type": None, "tensor_model_parallel_size": 2, "test_mode": False, + "thd_max_packed_sequences": 32, "timers": None, "tp_comm_atomic_ag": False, "tp_comm_atomic_rs": False, @@ -336,11 +338,12 @@ "use_transformer_engine_op_fuser": False, "moe_single_grouped_weight": False, "moe_single_grouped_bias": False, + "sequence_packing_scheduler": None, } # Fields to ignore entirely (ephemeral, environment-specific, very large). SKIP_FIELDS = set() # Fields that are allowed to appear in the live config even if not yet in the golden. -ALLOW_ADDED_FIELDS = set() +ALLOW_ADDED_FIELDS = {"pad_packed_seq_alignment", "pad_packed_seq_by_appending_dummy_seq"} def serialize_config(cfg: Any) -> Dict[str, Any]: diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py new file mode 100644 index 00000000000..fee987ae95a --- /dev/null +++ b/tests/unit_tests/test_sequence_packing.py @@ -0,0 +1,551 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import random +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.datasets.data_schedule import ( + _build_thd_padding_mask, + _get_scheduler_max_real_num_seqs, + _sanitize_thd_padding_values, + get_batch_on_this_rank_for_sequence_packing, + wrap_data_iterator, +) +from megatron.core.rerun_state_machine import RerunDataIterator +from megatron.training.global_vars import unset_global_variables +from tests.unit_tests.test_utilities import Utils + + +def test_scheduler_max_real_num_seqs_reserves_dummy_sequence(): + config = SimpleNamespace( + thd_max_packed_sequences=32, + pad_packed_seq_alignment="max", + pad_packed_seq_by_appending_dummy_seq=True, + ) + + assert _get_scheduler_max_real_num_seqs(config) == 31 + + config.pad_packed_seq_by_appending_dummy_seq = False + assert _get_scheduler_max_real_num_seqs(config) == 32 + + config.pad_packed_seq_alignment = None + config.pad_packed_seq_by_appending_dummy_seq = True + assert _get_scheduler_max_real_num_seqs(config) == 32 + + +def test_scheduler_max_real_num_seqs_rejects_dummy_without_capacity(): + config = SimpleNamespace( + thd_max_packed_sequences=1, + pad_packed_seq_alignment="max", + pad_packed_seq_by_appending_dummy_seq=True, + ) + + with pytest.raises(ValueError, match="includes that dummy sequence"): + _get_scheduler_max_real_num_seqs(config) + + +def test_scheduler_thd_padding_mask_from_cu_seqlens(): + cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32) + cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32) + + padding_mask = _build_thd_padding_mask(cu_seqlens, cu_seqlens_padded) + + assert torch.equal( + padding_mask, torch.tensor([False, False, False, True, False, False, True, True]) + ) + + +def test_scheduler_sanitizes_thd_padding_values(): + padding_mask = torch.tensor([False, False, True, False, True]) + batch = { + 'tokens': torch.tensor([11, 12, -1, 21, -1], dtype=torch.int64), + 'labels': torch.tensor([12, 13, -1, 22, -1], dtype=torch.int64), + 'loss_mask': torch.ones(5, dtype=torch.float32), + 'position_ids': torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64), + } + + _sanitize_thd_padding_values(batch, padding_mask) + + assert torch.equal(batch['tokens'], torch.tensor([11, 12, 0, 21, 0])) + assert torch.equal(batch['labels'], torch.tensor([12, 13, 0, 22, 0])) + assert torch.equal(batch['loss_mask'], torch.tensor([1.0, 1.0, 0.0, 1.0, 0.0])) + assert torch.equal(batch['position_ids'], torch.tensor([0, 1, 0, 0, 0])) + + +class MockVariableLengthSequencePackingDataIterator: + """ + Mock data iterator for testing get_batch_on_this_rank_for_sequence_packing. + + Generates variable-length (THD format) packed sequences with deterministic + data for verification across parallel ranks. + """ + + def __init__( + self, + total_seq_length: int, + sequence_lengths: list, + local_cp_size: int = None, + device: str = "cuda", + seed: int = 42, + ): + """ + Args: + total_seq_length: Total length of packed sequences + sequence_lengths: List of individual sequence lengths (variable-length). + If None, generates random variable lengths. + device: Device to create tensors on + seed: Random seed for reproducibility + """ + self.total_seq_length = total_seq_length + self.sequence_lengths = sequence_lengths + self.local_cp_size = local_cp_size + self.device = device + self.seed = seed + assert ( + sum(self.sequence_lengths) == total_seq_length + ), f"Sequence lengths sum {sum(self.sequence_lengths)} != total {total_seq_length}" + + def __iter__(self): + """Interface for the data iterator.""" + return self + + def __next__(self): + """Generate a mock batch with variable-length THD format.""" + dev = self.device + torch.manual_seed(self.seed) + torch.cuda.manual_seed(self.seed) + + tokens = torch.randint(0, 16384, (self.total_seq_length,), dtype=torch.int64, device=dev) + + # Create position_ids that reset for each sequence (THD format) + position_ids = [] + for seq_len in self.sequence_lengths: + position_ids.extend(range(seq_len)) + position_ids = torch.tensor(position_ids, dtype=torch.int64, device=dev) + + # Labels are tokens shifted by 1 for easy verification + labels = tokens + 1 + + # Loss mask: 1.0 for all positions except padding (none here) + loss_mask = torch.ones(self.total_seq_length, dtype=torch.float32, device=dev) + + # Create cu_seqlens for variable-length packed sequences + cu_seqlens = [0] + for seq_len in self.sequence_lengths: + cu_seqlens.append(cu_seqlens[-1] + seq_len) + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=dev) + cu_seqlens_padded = cu_seqlens.clone() + + max_seqlen = torch.tensor([max(self.sequence_lengths)], dtype=torch.int32, device=dev) + + batch = { + "tokens": tokens, + "position_ids": position_ids, + "labels": labels, + "loss_mask": loss_mask, + "cu_seqlens": cu_seqlens, + "cu_seqlens_padded": cu_seqlens_padded, + "max_seqlen": max_seqlen, + } + + if not ( + parallel_state.is_pipeline_first_stage(ignore_virtual=True) + or parallel_state.is_pipeline_last_stage(ignore_virtual=True) + ): + batch["tokens"] = None + batch["position_ids"] = None + batch["labels"] = None + batch["loss_mask"] = None + + if self.local_cp_size is not None: + batch["local_cp_size"] = torch.tensor( + [self.local_cp_size], dtype=torch.int32, device=dev + ) + + return batch + + +def _gather_tensor_from_tp_group(tensor): + """Gather tensors from all TP ranks for comparison.""" + assert tensor is not None, "Tensor should not be None" + tp_size = parallel_state.get_tensor_model_parallel_world_size() + gathered = [torch.zeros_like(tensor) for _ in range(tp_size)] + torch.distributed.all_gather( + gathered, tensor, group=parallel_state.get_tensor_model_parallel_group() + ) + return gathered + + +def _gather_tensor_from_all_ranks(tensor): + """Gather tensors from all PP ranks for comparison.""" + assert tensor is not None, "Tensor should not be None" + if type(tensor) is int: + tensor = torch.tensor(tensor, dtype=torch.int32, device=torch.cuda.current_device()) + gathered = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())] + torch.distributed.all_gather(gathered, tensor) + return gathered + + +@pytest.mark.parametrize( + ("tp", "pp", "cp"), + [ + (1, 1, 1), # Basic case: no parallelism + (2, 1, 1), # Tensor parallel only + (1, 2, 1), # Pipeline parallel only + (2, 2, 1), # TP + PP + (1, 1, 2), # CP only + (2, 1, 2), # TP + CP + (1, 2, 2), # PP + CP + (1, 4, 1), # Has middle pp stage + ], +) +def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp): + """ + Test get_batch_on_this_rank_for_sequence_packing function with variable-length THD format. + + This test verifies: + 1. TP ranks: All ranks within a TP group receive identical data after broadcast + 2. PP ranks: Middle PP ranks have the same packed_seq_params as first/last stages + 3. CP ranks: Data is correctly partitioned with proper shape and values + 4. Variable-length (THD) format: Different sequence lengths are handled correctly + """ + args = SimpleNamespace() + args.tensor_model_parallel_size = tp + args.pipeline_model_parallel_size = pp + args.context_parallel_size = cp + args.virtual_pipeline_model_parallel_size = None + args.data_parallel_size = 8 // (tp * pp * cp) + args.seq_length = 8192 + + # Skip invalid configurations + if args.data_parallel_size < 1: + raise ValueError(f"Invalid config: tp={tp}, pp={pp}, cp={cp} exceeds world size 8") + + # Initialize model parallel + Utils.initialize_model_parallel(tp, pp, None, context_parallel_size=cp) + + try: + # Create mock data iterator with variable-length sequences + # Only TP rank 0 needs the iterator; other TP ranks pass None + tp_rank = parallel_state.get_tensor_model_parallel_rank() + if tp_rank == 0: + # Use deterministic seed based on DP rank so same data within TP/PP/CP group + dp_rank = parallel_state.get_data_parallel_rank() + sequence_lengths = [1024, 2048, 512, 1536, 3072] + assert ( + sum(sequence_lengths) == args.seq_length + ), f"Sequence lengths sum {sum(sequence_lengths)} != total {args.seq_length}" + data_iterator = iter( + MockVariableLengthSequencePackingDataIterator( + total_seq_length=args.seq_length, + sequence_lengths=sequence_lengths, # Variable lengths, sum=8192 + seed=42 + dp_rank, # Same seed within PP/CP group + ) + ) + else: + # Non-TP-rank-0 ranks don't need the iterator + data_iterator = None + + # Call the function under test + result = get_batch_on_this_rank_for_sequence_packing( + data_iterator=data_iterator, mtp_on_this_rank=False, vp_stage=None + ) + + # Unpack the result. Scheduler THD always returns padding_mask. + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = ( + result + ) + + # Get parallel state info + tp_rank = parallel_state.get_tensor_model_parallel_rank() + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + cp_rank = parallel_state.get_context_parallel_rank() + is_first_stage = parallel_state.is_pipeline_first_stage(ignore_virtual=True) + is_last_stage = parallel_state.is_pipeline_last_stage(ignore_virtual=True) + is_first_or_last = is_first_stage or is_last_stage + + assert padding_mask is not None + assert padding_mask.dtype == torch.bool + assert padding_mask.dim() == 2 + assert padding_mask.size(0) == 1 + assert not padding_mask.any(), "Mock data has no per-sequence padding." + + # ===================================================================== + # TEST 1: Verify data based on pipeline stage + # ===================================================================== + if is_first_stage: + assert tokens is not None, "First stage should have tokens" + assert position_ids is not None, "First stage should have position_ids" + assert tokens.dim() == 2, "Tokens should be 2D (batch, seq)" + assert position_ids.dim() == 2, "Position IDs should be 2D (batch, seq)" + assert tokens.size(0) == 1, "batch should be 1 in THD format" + assert position_ids.size(0) == 1, "batch should be 1 in THD format" + else: + assert tokens is None, "Non-first stage should not have tokens" + assert position_ids is None, "Non-first stage should not have position_ids" + + if is_last_stage: + assert labels is not None, "Last stage should have labels" + assert loss_mask is not None, "Last stage should have loss_mask" + assert labels.dim() == 2, "Labels should be 2D (batch, seq)" + assert loss_mask.dim() == 2, "Loss mask should be 2D (batch, seq)" + assert labels.size(0) == 1, "batch should be 1 in THD format" + assert loss_mask.size(0) == 1, "batch should be 1 in THD format" + else: + assert labels is None, "Non-last stage should not have labels" + assert loss_mask is None, "Non-last stage should not have loss_mask" + + # ===================================================================== + # TEST 2: Verify all ranks have consistent packed_seq_params + # ===================================================================== + assert packed_seq_params is not None + assert packed_seq_params.qkv_format == "thd" + + test_keys = [ + "cu_seqlens_q", + "cu_seqlens_q_padded", + "max_seqlen_q", + "cu_seqlens_kv", + "cu_seqlens_kv_padded", + "max_seqlen_kv", + ] + for key in test_keys: + tensor = getattr(packed_seq_params, key) + assert tensor is not None + gathered_tensor = _gather_tensor_from_all_ranks(tensor) + for i in range(1, len(gathered_tensor)): + assert torch.equal( + gathered_tensor[0], gathered_tensor[i] + ), f"Rank 0 and rank {i} have different {key}" + + # ===================================================================== + # TEST 3: Verify TP ranks receive identical data after broadcast + # ===================================================================== + if tp > 1: + test_tensors = [padding_mask] + if is_first_stage: + test_tensors.extend([tokens, position_ids]) + if is_last_stage: + test_tensors.extend([labels, loss_mask]) + + for tensor in test_tensors: + gathered_tensors = _gather_tensor_from_tp_group(tensor) + for i in range(1, tp): + assert torch.equal( + gathered_tensors[0], gathered_tensors[i] + ), f"TP rank 0 and rank {i} have different data" + + # ===================================================================== + # TEST 4: Verify CP partitioning + # ===================================================================== + if cp > 1: + # With CP, the sequence should be partitioned + expected_seq_len = args.seq_length // cp + + if is_first_stage: + actual_seq_len = tokens.shape[1] + assert ( + actual_seq_len == expected_seq_len + ), f"CP partitioned tokens have wrong shape: {actual_seq_len} != {expected_seq_len}" + + # Verify labels only if all CP ranks are at last stage + if is_last_stage: + actual_seq_len = labels.shape[1] + assert ( + actual_seq_len == expected_seq_len + ), f"CP partitioned labels have wrong shape: {actual_seq_len} != {expected_seq_len}" + + actual_seq_len = padding_mask.shape[1] + assert ( + actual_seq_len == expected_seq_len + ), f"CP partitioned padding_mask has wrong shape: {actual_seq_len} != {expected_seq_len}" + + finally: + Utils.destroy_model_parallel() + unset_global_variables() + + +@pytest.mark.parametrize( + ("tp", "pp", "cp", "vpp", "scheduler_type"), + [ + (1, 1, 8, None, "dp_balanced"), + (2, 1, 4, None, "dp_balanced"), + (2, 4, 1, None, "dp_balanced"), + (2, 2, 1, None, "dp_balanced"), + (1, 4, 1, 4, "dp_balanced"), + ], +) +def test_wrap_dataloader(tp, pp, cp, vpp, scheduler_type): + ''' + Test wrap_dataloader function with different scheduler types. + ''' + args = SimpleNamespace() + args.tensor_model_parallel_size = tp + args.pipeline_model_parallel_size = pp + args.context_parallel_size = cp + args.virtual_pipeline_model_parallel_size = None + args.data_parallel_size = 8 // (tp * pp * cp) + args.seq_length = 8192 + args.max_seqlen_per_dp_cp_rank = 8192 + + # Skip invalid configurations + if args.data_parallel_size < 1: + raise ValueError(f"Invalid config: tp={tp}, pp={pp}, cp={cp} exceeds world size 8") + + def _create_single_sample(seq_len): + # hard code the padding size to 16 + pad_size = 16 + seq_len_padded = ((seq_len + pad_size - 1) // pad_size) * pad_size + device = torch.device("cuda", torch.cuda.current_device()) + tokens = torch.randint(0, 128, (seq_len_padded,), dtype=torch.int64, device=device) + labels = tokens + 1 + position_ids = torch.arange(seq_len_padded, dtype=torch.int64, device=device) + loss_mask = torch.ones(seq_len_padded, dtype=torch.float32, device=device) + loss_mask[0:seq_len] = 1 + loss_mask[seq_len:] = 0 + cu_seqlens = torch.tensor([0, seq_len_padded], dtype=torch.int32, device=device) + + return { + 'tokens': tokens, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': position_ids, + 'cu_seqlens': cu_seqlens, + } + + # Initialize model parallel + Utils.initialize_model_parallel(tp, pp, vpp, context_parallel_size=cp) + + global_batch_size = 64 + micro_batch_size = 1 + nums = [random.randint(2048, args.seq_length) for _ in range(global_batch_size)] # 64 sequences + + config = SimpleNamespace() + config.max_seqlen_per_dp_cp_rank = args.max_seqlen_per_dp_cp_rank + config.microbatch_group_size_per_vp_stage = pp + config.virtual_pipeline_model_parallel_size = vpp + config.sequence_packing_scheduler = scheduler_type + + dp_rank = parallel_state.get_data_parallel_rank() + dp_size = parallel_state.get_data_parallel_world_size() + + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + tp_rank = parallel_state.get_tensor_model_parallel_rank() + + is_pp_first = pp_rank == 0 + is_pp_last = pp_rank == pp - 1 + is_pp_first_or_last = is_pp_first or is_pp_last + is_tp_first = tp_rank == 0 + + num_micro_batches_old = global_batch_size // micro_batch_size // dp_size + + if is_tp_first and (is_pp_first or is_pp_last): + samples = [ + _create_single_sample(num) + for num in nums[dp_rank * num_micro_batches_old : (dp_rank + 1) * num_micro_batches_old] + ] + data_iterator = RerunDataIterator(iter(samples)) + else: + data_iterator = None + + if is_tp_first: + if vpp is not None and vpp > 1: + if is_pp_first: + data_iterator = [data_iterator] + [None for _ in range(vpp - 1)] + elif is_pp_last: + data_iterator = [None for _ in range(vpp - 1)] + [data_iterator] + else: + data_iterator = [None for _ in range(vpp)] + try: + # Call the function under test + ( + new_data_iterator, + num_micro_batches, + num_total_tokens_this_global_batch, + sequence_square_sum_this_global_batch, + ) = wrap_data_iterator(data_iterator, config, num_micro_batches_old) + + # check the result + assert type(num_micro_batches) is int + assert ( + type(num_total_tokens_this_global_batch) is float + or type(num_total_tokens_this_global_batch) is np.float32 + ) + assert ( + type(sequence_square_sum_this_global_batch) is float + or type(sequence_square_sum_this_global_batch) is np.float32 + ) + + def _check_batch(batch_all, batch_keys): + for batch in batch_all: + assert set(batch_keys) <= set( + batch.keys() + ), f"batch keys: {set(batch.keys())} missing {set(batch_keys) - set(batch.keys())}" + for key in batch_keys: + assert batch[key] is not None + + if is_tp_first: + # CHECK KEYS + batch_keys = ["cu_seqlens", "max_seqlen", "cu_seqlens_padded"] + if vpp is not None and vpp > 1: + # check metadata for all stages (save batches to avoid re-consuming iterators) + all_stage_batches = [] + for temp_data_iterator in new_data_iterator: + stage_batch = [next(temp_data_iterator) for _ in range(num_micro_batches)] + all_stage_batches.append(stage_batch) + _check_batch(stage_batch, batch_keys) + + # check for first or last stage on first or last pp rank + if is_pp_first_or_last: + batch_all = all_stage_batches[0] if is_pp_first else all_stage_batches[-1] + batch_keys += ["tokens", "position_ids", "labels", "loss_mask"] + _check_batch(batch_all, batch_keys) + else: + # non-VPP: single iterator + batch_all = [next(new_data_iterator) for _ in range(num_micro_batches)] + if is_pp_first_or_last: + batch_keys += ["tokens", "position_ids", "labels", "loss_mask"] + _check_batch(batch_all, batch_keys) + + # CHECK TOKEN SUM ON FIRST OR LAST PP RANK + # Note: data_iterator is consumed by wrap_data_iterator, new_data_iterator is consumed above. + # Use `samples` for before-wrap, reuse `batch_all` from the check above for after-wrap. + if is_pp_first_or_last: + # Compute token sum before wrap + token_sum_before = torch.tensor(0, dtype=torch.int64, device='cuda') + for sample in samples: + token_sum_before += sample['tokens'].long().sum() + + # Compute token sum after wrap (batch_all already collected above with tokens) + token_sum_after = torch.tensor(0, dtype=torch.int64, device='cuda') + for batch in batch_all: + token_sum_after += batch['tokens'].long().sum() + + # Reduce sum across dp_cp group and verify equality + dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=False) + torch.distributed.all_reduce( + token_sum_before, op=torch.distributed.ReduceOp.SUM, group=dp_cp_group + ) + torch.distributed.all_reduce( + token_sum_after, op=torch.distributed.ReduceOp.SUM, group=dp_cp_group + ) + + assert ( + token_sum_before == token_sum_after + ), f"Token sum mismatch: before={token_sum_before.item()}, after={token_sum_after.item()}" + + else: + if vpp is not None and vpp > 1: + assert type(new_data_iterator) is list and len(new_data_iterator) == vpp + for data_iterator in new_data_iterator: + assert data_iterator is None + else: + assert new_data_iterator is None + + finally: + Utils.destroy_model_parallel() + unset_global_variables() diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 52edbf9a264..edbc33d159c 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1124,6 +1124,62 @@ def test_get_cuda_graph_input_data(self, num_microbatches, pp_size, vpp_size): ), f"Order length mismatch: expected {expected_order_length}, got {len(order)}" +class TestRequiredNumMicrobatchSlots: + """Pure-Python tests for ``_get_required_num_microbatch_slots_from_order``. + + The method derives the smallest cuda-graph slot count that guarantees no + in-flight microbatch's static buffer is reused before its backward + completes. ``order`` is a 1F1B / interleaved-1F1B schedule transcript + where ``+chunk_id`` denotes a forward and ``-chunk_id`` a backward. + Non-integer entries (e.g. ``0.5`` for wgrad sub-steps) are skipped. + """ + + @staticmethod + def _slots(order, num_chunks): + return TECudaGraphHelper._get_required_num_microbatch_slots_from_order(order, num_chunks) + + def test_single_chunk_single_microbatch(self): + # F0 then B0: one slot is enough. + assert self._slots([1, -1], 1) == 1 + + def test_single_chunk_pp_pipeline_4_microbatches_pp2(self): + # PP=2 1F1B with 4 microbatches: warmup F-F, then F-B-F-B-..., then cooldown B-B. + # Max in-flight = 2. + order = [1, 1, -1, 1, -1, 1, -1, -1] + assert self._slots(order, 1) == 2 + + def test_two_chunks_independent(self): + # Two model chunks (VPP=2), each running a tiny PP=2-style 1F1B in turn. + # Per chunk max in-flight = 2 -> 2 slots. + order = [1, 1, -1, -1, 2, 2, -2, -2] + assert self._slots(order, 2) == 2 + + def test_two_chunks_interleaved(self): + # Worst case: forwards stack up across chunks before any backward. + # F0 F0 F1 F1 B1 B1 B0 B0 -> per-chunk max in-flight = 2. + order = [1, 1, 2, 2, -2, -2, -1, -1] + assert self._slots(order, 2) == 2 + + def test_skips_non_integer_entries(self): + # Float c_ids (e.g. 0.5 for wgrad sub-steps) must be ignored. + order = [1, 0.5, -0.5, -1] + assert self._slots(order, 1) == 1 + + def test_minimum_slot_is_one(self): + # Empty / no-op order still returns at least 1 (we always need a slot). + assert self._slots([], 1) == 1 + + def test_unbalanced_order_asserts(self): + # Forward without matching backward -> outstanding != 0 at end -> assert. + with pytest.raises(AssertionError): + self._slots([1], 1) + + def test_negative_outstanding_asserts(self): + # Backward before any forward for a chunk -> outstanding goes negative. + with pytest.raises(AssertionError): + self._slots([-1], 1) + + def is_deep_ep_available(): from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP diff --git a/tests/unit_tests/transformer/test_thd_correctness.py b/tests/unit_tests/transformer/test_thd_correctness.py index ccf70b8a885..09dcbb725c5 100644 --- a/tests/unit_tests/transformer/test_thd_correctness.py +++ b/tests/unit_tests/transformer/test_thd_correctness.py @@ -110,14 +110,14 @@ def pad_thd_to_max(self) -> bool: # TP/CP/SP: similarity checks (TE Attention) # ------------------------------------------------------------------------- TestCase("tp2_cp4_sp", 4096, 64, 4, 12288, [2039, 1013, 509], 2, 4, True, "similarity"), - TestCase("tp2_cp2_sp_longseq", 4096, 32, 8, 14336, [65536, 8191, 4096], 2, 2, True, "similarity"), + TestCase("tp2_cp2_sp_longseq", 4096, 32, 8, 14336, [16384, 4096, 2048], 2, 4, True, "similarity"), # ------------------------------------------------------------------------- # Edge cases # ------------------------------------------------------------------------- TestCase("short_seqs_parallel", 1024, 16, 4, 4096, [17, 31, 11], 2, 2, True, "similarity"), TestCase("extreme_mixed", 4096, 32, 8, 14336, [4093, 127, 257], 2, 2, True, "similarity"), - TestCase("long_short_mix", 4096, 32, 8, 14336, [65535, 512, 1024], 2, 2, True, "similarity"), + TestCase("long_short_mix", 4096, 32, 8, 14336, [16384, 512, 1024], 2, 4, True, "similarity"), ] # fmt: on diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py new file mode 100644 index 00000000000..1c41fe64def --- /dev/null +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -0,0 +1,1005 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +""" +Unit tests for THD format with CUDA Graph support. + +Padding helpers and dataclass round-trip (any GPU count, fast): + torchrun --nproc_per_node 1 -m pytest -xvs \ + tests/unit_tests/transformer/test_thd_cuda_graph.py \ + -k "Pad or Decompose" + +End-to-end no-graph vs graph bitwise loss/grad_norm match for +Moonlight-16B and Qwen3-8B with TP2_CP2_PP2 + sequence packing +(requires 8 GPUs, slow ~5 min per run, 4 runs total). Moonlight covers +MoE router/preprocess graph capture with router fusion; Qwen3 is dense and +covers attention graph capture: + pytest -xvs tests/unit_tests/transformer/test_thd_cuda_graph.py::TestE2EBitwise + +The E2E test directly subprocesses `torchrun pretrain_gpt.py` -- the same +command exercised by test_moonlight_qwen3_bitwise.sh -- with both +cuda_graph_impl=none and cuda_graph_impl=transformer_engine, then compares +the per-iteration loss / grad_norm lines. They must be exactly equal. +""" + +import os +import re +import socket +import subprocess +from pathlib import Path + +import pytest +import torch + +from megatron.core.packed_seq_params import ( + PackedSeqParams, + _resolve_thd_padding_lengths, + get_thd_padding_kwargs, + pad_sequence_for_thd, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import TransformerLayer +from tests.unit_tests.test_utilities import Utils + +os.environ.setdefault('NVTE_ALLOW_NONDETERMINISTIC_ALGO', '0') +os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8') + + +_REQUIRES_TWO_RANKS = pytest.mark.skipif( + int(os.environ.get("WORLD_SIZE", "1")) < 2 or torch.cuda.device_count() < 2, + reason="requires torchrun with at least 2 GPUs", +) + + +# ============================================================================= +# Helpers (shared by the lightweight unit tests) +# ============================================================================= + + +def _make_cu(seqlens, device="cuda"): + cu = torch.zeros(len(seqlens) + 1, dtype=torch.int32, device=device) + for i, s in enumerate(seqlens): + cu[i + 1] = cu[i] + s + return cu + + +def _make_psp(seqlens): + cu = _make_cu(seqlens) + return PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu, + cu_seqlens_kv=cu.clone(), + cu_seqlens_q_padded=cu.clone(), + cu_seqlens_kv_padded=cu.clone(), + max_seqlen_q=max(seqlens), + max_seqlen_kv=max(seqlens), + ) + + +def _build_layer(H, nh, nkv, ffn, max_seqlen, max_num_seqs, tp=1, sp=False): + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec + + config = TransformerConfig( + num_layers=1, + hidden_size=H, + num_attention_heads=nh, + num_query_groups=nkv, + ffn_hidden_size=ffn, + max_seqlen_per_dp_cp_rank=max_seqlen, + thd_max_packed_sequences=max_num_seqs, + tensor_model_parallel_size=tp, + sequence_parallel=sp, + bf16=True, + ) + model_parallel_cuda_manual_seed(42) + return ( + TransformerLayer( + config, get_gpt_layer_with_transformer_engine_spec().submodules, layer_number=1 + ) + .cuda() + .bfloat16() + ) + + +# ============================================================================= +# 1. pad_sequence_for_thd correctness +# ============================================================================= + + +@pytest.mark.internal +@pytest.mark.parametrize("cuda_graph_static,expected_max_num_seqs", [(False, None), (True, 32)]) +def test_pad_to_max_resolves_padding_kwargs(cuda_graph_static, expected_max_num_seqs): + alignment, target_len, max_num_seqs = get_thd_padding_kwargs( + pad_packed_seq_alignment="max", + max_seqlen_per_dp_cp_rank=8192, + thd_max_packed_sequences=32, + cuda_graph_static=cuda_graph_static, + ) + + assert alignment is None + assert target_len == 8192 + assert max_num_seqs == expected_max_num_seqs + + +class TestResolveThdPaddingLengths: + + def setup_method(self): + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + + def teardown_method(self): + Utils.destroy_model_parallel() + + @pytest.mark.internal + @pytest.mark.parametrize( + "source,target_len,alignment,expected", + [ + ("tokens", None, 64, (80, 80, 128, 128)), + ("labels", 256, None, (80, 80, 256, 256)), + ("metadata", None, 64, (80, 80, 128, 128)), + ], + ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_non_cp_length_resolution_contract(self, source, target_len, alignment, expected): + """Resolve lengths from local tensors when present, otherwise from THD metadata.""" + tokens, labels = None, None + psp = _make_psp([50, 30]) + + if source == "tokens": + tokens = torch.ones(1, 80, device="cuda") + psp = PackedSeqParams(qkv_format="thd") + expected_device = tokens.device + elif source == "labels": + labels = torch.ones(1, 80, device="cuda") + expected_device = labels.device + else: + expected_device = psp.cu_seqlens_q.device + + local_actual, global_actual, local_target, global_target, mask_device = ( + _resolve_thd_padding_lengths( + tokens, labels, None, None, psp, target_len=target_len, alignment=alignment + ) + ) + + assert (local_actual, global_actual, local_target, global_target) == expected + assert mask_device == expected_device + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_no_tensor_requires_cu_seqlens(self): + """All-None tensor inputs need cu_seqlens to build a padding mask.""" + psp = PackedSeqParams(qkv_format="thd") + + with pytest.raises(AssertionError, match="cu_seqlens_q must be available"): + _resolve_thd_padding_lengths( + None, None, None, None, psp, target_len=128, alignment=None + ) + + @pytest.mark.internal + @_REQUIRES_TWO_RANKS + def test_cp_tensor_alignment_uses_local_target_and_global_tail(self): + """CP-local padding tail determines the global padded endpoint.""" + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + + tokens = torch.ones(1, 1600, device="cuda") + psp = _make_psp([1600, 1600]) + + local_actual, global_actual, local_target, global_target, mask_device = ( + _resolve_thd_padding_lengths( + tokens, None, None, None, psp, target_len=None, alignment=128 + ) + ) + + assert (local_actual, global_actual, local_target, global_target) == ( + 1600, + 3200, + 1664, + 3328, + ) + assert mask_device == tokens.device + + @pytest.mark.internal + @_REQUIRES_TWO_RANKS + def test_cp_tensor_target_len_scales_global_target(self): + """Fixed target_len is CP-local and scales to a global endpoint.""" + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + + tokens = torch.ones(1, 80, device="cuda") + psp = _make_psp([140]) + + local_actual, global_actual, local_target, global_target, mask_device = ( + _resolve_thd_padding_lengths( + tokens, None, None, None, psp, target_len=128, alignment=None + ) + ) + + assert (local_actual, global_actual, local_target, global_target) == (80, 140, 128, 256) + assert mask_device == tokens.device + + @pytest.mark.internal + @pytest.mark.parametrize( + "alignment,target_len,expected_global_target", [(128, None, 256), (None, 128, 256)] + ) + @_REQUIRES_TWO_RANKS + def test_cp_no_tensor_partitions_actual_and_target_lengths( + self, alignment, target_len, expected_global_target + ): + """Without local tensors, CP-local lengths come from THD partition indices.""" + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + + from megatron.core import parallel_state + from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices + + psp = _make_psp([140]) + cp_size = parallel_state.get_context_parallel_world_size() + cp_rank = parallel_state.get_context_parallel_rank() + expected_local_actual = get_thd_partitioned_indices( + psp.cu_seqlens_q, 140, cp_size, cp_rank + ).numel() + expected_local_target = get_thd_partitioned_indices( + psp.cu_seqlens_q, expected_global_target, cp_size, cp_rank + ).numel() + + local_actual, global_actual, local_target, global_target, mask_device = ( + _resolve_thd_padding_lengths( + None, None, None, None, psp, target_len=target_len, alignment=alignment + ) + ) + + assert (local_actual, global_actual, local_target, global_target) == ( + expected_local_actual, + 140, + expected_local_target, + expected_global_target, + ) + assert mask_device == psp.cu_seqlens_q.device + + +class TestPadSequenceForThd: + + def setup_method(self): + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + + def teardown_method(self): + Utils.destroy_model_parallel() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_generic_alignment_appends_dummy_padding_sequence(self): + """Generic THD padding covers tail slots with an independent dummy sequence.""" + seqlens, total_T = [50, 30], 80 + psp = _make_psp(seqlens) + orig = psp.cu_seqlens_q.clone() + p_tok, _, _, _, p, mask = pad_sequence_for_thd( + torch.ones(1, total_T, device="cuda"), None, None, None, psp, alignment=64 + ) + assert p_tok.shape == (1, 128) + expected = torch.cat((orig, torch.tensor([128], dtype=orig.dtype, device=orig.device))) + assert torch.equal(p.cu_seqlens_q, expected) + assert torch.equal(p.cu_seqlens_q_padded, expected) + assert p.pad_between_seqs is False + assert mask.shape == (1, 128) + assert not mask[0, :total_T].any() and mask[0, total_T:].all() + + @pytest.mark.internal + @_REQUIRES_TWO_RANKS + def test_cp_alignment_uses_global_cu_seqlens_length(self): + """CP-local token length must not cap global packed-sequence padding.""" + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + + psp = _make_psp([140]) + local_T = 80 + p_tok, _, _, _, p, mask = pad_sequence_for_thd( + torch.ones(1, local_T, device="cuda"), None, None, None, psp, alignment=128 + ) + + assert p_tok.shape[-1] >= local_T + assert p.cu_seqlens_q[-1].item() == 256 + assert p.cu_seqlens_q_padded[-1].item() == 256 + assert p.max_seqlen_q == 140 + assert p.max_seqlen_kv == 140 + assert mask.shape[-1] == p_tok.shape[-1] + assert not mask[0, :local_T].any() + + @pytest.mark.internal + @_REQUIRES_TWO_RANKS + def test_cp_alignment_covers_local_padding_tail(self): + """CP-local padding can create a global tail even when global length is aligned.""" + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + + psp = _make_psp([1600, 1600]) + local_T = 1600 + p_tok, _, _, _, p, mask = pad_sequence_for_thd( + torch.ones(1, local_T, device="cuda"), None, None, None, psp, alignment=128 + ) + + assert p_tok.shape[-1] == 1664 + assert p.cu_seqlens_q[-1].item() == 3328 + assert p.cu_seqlens_q_padded[-1].item() == 3328 + assert mask.shape[-1] == p_tok.shape[-1] + assert not mask[0, :local_T].any() + assert mask[0, local_T:].all() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_padding_without_dummy_sequence_preserves_metadata(self): + """Disabling dummy sequence padding only pads token-like tensors.""" + seqlens, total_T = [50, 30], 80 + psp = _make_psp(seqlens) + psp.pad_between_seqs = False + orig = psp.cu_seqlens_q.clone() + p_tok, _, _, _, p, mask = pad_sequence_for_thd( + torch.ones(1, total_T, device="cuda"), + None, + None, + None, + psp, + alignment=64, + pad_by_appending_dummy_seq=False, + ) + assert p_tok.shape == (1, 128) + assert torch.equal(p.cu_seqlens_q, orig) + assert torch.equal(p.cu_seqlens_q_padded, orig) + assert p.max_seqlen_q == max(seqlens) + assert p.max_seqlen_kv == max(seqlens) + assert p.pad_between_seqs is False + assert mask.shape == (1, 128) + assert not mask[0, :total_T].any() and mask[0, total_T:].all() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_shapes_and_data_preservation(self): + """Shapes are static; original data intact; padding zero-filled.""" + seqlens, max_seqlen, max_num_seqs = [100, 50, 30], 256, 8 + total_T = sum(seqlens) + tokens = torch.arange(total_T, device="cuda").unsqueeze(0).float() + p_tok, p_lab, p_loss, p_pos, p_params, p_mask = pad_sequence_for_thd( + tokens, + tokens.clone(), + torch.ones(1, total_T, device="cuda"), + torch.arange(total_T, device="cuda").unsqueeze(0), + _make_psp(seqlens), + target_len=max_seqlen, + max_num_seqs=max_num_seqs, + ) + for t in (p_tok, p_lab, p_loss, p_pos): + assert t.shape == (1, max_seqlen) + for cu in ( + p_params.cu_seqlens_q, + p_params.cu_seqlens_kv, + p_params.cu_seqlens_q_padded, + p_params.cu_seqlens_kv_padded, + ): + assert cu.shape[0] == max_num_seqs + 1 + expected_cu = torch.tensor( + [0, 100, 150, 180, 256, 256, 256, 256, 256], dtype=torch.int32, device="cuda" + ) + assert torch.equal(p_params.cu_seqlens_q, expected_cu) + assert torch.equal(p_params.cu_seqlens_kv, expected_cu) + assert torch.equal(p_params.cu_seqlens_q_padded, expected_cu) + assert torch.equal(p_params.cu_seqlens_kv_padded, expected_cu) + assert p_params.max_seqlen_q == max_seqlen + assert p_params.max_seqlen_kv == max_seqlen + assert p_params.pad_between_seqs is False + assert p_mask.shape == (1, max_seqlen) and p_mask.dtype == torch.bool + assert torch.equal(p_tok[0, :total_T], tokens[0]) + assert (p_tok[0, total_T:] == 0).all() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_eager_pad_to_max_adds_dummy_padding_sequence(self): + """Eager pad-to-max represents the tail as an independent dummy sequence.""" + seqlens, total_T, target_len = [50, 30], 80, 8192 + psp = _make_psp(seqlens) + orig_cu = psp.cu_seqlens_q.clone() + alignment, pad_target_len, max_num_seqs = get_thd_padding_kwargs( + pad_packed_seq_alignment="max", + max_seqlen_per_dp_cp_rank=target_len, + thd_max_packed_sequences=32, + cuda_graph_static=False, + ) + + p_tok, _, _, _, p_params, p_mask = pad_sequence_for_thd( + torch.ones(1, total_T, device="cuda"), + None, + None, + None, + psp, + alignment=alignment, + target_len=pad_target_len, + max_num_seqs=max_num_seqs, + ) + + assert p_tok.shape == (1, target_len) + expected = torch.cat( + (orig_cu, torch.tensor([target_len], dtype=orig_cu.dtype, device=orig_cu.device)) + ) + assert torch.equal(p_params.cu_seqlens_q, expected) + assert torch.equal(p_params.cu_seqlens_q_padded, expected) + assert p_params.cu_seqlens_q.shape[0] == orig_cu.shape[0] + 1 + assert p_params.max_seqlen_q == target_len - total_T + assert p_params.max_seqlen_kv == target_len - total_T + assert p_params.total_tokens == target_len + assert p_params.pad_between_seqs is False + assert p_mask.shape == (1, target_len) + assert not p_mask[0, :total_T].any() + assert p_mask[0, total_T:].all() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_padding_mask_boundary(self): + """False at real positions, True at padding (MoE aux-loss contract).""" + seqlens, total_T, max_seqlen = [60, 40], 100, 128 + _, _, _, _, _, m = pad_sequence_for_thd( + torch.ones(1, total_T, device="cuda"), + None, + None, + None, + _make_psp(seqlens), + target_len=max_seqlen, + max_num_seqs=4, + ) + assert not m[0, :total_T].any() and m[0, total_T:].all() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_padding_mask_preserves_existing_padding(self): + """Existing THD padding and appended tail padding are merged in one helper.""" + seqlens, total_T, max_seqlen = [4, 4], 8, 10 + padding_mask = torch.tensor( + [[False, False, False, True, False, False, True, True]], dtype=torch.bool, device="cuda" + ) + + _, _, _, _, _, m = pad_sequence_for_thd( + torch.ones(1, total_T, device="cuda"), + None, + None, + None, + _make_psp(seqlens), + target_len=max_seqlen, + max_num_seqs=4, + padding_mask=padding_mask, + ) + + assert torch.equal( + m, + torch.tensor( + [[False, False, False, True, False, False, True, True, True, True]], + dtype=torch.bool, + device="cuda", + ), + ) + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_cu_seqlens_fill_value(self): + """Static cu padding repeats dummy valid/padded cumulative values.""" + seqlens, total_T = [50, 30], 80 + _, _, _, _, p, _ = pad_sequence_for_thd( + torch.ones(1, total_T, device="cuda"), + None, + None, + None, + _make_psp(seqlens), + target_len=128, + max_num_seqs=32, + ) + assert p.cu_seqlens_q[0] == 0 and p.cu_seqlens_q[2] == 80 + assert (p.cu_seqlens_q[3:] == 128).all() + assert p.cu_seqlens_q_padded[0] == 0 and p.cu_seqlens_q_padded[2] == 80 + assert (p.cu_seqlens_q_padded[3:] == 128).all() + assert p.pad_between_seqs is False + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_none_inputs(self): + """Non-pre_process PP: mask from cu_seqlens when all tensors None.""" + seqlens, total_T, max_seqlen = [50, 30], 80, 128 + _, _, _, _, _, mask = pad_sequence_for_thd( + None, None, None, None, _make_psp(seqlens), target_len=max_seqlen, max_num_seqs=4 + ) + assert mask.shape == (1, max_seqlen) + assert not mask[0, :total_T].any() and mask[0, total_T:].all() + + +# ============================================================================= +# 2. PackedSeqParams decompose / reconstruct +# ============================================================================= + + +class TestDecomposeReconstruct: + + def setup_method(self): + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + + def teardown_method(self): + Utils.destroy_model_parallel() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_round_trip(self): + """Decompose then reconstruct preserves cu_seqlens values.""" + psp = _make_psp([100, 50, 30]) + orig = { + k: getattr(psp, k).clone() + for k in ( + 'cu_seqlens_q', + 'cu_seqlens_kv', + 'cu_seqlens_q_padded', + 'cu_seqlens_kv_padded', + ) + } + layer = _build_layer(256, 4, 4, 1024, 128, 8) + kw = {'packed_seq_params': psp, 'other': 'kept'} + TransformerLayer._decompose_packed_seq_params_to_kwargs(kw) + assert 'packed_seq_params' not in kw and 'cu_seqlens_q' in kw + layer._reconstruct_packed_seq_params_from_kwargs(kw) + r = kw['packed_seq_params'] + assert r.qkv_format == 'thd' and r.max_seqlen_q == 128 + assert r.pad_between_seqs is False + for k, v in orig.items(): + assert torch.equal(getattr(r, k), v) + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_noop_without_packed_seq_params(self): + """No-ops on non-THD kwargs (SBHD path).""" + layer = _build_layer(256, 4, 4, 1024, 128, 8) + kw = {'hidden_states': torch.randn(10, 1, 256, device="cuda")} + keys = set(kw.keys()) + TransformerLayer._decompose_packed_seq_params_to_kwargs(kw) + assert set(kw.keys()) == keys + layer._reconstruct_packed_seq_params_from_kwargs(kw) + assert set(kw.keys()) == keys + + +class TestStaticInputs: + + def setup_method(self): + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + + def teardown_method(self): + Utils.destroy_model_parallel() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_static_padding_mask_is_unmasked_for_capture(self): + """Capture-time padding_mask must not mark every static token as padding.""" + layer = _build_layer(256, 4, 4, 1024, 128, 8) + layer.config.sequence_packing_scheduler = "dp_balanced" + layer.config.cuda_graph_impl = "transformer_engine" + + static_inputs = layer.get_layer_static_inputs(seq_length=128, micro_batch_size=1) + + assert static_inputs["padding_mask"].shape == (1, 128) + assert not static_inputs["padding_mask"].any() + + +class TestDynamicMicrobatchSlots: + + @pytest.mark.internal + def test_pp2_slots_track_max_outstanding_microbatches(self): + from megatron.core.transformer.cuda_graphs import TECudaGraphHelper + + order = [1, 1, -1, 1, -1, 1, -1, -1] + + assert TECudaGraphHelper._get_required_num_microbatch_slots_from_order(order, 1) == 2 + + @pytest.mark.internal + def test_vpp_slots_track_each_chunk_liveness(self): + from megatron.core.transformer.cuda_graphs import TECudaGraphHelper + + order = [1, 1, 1, 2, 2, 2, -2, 1, -2, 1, -2, 2, -1, 2, -1, -1, -2, -2, -1, -1] + + assert TECudaGraphHelper._get_required_num_microbatch_slots_from_order(order, 2) == 5 + + @pytest.mark.internal + def test_dp_balanced_thd_capture_upper_bound_uses_max_sequence_length(self): + from megatron.core.transformer.cuda_graphs import TECudaGraphHelper + + assert ( + TECudaGraphHelper._get_dp_balanced_thd_max_num_microbatches( + global_batch_size=64, + dp_size=1, + cp_size=1, + max_seqlen_per_dp_cp_rank=4096, + max_sequence_length=4096, + max_num_seqs=8, + ) + == 64 + ) + assert ( + TECudaGraphHelper._get_dp_balanced_thd_max_num_microbatches( + global_batch_size=64, + dp_size=1, + cp_size=2, + max_seqlen_per_dp_cp_rank=4096, + max_sequence_length=4096, + max_num_seqs=8, + ) + == 32 + ) + + @pytest.mark.internal + def test_dp_balanced_thd_capture_upper_bound_aligns_vpp_groups(self): + from megatron.core.transformer.cuda_graphs import TECudaGraphHelper + + assert ( + TECudaGraphHelper._get_dp_balanced_thd_max_num_microbatches( + global_batch_size=18, + dp_size=1, + cp_size=1, + max_seqlen_per_dp_cp_rank=4096, + max_sequence_length=2048, + microbatch_group_size_per_vp_stage=8, + max_num_seqs=8, + ) + == 16 + ) + + +# ============================================================================= +# 3. E2E no-graph vs graph bitwise loss/grad_norm match +# Subprocess-launches `torchrun pretrain_gpt.py` -- same recipe as +# test_moonlight_qwen3_bitwise.sh -- and asserts the per-iteration +# metric strings are byte-identical between the two runs. +# ============================================================================= + +# Common args shared across both models. +_REPO_ROOT = Path(__file__).resolve().parents[3] + +_VARLEN_JSON = ( + '{"mode":"distribution","type":"lognormal",' + '"format":"thd","min_seq_len":512,"max_seq_len":4096,' + '"mean_seq_len":3072,"lognormal_sigma":1.1}' +) + +_QWEN3_VARLEN_JSON = ( + '{"mode":"distribution","type":"lognormal",' + '"format":"thd","min_seq_len":128,"max_seq_len":1024,' + '"mean_seq_len":512,"lognormal_sigma":0.8}' +) + +_TRAIN_ITERS = 5 + +_COMMON_ARGS = [ + "--seq-length", + "4096", + "--max-position-embeddings", + "8192", + "--micro-batch-size", + "1", + "--global-batch-size", + "64", + "--train-iters", + str(_TRAIN_ITERS), + "--lr", + "1e-5", + "--min-lr", + "1e-6", + "--lr-decay-style", + "cosine", + "--lr-warmup-iters", + "1", + "--weight-decay", + "0.01", + "--clip-grad", + "1.0", + "--seed", + "1234", + "--te-rng-tracker", + "--bf16", + "--tensor-model-parallel-size", + "2", + "--pipeline-model-parallel-size", + "2", + "--context-parallel-size", + "2", + "--swiglu", + "--disable-bias-linear", + "--sequence-parallel", + "--use-varlen-dataset", + "--mock-data", + "--tokenizer-type", + "NullTokenizer", + "--varlen-mock-dataset-config-json", + _VARLEN_JSON, + "--sequence-packing-scheduler", + "dp_balanced", + "--max-seqlen-per-dp-cp-rank", + "4096", + "--pad-packed-seq-alignment", + "max", + "--no-pad-packed-seq-by-appending-dummy-seq", + "--calculate-per-token-loss", + "--transformer-impl", + "transformer_engine", + "--attention-dropout", + "0", + "--hidden-dropout", + "0", + "--no-bias-swiglu-fusion", + "--no-gradient-accumulation-fusion", + "--no-save-optim", + "--no-save-rng", + "--save-interval", + "999999", + "--eval-interval", + "999999", + "--eval-iters", + "1", + "--log-interval", + "1", + "--no-check-for-nan-in-loss-and-grad", + "--deterministic-mode", + "--thd-max-packed-sequences", + "8", +] + + +def _with_arg_replacements(args, replacements): + args = list(args) + for name, value in replacements.items(): + idx = args.index(name) + args[idx + 1] = value + return args + + +_QWEN3_COMMON_ARGS = _with_arg_replacements( + _COMMON_ARGS, + { + "--seq-length": "1024", + "--varlen-mock-dataset-config-json": _QWEN3_VARLEN_JSON, + "--max-seqlen-per-dp-cp-rank": "512", + }, +) + + +_MOONLIGHT_ARGS = _COMMON_ARGS + [ + "--num-layers", + "27", + "--hidden-size", + "2048", + "--ffn-hidden-size", + "11264", + "--num-attention-heads", + "16", + "--decoder-first-pipeline-num-layers", + "13", + "--decoder-last-pipeline-num-layers", + "14", + "--expert-model-parallel-size", + "4", + "--expert-tensor-parallel-size", + "1", + "--multi-latent-attention", + "--kv-lora-rank", + "512", + "--qk-head-dim", + "128", + "--qk-pos-emb-head-dim", + "64", + "--v-head-dim", + "128", + "--num-experts", + "64", + "--moe-ffn-hidden-size", + "1408", + "--moe-router-topk", + "6", + "--moe-shared-expert-intermediate-size", + "2816", + "--moe-layer-freq", + "([0]+[1]*26)", + "--moe-token-dispatcher-type", + "flex", + "--moe-flex-dispatcher-backend", + "hybridep", + "--moe-router-fusion", + "--moe-router-score-function", + "sigmoid", + "--moe-router-topk-scaling-factor", + "2.446", + "--moe-router-load-balancing-type", + "aux_loss", + "--moe-aux-loss-coeff", + "0.001", + "--normalization", + "RMSNorm", + "--norm-epsilon", + "1e-5", + "--rotary-base", + "50000", + "--vocab-size", + "163840", +] + +_QWEN3_ARGS = _QWEN3_COMMON_ARGS + [ + "--num-layers", + "36", + "--hidden-size", + "4096", + "--ffn-hidden-size", + "12288", + "--num-attention-heads", + "32", + "--group-query-attention", + "--num-query-groups", + "8", + "--max-position-embeddings", + "40960", + "--normalization", + "RMSNorm", + "--norm-epsilon", + "1e-6", + "--rotary-base", + "1000000", + "--untie-embeddings-and-output-weights", + "--vocab-size", + "151936", + "--moe-token-dispatcher-type", + "flex", + "--moe-flex-dispatcher-backend", + "hybridep", +] + +_ATTN_CUDA_GRAPH_ARGS = [ + "--cuda-graph-impl", + "transformer_engine", + "--cuda-graph-dynamic-microbatches", + "--cuda-graph-modules", + "attn", +] + +_MOE_CUDA_GRAPH_ARGS = _ATTN_CUDA_GRAPH_ARGS + ["moe_preprocess", "moe_router"] + + +def _get_available_port(preferred): + """Return preferred if free, otherwise ask the OS for an available localhost port.""" + for port in (preferred, 0): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + try: + sock.bind(("localhost", port)) + except OSError: + continue + return sock.getsockname()[1] + raise RuntimeError("Could not find an available localhost port") + + +def _run_pretrain(model_args, cuda_graph_args, master_port): + """Subprocess-launch `torchrun pretrain_gpt.py` once and capture stdout.""" + env = os.environ.copy() + env["PYTHONPATH"] = str(_REPO_ROOT) + ":" + env.get("PYTHONPATH", "") + env["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" + env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + env["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + env["NCCL_ALGO"] = "^NVLS" + # Strip any inherited torchrun env so this subprocess starts a fresh group. + for k in list(env.keys()): + if k.startswith( + ( + "TORCHELASTIC_", + "MASTER_", + "RANK", + "LOCAL_RANK", + "WORLD_SIZE", + "GROUP_RANK", + "LOCAL_WORLD_SIZE", + ) + ): + env.pop(k, None) + # Clear pytest-conftest env vars that disable TE attention backends + # (set by tests/unit_tests/conftest.py::set_env). Pretrain needs at + # least one of fused/flash attention to build the model. + env.pop("NVTE_FLASH_ATTN", None) + env.pop("NVTE_FUSED_ATTN", None) + + cmd = ( + [ + "torchrun", + "--nproc_per_node", + "8", + "--nnodes", + "1", + "--master_addr", + "localhost", + "--master_port", + str(_get_available_port(master_port)), + "pretrain_gpt.py", + ] + + model_args + + cuda_graph_args + ) + + result = subprocess.run( + cmd, cwd=_REPO_ROOT, env=env, capture_output=True, text=True, timeout=900 + ) + return result + + +_ITER_START_RE = re.compile(r"iteration\s+(\d+)/\s*\d+ \|") + + +def _extract_metrics(stdout): + """Extract deterministic per-iteration fields from a training log. + + Captured torchrun stdout interleaves writes from multiple ranks at the byte + level (no newline between rank-0's iter line and rank-7's "Number of + parameters" line, e.g.). So we cannot rely on full-line matching: we locate + each `iteration N/M |` marker and pull the deterministic fields by name + from a small window after it. Wall-clock `elapsed time per iteration` + is intentionally excluded. + """ + results = [] + for m in _ITER_START_RE.finditer(stdout): + window = stdout[m.start() : m.start() + 800] + lr = re.search(r"learning rate:\s*(\S+)", window) + lm_loss = re.search(r"lm loss:\s*(\S+)", window) + grad_norm = re.search(r"grad norm:\s*(\S+)", window) + if not (lr and lm_loss and grad_norm): + continue + parts = [f"iter={m.group(1)}", f"lr={lr.group(1)}", f"lm_loss={lm_loss.group(1)}"] + parts.append(f"grad_norm={grad_norm.group(1)}") + results.append(" | ".join(parts)) + return results + + +@pytest.mark.internal +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(torch.cuda.device_count() < 8, reason="requires 8 GPUs") +@pytest.mark.parametrize( + "model_name,model_args,cuda_graph_args,base_port", + [ + ("moonlight", _MOONLIGHT_ARGS, _MOE_CUDA_GRAPH_ARGS, 29660), + ("qwen3", _QWEN3_ARGS, _ATTN_CUDA_GRAPH_ARGS, 29662), + ], +) +class TestE2EBitwise: + """End-to-end bitwise comparison: pretrain_gpt.py noGraph vs cudaGraph. + + Each test launches `torchrun pretrain_gpt.py` twice -- once without CUDA + graphs and once with `cuda_graph_impl=transformer_engine` -- using the same + model/test settings as test_moonlight_qwen3_bitwise.sh. Moonlight covers + attn/moe_preprocess/moe_router graphs with router fusion; Qwen3 covers attn + graphs because this test's Qwen3 recipe is dense. + Asserts the per-iteration `lm loss / grad norm` lines are byte-identical. + + Slow (~5 min per model). Marked `internal` so CI can opt-in. + """ + + def test_no_graph_vs_graph(self, model_name, model_args, cuda_graph_args, base_port): + # No graph baseline. + r1 = _run_pretrain(model_args, cuda_graph_args=[], master_port=base_port) + assert r1.returncode == 0, ( + f"[{model_name}] noGraph pretrain failed (rc={r1.returncode})\n" + f"--- stdout (tail) ---\n{r1.stdout[-4000:]}\n" + f"--- stderr (tail) ---\n{r1.stderr[-2000:]}" + ) + metrics_eager = _extract_metrics(r1.stdout) + assert len(metrics_eager) == _TRAIN_ITERS, ( + f"[{model_name}] noGraph: expected {_TRAIN_ITERS} metric lines, " + f"got {len(metrics_eager)}\n" + f"--- stdout (tail) ---\n{r1.stdout[-2000:]}" + ) + + # CUDA graph capture. + r2 = _run_pretrain(model_args, cuda_graph_args=cuda_graph_args, master_port=base_port + 1) + assert r2.returncode == 0, ( + f"[{model_name}] cudaGraph pretrain failed (rc={r2.returncode})\n" + f"--- stdout (tail) ---\n{r2.stdout[-4000:]}\n" + f"--- stderr (tail) ---\n{r2.stderr[-2000:]}" + ) + metrics_graph = _extract_metrics(r2.stdout) + assert len(metrics_graph) == _TRAIN_ITERS, ( + f"[{model_name}] cudaGraph: expected {_TRAIN_ITERS} metric lines, " + f"got {len(metrics_graph)}\n" + f"--- stdout (tail) ---\n{r2.stdout[-2000:]}" + ) + + # Bitwise compare per iteration. + for i, (a, b) in enumerate(zip(metrics_eager, metrics_graph)): + assert a == b, f"[{model_name}] iter {i+1} differs:\n" f" eager: {a}\n" f" graph: {b}"