From e9265f514a8f72937d24e9ce09561865105d33e4 Mon Sep 17 00:00:00 2001 From: Haochen Yuan Date: Tue, 7 Apr 2026 22:25:23 -0700 Subject: [PATCH 01/25] add cuda graph support for thd format training Signed-off-by: HaochenYuan --- .../models/common/embeddings/rope_utils.py | 97 ++++++----- megatron/core/models/gpt/gpt_model.py | 18 ++- megatron/core/packed_seq_params.py | 153 ++++++++++++++++++ .../text/libraries/null_tokenizer.py | 4 +- megatron/core/transformer/cuda_graphs.py | 138 ++++++++++++++++ megatron/core/transformer/module.py | 76 ++++++--- megatron/core/transformer/moe/moe_layer.py | 23 +++ .../core/transformer/transformer_config.py | 38 +++++ .../core/transformer/transformer_layer.py | 140 ++++++++++++---- pretrain_gpt.py | 49 +++++- 10 files changed, 622 insertions(+), 114 deletions(-) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index c97f738771b..d46d31fdbc5 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -217,17 +217,19 @@ def _apply_rotary_pos_emb_thd( cp_group: torch.distributed.ProcessGroup = None, multi_latent_attention: Optional[bool] = 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 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 +242,45 @@ 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) + seq_idx = torch.searchsorted(cu_seqlens, token_pos, right=True) - 1 + seq_idx = seq_idx.clamp(min=0, max=cu_seqlens.shape[0] - 2) - 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) + seq_start = cu_seqlens[seq_idx] + local_pos = token_pos - seq_start + + full_seqlen = (cu_seqlens[seq_idx + 1] - seq_start) * cp_size + chunk_size = (cu_seqlens[seq_idx + 1] - seq_start) // 2 + + if cp_size > 1: + is_first_half = local_pos < chunk_size + freq_pos = torch.where( + is_first_half, + cp_rank * chunk_size + local_pos, + full_seqlen - (cp_rank + 1) * chunk_size + (local_pos - chunk_size), + ) + else: + freq_pos = local_pos + + if freqs.dim() >= 1 and freqs.size(0) > total_tokens: + freq_pos = freq_pos + seq_start * cp_size + + 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( diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 01df346c05c..7f284812179 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -331,7 +331,23 @@ def _preprocess( # Decoder embedding. if decoder_input is not None: - pass + # For non-pre_process PP stages that receive decoder_input, scatter padding_mask + # to match the sequence-parallel partitioned hidden_states if needed. + if ( + padding_mask is not None + and self.config.sequence_parallel + and padding_mask.shape[1] != decoder_input.shape[0] + and padding_mask.shape[1] % self.config.tensor_model_parallel_size == 0 + and padding_mask.shape[1] // self.config.tensor_model_parallel_size + == decoder_input.shape[0] + ): + padding_mask = ( + tensor_parallel.scatter_to_sequence_parallel_region( + padding_mask.transpose(0, 1).contiguous() + ) + .transpose(0, 1) + .contiguous() + ) elif self.pre_process: if padding_mask is not None: assert padding_mask.shape == input_ids.shape, ( diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index b1b4275fee1..60d118a060e 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 Optional, Tuple import torch import torch.distributed as dist +import torch.nn.functional as F from torch import Tensor @@ -78,3 +80,154 @@ 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_thd_for_cuda_graph( + tokens: Optional[Tensor], + labels: Optional[Tensor], + loss_mask: Optional[Tensor], + position_ids: Optional[Tensor], + packed_seq_params: PackedSeqParams, + max_seqlen: int, + max_num_seqs: int, +) -> Tuple[ + Optional[Tensor], + Optional[Tensor], + Optional[Tensor], + Optional[Tensor], + PackedSeqParams, + Optional[Tensor], +]: + """Pad THD batch data to fixed sizes for CUDA Graph compatibility. + + CUDA Graph requires static tensor shapes. This function pads: + - tokens, labels, loss_mask, position_ids along dim=-1 to max_seqlen + - cu_seqlens tensors to (max_num_seqs + 1) entries, filled with actual_T + - Generates padding_mask for MoE aux loss exclusion + + Returns: + Padded (tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask) + padding_mask: [1, max_seqlen] bool tensor, True at padding positions. + """ + + def _pad_seq_tensor(t, target_len): + if t is None: + return None + actual_len = t.shape[-1] + if actual_len >= target_len: + return t + return F.pad(t, (0, target_len - actual_len), value=0) + + def _pad_cu_seqlens(cu_seqlens, target_entries): + if cu_seqlens is None: + return None + actual_entries = cu_seqlens.shape[0] + 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 + + actual_T = None + mask_device = None + for candidate in (tokens, labels, loss_mask, position_ids): + if candidate is not None: + actual_T = candidate.shape[-1] + mask_device = candidate.device + break + actual_T_is_local = actual_T is not None + if actual_T is None: + assert packed_seq_params.cu_seqlens_q is not None, ( + "packed_seq_params.cu_seqlens_q must be available to derive padding_mask " + "when tokens/labels/loss_mask/position_ids are all None." + ) + actual_T = int(packed_seq_params.cu_seqlens_q[-1].item()) + mask_device = packed_seq_params.cu_seqlens_q.device + + if actual_T is not None and 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 + from megatron.core import parallel_state + + _cp_size = ( + packed_seq_params.local_cp_size + if packed_seq_params.local_cp_size is not None + else parallel_state.get_context_parallel_world_size() + ) + _global_max_seqlen = max_seqlen * _cp_size + assert _max_individual <= _global_max_seqlen, ( + f"Individual request length ({_max_individual}) exceeds the global max sequence length " + f"({_global_max_seqlen} = max_seqlen_per_dp_cp_rank {max_seqlen} * cp_size {_cp_size}). " + f"Each request must fit within the CUDA Graph static buffer after CP partitioning. " + f"Increase --max-seqlen-per-dp-cp-rank or --seq-length, or filter out overlong requests." + ) + + tokens = _pad_seq_tensor(tokens, max_seqlen) + labels = _pad_seq_tensor(labels, max_seqlen) + loss_mask = _pad_seq_tensor(loss_mask, max_seqlen) + position_ids = _pad_seq_tensor(position_ids, max_seqlen) + + target_cu_entries = max_num_seqs + 1 + padded_params = PackedSeqParams( + qkv_format=packed_seq_params.qkv_format, + cu_seqlens_q=_pad_cu_seqlens(packed_seq_params.cu_seqlens_q, target_cu_entries), + cu_seqlens_kv=_pad_cu_seqlens(packed_seq_params.cu_seqlens_kv, target_cu_entries), + cu_seqlens_q_padded=_pad_cu_seqlens( + packed_seq_params.cu_seqlens_q_padded, target_cu_entries + ), + cu_seqlens_kv_padded=_pad_cu_seqlens( + packed_seq_params.cu_seqlens_kv_padded, target_cu_entries + ), + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + local_cp_size=packed_seq_params.local_cp_size, + cp_group=packed_seq_params.cp_group, + ) + + from megatron.core import parallel_state + + cp_size = ( + packed_seq_params.local_cp_size + if packed_seq_params.local_cp_size is not None + else parallel_state.get_context_parallel_world_size() + ) + cp_rank = parallel_state.get_context_parallel_rank() if cp_size > 1 else 0 + + if cp_size > 1: + from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices + + if actual_T_is_local: + local_actual_T = int(actual_T) + local_max_seqlen = int(max_seqlen) + else: + local_actual_T = int( + get_thd_partitioned_indices( + 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, + int(actual_T), + cp_size, + cp_rank, + ).numel() + ) + local_max_seqlen = int( + get_thd_partitioned_indices( + padded_params.cu_seqlens_q_padded + if padded_params.cu_seqlens_q_padded is not None + else padded_params.cu_seqlens_q, + max_seqlen, + cp_size, + cp_rank, + ).numel() + ) + padding_mask = ( + torch.arange(local_max_seqlen, device=mask_device).unsqueeze(0) >= local_actual_T + ) + else: + padding_mask = torch.arange(max_seqlen, device=mask_device).unsqueeze(0) >= actual_T + + 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..47b43dbe577 100644 --- a/megatron/core/tokenizers/text/libraries/null_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/null_tokenizer.py @@ -93,8 +93,8 @@ def eod(self): @property def pad_id(self): - """Returns pad token.""" - return self._pad_id + """Returns id of padding token (same as eod for NullTokenizer).""" + return self._eod_id @property def additional_special_tokens_ids(self): diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 884444557a0..aa5ae94ddb8 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1993,6 +1993,39 @@ def get_rotary_pos_emb(transformer_module, transformer_input): static_inputs = layer.get_layer_static_inputs(self.seq_length, self.micro_batch_size) + # For the post_process stage (last PP/VPP chunk with labels), padding_mask + # arrives at full CP-local size (max_seqlen_per_dp_cp_rank) because: + # 1. labels are present -> actual_T_is_local=True -> no CP re-partition + # 2. pre_process=False -> _preprocess does not scatter + # Other non-pre_process chunks (intermediate VPP) have no data, so padding_mask + # is CP-partitioned to ~max_seqlen/CP ~ max_seqlen/TP (default static size). + if ( + 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_of_the_layer, "pre_process", True) + and getattr(chunk_of_the_layer, "post_process", False) + and "padding_mask" in static_inputs + ): + local_slen = self.config.max_seqlen_per_dp_cp_rank + static_inputs["padding_mask"] = torch.ones( + 1, local_slen, dtype=torch.bool, device=torch.cuda.current_device() + ) + + if os.getenv("THD_DEBUG_CG_IO", "0") == "1": + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + hs = static_inputs.get("hidden_states", None) + pm = static_inputs.get("padding_mask", None) + print( + f"[THD_DEBUG_CG_IO][capture] rank={rank} " + f"layer={getattr(layer, 'layer_number', 'na')} " + f"pre_process={getattr(chunk_of_the_layer, 'pre_process', 'na')} " + f"hidden_shape={tuple(hs.shape) if torch.is_tensor(hs) else 'na'} " + f"padding_mask_shape={tuple(pm.shape) if torch.is_tensor(pm) else 'na'}", + flush=True, + ) + from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.transformer_layer import TransformerLayer @@ -2164,6 +2197,59 @@ 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)) + + @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, + ) + def _get_cuda_graph_input_data(self): """ Create the CUDA Graph capturing input data. @@ -2182,6 +2268,58 @@ def _get_cuda_graph_input_data(self): 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, + get_schedule_table as _probe_get_st, + ) + _, _, _probe_warmup, _ = _probe_get_pp( + probe_num_microbatches, + self.num_model_chunks, + self.config.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, + self.config.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()) + requested_num_slots = self.config.cuda_graph_num_microbatch_slots + if requested_num_slots is not None: + assert requested_num_slots >= auto_num_slots, ( + "cuda_graph_num_microbatch_slots is smaller than the minimum safe number " + f"of slots for the current PP/VPP topology: requested={requested_num_slots}, " + f"required>={auto_num_slots}" + ) + self.num_microbatches = requested_num_slots + else: + self.num_microbatches = auto_num_slots + 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 enabled. ' + f'runtime_num_microbatches={get_num_microbatches()}, ' + f'auto_num_slots={auto_num_slots}, capture_num_microbatches={self.num_microbatches}', + ) else: self.num_microbatches = get_num_microbatches() diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index d19f6d094c0..5e3b540a7e9 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -1,6 +1,7 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. """Megatron Module.""" +import os from functools import partial from typing import Optional, Tuple @@ -222,6 +223,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,25 +237,43 @@ 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(), - ) + if self._is_thd_cuda_graph(): + 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." + ) + max_T = self.config.max_seqlen_per_dp_cp_rank + slen_per_cptp = ( + max_T // tensor_model_parallel_size if sequence_parallel else max_T + ) + static_inputs = {} + static_inputs["hidden_states"] = torch.ones( + (slen_per_cptp, 1, self.config.hidden_size), + dtype=torch.bfloat16, + requires_grad=True, + device=torch.cuda.current_device(), + ) + else: + slen_per_cp = seq_length // context_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 def setup_manual_hooks(self, make_hook_func): @@ -300,6 +326,23 @@ def _te_cuda_graph_replay(self, *args, **kwargs): cg_index = getattr(self, 'current_microbatch', 0) % len(self.cuda_graphs) cudagraph_args, cudagraph_kwargs = self._get_te_cuda_graph_replay_args(*args, **kwargs) + if os.getenv("THD_DEBUG_CG_IO", "0") == "1": + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + hidden_shape = ( + tuple(cudagraph_args[0].shape) + if len(cudagraph_args) > 0 and torch.is_tensor(cudagraph_args[0]) + else 'na' + ) + padding_mask = cudagraph_kwargs.get("padding_mask", None) + print( + f"[THD_DEBUG_CG_IO][replay] rank={rank} " + f"layer={getattr(self, 'layer_number', 'na')} " + f"microbatch={getattr(self, 'current_microbatch', 'na')} cg_index={cg_index} " + f"hidden_shape={hidden_shape} " + f"padding_mask_shape={tuple(padding_mask.shape) if torch.is_tensor(padding_mask) else 'na'}", + flush=True, + ) + for hook, hook_args in self.cuda_graph_manual_hooks: hook(*hook_args) return self.cuda_graphs[cg_index](*cudagraph_args, **cudagraph_kwargs) @@ -318,15 +361,6 @@ def _get_te_cuda_graph_replay_args(self, *args, **kwargs): cudagraph_kwargs = kwargs.copy() cudagraph_kwargs['is_first_microbatch'] = getattr(self, 'current_microbatch', 0) == 0 - if self.config.fine_grained_activation_offloading and getattr( - self, 'offload_module_in_cuda_graph', False - ): - from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( - FineGrainedActivationOffloadingInterface as off_interface, - ) - - cudagraph_kwargs['cuda_graph_stream'] = off_interface.cuda_graph_stream() - cudagraph_kwargs['cuda_graph_event'] = off_interface.cuda_graph_event() return cudagraph_args, cudagraph_kwargs def _should_call_local_cudagraph(self, *args, **kwargs): diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 359b7c4a4bd..2ecd4fb15b2 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -703,6 +703,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() diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 6687d572914..b757e7110f7 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1081,6 +1081,27 @@ 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_cuda_graph_max_num_seqs: int = 32 + """Maximum number of packed sequences per microbatch when using THD format with CUDA Graph. + cu_seqlens tensors will be padded to this size + 1.""" + + cuda_graph_dynamic_microbatches: bool = field( + default=False, + metadata={"argparse_meta": {"arg_names": ["--cuda-graph-dynamic-microbatches"]}}, + ) + """Enable CUDA graph slot reuse so the same captured graphs can be replayed for a dynamic + number of microbatches. This option is only meaningful for cuda_graph_impl=transformer_engine. + When enabled, capture builds a bounded number of graph slots and replay maps real + microbatch_id to slot_id by modulo.""" + + cuda_graph_num_microbatch_slots: Optional[int] = field( + default=None, + metadata={"argparse_meta": {"arg_names": ["--cuda-graph-num-microbatch-slots"]}}, + ) + """Number of CUDA graph slots to capture per layer for dynamic microbatch replay. + If None, an automatic slot count is derived from the PP/VPP schedule topology. + If set, the provided value must be >= the automatically derived safe minimum.""" + #################### # Hyper-Connection Configuration #################### @@ -2654,6 +2675,23 @@ 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": + if self.cuda_graph_dynamic_microbatches: + if self.cuda_graph_num_microbatch_slots is not None: + assert self.cuda_graph_num_microbatch_slots >= 1, ( + "cuda_graph_num_microbatch_slots must be >= 1 when " + "cuda_graph_dynamic_microbatches is enabled." + ) + else: + assert not self.cuda_graph_dynamic_microbatches, ( + "cuda_graph_dynamic_microbatches is only supported with " + "cuda_graph_impl=transformer_engine." + ) + assert self.cuda_graph_num_microbatch_slots is None, ( + "cuda_graph_num_microbatch_slots 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 diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 67bb04837ac..b896bdae9ef 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1079,37 +1079,66 @@ 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) - if not isinstance(self.self_attention, IdentityOp) and ( - not self.config.cuda_graph_modules - or CudaGraphModule.attn in self.config.cuda_graph_modules - ): - if not self.config.create_attention_mask_in_dataloader: - if self.self_attention.attn_mask_type not in ( - AttnMaskType.causal, - AttnMaskType.no_mask, - AttnMaskType.causal_bottom_right, - ): - log_single_rank( - logger, - logging.WARNING, - "TE CUDA graph capture is omitting attention_mask because " - "create_attention_mask_in_dataloader is False, but " - f"attn_mask_type={self.self_attention.attn_mask_type.name} may require " - "an explicit mask. Ensure this is intended for the current workload.", + if self._is_thd_cuda_graph(): + if not isinstance(self.self_attention, IdentityOp) and ( + not self.config.cuda_graph_modules + or CudaGraphModule.attn in self.config.cuda_graph_modules + ): + max_T = self.config.max_seqlen_per_dp_cp_rank + max_num_seqs = self.config.thd_cuda_graph_max_num_seqs + device = torch.cuda.current_device() + + cu_seqlens = torch.zeros(max_num_seqs + 1, dtype=torch.int32, device=device) + cu_seqlens[0] = 0 + cu_seqlens[1] = max_T + cu_seqlens[2:] = 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 = slen_for_mask // self.config.tensor_model_parallel_size + static_inputs["padding_mask"] = torch.ones( + 1, slen_for_mask, dtype=torch.bool, device=torch.cuda.current_device() + ) + else: + if not isinstance(self.self_attention, IdentityOp) and ( + not self.config.cuda_graph_modules + or CudaGraphModule.attn in self.config.cuda_graph_modules + ): + if not self.config.create_attention_mask_in_dataloader: + if self.self_attention.attn_mask_type not in ( + AttnMaskType.causal, + AttnMaskType.no_mask, + AttnMaskType.causal_bottom_right, + ): + log_single_rank( + logger, + logging.WARNING, + "TE CUDA graph capture is omitting attention_mask because " + "create_attention_mask_in_dataloader is False, but " + f"attn_mask_type={self.self_attention.attn_mask_type.name} may require " + "an explicit mask. Ensure this is intended for the current workload.", + ) + else: + 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()) + .reshape(1, 1, slen_per_cp, seq_length) + .tile(micro_batch_size, 1, 1, 1) ) - else: - 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()) - .reshape(1, 1, slen_per_cp, seq_length) - .tile(micro_batch_size, 1, 1, 1) - ) # Add input_ids for hash-based MoE routing under CUDA graphs. # Only add for layers that actually use hash routing, @@ -1122,7 +1151,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): @@ -1153,6 +1181,46 @@ 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 + 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, + ) + kwargs['packed_seq_params'] = packed_seq_params + def _te_cuda_graph_capture(self, *args, **kwargs): """ CUDA Graph capture for this layer using TE interface. @@ -1160,10 +1228,11 @@ 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. if self.offload_module_in_cuda_graph: if len(args) > 0: hidden_states = args[0] @@ -1197,7 +1266,9 @@ 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), ) if not isinstance(hidden_states, list) and not isinstance(hidden_states, tuple): cuda_graph_outputs = [hidden_states] @@ -1206,8 +1277,6 @@ def _te_cuda_graph_capture(self, *args, **kwargs): if context is not None: cuda_graph_outputs.append(context) # Record the forward event on cuda graph stream for cuda graph capture. - # This is to ensure the main stream waits for computing on cuda graph stream to complete, - # and overlaps with the D2H transfer on offloading stream. if self.offload_module_in_cuda_graph: self.off_interface.forward_record() return tuple(cuda_graph_outputs) @@ -1218,7 +1287,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. """ + self._decompose_packed_seq_params_to_kwargs(kwargs) + context = None if ( self.config.cuda_graph_modules @@ -1350,7 +1422,11 @@ 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)) + output = self._forward_mlp( + *cuda_graph_output, + padding_mask=kwargs.get("padding_mask", None), + input_ids=kwargs.get("input_ids", None), + ) return output, context def _get_te_cuda_graph_replay_args(self, *args, **kwargs): diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 9646e6b36ce..b004397d4c5 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -30,7 +30,7 @@ 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, pad_thd_for_cuda_graph 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,13 +125,26 @@ 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: - return get_batch_on_this_rank_for_sequence_packing( + result = 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, ) + # Pad THD batch for CUDA Graph compatibility when max_seqlen_per_dp_cp_rank is set. + padding_mask = None + if config.max_seqlen_per_dp_cp_rank is not None and len(result) >= 6: + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = result[:6] + if packed_seq_params is not None: + tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = \ + pad_thd_for_cuda_graph( + tokens, labels, loss_mask, position_ids, packed_seq_params, + max_seqlen=config.max_seqlen_per_dp_cp_rank, + max_num_seqs=config.thd_cuda_graph_max_num_seqs, + ) + return tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask + return (*result, padding_mask) # TODO: this is pretty hacky, find a better way is_packed_sequence = args.sft or (args.use_varlen_dataset and not args.varlen_sbhd_validation) @@ -140,7 +153,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 +189,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 +201,29 @@ 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 THD batch for CUDA Graph compatibility when max_seqlen_per_dp_cp_rank is set. + padding_mask = None + if config.max_seqlen_per_dp_cp_rank 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) + tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = \ + pad_thd_for_cuda_graph( + tokens, labels, loss_mask, position_ids, packed_seq_params, + max_seqlen=config.max_seqlen_per_dp_cp_rank, + max_num_seqs=config.thd_cuda_graph_max_num_seqs, + ) + 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 + + return (*batch.values(), packed_seq_params, padding_mask) # define spiky loss as a loss that's 10x the max loss observed @@ -272,8 +308,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() @@ -294,6 +330,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 From 81a89ed1687c1882c8ca4db787be237d148be5c4 Mon Sep 17 00:00:00 2001 From: haochen Yuan Date: Fri, 17 Apr 2026 02:31:06 -0700 Subject: [PATCH 02/25] add unit test Signed-off-by: HaochenYuan --- .../transformer/test_thd_cuda_graph.py | 391 ++++++++++++++++++ 1 file changed, 391 insertions(+) create mode 100644 tests/unit_tests/transformer/test_thd_cuda_graph.py 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..79f93091303 --- /dev/null +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -0,0 +1,391 @@ +# 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_EP4_ETP1 + sequence packing +(requires 8 GPUs, slow ~5 min per run, 4 runs total): + 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 subprocess + +import pytest +import torch + +from megatron.core.packed_seq_params import PackedSeqParams, pad_thd_for_cuda_graph +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') + + +# ============================================================================= +# 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_cuda_graph_max_num_seqs=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_thd_for_cuda_graph correctness +# ============================================================================= + +class TestPadThdForCudaGraph: + + 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_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_thd_for_cuda_graph( + tokens, tokens.clone(), torch.ones(1, total_T, device="cuda"), + torch.arange(total_T, device="cuda").unsqueeze(0), + _make_psp(seqlens), max_seqlen, 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 + 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_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_thd_for_cuda_graph( + torch.ones(1, total_T, device="cuda"), None, None, None, + _make_psp(seqlens), max_seqlen, 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_cu_seqlens_fill_value(self): + """Padded entries repeat last cumulative sum (prevents OOB reads).""" + seqlens, total_T = [50, 30], 80 + _, _, _, _, p, _ = pad_thd_for_cuda_graph( + torch.ones(1, total_T, device="cuda"), None, None, None, + _make_psp(seqlens), 128, 32) + assert p.cu_seqlens_q[0] == 0 and p.cu_seqlens_q[2] == 80 + assert (p.cu_seqlens_q[3:] == 80).all() + + @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_thd_for_cuda_graph( + None, None, None, None, _make_psp(seqlens), max_seqlen, 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 + 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 + + +# ============================================================================= +# 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 (matches test_moonlight_qwen3_bitwise.sh). +_MEGATRON_DIR = os.environ.get( + 'MEGATRON_DIR', + '/lustre/fsw/coreai_devtech_all/haocheny/migrate_to_TE_0415/Megatron-LM') +_MOONLIGHT_LOAD = os.environ.get( + 'MOONLIGHT_CKPT', + '/lustre/fsw/coreai_devtech_all/haocheny/mcore_models/Moonlight-16B-A3B-Instruct') + +_SFT_JSON = ( + '{"mode":"distribution","type":"lognormal",' + '"min_seq_len":128,"max_seq_len":2048,"mean_seq_len":1024,"lognormal_sigma":0.8}' +) + +_TRAIN_ITERS = 5 + +_COMMON_ARGS = [ + "--seq-length", "2048", "--max-position-embeddings", "8192", + "--micro-batch-size", "1", "--global-batch-size", "4", + "--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", + "--sft", "--mock-data", + "--tokenizer-type", "NullTokenizer", + "--sft-mock-dataset-config-json", _SFT_JSON, + "--sequence-packing-scheduler", "dp_balanced", + "--max-seqlen-per-dp-cp-rank", "1024", + "--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-cuda-graph-max-num-seqs", "32", +] + +_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", "alltoall", + "--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", + "--load", _MOONLIGHT_LOAD, + "--no-load-optim", "--no-load-rng", +] + +_QWEN3_ARGS = _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", "alltoall", +] + + +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"] = _MEGATRON_DIR + ":" + 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(master_port), + "pretrain_gpt.py", + ] + model_args + cuda_graph_args + + result = subprocess.run( + cmd, cwd=_MEGATRON_DIR, 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) + lb_loss = re.search(r"load_balancing_loss:\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)}", + ] + if lb_loss: + parts.append(f"lb_loss={lb_loss.group(1)}") + parts.append(f"grad_norm={grad_norm.group(1)}") + results.append(" | ".join(parts)) + return results + + +@pytest.mark.internal +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(torch.cuda.device_count() < 8, reason="requires 8 GPUs") +@pytest.mark.parametrize( + "model_name,model_args,base_port", + [ + ("moonlight", _MOONLIGHT_ARGS, 29660), + ("qwen3", _QWEN3_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 cuda_graph_scope=attn` + -- using the exact same args as test_moonlight_qwen3_bitwise.sh. + Asserts the per-iteration `lm loss / load_balancing_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, 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-impl", "transformer_engine", + "--cuda-graph-scope", "attn", + ], + 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}") From 6e3e681c6245eeeafcf7131cc3674c485e92e183 Mon Sep 17 00:00:00 2001 From: haochen Yuan Date: Wed, 29 Apr 2026 09:30:26 -0700 Subject: [PATCH 03/25] fix & refactor pad-thd logic Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 40 +++++++++++++++++-- megatron/core/transformer/module.py | 10 +++++ .../core/transformer/transformer_config.py | 14 +++++-- .../core/transformer/transformer_layer.py | 6 ++- pretrain_gpt.py | 21 +++------- tests/unit_tests/test_sequence_packing.py | 8 +++- .../transformer/test_thd_cuda_graph.py | 4 +- 7 files changed, 77 insertions(+), 26 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index b6a6a65dc3c..d7184d213d8 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -34,6 +34,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 +43,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 packed sequences per + microbatch. When set, the scheduler closes a pack as soon as it + reaches this many sequences in addition to the token-budget condition. """ 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 +123,10 @@ 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: @@ -459,6 +467,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=getattr(config, 'thd_max_num_seqs', None), **scheduler_kwargs, ) @@ -486,6 +495,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. @@ -493,8 +503,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 THD CUDA Graph padding. When None + or config.max_seqlen_per_dp_cp_rank 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: @@ -680,8 +693,29 @@ def get_batch_on_this_rank_for_sequence_packing( cp_group=cp_group, ) + # Pad to static shapes for THD + CUDA Graph when requested. + padding_mask = None + if ( + config is not None + and getattr(config, 'max_seqlen_per_dp_cp_rank', None) is not None + and packed_seq_params is not None + ): + from megatron.core.packed_seq_params import pad_thd_for_cuda_graph + + tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = ( + pad_thd_for_cuda_graph( + tokens, + labels, + loss_mask, + position_ids, + packed_seq_params, + max_seqlen=config.max_seqlen_per_dp_cp_rank, + max_num_seqs=config.thd_max_num_seqs, + ) + ) + # "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/transformer/module.py b/megatron/core/transformer/module.py index 5e3b540a7e9..b38e9d350ea 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -243,6 +243,7 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): 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 sequence_parallel = self.config.sequence_parallel tensor_model_parallel_size = self.config.tensor_model_parallel_size @@ -361,6 +362,15 @@ def _get_te_cuda_graph_replay_args(self, *args, **kwargs): cudagraph_kwargs = kwargs.copy() cudagraph_kwargs['is_first_microbatch'] = getattr(self, 'current_microbatch', 0) == 0 + if self.config.fine_grained_activation_offloading and getattr( + self, 'offload_module_in_cuda_graph', False + ): + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, + ) + + cudagraph_kwargs['cuda_graph_stream'] = off_interface.cuda_graph_stream() + cudagraph_kwargs['cuda_graph_event'] = off_interface.cuda_graph_event() return cudagraph_args, cudagraph_kwargs def _should_call_local_cudagraph(self, *args, **kwargs): diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index b757e7110f7..218c9c2d563 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1081,9 +1081,17 @@ 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_cuda_graph_max_num_seqs: int = 32 - """Maximum number of packed sequences per microbatch when using THD format with CUDA Graph. - cu_seqlens tensors will be padded to this size + 1.""" + thd_max_num_seqs: int = 32 + """Maximum number of packed sequences per microbatch in THD format. The packing + scheduler closes a pack as soon as it reaches this many sequences (in addition to + the existing token-budget condition). 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 = field( default=False, diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index b896bdae9ef..47973e7871d 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1093,7 +1093,7 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): or CudaGraphModule.attn in self.config.cuda_graph_modules ): max_T = self.config.max_seqlen_per_dp_cp_rank - max_num_seqs = self.config.thd_cuda_graph_max_num_seqs + max_num_seqs = self.config.thd_max_num_seqs device = torch.cuda.current_device() cu_seqlens = torch.zeros(max_num_seqs + 1, dtype=torch.int32, device=device) @@ -1233,6 +1233,8 @@ def _te_cuda_graph_capture(self, *args, **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. if self.offload_module_in_cuda_graph: if len(args) > 0: hidden_states = args[0] @@ -1277,6 +1279,8 @@ def _te_cuda_graph_capture(self, *args, **kwargs): if context is not None: cuda_graph_outputs.append(context) # Record the forward event on cuda graph stream for cuda graph capture. + # This is to ensure the main stream waits for computing on cuda graph stream to complete, + # and overlaps with the D2H transfer on offloading stream. if self.offload_module_in_cuda_graph: self.off_interface.forward_record() return tuple(cuda_graph_outputs) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index b004397d4c5..adbf1c26ec7 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -125,26 +125,17 @@ 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: - result = get_batch_on_this_rank_for_sequence_packing( + # `get_batch_on_this_rank_for_sequence_packing` applies THD + CUDA Graph + # padding internally when `config.max_seqlen_per_dp_cp_rank` is set, and + # returns a 7-tuple including `padding_mask` (None when no padding). + 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, ) - # Pad THD batch for CUDA Graph compatibility when max_seqlen_per_dp_cp_rank is set. - padding_mask = None - if config.max_seqlen_per_dp_cp_rank is not None and len(result) >= 6: - tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = result[:6] - if packed_seq_params is not None: - tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = \ - pad_thd_for_cuda_graph( - tokens, labels, loss_mask, position_ids, packed_seq_params, - max_seqlen=config.max_seqlen_per_dp_cp_rank, - max_num_seqs=config.thd_cuda_graph_max_num_seqs, - ) - return tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask - return (*result, padding_mask) # TODO: this is pretty hacky, find a better way is_packed_sequence = args.sft or (args.use_varlen_dataset and not args.varlen_sbhd_validation) @@ -212,7 +203,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): pad_thd_for_cuda_graph( tokens, labels, loss_mask, position_ids, packed_seq_params, max_seqlen=config.max_seqlen_per_dp_cp_rank, - max_num_seqs=config.thd_cuda_graph_max_num_seqs, + max_num_seqs=config.thd_max_num_seqs, ) if 'tokens' in batch: batch['tokens'] = tokens diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index bf929c374d4..cf535e5a7db 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -210,8 +210,12 @@ 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 + # Unpack the result. The helper now always returns a 7-tuple; the 7th + # value is `padding_mask` (None when THD CUDA Graph padding is not in use). + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = result + assert padding_mask is None, ( + "padding_mask should be None when config is not passed (legacy behavior)." + ) # Get parallel state info tp_rank = parallel_state.get_tensor_model_parallel_rank() diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 79f93091303..bd7e220cda7 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -63,7 +63,7 @@ def _build_layer(H, nh, nkv, ffn, max_seqlen, max_num_seqs, tp=1, sp=False): 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_cuda_graph_max_num_seqs=max_num_seqs, + thd_max_num_seqs=max_num_seqs, tensor_model_parallel_size=tp, sequence_parallel=sp, bf16=True) model_parallel_cuda_manual_seed(42) return TransformerLayer( @@ -222,7 +222,7 @@ def test_noop_without_packed_seq_params(self): "--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-cuda-graph-max-num-seqs", "32", + "--thd-max-num-seqs", "32", ] _MOONLIGHT_ARGS = _COMMON_ARGS + [ From 00c51e13247685110aec81e60aa3db91cbdbff7d Mon Sep 17 00:00:00 2001 From: haochen Yuan Date: Sat, 9 May 2026 02:19:09 -0700 Subject: [PATCH 04/25] refactor Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 15 ++-- megatron/core/datasets/data_schedule_utils.py | 13 ++- .../models/common/embeddings/rope_utils.py | 28 +++++-- megatron/core/packed_seq_params.py | 67 +++++++++++----- megatron/core/transformer/cuda_graphs.py | 53 ++++++------- megatron/core/transformer/module.py | 62 ++++++--------- .../core/transformer/transformer_config.py | 10 +-- .../core/transformer/transformer_layer.py | 79 ++++++++++--------- pretrain_gpt.py | 19 ++++- .../transformer/test_cuda_graphs.py | 58 ++++++++++++++ 10 files changed, 249 insertions(+), 155 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index d7184d213d8..a7181e58dc6 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -580,13 +580,14 @@ def get_batch_on_this_rank_for_sequence_packing( if is_first_or_last_stage or mtp_on_this_rank: if is_tp_rank_0: - # 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) + # 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 tokens/labels length. + 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) diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 51be6282ffe..eaeaa1b8ac5 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -25,20 +25,17 @@ def get_cp_slice_for_thd(batch, cp_group): if cp_size <= 1: return cp_rank = cp_group.rank() - # Use whichever data field is available to determine total_tokens - for _key in ['tokens', 'labels', 'loss_mask', 'position_ids']: - if _key in batch and batch[_key] is not None: - total_tokens = batch[_key].size(0) - break - else: - raise ValueError("Cannot determine total_tokens: no data field found in batch") # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as # cu_seqlens to get the correct result. # TODO: Revert this workaround once TE fixes the issue. cu_seqlens = batch["cu_seqlens_padded"] + # 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 key in batch and batch[key] is not None: batch[key] = batch[key].index_select(0, index) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index d46d31fdbc5..afb34a63c68 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -247,14 +247,24 @@ def _apply_rotary_pos_emb_thd( device = t.device token_pos = torch.arange(total_tokens, device=device) + # `searchsorted(..., right=True) - 1` returns the index of the sequence each + # token belongs to. The `clamp` here guards padding tokens whose + # `searchsorted` position can fall outside the valid `[0, num_seqs)` range + # (last cu_seqlens entry); they get any wrong-but-harmless freq because + # their RoPE output is masked away later by `padding_mask` / `loss_mask`. seq_idx = torch.searchsorted(cu_seqlens, token_pos, right=True) - 1 seq_idx = seq_idx.clamp(min=0, max=cu_seqlens.shape[0] - 2) seq_start = cu_seqlens[seq_idx] local_pos = token_pos - seq_start - full_seqlen = (cu_seqlens[seq_idx + 1] - seq_start) * cp_size - chunk_size = (cu_seqlens[seq_idx + 1] - seq_start) // 2 + # int64 for the multiplications: `cu_seqlens` are int32 (per TE convention), + # but `(seq_len * cp_size)` and downstream arithmetic can overflow int32 at + # very long contexts (e.g. >65k tokens × cp_size). Cast once here and the + # rest of the computation stays in int64. + seq_len_i64 = (cu_seqlens[seq_idx + 1] - seq_start).to(torch.int64) + full_seqlen = seq_len_i64 * cp_size + chunk_size = seq_len_i64 // 2 if cp_size > 1: is_first_half = local_pos < chunk_size @@ -264,11 +274,19 @@ def _apply_rotary_pos_emb_thd( full_seqlen - (cp_rank + 1) * chunk_size + (local_pos - chunk_size), ) else: - freq_pos = local_pos + freq_pos = local_pos.to(torch.int64) if freqs.dim() >= 1 and freqs.size(0) > total_tokens: - freq_pos = freq_pos + seq_start * cp_size - + # `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 + seq_start.to(torch.int64) * cp_size + + # Same rationale as the seq_idx clamp above: padded positions can index + # past `freqs`; they receive a known wrong-but-harmless freq that gets + # masked away. If you suspect a real out-of-range bug, swap clamp for an + # assert during development. freq_pos = freq_pos.clamp(min=0, max=freqs.shape[0] - 1) freqs_packed = freqs[freq_pos] diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 60d118a060e..698dd752532 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -82,6 +82,52 @@ def resolve_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_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``: this is + the reviewer-flagged overflow case and corresponds to "too many packed + sequences in this microbatch for thd_max_num_seqs". Failing fast prevents + a silent CUDA-graph shape mismatch at replay. + """ + 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_num_seqs " + f"({target_entries - 1}). Increase --thd-max-num-seqs, 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 pad_thd_for_cuda_graph( tokens: Optional[Tensor], labels: Optional[Tensor], @@ -110,27 +156,6 @@ def pad_thd_for_cuda_graph( padding_mask: [1, max_seqlen] bool tensor, True at padding positions. """ - def _pad_seq_tensor(t, target_len): - if t is None: - return None - actual_len = t.shape[-1] - if actual_len >= target_len: - return t - return F.pad(t, (0, target_len - actual_len), value=0) - - def _pad_cu_seqlens(cu_seqlens, target_entries): - if cu_seqlens is None: - return None - actual_entries = cu_seqlens.shape[0] - 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 - actual_T = None mask_device = None for candidate in (tokens, labels, loss_mask, position_ids): diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index aa5ae94ddb8..889a2c12e7e 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1993,39 +1993,12 @@ def get_rotary_pos_emb(transformer_module, transformer_input): static_inputs = layer.get_layer_static_inputs(self.seq_length, self.micro_batch_size) - # For the post_process stage (last PP/VPP chunk with labels), padding_mask - # arrives at full CP-local size (max_seqlen_per_dp_cp_rank) because: - # 1. labels are present -> actual_T_is_local=True -> no CP re-partition - # 2. pre_process=False -> _preprocess does not scatter - # Other non-pre_process chunks (intermediate VPP) have no data, so padding_mask - # is CP-partitioned to ~max_seqlen/CP ~ max_seqlen/TP (default static size). - if ( - 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_of_the_layer, "pre_process", True) - and getattr(chunk_of_the_layer, "post_process", False) - and "padding_mask" in static_inputs - ): + 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.ones( 1, local_slen, dtype=torch.bool, device=torch.cuda.current_device() ) - if os.getenv("THD_DEBUG_CG_IO", "0") == "1": - rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 - hs = static_inputs.get("hidden_states", None) - pm = static_inputs.get("padding_mask", None) - print( - f"[THD_DEBUG_CG_IO][capture] rank={rank} " - f"layer={getattr(layer, 'layer_number', 'na')} " - f"pre_process={getattr(chunk_of_the_layer, 'pre_process', 'na')} " - f"hidden_shape={tuple(hs.shape) if torch.is_tensor(hs) else 'na'} " - f"padding_mask_shape={tuple(pm.shape) if torch.is_tensor(pm) else 'na'}", - flush=True, - ) - from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.transformer_layer import TransformerLayer @@ -2201,6 +2174,30 @@ 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. diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index b38e9d350ea..f981d22e719 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -1,7 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. """Megatron Module.""" -import os from functools import partial from typing import Optional, Tuple @@ -249,33 +248,37 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): tensor_model_parallel_size = self.config.tensor_model_parallel_size 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." ) - max_T = self.config.max_seqlen_per_dp_cp_rank - slen_per_cptp = ( - max_T // tensor_model_parallel_size if sequence_parallel else max_T - ) - static_inputs = {} - static_inputs["hidden_states"] = torch.ones( - (slen_per_cptp, 1, self.config.hidden_size), - dtype=torch.bfloat16, - requires_grad=True, - device=torch.cuda.current_device(), - ) + slen_full = self.config.max_seqlen_per_dp_cp_rank + batch = 1 else: - slen_per_cp = seq_length // context_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, + # 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. Hardcoding bfloat16 silently breaks --fp16. + 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(), ) - return static_inputs + } def setup_manual_hooks(self, make_hook_func): """ @@ -327,23 +330,6 @@ def _te_cuda_graph_replay(self, *args, **kwargs): cg_index = getattr(self, 'current_microbatch', 0) % len(self.cuda_graphs) cudagraph_args, cudagraph_kwargs = self._get_te_cuda_graph_replay_args(*args, **kwargs) - if os.getenv("THD_DEBUG_CG_IO", "0") == "1": - rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 - hidden_shape = ( - tuple(cudagraph_args[0].shape) - if len(cudagraph_args) > 0 and torch.is_tensor(cudagraph_args[0]) - else 'na' - ) - padding_mask = cudagraph_kwargs.get("padding_mask", None) - print( - f"[THD_DEBUG_CG_IO][replay] rank={rank} " - f"layer={getattr(self, 'layer_number', 'na')} " - f"microbatch={getattr(self, 'current_microbatch', 'na')} cg_index={cg_index} " - f"hidden_shape={hidden_shape} " - f"padding_mask_shape={tuple(padding_mask.shape) if torch.is_tensor(padding_mask) else 'na'}", - flush=True, - ) - for hook, hook_args in self.cuda_graph_manual_hooks: hook(*hook_args) return self.cuda_graphs[cg_index](*cudagraph_args, **cudagraph_kwargs) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 218c9c2d563..90b83e106a8 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1093,19 +1093,13 @@ class TransformerConfig(ModelParallelConfig): token budget); setting it too large just allocates a slightly larger cu_seqlens buffer.""" - cuda_graph_dynamic_microbatches: bool = field( - default=False, - metadata={"argparse_meta": {"arg_names": ["--cuda-graph-dynamic-microbatches"]}}, - ) + cuda_graph_dynamic_microbatches: bool = False """Enable CUDA graph slot reuse so the same captured graphs can be replayed for a dynamic number of microbatches. This option is only meaningful for cuda_graph_impl=transformer_engine. When enabled, capture builds a bounded number of graph slots and replay maps real microbatch_id to slot_id by modulo.""" - cuda_graph_num_microbatch_slots: Optional[int] = field( - default=None, - metadata={"argparse_meta": {"arg_names": ["--cuda-graph-num-microbatch-slots"]}}, - ) + cuda_graph_num_microbatch_slots: Optional[int] = None """Number of CUDA graph slots to capture per layer for dynamic microbatch replay. If None, an automatic slot count is derived from the PP/VPP schedule topology. If set, the provided value must be >= the automatically derived safe minimum.""" diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 47973e7871d..1199ded9f77 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1086,20 +1086,21 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): 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() + + # 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 not isinstance(self.self_attention, IdentityOp) and ( - not self.config.cuda_graph_modules - or CudaGraphModule.attn in self.config.cuda_graph_modules - ): + if attn_in_graph: max_T = self.config.max_seqlen_per_dp_cp_rank max_num_seqs = self.config.thd_max_num_seqs - device = torch.cuda.current_device() - cu_seqlens = torch.zeros(max_num_seqs + 1, dtype=torch.int32, device=device) - cu_seqlens[0] = 0 - cu_seqlens[1] = max_T - cu_seqlens[2:] = max_T + cu_seqlens[1:] = max_T static_inputs["cu_seqlens_q"] = cu_seqlens static_inputs["cu_seqlens_kv"] = cu_seqlens.clone() @@ -1108,37 +1109,33 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): slen_for_mask = self.config.max_seqlen_per_dp_cp_rank if self.config.sequence_parallel: - slen_for_mask = slen_for_mask // self.config.tensor_model_parallel_size + slen_for_mask //= self.config.tensor_model_parallel_size static_inputs["padding_mask"] = torch.ones( - 1, slen_for_mask, dtype=torch.bool, device=torch.cuda.current_device() + 1, slen_for_mask, dtype=torch.bool, device=device ) - else: - if not isinstance(self.self_attention, IdentityOp) and ( - not self.config.cuda_graph_modules - or CudaGraphModule.attn in self.config.cuda_graph_modules - ): - if not self.config.create_attention_mask_in_dataloader: - if self.self_attention.attn_mask_type not in ( - AttnMaskType.causal, - AttnMaskType.no_mask, - AttnMaskType.causal_bottom_right, - ): - log_single_rank( - logger, - logging.WARNING, - "TE CUDA graph capture is omitting attention_mask because " - "create_attention_mask_in_dataloader is False, but " - f"attn_mask_type={self.self_attention.attn_mask_type.name} may require " - "an explicit mask. Ensure this is intended for the current workload.", - ) - else: - 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()) - .reshape(1, 1, slen_per_cp, seq_length) - .tile(micro_batch_size, 1, 1, 1) + elif attn_in_graph: + if not self.config.create_attention_mask_in_dataloader: + if self.self_attention.attn_mask_type not in ( + AttnMaskType.causal, + AttnMaskType.no_mask, + AttnMaskType.causal_bottom_right, + ): + log_single_rank( + logger, + logging.WARNING, + "TE CUDA graph capture is omitting attention_mask because " + "create_attention_mask_in_dataloader is False, but " + f"attn_mask_type={self.self_attention.attn_mask_type.name} may require " + "an explicit mask. Ensure this is intended for the current workload.", ) + else: + 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(device) + .reshape(1, 1, slen_per_cp, seq_length) + .tile(micro_batch_size, 1, 1, 1) + ) # Add input_ids for hash-based MoE routing under CUDA graphs. # Only add for layers that actually use hash routing, @@ -1426,8 +1423,14 @@ 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. + # Pull hidden_states explicitly from cuda_graph_output rather than + # using `*cuda_graph_output, padding_mask=...`: the latter would + # collide if cuda_graph_output ever included a `padding_mask` + # positional element. + 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( - *cuda_graph_output, + hidden_states, padding_mask=kwargs.get("padding_mask", None), input_ids=kwargs.get("input_ids", None), ) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index adbf1c26ec7..494624690c1 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -214,7 +214,16 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): if 'position_ids' in batch: batch['position_ids'] = position_ids - return (*batch.values(), packed_seq_params, padding_mask) + # 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 @@ -310,7 +319,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: diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 4d42f47ce62..71dd7065d14 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1065,6 +1065,64 @@ 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 From b05683100e8c3fd3b3e86b2641f5ce8fb1072c6c Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Wed, 20 May 2026 06:02:29 -0700 Subject: [PATCH 05/25] refactor Signed-off-by: HaochenYuan --- .../models/common/embeddings/rope_utils.py | 56 +++++++++++-------- megatron/core/packed_seq_params.py | 37 +++++------- megatron/core/transformer/attention.py | 5 ++ .../absorbed_mla.py | 5 ++ .../transformer/multi_latent_attention.py | 5 ++ .../core/transformer/transformer_layer.py | 18 ++++-- 6 files changed, 78 insertions(+), 48 deletions(-) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index afb34a63c68..72a98e14c96 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -216,6 +216,7 @@ 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: """Apply RoPE for `thd` format using pure CUDA ops (CUDA Graph compatible). @@ -227,6 +228,7 @@ def _apply_rotary_pos_emb_thd( 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 [total_tokens, h, d]. Input with RoPE applied. @@ -246,42 +248,50 @@ def _apply_rotary_pos_emb_thd( total_tokens = t.shape[0] device = t.device - token_pos = torch.arange(total_tokens, device=device) - # `searchsorted(..., right=True) - 1` returns the index of the sequence each - # token belongs to. The `clamp` here guards padding tokens whose - # `searchsorted` position can fall outside the valid `[0, num_seqs)` range - # (last cu_seqlens entry); they get any wrong-but-harmless freq because - # their RoPE output is masked away later by `padding_mask` / `loss_mask`. - seq_idx = torch.searchsorted(cu_seqlens, token_pos, right=True) - 1 + 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) - seq_start = cu_seqlens[seq_idx] - local_pos = token_pos - seq_start - - # int64 for the multiplications: `cu_seqlens` are int32 (per TE convention), - # but `(seq_len * cp_size)` and downstream arithmetic can overflow int32 at - # very long contexts (e.g. >65k tokens × cp_size). Cast once here and the - # rest of the computation stays in int64. - seq_len_i64 = (cu_seqlens[seq_idx + 1] - seq_start).to(torch.int64) - full_seqlen = seq_len_i64 * cp_size - chunk_size = seq_len_i64 // 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: - is_first_half = local_pos < chunk_size + 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 * chunk_size + local_pos, - full_seqlen - (cp_rank + 1) * chunk_size + (local_pos - chunk_size), + 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) - if freqs.dim() >= 1 and freqs.size(0) > total_tokens: + if max_seqlen is None: + exact_packed_freqs = freqs.dim() >= 1 and cp_size == 1 and freqs.size(0) > total_tokens + else: + 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 + seq_start.to(torch.int64) * cp_size + freq_pos = freq_pos + global_seq_start # Same rationale as the seq_idx clamp above: padded positions can index # past `freqs`; they receive a known wrong-but-harmless freq that gets @@ -311,6 +321,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 @@ -386,6 +397,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 698dd752532..c2dcdbabb7a 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -172,21 +172,23 @@ def pad_thd_for_cuda_graph( actual_T = int(packed_seq_params.cu_seqlens_q[-1].item()) mask_device = packed_seq_params.cu_seqlens_q.device + from megatron.core import parallel_state + + cp_size = ( + packed_seq_params.local_cp_size + if packed_seq_params.local_cp_size is not None + else parallel_state.get_context_parallel_world_size() + ) + cp_rank = parallel_state.get_context_parallel_rank() if cp_size > 1 else 0 + max_seqlen_global = max_seqlen * cp_size + if actual_T is not None and 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 - from megatron.core import parallel_state - - _cp_size = ( - packed_seq_params.local_cp_size - if packed_seq_params.local_cp_size is not None - else parallel_state.get_context_parallel_world_size() - ) - _global_max_seqlen = max_seqlen * _cp_size - assert _max_individual <= _global_max_seqlen, ( + assert _max_individual <= max_seqlen_global, ( f"Individual request length ({_max_individual}) exceeds the global max sequence length " - f"({_global_max_seqlen} = max_seqlen_per_dp_cp_rank {max_seqlen} * cp_size {_cp_size}). " + f"({max_seqlen_global} = max_seqlen_per_dp_cp_rank {max_seqlen} * cp_size {cp_size}). " f"Each request must fit within the CUDA Graph static buffer after CP partitioning. " f"Increase --max-seqlen-per-dp-cp-rank or --seq-length, or filter out overlong requests." ) @@ -207,21 +209,12 @@ def pad_thd_for_cuda_graph( cu_seqlens_kv_padded=_pad_cu_seqlens( packed_seq_params.cu_seqlens_kv_padded, target_cu_entries ), - max_seqlen_q=max_seqlen, - max_seqlen_kv=max_seqlen, + max_seqlen_q=max_seqlen_global, + max_seqlen_kv=max_seqlen_global, local_cp_size=packed_seq_params.local_cp_size, cp_group=packed_seq_params.cp_group, ) - from megatron.core import parallel_state - - cp_size = ( - packed_seq_params.local_cp_size - if packed_seq_params.local_cp_size is not None - else parallel_state.get_context_parallel_world_size() - ) - cp_rank = parallel_state.get_context_parallel_rank() if cp_size > 1 else 0 - if cp_size > 1: from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices @@ -244,7 +237,7 @@ def pad_thd_for_cuda_graph( padded_params.cu_seqlens_q_padded if padded_params.cu_seqlens_q_padded is not None else padded_params.cu_seqlens_q, - max_seqlen, + max_seqlen_global, cp_size, cp_rank, ).numel() diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 3e61eb12a5f..ea636633647 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -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/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/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 07ad6e6b637..d626c84ad7d 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_layer.py b/megatron/core/transformer/transformer_layer.py index 1199ded9f77..688d08f1bdd 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1097,7 +1097,14 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): if self._is_thd_cuda_graph(): if attn_in_graph: - max_T = self.config.max_seqlen_per_dp_cp_rank + # Static cu_seqlens shaped [thd_max_num_seqs + 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_num_seqs cu_seqlens = torch.zeros(max_num_seqs + 1, dtype=torch.int32, device=device) cu_seqlens[1:] = max_T @@ -1206,7 +1213,7 @@ def _reconstruct_packed_seq_params_from_kwargs(self, kwargs): """ if 'cu_seqlens_q' not in kwargs: return - max_seqlen = self.config.max_seqlen_per_dp_cp_rank + 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'), @@ -1290,9 +1297,8 @@ def _te_cuda_graph_replay(self, *args, **kwargs): Hence, `inference_context` and `packed_seq_params` are excluded from input list. For THD format, PackedSeqParams is decomposed into individual tensor kwargs. """ - self._decompose_packed_seq_params_to_kwargs(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 @@ -1300,6 +1306,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 From 39fbb0e42cfeaa377a07ccbc424688c947a10f0b Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Tue, 26 May 2026 00:53:51 -0700 Subject: [PATCH 06/25] fix linting Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 9 +- megatron/core/packed_seq_params.py | 19 +- megatron/core/transformer/attention.py | 8 +- megatron/core/transformer/cuda_graphs.py | 6 +- megatron/core/transformer/module.py | 10 +- .../core/transformer/transformer_layer.py | 8 +- tests/unit_tests/test_sequence_packing.py | 8 +- .../transformer/test_cuda_graphs.py | 4 +- .../transformer/test_thd_cuda_graph.py | 363 +++++++++++++----- 9 files changed, 297 insertions(+), 138 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index a7181e58dc6..5fd5e4e97ac 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -123,9 +123,8 @@ 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 - and (self.max_num_seqs is None or len(single_microbatch) < self.max_num_seqs) + 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] @@ -585,9 +584,7 @@ def get_batch_on_this_rank_for_sequence_packing( # stage. cu_seqlens_padded keeps the pre-CP packed length; divide # by cp_size to match the already CP-sliced tokens/labels length. cp_world = cp_group.size() - total_tokens = ( - batch['cu_seqlens_padded'][-1].to(torch.int32) // cp_world - ).reshape(1) + 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) diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index c2dcdbabb7a..6a9e00cc7d3 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -190,7 +190,8 @@ def pad_thd_for_cuda_graph( f"Individual request length ({_max_individual}) exceeds the global max sequence length " f"({max_seqlen_global} = max_seqlen_per_dp_cp_rank {max_seqlen} * cp_size {cp_size}). " f"Each request must fit within the CUDA Graph static buffer after CP partitioning. " - f"Increase --max-seqlen-per-dp-cp-rank or --seq-length, or filter out overlong requests." + f"Increase --max-seqlen-per-dp-cp-rank or --seq-length, or filter out overlong " + f"requests." ) tokens = _pad_seq_tensor(tokens, max_seqlen) @@ -224,9 +225,11 @@ def pad_thd_for_cuda_graph( else: local_actual_T = int( get_thd_partitioned_indices( - 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, + ( + 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 + ), int(actual_T), cp_size, cp_rank, @@ -234,9 +237,11 @@ def pad_thd_for_cuda_graph( ) local_max_seqlen = int( get_thd_partitioned_indices( - padded_params.cu_seqlens_q_padded - if padded_params.cu_seqlens_q_padded is not None - else padded_params.cu_seqlens_q, + ( + padded_params.cu_seqlens_q_padded + if padded_params.cu_seqlens_q_padded is not None + else padded_params.cu_seqlens_q + ), max_seqlen_global, cp_size, cp_rank, diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index ea636633647..02483307d00 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -68,7 +68,9 @@ rearrange = None try: - from flash_attn_3.flash_attn_interface import _flash_attn_forward + 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, ) @@ -79,7 +81,9 @@ if not HAVE_FA3: try: - from flashattn_hopper.flash_attn_interface import _flash_attn_forward + from flashattn_hopper.flash_attn_interface import ( + _flash_attn_forward, + ) from flashattn_hopper.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 889a2c12e7e..3c7767d966e 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -2269,8 +2269,11 @@ def _get_cuda_graph_input_data(self): 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, @@ -2315,7 +2318,8 @@ def _get_cuda_graph_input_data(self): level=logging.INFO, msg=f'Rank {torch.distributed.get_rank()}: dynamic CUDA graph slots enabled. ' f'runtime_num_microbatches={get_num_microbatches()}, ' - f'auto_num_slots={auto_num_slots}, capture_num_microbatches={self.num_microbatches}', + f'auto_num_slots={auto_num_slots}, ' + f'capture_num_microbatches={self.num_microbatches}', ) else: self.num_microbatches = get_num_microbatches() diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index f981d22e719..39ab996e061 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -249,18 +249,16 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): 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." - ) + 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 - ) + 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. Hardcoding bfloat16 silently breaks --fp16. diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 688d08f1bdd..e616a0aff7d 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1102,9 +1102,7 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): # 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_T = self.config.max_seqlen_per_dp_cp_rank * self.config.context_parallel_size max_num_seqs = self.config.thd_max_num_seqs cu_seqlens = torch.zeros(max_num_seqs + 1, dtype=torch.int32, device=device) cu_seqlens[1:] = max_T @@ -1437,7 +1435,9 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): # using `*cuda_graph_output, padding_mask=...`: the latter would # collide if cuda_graph_output ever included a `padding_mask` # positional element. - assert len(cuda_graph_output) >= 1, "expected at least hidden_states in cuda_graph_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, diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index cf535e5a7db..fbb3028209d 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -212,10 +212,12 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc # Unpack the result. The helper now always returns a 7-tuple; the 7th # value is `padding_mask` (None when THD CUDA Graph padding is not in use). - tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = result - assert padding_mask is None, ( - "padding_mask should be None when config is not passed (legacy behavior)." + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = ( + result ) + assert ( + padding_mask is None + ), "padding_mask should be None when config is not passed (legacy behavior)." # Get parallel state info tp_rank = parallel_state.get_tensor_model_parallel_rank() diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 71dd7065d14..25e78fd0c91 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1077,9 +1077,7 @@ class TestRequiredNumMicrobatchSlots: @staticmethod def _slots(order, num_chunks): - return TECudaGraphHelper._get_required_num_microbatch_slots_from_order( - 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. diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index bd7e220cda7..e7a85d795ac 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -40,6 +40,7 @@ # 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): @@ -50,31 +51,46 @@ def _make_cu(seqlens, device="cuda"): 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)) + 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, - ) + 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, + 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_num_seqs=max_num_seqs, - tensor_model_parallel_size=tp, sequence_parallel=sp, bf16=True) + 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() + return ( + TransformerLayer( + config, get_gpt_layer_with_transformer_engine_spec().submodules, layer_number=1 + ) + .cuda() + .bfloat16() + ) # ============================================================================= # 1. pad_thd_for_cuda_graph correctness # ============================================================================= + class TestPadThdForCudaGraph: def setup_method(self): @@ -91,13 +107,22 @@ def test_shapes_and_data_preservation(self): 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_thd_for_cuda_graph( - tokens, tokens.clone(), torch.ones(1, total_T, device="cuda"), + tokens, + tokens.clone(), + torch.ones(1, total_T, device="cuda"), torch.arange(total_T, device="cuda").unsqueeze(0), - _make_psp(seqlens), max_seqlen, max_num_seqs) + _make_psp(seqlens), + max_seqlen, + 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): + 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 assert p_mask.shape == (1, max_seqlen) and p_mask.dtype == torch.bool assert torch.equal(p_tok[0, :total_T], tokens[0]) @@ -109,8 +134,14 @@ 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_thd_for_cuda_graph( - torch.ones(1, total_T, device="cuda"), None, None, None, - _make_psp(seqlens), max_seqlen, 4) + torch.ones(1, total_T, device="cuda"), + None, + None, + None, + _make_psp(seqlens), + max_seqlen, + 4, + ) assert not m[0, :total_T].any() and m[0, total_T:].all() @pytest.mark.internal @@ -119,8 +150,8 @@ def test_cu_seqlens_fill_value(self): """Padded entries repeat last cumulative sum (prevents OOB reads).""" seqlens, total_T = [50, 30], 80 _, _, _, _, p, _ = pad_thd_for_cuda_graph( - torch.ones(1, total_T, device="cuda"), None, None, None, - _make_psp(seqlens), 128, 32) + torch.ones(1, total_T, device="cuda"), None, None, None, _make_psp(seqlens), 128, 32 + ) assert p.cu_seqlens_q[0] == 0 and p.cu_seqlens_q[2] == 80 assert (p.cu_seqlens_q[3:] == 80).all() @@ -130,7 +161,8 @@ 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_thd_for_cuda_graph( - None, None, None, None, _make_psp(seqlens), max_seqlen, 4) + None, None, None, None, _make_psp(seqlens), max_seqlen, 4 + ) assert mask.shape == (1, max_seqlen) assert not mask[0, :total_T].any() and mask[0, total_T:].all() @@ -139,6 +171,7 @@ def test_none_inputs(self): # 2. PackedSeqParams decompose / reconstruct # ============================================================================= + class TestDecomposeReconstruct: def setup_method(self): @@ -152,8 +185,15 @@ def teardown_method(self): 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')} + 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) @@ -186,11 +226,12 @@ def test_noop_without_packed_seq_params(self): # Common args shared across both models (matches test_moonlight_qwen3_bitwise.sh). _MEGATRON_DIR = os.environ.get( - 'MEGATRON_DIR', - '/lustre/fsw/coreai_devtech_all/haocheny/migrate_to_TE_0415/Megatron-LM') + 'MEGATRON_DIR', '/lustre/fsw/coreai_devtech_all/haocheny/migrate_to_TE_0415/Megatron-LM' +) _MOONLIGHT_LOAD = os.environ.get( 'MOONLIGHT_CKPT', - '/lustre/fsw/coreai_devtech_all/haocheny/mcore_models/Moonlight-16B-A3B-Instruct') + '/lustre/fsw/coreai_devtech_all/haocheny/mcore_models/Moonlight-16B-A3B-Instruct', +) _SFT_JSON = ( '{"mode":"distribution","type":"lognormal",' @@ -200,66 +241,161 @@ def test_noop_without_packed_seq_params(self): _TRAIN_ITERS = 5 _COMMON_ARGS = [ - "--seq-length", "2048", "--max-position-embeddings", "8192", - "--micro-batch-size", "1", "--global-batch-size", "4", - "--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", - "--sft", "--mock-data", - "--tokenizer-type", "NullTokenizer", - "--sft-mock-dataset-config-json", _SFT_JSON, - "--sequence-packing-scheduler", "dp_balanced", - "--max-seqlen-per-dp-cp-rank", "1024", + "--seq-length", + "2048", + "--max-position-embeddings", + "8192", + "--micro-batch-size", + "1", + "--global-batch-size", + "4", + "--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", + "--sft", + "--mock-data", + "--tokenizer-type", + "NullTokenizer", + "--sft-mock-dataset-config-json", + _SFT_JSON, + "--sequence-packing-scheduler", + "dp_balanced", + "--max-seqlen-per-dp-cp-rank", + "1024", "--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-num-seqs", "32", + "--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-num-seqs", + "32", ] _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", + "--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", "alltoall", - "--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", - "--load", _MOONLIGHT_LOAD, - "--no-load-optim", "--no-load-rng", + "--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", + "alltoall", + "--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", + "--load", + _MOONLIGHT_LOAD, + "--no-load-optim", + "--no-load-rng", ] _QWEN3_ARGS = _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", + "--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", "alltoall", + "--vocab-size", + "151936", + "--moe-token-dispatcher-type", + "alltoall", ] @@ -273,8 +409,17 @@ def _run_pretrain(model_args, cuda_graph_args, master_port): 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")): + 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 @@ -282,15 +427,25 @@ def _run_pretrain(model_args, cuda_graph_args, master_port): 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(master_port), - "pretrain_gpt.py", - ] + model_args + cuda_graph_args + cmd = ( + [ + "torchrun", + "--nproc_per_node", + "8", + "--nnodes", + "1", + "--master_addr", + "localhost", + "--master_port", + str(master_port), + "pretrain_gpt.py", + ] + + model_args + + cuda_graph_args + ) result = subprocess.run( - cmd, cwd=_MEGATRON_DIR, env=env, capture_output=True, text=True, - timeout=900, + cmd, cwd=_MEGATRON_DIR, env=env, capture_output=True, text=True, timeout=900 ) return result @@ -310,18 +465,14 @@ def _extract_metrics(stdout): """ results = [] for m in _ITER_START_RE.finditer(stdout): - window = stdout[m.start():m.start() + 800] + 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) lb_loss = re.search(r"load_balancing_loss:\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 = [f"iter={m.group(1)}", f"lr={lr.group(1)}", f"lm_loss={lm_loss.group(1)}"] if lb_loss: parts.append(f"lb_loss={lb_loss.group(1)}") parts.append(f"grad_norm={grad_norm.group(1)}") @@ -334,10 +485,7 @@ def _extract_metrics(stdout): @pytest.mark.skipif(torch.cuda.device_count() < 8, reason="requires 8 GPUs") @pytest.mark.parametrize( "model_name,model_args,base_port", - [ - ("moonlight", _MOONLIGHT_ARGS, 29660), - ("qwen3", _QWEN3_ARGS, 29662), - ], + [("moonlight", _MOONLIGHT_ARGS, 29660), ("qwen3", _QWEN3_ARGS, 29662)], ) class TestE2EBitwise: """End-to-end bitwise comparison: pretrain_gpt.py noGraph vs cudaGraph. @@ -357,35 +505,38 @@ def test_no_graph_vs_graph(self, model_name, model_args, 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:]}") + 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:]}") + f"--- stdout (tail) ---\n{r1.stdout[-2000:]}" + ) # CUDA graph capture. r2 = _run_pretrain( model_args, cuda_graph_args=[ - "--cuda-graph-impl", "transformer_engine", - "--cuda-graph-scope", "attn", + "--cuda-graph-impl", + "transformer_engine", + "--cuda-graph-scope", + "attn", ], 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:]}") + 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:]}") + 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}") + assert a == b, f"[{model_name}] iter {i+1} differs:\n" f" eager: {a}\n" f" graph: {b}" From 258bd1e30b377f9bf2cffa5eec2b6df0f2719b6d Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Tue, 26 May 2026 01:00:02 -0700 Subject: [PATCH 07/25] fix linting Signed-off-by: HaochenYuan --- megatron/core/transformer/attention.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 02483307d00..d367742f672 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -68,12 +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: From cc928a92c101a2b9261cc047bae68af15fc12a39 Mon Sep 17 00:00:00 2001 From: HaochenYuan <106647990+HaochenYuan@users.noreply.github.com> Date: Tue, 26 May 2026 16:05:52 +0800 Subject: [PATCH 08/25] fix linting Removed redundant import of _flash_attn_forward. Signed-off-by: HaochenYuan --- megatron/core/transformer/attention.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index d367742f672..e712122ef69 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -79,9 +79,7 @@ if not HAVE_FA3: try: - from flashattn_hopper.flash_attn_interface import ( - _flash_attn_forward, - ) + from flashattn_hopper.flash_attn_interface import _flash_attn_forward from flashattn_hopper.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) From 4e016666786b5fbc1131e0636bdaaf7fe6c85558 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Wed, 27 May 2026 00:52:32 -0700 Subject: [PATCH 09/25] fix CI Signed-off-by: HaochenYuan --- .../text/libraries/null_tokenizer.py | 4 +-- .../models/test_hybrid_moe_model.py | 3 +++ .../transformer/test_thd_cuda_graph.py | 25 ++++++------------- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/megatron/core/tokenizers/text/libraries/null_tokenizer.py b/megatron/core/tokenizers/text/libraries/null_tokenizer.py index 47b43dbe577..160aaa8bcb3 100644 --- a/megatron/core/tokenizers/text/libraries/null_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/null_tokenizer.py @@ -93,8 +93,8 @@ def eod(self): @property def pad_id(self): - """Returns id of padding token (same as eod for NullTokenizer).""" - return self._eod_id + """Returns id of padding token.""" + return self._pad_id @property def additional_special_tokens_ids(self): diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index f3d5e47a103..16c1de92122 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -79,6 +79,8 @@ "cuda_graph_retain_backward_graph": False, "cuda_graph_modules": [], "cuda_graph_use_single_mempool": True, + "cuda_graph_dynamic_microbatches": False, + "cuda_graph_num_microbatch_slots": None, "cuda_graph_scope": None, "cuda_graph_warmup_steps": 3, "deallocate_pipeline_outputs": True, @@ -281,6 +283,7 @@ "symmetric_ar_type": None, "tensor_model_parallel_size": 2, "test_mode": False, + "thd_max_num_seqs": 32, "timers": None, "tp_comm_atomic_ag": False, "tp_comm_atomic_rs": False, diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index e7a85d795ac..73477c6a54d 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -22,6 +22,7 @@ import os import re import subprocess +from pathlib import Path import pytest import torch @@ -224,14 +225,8 @@ def test_noop_without_packed_seq_params(self): # metric strings are byte-identical between the two runs. # ============================================================================= -# Common args shared across both models (matches test_moonlight_qwen3_bitwise.sh). -_MEGATRON_DIR = os.environ.get( - 'MEGATRON_DIR', '/lustre/fsw/coreai_devtech_all/haocheny/migrate_to_TE_0415/Megatron-LM' -) -_MOONLIGHT_LOAD = os.environ.get( - 'MOONLIGHT_CKPT', - '/lustre/fsw/coreai_devtech_all/haocheny/mcore_models/Moonlight-16B-A3B-Instruct', -) +# Common args shared across both models. +_REPO_ROOT = Path(__file__).resolve().parents[3] _SFT_JSON = ( '{"mode":"distribution","type":"lognormal",' @@ -365,10 +360,6 @@ def test_noop_without_packed_seq_params(self): "50000", "--vocab-size", "163840", - "--load", - _MOONLIGHT_LOAD, - "--no-load-optim", - "--no-load-rng", ] _QWEN3_ARGS = _COMMON_ARGS + [ @@ -402,7 +393,7 @@ def test_noop_without_packed_seq_params(self): 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"] = _MEGATRON_DIR + ":" + env.get("PYTHONPATH", "") + 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" @@ -445,7 +436,7 @@ def _run_pretrain(model_args, cuda_graph_args, master_port): ) result = subprocess.run( - cmd, cwd=_MEGATRON_DIR, env=env, capture_output=True, text=True, timeout=900 + cmd, cwd=_REPO_ROOT, env=env, capture_output=True, text=True, timeout=900 ) return result @@ -491,8 +482,8 @@ 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 cuda_graph_scope=attn` - -- using the exact same args as test_moonlight_qwen3_bitwise.sh. + graphs and once with `cuda_graph_impl=transformer_engine cuda_graph_modules=attn` + -- using the same model/test settings as test_moonlight_qwen3_bitwise.sh. Asserts the per-iteration `lm loss / load_balancing_loss / grad norm` lines are byte-identical. @@ -520,7 +511,7 @@ def test_no_graph_vs_graph(self, model_name, model_args, base_port): cuda_graph_args=[ "--cuda-graph-impl", "transformer_engine", - "--cuda-graph-scope", + "--cuda-graph-modules", "attn", ], master_port=base_port + 1, From 60612bab1ec5dc11eb5478c7219c8d7b08bcc780 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Wed, 27 May 2026 03:53:59 -0700 Subject: [PATCH 10/25] change UT cp size to avoid OOM Signed-off-by: HaochenYuan --- tests/unit_tests/transformer/test_thd_correctness.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/transformer/test_thd_correctness.py b/tests/unit_tests/transformer/test_thd_correctness.py index 533f64081f4..7ee3ff1c537 100644 --- a/tests/unit_tests/transformer/test_thd_correctness.py +++ b/tests/unit_tests/transformer/test_thd_correctness.py @@ -111,7 +111,7 @@ 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_cp4_sp_longseq", 4096, 32, 8, 14336, [65536, 8191, 4096], 2, 4, True, "similarity"), # ------------------------------------------------------------------------- # Edge cases From b061b7509709331d895611f02acf6954bdb32edc Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Wed, 27 May 2026 10:46:29 -0700 Subject: [PATCH 11/25] shorten the UT seqlen Signed-off-by: HaochenYuan --- tests/unit_tests/transformer/test_thd_correctness.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/transformer/test_thd_correctness.py b/tests/unit_tests/transformer/test_thd_correctness.py index 7ee3ff1c537..560f61a4ede 100644 --- a/tests/unit_tests/transformer/test_thd_correctness.py +++ b/tests/unit_tests/transformer/test_thd_correctness.py @@ -111,7 +111,7 @@ 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_cp4_sp_longseq", 4096, 32, 8, 14336, [65536, 8191, 4096], 2, 4, True, "similarity"), + TestCase("tp2_cp2_sp_longseq", 4096, 32, 8, 14336, [8192, 4096, 2048], 2, 2, True, "similarity"), # ------------------------------------------------------------------------- # Edge cases From 23d88387343a4ee5ff7949be09f54195554828f0 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Wed, 27 May 2026 23:47:53 -0700 Subject: [PATCH 12/25] fix UT Signed-off-by: HaochenYuan --- tests/unit_tests/transformer/test_thd_correctness.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/transformer/test_thd_correctness.py b/tests/unit_tests/transformer/test_thd_correctness.py index 560f61a4ede..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, [8192, 4096, 2048], 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 From 685e986da9709483562f51cce9c2ef8b7e56db22 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 8 Jun 2026 07:26:37 -0700 Subject: [PATCH 13/25] refactor padding Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 44 ++-- megatron/core/datasets/data_schedule_utils.py | 5 +- megatron/core/model_parallel_config.py | 27 +++ .../models/common/embeddings/rope_utils.py | 6 +- megatron/core/packed_seq_params.py | 205 ++++++++++++------ megatron/core/transformer/module.py | 2 +- .../core/transformer/transformer_config.py | 15 ++ .../core/transformer/transformer_layer.py | 7 +- megatron/training/arguments.py | 33 +++ pretrain_gpt.py | 45 +++- tests/unit_tests/test_sequence_packing.py | 5 +- .../transformer/test_thd_cuda_graph.py | 59 +++-- 12 files changed, 325 insertions(+), 128 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 5fd5e4e97ac..f75daa675be 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -19,7 +19,7 @@ next_hdp_group, reroute_samples_to_dcp_ranks, ) -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, 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 @@ -502,8 +502,8 @@ 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 THD CUDA Graph padding. When None - or config.max_seqlen_per_dp_cp_rank is None, no padding is applied. + config: Model parallel config used for optional THD packed-sequence padding. + When None or config.pad_packed_seq_alignment is None, no padding is applied. Returns: tuple of (tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask) @@ -676,9 +676,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, @@ -691,24 +690,35 @@ def get_batch_on_this_rank_for_sequence_packing( cp_group=cp_group, ) - # Pad to static shapes for THD + CUDA Graph when requested. + # Pad the already-packed THD tensors at the end when requested. CUDA Graph + # additionally pads cu_seqlens tensors to thd_max_num_seqs + 1 entries. padding_mask = None - if ( - config is not None - and getattr(config, 'max_seqlen_per_dp_cp_rank', None) is not None - and packed_seq_params is not None - ): - from megatron.core.packed_seq_params import pad_thd_for_cuda_graph - + 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: + cuda_graph_static = getattr(config, 'cuda_graph_impl', 'none') != 'none' + static_target = ( + pad_alignment == 0 and getattr(config, 'max_seqlen_per_dp_cp_rank', None) is not None + ) + if cuda_graph_static or static_target: + target_len = config.max_seqlen_per_dp_cp_rank + max_num_seqs = config.thd_max_num_seqs + alignment = None + else: + target_len = None + max_num_seqs = None + alignment = max_seqlen if pad_alignment == 0 else pad_alignment tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = ( - pad_thd_for_cuda_graph( + pad_sequence_for_thd( tokens, labels, loss_mask, position_ids, packed_seq_params, - max_seqlen=config.max_seqlen_per_dp_cp_rank, - max_num_seqs=config.thd_max_num_seqs, + alignment=alignment, + target_len=target_len, + max_num_seqs=max_num_seqs, ) ) diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index eaeaa1b8ac5..fc906ebf725 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -25,9 +25,8 @@ def get_cp_slice_for_thd(batch, cp_group): if cp_size <= 1: return cp_rank = cp_group.rank() - # 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 diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index d5e8f721297..541faa2ce68 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -87,6 +87,25 @@ class ModelParallelConfig: default_dynamic_cp: Dynamic-CP scheduler for packed sequence balancing. """ + pad_packed_seq_alignment: Optional[int] = field( + default=None, + metadata={ + "argparse_meta": { + "arg_names": ["--pad-packed-seq-alignment"], + "nargs": "?", + "const": 0, + "type": int, + } + }, + ) + """Pad THD packed sequence tensors after packing. + + If set without a value, the caller chooses a static target; the standard THD + training paths use max_seqlen_per_dp_cp_rank and pad cu_seqlens tensors to + thd_max_num_seqs + 1 entries. If set to a positive integer N, token-like + tensors are padded to a multiple of N and cu_seqlens metadata is preserved. + """ + expert_model_parallel_size: int = 1 """Distributes Moe Experts across sub data parallel dimension.""" @@ -469,6 +488,14 @@ def __post_init__(self): f"got {self.min_dynamic_context_parallel_size}" ) + if self.pad_packed_seq_alignment is not None: + if self.pad_packed_seq_alignment < 0: + raise ValueError( + "pad_packed_seq_alignment must be >= 0. Use the flag without a value " + "to let the caller choose a static target, or pass a positive integer " + "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 72a98e14c96..aa6d9c8a634 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -293,10 +293,8 @@ def _apply_rotary_pos_emb_thd( # sequence, no shift is needed. freq_pos = freq_pos + global_seq_start - # Same rationale as the seq_idx clamp above: padded positions can index - # past `freqs`; they receive a known wrong-but-harmless freq that gets - # masked away. If you suspect a real out-of-range bug, swap clamp for an - # assert during development. + # 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] diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 6a9e00cc7d3..787016c8796 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -104,10 +104,9 @@ def _pad_seq_tensor(t: Optional[Tensor], target_len: int) -> Optional[Tensor]: 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``: this is - the reviewer-flagged overflow case and corresponds to "too many packed - sequences in this microbatch for thd_max_num_seqs". Failing fast prevents - a silent CUDA-graph shape mismatch at replay. + 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 @@ -128,33 +127,21 @@ def _pad_cu_seqlens(cu_seqlens: Optional[Tensor], target_entries: int) -> Option return padded -def pad_thd_for_cuda_graph( +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 _resolve_thd_padding_lengths( tokens: Optional[Tensor], labels: Optional[Tensor], loss_mask: Optional[Tensor], position_ids: Optional[Tensor], packed_seq_params: PackedSeqParams, - max_seqlen: int, - max_num_seqs: int, -) -> Tuple[ - Optional[Tensor], - Optional[Tensor], - Optional[Tensor], - Optional[Tensor], - PackedSeqParams, - Optional[Tensor], -]: - """Pad THD batch data to fixed sizes for CUDA Graph compatibility. - - CUDA Graph requires static tensor shapes. This function pads: - - tokens, labels, loss_mask, position_ids along dim=-1 to max_seqlen - - cu_seqlens tensors to (max_num_seqs + 1) entries, filled with actual_T - - Generates padding_mask for MoE aux loss exclusion - - Returns: - Padded (tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask) - padding_mask: [1, max_seqlen] bool tensor, True at padding positions. - """ + target_len: Optional[int], + alignment: Optional[int], +) -> Tuple[int, int, int, int, torch.device, bool]: + """Resolve local/global THD padding lengths without changing tensors.""" actual_T = None mask_device = None @@ -180,48 +167,23 @@ def pad_thd_for_cuda_graph( else parallel_state.get_context_parallel_world_size() ) cp_rank = parallel_state.get_context_parallel_rank() if cp_size > 1 else 0 - max_seqlen_global = max_seqlen * cp_size - - if actual_T is not None and 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 <= max_seqlen_global, ( - f"Individual request length ({_max_individual}) exceeds the global max sequence length " - f"({max_seqlen_global} = max_seqlen_per_dp_cp_rank {max_seqlen} * cp_size {cp_size}). " - f"Each request must fit within the CUDA Graph static buffer after CP partitioning. " - f"Increase --max-seqlen-per-dp-cp-rank or --seq-length, or filter out overlong " - f"requests." - ) - - tokens = _pad_seq_tensor(tokens, max_seqlen) - labels = _pad_seq_tensor(labels, max_seqlen) - loss_mask = _pad_seq_tensor(loss_mask, max_seqlen) - position_ids = _pad_seq_tensor(position_ids, max_seqlen) - target_cu_entries = max_num_seqs + 1 - padded_params = PackedSeqParams( - qkv_format=packed_seq_params.qkv_format, - cu_seqlens_q=_pad_cu_seqlens(packed_seq_params.cu_seqlens_q, target_cu_entries), - cu_seqlens_kv=_pad_cu_seqlens(packed_seq_params.cu_seqlens_kv, target_cu_entries), - cu_seqlens_q_padded=_pad_cu_seqlens( - packed_seq_params.cu_seqlens_q_padded, target_cu_entries - ), - cu_seqlens_kv_padded=_pad_cu_seqlens( - packed_seq_params.cu_seqlens_kv_padded, target_cu_entries - ), - max_seqlen_q=max_seqlen_global, - max_seqlen_kv=max_seqlen_global, - local_cp_size=packed_seq_params.local_cp_size, - cp_group=packed_seq_params.cp_group, - ) + if target_len is None: + assert alignment is not None, "Either target_len or alignment must be provided." + global_target_len = _round_up_to_alignment(int(actual_T), alignment) + else: + global_target_len = int(target_len) * cp_size if cp_size > 1: from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices if actual_T_is_local: local_actual_T = int(actual_T) - local_max_seqlen = int(max_seqlen) + local_target_len = ( + int(target_len) + if target_len is not None + else _round_up_to_alignment(local_actual_T, alignment) + ) else: local_actual_T = int( get_thd_partitioned_indices( @@ -235,22 +197,125 @@ def pad_thd_for_cuda_graph( cp_rank, ).numel() ) - local_max_seqlen = int( + local_target_len = int( get_thd_partitioned_indices( ( - padded_params.cu_seqlens_q_padded - if padded_params.cu_seqlens_q_padded is not None - else padded_params.cu_seqlens_q + 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 ), - max_seqlen_global, + global_target_len, cp_size, cp_rank, ).numel() ) - padding_mask = ( - torch.arange(local_max_seqlen, device=mask_device).unsqueeze(0) >= local_actual_T - ) else: - padding_mask = torch.arange(max_seqlen, device=mask_device).unsqueeze(0) >= actual_T + local_actual_T = int(actual_T) + local_target_len = global_target_len + + return ( + int(actual_T), + local_actual_T, + local_target_len, + global_target_len, + mask_device, + actual_T_is_local, + ) + + +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, +) -> 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. When ``max_num_seqs`` is provided, the + four cu_seqlens tensors are also padded to ``max_num_seqs + 1`` entries; + 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." + ) + + actual_T, local_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, + ) + ) + + if actual_T is not None and 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." + ) + + 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) + + target_cu_entries = None if max_num_seqs is None else max_num_seqs + 1 + padded_params = PackedSeqParams( + qkv_format=packed_seq_params.qkv_format, + cu_seqlens_q=( + packed_seq_params.cu_seqlens_q + if target_cu_entries is None + else _pad_cu_seqlens(packed_seq_params.cu_seqlens_q, target_cu_entries) + ), + cu_seqlens_kv=( + packed_seq_params.cu_seqlens_kv + if target_cu_entries is None + else _pad_cu_seqlens(packed_seq_params.cu_seqlens_kv, target_cu_entries) + ), + cu_seqlens_q_padded=( + packed_seq_params.cu_seqlens_q_padded + if target_cu_entries is None + else _pad_cu_seqlens(packed_seq_params.cu_seqlens_q_padded, target_cu_entries) + ), + cu_seqlens_kv_padded=( + packed_seq_params.cu_seqlens_kv_padded + if target_cu_entries is None + else _pad_cu_seqlens(packed_seq_params.cu_seqlens_kv_padded, target_cu_entries) + ), + max_seqlen_q=( + global_target_len if target_cu_entries is not None else packed_seq_params.max_seqlen_q + ), + max_seqlen_kv=( + global_target_len if target_cu_entries is not None else packed_seq_params.max_seqlen_kv + ), + 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, + ) + + padding_mask = torch.arange(local_target_len, device=mask_device).unsqueeze(0) >= local_actual_T return tokens, labels, loss_mask, position_ids, padded_params, padding_mask diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index 39ab996e061..c5211e3e6d6 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -261,7 +261,7 @@ def get_layer_static_inputs(self, seq_length, 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. Hardcoding bfloat16 silently breaks --fp16. + # through the captured graph. if self.config.bf16: dtype = torch.bfloat16 elif self.config.fp16: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 90b83e106a8..fa809c4ba88 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -3021,6 +3021,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 == 0 + or self.pad_packed_seq_alignment == self.max_seqlen_per_dp_cp_rank + ), ( + "THD CUDA Graph requires --pad-packed-seq-alignment without a value " + "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 e616a0aff7d..54f280efdbe 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1431,10 +1431,9 @@ 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. - # Pull hidden_states explicitly from cuda_graph_output rather than - # using `*cuda_graph_output, padding_mask=...`: the latter would - # collide if cuda_graph_output ever included a `padding_mask` - # positional element. + # 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" diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 29e21544599..66199fb7cf8 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1571,6 +1571,39 @@ def validate_args(args, defaults={}): f"to {args.data_parallel_size * args.context_parallel_size}." ) + if args.sequence_packing_scheduler is not None: + if args.sequence_packing_scheduler == 'dp_balanced': + total_cp_ranks = args.context_parallel_size + else: + total_cp_ranks = args.data_parallel_size * args.context_parallel_size + assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, ( + f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) ' + f'must be >= single sequence max length ({args.seq_length})' + ) + + if getattr(args, 'pad_packed_seq_alignment', None) is not None: + if args.pad_packed_seq_alignment < 0: + raise ValueError( + '--pad-packed-seq-alignment must be used without a value or with a ' + 'non-negative integer 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 != 0 + and args.pad_packed_seq_alignment != args.max_seqlen_per_dp_cp_rank + ): + raise ValueError( + 'THD CUDA Graph requires --pad-packed-seq-alignment without a value ' + '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/pretrain_gpt.py b/pretrain_gpt.py index 494624690c1..6fecc802855 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -30,7 +30,7 @@ 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, pad_thd_for_cuda_graph +from megatron.core.packed_seq_params import PackedSeqParams, 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,9 +125,8 @@ 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` applies THD + CUDA Graph - # padding internally when `config.max_seqlen_per_dp_cp_rank` is set, and - # returns a 7-tuple including `padding_mask` (None when no padding). + # `get_batch_on_this_rank_for_sequence_packing` owns optional THD padding + # and returns a 7-tuple including `padding_mask` (None when no padding). return get_batch_on_this_rank_for_sequence_packing( data_iterator, vpp_size=config.virtual_pipeline_model_parallel_size, @@ -192,19 +191,43 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): batch, cu_seqlens, cu_seqlens_padded, max_seqlen ) - # Pad THD batch for CUDA Graph compatibility when max_seqlen_per_dp_cp_rank is set. + # Pad the already-packed THD tensors at the end when requested. CUDA Graph + # additionally pads cu_seqlens tensors to thd_max_num_seqs + 1 entries. padding_mask = None - if config.max_seqlen_per_dp_cp_rank is not None and packed_seq_params is not 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) - tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = \ - pad_thd_for_cuda_graph( - tokens, labels, loss_mask, position_ids, packed_seq_params, - max_seqlen=config.max_seqlen_per_dp_cp_rank, - max_num_seqs=config.thd_max_num_seqs, + cuda_graph_static = config.cuda_graph_impl != "none" + static_target = ( + config.pad_packed_seq_alignment == 0 + and config.max_seqlen_per_dp_cp_rank is not None + ) + if cuda_graph_static or static_target: + target_len = config.max_seqlen_per_dp_cp_rank + max_num_seqs = config.thd_max_num_seqs + alignment = None + else: + target_len = None + max_num_seqs = None + alignment = ( + int(max_seqlen[0].item()) + if config.pad_packed_seq_alignment == 0 + else config.pad_packed_seq_alignment + ) + 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, ) + ) if 'tokens' in batch: batch['tokens'] = tokens if 'labels' in batch: diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index fbb3028209d..95bc5d2ea29 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -210,14 +210,13 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc dynamic_cp=dynamic_cp, ) - # Unpack the result. The helper now always returns a 7-tuple; the 7th - # value is `padding_mask` (None when THD CUDA Graph padding is not in use). + # The helper returns a 7-tuple; padding_mask is None when THD padding is disabled. tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = ( result ) assert ( padding_mask is None - ), "padding_mask should be None when config is not passed (legacy behavior)." + ), "padding_mask should be None when no padding config is provided." # Get parallel state info tp_rank = parallel_state.get_tensor_model_parallel_rank() diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 73477c6a54d..4f11b56d731 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -27,7 +27,10 @@ import pytest import torch -from megatron.core.packed_seq_params import PackedSeqParams, pad_thd_for_cuda_graph +from megatron.core.packed_seq_params import ( + PackedSeqParams, + 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 @@ -88,11 +91,11 @@ def _build_layer(H, nh, nkv, ffn, max_seqlen, max_num_seqs, tp=1, sp=False): # ============================================================================= -# 1. pad_thd_for_cuda_graph correctness +# 1. pad_sequence_for_thd correctness # ============================================================================= -class TestPadThdForCudaGraph: +class TestPadSequenceForThd: def setup_method(self): Utils.initialize_model_parallel(tensor_model_parallel_size=1) @@ -100,6 +103,21 @@ def setup_method(self): 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_preserves_cu_seqlens(self): + """Generic THD padding aligns token tensors while preserving sequence metadata.""" + 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) + assert torch.equal(p.cu_seqlens_q, orig) + 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): @@ -107,14 +125,14 @@ def test_shapes_and_data_preservation(self): 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_thd_for_cuda_graph( + 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), - max_seqlen, - max_num_seqs, + 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) @@ -134,14 +152,14 @@ def test_shapes_and_data_preservation(self): 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_thd_for_cuda_graph( + _, _, _, _, _, m = pad_sequence_for_thd( torch.ones(1, total_T, device="cuda"), None, None, None, _make_psp(seqlens), - max_seqlen, - 4, + target_len=max_seqlen, + max_num_seqs=4, ) assert not m[0, :total_T].any() and m[0, total_T:].all() @@ -150,8 +168,14 @@ def test_padding_mask_boundary(self): def test_cu_seqlens_fill_value(self): """Padded entries repeat last cumulative sum (prevents OOB reads).""" seqlens, total_T = [50, 30], 80 - _, _, _, _, p, _ = pad_thd_for_cuda_graph( - torch.ones(1, total_T, device="cuda"), None, None, None, _make_psp(seqlens), 128, 32 + _, _, _, _, 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:] == 80).all() @@ -161,8 +185,8 @@ def test_cu_seqlens_fill_value(self): 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_thd_for_cuda_graph( - None, None, None, None, _make_psp(seqlens), max_seqlen, 4 + _, _, _, _, _, 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() @@ -281,6 +305,7 @@ def test_noop_without_packed_seq_params(self): "dp_balanced", "--max-seqlen-per-dp-cp-rank", "1024", + "--pad-packed-seq-alignment", "--calculate-per-token-loss", "--transformer-impl", "transformer_engine", @@ -343,7 +368,9 @@ def test_noop_without_packed_seq_params(self): "--moe-layer-freq", "([0]+[1]*26)", "--moe-token-dispatcher-type", - "alltoall", + "flex", + "--moe-flex-dispatcher-backend", + "hybridep", "--moe-router-score-function", "sigmoid", "--moe-router-topk-scaling-factor", @@ -386,7 +413,9 @@ def test_noop_without_packed_seq_params(self): "--vocab-size", "151936", "--moe-token-dispatcher-type", - "alltoall", + "flex", + "--moe-flex-dispatcher-backend", + "hybridep", ] From 0b0e9fc65142c4b56ba0246a48f7f42ca928cdb8 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Wed, 10 Jun 2026 07:25:36 -0700 Subject: [PATCH 14/25] skip length align when use cuda graph Signed-off-by: HaochenYuan --- .../core/transformer/moe/token_dispatcher.py | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index e8f2044650e..acfb9eaeacc 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1070,15 +1070,27 @@ 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() + ): + # Capture 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() (forbidden during 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 From 597f1d742c343500d4db3c664d7c2ecdd972c387 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Thu, 11 Jun 2026 01:13:35 -0700 Subject: [PATCH 15/25] remove redundancy Signed-off-by: HaochenYuan --- megatron/core/models/gpt/gpt_model.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 7f284812179..01df346c05c 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -331,23 +331,7 @@ def _preprocess( # Decoder embedding. if decoder_input is not None: - # For non-pre_process PP stages that receive decoder_input, scatter padding_mask - # to match the sequence-parallel partitioned hidden_states if needed. - if ( - padding_mask is not None - and self.config.sequence_parallel - and padding_mask.shape[1] != decoder_input.shape[0] - and padding_mask.shape[1] % self.config.tensor_model_parallel_size == 0 - and padding_mask.shape[1] // self.config.tensor_model_parallel_size - == decoder_input.shape[0] - ): - padding_mask = ( - tensor_parallel.scatter_to_sequence_parallel_region( - padding_mask.transpose(0, 1).contiguous() - ) - .transpose(0, 1) - .contiguous() - ) + pass elif self.pre_process: if padding_mask is not None: assert padding_mask.shape == input_ids.shape, ( From 6224b98fcaa6892b1d79763c6685945de6a04954 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Thu, 11 Jun 2026 23:58:33 -0700 Subject: [PATCH 16/25] Apply THD CUDA graph refactor fixes Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 70 +++- megatron/core/model_parallel_config.py | 62 +++- megatron/core/packed_seq_params.py | 281 +++++++++----- megatron/core/transformer/cuda_graphs.py | 11 +- megatron/core/transformer/moe/moe_utils.py | 10 +- .../core/transformer/transformer_config.py | 34 +- .../core/transformer/transformer_layer.py | 1 + megatron/training/arguments.py | 26 +- pretrain_gpt.py | 28 +- .../models/test_hybrid_moe_model.py | 1 - tests/unit_tests/test_sequence_packing.py | 29 ++ .../transformer/test_thd_cuda_graph.py | 344 +++++++++++++++++- 12 files changed, 702 insertions(+), 195 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index f75daa675be..c65673740a3 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -19,7 +19,11 @@ next_hdp_group, reroute_samples_to_dcp_ranks, ) -from megatron.core.packed_seq_params import PackedSeqParams, pad_sequence_for_thd +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 @@ -43,9 +47,9 @@ 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 packed sequences per - microbatch. When set, the scheduler closes a pack as soon as it - reaches this many sequences in addition to the token-budget condition. + 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 @@ -415,6 +419,36 @@ 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_num_seqs`` 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_num_seqs', 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_num_seqs 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_num_seqs must be >= 2 when THD padding appends a dummy " + "sequence, because thd_max_num_seqs 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 ): @@ -457,6 +491,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_num_seqs', None) + ) + scheduler = scheduler_map[scheduler_type]( config.max_seqlen_per_dp_cp_rank, cp_size, @@ -466,7 +506,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=getattr(config, 'thd_max_num_seqs', None), + max_num_seqs=scheduler_max_num_seqs, **scheduler_kwargs, ) @@ -688,6 +728,7 @@ 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 @@ -697,18 +738,12 @@ def get_batch_on_this_rank_for_sequence_packing( 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: - cuda_graph_static = getattr(config, 'cuda_graph_impl', 'none') != 'none' - static_target = ( - pad_alignment == 0 and getattr(config, 'max_seqlen_per_dp_cp_rank', None) 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_num_seqs', None), + getattr(config, 'cuda_graph_impl', 'none') != 'none', ) - if cuda_graph_static or static_target: - target_len = config.max_seqlen_per_dp_cp_rank - max_num_seqs = config.thd_max_num_seqs - alignment = None - else: - target_len = None - max_num_seqs = None - alignment = max_seqlen if pad_alignment == 0 else pad_alignment tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = ( pad_sequence_for_thd( tokens, @@ -719,6 +754,9 @@ def get_batch_on_this_rank_for_sequence_packing( 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 + ), ) ) diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 541faa2ce68..d65adaa966e 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,23 +102,29 @@ class ModelParallelConfig: default_dynamic_cp: Dynamic-CP scheduler for packed sequence balancing. """ - pad_packed_seq_alignment: Optional[int] = field( + pad_packed_seq_alignment: Optional[Union[int, Literal["max"]]] = field( default=None, metadata={ "argparse_meta": { "arg_names": ["--pad-packed-seq-alignment"], - "nargs": "?", - "const": 0, - "type": int, + "type": _parse_pad_packed_seq_alignment, } }, ) """Pad THD packed sequence tensors after packing. - If set without a value, the caller chooses a static target; the standard THD - training paths use max_seqlen_per_dp_cp_rank and pad cu_seqlens tensors to - thd_max_num_seqs + 1 entries. If set to a positive integer N, token-like - tensors are padded to a multiple of N and cu_seqlens metadata is preserved. + 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_num_seqs + 1 entries. """ expert_model_parallel_size: int = 1 @@ -489,12 +510,27 @@ def __post_init__(self): ) if self.pad_packed_seq_alignment is not None: - if self.pad_packed_seq_alignment < 0: + 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( - "pad_packed_seq_alignment must be >= 0. Use the flag without a value " - "to let the caller choose a static target, or pass a positive integer " - "alignment." + "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: diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 787016c8796..95274bc51f6 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -1,6 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. from dataclasses import dataclass -from typing import Optional, Tuple +from typing import Literal, Optional, Tuple, Union import torch import torch.distributed as dist @@ -26,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. @@ -127,11 +128,52 @@ def _pad_cu_seqlens(cu_seqlens: Optional[Tensor], target_entries: int) -> Option 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_num_seqs: 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_num_seqs + 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_num_seqs + + 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], @@ -140,27 +182,19 @@ def _resolve_thd_padding_lengths( packed_seq_params: PackedSeqParams, target_len: Optional[int], alignment: Optional[int], -) -> Tuple[int, int, int, int, torch.device, bool]: - """Resolve local/global THD padding lengths without changing tensors.""" +) -> Tuple[int, int, int, int, torch.device]: + """Resolve local/global THD padding lengths without changing tensors. - actual_T = None - mask_device = None - for candidate in (tokens, labels, loss_mask, position_ids): - if candidate is not None: - actual_T = candidate.shape[-1] - mask_device = candidate.device - break - actual_T_is_local = actual_T is not None - if actual_T is None: - assert packed_seq_params.cu_seqlens_q is not None, ( - "packed_seq_params.cu_seqlens_q must be available to derive padding_mask " - "when tokens/labels/loss_mask/position_ids are all None." - ) - actual_T = int(packed_seq_params.cu_seqlens_q[-1].item()) - mask_device = packed_seq_params.cu_seqlens_q.device + 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. + """ + # Resolve the CP geometry used to translate local lengths to global endpoints. from megatron.core import parallel_state - cp_size = ( packed_seq_params.local_cp_size if packed_seq_params.local_cp_size is not None @@ -168,59 +202,79 @@ def _resolve_thd_padding_lengths( ) cp_rank = parallel_state.get_context_parallel_rank() if cp_size > 1 else 0 - if target_len is None: - assert alignment is not None, "Either target_len or alignment must be provided." - global_target_len = _round_up_to_alignment(int(actual_T), alignment) + # 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: - global_target_len = int(target_len) * cp_size + 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 - if actual_T_is_local: - local_actual_T = int(actual_T) - local_target_len = ( - int(target_len) - if target_len is not None - else _round_up_to_alignment(local_actual_T, alignment) - ) - else: - local_actual_T = int( - get_thd_partitioned_indices( - ( - 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 - ), - int(actual_T), - cp_size, - cp_rank, - ).numel() - ) - local_target_len = int( - get_thd_partitioned_indices( - ( - 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 - ), - global_target_len, - cp_size, - cp_rank, - ).numel() - ) + 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: - local_actual_T = int(actual_T) + # Without CP, local and global metadata lengths are identical. + local_actual_T = global_actual_T local_target_len = global_target_len - return ( - int(actual_T), - local_actual_T, - local_target_len, - global_target_len, - mask_device, - actual_T_is_local, - ) + return local_actual_T, global_actual_T, local_target_len, global_target_len, mask_device def pad_sequence_for_thd( @@ -232,6 +286,7 @@ def pad_sequence_for_thd( alignment: Optional[int] = None, target_len: Optional[int] = None, max_num_seqs: Optional[int] = None, + pad_by_appending_dummy_seq: bool = True, ) -> Tuple[ Optional[Tensor], Optional[Tensor], @@ -243,9 +298,37 @@ def pad_sequence_for_thd( """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. When ``max_num_seqs`` is provided, the - four cu_seqlens tensors are also padded to ``max_num_seqs + 1`` entries; - this is required by CUDA Graph replay because those tensors are graph inputs. + 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. + + 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) @@ -255,7 +338,7 @@ def pad_sequence_for_thd( "Exactly one of alignment or target_len must be provided for THD padding." ) - actual_T, local_actual_T, local_target_len, global_target_len, mask_device, _ = ( + local_actual_T, global_actual_T, local_target_len, global_target_len, mask_device = ( _resolve_thd_padding_lengths( tokens, labels, @@ -267,7 +350,8 @@ def pad_sequence_for_thd( ) ) - if actual_T is not None and packed_seq_params.cu_seqlens_q is not None: + # 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 @@ -277,45 +361,60 @@ def pad_sequence_for_thd( 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=( - packed_seq_params.cu_seqlens_q - if target_cu_entries is None - else _pad_cu_seqlens(packed_seq_params.cu_seqlens_q, target_cu_entries) - ), - cu_seqlens_kv=( - packed_seq_params.cu_seqlens_kv - if target_cu_entries is None - else _pad_cu_seqlens(packed_seq_params.cu_seqlens_kv, target_cu_entries) - ), - cu_seqlens_q_padded=( - packed_seq_params.cu_seqlens_q_padded - if target_cu_entries is None - else _pad_cu_seqlens(packed_seq_params.cu_seqlens_q_padded, target_cu_entries) - ), - cu_seqlens_kv_padded=( - packed_seq_params.cu_seqlens_kv_padded - if target_cu_entries is None - else _pad_cu_seqlens(packed_seq_params.cu_seqlens_kv_padded, target_cu_entries) - ), + 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 packed_seq_params.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 packed_seq_params.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. padding_mask = torch.arange(local_target_len, device=mask_device).unsqueeze(0) >= local_actual_T return tokens, labels, loss_mask, position_ids, padded_params, padding_mask diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 3c7767d966e..9f3cbfe8752 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -2301,16 +2301,7 @@ def _get_cuda_graph_input_data(self): auto_num_slots_tensor, op=torch.distributed.ReduceOp.MAX, group=pp_group ) auto_num_slots = int(auto_num_slots_tensor.item()) - requested_num_slots = self.config.cuda_graph_num_microbatch_slots - if requested_num_slots is not None: - assert requested_num_slots >= auto_num_slots, ( - "cuda_graph_num_microbatch_slots is smaller than the minimum safe number " - f"of slots for the current PP/VPP topology: requested={requested_num_slots}, " - f"required>={auto_num_slots}" - ) - self.num_microbatches = requested_num_slots - else: - self.num_microbatches = auto_num_slots + self.num_microbatches = auto_num_slots log_on_each_pipeline_stage( logger=logger, tp_group=None, diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 44675062d42..d2b17876292 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -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. @@ -228,7 +228,7 @@ def get_tokens_per_expert_and_token_count( reduce_group: torch.distributed.ProcessGroup, 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. """ @@ -237,8 +237,8 @@ def get_tokens_per_expert_and_token_count( local_tokens_per_expert, reduce_group ) 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() diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index fa809c4ba88..741d0461ba8 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1082,9 +1082,9 @@ class TransformerConfig(ModelParallelConfig): string names before normalization so existing CUDA_GRAPH_MODULES_DEPRECATIONS handles them.""" thd_max_num_seqs: int = 32 - """Maximum number of packed sequences per microbatch in THD format. The packing - scheduler closes a pack as soon as it reaches this many sequences (in addition to - the existing token-budget condition). When CUDA Graph is enabled, cu_seqlens + """Maximum number of THD sequence entries 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, @@ -1099,11 +1099,6 @@ class TransformerConfig(ModelParallelConfig): When enabled, capture builds a bounded number of graph slots and replay maps real microbatch_id to slot_id by modulo.""" - cuda_graph_num_microbatch_slots: Optional[int] = None - """Number of CUDA graph slots to capture per layer for dynamic microbatch replay. - If None, an automatic slot count is derived from the PP/VPP schedule topology. - If set, the provided value must be >= the automatically derived safe minimum.""" - #################### # Hyper-Connection Configuration #################### @@ -1450,6 +1445,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 @@ -2677,22 +2678,11 @@ 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": - if self.cuda_graph_dynamic_microbatches: - if self.cuda_graph_num_microbatch_slots is not None: - assert self.cuda_graph_num_microbatch_slots >= 1, ( - "cuda_graph_num_microbatch_slots must be >= 1 when " - "cuda_graph_dynamic_microbatches is enabled." - ) - else: + 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 self.cuda_graph_num_microbatch_slots is None, ( - "cuda_graph_num_microbatch_slots is only supported with " - "cuda_graph_impl=transformer_engine." - ) assert ( CudaGraphModule.moe not in self.cuda_graph_modules @@ -3028,10 +3018,10 @@ def _scope_to_str(s): "THD CUDA Graph requires --pad-packed-seq-alignment to be set." ) assert ( - self.pad_packed_seq_alignment == 0 + 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 without a value " + "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}." ) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 54f280efdbe..4edcbbbe02e 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1220,6 +1220,7 @@ def _reconstruct_packed_seq_params_from_kwargs(self, kwargs): 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 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 66199fb7cf8..1cbe796b5cd 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, @@ -1582,11 +1583,26 @@ def validate_args(args, defaults={}): ) if getattr(args, 'pad_packed_seq_alignment', None) is not None: - if args.pad_packed_seq_alignment < 0: + 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( - '--pad-packed-seq-alignment must be used without a value or with a ' - 'non-negative integer alignment.' + '--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 @@ -1594,11 +1610,11 @@ def validate_args(args, defaults={}): 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 != 0 + 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 without a value ' + "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}.' diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 6fecc802855..0cfa8f4992e 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, pad_sequence_for_thd +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 @@ -199,23 +203,12 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): labels = batch.get('labels', None) loss_mask = batch.get('loss_mask', None) position_ids = batch.get('position_ids', None) - cuda_graph_static = config.cuda_graph_impl != "none" - static_target = ( - config.pad_packed_seq_alignment == 0 - and config.max_seqlen_per_dp_cp_rank is not None + alignment, target_len, max_num_seqs = get_thd_padding_kwargs( + config.pad_packed_seq_alignment, + config.max_seqlen_per_dp_cp_rank, + config.thd_max_num_seqs, + config.cuda_graph_impl != "none", ) - if cuda_graph_static or static_target: - target_len = config.max_seqlen_per_dp_cp_rank - max_num_seqs = config.thd_max_num_seqs - alignment = None - else: - target_len = None - max_num_seqs = None - alignment = ( - int(max_seqlen[0].item()) - if config.pad_packed_seq_alignment == 0 - else config.pad_packed_seq_alignment - ) tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = ( pad_sequence_for_thd( tokens, @@ -226,6 +219,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): 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: diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 16c1de92122..4ccea5e10d7 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -80,7 +80,6 @@ "cuda_graph_modules": [], "cuda_graph_use_single_mempool": True, "cuda_graph_dynamic_microbatches": False, - "cuda_graph_num_microbatch_slots": None, "cuda_graph_scope": None, "cuda_graph_warmup_steps": 3, "deallocate_pipeline_outputs": True, diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index 95bc5d2ea29..5a8c5d803ec 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -9,6 +9,7 @@ from megatron.core import parallel_state from megatron.core.datasets.data_schedule import ( + _get_scheduler_max_real_num_seqs, get_batch_on_this_rank_for_sequence_packing, wrap_data_iterator, ) @@ -17,6 +18,34 @@ from tests.unit_tests.test_utilities import Utils +def test_scheduler_max_real_num_seqs_reserves_dummy_sequence(): + config = SimpleNamespace( + thd_max_num_seqs=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_num_seqs=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) + + class MockVariableLengthSequencePackingDataIterator: """ Mock data iterator for testing get_batch_on_this_rank_for_sequence_packing. diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 4f11b56d731..ad0cc36af1c 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -9,8 +9,10 @@ -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_EP4_ETP1 + sequence packing -(requires 8 GPUs, slow ~5 min per run, 4 runs total): +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 @@ -29,6 +31,8 @@ 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 @@ -95,6 +99,172 @@ def _build_layer(H, nh, nkv, ffn, max_seqlen, max_num_seqs, tp=1, sp=False): # ============================================================================= +@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_num_seqs=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 + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + 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 + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + 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)], + ) + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + 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): @@ -105,8 +275,8 @@ def teardown_method(self): @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_generic_alignment_preserves_cu_seqlens(self): - """Generic THD padding aligns token tensors while preserving sequence metadata.""" + 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() @@ -114,7 +284,80 @@ def test_generic_alignment_preserves_cu_seqlens(self): 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 + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + 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 + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + 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() @@ -143,10 +386,65 @@ def test_shapes_and_data_preservation(self): 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_num_seqs=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): @@ -166,7 +464,7 @@ def test_padding_mask_boundary(self): @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_cu_seqlens_fill_value(self): - """Padded entries repeat last cumulative sum (prevents OOB reads).""" + """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"), @@ -178,7 +476,10 @@ def test_cu_seqlens_fill_value(self): max_num_seqs=32, ) assert p.cu_seqlens_q[0] == 0 and p.cu_seqlens_q[2] == 80 - assert (p.cu_seqlens_q[3:] == 80).all() + 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") @@ -226,6 +527,7 @@ def test_round_trip(self): 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) @@ -306,6 +608,7 @@ def test_noop_without_packed_seq_params(self): "--max-seqlen-per-dp-cp-rank", "1024", "--pad-packed-seq-alignment", + "max", "--calculate-per-token-loss", "--transformer-impl", "transformer_engine", @@ -371,6 +674,7 @@ def test_noop_without_packed_seq_params(self): "flex", "--moe-flex-dispatcher-backend", "hybridep", + "--moe-router-fusion", "--moe-router-score-function", "sigmoid", "--moe-router-topk-scaling-factor", @@ -418,6 +722,18 @@ def test_noop_without_packed_seq_params(self): "hybridep", ] +_ATTN_CUDA_GRAPH_ARGS = [ + "--cuda-graph-impl", + "transformer_engine", + "--cuda-graph-modules", + "attn", +] + +_MOE_CUDA_GRAPH_ARGS = _ATTN_CUDA_GRAPH_ARGS + [ + "moe_preprocess", + "moe_router", +] + def _run_pretrain(model_args, cuda_graph_args, master_port): """Subprocess-launch `torchrun pretrain_gpt.py` once and capture stdout.""" @@ -504,8 +820,11 @@ def _extract_metrics(stdout): @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,base_port", - [("moonlight", _MOONLIGHT_ARGS, 29660), ("qwen3", _QWEN3_ARGS, 29662)], + "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. @@ -519,7 +838,7 @@ class TestE2EBitwise: Slow (~5 min per model). Marked `internal` so CI can opt-in. """ - def test_no_graph_vs_graph(self, model_name, model_args, base_port): + 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, ( @@ -537,12 +856,7 @@ def test_no_graph_vs_graph(self, model_name, model_args, base_port): # CUDA graph capture. r2 = _run_pretrain( model_args, - cuda_graph_args=[ - "--cuda-graph-impl", - "transformer_engine", - "--cuda-graph-modules", - "attn", - ], + cuda_graph_args=cuda_graph_args, master_port=base_port + 1, ) assert r2.returncode == 0, ( From aa0cfaecacb633b643231943822f8e1954af16f2 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Fri, 12 Jun 2026 00:36:36 -0700 Subject: [PATCH 17/25] fix linting Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 5 +- megatron/core/model_parallel_config.py | 3 +- megatron/core/packed_seq_params.py | 17 ++---- .../core/transformer/transformer_config.py | 6 +- .../transformer/test_thd_cuda_graph.py | 55 ++++--------------- 5 files changed, 22 insertions(+), 64 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index c65673740a3..268934b2316 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -435,9 +435,8 @@ def _get_scheduler_max_real_num_seqs(config) -> Optional[int]: if max_num_seqs < 1: raise ValueError(f"thd_max_num_seqs 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 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( diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index d65adaa966e..ce57898c8b6 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -521,8 +521,7 @@ def __post_init__(self): 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." + "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( diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 95274bc51f6..3f499db120d 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -195,6 +195,7 @@ def _resolve_thd_padding_lengths( # Resolve the CP geometry used to translate local lengths to global endpoints. from megatron.core import parallel_state + cp_size = ( packed_seq_params.local_cp_size if packed_seq_params.local_cp_size is not None @@ -254,19 +255,13 @@ def _resolve_thd_padding_lengths( # 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, + 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, + partition_cu_seqlens, global_target_len, cp_size, cp_rank ).numel() ) else: @@ -334,9 +329,9 @@ def pad_sequence_for_thd( 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." - ) + 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( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 741d0461ba8..2cf616f3f3c 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -3014,9 +3014,9 @@ def _scope_to_str(s): 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 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 diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index ad0cc36af1c..6bdc675507b 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -100,10 +100,7 @@ def _build_layer(H, nh, nkv, ffn, max_seqlen, max_num_seqs, tp=1, sp=False): @pytest.mark.internal -@pytest.mark.parametrize( - "cuda_graph_static,expected_max_num_seqs", - [(False, None), (True, 32)], -) +@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", @@ -210,18 +207,12 @@ def test_cp_tensor_target_len_scales_global_target(self): ) ) - assert (local_actual, global_actual, local_target, global_target) == ( - 80, - 140, - 128, - 256, - ) + 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)], + "alignment,target_len,expected_global_target", [(128, None, 256), (None, 128, 256)] ) @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") def test_cp_no_tensor_partitions_actual_and_target_lengths( @@ -246,13 +237,7 @@ def test_cp_no_tensor_partitions_actual_and_target_lengths( 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, + None, None, None, None, psp, target_len=target_len, alignment=alignment ) ) @@ -286,10 +271,7 @@ def test_generic_alignment_appends_dummy_padding_sequence(self): 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 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() @@ -387,9 +369,7 @@ def test_shapes_and_data_preservation(self): ): 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", + [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) @@ -432,10 +412,7 @@ def test_eager_pad_to_max_adds_dummy_padding_sequence(self): (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 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 @@ -722,17 +699,9 @@ def test_noop_without_packed_seq_params(self): "hybridep", ] -_ATTN_CUDA_GRAPH_ARGS = [ - "--cuda-graph-impl", - "transformer_engine", - "--cuda-graph-modules", - "attn", -] +_ATTN_CUDA_GRAPH_ARGS = ["--cuda-graph-impl", "transformer_engine", "--cuda-graph-modules", "attn"] -_MOE_CUDA_GRAPH_ARGS = _ATTN_CUDA_GRAPH_ARGS + [ - "moe_preprocess", - "moe_router", -] +_MOE_CUDA_GRAPH_ARGS = _ATTN_CUDA_GRAPH_ARGS + ["moe_preprocess", "moe_router"] def _run_pretrain(model_args, cuda_graph_args, master_port): @@ -854,11 +823,7 @@ def test_no_graph_vs_graph(self, model_name, model_args, cuda_graph_args, base_p ) # CUDA graph capture. - r2 = _run_pretrain( - model_args, - cuda_graph_args=cuda_graph_args, - master_port=base_port + 1, - ) + 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" From 64ae18a96bbc892532df27b2176d6ea5195839e6 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 15 Jun 2026 21:26:07 -0700 Subject: [PATCH 18/25] fix graph capture slot for dynamic num_microbatch Signed-off-by: HaochenYuan --- megatron/core/transformer/cuda_graphs.py | 131 ++++++++++++-- .../core/transformer/moe/token_dispatcher.py | 11 +- .../core/transformer/transformer_config.py | 8 +- .../core/transformer/transformer_layer.py | 2 +- megatron/training/training.py | 38 ++++ .../transformer/test_thd_cuda_graph.py | 167 ++++++++++++++++-- 6 files changed, 317 insertions(+), 40 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 9f3cbfe8752..7da996a9bd0 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 ( @@ -1767,7 +1768,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 ( @@ -1783,12 +1791,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 @@ -1995,7 +2005,7 @@ def get_rotary_pos_emb(transformer_module, transformer_input): 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.ones( + static_inputs["padding_mask"] = torch.zeros( 1, local_slen, dtype=torch.bool, device=torch.cuda.current_device() ) @@ -2247,6 +2257,83 @@ def _get_probe_num_microbatches_for_dynamic_slots(self): 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_num_seqs', 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. @@ -2259,6 +2346,12 @@ 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 ( @@ -2277,14 +2370,14 @@ def _get_cuda_graph_input_data(self): _, _, _probe_warmup, _ = _probe_get_pp( probe_num_microbatches, self.num_model_chunks, - self.config.microbatch_group_size_per_vp_stage, + 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, - self.config.microbatch_group_size_per_vp_stage, + microbatch_group_size_per_vp_stage, ) _probe_order = convert_schedule_table_to_order( _probe_warmup, self.num_model_chunks, _probe_st @@ -2301,16 +2394,34 @@ def _get_cuda_graph_input_data(self): auto_num_slots_tensor, op=torch.distributed.ReduceOp.MAX, group=pp_group ) auto_num_slots = int(auto_num_slots_tensor.item()) - self.num_microbatches = auto_num_slots + 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 enabled. ' - f'runtime_num_microbatches={get_num_microbatches()}, ' + 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'capture_num_microbatches={self.num_microbatches}', + 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() @@ -2318,14 +2429,14 @@ def _get_cuda_graph_input_data(self): _, _, 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, + 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/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index acfb9eaeacc..4586b1bfa15 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1070,16 +1070,15 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): or self.config.moe_hybridep_pad_variable_tokens ) if equalize_thd_token_counts: - if ( - self.config.sequence_packing_scheduler is not None - and torch.cuda.is_current_stream_capturing() + if self.config.sequence_packing_scheduler is not None and ( + torch.cuda.is_current_stream_capturing() or torch.compiler.is_compiling() ): - # Capture path: routing_map has already been padded to a static + # 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() (forbidden during stream capture) and use - # the local value directly. + # 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 diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2cf616f3f3c..fe9b1eec0ec 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1094,10 +1094,10 @@ class TransformerConfig(ModelParallelConfig): buffer.""" cuda_graph_dynamic_microbatches: bool = False - """Enable CUDA graph slot reuse so the same captured graphs can be replayed for a dynamic - number of microbatches. This option is only meaningful for cuda_graph_impl=transformer_engine. - When enabled, capture builds a bounded number of graph slots and replay maps real - microbatch_id to slot_id by modulo.""" + """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 diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 4edcbbbe02e..ed87346d170 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1115,7 +1115,7 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): 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.ones( + static_inputs["padding_mask"] = torch.zeros( 1, slen_for_mask, dtype=torch.bool, device=device ) elif attn_in_graph: diff --git a/megatron/training/training.py b/megatron/training/training.py index 21425422ffb..30260782e41 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3562,6 +3562,7 @@ def trace_handler(p): seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, optimizers=[optimizer], + thd_sequence_length_upper_bound=_get_thd_sequence_length_upper_bound(args), ) # Run training iterations till done. @@ -4595,3 +4596,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/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 6bdc675507b..31292f1da3c 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -23,6 +23,7 @@ import os import re +import socket import subprocess from pathlib import Path @@ -521,6 +522,91 @@ def test_noop_without_packed_seq_params(self): 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 @@ -531,22 +617,29 @@ def test_noop_without_packed_seq_params(self): # Common args shared across both models. _REPO_ROOT = Path(__file__).resolve().parents[3] -_SFT_JSON = ( +_VARLEN_JSON = ( '{"mode":"distribution","type":"lognormal",' - '"min_seq_len":128,"max_seq_len":2048,"mean_seq_len":1024,"lognormal_sigma":0.8}' + '"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", - "2048", + "4096", "--max-position-embeddings", "8192", "--micro-batch-size", "1", "--global-batch-size", - "4", + "64", "--train-iters", str(_TRAIN_ITERS), "--lr", @@ -574,18 +667,19 @@ def test_noop_without_packed_seq_params(self): "--swiglu", "--disable-bias-linear", "--sequence-parallel", - "--sft", + "--use-varlen-dataset", "--mock-data", "--tokenizer-type", "NullTokenizer", - "--sft-mock-dataset-config-json", - _SFT_JSON, + "--varlen-mock-dataset-config-json", + _VARLEN_JSON, "--sequence-packing-scheduler", "dp_balanced", "--max-seqlen-per-dp-cp-rank", - "1024", + "4096", "--pad-packed-seq-alignment", "max", + "--no-pad-packed-seq-by-appending-dummy-seq", "--calculate-per-token-loss", "--transformer-impl", "transformer_engine", @@ -608,9 +702,28 @@ def test_noop_without_packed_seq_params(self): "--no-check-for-nan-in-loss-and-grad", "--deterministic-mode", "--thd-max-num-seqs", - "32", + "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", @@ -670,7 +783,7 @@ def test_noop_without_packed_seq_params(self): "163840", ] -_QWEN3_ARGS = _COMMON_ARGS + [ +_QWEN3_ARGS = _QWEN3_COMMON_ARGS + [ "--num-layers", "36", "--hidden-size", @@ -699,11 +812,29 @@ def test_noop_without_packed_seq_params(self): "hybridep", ] -_ATTN_CUDA_GRAPH_ARGS = ["--cuda-graph-impl", "transformer_engine", "--cuda-graph-modules", "attn"] +_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() @@ -742,7 +873,7 @@ def _run_pretrain(model_args, cuda_graph_args, master_port): "--master_addr", "localhost", "--master_port", - str(master_port), + str(_get_available_port(master_port)), "pretrain_gpt.py", ] + model_args @@ -774,12 +905,9 @@ def _extract_metrics(stdout): 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) - lb_loss = re.search(r"load_balancing_loss:\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)}"] - if lb_loss: - parts.append(f"lb_loss={lb_loss.group(1)}") parts.append(f"grad_norm={grad_norm.group(1)}") results.append(" | ".join(parts)) return results @@ -799,10 +927,11 @@ 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 cuda_graph_modules=attn` - -- using the same model/test settings as test_moonlight_qwen3_bitwise.sh. - Asserts the per-iteration `lm loss / load_balancing_loss / grad norm` - lines are byte-identical. + 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. """ From d98e4320b12008ef8f31ba31cc84d8c97d3a6980 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 15 Jun 2026 22:00:52 -0700 Subject: [PATCH 19/25] Add THD padding mask and aux loss fixes from Tailai Ma Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 108 ++++++-- megatron/core/datasets/data_schedule_utils.py | 9 +- megatron/core/packed_seq_params.py | 27 +- megatron/core/transformer/moe/moe_layer.py | 17 +- megatron/core/transformer/moe/moe_utils.py | 19 +- megatron/core/transformer/moe/router.py | 232 ++++++++++++------ .../core/transformer/transformer_layer.py | 51 +++- megatron/training/training.py | 11 +- pretrain_gpt.py | 4 +- tests/unit_tests/test_sequence_packing.py | 47 +++- .../transformer/test_thd_cuda_graph.py | 31 +++ 11 files changed, 427 insertions(+), 129 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 268934b2316..b83a77c96aa 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -29,6 +29,51 @@ 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.""" @@ -598,10 +643,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. @@ -612,22 +668,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: - # 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 tokens/labels length. - 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() + # 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: @@ -659,7 +712,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 @@ -675,7 +735,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: @@ -695,6 +755,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) @@ -705,6 +766,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() @@ -732,7 +794,6 @@ def get_batch_on_this_rank_for_sequence_packing( # Pad the already-packed THD tensors at the end when requested. CUDA Graph # additionally pads cu_seqlens tensors to thd_max_num_seqs + 1 entries. - padding_mask = None pad_alignment = ( getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None ) @@ -756,6 +817,7 @@ def get_batch_on_this_rank_for_sequence_packing( pad_by_appending_dummy_seq=getattr( config, 'pad_packed_seq_by_appending_dummy_seq', True ), + padding_mask=padding_mask, ) ) diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index fc906ebf725..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,6 +20,7 @@ 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: @@ -33,7 +34,9 @@ def get_cp_slice_for_thd(batch, cp_group): # 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 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/packed_seq_params.py b/megatron/core/packed_seq_params.py index 3f499db120d..8a6e42906f2 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -102,6 +102,22 @@ def _pad_seq_tensor(t: Optional[Tensor], target_len: int) -> Optional[Tensor]: 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. @@ -282,6 +298,7 @@ def pad_sequence_for_thd( target_len: Optional[int] = None, max_num_seqs: Optional[int] = None, pad_by_appending_dummy_seq: bool = True, + padding_mask: Optional[Tensor] = None, ) -> Tuple[ Optional[Tensor], Optional[Tensor], @@ -314,6 +331,8 @@ def pad_sequence_for_thd( ``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. Notes: - THD CP slicing is defined by Transformer Engine. On metadata-only @@ -410,6 +429,12 @@ def pad_sequence_for_thd( ) # True marks padded local token slots for routing/loss paths. - padding_mask = torch.arange(local_target_len, device=mask_device).unsqueeze(0) >= local_actual_T + 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/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 2ecd4fb15b2..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: @@ -739,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 d2b17876292..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 @@ -226,6 +226,7 @@ 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, ) -> Tuple[torch.Tensor, Union[int, torch.Tensor], Union[int, torch.Tensor]]: @@ -233,15 +234,23 @@ def get_tokens_per_expert_and_token_count( 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 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 01b0adea88d..53112880be0 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""" @@ -319,16 +335,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, ) @@ -349,8 +368,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 @@ -362,6 +386,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. @@ -377,10 +402,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, ) @@ -404,8 +431,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 @@ -421,7 +453,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, @@ -452,18 +485,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. @@ -472,7 +538,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). @@ -501,44 +570,45 @@ 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) @@ -550,48 +620,48 @@ 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: @@ -603,7 +673,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 @@ -695,14 +765,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. @@ -770,6 +841,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, @@ -778,6 +850,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, @@ -802,15 +875,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. """ @@ -830,7 +904,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 @@ -939,6 +1018,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. @@ -954,6 +1034,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/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ed87346d170..ebff6fcbb2b 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -776,6 +776,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 @@ -799,6 +800,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. @@ -858,6 +860,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: @@ -1274,6 +1278,7 @@ def _te_cuda_graph_capture(self, *args, **kwargs): 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] @@ -1443,6 +1448,7 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): 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 @@ -1788,6 +1794,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 @@ -1904,6 +1911,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.""" @@ -1950,6 +1958,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: @@ -2148,7 +2158,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 @@ -2263,7 +2277,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. @@ -2292,6 +2308,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: @@ -2345,7 +2362,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. @@ -2362,10 +2384,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 @@ -2393,6 +2422,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( @@ -2400,15 +2430,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/training.py b/megatron/training/training.py index 30260782e41..24868a5de7a 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2276,6 +2276,7 @@ def train_step( """Single training step.""" args = get_args() timers = get_timers() + num_microbatches = get_num_microbatches() rerun_state_machine = get_rerun_state_machine() save_params_in_this_iteration = ( @@ -2411,7 +2412,7 @@ def _save_state_dict(attr_name, label): should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit() if should_exit: - return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0 + return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0, num_microbatches # Empty unused memory. if args.empty_unused_memory_level >= 1: @@ -2494,6 +2495,7 @@ def _save_state_dict(attr_name, label): grad_norm, num_zeros_in_grad, log_max_attention_logit, + num_microbatches, ) return ( {}, @@ -2504,6 +2506,7 @@ def _save_state_dict(attr_name, label): grad_norm, num_zeros_in_grad, log_max_attention_logit, + num_microbatches, ) @@ -2523,6 +2526,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() @@ -2697,7 +2701,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") @@ -3696,6 +3700,7 @@ def trace_handler(p): grad_norm = 0.0 num_zeros_in_grad = 0 max_attention_logit = None + num_microbatches = get_num_microbatches() else: ft_integration.on_training_step_start() ( @@ -3707,6 +3712,7 @@ def trace_handler(p): grad_norm, num_zeros_in_grad, max_attention_logit, + num_microbatches, ) = train_step( forward_step_func, train_data_iterator, @@ -3858,6 +3864,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 diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 0cfa8f4992e..3e13b36aa13 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -129,8 +129,8 @@ 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 optional THD padding - # and returns a 7-tuple including `padding_mask` (None when no padding). + # `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, diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index 5a8c5d803ec..51d6f1fdc58 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -9,7 +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, ) @@ -46,6 +48,35 @@ def test_scheduler_max_real_num_seqs_rejects_dummy_without_capacity(): _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. @@ -239,13 +270,14 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc dynamic_cp=dynamic_cp, ) - # The helper returns a 7-tuple; padding_mask is None when THD padding is disabled. + # 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 None - ), "padding_mask should be None when no padding config is provided." + 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() @@ -322,7 +354,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: @@ -354,6 +386,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_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 31292f1da3c..0b094c619f5 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -439,6 +439,37 @@ def test_padding_mask_boundary(self): ) 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): From 0d41b56f751eda5c95fb4a59f8669cc71594acb3 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 15 Jun 2026 22:17:04 -0700 Subject: [PATCH 20/25] fix CI Signed-off-by: HaochenYuan --- megatron/training/arguments.py | 10 ---------- .../unit_tests/models/test_hybrid_moe_model.py | 5 ++++- .../transformer/test_thd_cuda_graph.py | 17 ++++++++++++----- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 1cbe796b5cd..e11fb472b27 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1572,16 +1572,6 @@ def validate_args(args, defaults={}): f"to {args.data_parallel_size * args.context_parallel_size}." ) - if args.sequence_packing_scheduler is not None: - if args.sequence_packing_scheduler == 'dp_balanced': - total_cp_ranks = args.context_parallel_size - else: - total_cp_ranks = args.data_parallel_size * args.context_parallel_size - assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, ( - f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) ' - f'must be >= single sequence max length ({args.seq_length})' - ) - 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 diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 4ccea5e10d7..e190ba7f08c 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -351,7 +351,10 @@ # 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/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 0b094c619f5..1dffe8c136d 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -45,6 +45,12 @@ 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) # ============================================================================= @@ -169,7 +175,7 @@ def test_no_tensor_requires_cu_seqlens(self): ) @pytest.mark.internal - @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + @_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() @@ -193,7 +199,7 @@ def test_cp_tensor_alignment_uses_local_target_and_global_tail(self): assert mask_device == tokens.device @pytest.mark.internal - @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + @_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() @@ -215,7 +221,7 @@ def test_cp_tensor_target_len_scales_global_target(self): @pytest.mark.parametrize( "alignment,target_len,expected_global_target", [(128, None, 256), (None, 128, 256)] ) - @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + @_REQUIRES_TWO_RANKS def test_cp_no_tensor_partitions_actual_and_target_lengths( self, alignment, target_len, expected_global_target ): @@ -278,7 +284,7 @@ def test_generic_alignment_appends_dummy_padding_sequence(self): assert not mask[0, :total_T].any() and mask[0, total_T:].all() @pytest.mark.internal - @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + @_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() @@ -299,7 +305,7 @@ def test_cp_alignment_uses_global_cu_seqlens_length(self): assert not mask[0, :local_T].any() @pytest.mark.internal - @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least 2 GPUs") + @_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() @@ -945,6 +951,7 @@ def _extract_metrics(stdout): @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( From 3904a9319e530351e51f9e19b7557017babc75e4 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 15 Jun 2026 22:21:32 -0700 Subject: [PATCH 21/25] fix linting Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 11 ++++------- megatron/core/transformer/cuda_graphs.py | 8 ++------ megatron/core/transformer/moe/router.py | 9 ++++----- tests/unit_tests/models/test_hybrid_moe_model.py | 5 +---- tests/unit_tests/test_sequence_packing.py | 3 +-- tests/unit_tests/transformer/test_thd_cuda_graph.py | 4 +--- 6 files changed, 13 insertions(+), 27 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index b83a77c96aa..6a37834f124 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -29,7 +29,9 @@ 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: +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 @@ -56,12 +58,7 @@ def _build_thd_padding_mask(cu_seqlens: torch.Tensor, cu_seqlens_padded: torch.T 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, - } + 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: diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 7da996a9bd0..2ebbe5f9d51 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -2375,9 +2375,7 @@ def _get_cuda_graph_input_data(self): 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_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 @@ -2434,9 +2432,7 @@ def _get_cuda_graph_input_data(self): p2p_communicator=self.p2p_communicator, ) schedule_table = get_schedule_table( - self.num_microbatches, - self.num_model_chunks, - 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/moe/router.py b/megatron/core/transformer/moe/router.py index 53112880be0..fc1cf7d5cf8 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -597,7 +597,9 @@ def attach_and_log_load_balancing_loss( 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) + 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 @@ -654,10 +656,7 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None): # 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, - ) + 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 diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index e190ba7f08c..269e47db972 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -351,10 +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 = { - "pad_packed_seq_alignment", - "pad_packed_seq_by_appending_dummy_seq", -} +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 51d6f1fdc58..4368a770c64 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -55,8 +55,7 @@ def test_scheduler_thd_padding_mask_from_cu_seqlens(): 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]), + padding_mask, torch.tensor([False, False, False, True, False, False, True, True]) ) diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 1dffe8c136d..61ed9392d41 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -451,9 +451,7 @@ 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", + [[False, False, False, True, False, False, True, True]], dtype=torch.bool, device="cuda" ) _, _, _, _, _, m = pad_sequence_for_thd( From 5e0d88e45bd4dbb8716ac85d1c562e2bc6e03dc0 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 15 Jun 2026 23:35:30 -0700 Subject: [PATCH 22/25] - Use thd_max_packed_sequences to clarify dummy sequence capacity semantics - Remove THD RoPE packed-frequency shape heuristic by requiring max_seqlen - Prefer explicit CP group when resolving THD padding lengths Signed-off-by: HaochenYuan --- megatron/core/datasets/data_schedule.py | 17 +++-- .../inference/contexts/dynamic_context.py | 1 + megatron/core/model_parallel_config.py | 2 +- .../models/common/embeddings/rope_utils.py | 9 ++- megatron/core/packed_seq_params.py | 76 +++++++++++++++---- megatron/core/transformer/attention.py | 4 +- megatron/core/transformer/cuda_graphs.py | 2 +- .../deepseek_v4_hybrid_attention.py | 8 ++ .../core/transformer/transformer_config.py | 6 +- .../core/transformer/transformer_layer.py | 4 +- pretrain_gpt.py | 4 +- .../models/test_hybrid_moe_model.py | 2 +- tests/unit_tests/test_sequence_packing.py | 4 +- .../transformer/test_thd_cuda_graph.py | 8 +- 14 files changed, 106 insertions(+), 41 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index a2f1163b988..0b2039c60eb 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -468,26 +468,26 @@ 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_num_seqs`` is the final static THD capacity, including the + ``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_num_seqs', None) + 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_num_seqs must be >= 1, got {max_num_seqs}.") + 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_num_seqs must be >= 2 when THD padding appends a dummy " - "sequence, because thd_max_num_seqs includes that dummy sequence." + "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 @@ -539,7 +539,7 @@ def wrap_data_iterator( scheduler_max_num_seqs = ( _get_scheduler_max_real_num_seqs(config) if scheduler_type == 'dp_balanced' - else getattr(config, 'thd_max_num_seqs', None) + else getattr(config, 'thd_max_packed_sequences', None) ) scheduler = scheduler_map[scheduler_type]( @@ -794,7 +794,7 @@ def get_batch_on_this_rank_for_sequence_packing( ) # Pad the already-packed THD tensors at the end when requested. CUDA Graph - # additionally pads cu_seqlens tensors to thd_max_num_seqs + 1 entries. + # 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 ) @@ -802,7 +802,7 @@ def get_batch_on_this_rank_for_sequence_packing( 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_num_seqs', 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 = ( @@ -819,6 +819,7 @@ def get_batch_on_this_rank_for_sequence_packing( config, 'pad_packed_seq_by_appending_dummy_seq', True ), padding_mask=padding_mask, + cp_group=cp_group, ) ) 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 ce57898c8b6..ba2f4715fc7 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -124,7 +124,7 @@ class ModelParallelConfig: 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_num_seqs + 1 entries. + cu_seqlens tensors to thd_max_packed_sequences + 1 entries. """ expert_model_parallel_size: int = 1 diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index aa6d9c8a634..ad790e5c5ad 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -282,10 +282,11 @@ def _apply_rotary_pos_emb_thd( else: freq_pos = local_pos.to(torch.int64) - if max_seqlen is None: - exact_packed_freqs = freqs.dim() >= 1 and cp_size == 1 and freqs.size(0) > total_tokens - else: - exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) > max_seqlen + 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 diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 8a6e42906f2..3095b1b8464 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -129,8 +129,8 @@ def _pad_cu_seqlens(cu_seqlens: Optional[Tensor], target_entries: int) -> Option return None actual_entries = cu_seqlens.shape[0] assert actual_entries <= target_entries, ( - f"Actual num_seqs ({actual_entries - 1}) exceeds thd_max_num_seqs " - f"({target_entries - 1}). Increase --thd-max-num-seqs, decrease " + 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." ) @@ -167,7 +167,7 @@ def _round_up_to_alignment(value: int, alignment: int) -> int: def get_thd_padding_kwargs( pad_packed_seq_alignment: Union[int, Literal["max"]], max_seqlen_per_dp_cp_rank: Optional[int], - thd_max_num_seqs: 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. @@ -177,12 +177,12 @@ def get_thd_padding_kwargs( - ``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_num_seqs + 1`` is a CUDA Graph static-input + 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_num_seqs + 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 @@ -198,6 +198,9 @@ def _resolve_thd_padding_lengths( 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. @@ -209,15 +212,9 @@ def _resolve_thd_padding_lengths( mask_device: Device used to build the returned padding mask. """ - # Resolve the CP geometry used to translate local lengths to global endpoints. - from megatron.core import parallel_state - - cp_size = ( - packed_seq_params.local_cp_size - if packed_seq_params.local_cp_size is not None - else parallel_state.get_context_parallel_world_size() + cp_size, cp_rank = _resolve_thd_cp_geometry( + packed_seq_params, cp_group=cp_group, cp_size=cp_size, cp_rank=cp_rank ) - cp_rank = parallel_state.get_context_parallel_rank() if cp_size > 1 else 0 # Find the first token-like tensor that carries this rank's local length. local_tensor_T = None @@ -288,6 +285,47 @@ def _resolve_thd_padding_lengths( 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], @@ -299,6 +337,9 @@ def pad_sequence_for_thd( 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], @@ -333,6 +374,12 @@ def pad_sequence_for_thd( 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 @@ -361,6 +408,9 @@ def pad_sequence_for_thd( packed_seq_params, target_len=target_len, alignment=alignment, + cp_group=cp_group, + cp_size=cp_size, + cp_rank=cp_rank, ) ) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index e712122ef69..d367742f672 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -79,7 +79,9 @@ if not HAVE_FA3: try: - from flashattn_hopper.flash_attn_interface import _flash_attn_forward + from flashattn_hopper.flash_attn_interface import ( + _flash_attn_forward, + ) from flashattn_hopper.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 7eefc2409c7..d59d1fbf5b0 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -2340,7 +2340,7 @@ def _get_thd_varlen_max_num_microbatches( else self.seq_length ) - max_num_seqs = getattr(self.config, 'thd_max_num_seqs', None) + 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( 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..be533fcfe55 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 @@ -327,9 +327,11 @@ def forward( else packed_seq_params.cu_seqlens_kv ) rope_seqlen = cu_seqlens_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/transformer_config.py b/megatron/core/transformer/transformer_config.py index 8459303b73e..9acf08f55de 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1082,8 +1082,10 @@ 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_num_seqs: int = 32 - """Maximum number of THD sequence entries per microbatch, including any dummy + 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. diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 913bb1886bb..546d7146039 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1204,13 +1204,13 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): if self._is_thd_cuda_graph(): if attn_in_graph: - # Static cu_seqlens shaped [thd_max_num_seqs + 1]. We seed it as + # 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_num_seqs + 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 diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 3e13b36aa13..9265e8a832a 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -196,7 +196,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): ) # Pad the already-packed THD tensors at the end when requested. CUDA Graph - # additionally pads cu_seqlens tensors to thd_max_num_seqs + 1 entries. + # 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) @@ -206,7 +206,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = 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_num_seqs, + config.thd_max_packed_sequences, config.cuda_graph_impl != "none", ) tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = ( diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index beaa17fbcc5..44f0ec8d2fd 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -282,7 +282,7 @@ "symmetric_ar_type": None, "tensor_model_parallel_size": 2, "test_mode": False, - "thd_max_num_seqs": 32, + "thd_max_packed_sequences": 32, "timers": None, "tp_comm_atomic_ag": False, "tp_comm_atomic_rs": False, diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index 4368a770c64..f1fbcb53dea 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -22,7 +22,7 @@ def test_scheduler_max_real_num_seqs_reserves_dummy_sequence(): config = SimpleNamespace( - thd_max_num_seqs=32, + thd_max_packed_sequences=32, pad_packed_seq_alignment="max", pad_packed_seq_by_appending_dummy_seq=True, ) @@ -39,7 +39,7 @@ def test_scheduler_max_real_num_seqs_reserves_dummy_sequence(): def test_scheduler_max_real_num_seqs_rejects_dummy_without_capacity(): config = SimpleNamespace( - thd_max_num_seqs=1, + thd_max_packed_sequences=1, pad_packed_seq_alignment="max", pad_packed_seq_by_appending_dummy_seq=True, ) diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 61ed9392d41..92a81b2fcdc 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -86,7 +86,7 @@ def _build_layer(H, nh, nkv, ffn, max_seqlen, max_num_seqs, tp=1, sp=False): num_query_groups=nkv, ffn_hidden_size=ffn, max_seqlen_per_dp_cp_rank=max_seqlen, - thd_max_num_seqs=max_num_seqs, + thd_max_packed_sequences=max_num_seqs, tensor_model_parallel_size=tp, sequence_parallel=sp, bf16=True, @@ -112,7 +112,7 @@ def test_pad_to_max_resolves_padding_kwargs(cuda_graph_static, expected_max_num_ alignment, target_len, max_num_seqs = get_thd_padding_kwargs( pad_packed_seq_alignment="max", max_seqlen_per_dp_cp_rank=8192, - thd_max_num_seqs=32, + thd_max_packed_sequences=32, cuda_graph_static=cuda_graph_static, ) @@ -399,7 +399,7 @@ def test_eager_pad_to_max_adds_dummy_padding_sequence(self): 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_num_seqs=32, + thd_max_packed_sequences=32, cuda_graph_static=False, ) @@ -736,7 +736,7 @@ def test_dp_balanced_thd_capture_upper_bound_aligns_vpp_groups(self): "1", "--no-check-for-nan-in-loss-and-grad", "--deterministic-mode", - "--thd-max-num-seqs", + "--thd-max-packed-sequences", "8", ] From 07a24f23b928a415f0927409ab8766490ac2acaf Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 15 Jun 2026 23:41:13 -0700 Subject: [PATCH 23/25] fix linting Signed-off-by: HaochenYuan --- megatron/core/transformer/attention.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index d367742f672..e712122ef69 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -79,9 +79,7 @@ if not HAVE_FA3: try: - from flashattn_hopper.flash_attn_interface import ( - _flash_attn_forward, - ) + from flashattn_hopper.flash_attn_interface import _flash_attn_forward from flashattn_hopper.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) From 00fa6cbf11448b8cf70d56b285e35d36a3f78034 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Tue, 16 Jun 2026 01:17:33 -0700 Subject: [PATCH 24/25] fix CI & dsv4 Signed-off-by: HaochenYuan --- .../deepseek_v4_hybrid_attention.py | 2 +- tests/unit_tests/fusions/test_mla_yarn_rope_apply.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) 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 be533fcfe55..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,7 +326,7 @@ 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 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, From 30a727dc6f1dc8e6e0bc3770bc4a7070d3f43090 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Tue, 16 Jun 2026 06:56:40 -0700 Subject: [PATCH 25/25] remove microbatch variance checking when using sequence packing Signed-off-by: HaochenYuan --- megatron/training/training.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index facfcf55f07..cac829727d7 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3483,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()})"