From 9b4db6ac074204079de25aa8f918137359a897a7 Mon Sep 17 00:00:00 2001 From: xiaoyao0115 <1804647152@qq.com> Date: Mon, 29 Jun 2026 00:51:15 -0700 Subject: [PATCH] several fixes for thd e2e Signed-off-by: xiaoyao0115 <1804647152@qq.com> --- megatron/core/datasets/data_schedule.py | 5 +- megatron/core/model_parallel_config.py | 10 +- megatron/core/packed_seq_params.py | 80 +++++++-- megatron/core/transformer/cuda_graphs.py | 2 +- megatron/core/transformer/moe/router.py | 10 +- .../transformer/multi_token_prediction.py | 43 +++-- .../core/transformer/transformer_layer.py | 1 - megatron/training/datasets/varlen_dataset.py | 5 + megatron/training/training.py | 8 +- megatron/training/utils/common_utils.py | 56 +++++- pretrain_gpt.py | 13 +- tests/unit_tests/data/test_varlen_dataset.py | 79 ++++++++- tests/unit_tests/test_sequence_packing.py | 47 +++-- .../transformer/moe/test_routers.py | 10 +- .../transformer/moe/test_token_dispatcher.py | 9 +- .../test_multi_token_prediction.py | 62 ++++--- .../transformer/test_thd_cuda_graph.py | 164 ++++++++++++------ 17 files changed, 456 insertions(+), 148 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index d5bf07053cc..d6724bf67e4 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -468,7 +468,7 @@ def _get_scheduler_max_real_num_seqs(config) -> Optional[int]: 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 + config, 'pad_packed_seq_by_appending_dummy_seq', False ): if max_num_seqs < 2: raise ValueError( @@ -778,7 +778,6 @@ 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 @@ -804,7 +803,7 @@ def get_batch_on_this_rank_for_sequence_packing( 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 + config, 'pad_packed_seq_by_appending_dummy_seq', False ), padding_mask=padding_mask, cp_group=cp_group, diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 5c2786285b1..2dfc4800d2f 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -118,13 +118,13 @@ class ModelParallelConfig: tensors are padded to a multiple of N. """ - pad_packed_seq_by_appending_dummy_seq: bool = True + pad_packed_seq_by_appending_dummy_seq: bool = False """Represent a THD packed-sequence padding tail by appending a dummy sequence. - When disabled, token-like tensors are still padded according to - pad_packed_seq_alignment, but cu_seqlens sequence boundaries are not extended - for the padding tail. CUDA Graph static-input padding may still pad the - cu_seqlens tensors to thd_max_packed_sequences + 1 entries. + By default, token-like tensors are padded according to + pad_packed_seq_alignment, cu_seqlens is unchanged, and the last + cu_seqlens_padded endpoint is extended over the padding tail. When enabled, + the tail is represented as a separate zero-valid-token dummy sequence. """ expert_model_parallel_size: int = 1 diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 3095b1b8464..e96f23ae2a1 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -159,6 +159,16 @@ def _append_dummy_seq(cu_seqlens: Optional[Tensor], dummy_end: int) -> Optional[ return torch.cat((cu_seqlens, dummy), dim=0) +def _replace_last_cu_seqlen(cu_seqlens: Optional[Tensor], padded_end: int) -> Optional[Tensor]: + """Return cu_seqlens with its final physical endpoint replaced.""" + if cu_seqlens is None: + return None + + result = cu_seqlens.clone() + result[-1] = int(padded_end) + return result + + 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 @@ -225,9 +235,14 @@ def _resolve_thd_padding_lengths( mask_device = candidate.device break - # Prefer THD metadata for the global packed length when it is available. + # The padded endpoint describes tensor storage. The unpadded endpoint is + # only the compact valid-token count when gaps exist between sequences. has_local_tensor = local_tensor_T is not None - if packed_seq_params.cu_seqlens_q is not None: + if packed_seq_params.cu_seqlens_q_padded is not None: + global_actual_T = int(packed_seq_params.cu_seqlens_q_padded[-1].item()) + if mask_device is None: + mask_device = packed_seq_params.cu_seqlens_q_padded.device + elif 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 @@ -335,7 +350,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, + pad_by_appending_dummy_seq: bool = False, padding_mask: Optional[Tensor] = None, cp_group: Optional[dist.ProcessGroup] = None, cp_size: Optional[int] = None, @@ -371,7 +386,8 @@ def pad_sequence_for_thd( 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. + tail as an extra dummy sequence. Otherwise leave cu_seqlens unchanged + and extend the final cu_seqlens_padded endpoint over the tail. 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 @@ -386,8 +402,10 @@ def pad_sequence_for_thd( 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. + - By default, post-pack padding extends only the last padded sequence; + original cu_seqlens continue to describe valid-token counts. + - When ``pad_by_appending_dummy_seq`` is true, the padding tail is + represented as a separate zero-valid-token dummy sequence. - ``max_num_seqs`` pads all four cu_seqlens tensors; this is required by CUDA Graph replay because those tensors are graph inputs. @@ -437,16 +455,40 @@ def pad_sequence_for_thd( 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. + # Cover post-pack padding by extending the last physical sequence, or by + # representing it as a separate zero-valid-token dummy sequence. 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 + has_padding_tail = global_target_len > global_actual_T + has_dummy_padding_seq = pad_by_appending_dummy_seq and has_padding_tail 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) + # None is TE's auto-detect mode. Keep real and physical boundaries + # separate unless the caller explicitly disables padding between sequences. + if ( + packed_seq_params.pad_between_seqs is not False + and cu_seqlens_q is not None + and cu_seqlens_q_padded is not None + ): + cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, int(cu_seqlens_q[-1].item())) + cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len) + else: + cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, global_target_len) + cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len) + if ( + packed_seq_params.pad_between_seqs is not False + and cu_seqlens_kv is not None + and cu_seqlens_kv_padded is not None + ): + cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, int(cu_seqlens_kv[-1].item())) + cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len) + else: + cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, global_target_len) + cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len) + elif has_padding_tail: + q_physical_cu = cu_seqlens_q_padded if cu_seqlens_q_padded is not None else cu_seqlens_q + kv_physical_cu = cu_seqlens_kv_padded if cu_seqlens_kv_padded is not None else cu_seqlens_kv + cu_seqlens_q_padded = _replace_last_cu_seqlen(q_physical_cu, global_target_len) + cu_seqlens_kv_padded = _replace_last_cu_seqlen(kv_physical_cu, global_target_len) # Pad cu_seqlens entry counts for static CUDA Graph inputs. if target_cu_entries is not None: @@ -455,6 +497,12 @@ def pad_sequence_for_thd( 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) + def _max_physical_seqlen(cu_seqlens: Optional[Tensor], fallback: int) -> int: + if cu_seqlens is None or cu_seqlens.numel() < 2: + return fallback + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + return max(fallback, int(lengths.max().item())) + # Rebuild PackedSeqParams with the padded tensor and metadata shapes. padded_params = PackedSeqParams( qkv_format=packed_seq_params.qkv_format, @@ -465,17 +513,17 @@ def pad_sequence_for_thd( max_seqlen_q=( global_target_len if target_cu_entries is not None - else max(packed_seq_params.max_seqlen_q, dummy_seq_len) + else _max_physical_seqlen(cu_seqlens_q_padded, packed_seq_params.max_seqlen_q) ), max_seqlen_kv=( global_target_len if target_cu_entries is not None - else max(packed_seq_params.max_seqlen_kv, dummy_seq_len) + else _max_physical_seqlen(cu_seqlens_kv_padded, 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, - pad_between_seqs=False if has_dummy_padding_seq else packed_seq_params.pad_between_seqs, + pad_between_seqs=packed_seq_params.pad_between_seqs, ) # True marks padded local token slots for routing/loss paths. diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index d59d1fbf5b0..c537bafc942 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -2344,7 +2344,7 @@ def _get_thd_varlen_max_num_microbatches( 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 + self.config, 'pad_packed_seq_by_appending_dummy_seq', False ): max_num_seqs -= 1 diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 5aa03c19aa4..b2ff796e450 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -827,7 +827,15 @@ def routing( router_replay=self.router_replay, ) - # Apply token dropping to probs and routing_map. + # Padding tokens must not enter the dispatcher. Masking only auxiliary + # losses still lets padding alter expert capacity and grouped-GEMM shapes. + if padding_mask is not None: + valid_tokens = (~padding_mask).unsqueeze(-1) + probs = probs * valid_tokens + routing_map = routing_map & valid_tokens + + # Apply token dropping after removing padding so padding tokens cannot + # consume expert capacity. if self.config.moe_expert_capacity_factor is not None: probs, routing_map = apply_router_token_dropping( probs, diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index c0d5ef01e25..a9c50bc0d9f 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -371,8 +371,8 @@ def save_metrics_to_tracker( """Save normalized MTP loss and acceptance counts for logging. This compatibility path is used by tests and callers that already - computed a normalized per-layer loss. Dynamic-CP code should use - ``save_loss_to_tracker`` so loss is weighted by token counts. + computed a normalized per-layer loss. Dynamic-CP code uses + ``save_loss_to_tracker`` to normalize each local contribution safely. """ if layer_number is None: return @@ -402,11 +402,12 @@ def save_loss_to_tracker( reduce_group: Optional[torch.distributed.ProcessGroup] = None, avg_group: Optional[torch.distributed.ProcessGroup] = None, ): - """Save the mtp loss sum and token count for logging. + """Normalize and accumulate a local MTP loss for logging. - Stores raw sums so that the global per-token loss can be computed - correctly after all-reduce, even when token counts differ across - ranks (e.g. Dynamic CP) or microbatches. + MTP is normalized independently for each microbatch. The tracker + accumulates those normalized losses, then reduction combines the sums + across ranks. This intentionally preserves sequence-packing semantics + instead of changing the metric to global ``sum(loss) / sum(tokens)``. Args: loss_sum (torch.Tensor): Sum of per-element losses on this rank. @@ -415,8 +416,8 @@ def save_loss_to_tracker( num_layers (int): The number of total layers. correct (Optional[torch.Tensor]): Number of correct MTP predictions. total (Optional[torch.Tensor]): Total number of MTP predictions. - reduce_group (torch.distributed.ProcessGroup): The group for sum-reducing losses. - avg_group (torch.distributed.ProcessGroup): The group for sum-reducing before averaging. + reduce_group (torch.distributed.ProcessGroup): Group for summing losses. + avg_group (torch.distributed.ProcessGroup): Group for averaging losses. """ if layer_number is None: return @@ -424,9 +425,8 @@ def save_loss_to_tracker( tracker = MTPLossLoggingHelper.tracker if "loss_sums" not in tracker: tracker["loss_sums"] = torch.zeros(num_layers, device=torch.cuda.current_device()) - tracker["num_tokens"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + loss_sum = (loss_sum * (num_tokens > 0).to(loss_sum.dtype)) / num_tokens.clamp(min=1) tracker["loss_sums"][layer_number] += loss_sum.detach() - tracker["num_tokens"][layer_number] += num_tokens.detach() if correct is not None and total is not None: if "correct_values" not in tracker: tracker["correct_values"] = torch.zeros( @@ -485,7 +485,6 @@ def clean_loss_in_tracker(): tracker = MTPLossLoggingHelper.tracker if "loss_sums" in tracker: tracker["loss_sums"].zero_() - tracker["num_tokens"].zero_() if "values" in tracker: tracker["values"].zero_() if "correct_values" in tracker: @@ -499,21 +498,21 @@ def clean_loss_in_tracker(): def reduce_loss_in_tracker(): """Collect and reduce the mtp losses across ranks. - Packs loss sums and token counts into a single tensor for one - all-reduce, then computes per-token loss. This produces correct - weighted-average results even when ranks hold different numbers - of tokens (e.g. Dynamic CP with variable CP sizes). + Each element is already a sum of normalized microbatch losses. Sum + reductions preserve additive groups, while the DP+CP average keeps the + legacy logging contract. """ tracker = MTPLossLoggingHelper.tracker if "loss_sums" not in tracker: return - packed = torch.cat([tracker["loss_sums"], tracker["num_tokens"]]) - for group_key in ('reduce_group', 'avg_group'): - group = tracker.get(group_key) - if group is not None: - torch.distributed.all_reduce(packed, group=group) - loss_sums, num_tokens = packed.chunk(2) - tracker["values"] = loss_sums / num_tokens.clamp(min=1) + values = tracker["loss_sums"] + if tracker.get('reduce_group') is not None: + torch.distributed.all_reduce(values, group=tracker['reduce_group']) + if tracker.get('avg_group') is not None: + torch.distributed.all_reduce( + values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG + ) + tracker["values"] = values @staticmethod def track_mtp_metrics(loss_scale, iteration, writer, wandb_writer=None, total_loss_dict=None): diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index d3571152966..29990e9cae7 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1327,7 +1327,6 @@ 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/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py index 8b334a0e73f..eeb3e82c541 100644 --- a/megatron/training/datasets/varlen_dataset.py +++ b/megatron/training/datasets/varlen_dataset.py @@ -406,10 +406,15 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: loss_mask = torch.ones(max_len, dtype=torch.float32) loss_mask[valid_len:] = 0.0 # mask the right-padded tail by position loss_mask[labels == IGNORE_INDEX] = 0.0 + # Keep physical padding separate from the LM loss mask: prompt + # tokens may be loss-masked but must still participate in MoE. + padding_mask = torch.zeros(max_len, dtype=torch.bool) + padding_mask[valid_len:] = True return { 'tokens': input_ids, 'labels': labels, 'loss_mask': loss_mask, + 'padding_mask': padding_mask, 'position_ids': torch.arange(max_len, dtype=torch.int64), } diff --git a/megatron/training/training.py b/megatron/training/training.py index ae6216c260b..f1e9717508a 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2954,10 +2954,10 @@ def training_log( # Log MTP metrics. if args.mtp_num_layers is not None: - # MTP tracker stores raw loss sums and token counts, so after reduction - # tracker["values"] already equals the per-token loss (loss_sum / num_tokens) - # aggregated across all ranks and microbatches. No further scaling needed. - mtp_loss_scale = 1.0 + # The tracker stores a sum of normalized microbatch losses. + # Sequence-packing schedulers may change the number of microbatches for + # this step, so use the scheduled count passed to training_log. + mtp_loss_scale = 1 / (num_microbatches or get_num_microbatches()) MTPLossLoggingHelper.track_mtp_metrics( mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict ) diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py index 7b5e3b46fe1..30a22acd77a 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -545,6 +545,9 @@ def _broadcast(item): if "attention_mask" not in data else data["attention_mask"].cuda(non_blocking=True) ), + 'padding_mask': ( + None if "padding_mask" not in data else data["padding_mask"].cuda(non_blocking=True) + ), 'position_ids': data["position_ids"].cuda(non_blocking=True), 'cu_seqlens': ( None if "cu_seqlens" not in data else data["cu_seqlens"].cuda(non_blocking=True) @@ -559,11 +562,18 @@ def _broadcast(item): ), } + has_padding_mask = torch.tensor( + [batch['padding_mask'] is not None], + dtype=torch.int64, + device=torch.cuda.current_device(), + ) + _broadcast(has_padding_mask) + def _broadcast_cu_seqlens(cu_seqlens): if getattr(args, 'cuda_graph_impl', 'none') == 'full_iteration': - assert cu_seqlens is None, ( - "cu_seqlens is not supported with cuda_graph_impl=full_iteration" - ) + assert ( + cu_seqlens is None + ), "cu_seqlens is not supported with cuda_graph_impl=full_iteration" return dev = torch.cuda.current_device() n = 0 if cu_seqlens is None else int(cu_seqlens.numel()) @@ -589,6 +599,7 @@ def _broadcast_cu_seqlens(cu_seqlens): _broadcast(batch['tokens']) _broadcast(batch['labels']) _broadcast(batch['loss_mask']) + _broadcast(batch['padding_mask']) _broadcast(batch['attention_mask']) _broadcast(batch['position_ids']) _broadcast_cu_seqlens(batch['cu_seqlens']) @@ -597,6 +608,7 @@ def _broadcast_cu_seqlens(cu_seqlens): elif mpu.is_pipeline_first_stage(): _broadcast(batch['tokens']) + _broadcast(batch['padding_mask']) _broadcast(batch['attention_mask']) _broadcast(batch['position_ids']) _broadcast_cu_seqlens(batch['cu_seqlens']) @@ -608,9 +620,26 @@ def _broadcast_cu_seqlens(cu_seqlens): # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage. _broadcast(batch['labels']) _broadcast(batch['loss_mask']) + _broadcast(batch['padding_mask']) _broadcast(batch['attention_mask']) + else: + # SBHD validation needs physical padding metadata on intermediate + # stages because those stages may also contain MoE layers. + _broadcast(batch['padding_mask']) + batch['tokens'] = None + batch['labels'] = None + batch['loss_mask'] = None + batch['attention_mask'] = None + batch['position_ids'] = None + batch['cu_seqlens'] = None + batch['max_seqlen'] = None + batch['local_cp_size'] = None + else: + has_padding_mask = torch.empty(1, dtype=torch.int64, device=torch.cuda.current_device()) + _broadcast(has_padding_mask) + if args.dynamic_context_parallel: seq_len = torch.tensor(0, dtype=torch.int32, device=torch.cuda.current_device()) _broadcast(seq_len) @@ -621,6 +650,11 @@ def _broadcast_cu_seqlens(cu_seqlens): tokens = torch.empty(shape, dtype=torch.int64, device=torch.cuda.current_device()) labels = torch.empty(shape, dtype=torch.int64, device=torch.cuda.current_device()) loss_mask = torch.empty(shape, dtype=torch.float32, device=torch.cuda.current_device()) + padding_mask = ( + torch.empty(shape, dtype=torch.bool, device=torch.cuda.current_device()) + if bool(has_padding_mask.item()) + else None + ) if args.create_attention_mask_in_dataloader: shape_attention_mask = ( (args.micro_batch_size, 1, args.seq_length, args.seq_length) @@ -666,6 +700,7 @@ def _broadcast_cu_seqlens(): _broadcast(tokens) _broadcast(labels) _broadcast(loss_mask) + _broadcast(padding_mask) _broadcast(attention_mask) _broadcast(position_ids) cu_seqlens = _broadcast_cu_seqlens() @@ -677,6 +712,7 @@ def _broadcast_cu_seqlens(): loss_mask = None _broadcast(tokens) + _broadcast(padding_mask) _broadcast(attention_mask) _broadcast(position_ids) cu_seqlens = _broadcast_cu_seqlens() @@ -693,12 +729,26 @@ def _broadcast_cu_seqlens(): _broadcast(labels) _broadcast(loss_mask) + _broadcast(padding_mask) _broadcast(attention_mask) + else: + tokens = None + labels = None + loss_mask = None + attention_mask = None + position_ids = None + cu_seqlens = None + max_seqlen = None + local_cp_size = None + + _broadcast(padding_mask) + batch = { 'tokens': tokens, 'labels': labels, 'loss_mask': loss_mask, + 'padding_mask': padding_mask, 'attention_mask': attention_mask, 'position_ids': position_ids, 'cu_seqlens': cu_seqlens, diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 9265e8a832a..16e9d4595b7 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -142,9 +142,11 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): # 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) + needs_padding_mask = args.use_varlen_dataset and args.varlen_sbhd_validation if ( not is_first_or_last_pipeline_stage(vp_stage) and not is_packed_sequence + and not needs_padding_mask and ((not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage))) ): return None, None, None, None, None, None, None @@ -197,7 +199,9 @@ 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_packed_sequences + 1 entries. - padding_mask = None + # SBHD validation samples carry physical right-padding metadata. CP has + # already partitioned it with the other sequence-dimension tensors. + padding_mask = batch.get('padding_mask') 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) @@ -365,7 +369,12 @@ def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): config = core_transformer_config_from_args(args) if mpu.get_tensor_model_parallel_rank() != 0: return False - elif is_packed_sequence: + elif is_packed_sequence or ( + getattr(args, 'use_varlen_dataset', False) + and getattr(args, 'varlen_sbhd_validation', False) + ): + # Packed THD and SBHD validation both need padding metadata on every + # pipeline stage so each MoE layer excludes physical padding. return True return is_first_or_last_pipeline_stage(vp_stage) or mtp_on_this_rank( config, ignore_virtual=False, vp_stage=vp_stage diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py index c158a365b46..2015d393445 100644 --- a/tests/unit_tests/data/test_varlen_dataset.py +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -569,12 +569,87 @@ def test_getitem_sbhd_pads_to_seq_length_and_masks_tail(): ds = _make_varlen(["abc"], _make_config(tok, seq_length=8, sbhd=True)) out = ds[0] # SBHD emits fixed [seq_length] samples with no packing metadata. - assert set(out) == {"tokens", "labels", "loss_mask", "position_ids"} + assert set(out) == {"tokens", "labels", "loss_mask", "padding_mask", "position_ids"} assert out["tokens"].numel() == 8 loss_mask = out["loss_mask"].tolist() + padding_mask = out["padding_mask"].tolist() # tokens=[a,b,c,eod]: valid_len=3 -> first 3 kept (incl. real eod), rest masked. assert loss_mask[0:3] == [1.0, 1.0, 1.0] assert all(v == 0.0 for v in loss_mask[3:]) + assert padding_mask[0:3] == [False, False, False] + assert all(padding_mask[3:]) + assert out["padding_mask"].dtype == torch.bool + + +def test_sbhd_padding_mask_is_partitioned_with_tokens(monkeypatch): + """CP zigzag slicing must select identical token and padding-mask positions.""" + from megatron.core.utils import get_pretrain_batch_on_this_cp_rank + + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group: 4) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group: 1) + + tokens = torch.arange(16, dtype=torch.int64).view(1, 16) + padding_mask = tokens >= 10 + batch = {"tokens": tokens.clone(), "padding_mask": padding_mask.clone()} + result = get_pretrain_batch_on_this_cp_rank(batch, cp_group=object()) + + expected_indices = torch.tensor([2, 3, 12, 13]) + assert torch.equal(result["tokens"], tokens.index_select(1, expected_indices)) + assert torch.equal(result["padding_mask"], padding_mask.index_select(1, expected_indices)) + + +def test_sbhd_get_batch_returns_dataset_padding_mask(monkeypatch): + """The dataset padding mask must survive the pretrain_gpt batch handoff.""" + import pretrain_gpt + + padding_mask = torch.tensor([[False, False, True, True]], dtype=torch.bool) + source_batch = { + "tokens": torch.tensor([[1, 2, 0, 0]], dtype=torch.int64), + "labels": torch.tensor([[2, 0, 0, 0]], dtype=torch.int64), + "loss_mask": torch.tensor([[1.0, 1.0, 0.0, 0.0]]), + "padding_mask": padding_mask, + "attention_mask": None, + "position_ids": torch.arange(4, dtype=torch.int64).view(1, 4), + } + args = SimpleNamespace( + sequence_packing_scheduler=None, + sft=False, + use_varlen_dataset=True, + varlen_sbhd_validation=True, + dynamic_context_parallel=False, + ) + config = SimpleNamespace( + virtual_pipeline_model_parallel_size=None, pad_packed_seq_alignment=None + ) + + monkeypatch.setattr(pretrain_gpt, "get_args", lambda: args) + monkeypatch.setattr(pretrain_gpt, "core_transformer_config_from_args", lambda _: config) + # Exercise an intermediate PP stage: it must not take the legacy early return. + monkeypatch.setattr(pretrain_gpt, "is_first_or_last_pipeline_stage", lambda _: False) + monkeypatch.setattr(pretrain_gpt, "mtp_on_this_rank", lambda *args, **kwargs: False) + monkeypatch.setattr( + pretrain_gpt, "get_batch_on_this_tp_rank", lambda *args, **kwargs: source_batch.copy() + ) + monkeypatch.setattr(pretrain_gpt, "get_batch_on_this_cp_rank", lambda batch: batch) + + *_, returned_padding_mask = pretrain_gpt.get_batch(iter(())) + assert torch.equal(returned_padding_mask, padding_mask) + + +def test_sbhd_dataset_is_built_on_intermediate_pipeline_stage(monkeypatch): + """Every PP stage needs SBHD padding metadata for its local MoE layers.""" + import pretrain_gpt + + args = SimpleNamespace(use_varlen_dataset=True, varlen_sbhd_validation=True) + monkeypatch.setattr(pretrain_gpt, "get_args", lambda: args) + monkeypatch.setattr( + pretrain_gpt, "core_transformer_config_from_args", lambda _: SimpleNamespace() + ) + monkeypatch.setattr(pretrain_gpt.mpu, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(pretrain_gpt, "is_first_or_last_pipeline_stage", lambda _: False) + monkeypatch.setattr(pretrain_gpt, "mtp_on_this_rank", lambda *args, **kwargs: False) + + assert pretrain_gpt.is_dataset_built_on_rank() is True def test_mock_getitem_thd_keys_and_pad_fallback(): @@ -724,6 +799,8 @@ def test_sbhd_validation_dataloader_uses_default_collate(): assert batch["tokens"].shape == (mbs, seq_len) assert batch["labels"].shape == (mbs, seq_len) assert batch["loss_mask"].shape == (mbs, seq_len) + assert batch["padding_mask"].shape == (mbs, seq_len) + assert batch["padding_mask"].dtype == torch.bool finally: destroy_global_vars() Utils.destroy_model_parallel() diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index fe85b3ad8f9..bf6f4dd5037 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -90,6 +90,7 @@ def __init__( self, total_seq_length: int, sequence_lengths: list, + padded_sequence_lengths: list = None, local_cp_size: int = None, device: str = "cuda", seed: int = 42, @@ -104,12 +105,19 @@ def __init__( """ self.total_seq_length = total_seq_length self.sequence_lengths = sequence_lengths + self.padded_sequence_lengths = padded_sequence_lengths or sequence_lengths self.local_cp_size = local_cp_size self.device = device self.seed = seed - assert ( - sum(self.sequence_lengths) == total_seq_length - ), f"Sequence lengths sum {sum(self.sequence_lengths)} != total {total_seq_length}" + assert len(self.sequence_lengths) == len(self.padded_sequence_lengths) + assert all( + real <= padded + for real, padded in zip(self.sequence_lengths, self.padded_sequence_lengths) + ) + assert sum(self.padded_sequence_lengths) == total_seq_length, ( + f"Padded sequence lengths sum {sum(self.padded_sequence_lengths)} " + f"!= total {total_seq_length}" + ) def __iter__(self): """Interface for the data iterator.""" @@ -125,24 +133,34 @@ def __next__(self): # Create position_ids that reset for each sequence (THD format) position_ids = [] - for seq_len in self.sequence_lengths: + for seq_len, padded_seq_len in zip(self.sequence_lengths, self.padded_sequence_lengths): position_ids.extend(range(seq_len)) + position_ids.extend([0] * (padded_seq_len - seq_len)) position_ids = torch.tensor(position_ids, dtype=torch.int64, device=dev) # Labels are tokens shifted by 1 for easy verification labels = tokens + 1 # Loss mask: 1.0 for all positions except padding (none here) - loss_mask = torch.ones(self.total_seq_length, dtype=torch.float32, device=dev) + loss_mask = [] + for seq_len, padded_seq_len in zip(self.sequence_lengths, self.padded_sequence_lengths): + loss_mask.extend([1.0] * seq_len) + loss_mask.extend([0.0] * (padded_seq_len - seq_len)) + loss_mask = torch.tensor(loss_mask, dtype=torch.float32, device=dev) # Create cu_seqlens for variable-length packed sequences cu_seqlens = [0] for seq_len in self.sequence_lengths: cu_seqlens.append(cu_seqlens[-1] + seq_len) cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=dev) - cu_seqlens_padded = cu_seqlens.clone() + cu_seqlens_padded = [0] + for seq_len in self.padded_sequence_lengths: + cu_seqlens_padded.append(cu_seqlens_padded[-1] + seq_len) + cu_seqlens_padded = torch.tensor(cu_seqlens_padded, dtype=torch.int32, device=dev) - max_seqlen = torch.tensor([max(self.sequence_lengths)], dtype=torch.int32, device=dev) + max_seqlen = torch.tensor( + [max(self.padded_sequence_lengths)], dtype=torch.int32, device=dev + ) batch = { "tokens": tokens, @@ -283,14 +301,14 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc if tp_rank == 0: # Use deterministic seed based on DP rank so same data within TP/PP/CP group dp_rank = parallel_state.get_data_parallel_rank() - sequence_lengths = [1024, 2048, 512, 1536, 3072] - assert ( - sum(sequence_lengths) == args.seq_length - ), f"Sequence lengths sum {sum(sequence_lengths)} != total {args.seq_length}" + sequence_lengths = [1000, 2040, 500, 1500, 3000] + padded_sequence_lengths = [1024, 2048, 512, 1536, 3072] + assert sum(padded_sequence_lengths) == args.seq_length data_iterator = iter( MockVariableLengthSequencePackingDataIterator( total_seq_length=args.seq_length, sequence_lengths=sequence_lengths, + padded_sequence_lengths=padded_sequence_lengths, local_cp_size=local_cp_size, seed=42 + dp_rank, ) @@ -313,7 +331,7 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc 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." + assert padding_mask.any(), "Mock data intentionally has per-sequence padding." # Get parallel state info tp_rank = parallel_state.get_tensor_model_parallel_rank() @@ -352,6 +370,11 @@ def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp, dynamic_cp, loc # ===================================================================== assert packed_seq_params is not None assert packed_seq_params.qkv_format == "thd" + assert packed_seq_params.pad_between_seqs is None + assert not torch.equal( + packed_seq_params.cu_seqlens_q, packed_seq_params.cu_seqlens_q_padded + ) + assert packed_seq_params.cu_seqlens_q[-1] < packed_seq_params.cu_seqlens_q_padded[-1] test_keys = [ "cu_seqlens_q", diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index def3ebc04d6..d5f89a44bca 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -134,9 +134,13 @@ def test_aux_loss(self): @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_router_with_padding_mask(self): + @pytest.mark.parametrize("router_fusion", [False, True]) + def test_router_with_padding_mask(self, router_fusion): """Test that padding mask correctly excludes padding tokens from routing.""" + if router_fusion and not HAVE_ROUTER_FUSION: + pytest.skip("TE fused router ops not available") self.router = self.router.cuda() + self.router.config.moe_router_fusion = router_fusion seq_len = 32 batch_size = 2 hidden_size = self.router.config.hidden_size @@ -176,6 +180,10 @@ def test_router_with_padding_mask(self): self.router.config.num_moe_experts, ) + padding_rows = padding_mask.reshape(-1) + assert torch.count_nonzero(probs_with_mask[padding_rows]) == 0 + assert not routing_map_with_mask[padding_rows].any() + # Verify that probs for valid tokens are similar assert torch.equal(probs_valid_part, probs_without_mask) diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py index e769dab664f..bd20281ac41 100644 --- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py +++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py @@ -502,14 +502,15 @@ def _to_cu_seqlens(seqlens): def _make_thd_packed_seq_params(seqlens, cp_size, tp_size): padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size) + cu_seqlens = _to_cu_seqlens(seqlens) cu_seqlens_padded = _to_cu_seqlens(padded_seqlens) max_seqlen = max(padded_seqlens) - # Match get_batch_on_this_rank_for_sequence_packing(): TE consumes padded - # cumulative lengths as both cu_seqlens and cu_seqlens_padded for THD. + # Match the runtime contract: original boundaries count valid tokens while + # padded boundaries describe physical THD storage consumed by attention. return PackedSeqParams( qkv_format="thd", - cu_seqlens_q=cu_seqlens_padded, - cu_seqlens_kv=cu_seqlens_padded, + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, cu_seqlens_q_padded=cu_seqlens_padded, cu_seqlens_kv_padded=cu_seqlens_padded, max_seqlen_q=max_seqlen, diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 849d8a4d42c..2a1b4b7bc39 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -685,8 +685,8 @@ def test_forward_backward(self, tmp_path_dist_ckpt, tp, cp, full_recompute): labels=labels, loss_mask=loss_mask, ) - # forward only fills raw loss_sums / num_tokens. Trigger the reduction - # so tracker["values"] (per-token loss across DP+CP) becomes available. + # Forward accumulates normalized losses. Trigger the DP+CP + # reduction so tracker["values"] becomes available. MTPLossLoggingHelper.reduce_loss_in_tracker() tracker = MTPLossLoggingHelper.tracker assert "values" in tracker @@ -741,9 +741,7 @@ def set_ckpt_path(ckpt_path): labels=labels, loss_mask=loss_mask, ) - # reduce_loss_in_tracker performs sum-reduce of loss_sums and - # num_tokens across DP+CP, then computes sum/sum -- already the - # correct global per-token loss, no extra CP averaging needed. + # Combine normalized loss contributions across DP+CP. MTPLossLoggingHelper.reduce_loss_in_tracker() tracker = MTPLossLoggingHelper.tracker assert "values" in tracker @@ -845,8 +843,7 @@ def test_packed_sequences(self, tp, cp): assert output.shape[0] == 1 # batch size assert output.shape[1] == total_seq_length - # Verify MTP loss was computed; reduce raw loss_sums/num_tokens into - # tracker["values"] (per-token loss) first. + # Verify MTP loss was computed; reduce local contributions first. MTPLossLoggingHelper.reduce_loss_in_tracker() tracker = MTPLossLoggingHelper.tracker assert "values" in tracker @@ -1234,7 +1231,7 @@ def test_save_metrics_to_tracker(self): assert tracker["avg_group"] is None def test_save_loss_to_tracker(self): - """Test saving loss sum and token count to tracker.""" + """Test saving a normalized loss to the tracker.""" loss_sum = torch.tensor(1.3) num_tokens = torch.tensor(5.0) layer_number = 2 @@ -1247,14 +1244,11 @@ def test_save_loss_to_tracker(self): num_layers=num_layers, ) - # Tracker now stores raw loss sums and token counts; per-token loss - # is computed in reduce_loss_in_tracker. assert "loss_sums" in MTPLossLoggingHelper.tracker - assert "num_tokens" in MTPLossLoggingHelper.tracker assert MTPLossLoggingHelper.tracker["loss_sums"].shape == (num_layers,) - assert MTPLossLoggingHelper.tracker["num_tokens"].shape == (num_layers,) - assert MTPLossLoggingHelper.tracker["loss_sums"][layer_number] == loss_sum - assert MTPLossLoggingHelper.tracker["num_tokens"][layer_number] == num_tokens + assert torch.isclose( + MTPLossLoggingHelper.tracker["loss_sums"][layer_number], loss_sum / num_tokens + ) assert MTPLossLoggingHelper.tracker["reduce_group"] is None assert MTPLossLoggingHelper.tracker["avg_group"] is None @@ -1271,7 +1265,7 @@ def __init__(self, gather_output): assert _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=True), False) is True def test_track_mtp_metrics(self): - """Test tracking MTP metrics including token-weighted loss and acceptance rate.""" + """Test tracking normalized MTP loss and acceptance rate.""" loss_sum = torch.tensor(2.3) num_tokens = torch.tensor(1.0) num_layers = self.num_layers @@ -1371,10 +1365,36 @@ def log(self, metrics, iteration): # Verify tracker is cleaned assert torch.all(MTPLossLoggingHelper.tracker["loss_sums"] == 0) - assert torch.all(MTPLossLoggingHelper.tracker["num_tokens"] == 0) assert MTPLossLoggingHelper.tracker["reduce_group"] is None assert MTPLossLoggingHelper.tracker["avg_group"] is None + def test_microbatch_means_are_not_globally_token_weighted(self): + """MTP logging preserves the pre-#4226 microbatch-normalized semantics.""" + MTPLossLoggingHelper.save_loss_to_tracker( + loss_sum=torch.tensor(8.0), num_tokens=torch.tensor(2.0), layer_number=0, num_layers=1 + ) + MTPLossLoggingHelper.save_loss_to_tracker( + loss_sum=torch.tensor(4.0), num_tokens=torch.tensor(4.0), layer_number=0, num_layers=1 + ) + + class DummyWriter: + def __init__(self): + self.scalars = {} + + def add_scalar(self, name, value, iteration): + self.scalars[name] = value + + writer = DummyWriter() + MTPLossLoggingHelper.track_mtp_metrics( + loss_scale=0.5, iteration=1, writer=writer, total_loss_dict={} + ) + + logged_loss = torch.as_tensor(writer.scalars["mtp_1 loss"]) + microbatch_mean_average = torch.tensor(((8.0 / 2.0) + (4.0 / 4.0)) / 2.0) + global_token_weighted = torch.tensor((8.0 + 4.0) / (2.0 + 4.0)) + assert torch.isclose(logged_loss, microbatch_mean_average) + assert not torch.isclose(logged_loss, global_token_weighted) + def test_track_mtp_loss_preserves_legacy_normalized_loss_semantics(self): """MTP loss logging should not become token-weighted when acceptance counters are added.""" first_loss = torch.tensor(10.0) @@ -1566,8 +1586,8 @@ def test_forward_backward_mamba(self, tmp_path_dist_ckpt, tp, cp): labels=labels, loss_mask=loss_mask, ) - # forward only fills raw loss_sums / num_tokens. Reduce them first so - # tracker["values"] (per-token loss across DP+CP) becomes available. + # Forward accumulates normalized losses. Reduce them first so + # tracker["values"] becomes available. MTPLossLoggingHelper.reduce_loss_in_tracker() tracker = MTPLossLoggingHelper.tracker assert "values" in tracker @@ -1617,8 +1637,7 @@ def set_ckpt_path(ckpt_path): labels=labels, loss_mask=loss_mask, ) - # reduce_loss_in_tracker already computes the cross-DP+CP per-token - # loss (sum/sum), no extra CP averaging needed. + # Combine normalized loss contributions across DP+CP. MTPLossLoggingHelper.reduce_loss_in_tracker() tracker = MTPLossLoggingHelper.tracker assert "values" in tracker @@ -1922,8 +1941,7 @@ def model_provider( ) assert torch.isfinite(output).all(), f"Non-finite output (TP={tp})" - # Reduce raw loss_sums/num_tokens into tracker["values"] (per-token - # loss across DP+CP) before reading. + # Reduce normalized loss contributions before reading. MTPLossLoggingHelper.reduce_loss_in_tracker() tracker = MTPLossLoggingHelper.tracker assert "values" in tracker, f"MTP loss not logged (TP={tp})" diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 92a81b2fcdc..b619b3d2643 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -268,18 +268,27 @@ def teardown_method(self): @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_generic_alignment_appends_dummy_padding_sequence(self): - """Generic THD padding covers tail slots with an independent dummy sequence.""" + """The optional dummy mode adds a zero-valid-token sequence.""" seqlens, total_T = [50, 30], 80 psp = _make_psp(seqlens) orig = psp.cu_seqlens_q.clone() p_tok, _, _, _, p, mask = pad_sequence_for_thd( - torch.ones(1, total_T, device="cuda"), None, None, None, psp, alignment=64 + torch.ones(1, total_T, device="cuda"), + None, + None, + None, + psp, + alignment=64, + pad_by_appending_dummy_seq=True, ) 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 + expected_valid = torch.cat((orig, torch.tensor([80], dtype=orig.dtype, device=orig.device))) + expected_padded = torch.cat( + (orig, torch.tensor([128], dtype=orig.dtype, device=orig.device)) + ) + assert torch.equal(p.cu_seqlens_q, expected_valid) + assert torch.equal(p.cu_seqlens_q_padded, expected_padded) + assert p.pad_between_seqs is None assert mask.shape == (1, 128) assert not mask[0, :total_T].any() and mask[0, total_T:].all() @@ -297,10 +306,10 @@ def test_cp_alignment_uses_global_cu_seqlens_length(self): ) assert p_tok.shape[-1] >= local_T - assert p.cu_seqlens_q[-1].item() == 256 + assert p.cu_seqlens_q[-1].item() == 140 assert p.cu_seqlens_q_padded[-1].item() == 256 - assert p.max_seqlen_q == 140 - assert p.max_seqlen_kv == 140 + assert p.max_seqlen_q == 256 + assert p.max_seqlen_kv == 256 assert mask.shape[-1] == p_tok.shape[-1] assert not mask[0, :local_T].any() @@ -318,37 +327,58 @@ def test_cp_alignment_covers_local_padding_tail(self): ) assert p_tok.shape[-1] == 1664 - assert p.cu_seqlens_q[-1].item() == 3328 + assert p.cu_seqlens_q[-1].item() == 3200 assert p.cu_seqlens_q_padded[-1].item() == 3328 + assert p.max_seqlen_q == 1728 + assert p.max_seqlen_kv == 1728 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() + def test_padding_without_dummy_extends_only_last_padded_endpoint(self): + """Default tail padding leaves real cu_seqlens unchanged.""" + cu_valid = torch.tensor([0, 3, 5], dtype=torch.int32, device="cuda") + cu_padded = torch.tensor([0, 4, 8], dtype=torch.int32, device="cuda") + psp = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_valid, + cu_seqlens_kv=cu_valid.clone(), + cu_seqlens_q_padded=cu_padded, + cu_seqlens_kv_padded=cu_padded.clone(), + max_seqlen_q=4, + max_seqlen_kv=4, + ) + initial_padding_mask = torch.tensor( + [[False, False, False, True, False, False, True, True]], dtype=torch.bool, device="cuda" + ) p_tok, _, _, _, p, mask = pad_sequence_for_thd( - torch.ones(1, total_T, device="cuda"), + torch.ones(1, 8, device="cuda"), None, None, None, psp, - alignment=64, - pad_by_appending_dummy_seq=False, + target_len=10, + padding_mask=initial_padding_mask, + ) + expected_padded = torch.tensor([0, 4, 10], dtype=torch.int32, device="cuda") + assert p_tok.shape == (1, 10) + assert torch.equal(p.cu_seqlens_q, cu_valid) + assert torch.equal(p.cu_seqlens_kv, cu_valid) + assert torch.equal(p.cu_seqlens_q_padded, expected_padded) + assert torch.equal(p.cu_seqlens_kv_padded, expected_padded) + assert p.max_seqlen_q == 6 + assert p.max_seqlen_kv == 6 + assert p.pad_between_seqs is None + assert torch.equal( + mask, + torch.tensor( + [[False, False, False, True, False, False, True, True, True, True]], + dtype=torch.bool, + device="cuda", + ), ) - assert p_tok.shape == (1, 128) - assert torch.equal(p.cu_seqlens_q, orig) - assert torch.equal(p.cu_seqlens_q_padded, orig) - assert p.max_seqlen_q == max(seqlens) - assert p.max_seqlen_kv == max(seqlens) - assert p.pad_between_seqs is False - assert mask.shape == (1, 128) - assert not mask[0, :total_T].any() and mask[0, total_T:].all() @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -375,24 +405,27 @@ 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" + expected_valid_cu = torch.tensor( + [0, 100, 150, 180, 180, 180, 180, 180, 180], 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) + expected_padded_cu = torch.tensor( + [0, 100, 150, 256, 256, 256, 256, 256, 256], dtype=torch.int32, device="cuda" + ) + assert torch.equal(p_params.cu_seqlens_q, expected_valid_cu) + assert torch.equal(p_params.cu_seqlens_kv, expected_valid_cu) + assert torch.equal(p_params.cu_seqlens_q_padded, expected_padded_cu) + assert torch.equal(p_params.cu_seqlens_kv_padded, expected_padded_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_params.pad_between_seqs is None 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.""" + def test_eager_pad_to_max_extends_last_padded_sequence(self): + """Eager pad-to-max leaves real boundaries unchanged.""" seqlens, total_T, target_len = [50, 30], 80, 8192 psp = _make_psp(seqlens) orig_cu = psp.cu_seqlens_q.clone() @@ -415,16 +448,15 @@ def test_eager_pad_to_max_adds_dummy_padding_sequence(self): ) 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 + expected_padded = orig_cu.clone() + expected_padded[-1] = target_len + assert torch.equal(p_params.cu_seqlens_q, orig_cu) + assert torch.equal(p_params.cu_seqlens_q_padded, expected_padded) + assert p_params.cu_seqlens_q.shape == orig_cu.shape + assert p_params.max_seqlen_q == target_len - orig_cu[-2].item() + assert p_params.max_seqlen_kv == target_len - orig_cu[-2].item() assert p_params.total_tokens == target_len - assert p_params.pad_between_seqs is False + assert p_params.pad_between_seqs is None assert p_mask.shape == (1, target_len) assert not p_mask[0, :total_T].any() assert p_mask[0, total_T:].all() @@ -477,7 +509,7 @@ def test_padding_mask_preserves_existing_padding(self): @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_cu_seqlens_fill_value(self): - """Static cu padding repeats dummy valid/padded cumulative values.""" + """Static cu padding repeats distinct valid and physical endpoints.""" seqlens, total_T = [50, 30], 80 _, _, _, _, p, _ = pad_sequence_for_thd( torch.ones(1, total_T, device="cuda"), @@ -489,10 +521,42 @@ 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:] == 128).all() - assert p.cu_seqlens_q_padded[0] == 0 and p.cu_seqlens_q_padded[2] == 80 + assert (p.cu_seqlens_q[3:] == 80).all() + assert p.cu_seqlens_q_padded[0] == 0 and p.cu_seqlens_q_padded[2] == 128 assert (p.cu_seqlens_q_padded[3:] == 128).all() - assert p.pad_between_seqs is False + assert p.pad_between_seqs is None + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_tail_padding_updates_only_last_padded_endpoint(self): + """A physical padding tail must not increase the valid-token endpoint.""" + cu_valid = torch.tensor([0, 18, 44, 52, 96, 118], dtype=torch.int32, device="cuda") + cu_padded = torch.tensor([0, 24, 56, 64, 112, 144], dtype=torch.int32, device="cuda") + psp = PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_valid, + cu_seqlens_kv=cu_valid.clone(), + cu_seqlens_q_padded=cu_padded, + cu_seqlens_kv_padded=cu_padded.clone(), + max_seqlen_q=48, + max_seqlen_kv=48, + ) + + _, _, _, _, padded, _ = pad_sequence_for_thd( + torch.ones(1, 144, device="cuda"), None, None, None, psp, target_len=160, max_num_seqs=8 + ) + + expected_valid = torch.tensor( + [0, 18, 44, 52, 96, 118, 118, 118, 118], dtype=torch.int32, device="cuda" + ) + expected_padded = torch.tensor( + [0, 24, 56, 64, 112, 160, 160, 160, 160], dtype=torch.int32, device="cuda" + ) + assert torch.equal(padded.cu_seqlens_q, expected_valid) + assert torch.equal(padded.cu_seqlens_kv, expected_valid) + assert torch.equal(padded.cu_seqlens_q_padded, expected_padded) + assert torch.equal(padded.cu_seqlens_kv_padded, expected_padded) + assert padded.pad_between_seqs is None @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -540,7 +604,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 + assert r.pad_between_seqs is None for k, v in orig.items(): assert torch.equal(getattr(r, k), v)