diff --git a/src/mcore_bridge/model/hybrid_model.py b/src/mcore_bridge/model/hybrid_model.py index 9e6424d..75975c8 100644 --- a/src/mcore_bridge/model/hybrid_model.py +++ b/src/mcore_bridge/model/hybrid_model.py @@ -33,18 +33,15 @@ def __init__( post_process: bool = True, vp_stage: Optional[int] = None, ): - # `ModelLoader.build_model` passes the stack spec positionally as - # `transformer_layer_spec`; upstream names the same argument `hybrid_stack_spec`. - # MTP needs no separate spec here: upstream derives it from the `/` suffix of - # `hybrid_layer_pattern`. vocab_size = math.ceil( config.padded_vocab_size / config.tensor_model_parallel_size) * config.tensor_model_parallel_size + hybrid_layer_pattern = self._resolve_hybrid_layer_pattern(config) super().__init__( config, transformer_layer_spec, vocab_size, config.max_position_embeddings, - hybrid_layer_pattern=config.hybrid_layer_pattern, + hybrid_layer_pattern=hybrid_layer_pattern, pre_process=pre_process, post_process=post_process, share_embeddings_and_output_weights=not config.untie_embeddings_and_output_weights, @@ -53,6 +50,15 @@ def __init__( vp_stage=vp_stage, ) + @staticmethod + def _resolve_hybrid_layer_pattern(config: ModelConfig) -> Optional[str]: + pattern = config.hybrid_layer_pattern + mtp_pattern = getattr(config, 'mtp_hybrid_override_pattern', None) + mtp_num_layers = config.mtp_num_layers + if (pattern and mtp_pattern and mtp_num_layers and '/' not in pattern): + pattern = pattern + '/' + '/'.join([mtp_pattern] * mtp_num_layers) + return pattern + def _get_padding_mask(self, attention_mask) -> Optional[torch.Tensor]: """Mark fully-padded sequence positions, sharded to match the hidden states.""" if isinstance(attention_mask, dict):