diff --git a/src/mcore_bridge/model/gpts/qwen3_next.py b/src/mcore_bridge/model/gpts/qwen3_next.py index 6fcf1ea5..f7219b40 100644 --- a/src/mcore_bridge/model/gpts/qwen3_next.py +++ b/src/mcore_bridge/model/gpts/qwen3_next.py @@ -20,7 +20,7 @@ from mcore_bridge.bridge import GPTBridge from mcore_bridge.config import ModelConfig -from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger +from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger, get_num_samples from ..constant import ModelType from ..register import ModelLoader, ModelMeta, register_model @@ -470,12 +470,12 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): # Note: for packed inputs, we do not perform padding_free unpadding. # Doing so would allow different sequences to see each other; for efficiency we keep this implementation. if thd_format: + num_samples = get_num_samples(packed_seq_params) max_seqlen_q = int(packed_seq_params.max_seqlen_q) - new_hidden_states = hidden_states.new_zeros( - (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1])) - attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) + new_hidden_states = hidden_states.new_zeros((num_samples, max_seqlen_q, hidden_states.shape[-1])) + attention_mask = hidden_states.new_zeros((num_samples, max_seqlen_q), dtype=torch.bool) cu_seqlens_q = packed_seq_params.cu_seqlens_q - for i in range(packed_seq_params.num_samples): + for i in range(num_samples): start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1] attention_mask[i, :end - start] = True new_hidden_states[i, :end - start] = hidden_states[start:end, 0] diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_5.py b/src/mcore_bridge/model/mm_gpts/qwen3_5.py index b673bf73..96a01129 100644 --- a/src/mcore_bridge/model/mm_gpts/qwen3_5.py +++ b/src/mcore_bridge/model/mm_gpts/qwen3_5.py @@ -7,7 +7,7 @@ from megatron.core.transformer.attention import SelfAttentionSubmodules from megatron.core.transformer.transformer_config import TransformerConfig -from mcore_bridge.utils import get_env_args +from mcore_bridge.utils import get_env_args, get_num_samples from ..constant import ModelType from ..gpts.qwen3_next import Qwen3NextBridge, Qwen3NextLoader, resolve_gdn_attention_mask @@ -40,12 +40,12 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): # Note: for packed inputs, we do not perform padding_free unpadding. # Doing so would allow different sequences to see each other; for efficiency we keep this implementation. if thd_format: + num_samples = get_num_samples(packed_seq_params) max_seqlen_q = int(packed_seq_params.max_seqlen_q) - new_hidden_states = hidden_states.new_zeros( - (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1])) - attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) + new_hidden_states = hidden_states.new_zeros((num_samples, max_seqlen_q, hidden_states.shape[-1])) + attention_mask = hidden_states.new_zeros((num_samples, max_seqlen_q), dtype=torch.bool) cu_seqlens_q = packed_seq_params.cu_seqlens_q - for i in range(packed_seq_params.num_samples): + for i in range(num_samples): start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1] attention_mask[i, :end - start] = True new_hidden_states[i, :end - start] = hidden_states[start:end, 0] diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py index 1121bd2b..d2b60ce7 100644 --- a/src/mcore_bridge/model/modules/ple.py +++ b/src/mcore_bridge/model/modules/ple.py @@ -9,7 +9,7 @@ from torch import nn from typing import List, Optional -from ...utils.megatron_utils import reconstruct_tensor_cp, split_cp_inputs +from ...utils.megatron_utils import get_num_samples, reconstruct_tensor_cp, split_cp_inputs from .hyper_connection_gated import Qwen4ExpTextGroupedRMSNorm _MASK64 = (1 << 64) - 1 @@ -382,7 +382,7 @@ def _forward_impl( """hidden_states: [s, b, nH] (bsh) or thd [T, 1, nH]; input_ids: [b, s] or [1, T].""" thd = packed_seq_params is not None and getattr(packed_seq_params, 'qkv_format', 'bshd') == 'thd' if thd: - num_samples = packed_seq_params.num_samples + num_samples = get_num_samples(packed_seq_params) # PackedSeqParams.max_seqlen_q is declared `int` in mcore and swift # normalizes it to int, so `.item()` would raise AttributeError; # tolerate a 0-d tensor from other callers. diff --git a/src/mcore_bridge/utils/__init__.py b/src/mcore_bridge/utils/__init__.py index f4fffcdc..9f2ddeeb 100644 --- a/src/mcore_bridge/utils/__init__.py +++ b/src/mcore_bridge/utils/__init__.py @@ -3,8 +3,8 @@ from .env import get_dist_setting, get_node_setting, is_dist, is_last_rank, is_local_master, is_master from .import_utils import _LazyModule, is_flash_attn_3_available from .logger import get_logger -from .megatron_utils import (get_local_layer_specs, reconstruct_tensor_cp, roll_tensor, set_random_seed, - split_cp_inputs, unwrap_model) +from .megatron_utils import (get_local_layer_specs, get_num_samples, reconstruct_tensor_cp, roll_tensor, + set_random_seed, split_cp_inputs, unwrap_model) from .safetensors import SafetensorLazyLoader, StreamingSafetensorSaver from .torch_utils import gc_collect, get_current_device, safe_ddp_context, to_device from .utils import deep_getattr, get_env_args, json_parse_to_dict, patch_deepcopy diff --git a/src/mcore_bridge/utils/megatron_utils.py b/src/mcore_bridge/utils/megatron_utils.py index c7368b65..419ed539 100644 --- a/src/mcore_bridge/utils/megatron_utils.py +++ b/src/mcore_bridge/utils/megatron_utils.py @@ -83,6 +83,30 @@ def split_cp_inputs(inputs: torch.Tensor, return torch.cat(new_inputs, dim=dim) +def get_num_samples(packed_seq_params) -> int: + """Number of real (unpadded) samples packed into a THD micro-batch. + + ``cu_seqlens_q.numel() - 1`` is *not* this number: the data pipeline pads the packed + sequence (see swift ``get_padding_to``) and the padding tokens carry ``position_ids == 0``, + so the padded tail opens extra segments (one per token when cp_size == 1, one per + ``2 * cp_size`` tokens otherwise). Using the segment count would inflate the + ``[num_samples, max_seqlen_q, h]`` buffers built by the GDN/PLE THD paths by up to + ``padding_to - 1`` extra rows. + + ``num_samples``/``seq_lens`` are attached to ``PackedSeqParams`` by the data pipeline; + they are not mcore dataclass fields. Fall back to the segment count only when neither is + available, i.e. for callers that build ``PackedSeqParams`` directly (the convert-precision + check and the README example), which run tiny batches. + """ + num_samples = getattr(packed_seq_params, 'num_samples', None) + if num_samples is not None: + return int(num_samples) + seq_lens = getattr(packed_seq_params, 'seq_lens', None) + if seq_lens is not None: + return int(seq_lens.shape[0]) + return int(packed_seq_params.cu_seqlens_q.numel()) - 1 + + def reconstruct_tensor_cp(tensor, packed_seq_params, dim: int) -> torch.Tensor: """In CP mode, all-gather and undo the load-balanced (zigzag) chunking produced by ``split_cp_inputs``, restoring the full sequence in original