diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 5e3c1084792..0b2039c60eb 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -19,12 +19,58 @@ next_hdp_group, reroute_samples_to_dcp_ranks, ) -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.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank +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 BasePackingScheduler: """Base class for sequence packing schedulers.""" @@ -34,6 +80,7 @@ def __init__( cp_size: int, dp_size: int, microbatch_group_size_per_vp_stage: Optional[int], + max_num_seqs: Optional[int] = None, ): """ Args: @@ -42,11 +89,15 @@ def __init__( 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.""" @@ -118,7 +169,9 @@ def get_groups_and_subsamples(self, sample_id_seqlens): single_microbatch = [] for i in range(len(sample_id_seqlens)): - if sum_seqlen + sample_id_seqlens[i][1] <= self.max_seq_len_all_ranks: + 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: @@ -412,6 +465,35 @@ def get_groups_and_subsamples(self, sample_id_seqlens): } +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 ): @@ -454,6 +536,12 @@ def wrap_data_iterator( if scheduler_type == 'default_dynamic_cp': scheduler_kwargs['min_cp_size'] = config.min_dynamic_context_parallel_size + scheduler_max_num_seqs = ( + _get_scheduler_max_real_num_seqs(config) + if scheduler_type == 'dp_balanced' + else getattr(config, 'thd_max_packed_sequences', None) + ) + scheduler = scheduler_map[scheduler_type]( config.max_seqlen_per_dp_cp_rank, cp_size, @@ -463,6 +551,7 @@ def wrap_data_iterator( if config.virtual_pipeline_model_parallel_size is None else config.microbatch_group_size_per_vp_stage ), + max_num_seqs=scheduler_max_num_seqs, **scheduler_kwargs, ) @@ -490,6 +579,7 @@ def get_batch_on_this_rank_for_sequence_packing( vp_stage: Optional[int] = None, dynamic_cp: bool = False, pg_collection: Optional[ProcessGroupCollection] = None, + config=None, ): """ Get a batch of data for sequence packing. @@ -497,8 +587,11 @@ def get_batch_on_this_rank_for_sequence_packing( 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) + tuple of (tokens, labels, loss_mask, attention_mask, position_ids, + packed_seq_params, padding_mask) """ if pg_collection is None: @@ -551,10 +644,21 @@ def get_batch_on_this_rank_for_sequence_packing( group_size=local_cp_size_val ) - # Partition tokens, position_ids, labels, loss_mask for context parallel. - # Only TP rank 0 on stages that have data (first/last PP stage or MTP stage) needs this. - if is_tp_rank_0 and (is_first_or_last_stage or mtp_on_this_rank): - get_cp_slice_for_thd(batch, cp_group) + # 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 sequence tensors for context parallelism. Padding mask is needed + # on every PP stage, while data tensors are only needed on first/last/MTP stages. + if is_tp_rank_0: + 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']) + get_cp_slice_for_thd(batch, cp_group, keys=cp_slice_keys) # Broadcast cu_seqlens_size because we need it to create placeholder for cu_seqlens and # cu_seqlens_padded for non TP 0 ranks. @@ -565,23 +669,19 @@ def get_batch_on_this_rank_for_sequence_packing( broadcast_tensor(cu_seqlen_size, tp_src_rank, tp_group) cu_seqlen_size = cu_seqlen_size.item() - # Broadcast total_tokens because we need it to create placeholder for tokens, position_ids, - # labels, loss_mask for non TP 0 ranks. Only first stage, last stage, - # and stage with mtp need this. - - if is_first_or_last_stage or mtp_on_this_rank: - if is_tp_rank_0: - # Use whichever data field is available (first stage has tokens, last has labels). - # Avoid `tokens or labels`: PyTorch tensors raise on truthiness when they have - # more than one element ("Boolean value of Tensor ... is ambiguous"). - _data_field = batch.get('tokens') - if _data_field is None: - _data_field = batch.get('labels') - total_tokens = torch.tensor(_data_field.size(0), dtype=torch.int32, device=dev) - else: - total_tokens = torch.empty(1, dtype=torch.int32, device=dev) - broadcast_tensor(total_tokens, tp_src_rank, tp_group) - total_tokens = total_tokens.item() + # 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: + # Under VPP, the last PP stage has labels but no tokens, so derive + # total_tokens from cu_seqlens_padded, which is present on every + # stage. cu_seqlens_padded keeps the pre-CP packed length; divide + # by cp_size to match the already CP-sliced sequence tensors. + cp_world = cp_group.size() + total_tokens = (batch['cu_seqlens_padded'][-1].to(torch.int32) // cp_world).reshape(1) + else: + total_tokens = torch.empty(1, dtype=torch.int32, device=dev) + broadcast_tensor(total_tokens, tp_src_rank, tp_group) + total_tokens = total_tokens.item() # Step1: Prepare "tokens", "position_ids" for first stage and stage with mtp on all TP ranks. if is_first_stage or mtp_on_this_rank: @@ -613,7 +713,14 @@ def get_batch_on_this_rank_for_sequence_packing( batch['labels'] = None batch['loss_mask'] = None - # Step3: Prepare "cu_seqlens", "cu_seqlens_padded", "max_seqlen" on all ranks. + # 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 @@ -629,7 +736,7 @@ def get_batch_on_this_rank_for_sequence_packing( batch['cu_seqlens_padded'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev) batch['max_seqlen'] = torch.empty(1, dtype=torch.int32, device=dev) - # Step4: Prepare "local_cp_size" if dynamic context parallel is enabled. + # Step5: Prepare "local_cp_size" if dynamic context parallel is enabled. if dynamic_cp: if is_tp_rank_0: if type(batch['local_cp_size']) == int: @@ -649,6 +756,7 @@ def get_batch_on_this_rank_for_sequence_packing( 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) @@ -659,6 +767,7 @@ def get_batch_on_this_rank_for_sequence_packing( 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() @@ -669,9 +778,8 @@ def get_batch_on_this_rank_for_sequence_packing( else None ) - # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as cu_seqlens to - # get the correct result. - # TODO: Revert this workaround once TE fixes the issue. + # 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, @@ -682,10 +790,41 @@ def get_batch_on_this_rank_for_sequence_packing( max_seqlen_kv=max_seqlen, local_cp_size=local_cp_size, cp_group=cp_group, + 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 + return tokens, labels, loss_mask, None, position_ids, packed_seq_params, padding_mask class HybridCPDataLoaderWrapper: diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 51be6282ffe..5f8bb6b3e01 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -3,7 +3,7 @@ from collections import deque from functools import lru_cache from math import ceil, log2 -from typing import Callable, Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Sequence, Tuple import torch @@ -11,7 +11,7 @@ from megatron.core.rerun_state_machine import RerunDataIterator -def get_cp_slice_for_thd(batch, cp_group): +def get_cp_slice_for_thd(batch, cp_group, keys: Optional[Sequence[str]] = None): """Partition sequence data for context parallelism in THD format. Uses TE's THD partitioned indices to split the packed sequence across CP ranks. @@ -20,25 +20,24 @@ def get_cp_slice_for_thd(batch, cp_group): Args: batch: Dict with packed sequence data. cp_group: Context parallel process group. + keys: Sequence data keys to slice. Defaults to the original THD data tensors. """ cp_size = cp_group.size() if cp_size <= 1: return cp_rank = cp_group.rank() - # Use whichever data field is available to determine total_tokens - for _key in ['tokens', 'labels', 'loss_mask', 'position_ids']: - if _key in batch and batch[_key] is not None: - total_tokens = batch[_key].size(0) - break - else: - raise ValueError("Cannot determine total_tokens: no data field found in batch") - # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as - # cu_seqlens to get the correct result. - # TODO: Revert this workaround once TE fixes the issue. + # Partition with padded cumulative lengths so CP slices match the THD + # sequence boundaries consumed by attention kernels. cu_seqlens = batch["cu_seqlens_padded"] + # Use cu_seqlens_padded[-1] for total_tokens instead of batch['tokens'].size(0): + # under VPP, the last PP stage has labels/loss_mask but no tokens, so + # batch['tokens'] is None on that stage. cu_seqlens_padded is always populated. + total_tokens = int(cu_seqlens[-1].item()) index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) - for key in ['tokens', 'position_ids', 'labels', 'loss_mask']: - if key in batch: + if keys is None: + keys = ('tokens', 'position_ids', 'labels', 'loss_mask') + for key in keys: + if key in batch and batch[key] is not None: batch[key] = batch[key].index_select(0, index) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index c3e8b4f5e1c..1aa460f8fab 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1686,6 +1686,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 d5e8f721297..ba2f4715fc7 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -2,13 +2,28 @@ import warnings from dataclasses import dataclass, field -from typing import Callable, ContextManager, Literal, Optional +from typing import Callable, ContextManager, Literal, Optional, Union import torch from megatron.core.utils import experimental_api +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 @experimental_api class ModelParallelConfig: @@ -87,6 +102,31 @@ class ModelParallelConfig: default_dynamic_cp: Dynamic-CP scheduler for packed sequence balancing. """ + 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.""" @@ -469,6 +509,28 @@ def __post_init__(self): f"got {self.min_dynamic_context_parallel_size}" ) + 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 c97f738771b..ad790e5c5ad 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -216,18 +216,22 @@ def _apply_rotary_pos_emb_thd( mla_output_remove_interleaving: bool = False, 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( @@ -240,54 +244,70 @@ 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() - - # Handle two different frequency tensor formats: - # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains all positions across all sequences - # -> Use offset-based mapping for exact positional correspondence - # 2. Otherwise: freqs contains only max sequence length positions - # -> Use traditional mapping without offsets (map first :seqlen part) - if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens[-1]: - # CASE 1: Exact mapping with offsets - # Build packed freqs in one pass, then apply once to the whole packed tensor - sequence_splits = torch.split(t, seqlens) - freq_slices = [] - 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() - freq_slices.append( - _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) - ) - freqs_packed = torch.cat(freq_slices, dim=0) + total_tokens = t.shape[0] + device = t.device - return _apply_rotary_pos_emb_bshd( - t.unsqueeze(1), - freqs_packed, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - inverse=inverse, - mla_output_remove_interleaving=mla_output_remove_interleaving, - ).squeeze(1) - else: - # CASE 2: Traditional mapping without offsets - # Build packed freqs for all sequences using the standard mapping, then apply once - sequence_splits = torch.split(t, seqlens) - freqs_packed = torch.cat( - [_get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs) for x in sequence_splits], - dim=0, + 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) - return _apply_rotary_pos_emb_bshd( - t.unsqueeze(1), - freqs_packed, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - inverse=inverse, - mla_output_remove_interleaving=mla_output_remove_interleaving, - ).squeeze(1) + 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, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ).squeeze(1) def apply_rotary_pos_emb( @@ -300,6 +320,7 @@ def apply_rotary_pos_emb( mla_rotary_interleaved: bool = False, inverse: bool = False, mla_output_remove_interleaving: bool = False, + max_seqlen: Optional[int] = None, ): """ Reroute to the appropriate apply_rotary_pos_emb function depending on @@ -375,6 +396,7 @@ def apply_rotary_pos_emb( cp_group=cp_group, inverse=inverse, mla_output_remove_interleaving=mla_output_remove_interleaving, + max_seqlen=max_seqlen, ) diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index b1b4275fee1..3095b1b8464 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 @@ -24,6 +26,7 @@ class PackedSeqParams: cp_group: dist.ProcessGroup = None total_tokens: int = None seq_idx: Tensor = None + pad_between_seqs: Optional[bool] = None def __post_init__(self): """Pre-compute seq_idx for Mamba mixer CUDA graph compatibility. @@ -78,3 +81,410 @@ def resolve_cp_group( if packed_seq_params is not None and packed_seq_params.cp_group is not None: return packed_seq_params.cp_group return static_cp_group + + +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, + 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/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 3e61eb12a5f..e712122ef69 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -68,10 +68,10 @@ rearrange = None try: - from flash_attn_3.flash_attn_interface import _flash_attn_forward from flash_attn_3.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) + from flashattn_hopper.flash_attn_interface import _flash_attn_forward HAVE_FA3 = True except ImportError as e: @@ -1324,8 +1324,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: @@ -1338,6 +1341,7 @@ def forward( cu_seqlens=cu_seqlens_q, mscale=_yarn_get_concentration_factor_from_config(self.config), cp_group=self.pg_collection.cp, + max_seqlen=rope_max_seqlen_q, ) else: query = inference_context.apply_rotary_emb_query( @@ -1351,6 +1355,7 @@ def forward( cu_seqlens=cu_seqlens_kv, mscale=_yarn_get_concentration_factor_from_config(self.config), 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 3ec40c3b2d3..d59d1fbf5b0 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 ( @@ -1798,7 +1799,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 ( @@ -1814,12 +1822,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 @@ -2024,6 +2034,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 @@ -2195,6 +2211,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. @@ -2207,26 +2377,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 5c591f7b534..1b7ad6982a9 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -406,8 +406,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 @@ -605,6 +608,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( @@ -615,6 +619,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/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index a8f5a65a185..a6758a67070 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -326,10 +326,12 @@ def forward( if packed_seq_params.cu_seqlens_kv_padded is not None else packed_seq_params.cu_seqlens_kv ) - rope_seqlen = cu_seqlens_kv + rope_seqlen = packed_seq_params.max_seqlen_kv + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv else: cu_seqlens_kv = None rope_seqlen = seq_len + rope_max_seqlen_kv = None # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's # concentration factor (mscale) is NOT part of the DSv4 model contract -- # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0. @@ -379,6 +381,7 @@ def forward( mla_rotary_interleaved=True, inverse=True, mla_output_remove_interleaving=True, + max_seqlen=rope_max_seqlen_kv, ) core_attn_out = torch.cat([content_part, rot_part], dim=-1) core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), -1) @@ -569,8 +572,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 @@ -679,6 +685,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, mla_output_remove_interleaving=True, + max_seqlen=rope_max_seqlen_q, ) # query: [num_tokens, n, (qk_head_dim + v_head_dim)] query = torch.cat([q_no_pe, q_pos_emb], dim=-1) @@ -696,6 +703,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, mla_output_remove_interleaving=True, + max_seqlen=rope_max_seqlen_kv, ) # Single head: key = value = [num_tokens, 1, v_head_dim] diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index d19f6d094c0..c5211e3e6d6 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -222,6 +222,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. @@ -229,26 +236,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 359b7c4a4bd..59684a34b0d 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -12,6 +12,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_logging import get_moe_overload_factor_tracker @@ -452,13 +453,16 @@ def route( hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None, input_ids: 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, input_ids) + probs, routing_map = apply_module(self.router)( + hidden_states, padding_mask, input_ids, packed_seq_params + ) return probs, routing_map @maybe_skip_or_early_return_by_cudagraph("preprocess") @@ -667,6 +671,7 @@ def forward( intermediate_tensors=None, padding_mask: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """Forward pass for the MoE layer. @@ -678,9 +683,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. input_ids (torch.Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. Defaults to None. Returns: @@ -703,6 +708,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() @@ -716,7 +744,9 @@ def custom_forward(hidden_states, intermediate_tensors=None, padding_mask=None): self._overload_log_num_local_tokens = ( self._num_token_rows_from_moe_hidden_states(hidden_states) ) - probs, routing_map = self.route(hidden_states, padding_mask, input_ids) + probs, routing_map = self.route( + hidden_states, padding_mask, input_ids, 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 44675062d42..087d32c10e0 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -2,7 +2,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 @@ -58,7 +58,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, @@ -108,7 +108,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. @@ -226,22 +226,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 580c8a5a650..548a136d543 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""" @@ -320,16 +336,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, ) @@ -350,8 +369,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 @@ -363,6 +387,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. @@ -378,10 +403,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, ) @@ -405,8 +432,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 @@ -422,7 +454,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, @@ -453,18 +486,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. @@ -473,7 +539,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). @@ -507,44 +576,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) @@ -556,48 +628,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: @@ -609,7 +678,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 @@ -701,14 +770,15 @@ def routing( logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None, input_ids: 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. input_ids (torch.Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. Defaults to None. @@ -776,6 +846,7 @@ def routing( 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, @@ -784,6 +855,7 @@ def routing( seq_length, bsz, with_padding_mask=padding_mask is not None, + packed_seq_params=packed_seq_params, ) probs = self._apply_global_aux_loss( probs, @@ -808,15 +880,16 @@ def forward( input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None, input_ids: 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. input_ids (torch.Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. Defaults to None. """ @@ -836,7 +909,12 @@ def forward( logits, self.config.moe_router_force_biased, self.layer_number ) - probs, routing_map = self.routing(logits, padding_mask=padding_mask, input_ids=input_ids) + probs, routing_map = self.routing( + logits, + padding_mask=padding_mask, + input_ids=input_ids, + packed_seq_params=packed_seq_params, + ) return probs, routing_map @@ -945,6 +1023,7 @@ def forward( input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, ): """Simplified forward pass for inference - returns dense tensors only. @@ -960,6 +1039,6 @@ def forward( """ if not InferenceMode.is_active(): - return super().forward(input, padding_mask, input_ids) + return super().forward(input, padding_mask, input_ids, packed_seq_params) return self._forward(input, padding_mask) diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index ed18b6df511..61bd7a6f94c 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1073,15 +1073,26 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): or self.config.moe_hybridep_pad_variable_tokens ) if equalize_thd_token_counts: - # Use the actual tp_ep max so all ranks in the MoE communication - # group pass the same token count to HybridEP. - max_num_tokens_across_ep = torch.tensor( - [num_tokens], device=routing_map.device, dtype=torch.long - ) - torch.distributed.all_reduce( - max_num_tokens_across_ep, op=torch.distributed.ReduceOp.MAX, group=self.group - ) - padded_num_tokens = int(max_num_tokens_across_ep.item()) + if self.config.sequence_packing_scheduler is not None and ( + torch.cuda.is_current_stream_capturing() or torch.compiler.is_compiling() + ): + # CUDA graph path: routing_map has already been padded to a static + # length upstream (CUDA graph + sequence packing implies + # cu_seqlens_q_padded -> max_seqlen_per_dp_cp_rank), so num_tokens + # is identical across the EP communication group. Skip the + # all_reduce + .item() during both dynamo tracing and stream + # capture, and use the local value directly. + padded_num_tokens = num_tokens + else: + # Use the actual tp_ep max so all ranks in the MoE communication + # group pass the same token count to HybridEP. + max_num_tokens_across_ep = torch.tensor( + [num_tokens], device=routing_map.device, dtype=torch.long + ) + torch.distributed.all_reduce( + max_num_tokens_across_ep, op=torch.distributed.ReduceOp.MAX, group=self.group + ) + padded_num_tokens = int(max_num_tokens_across_ep.item()) padded_num_tokens += -padded_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT self._padded_num_tokens = padded_num_tokens diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index a1851d479b2..0523096bea7 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -728,8 +728,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 @@ -941,6 +944,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( @@ -951,6 +955,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 ff9cffa8fef..e1ef8fb344f 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1085,6 +1085,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.""" + #################### # Hyper-Connection Configuration #################### @@ -1435,6 +1455,12 @@ def __post_init__(self): ), f"linear_attention_freq must be set for linear attention." if self.experimental_attention_variant == "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 @@ -2660,6 +2686,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 @@ -2987,6 +3019,21 @@ 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.dynamic_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: diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index b92efdfb8e7..546d7146039 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -862,6 +862,7 @@ def forward(self, *args, **kwargs): kwargs.get("inference_context", None), padding_mask=kwargs.get("padding_mask", None), input_ids=kwargs.get("input_ids", None), + packed_seq_params=kwargs.get("packed_seq_params", None), ) return output, context @@ -885,6 +886,7 @@ def _forward_mlp_output_with_bias( inference_context: BaseInferenceContext | None = None, padding_mask: Tensor | None = None, input_ids: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, ) -> tuple[tuple[Tensor, Tensor | None], Tensor]: """Run pre-MLP norm + MLP/MoE and return the raw output before BDA.""" pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) @@ -927,6 +929,8 @@ def _forward_mlp_output_with_bias( moe_kwargs = {} if self.is_moe_layer and input_ids is not None: moe_kwargs["input_ids"] = input_ids + 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: @@ -990,6 +994,7 @@ def _forward_mlp( inference_context: BaseInferenceContext | None = None, padding_mask: Tensor | None = None, input_ids: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, ) -> Tensor | list[Tensor | None]: """ Perform a forward pass through the feed-forward layer. @@ -1013,6 +1018,7 @@ def _forward_mlp( inference_context=inference_context, padding_mask=padding_mask, input_ids=input_ids, + packed_seq_params=packed_seq_params, ) if ( @@ -1180,15 +1186,46 @@ 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: if not self.config.create_attention_mask_in_dataloader: if self.self_attention.attn_mask_type not in ( AttnMaskType.causal, @@ -1207,7 +1244,7 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): 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) ) @@ -1223,7 +1260,6 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): static_inputs["input_ids"] = torch.zeros( (micro_batch_size, seq_length), dtype=torch.long, device=torch.cuda.current_device() ) - return static_inputs def _get_submodules_under_cudagraphs(self): @@ -1254,6 +1290,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. @@ -1261,7 +1338,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. @@ -1298,7 +1378,10 @@ def _te_cuda_graph_capture(self, *args, **kwargs): ) ): hidden_states = self._forward_mlp( - hidden_states, input_ids=kwargs.get("input_ids", None) + hidden_states, + padding_mask=kwargs.get("padding_mask", None), + input_ids=kwargs.get("input_ids", 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] @@ -1319,8 +1402,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 @@ -1328,6 +1413,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 @@ -1503,7 +1592,19 @@ 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, input_ids=kwargs.get("input_ids", None)) + # 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), + input_ids=kwargs.get("input_ids", None), + packed_seq_params=kwargs.get("packed_seq_params", None), + ) return output, context def _get_te_cuda_graph_replay_args(self, *args, **kwargs): @@ -1848,6 +1949,7 @@ def forward(self, *args, **kwargs): kwargs.get("inference_context", None), padding_mask=kwargs.get("padding_mask", None), input_ids=kwargs.get("input_ids", None), + packed_seq_params=kwargs.get("packed_seq_params", None), mhc_recompute_manager=mhc_recompute_manager, ) return output, context @@ -1964,6 +2066,7 @@ def _forward_mlp( inference_context=None, padding_mask=None, input_ids=None, + packed_seq_params: Optional[PackedSeqParams] = None, mhc_recompute_manager: Optional['CheckpointManager'] = None, ): """Forward MLP with hyper connection pre/post processing.""" @@ -2008,6 +2111,8 @@ def _forward_mlp( moe_kwargs = {} if self.is_moe_layer and input_ids is not None: moe_kwargs['input_ids'] = input_ids + 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: @@ -2206,7 +2311,11 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): ) self.recompute_pre_mlp_layernorm = recompute_pre_mlp_layernorm else: - output = self._forward_mlp(*cuda_graph_output, input_ids=kwargs.get("input_ids", None)) + output = self._forward_mlp( + *cuda_graph_output, + input_ids=kwargs.get("input_ids", None), + packed_seq_params=kwargs.get("packed_seq_params", None), + ) return output, context @@ -2321,7 +2430,9 @@ 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, input_ids=None): + def _forward_mlp_router( + self, hidden_states, padding_mask=None, input_ids=None, packed_seq_params=None + ): """ Executes the router phase of the MoE block. @@ -2350,6 +2461,7 @@ def _forward_mlp_router(self, hidden_states, padding_mask=None, input_ids=None): intermediate_tensors=(), padding_mask=padding_mask, input_ids=input_ids, + packed_seq_params=packed_seq_params, ) for attr_name in self.mlp.token_dispatcher.cudagraph_attrs: @@ -2403,7 +2515,12 @@ def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_b return self._forward_post_mlp((output, mlp_bias), residual) def _forward_mlp( - self, hidden_states, inference_context=None, padding_mask=None, input_ids=None + self, + hidden_states, + inference_context=None, + padding_mask=None, + input_ids=None, + packed_seq_params=None, ): """ Orchestrates the MLP forward pass, handling partial CUDA graph execution logic. @@ -2420,10 +2537,17 @@ def _forward_mlp( ) def _forward_mlp_partial_cudagraphs( - hidden_states, inference_context=None, padding_mask=None, input_ids=None + hidden_states, + inference_context=None, + padding_mask=None, + input_ids=None, + packed_seq_params=None, ): residual, hidden_states, probs, shared_expert_output = self._forward_mlp_router( - hidden_states, padding_mask=padding_mask, input_ids=input_ids + hidden_states, + padding_mask=padding_mask, + input_ids=input_ids, + packed_seq_params=packed_seq_params, ) # After the router graph replays, the captured .copy_() operations that update @@ -2451,6 +2575,7 @@ def _forward_mlp_partial_cudagraphs( hidden_states, padding_mask=padding_mask, input_ids=input_ids, + packed_seq_params=packed_seq_params, ) else: return tensor_parallel.checkpoint( @@ -2458,15 +2583,22 @@ def _forward_mlp_partial_cudagraphs( _forward_mlp_partial_cudagraphs, padding_mask=padding_mask, input_ids=input_ids, + packed_seq_params=packed_seq_params, ), False, hidden_states, ) else: return _forward_mlp_partial_cudagraphs( - hidden_states, padding_mask=padding_mask, input_ids=input_ids + hidden_states, + padding_mask=padding_mask, + input_ids=input_ids, + packed_seq_params=packed_seq_params, ) else: return super()._forward_mlp( - hidden_states, padding_mask=padding_mask, input_ids=input_ids + hidden_states, + padding_mask=padding_mask, + input_ids=input_ids, + packed_seq_params=packed_seq_params, ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index ea0fce46c58..cf7b55c88bf 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -17,6 +17,7 @@ from megatron.core.activations import squared_relu from megatron.core.dist_checkpointing.validation import StrictHandling from megatron.core.fusions.fused_bias_geglu import quick_gelu +from megatron.core.model_parallel_config import _parse_pad_packed_seq_alignment from megatron.core.msc_utils import MultiStorageClientFeature from megatron.core.quantization.utils import ( kitchen_quantization_recipe_config, @@ -1571,6 +1572,44 @@ def validate_args(args, defaults={}): f"to {args.data_parallel_size * args.context_parallel_size}." ) + 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.dynamic_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 ( diff --git a/megatron/training/training.py b/megatron/training/training.py index e7cad5ce340..cac829727d7 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2167,6 +2167,7 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch """ args = get_args() timers = get_timers() + num_microbatches = get_num_microbatches() rerun_state_machine = get_rerun_state_machine() save_params_in_this_iteration = (args.save_params_interval is not None and @@ -2298,6 +2299,7 @@ def _save_state_dict(attr_name, label): None, None, 0, + num_microbatches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch, ) @@ -2391,6 +2393,7 @@ def _save_state_dict(attr_name, label): grad_norm, num_zeros_in_grad, log_max_attention_logit, + num_microbatches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch, ) @@ -2403,6 +2406,7 @@ def _save_state_dict(attr_name, label): grad_norm, num_zeros_in_grad, log_max_attention_logit, + num_microbatches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch, ) @@ -2424,6 +2428,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() @@ -2588,7 +2593,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") @@ -2872,11 +2877,11 @@ def save_checkpoint_and_time( train_data_iterator=train_data_iterator, preprocess_common_state_dict_fn=preprocess_common_state_dict, ) - + # Stop timer and compute time elapsed to save checkpoint. Stop timer before timers.log() call as it resets the timer. timers(timer_key).stop(barrier=True) save_checkpoint_duration = timers(timer_key).elapsed(reset=False) - + if should_report_memory: # Track memory after checkpoint save. report_memory(f"(after save_checkpoint for iteration {iteration})") @@ -3439,12 +3444,13 @@ 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. buffered_rollouts = None while iteration < args.train_iters: - if (args.profile + if (args.profile and (len(args.profile_ranks) == 0 or torch.distributed.get_rank() in args.profile_ranks)): # Enable NVTX range when profiling starts and nvtx_ranges is set. @@ -3477,9 +3483,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()})" @@ -3563,6 +3569,7 @@ def trace_handler(p): grad_norm = 0.0 num_zeros_in_grad = 0 max_attention_logit = None + num_microbatches = get_num_microbatches() seqlen_sum_this_global_batch = None seqlen_squared_sum_this_global_batch = None else: @@ -3576,6 +3583,7 @@ def trace_handler(p): grad_norm, num_zeros_in_grad, max_attention_logit, + num_microbatches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch, ) = train_step( @@ -3730,6 +3738,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 @@ -4413,3 +4422,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, "dynamic_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 9646e6b36ce..9265e8a832a 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -30,7 +30,11 @@ 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.rerun_state_machine import get_rerun_state_machine from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer from megatron.core.transformer.multi_token_prediction import get_mtp_ranks, mtp_on_this_rank @@ -125,12 +129,15 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): 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(config, ignore_virtual=False, vp_stage=vp_stage), vp_stage=vp_stage, dynamic_cp=args.dynamic_context_parallel, + config=config, ) # TODO: this is pretty hacky, find a better way @@ -140,7 +147,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): and not is_packed_sequence and ((not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage))) ): - return None, None, None, None, None, None + return None, None, None, None, None, None, None # get batches based on the TP rank you are on batch = get_batch_on_this_tp_rank( @@ -176,6 +183,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): max_seqlen_kv=int(max_seqlen[0].item()), qkv_format='thd', ), + None, ) if cu_seqlens is None: @@ -187,7 +195,52 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): batch, cu_seqlens, cu_seqlens_padded, max_seqlen ) - return (*batch.values(), packed_seq_params) + # 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. + padding_mask = None + if config.pad_packed_seq_alignment is not None and packed_seq_params is not None: + tokens = batch.get('tokens', None) + labels = batch.get('labels', None) + loss_mask = batch.get('loss_mask', None) + position_ids = batch.get('position_ids', 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, + ) + ) + if 'tokens' in batch: + batch['tokens'] = tokens + if 'labels' in batch: + batch['labels'] = labels + if 'loss_mask' in batch: + batch['loss_mask'] = loss_mask + if 'position_ids' in batch: + batch['position_ids'] = position_ids + + # Unpack explicitly to avoid relying on dict insertion order. + return ( + batch.get('tokens'), + batch.get('labels'), + batch.get('loss_mask'), + batch.get('attention_mask'), + batch.get('position_ids'), + packed_seq_params, + padding_mask, + ) # define spiky loss as a loss that's 10x the max loss observed @@ -272,8 +325,8 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa global stimer with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") - tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = get_batch( - data_iterator, vp_stage + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = ( + get_batch(data_iterator, vp_stage) ) timers('batch-generator').stop() @@ -283,7 +336,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: @@ -294,6 +353,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 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 968e07e9201..92119bb2a5f 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -56,6 +56,7 @@ def _test_fused_mla_rope_inplace(input_format, inverse=False, remove_interleavin multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -101,6 +102,7 @@ def _test_fused_mla_rope_inplace(input_format, inverse=False, remove_interleavin freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, @@ -151,6 +153,7 @@ def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -209,6 +212,7 @@ def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False): 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 4c1837cd900..44f0ec8d2fd 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -79,6 +79,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, @@ -349,7 +351,7 @@ # 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 index bf929c374d4..f1fbcb53dea 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -9,6 +9,9 @@ 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, ) @@ -17,6 +20,62 @@ 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. @@ -210,8 +269,14 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc dynamic_cp=dynamic_cp, ) - # Unpack the result - tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = result + # The helper returns a 7-tuple; scheduler THD always provides padding_mask. + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = ( + result + ) + assert padding_mask is not None + assert padding_mask.dtype == torch.bool + assert padding_mask.dim() == 2 + assert not padding_mask.any(), "Mock data has no per-sequence padding." # Get parallel state info tp_rank = parallel_state.get_tensor_model_parallel_rank() @@ -288,7 +353,7 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc # TEST 3: Verify TP ranks receive identical data after broadcast # ===================================================================== if tp > 1: - test_tensors = [] + test_tensors = [padding_mask] if is_first_stage: test_tensors.extend([tokens, position_ids]) if is_last_stage: @@ -320,6 +385,11 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc 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() diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 4c6844f9b53..a19d2246b98 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1110,6 +1110,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 533f64081f4..ac759c2bdce 100644 --- a/tests/unit_tests/transformer/test_thd_correctness.py +++ b/tests/unit_tests/transformer/test_thd_correctness.py @@ -111,14 +111,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..92a81b2fcdc --- /dev/null +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -0,0 +1,1006 @@ +# 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.skip(reason="Temporarily disabled until the required Transformer Engine PR lands.") +@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}"