diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index a472003..c523140 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/src/mcore_bridge/config/model_config.py @@ -196,6 +196,26 @@ class ModelConfig(TransformerConfig): attention_output_gate: bool = False linear_decoupled_in_proj: bool = False + # qwen3.8-flash-next (HC + PLE + QSA) + hc_count: Optional[int] = None + hc_lowrank: Optional[int] = None + ple_layer_ids: Optional[List[int]] = None + ple_embed_dim: Optional[int] = None + ple_conv_kernel_size: Optional[int] = None + ngram_size: Optional[int] = None + heads_per_ngram: Optional[int] = None + ngram_vocab_size_base: Optional[int] = None + make_ngram_vocab_size_divisible_by: Optional[int] = None + split_ngram_parts: Optional[int] = None + ple_seed: Optional[int] = None + eos_token_id: Optional[int] = None + indexer_n_heads: Optional[int] = None + indexer_kv_heads: Optional[int] = None + indexer_head_dim: Optional[int] = None + indexer_budget: Optional[int] = None + indexer_compress_ratio: Optional[int] = None + output_gate_type: Optional[str] = None + # nemotron_h (hybrid mamba2 + attention + moe) hybrid_layer_pattern: Optional[str] = None diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py index ebd3086..ef55470 100644 --- a/src/mcore_bridge/config/parser.py +++ b/src/mcore_bridge/config/parser.py @@ -52,6 +52,23 @@ 'linear_key_head_dim': ['linear_key_head_dim'], 'linear_value_head_dim': ['linear_value_head_dim'], 'linear_conv_kernel_dim': ['linear_conv_kernel_dim'], + # qwen4_exp + 'hc_count': ['hc_count'], + 'hc_lowrank': ['hc_lowrank'], + 'ple_layer_ids': ['ple_layer_ids'], + 'ple_embed_dim': ['ple_embed_dim'], + 'ple_conv_kernel_size': ['ple_conv_kernel_size'], + 'ngram_size': ['ngram_size'], + 'heads_per_ngram': ['heads_per_ngram'], + 'ngram_vocab_size_base': ['ngram_vocab_size_base'], + 'make_ngram_vocab_size_divisible_by': ['make_ngram_vocab_size_divisible_by'], + 'split_ngram_parts': ['split_ngram_parts'], + 'indexer_n_heads': ['indexer_n_heads'], + 'indexer_kv_heads': ['indexer_kv_heads'], + 'indexer_head_dim': ['indexer_head_dim'], + 'indexer_budget': ['indexer_budget'], + 'indexer_compress_ratio': ['indexer_compress_ratio'], + 'output_gate_type': ['output_gate_type'], # dsa 'dsa_indexer_n_heads': ['index_n_heads'], 'dsa_indexer_head_dim': ['index_head_dim'], @@ -244,6 +261,39 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]: if use_mcore_gdn: res['experimental_attention_variant'] = 'gated_delta_net' res.setdefault('linear_attention_freq', 4) + elif hf_model_type == 'qwen4_exp': + use_mcore_gdn = get_env_args('USE_MCORE_GDN', bool, True) + res['layernorm_zero_centered_gamma'] = True + res['attention_output_gate'] = True + res['qk_layernorm'] = True + res['linear_decoupled_in_proj'] = True + res['moe_shared_expert_gate'] = True + if use_mcore_gdn: + res['experimental_attention_variant'] = 'gated_delta_net' + text_config = getattr(hf_config, 'text_config', hf_config) + num_layers = res['num_layers'] + linear_pattern = ['1' if t == 'linear_attention' else '0' for t in layer_types] + res['linear_attention_freq'] = f"[{','.join(linear_pattern)}]" + if res.get('num_moe_experts'): + res['moe_layer_freq'] = f"[{','.join(['1'] * num_layers)}]" + # seed is hardcoded in transformers, not in config + res['ple_seed'] = int(getattr(text_config, 'seed', 1234)) + eos_token_id = getattr(text_config, 'eos_token_id', None) + if eos_token_id is not None: + res['eos_token_id'] = eos_token_id + # These fields must come from the model config: ModelConfig carries no + # defaults for them, and a silently substituted value would corrupt the + # n-gram hash-table sharding/math at checkpoint conversion. + _required = [ + 'hc_count', 'hc_lowrank', 'ple_layer_ids', 'ple_embed_dim', 'ple_conv_kernel_size', 'ngram_size', + 'heads_per_ngram', 'ngram_vocab_size_base', 'make_ngram_vocab_size_divisible_by', 'split_ngram_parts', + 'eos_token_id', 'indexer_n_heads', 'indexer_kv_heads', 'indexer_head_dim', 'indexer_budget', + 'indexer_compress_ratio' + ] + _missing = [k for k in _required if res.get(k) is None] + if _missing: + raise ValueError(f'qwen4_exp config is missing required fields: {_missing}. ' + 'They must be provided by the model config.json.') elif llm_model_type == 'minimax_m2': res['add_qkv_bias'] = False elif llm_model_type == 'olmoe': diff --git a/src/mcore_bridge/model/constant.py b/src/mcore_bridge/model/constant.py index d05ad18..1296f34 100644 --- a/src/mcore_bridge/model/constant.py +++ b/src/mcore_bridge/model/constant.py @@ -25,6 +25,7 @@ class MLLMModelType: qwen3_omni = 'qwen3_omni' qwen3_asr = 'qwen3_asr' qwen3_5 = 'qwen3_5' + qwen4_exp = 'qwen4_exp' ovis2_5 = 'ovis2_5' internvl_chat = 'internvl_chat' diff --git a/src/mcore_bridge/model/gpt_model.py b/src/mcore_bridge/model/gpt_model.py index d6d1c10..367a91f 100644 --- a/src/mcore_bridge/model/gpt_model.py +++ b/src/mcore_bridge/model/gpt_model.py @@ -309,7 +309,7 @@ def forward( padding_mask = torch.chunk(padding_mask, tp_size, dim=1)[mpu.get_tensor_model_parallel_rank()] extra_block_kwargs['padding_mask'] = padding_mask.contiguous() - if self.config.moe_n_hash_layers > 0: + if self.config.moe_n_hash_layers > 0 or getattr(self.config, 'ple_layer_ids', None): extra_block_kwargs['input_ids'] = input_ids # Run decoder. diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py index 6b4c034..655415b 100644 --- a/src/mcore_bridge/model/gpts/__init__.py +++ b/src/mcore_bridge/model/gpts/__init__.py @@ -1,3 +1,3 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, nemotron_h, olmoe, - qwen3_emb, qwen3_next) + qwen3_emb, qwen3_next, qwen4_exp) diff --git a/src/mcore_bridge/model/gpts/qwen3_next.py b/src/mcore_bridge/model/gpts/qwen3_next.py index 0d4e290..69df8f9 100644 --- a/src/mcore_bridge/model/gpts/qwen3_next.py +++ b/src/mcore_bridge/model/gpts/qwen3_next.py @@ -79,8 +79,7 @@ def __init__(self, config: ModelConfig, hidden_size: int, eps: float = 1e-5): super().__init__() self.config = config self.eps = eps - # Initialize weight to zeros (Zero-Centered), matching HuggingFace Qwen3NextRMSNorm - self.weight = torch.nn.Parameter(torch.zeros(hidden_size)) + self.weight = torch.nn.Parameter(torch.zeros(hidden_size, dtype=config.params_dtype)) # Mark weight for SP gradient AllReduce across TP domain (consistent with TENorm/MCoreRMSNorm) setattr(self.weight, 'sequence_parallel', config.sequence_parallel) diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py new file mode 100644 index 0000000..7ce61f8 --- /dev/null +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -0,0 +1,482 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import copy +import torch +import torch.distributed as dist +import torch.nn.functional as F +from contextlib import contextmanager +from copy import deepcopy +from megatron.core import mpu +from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TENorm, TERowParallelLinear +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.ssm.gated_delta_net import GatedDeltaNetSubmodules +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import TransformerBlockSubmodules +from transformers.utils import is_torch_npu_available +from typing import List, Optional + +from mcore_bridge.utils import get_local_layer_specs, get_logger + +from ..modules import (GatedDeltaNet, QSAIndexer, Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, TransformerBlock, + TransformerLayer) +from ..register import ModelLoader +from .qwen3_next import Qwen3NextBridge, Qwen3NextRMSNorm, Qwen3NextSelfAttention + +logger = get_logger() + +_HC_WEIGHT_KEYS = ( + 'hc_norm.weight', + 'input_mix_weight_down.weight', + 'input_mix_weight_up.weight', + 'block_inject_weight.weight', +) + + +class Qwen4ExpGDN(GatedDeltaNet): + # upstream uses config.activation_func as the act_fn for both the gated output + # norm and the conv1d; but the conv1d path asserts act_fn in ['silu', 'swish'], + # so setting it to sigmoid would be rejected. Override only the output gate here. + def _apply_gated_norm(self, x, gate): + x_dtype = x.dtype + x = x.reshape(-1, x.shape[-1]) + y = self.out_norm(x) + gate = gate.reshape(-1, gate.shape[-1]) + output_gate_type = getattr(self.config, 'output_gate_type', None) + gate_act = torch.sigmoid if output_gate_type == 'sigmoid' else F.silu + y = y * gate_act(gate.float()) + return y.to(x_dtype) + + +class Qwen4ExpLayer(TransformerLayer): + # refer: transformers Qwen4ExpTextDecoderLayer + def __init__(self, config, submodules, layer_number: int = 1, **kwargs): + super().__init__(config, submodules, layer_number, **kwargs) + self.ple = None + if self.layer_number in config.ple_layer_ids: + self.ple = Qwen4ExpTextPLELayer( + config, config.ple_layer_ids.index(self.layer_number), pg_collection=self.pg_collection) + is_linear_attention = config.linear_attention_freq[self.layer_number - 1] + if not is_linear_attention and getattr(config, 'indexer_n_heads', None) is not None: + self.self_attention.indexer = QSAIndexer(config) + self.attn_hyper_connection = Qwen4ExpTextGatedResidual(config) + self.mlp_hyper_connection = Qwen4ExpTextGatedResidual(config) + + def forward(self, hidden_states: torch.Tensor, **kwargs): + attention_mask = kwargs.get('attention_mask') + packed_seq_params = kwargs.get('packed_seq_params') + attn_kwargs = dict( + attention_mask=attention_mask, + inference_context=kwargs.get('inference_context'), + rotary_pos_emb=kwargs.get('rotary_pos_emb'), + rotary_pos_cos=kwargs.get('rotary_pos_cos'), + rotary_pos_sin=kwargs.get('rotary_pos_sin'), + attention_bias=kwargs.get('attention_bias'), + packed_seq_params=packed_seq_params, + sequence_len_offset=kwargs.get('sequence_len_offset'), + ) + if self.ple is not None: + input_ids = kwargs.get('input_ids') + assert input_ids is not None, 'PLE layers require input_ids in extra_block_kwargs' + hidden_states = hidden_states + self.ple(hidden_states, input_ids, packed_seq_params) + + # attention sub-block (mirrors transformers Qwen4ExpTextDecoderLayer.forward) + hidden_states, hyper_input, injection_weights = self.attn_hyper_connection(hidden_states) + qsa_mask = self._qsa_select_mask(hidden_states, attn_kwargs) + if qsa_mask is not None: + # failed to full atention + attn_kwargs = dict(attn_kwargs, attention_mask=qsa_mask) + with self._patch_apply_rotary_pos_emb(), self._qsa_arbitrary_mask(qsa_mask is not None): + hidden_states, _ = self.self_attention(hidden_states=hidden_states, **attn_kwargs) + injection = hidden_states.unsqueeze(-2) * injection_weights.unsqueeze(-1) + hidden_states = hyper_input + injection.flatten(-2) + + # mlp sub-block + hidden_states, hyper_input, injection_weights = self.mlp_hyper_connection(hidden_states) + hidden_states = self.mlp(hidden_states) + if isinstance(hidden_states, tuple): + hidden_states = hidden_states[0] + injection = hidden_states.unsqueeze(-2) * injection_weights.unsqueeze(-1) + hidden_states = hyper_input + injection.flatten(-2) + return hidden_states, None + + @contextmanager + def _qsa_arbitrary_mask(self, enabled: bool): + """Temporarily switch the attention to `arbitrary` mask type. + + `Attention.forward` takes no `attn_mask_type` argument -- it reads + `self.attn_mask_type` (and core_attention's) -- so a custom mask is + silently ignored unless the type is flipped. Restored in `finally` so a + raising forward cannot leave the layer stuck in the slower unfused mode. + """ + if not enabled: + yield + return + attn = self.self_attention + targets = [attn] + core = getattr(attn, 'core_attention', None) + if core is not None: + targets.append(core) + saved = [(t, t.attn_mask_type) for t in targets if hasattr(t, 'attn_mask_type')] + for t, _ in saved: + t.attn_mask_type = AttnMaskType.arbitrary + try: + yield + finally: + for t, old in saved: + t.attn_mask_type = old + + def _warn_qsa_fallback_once(self, reason: str) -> None: + # warning_once dedupes on the message, so every QSA layer can call this and + # the user still sees it exactly once per distinct reason. + get_logger().warning_once(f'QSA sparse selection disabled: {reason}') + + def _qsa_select_mask(self, hidden_states, attn_kwargs): + # return None means full attention + # TODO: support padding_free & cp + indexer = getattr(self.self_attention, 'indexer', None) + if indexer is None: + return None + if attn_kwargs.get('packed_seq_params') is not None: + self._warn_qsa_fallback_once( + 'packing/padding_free is enabled (qkv_format=thd), which TE cannot combine with a ' + 'custom attention mask. QSA layers fall back to full attention -- training will ' + 'differ from sparse inference beyond the indexer budget. Pass `--padding_free false` ' + 'to enable QSA sparse selection.') + return None + if self.config.context_parallel_size > 1: + self._warn_qsa_fallback_once( + f'context_parallel_size={self.config.context_parallel_size} > 1 is not supported by ' + 'the QSA indexer yet (block pooling needs keys from other CP ranks). QSA layers fall ' + 'back to full attention -- training will differ from sparse inference beyond the ' + 'indexer budget.') + return None + rotary_pos_emb = attn_kwargs.get('rotary_pos_emb') + if rotary_pos_emb is None: + return None + return indexer.select_mask(hidden_states, rotary_pos_emb) + + +class Qwen4ExpTransformerBlock(TransformerBlock): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + config = self.config + hc_count = getattr(config, 'hc_count', 0) or 0 + if hc_count > 1 and self.has_final_layernorm_in_this_stage(): + # Final contraction (use_combine=False matches the checkpoint: + # hyper_connection_mixer has no block_inject_weight). + self.hyper_connection_mixer = Qwen4ExpTextGatedResidual(config, use_combine=False) + + +class Qwen4ExpBridge(Qwen3NextBridge): + hf_mixer_prefix = 'model.' + + def _get_hf_experts_attr(self, is_mtp: bool = False): + # The checkpoint stores experts as packed per-layer tensors + # (`mlp.experts.gate_up_proj` / `mlp.experts.down_proj`). + return True, True + + def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool): + mg_attn = None if mg_layer is None else mg_layer.self_attention + is_linear_attention = self.config.linear_attention_freq[layer_idx] + if is_linear_attention: + # GDN weights; this model has no input_layernorm for GDN layers. + hf_state_dict.update( + self._set_linear_attn_state(mg_attn, hf_state_dict, 'linear_attn.', layer_idx, to_mcore)) + else: + # Dense QSA-equivalent attention (qkv + output gate + q/k norms). + hf_state_dict.update(self._set_attn_state(mg_attn, hf_state_dict, 'self_attn.', layer_idx, to_mcore)) + has_indexer = mg_attn is not None and getattr(mg_attn, 'indexer', None) is not None + has_indexer = self._reduce_tensor_pp_group(has_indexer, to_mcore) + if has_indexer: + indexer = None if mg_attn is None else mg_attn.indexer + for mg_key, hf_key in [('index_qk_proj.weight', 'self_attn.indexer.index_qk_proj.weight'), + ('q_layernorm.weight', 'self_attn.indexer.q_layernorm.weight'), + ('k_layernorm.weight', 'self_attn.indexer.k_layernorm.weight')]: + self._set_state_dict(indexer, mg_key, hf_state_dict, hf_key, to_mcore) + return hf_state_dict + + def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool, is_mtp: bool = False): + mg_mlp = None if mg_layer is None else mg_layer.mlp + is_moe = mg_mlp is not None and hasattr(mg_mlp, 'experts') + if not to_mcore: + is_moe = torch.tensor([is_moe], dtype=torch.bool, device='cuda') + if self.pp_size > 1: + dist.all_reduce(is_moe, group=self.pp_group) + if is_moe: + hf_state_dict.update( + self._set_moe_state( + mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore, is_mtp=is_mtp)) + else: + hf_state_dict.update( + self._set_mlp_state(mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore)) + # No post_attention_layernorm in this model (HC norms replace it). + return hf_state_dict + + def _set_layer_hc(self, mg_layer, hf_state_dict, to_mcore: bool): + for key in ['attn_hyper_connection', 'mlp_hyper_connection']: + hyper_connection = None if mg_layer is None else getattr(mg_layer, key) + for weight_key in _HC_WEIGHT_KEYS: + self._set_state_dict(hyper_connection, weight_key, hf_state_dict, f'{key}.{weight_key}', to_mcore) + + # --- PLE ----------------------------------------------------------------- + # (mcore attribute name, hf checkpoint suffix). Both sides now use the + # transformers buffer names; the table is kept so these buffers stay on + # the dedicated PLE conversion path (pp broadcast, no TP split). + _PLE_NGRAM_BUFFERS = ( + ('layer_multipliers', 'layer_multipliers'), + ('ngram_heads_offsets', 'ngram_heads_offsets'), + ('ngram_heads_vocab_sizes', 'ngram_heads_vocab_sizes'), + ) + + def _get_tp_split_dim(self, mg_key): + # PLE weights are replicated across TP; in particular `conv1d.weight` + # must not use the dim-0 split that applies to the GDN conv1d. + if getattr(self, '_converting_ple', False): + return None + return super()._get_tp_split_dim(mg_key) + + def _get_pp_src_rank(self, has_module: bool) -> int: + """Global rank of the PP stage holding the module (all-reduce MAX).""" + holder = torch.tensor([dist.get_rank() if has_module else -1], dtype=torch.long, device='cuda') + if self.pp_size > 1: + dist.all_reduce(holder, op=dist.ReduceOp.MAX, group=self.pp_group) + return int(holder.item()) + + def _set_ple_ngram_embedding(self, ple, hf_state_dict, to_mcore: bool, pp_src_rank: int): + # The checkpoint shards the (padded) table into `parts` uniform row + # blocks, so shard boundaries must be derived from the padded size. + total = parts = dim = None + if ple is not None: + total = ple.ple_embedding.ngram_embedding.num_embeddings + parts = ple.ple_embedding.split_ngram_parts + dim = ple.ple_embedding.head_dim + if not to_mcore and self.pp_size > 1: + obj = [(total, parts, dim)] + dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) + total, parts, dim = obj[0] + if to_mcore and ple is None: + return + shard_size = (total + parts - 1) // parts + tp_size = mpu.get_tensor_model_parallel_world_size() + tp_rank = mpu.get_tensor_model_parallel_rank() + tp_ranks = dist.get_process_group_ranks(self.tp_group) + emb = ple.ple_embedding.ngram_embedding if ple is not None else None + dtype = emb.weight.dtype if emb is not None else self.config.params_dtype + device = emb.weight.device if emb is not None else torch.cuda.current_device() + per_partition = (emb.num_embeddings_per_partition if emb is not None else (total + tp_size - 1) // tp_size) + tp_start = tp_rank * per_partition if emb is not None else 0 + tp_end = min((tp_rank + 1) * per_partition, total) if emb is not None else 0 + if to_mcore: + for i in range(parts): + key = f'ple.ple_embedding.ngram_embedding.shard_{i}.weight' + if key not in hf_state_dict: + continue + cs, ce = i * shard_size, min((i + 1) * shard_size, total) + s, e = max(cs, tp_start), min(ce, tp_end) + if s < e: + weight = hf_state_dict[key].load() + emb.weight.data[s - tp_start:e - tp_start] = weight[s - cs:e - cs].to(emb.weight.dtype) + else: + for i in range(parts): + cs, ce = i * shard_size, min((i + 1) * shard_size, total) + pieces = [] + for r in range(tp_size): + r_start = r * per_partition + r_end = min((r + 1) * per_partition, total) + s, e = max(cs, r_start), min(ce, r_end) + if s >= e: + continue + if emb is not None and r == tp_rank: + piece = emb.weight.data[s - tp_start:e - tp_start].clone() + else: + piece = torch.empty(e - s, dim, dtype=dtype, device=device) + dist.broadcast(piece, src=tp_ranks[r], group=self.tp_group) + pieces.append(piece) + shard = torch.cat(pieces, dim=0) + if self.pp_size > 1: + dist.broadcast(shard, src=pp_src_rank, group=self.pp_group) + # Written directly into the state dict (bypasses _get_weight, + # which normally applies _target_device). + if self._target_device is not None: + shard = shard.to(self._target_device) + hf_state_dict[f'ple.ple_embedding.ngram_embedding.shard_{i}.weight'] = shard + + def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): + ple = None if mg_layer is None else getattr(mg_layer, 'ple', None) + if to_mcore: + # Only the stage owning the PLE layer reaches this path, so it + # must not run pp collectives (other pp ranks never enter here). + if ple is None: + return + pp_src_rank = None + else: + # to_hf: every pp rank calls this for every layer, so the pp + # collectives below stay in sync across stages. + pp_src_rank = self._get_pp_src_rank(ple is not None) + has_ple = self._reduce_tensor_pp_group(ple is not None, to_mcore) + if not has_ple: + return + for mg_buf, hf_buf in self._PLE_NGRAM_BUFFERS: + if to_mcore: + buffer = getattr(ple.ple_embedding, mg_buf) + buffer.copy_(hf_state_dict[f'ple.ple_embedding.{hf_buf}'].load().to(buffer.device)) + else: + tensor = getattr(ple.ple_embedding, mg_buf).data.clone() if ple is not None else None + if self.pp_size > 1: + obj = [tensor] + dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) + tensor = obj[0] + # Written directly into the state dict (bypasses _get_weight, + # which normally applies _target_device). + if tensor is not None and self._target_device is not None: + tensor = tensor.to(self._target_device) + hf_state_dict[f'ple.ple_embedding.{hf_buf}'] = tensor + self._set_ple_ngram_embedding(ple, hf_state_dict, to_mcore, pp_src_rank) + self._converting_ple = True + try: + for mg_key, hf_key in [('key_proj.weight', 'ple.key_proj.weight'), + ('value_proj.weight', 'ple.value_proj.weight'), + ('norm_key.weight', 'ple.norm_key.weight'), + ('norm_query.weight', 'ple.norm_query.weight'), + ('norm_conv.weight', 'ple.norm_conv.weight'), + ('conv1d.weight', 'ple.conv1d.weight')]: + self._set_state_dict(ple, mg_key, hf_state_dict, hf_key, to_mcore) + finally: + self._converting_ple = False + + def _set_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool): + hf_prefix = f'{hf_prefix}{layer_idx}.' + if to_mcore: + hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) + else: + hf_state_dict = {} + hf_state_dict.update(self._set_layer_attn(mg_layer, hf_state_dict, layer_idx, to_mcore)) + hf_state_dict.update(self._set_layer_mlp(mg_layer, hf_state_dict, layer_idx, to_mcore)) + self._set_layer_hc(mg_layer, hf_state_dict, to_mcore) + if (layer_idx + 1) in (self.config.ple_layer_ids or []): + self._set_layer_ple(mg_layer, hf_state_dict, to_mcore) + if to_mcore: + hf_state_dict = {} + else: + hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix) + return hf_state_dict + + def _set_final_layernorm(self, lm_model, hf_state_dict, to_mcore): + # This architecture has no final layernorm: the HC norms and the + # hyper_connection_mixer contraction replace it, and the checkpoint + # carries no `norm` weight. + pass + + def _convert_post_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore): + res = super()._convert_post_process(mg_model, hf_state_dict, hf_prefix, to_mcore) + lm_model = getattr(mg_model, 'language_model') if self.is_multimodal else mg_model + hc_count = getattr(self.config, 'hc_count', 0) or 0 + if hc_count > 1: + # The mixer only exists on the stage holding the final layernorm. + mixer_keys = ['hc_norm.weight', 'input_mix_weight_down.weight', 'input_mix_weight_up.weight'] + # super() returns {} in to_mcore mode: read from the incoming full + # state dict; in to_hf mode write into the dict super() returns. + mixer_sd = hf_state_dict if to_mcore else res + for key in mixer_keys: + self._set_state_dict(lm_model, f'decoder.hyper_connection_mixer.{key}', mixer_sd, + f'{self.hf_mixer_prefix}hyper_connection_mixer.{key}', to_mcore) + return res + + +class Qwen4ExpLoader(ModelLoader): + transformer_block = Qwen4ExpTransformerBlock + + def _get_moe_layer_pattern(self) -> List[bool]: + config = self.config + freq = config.moe_layer_freq + if isinstance(freq, list): + return [bool(x) for x in freq] + # int N: one MoE every N layers (mcore convention: i % N == N - 1). + return [i % freq == freq - 1 for i in range(config.num_layers)] + + def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): + config = self.config + config.hetereogenous_dist_checkpoint = True + # Context parallelism: PLE gathers the full sequence internally (undoing the + # CP zigzag) and GDN carries its own CP handling (a2a CP<->HP plus CP-aware + # cu_seqlens), so CP is no longer blanket-rejected here. Left unasserted so + # it can be exercised; QSA layers run dense attention, which mcore's + # attention already supports under CP. + if getattr(config, 'mtp_num_layers', None): + raise NotImplementedError('Qwen4-Exp MTP is not supported yet') + moe_spec = get_gpt_layer_with_transformer_engine_spec( + num_experts=config.num_moe_experts, + moe_grouped_gemm=config.moe_grouped_gemm, + qk_layernorm=config.qk_layernorm, + multi_latent_attention=config.multi_latent_attention, + use_kitchen=config.use_kitchen, + ) + if config.num_moe_experts is not None: + dense_spec = get_gpt_layer_with_transformer_engine_spec( + num_experts=None, + moe_grouped_gemm=config.moe_grouped_gemm, + qk_layernorm=config.qk_layernorm, + multi_latent_attention=config.multi_latent_attention, + use_kitchen=config.use_kitchen, + ) + else: + dense_spec = moe_spec + gdn_spec = ModuleSpec( + module=Qwen4ExpGDN, + submodules=GatedDeltaNetSubmodules( + in_proj=TEColumnParallelLinear, + out_norm=TENorm, + out_proj=TERowParallelLinear, + ), + ) + moe_pattern = self._get_moe_layer_pattern() + layer_specs = [] + for layer_idx, is_linear_attention in enumerate(config.linear_attention_freq): + layer_spec = deepcopy(moe_spec if moe_pattern[layer_idx] else dense_spec) + if is_linear_attention: + layer_spec.submodules.self_attention = deepcopy(gdn_spec) + else: + layer_spec.submodules.self_attention.submodules.linear_qkv = TEColumnParallelLinear + layer_spec.submodules.self_attention.module = Qwen3NextSelfAttention + if hasattr(layer_spec.submodules.self_attention.submodules, 'q_layernorm'): + layer_spec.submodules.self_attention.submodules.q_layernorm = Qwen3NextRMSNorm + if hasattr(layer_spec.submodules.self_attention.submodules, 'k_layernorm'): + layer_spec.submodules.self_attention.submodules.k_layernorm = Qwen3NextRMSNorm + # This model has no per-layer layernorms (HC norms replace them). + layer_spec.submodules.input_layernorm = IdentityOp + if hasattr(layer_spec.submodules, 'pre_mlp_layernorm'): + layer_spec.submodules.pre_mlp_layernorm = IdentityOp + layer_specs.append(layer_spec) + + local_layer_specs = get_local_layer_specs(config, layer_specs, vp_stage=vp_stage) + # No final layernorm in this model; keep the slot so the HC mixer + # stage logic (has_final_layernorm_in_this_stage) still triggers. + block_spec = TransformerBlockSubmodules(layer_specs=local_layer_specs, layer_norm=IdentityOp) + return block_spec + + def _set_transformer_layer(self, transformer_layer_spec): + for layer_spec in transformer_layer_spec.layer_specs: + layer_spec.module = Qwen4ExpLayer + + def build_model( + self, + pre_process=True, + post_process=True, + vp_stage: Optional[int] = None, + ): + model = super().build_model(pre_process, post_process, vp_stage) + lm_model = model.language_model if hasattr(model, 'language_model') else model + # The GDN out_norm uses ones-style weights, unlike the zero-centered + # HC norms, so opt it out of layernorm_zero_centered_gamma. + for layer in lm_model.decoder.layers: + if hasattr(layer.self_attention, 'out_norm'): + out_norm = layer.self_attention.out_norm + out_norm.zero_centered_gamma = False + if not is_torch_npu_available(): + assert hasattr(out_norm, 'zero_centered_gamma') + if hasattr(out_norm, 'config'): + out_norm.config = copy.copy(out_norm.config) + out_norm.config.layernorm_zero_centered_gamma = False + return model diff --git a/src/mcore_bridge/model/mm_gpts/__init__.py b/src/mcore_bridge/model/mm_gpts/__init__.py index db3e047..a931bb9 100644 --- a/src/mcore_bridge/model/mm_gpts/__init__.py +++ b/src/mcore_bridge/model/mm_gpts/__init__.py @@ -1,3 +1,3 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from . import (gemma4, glm, internvl, kimi_vl, llama4, llava, minicpmv4_6, muse_glimmer, qwen, qwen3_5, qwen3_5_gdn, - qwen3_asr, qwen3_omni, qwen3_vl) + qwen3_asr, qwen3_omni, qwen3_vl, qwen4_exp) diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_vl.py b/src/mcore_bridge/model/mm_gpts/qwen3_vl.py index 8a5a636..a6bb82d 100644 --- a/src/mcore_bridge/model/mm_gpts/qwen3_vl.py +++ b/src/mcore_bridge/model/mm_gpts/qwen3_vl.py @@ -80,7 +80,7 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config): deepstack_visual_embeds = visual_res.deepstack_features else: image_embeds, deepstack_visual_embeds = visual_res - deepstack_visual_embeds = torch.stack(deepstack_visual_embeds, dim=0) + deepstack_visual_embeds = (torch.stack(deepstack_visual_embeds, dim=0) if deepstack_visual_embeds else None) inputs_embeds = inputs_embeds + image_embeds.mean().to(device=inputs_embeds.device) * 0. visual_pos_masks = None else: @@ -125,7 +125,7 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config): inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) image_mask, video_mask = image_mask[..., 0], video_mask[..., 0] visual_pos_masks = image_mask | video_mask - if image_embeds is not None and video_embeds is not None: + if image_embeds is not None and video_embeds is not None and deepstack_visual_embeds: deepstack_image_embeds = [tensor[:image_tokens] for tensor in deepstack_visual_embeds] deepstack_video_embeds = [tensor[image_tokens:] for tensor in deepstack_visual_embeds] deepstack_visual_embeds = [] @@ -137,7 +137,7 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config): embed_joint[video_mask_joint, :] = vid_embed deepstack_visual_embeds.append(embed_joint) - deepstack_visual_embeds = torch.stack(deepstack_visual_embeds, dim=0) + deepstack_visual_embeds = (torch.stack(deepstack_visual_embeds, dim=0) if deepstack_visual_embeds else None) visual_pos_masks = visual_pos_masks.transpose(0, 1) # compat cp if self.config.context_parallel_size > 1: @@ -147,7 +147,8 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config): cu_seqlens = getattr(packed_seq_params, 'cu_seqlens_q', None) cp_mask = split_cp_inputs(cp_mask, cu_seqlens, 0) visual_pos_masks = split_cp_inputs(visual_pos_masks, cu_seqlens, 0) - deepstack_visual_embeds = deepstack_visual_embeds[:, cp_mask[(cp_mask != -1)]] + if deepstack_visual_embeds is not None: + deepstack_visual_embeds = deepstack_visual_embeds[:, cp_mask[(cp_mask != -1)]] # compat sp tp_world_size = parallel_state.get_tensor_model_parallel_world_size() tp_rank = parallel_state.get_tensor_model_parallel_rank() @@ -157,7 +158,8 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config): visual_start = 0 if tp_rank == 0 else sum(mask_tokens[:tp_rank]) visual_end = visual_start + mask_tokens[tp_rank] visual_pos_masks = visual_pos_masks[tp_rank] - deepstack_visual_embeds = deepstack_visual_embeds[:, visual_start:visual_end] + if deepstack_visual_embeds is not None: + deepstack_visual_embeds = deepstack_visual_embeds[:, visual_start:visual_end] return { 'inputs_embeds': inputs_embeds, 'visual_pos_masks': visual_pos_masks, diff --git a/src/mcore_bridge/model/mm_gpts/qwen4_exp.py b/src/mcore_bridge/model/mm_gpts/qwen4_exp.py new file mode 100644 index 0000000..1bd6903 --- /dev/null +++ b/src/mcore_bridge/model/mm_gpts/qwen4_exp.py @@ -0,0 +1,30 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from mcore_bridge.utils import get_env_args + +from ..constant import ModelType +from ..gpts.qwen4_exp import Qwen4ExpBridge, Qwen4ExpLoader, Qwen4ExpTransformerBlock +from ..register import ModelMeta, register_model +from .qwen3_5 import Qwen3_5Vit + + +class Qwen4ExpMMBridge(Qwen4ExpBridge): + hf_layers_prefix = 'model.language_model.layers' + hf_embed_key = 'model.language_model.embed_tokens.weight' + hf_mixer_prefix = 'model.language_model.' + + +class Qwen4ExpMMLoader(Qwen4ExpLoader): + transformer_block = Qwen4ExpTransformerBlock + + +use_mcore_gdn = get_env_args('USE_MCORE_GDN', bool, True) + +if use_mcore_gdn: + register_model( + ModelMeta( + ModelType.qwen4_exp, + ['qwen4_exp'], + bridge_cls=Qwen4ExpMMBridge, + visual_cls=Qwen3_5Vit, + loader=Qwen4ExpMMLoader, + )) diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py index 7269abd..7f5f2a2 100644 --- a/src/mcore_bridge/model/modules/__init__.py +++ b/src/mcore_bridge/model/modules/__init__.py @@ -4,8 +4,11 @@ from .dsa_indexer import DSAIndexer from .gated_delta_net import GatedDeltaNet from .gated_self_attention import GatedSelfAttention +from .hyper_connection_gated import Qwen4ExpTextGatedResidual, Qwen4ExpTextGroupedRMSNorm from .mtp_layer import MultiTokenPredictionLayer from .multi_latent_attention import MLASelfAttention +from .ple import Qwen4ExpTextNGramEmbedding, Qwen4ExpTextPLELayer +from .qsa_indexer import QSAIndexer from .topk_router import TopKRouter from .transformer_block import TransformerBlock from .transformer_layer import TransformerLayer diff --git a/src/mcore_bridge/model/modules/hyper_connection_gated.py b/src/mcore_bridge/model/modules/hyper_connection_gated.py new file mode 100644 index 0000000..78d373e --- /dev/null +++ b/src/mcore_bridge/model/modules/hyper_connection_gated.py @@ -0,0 +1,69 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +import torch.nn.functional as F +from torch import nn +from typing import Tuple + + +class Qwen4ExpTextGroupedRMSNorm(nn.Module): + + def __init__(self, dim: int, group_size: int, eps: float = 1e-6, dtype=None, sequence_parallel: bool = False): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + self.group_size = group_size + if dim % group_size != 0: + raise ValueError(f'hidden_size ({dim}) must be divisible by group_size ({group_size}).') + # mcore-specific: params_dtype weight + SP grad flag (HF builds fp32). + if dtype is not None: + self.weight.data = self.weight.data.to(dtype) + setattr(self.weight, 'sequence_parallel', sequence_parallel) + + def _norm(self, x: torch.Tensor) -> torch.Tensor: + x = x.reshape(*x.shape[:-1], -1, self.group_size) + out = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return out.flatten(-2) + + def forward(self, x): + output = self._norm(x.float()) + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + def extra_repr(self): + return f'{tuple(self.weight.shape)}, eps={self.eps}' + + +class Qwen4ExpTextGatedResidual(nn.Module): + + def __init__(self, config, use_combine: bool = True): + super().__init__() + self.hc_count = config.hc_count + self.hidden_size = config.hidden_size + hc_hidden_size = self.hc_count * self.hidden_size + # (transformers: Qwen4ExpTextRMSNorm(hc_hidden_size, group_size=self.hidden_size, + # eps=config.rms_norm_eps); layernorm_epsilon is mcore's rms_norm_eps.) + self.hc_norm = Qwen4ExpTextGroupedRMSNorm( + hc_hidden_size, group_size=self.hidden_size, eps=config.layernorm_epsilon, dtype=config.params_dtype) + self.input_mix_weight_down = nn.Linear(hc_hidden_size, config.hc_lowrank, bias=False, dtype=config.params_dtype) + self.input_mix_weight_up = nn.Linear(config.hc_lowrank, hc_hidden_size, bias=False, dtype=config.params_dtype) + self.block_inject_weight = nn.Linear( + hc_hidden_size, self.hc_count, bias=False, dtype=config.params_dtype) if use_combine else None + # mcore-specific: SP grad flag on the replicated weights. + for param in self.parameters(): + setattr(param, 'sequence_parallel', config.sequence_parallel) + + def forward(self, hyper_input: torch.Tensor) -> Tuple[torch.Tensor, ...]: + # Mirrors transformers `Qwen4ExpTextGatedResidual.forward` line by line. + if hyper_input.shape[-1] != self.hc_count * self.hidden_size: + raise ValueError(f'Expected {self.hc_count * self.hidden_size} hyper-connection features, ' + f'got {hyper_input.shape[-1]}.') + hyper_input_normed = self.hc_norm(hyper_input) + input_mix_weight = F.silu(self.input_mix_weight_down(hyper_input_normed) / self.hc_count) + input_mix_weight = torch.sigmoid(self.input_mix_weight_up(input_mix_weight)) + input_mix_weight = input_mix_weight.unflatten(-1, (self.hc_count, self.hidden_size)) + mixed_input = (input_mix_weight * hyper_input_normed.unflatten(-1, + (self.hc_count, self.hidden_size))).mean(dim=-2) + if self.block_inject_weight is None: + return mixed_input + injection_weights = 2 * torch.sigmoid(self.block_inject_weight(hyper_input_normed) / self.hc_count) + return mixed_input, hyper_input, injection_weights diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py new file mode 100644 index 0000000..1121bd2 --- /dev/null +++ b/src/mcore_bridge/model/modules/ple.py @@ -0,0 +1,409 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import copy +import math +import torch +import torch.nn.functional as F +from megatron.core.tensor_parallel import VocabParallelEmbedding +from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region, + scatter_to_sequence_parallel_region) +from torch import nn +from typing import List, Optional + +from ...utils.megatron_utils import reconstruct_tensor_cp, split_cp_inputs +from .hyper_connection_gated import Qwen4ExpTextGroupedRMSNorm + +_MASK64 = (1 << 64) - 1 +_SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 +_SPLITMIX_M1 = 0xBF58476D1CE4E5B9 +_SPLITMIX_M2 = 0x94D049BB133111EB +_PRIME_1 = 10007 + + +def _splitmix64(value: int) -> int: + value = (value + _SPLITMIX_GAMMA) & _MASK64 + value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64 + value = ((value ^ (value >> 27)) * _SPLITMIX_M2) & _MASK64 + return (value ^ (value >> 31)) & _MASK64 + + +def _build_layer_multipliers(unigram_vocab_size: int, ngram_size: int, ple_layer_index: int, seed: int) -> List[int]: + max_long = (1 << 63) - 1 + multiplier_max = max_long // max(unigram_vocab_size, 1) + half_bound = max(1, multiplier_max // 2) + base_seed = seed + _PRIME_1 * ple_layer_index + multipliers = [] + for index in range(ngram_size): + value = (base_seed + _SPLITMIX_GAMMA * (index + 1)) & _MASK64 + multipliers.append(2 * (_splitmix64(value) % half_bound) + 1) + return multipliers + + +def _is_prime_64(value: int) -> bool: + if value < 2: + return False + for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): + if value % prime == 0: + return value == prime + exponent = value - 1 + shifts = 0 + while exponent % 2 == 0: + exponent //= 2 + shifts += 1 + for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): + if base % value == 0: + continue + witness = pow(base, exponent, value) + if witness in (1, value - 1): + continue + for _ in range(shifts - 1): + witness = pow(witness, 2, value) + if witness == value - 1: + break + else: + return False + return True + + +def _nth_prime_after(start: int, count: int) -> int: + # Mirrors transformers `_find_nth_prime_after`. + prime = int(start) + for _ in range(count): + candidate = prime + 1 + if candidate <= 2: + prime = 2 + continue + if candidate % 2 == 0: + candidate += 1 + while not _is_prime_64(candidate): + candidate += 2 + prime = candidate + return prime + + +class Qwen4ExpTextNGramEmbedding(nn.Module): + + def __init__(self, config, ple_layer_index: int): + super().__init__() + self.ngram_size = config.ngram_size + self.context_len = self.ngram_size - 1 + self.heads_per_ngram = config.heads_per_ngram + self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram + self.ple_layer_index = ple_layer_index + # (transformers also stores unigram_vocab_size/ngram_vocab_size_base/ + # seed as attributes, consumed only by its deferred _init_weights + # buffer population; here the buffers are registered at construction, + # so no copies are kept.) + head_dim_per_ngram = config.ple_embed_dim // self.ngram_heads + # mcore-specific fail-loud: these drive the hash math and the checkpoint + # shard layout, so they must come from the model config (the parser + # validates them for qwen4_exp); a silently substituted default would + # corrupt the weight conversion. (transformers reads eos_token_id + # directly and has no split_ngram_parts: its table is replicated.) + eos_token_id = getattr(config, 'eos_token_id', None) + ple_seed = getattr(config, 'ple_seed', None) + split_ngram_parts = getattr(config, 'split_ngram_parts', None) + if eos_token_id is None or ple_seed is None or split_ngram_parts is None: + raise ValueError(f'eos_token_id/ple_seed/split_ngram_parts must be provided by the model ' + f'config (got {eos_token_id!r}/{ple_seed!r}/{split_ngram_parts!r}).') + self.eos_token_id = int(eos_token_id) + self.split_ngram_parts = int(split_ngram_parts) + self.head_dim = head_dim_per_ngram # mcore-specific: the bridge weight conversion reads it off the module. + + # Multipliers (splitmix64 derived, checkpoint-persistent). + multipliers = _build_layer_multipliers(config.padded_vocab_size, self.ngram_size, self.ple_layer_index, + int(ple_seed)) + self.register_buffer('layer_multipliers', torch.tensor(multipliers, dtype=torch.long), persistent=True) + + # Per-head prime table sizes/offsets (checkpoint-persistent), named as + # in transformers. + self.head_vocab_sizes = [] + self.head_offsets = [] + self.total_vocab_size = 0 + for head_idx in range(self.ngram_heads): + global_head_idx = self.ple_layer_index * self.ngram_heads + head_idx + size = _nth_prime_after(config.ngram_vocab_size_base - 1, global_head_idx + 1) + self.head_vocab_sizes.append(size) + self.head_offsets.append(self.total_vocab_size) + self.total_vocab_size += size + self.register_buffer( + 'ngram_heads_vocab_sizes', torch.tensor(self.head_vocab_sizes, dtype=torch.long), persistent=True) + self.register_buffer('ngram_heads_offsets', torch.tensor(self.head_offsets, dtype=torch.long), persistent=True) + ngram_vocab_divisor = config.make_ngram_vocab_size_divisible_by + padded_vocab_size = math.ceil(self.total_vocab_size / ngram_vocab_divisor) * ngram_vocab_divisor + # mcore-specific: TP-sharded table (a replicated nn.Embedding would be ~80GB). + self.ngram_embedding = VocabParallelEmbedding( + padded_vocab_size, + head_dim_per_ngram, + init_method=torch.nn.init.normal_, + config=config, + ) + + def _shift_right_ignore_eos(self, token_ids: torch.Tensor, shift: int) -> torch.Tensor: + # Mirrors transformers `_shift_right_ignore_eos`: segment-aware shift, + # segments are reset after every eos token (request boundary logic). + if shift == 0: + return token_ids + batch_size, seq_len = token_ids.shape + positions = torch.arange(seq_len, device=token_ids.device, dtype=torch.long) + eos_positions = torch.where(token_ids == self.eos_token_id, positions, -1) + previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values + previous_eos = torch.cat([eos_positions.new_full((batch_size, 1), -1), previous_eos_inclusive[:, :-1]], dim=1) + segment_start = previous_eos + 1 + position_in_segment = positions.unsqueeze(0) - segment_start + source_positions = positions - shift + gather_positions = source_positions.clamp_min(0).unsqueeze(0).expand(batch_size, -1) + shifted = token_ids.gather(dim=1, index=gather_positions) + valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0) + return torch.where(valid, shifted, token_ids.new_full((), self.eos_token_id)) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + """input_ids: [rows, L] -> embeddings [rows, L, embedding_dim]. + + Training variant of transformers ``forward(input_ids, past_key_values, + layer_idx)``: the previous context is eos padding (fresh request), + prepended here and dropped from the ids afterwards. + """ + input_ids = input_ids.long() + # mcore-specific: training has no conv-state cache; a fresh request + # starts from eos context, so prepend it (the no-previous-state branch + # of transformers `forward` returns exactly this). The latest HF + # constructor also takes a cache-only `layer_idx`, which has no + # training counterpart. + previous_context = input_ids.new_full((input_ids.shape[0], self.context_len), self.eos_token_id) + token_history = torch.cat([previous_context, input_ids], dim=-1) + shifted_tokens = [self._shift_right_ignore_eos(token_history, shift) for shift in range(self.ngram_size)] + + blocks = [] + for ngram in range(2, self.ngram_size + 1): + start_idx = (ngram - 2) * self.heads_per_ngram + end_idx = start_idx + self.heads_per_ngram + mixed_ids = shifted_tokens[0] * self.layer_multipliers[0] + for position in range(1, ngram): + mixed_ids = torch.bitwise_xor( + mixed_ids, + shifted_tokens[position] * self.layer_multipliers[position], + ) + head_vocab_sizes = self.ngram_heads_vocab_sizes[start_idx:end_idx] + head_offsets = self.ngram_heads_offsets[start_idx:end_idx] + ngram_ids = torch.remainder(mixed_ids.unsqueeze(-1), head_vocab_sizes.view(1, 1, -1)) + blocks.append(ngram_ids + head_offsets.view(1, 1, -1)) + + ngram_ids = torch.cat(blocks, dim=-1)[:, -input_ids.shape[1]:] + return self.ngram_embedding(ngram_ids).flatten(-2) + + +class Qwen4ExpTextPLELayer(nn.Module): + """Inject hashed n-gram features into every hyper-connection stream; + mirrors transformers ``Qwen4ExpTextPLELayer`` (training variant without + the inference cache and ``conv_mask``). + + PLE projects each token's concatenated n-gram embedding to a shared value + and one key per residual stream. The normalized stream activations gate + those values, then a dilated depthwise convolution adds local lexical + context. The returned tensor has width ``hc_count * hidden_size``. + + Checkpoint names under the layer prefix ``ple.``: + ple_embedding.{layer_multipliers,ngram_heads_offsets,ngram_heads_vocab_sizes} + ple_embedding.ngram_embedding.shard_{i}.weight + key_proj.weight / value_proj.weight + norm_key.weight / norm_query.weight / norm_conv.weight + conv1d.weight + """ + + def __init__(self, config, ple_layer_index: int, pg_collection=None): + super().__init__() + self.config = config + self.pg_collection = pg_collection + self.hidden_size = int(config.hidden_size) + self.hc_count = int(config.hc_count) + ple_embed_dim = int(config.ple_embed_dim) + hc_hidden_size = self.hidden_size * self.hc_count + self.ple_embedding = Qwen4ExpTextNGramEmbedding(config, ple_layer_index) + conv_kernel_size = int(config.ple_conv_kernel_size) + conv_dilation = int(config.ngram_size) + self.short_conv_state_len = (conv_kernel_size - 1) * conv_dilation + # Replicated projections in params_dtype. + self.key_proj = nn.Linear(ple_embed_dim, hc_hidden_size, bias=False, dtype=config.params_dtype) + self.value_proj = nn.Linear(ple_embed_dim, self.hidden_size, bias=False, dtype=config.params_dtype) + # mcore's config field layernorm_epsilon corresponds to HF's rms_norm_eps; + # the grouped norm is the mcore subclass adding dtype/SP-flag construction. + self.norm_key = Qwen4ExpTextGroupedRMSNorm( + hc_hidden_size, + group_size=self.hidden_size, + eps=config.layernorm_epsilon, + dtype=config.params_dtype, + sequence_parallel=config.sequence_parallel) + self.norm_query = Qwen4ExpTextGroupedRMSNorm( + hc_hidden_size, + group_size=self.hidden_size, + eps=config.layernorm_epsilon, + dtype=config.params_dtype, + sequence_parallel=config.sequence_parallel) + self.norm_conv = Qwen4ExpTextGroupedRMSNorm( + hc_hidden_size, + group_size=self.hidden_size, + eps=config.layernorm_epsilon, + dtype=config.params_dtype, + sequence_parallel=config.sequence_parallel) + self.conv1d = nn.Conv1d( + hc_hidden_size, + hc_hidden_size, + kernel_size=conv_kernel_size, + groups=hc_hidden_size, + dilation=conv_dilation, + bias=False, + dtype=config.params_dtype, + ) + nn.init.zeros_(self.conv1d.weight) + replicated_prefixes = ('key_proj', 'value_proj', 'norm_key', 'norm_query', 'norm_conv', 'conv1d') + for name, param in self.named_parameters(): + if name.startswith(replicated_prefixes) or name.endswith('norm.weight'): + # Replicated across TP; reduce grads across TP when SP is on. + setattr(param, 'sequence_parallel', config.sequence_parallel) + + def _short_conv(self, hidden_states: torch.Tensor) -> torch.Tensor: + # Training variant of transformers `_short_conv` (no conv state cache): + # causal alignment is achieved by left-padding the conv input with + # zeros (equivalent to trimming the padded conv output), which is a + # no-op on fresh-request context. + hidden_states = hidden_states.transpose(1, 2) + sequence_length = hidden_states.shape[-1] + conv_input = F.pad(hidden_states, (self.short_conv_state_len, 0)) + conv_input = conv_input[..., -(self.short_conv_state_len + sequence_length):] + return F.silu(self.conv1d(conv_input)).transpose(1, 2) + + def compute(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + """Mirrors transformers ``Qwen4ExpTextPLELayer.forward`` (training + variant); hidden_states/input_ids: [rows, L, nH]/[rows, L].""" + embeddings = self.ple_embedding(input_ids) # mcore-specific: no past_key_values cache arg + key_normed = self.norm_key(self.key_proj(embeddings)).unflatten(-1, (self.hc_count, self.hidden_size)) + value = self.value_proj(embeddings) + query_normed = self.norm_query(hidden_states).unflatten(-1, (self.hc_count, self.hidden_size)) + gate = (key_normed * query_normed).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated_value = torch.sigmoid(gate) * value.unsqueeze(-2) + gated_value_normed = self.norm_conv(gated_value.flatten(-2)) + gated_value = gated_value.flatten(-2) + # (transformers also applies conv_mask here; training relies on the + # loss mask instead.) + output = gated_value + self._short_conv(gated_value_normed) + return output + + @staticmethod + def _normalize_cu_seqlens(cu: Optional[torch.Tensor], total: int) -> Optional[torch.Tensor]: + """Align a (possibly padded/offset) cu_seqlens against the gathered full length. + + Mirrors GatedDeltaNet._resolve_cu_seqlens for the SP-only case: accept a + global cumulative layout (possibly missing its leading 0 or carrying a + per-batch offset), and return a cu whose last entry equals ``total`` so it + indexes the gathered full-sequence tensor directly. + """ + if cu is None: + return None + cu = cu.reshape(-1) + if cu.numel() > 0 and int(cu[0]) != 0: + if int(cu[-1]) == total: + cu = torch.cat([cu.new_zeros(1), cu]) + elif int(cu[-1]) - int(cu[0]) == total: + cu = cu - cu[0] + if not (cu.numel() > 0 and int(cu[-1]) == total): + raise ValueError(f'PLE cannot align cu_seqlens (last={int(cu[-1]) if cu.numel() else None}, ' + f'first={int(cu[0]) if cu.numel() else None}) with the gathered sequence ' + f'length {total} under sequence parallelism.') + return cu + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + packed_seq_params=None, + ) -> torch.Tensor: + """hidden_states: [s, b, nH] (bsh) or thd [T, 1, nH]; input_ids: [b, s] or [1, T]. + + mcore-specific wrapper (no counterpart in transformers): TP needs no + special handling (replicated weights, full-width hidden between + blocks); under SP/CP the shard is gathered to the full sequence + (undoing SP first, then the CP zigzag), ``compute`` runs on the full + sequence, and the additive output is scattered back. + """ + sp_on = ( + self.pg_collection is not None and getattr(self.config, 'sequence_parallel', False) + and getattr(self.config, 'tensor_model_parallel_size', 1) > 1) + cp_on = getattr(self.config, 'context_parallel_size', 1) > 1 + if not (sp_on or cp_on): + return self._forward_impl(hidden_states, input_ids, packed_seq_params) + + thd = packed_seq_params is not None and getattr(packed_seq_params, 'qkv_format', 'bshd') == 'thd' + + # ---- gather the shard into the full sequence (SP shard is innermost) ---- + if sp_on: + hidden_states = gather_from_sequence_parallel_region( + hidden_states, tensor_parallel_output_grad=False, group=self.pg_collection.tp) + if cp_on: + # For packed (thd) inputs the data pipeline (swift get_batch_on_this_cp_rank) + # zigzag-splits hidden/input_ids per sample via cu_seqlens_q while keeping + # packed_seq_params itself global, so the undo uses the same cu. + psp_for_cp = packed_seq_params if thd else None + hidden_states = reconstruct_tensor_cp(hidden_states, psp_for_cp, dim=0) + # The data pipeline may hand us either a CP-sharded or a full copy of + # input_ids; re-align only when the lengths disagree. + if input_ids.shape[-1] != hidden_states.shape[0]: + input_ids = reconstruct_tensor_cp(input_ids, psp_for_cp, dim=1) + elif input_ids.shape[-1] != hidden_states.shape[0]: + raise ValueError(f'PLE input_ids length {input_ids.shape[-1]} does not match the gathered ' + f'sequence length {hidden_states.shape[0]} under sequence parallelism; ' + 'gpt_model is expected to pass the full, unsharded input_ids.') + + if thd: + # SP keeps one global cu_seqlens copy per rank; normalize padded/offset + # forms against the gathered total so cu indexes the full sequence. + # copy.copy (not dataclasses.replace) preserves dynamically attached + # fields such as `num_samples` that the data pipeline relies on. + cu = self._normalize_cu_seqlens(getattr(packed_seq_params, 'cu_seqlens_q', None), hidden_states.shape[0]) + psp = copy.copy(packed_seq_params) + psp.cu_seqlens_q = cu + packed_seq_params = psp + + out = self._forward_impl(hidden_states, input_ids, packed_seq_params) + + # ---- scatter the additive output back to the caller's shard ---- + if cp_on: + out = split_cp_inputs(out, getattr(packed_seq_params, 'cu_seqlens_q', None) if thd else None, dim=0) + if sp_on: + out = scatter_to_sequence_parallel_region(out, group=self.pg_collection.tp) + return out + + def _forward_impl( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + packed_seq_params=None, + ) -> torch.Tensor: + """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 + # 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. + max_seqlen_q = packed_seq_params.max_seqlen_q + max_len = int(max_seqlen_q.item() if torch.is_tensor(max_seqlen_q) else max_seqlen_q) + cu = packed_seq_params.cu_seqlens_q + total = hidden_states.shape[0] + hid = hidden_states.new_zeros((num_samples, max_len, hidden_states.shape[-1])) + toks = input_ids.new_full((num_samples, max_len), self.ple_embedding.eos_token_id) + for i in range(num_samples): + start, end = int(cu[i]), int(cu[i + 1]) + hid[i, :end - start] = hidden_states[start:end, 0] + toks[i, :end - start] = input_ids[0, start:end] + res = self.compute(hid, toks) + out = res.new_zeros((total, 1, res.shape[-1])) + for i in range(num_samples): + start, end = int(cu[i]), int(cu[i + 1]) + out[start:end, 0] = res[i, :end - start] + return out + else: + # [s, b, nH] -> [b, s, nH]; input_ids [b, s] + hid = hidden_states.transpose(0, 1) + res = self.compute(hid, input_ids) + return res.transpose(0, 1).contiguous() diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py new file mode 100644 index 0000000..193bcbb --- /dev/null +++ b/src/mcore_bridge/model/modules/qsa_indexer.py @@ -0,0 +1,154 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import math +import torch +from torch import nn + + +class Qwen4ExpTextRMSNorm(nn.Module): + + def __init__(self, dim: int, eps: float = 1e-6, dtype=None, sequence_parallel: bool = False): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim, dtype=dtype)) + # Replicated across TP; reduce grads across TP when SP is on. + setattr(self.weight, 'sequence_parallel', sequence_parallel) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = x.float() + out = out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps) + # zero-centered: (1 + w), and Qwen4ExpText casts after scaling + return (out * (1.0 + self.weight.float())).type_as(x) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., :x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] + return torch.cat((-x2, x1), dim=-1) + + +class QSAIndexer(nn.Module): + # refer: transformers Qwen4ExpTextQSAIndexer + def __init__(self, config): + super().__init__() + self.config = config + self.index_n_heads = config.indexer_n_heads + self.index_kv_heads = config.indexer_kv_heads + self.index_head_dim = config.indexer_head_dim + self.compress_ratio = config.indexer_compress_ratio + self.token_budget = config.indexer_budget + self.block_topk = self.token_budget // self.compress_ratio + # Replicated projection (reference uses ReplicatedLinear). + self.index_qk_proj = nn.Linear( + config.hidden_size, (self.index_n_heads + self.index_kv_heads) * self.index_head_dim, + bias=False, + dtype=config.params_dtype) + self.q_layernorm = Qwen4ExpTextRMSNorm( + self.index_head_dim, + eps=config.layernorm_epsilon, + dtype=config.params_dtype, + sequence_parallel=config.sequence_parallel) + self.k_layernorm = Qwen4ExpTextRMSNorm( + self.index_head_dim, + eps=config.layernorm_epsilon, + dtype=config.params_dtype, + sequence_parallel=config.sequence_parallel) + setattr(self.index_qk_proj.weight, 'sequence_parallel', config.sequence_parallel) + + def forward(self, *args, **kwargs): + raise RuntimeError('QSAIndexer performs selection via `select_mask`, not `forward`.') + + @torch.no_grad() + def select_mask(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Return the QSA selection as a bool mask, or ``None`` when it is a no-op. + + Args: + hidden_states: ``[s, b, h]`` (mcore layout), pre-attention input -- + the same tensor the reference indexer consumes. + freqs: mcore rotary frequencies ``[s, 1, 1, rot_dim]``. mcore stores + angles rather than cos/sin, so they are materialized here the way + ``_patch_apply_rotary_pos_emb`` does (``cos(freqs) * mscale``), + keeping the indexer's RoPE identical to the attention's. + + Returns: + ``[b, 1, s, s]`` bool mask where True marks a *masked-out* key (TE's + ``arbitrary`` convention), or ``None`` if every visible key is + selected -- in which case the caller keeps the plain causal path and + pays nothing. + + Only the causal, unpacked layout is handled; callers must not invoke this + for packed/THD or context-parallel inputs (see the guard in the layer). + """ + s, b, _ = hidden_states.shape + R = self.compress_ratio + max_blocks = s // R + # Selection is a no-op while the causal prefix never exceeds the budget: + # `topk(min(block_topk, num_blocks))` then keeps every block and the tail + # re-adds the remainder, so the mask would be exactly causal. Skipping it + # keeps short sequences on TE's fused causal kernel. + if max_blocks <= self.block_topk: + return None + + device = hidden_states.device + # ---- project to indexer q/k ---- + # [s, b, h] -> [s, b, (nh + nkv) * d] + qk = self.index_qk_proj(hidden_states) + q, token_k = torch.split( + qk, [self.index_n_heads * self.index_head_dim, self.index_kv_heads * self.index_head_dim], dim=-1) + # -> [b, s, nh, d] / [b, s, d] + q = q.view(s, b, self.index_n_heads, self.index_head_dim).permute(1, 0, 2, 3) + raw_keys = token_k.view(s, b, self.index_kv_heads, self.index_head_dim).permute(1, 0, 2, 3).squeeze(2) + q = self.q_layernorm(q) + + # ---- materialize cos/sin from mcore freqs ---- + # freqs: [s, 1, 1, rot_dim] -> [s, rot_dim]; mscale mirrors the attention path. + mscale = getattr(self.config, 'attention_scaling', 1.0) or 1.0 + f = freqs.reshape(freqs.shape[0], -1)[:s] + cos = (torch.cos(f) * mscale).to(q.dtype) + sin = (torch.sin(f) * mscale).to(q.dtype) + rot = cos.shape[-1] + + def apply_rope(t, cos_, sin_): + t_rope, t_pass = t[..., :rot], t[..., rot:] + t_rope = (t_rope * cos_) + (_rotate_half(t_rope) * sin_) + return torch.cat((t_rope, t_pass), dim=-1) + + # queries rotate at their own position: cos [s, rot] -> [1, s, 1, rot] + q = apply_rope(q, cos[None, :, None, :], sin[None, :, None, :]) + + # ---- pool every block once (shared across queries) ---- + usable = max_blocks * R + key_groups = raw_keys[:, :usable].view(b, max_blocks, R, self.index_head_dim) + pooled = key_groups.float().mean(dim=2).to(raw_keys.dtype) + pooled = self.k_layernorm(pooled) + # blocks rotate at their first token's position + starts = torch.arange(max_blocks, device=device) * R + block_keys = apply_rope(pooled, cos[starts][None], sin[starts][None]) # [b, nb, d] + + # ---- score all (query, block) pairs ---- + scores = torch.einsum('bqhd,bkd->bqhk', q.float(), block_keys.float()) + scores = torch.relu(scores).sum(dim=2) / math.sqrt(self.index_head_dim) # [b, s, nb] + + # ---- restrict to blocks fully inside the causal prefix ---- + n_blocks = (torch.arange(s, device=device) + 1) // R # [s] + block_ids = torch.arange(max_blocks, device=device) + scores = scores.masked_fill((block_ids[None, :] >= n_blocks[:, None])[None], float('-inf')) + + # ---- top-k blocks -> token mask ---- + k = min(self.block_topk, max_blocks) + top_blocks = scores.topk(k, dim=-1).indices # [b, s, k] + keep = top_blocks < n_blocks[None, :, None] # drop the -inf padding slots + tok = (top_blocks.unsqueeze(-1) * R + torch.arange(R, device=device)).flatten(-2) + keep_tok = keep.unsqueeze(-1).expand(-1, -1, -1, R).flatten(-2) + + allowed = torch.zeros((b, s, s + 1), dtype=torch.bool, device=device) + allowed.scatter_(-1, torch.where(keep_tok, tok, torch.full_like(tok, s)).long(), True) + allowed = allowed[..., :s] + # tail: visible tokens after the last complete block are always attended + pos = torch.arange(s, device=device) + tail = (pos[None, :] >= (n_blocks * R)[:, None]) & (pos[None, :] <= pos[:, None]) + allowed |= tail[None] + # causal safety net (the block expansion never crosses it, but keep the + # invariant explicit so a future layout change fails loudly instead of + # silently attending to the future) + allowed &= (pos[None, :] <= pos[:, None])[None] + return ~allowed.unsqueeze(1) # True == masked out diff --git a/src/mcore_bridge/model/modules/transformer_block.py b/src/mcore_bridge/model/modules/transformer_block.py index 0beb20d..05a333f 100644 --- a/src/mcore_bridge/model/modules/transformer_block.py +++ b/src/mcore_bridge/model/modules/transformer_block.py @@ -358,11 +358,15 @@ def forward( # is called here to be future-proof and corner-case-proof. hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) enable_hyper_connections = getattr(self.config, 'enable_hyper_connections', False) + hc_count = getattr(self.config, 'hc_count', 0) or 0 + enable_gated_hc = hc_count > 1 and not enable_hyper_connections # Expand hidden states for hyper connections at the start of the block # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage if enable_hyper_connections and self.pre_process: hidden_states = HyperConnectionModule.input_expand(hidden_states, self.num_residual_streams) # [s, b, C] -> [s, b, n*C] + elif enable_gated_hc and self.pre_process: + hidden_states = hidden_states.repeat(1, 1, hc_count) # [s, b, C] -> [s, b, n*C] if self.config.sequence_parallel: rng_context = tensor_parallel.get_cuda_rng_tracker().fork() @@ -490,6 +494,11 @@ def forward( self.config.num_residual_streams, self.config.layernorm_epsilon, ) + elif enable_gated_hc and self.has_final_layernorm_in_this_stage(): + # Gated low-rank contraction (hyper_connection_mixer, use_combine=False + # so forward returns only the mixed stream). + # [s, b, n*C] -> [s, b, C] + hidden_states = self.hyper_connection_mixer(hidden_states) # Final layer norm. if self.final_layernorm is not None: @@ -515,4 +524,14 @@ def forward( if mhc_multistream is not None: return hidden_states, mhc_multistream + # A non-last PP stage hands its hidden_states (still multi-stream under + # mHC/gated-HC, since contraction happens on the last stage) straight to + # the next stage, and pipeline schedules' deallocate_output_tensor() + # asserts the tensor is not a view. The last-stage path guards this + # after the final layernorm (make_viewless_tensor above); the same + # guard is needed here because the last layer's output can itself be a + # view (e.g. TENorm). + if (enable_hyper_connections or enable_gated_hc) and not self.has_final_layernorm_in_this_stage(): + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + return hidden_states