From 8d7f3eb7ea006c0d7f288203ddd00142453762dd Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 19:54:22 +0800 Subject: [PATCH 1/9] support language_model_only --- requirements.txt | 2 +- src/mcore_bridge/bridge/gpt_bridge.py | 2 +- src/mcore_bridge/config/model_config.py | 3 ++- src/mcore_bridge/model/mm_gpt_model.py | 4 ++-- 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/requirements.txt b/requirements.txt index d947546..a664f53 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,4 @@ modelscope peft>=0.11,<0.20 safetensors tqdm -transformers>=4.33,<5.10.0 +transformers>=4.33,<5.11.0 diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index f3111ff..b77ee4c 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -1677,7 +1677,7 @@ def _convert_pre_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore else: hf_state_dict = {} self._set_word_embeddings(mg_model, hf_state_dict, to_mcore) - if self.is_multimodal: + if self.is_multimodal and not self.config.language_model_only: for prefix, mg_prefix in self.module_mapping.items(): mg_module = deep_getattr(mg_model, f'visual.{mg_prefix}') hf_state_dict.update(self._set_module(mg_module, hf_state_dict, f'{hf_prefix}{prefix}.', to_mcore)) diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index 79f0aee..257f16b 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/src/mcore_bridge/config/model_config.py @@ -222,6 +222,7 @@ class ModelConfig(TransformerConfig): mtp_shared_weights: bool = False # visual + language_model_only: bool = False hf_config: Optional[PretrainedConfig] = None vit_attn_impl: Optional[str] = None # e.g. 'flash_attention_2' @@ -343,7 +344,7 @@ def __post_init__(self): self.mcore_model_type = get_mcore_model_type(self.hf_model_type) self.model_meta = get_model_meta(self.mcore_model_type) self.is_multimodal = self.model_meta.visual_cls is not None - if self.is_multimodal: + if self.is_multimodal and not self.language_model_only: self.test_mm_type = getattr(self.model_meta.visual_cls, 'test_mm_type', 'image') else: self.test_mm_type = 'text' diff --git a/src/mcore_bridge/model/mm_gpt_model.py b/src/mcore_bridge/model/mm_gpt_model.py index 2a80c04..f2cb868 100644 --- a/src/mcore_bridge/model/mm_gpt_model.py +++ b/src/mcore_bridge/model/mm_gpt_model.py @@ -32,7 +32,7 @@ def __init__(self, self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights self.model_meta = config.model_meta self.visual = None - if pre_process and self.model_meta.visual_cls is not None: + if pre_process and self.model_meta.visual_cls is not None and not config.language_model_only: self.visual = self.model_meta.visual_cls(config) @contextmanager @@ -82,7 +82,7 @@ def forward( extra_kwargs = {k: kwargs[k] for k in self.language_model.extra_forward_keys} if decoder_input is not None: pass - elif self.pre_process: + elif self.pre_process and not self.config.language_model_only: kwargs.update({'input_ids': input_ids, 'packed_seq_params': packed_seq_params}) with self._patch_word_embeddings(kwargs): decoder_input = self.language_model.embedding(input_ids=input_ids, position_ids=position_ids) From 81fa46a812556a91e56a0d01454805fdec1916b1 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 23:04:53 +0800 Subject: [PATCH 2/9] update --- src/mcore_bridge/model/mm_gpt_model.py | 19 +++++++----- src/mcore_bridge/model/mm_gpts/gemma4.py | 37 +++++++++++++++++------- src/mcore_bridge/model/mm_gpts/utils.py | 12 +++++++- 3 files changed, 49 insertions(+), 19 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpt_model.py b/src/mcore_bridge/model/mm_gpt_model.py index f2cb868..eb274b8 100644 --- a/src/mcore_bridge/model/mm_gpt_model.py +++ b/src/mcore_bridge/model/mm_gpt_model.py @@ -32,7 +32,7 @@ def __init__(self, self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights self.model_meta = config.model_meta self.visual = None - if pre_process and self.model_meta.visual_cls is not None and not config.language_model_only: + if pre_process and self.model_meta.visual_cls is not None: self.visual = self.model_meta.visual_cls(config) @contextmanager @@ -48,12 +48,15 @@ def forward(_self, input_): packed_seq_params = kwargs.get('packed_seq_params') if self.visual is not None: res = self.visual.get_inputs_embeds(res, **kwargs) - kwargs.clear() - if isinstance(res, dict): - # compat dict - inputs_embeds = res.pop('inputs_embeds') - kwargs.update(res) - res = inputs_embeds + else: + assert self.config.language_model_only + res = self.visual.get_inputs_embeds_language_model(res, **kwargs) + kwargs.clear() + if isinstance(res, dict): + # compat dict + inputs_embeds = res.pop('inputs_embeds') + kwargs.update(res) + res = inputs_embeds if self.config.context_parallel_size > 1: res = split_cp_inputs(res, getattr(packed_seq_params, 'cu_seqlens_q', None), 1) if reduce_scatter_embeddings: @@ -82,7 +85,7 @@ def forward( extra_kwargs = {k: kwargs[k] for k in self.language_model.extra_forward_keys} if decoder_input is not None: pass - elif self.pre_process and not self.config.language_model_only: + elif self.pre_process: kwargs.update({'input_ids': input_ids, 'packed_seq_params': packed_seq_params}) with self._patch_word_embeddings(kwargs): decoder_input = self.language_model.embedding(input_ids=input_ids, position_ids=position_ids) diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index 5c6bd91..4d8c40e 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -74,14 +74,34 @@ def prepare_model(self, hf_config: PretrainedConfig): self.embed_audio = ( Gemma4MultimodalEmbedder(hf_config.audio_config, hf_config.text_config).to(dtype) if hf_config.audio_config is not None else None) - self.register_buffer('embed_scale', torch.tensor(hf_config.hidden_size**0.5).to(dtype), persistent=False) + self.prepare_language_model(hf_config) self.model_cls = Gemma4Model - def get_inputs_embeds(self, inputs_embeds, **kwargs): + def prepare_language_model(self, hf_config: PretrainedConfig): + self.register_buffer( + 'embed_scale', torch.tensor(hf_config.hidden_size**0.5).to(hf_config.torch_dtype), persistent=False) + + def get_inputs_embeds_language_model(self, inputs_embeds, **kwargs): input_ids = kwargs.get('input_ids') + hf_config = self.hf_config inputs_embeds = inputs_embeds * self.embed_scale.to(inputs_embeds.dtype) + image_mask = input_ids == hf_config.image_token_id + video_mask = input_ids == hf_config.video_token_id + audio_mask = input_ids == hf_config.audio_token_id + multimodal_mask = image_mask | video_mask | audio_mask + llm_input_ids = input_ids.clone() + llm_input_ids[multimodal_mask] = hf_config.text_config.pad_token_id + return {'inputs_embeds': inputs_embeds, 'llm_input_ids': llm_input_ids} + + def get_inputs_embeds(self, inputs_embeds, **kwargs): hf_config = self.hf_config + inputs_embeds = get_inputs_embeds_language_model(inputs_embeds, **kwargs) + input_ids = kwargs.get('input_ids') + image_mask = input_ids == hf_config.image_token_id + video_mask = input_ids == hf_config.video_token_id + audio_mask = input_ids == hf_config.audio_token_id + pixel_values = kwargs.get('pixel_values') pixel_values_videos = kwargs.get('pixel_values_videos') input_features = kwargs.get('input_features') @@ -89,12 +109,8 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs): image_position_ids = kwargs.get('image_position_ids') video_position_ids = kwargs.get('video_position_ids') - image_mask = input_ids == hf_config.image_token_id - video_mask = input_ids == hf_config.video_token_id - audio_mask = input_ids == hf_config.audio_token_id - multimodal_mask = image_mask | video_mask | audio_mask - llm_input_ids = input_ids.clone() - llm_input_ids[multimodal_mask] = hf_config.text_config.pad_token_id + res = self.get_inputs_embeds_language_model(inputs_embeds, **kwargs) + inputs_embeds = res['inputs_embeds'] if pixel_values is not None: with self.patch_hf_config(): @@ -121,7 +137,8 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs): audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype) audio_mask_e = audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) inputs_embeds = inputs_embeds.masked_scatter(audio_mask_e, audio_features) - return {'inputs_embeds': inputs_embeds, 'llm_input_ids': llm_input_ids} + res['inputs_embeds'] = inputs_embeds + return res class Gemma4SelfAttention(SelfAttention): @@ -878,7 +895,7 @@ def prepare_model(self, hf_config: PretrainedConfig): self.embed_audio = ( Gemma4UnifiedMultimodalEmbedder(hf_config.audio_config, hf_config.text_config).to(dtype) if hf_config.audio_config is not None else None) - self.register_buffer('embed_scale', torch.tensor(hf_config.hidden_size**0.5).to(dtype), persistent=False) + self.prepare_language_model(hf_config) self.model_cls = Gemma4UnifiedModel diff --git a/src/mcore_bridge/model/mm_gpts/utils.py b/src/mcore_bridge/model/mm_gpts/utils.py index f333ced..b610500 100644 --- a/src/mcore_bridge/model/mm_gpts/utils.py +++ b/src/mcore_bridge/model/mm_gpts/utils.py @@ -52,13 +52,20 @@ def __init__(self, config: ModelConfig): self.hf_config = hf_config self.prepare_attn_impl() with patch_get_dynamic_module(): - self.prepare_model(hf_config) + if config.language_model_only: + self.prepare_language_model(hf_config) + else: + self.prepare_model(hf_config) + self.to(device='cuda') @abstractmethod def prepare_model(self, hf_config: PretrainedConfig): pass + def prepare_language_model(self, hf_config: PretrainedConfig): + pass + def prepare_attn_impl(self): vit_attn_impl = self.config.vit_attn_impl or 'flash_attention_2' if self.config.attention_backend.name == 'flash': @@ -68,6 +75,9 @@ def prepare_attn_impl(self): def get_inputs_embeds(self, inputs_embeds, **kwargs): pass + def get_inputs_embeds_language_model(self, inputs_embeds, **kwargs): + return inputs_embeds + @staticmethod def _get_vision_config(hf_config): for k in ['vision_config', 'vit_config']: From 2d84379bdd1f5f772008595b6b00fb463567b599 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 23:06:28 +0800 Subject: [PATCH 3/9] fix --- src/mcore_bridge/model/mm_gpts/gemma4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index 4d8c40e..c4b7c0b 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -96,7 +96,7 @@ def get_inputs_embeds_language_model(self, inputs_embeds, **kwargs): def get_inputs_embeds(self, inputs_embeds, **kwargs): hf_config = self.hf_config - inputs_embeds = get_inputs_embeds_language_model(inputs_embeds, **kwargs) + inputs_embeds = self.get_inputs_embeds_language_model(inputs_embeds, **kwargs) input_ids = kwargs.get('input_ids') image_mask = input_ids == hf_config.image_token_id video_mask = input_ids == hf_config.video_token_id From 8df2d30d05b029734474d4278550b9ec4c4c7063 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 23:09:30 +0800 Subject: [PATCH 4/9] fix --- src/mcore_bridge/model/mm_gpts/gemma4.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index c4b7c0b..0b65390 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -96,7 +96,7 @@ def get_inputs_embeds_language_model(self, inputs_embeds, **kwargs): def get_inputs_embeds(self, inputs_embeds, **kwargs): hf_config = self.hf_config - inputs_embeds = self.get_inputs_embeds_language_model(inputs_embeds, **kwargs) + res = self.get_inputs_embeds_language_model(inputs_embeds, **kwargs) input_ids = kwargs.get('input_ids') image_mask = input_ids == hf_config.image_token_id video_mask = input_ids == hf_config.video_token_id @@ -109,7 +109,6 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs): image_position_ids = kwargs.get('image_position_ids') video_position_ids = kwargs.get('video_position_ids') - res = self.get_inputs_embeds_language_model(inputs_embeds, **kwargs) inputs_embeds = res['inputs_embeds'] if pixel_values is not None: From 775697a83d5053bcded084cea7372f4ee1e2fc21 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 23:19:24 +0800 Subject: [PATCH 5/9] fix --- src/mcore_bridge/model/mm_gpt_model.py | 7 +++---- src/mcore_bridge/model/mm_gpts/gemma4.py | 12 ++++++++---- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpt_model.py b/src/mcore_bridge/model/mm_gpt_model.py index eb274b8..9971a3c 100644 --- a/src/mcore_bridge/model/mm_gpt_model.py +++ b/src/mcore_bridge/model/mm_gpt_model.py @@ -46,11 +46,10 @@ def forward(_self, input_): res = origin_forward(_self, input_) _self.reduce_scatter_embeddings = reduce_scatter_embeddings packed_seq_params = kwargs.get('packed_seq_params') - if self.visual is not None: - res = self.visual.get_inputs_embeds(res, **kwargs) - else: - assert self.config.language_model_only + if self.config.language_model_only: res = self.visual.get_inputs_embeds_language_model(res, **kwargs) + else: + res = self.visual.get_inputs_embeds(res, **kwargs) kwargs.clear() if isinstance(res, dict): # compat dict diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index 0b65390..0b0fe3c 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -97,10 +97,6 @@ def get_inputs_embeds_language_model(self, inputs_embeds, **kwargs): def get_inputs_embeds(self, inputs_embeds, **kwargs): hf_config = self.hf_config res = self.get_inputs_embeds_language_model(inputs_embeds, **kwargs) - input_ids = kwargs.get('input_ids') - image_mask = input_ids == hf_config.image_token_id - video_mask = input_ids == hf_config.video_token_id - audio_mask = input_ids == hf_config.audio_token_id pixel_values = kwargs.get('pixel_values') pixel_values_videos = kwargs.get('pixel_values_videos') @@ -109,6 +105,14 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs): image_position_ids = kwargs.get('image_position_ids') video_position_ids = kwargs.get('video_position_ids') + input_ids = kwargs.get('input_ids') + image_mask = input_ids == hf_config.image_token_id + video_mask = input_ids == hf_config.video_token_id + audio_mask = input_ids == hf_config.audio_token_id + multimodal_mask = image_mask | video_mask | audio_mask + llm_input_ids = input_ids.clone() + llm_input_ids[multimodal_mask] = hf_config.text_config.pad_token_id + inputs_embeds = res['inputs_embeds'] if pixel_values is not None: From 50eb37a78d734db7c1e73f3da6478e8737811086 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 23:20:30 +0800 Subject: [PATCH 6/9] fix --- src/mcore_bridge/model/mm_gpts/gemma4.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index 0b0fe3c..fe4d8e5 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -82,17 +82,8 @@ def prepare_language_model(self, hf_config: PretrainedConfig): 'embed_scale', torch.tensor(hf_config.hidden_size**0.5).to(hf_config.torch_dtype), persistent=False) def get_inputs_embeds_language_model(self, inputs_embeds, **kwargs): - input_ids = kwargs.get('input_ids') - hf_config = self.hf_config inputs_embeds = inputs_embeds * self.embed_scale.to(inputs_embeds.dtype) - - image_mask = input_ids == hf_config.image_token_id - video_mask = input_ids == hf_config.video_token_id - audio_mask = input_ids == hf_config.audio_token_id - multimodal_mask = image_mask | video_mask | audio_mask - llm_input_ids = input_ids.clone() - llm_input_ids[multimodal_mask] = hf_config.text_config.pad_token_id - return {'inputs_embeds': inputs_embeds, 'llm_input_ids': llm_input_ids} + return {'inputs_embeds': inputs_embeds, 'llm_input_ids': kwargs.get('input_ids')} def get_inputs_embeds(self, inputs_embeds, **kwargs): hf_config = self.hf_config From 0834563ad7d966bb68e81f6b790f195d678c39f5 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 23:21:51 +0800 Subject: [PATCH 7/9] fix --- src/mcore_bridge/model/mm_gpts/gemma4.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index fe4d8e5..08b8eb6 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -132,6 +132,7 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs): audio_mask_e = audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) inputs_embeds = inputs_embeds.masked_scatter(audio_mask_e, audio_features) res['inputs_embeds'] = inputs_embeds + res['llm_input_ids'] = input_ids return res From 72352fe0e0d89a27ae9160fd02470681e82e8308 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 23:32:45 +0800 Subject: [PATCH 8/9] fix --- src/mcore_bridge/model/mm_gpt_model.py | 9 +++++---- src/mcore_bridge/model/mm_gpts/gemma4.py | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpt_model.py b/src/mcore_bridge/model/mm_gpt_model.py index 9971a3c..309a898 100644 --- a/src/mcore_bridge/model/mm_gpt_model.py +++ b/src/mcore_bridge/model/mm_gpt_model.py @@ -46,10 +46,11 @@ def forward(_self, input_): res = origin_forward(_self, input_) _self.reduce_scatter_embeddings = reduce_scatter_embeddings packed_seq_params = kwargs.get('packed_seq_params') - if self.config.language_model_only: - res = self.visual.get_inputs_embeds_language_model(res, **kwargs) - else: - res = self.visual.get_inputs_embeds(res, **kwargs) + if self.visual is not None: + if self.config.language_model_only: + res = self.visual.get_inputs_embeds_language_model(res, **kwargs) + else: + res = self.visual.get_inputs_embeds(res, **kwargs) kwargs.clear() if isinstance(res, dict): # compat dict diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index 08b8eb6..94b6e12 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -132,7 +132,7 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs): audio_mask_e = audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) inputs_embeds = inputs_embeds.masked_scatter(audio_mask_e, audio_features) res['inputs_embeds'] = inputs_embeds - res['llm_input_ids'] = input_ids + res['llm_input_ids'] = llm_input_ids return res From fd6a2f0d5a28c32157eb8bc0e1f68e697f304ea0 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 4 Jun 2026 23:33:35 +0800 Subject: [PATCH 9/9] fix --- src/mcore_bridge/model/mm_gpt_model.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpt_model.py b/src/mcore_bridge/model/mm_gpt_model.py index 309a898..b535bfb 100644 --- a/src/mcore_bridge/model/mm_gpt_model.py +++ b/src/mcore_bridge/model/mm_gpt_model.py @@ -51,12 +51,12 @@ def forward(_self, input_): res = self.visual.get_inputs_embeds_language_model(res, **kwargs) else: res = self.visual.get_inputs_embeds(res, **kwargs) - kwargs.clear() - if isinstance(res, dict): - # compat dict - inputs_embeds = res.pop('inputs_embeds') - kwargs.update(res) - res = inputs_embeds + kwargs.clear() + if isinstance(res, dict): + # compat dict + inputs_embeds = res.pop('inputs_embeds') + kwargs.update(res) + res = inputs_embeds if self.config.context_parallel_size > 1: res = split_cp_inputs(res, getattr(packed_seq_params, 'cu_seqlens_q', None), 1) if reduce_scatter_embeddings: