diff --git a/.gitignore b/.gitignore
index 2d3acbd..19224f1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,7 @@
/.vscode
*.egg-info/
+/.idea/
+/.qoder/
build/
dist/
__pycache__/
diff --git a/README.md b/README.md
index 8c43082..233f8ed 100644
--- a/README.md
+++ b/README.md
@@ -86,6 +86,7 @@ The following is the list of models supported by MCore-Bridge:
| GLM | glm4, glm4_moe, glm4_moe_lite
glm4v, glm4v_moe,
glm_moe_dsa |
| MiniMax | minimax_m2 |
| Kimi | kimi_k2, kimi_vl, kimi_k25 |
+| Bailing | bailing_moe |
| InternLM | internlm3, internvl_chat, internvl |
| Ovis | ovis2_5 |
| Llama | llama, llama4 |
diff --git a/README_zh.md b/README_zh.md
index 00a0e2d..3518b73 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -85,6 +85,7 @@ uv pip install -e . --torch-backend=auto
| GLM | glm4, glm4_moe, glm4_moe_lite
glm4v, glm4v_moe,
glm_moe_dsa |
| MiniMax | minimax_m2 |
| Kimi | kimi_k2, kimi_vl, kimi_k25 |
+| Bailing | bailing_moe |
| InternLM | internlm3, internvl_chat, internvl |
| Ovis | ovis2_5 |
| Llama | llama, llama4 |
diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py
index b171f20..c380328 100644
--- a/src/mcore_bridge/bridge/gpt_bridge.py
+++ b/src/mcore_bridge/bridge/gpt_bridge.py
@@ -37,6 +37,8 @@ class GPTBridge:
# HF Keys
hf_q_norm_key = 'q_norm.weight'
hf_k_norm_key = 'k_norm.weight'
+ hf_o_proj_key = 'o_proj'
+ hf_attn_prefix = 'self_attn'
hf_mlp_prefix = 'mlp'
hf_gate_key = 'gate.weight'
hf_shared_expert_key = None
@@ -523,11 +525,7 @@ def _filter_prefix(state_dict, prefix: str):
return state_dict
return {k: v for k, v in state_dict.items() if k.startswith(prefix)}
- def _set_attn_state(self, mg_attn, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool):
- if to_mcore:
- hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix)
- else:
- hf_state_dict = {}
+ def _set_qkv(self, mg_attn, hf_state_dict, to_mcore: bool):
config = self.config
num_query_groups = (
config.num_query_groups if config.num_query_groups is not None else config.num_attention_heads)
@@ -618,9 +616,6 @@ def _set_attn_state(self, mg_attn, hf_state_dict, hf_prefix: str, layer_idx: int
hf_state_dict['v_proj.weight_scale_inv'] = scale_inv[:, -kv_block:, :].reshape(
-1, hidden_size_block).clone()
del mg_attn_weight
- self._set_state_dict(mg_attn, 'linear_proj.weight', hf_state_dict, 'o_proj.weight', to_mcore)
- if config.add_bias_linear:
- self._set_state_dict(mg_attn, 'linear_proj.bias', hf_state_dict, 'o_proj.bias', to_mcore)
# Copy bias
if (config.add_bias_linear or config.add_qkv_bias) and not self._peft_format:
@@ -640,6 +635,18 @@ def _set_attn_state(self, mg_attn, hf_state_dict, hf_prefix: str, layer_idx: int
hf_state_dict['q_proj.bias'] = mg_attn_bias[:, :q_dim].reshape(-1).clone()
hf_state_dict['k_proj.bias'] = mg_attn_bias[:, q_dim:-kv_dim].reshape(-1).clone()
hf_state_dict['v_proj.bias'] = mg_attn_bias[:, -kv_dim:].reshape(-1).clone()
+ return hf_state_dict
+
+ def _set_attn_state(self, mg_attn, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool):
+ if to_mcore:
+ hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix)
+ else:
+ hf_state_dict = {}
+ config = self.config
+ hf_state_dict.update(self._set_qkv(mg_attn, hf_state_dict, to_mcore))
+ self._set_state_dict(mg_attn, 'linear_proj.weight', hf_state_dict, f'{self.hf_o_proj_key}.weight', to_mcore)
+ if config.add_bias_linear:
+ self._set_state_dict(mg_attn, 'linear_proj.bias', hf_state_dict, f'{self.hf_o_proj_key}.bias', to_mcore)
if getattr(config, 'softmax_type', 'vanilla') == 'learnable':
self._set_state_dict(mg_attn, 'core_attention.softmax_offset', hf_state_dict, 'sinks', to_mcore)
if config.qk_layernorm:
@@ -1559,10 +1566,12 @@ def _set_mla_attn_state(
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
if self.config.multi_latent_attention:
- hf_state_dict.update(self._set_mla_attn_state(mg_attn, hf_state_dict, 'self_attn.', layer_idx, to_mcore))
+ hf_state_dict.update(
+ self._set_mla_attn_state(mg_attn, hf_state_dict, f'{self.hf_attn_prefix}.', layer_idx, to_mcore))
self._set_state_dict(mg_layer, 'input_layernorm.weight', hf_state_dict, 'input_layernorm.weight', to_mcore)
else:
- hf_state_dict.update(self._set_attn_state(mg_attn, hf_state_dict, 'self_attn.', layer_idx, to_mcore))
+ hf_state_dict.update(
+ self._set_attn_state(mg_attn, hf_state_dict, f'{self.hf_attn_prefix}.', layer_idx, to_mcore))
self._set_state_dict(mg_layer, 'self_attention.linear_qkv.layer_norm_weight', hf_state_dict,
'input_layernorm.weight', to_mcore)
return hf_state_dict
diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py
index c498c28..877b63c 100644
--- a/src/mcore_bridge/config/parser.py
+++ b/src/mcore_bridge/config/parser.py
@@ -30,10 +30,12 @@
'moe_router_group_topk': ['topk_group'],
'num_moe_experts': ['num_experts', 'n_routed_experts', 'moe_num_experts', 'num_local_experts'],
'moe_router_pre_softmax': ['norm_topk_prob'],
+ 'moe_router_enable_expert_bias': ['moe_router_enable_expert_bias'],
+ 'rotary_interleaved': ['rope_interleave'],
# deepseek
'q_lora_rank': ['q_lora_rank'],
'kv_lora_rank': ['kv_lora_rank'],
- 'moe_router_score_function': ['scoring_func', 'moe_router_use_sigmoid'],
+ 'moe_router_score_function': ['scoring_func', 'moe_router_use_sigmoid', 'score_function'],
'moe_router_bias_update_rate': ['aux_loss_alpha'],
'qk_head_dim': ['qk_nope_head_dim'],
'qk_pos_emb_head_dim': ['qk_rope_head_dim'],
diff --git a/src/mcore_bridge/model/constant.py b/src/mcore_bridge/model/constant.py
index 6c61f09..ffbeaba 100644
--- a/src/mcore_bridge/model/constant.py
+++ b/src/mcore_bridge/model/constant.py
@@ -8,6 +8,7 @@ class LLMModelType:
glm4 = 'glm4'
minimax_m2 = 'minimax_m2'
hy_v3 = 'hy_v3'
+ bailing_moe = 'bailing_moe'
qwen3_emb = 'qwen3_emb'
diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py
index b989f1e..52b007f 100644
--- a/src/mcore_bridge/model/gpts/__init__.py
+++ b/src/mcore_bridge/model/gpts/__init__.py
@@ -1,2 +1,2 @@
# Copyright (c) ModelScope Contributors. All rights reserved.
-from . import glm4, hunyuan, llm, minimax_m2, olmoe, qwen3_emb, qwen3_next
+from . import bailing_moe, glm4, hunyuan, llm, minimax_m2, olmoe, qwen3_emb, qwen3_next
diff --git a/src/mcore_bridge/model/gpts/bailing_moe.py b/src/mcore_bridge/model/gpts/bailing_moe.py
new file mode 100644
index 0000000..f3c383a
--- /dev/null
+++ b/src/mcore_bridge/model/gpts/bailing_moe.py
@@ -0,0 +1,86 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from megatron.core.transformer.attention import SelfAttention
+from torch import Tensor
+from typing import Optional
+
+from mcore_bridge.bridge import GPTBridge
+
+from ..constant import ModelType
+from ..register import ModelLoader, ModelMeta, register_model
+
+
+class BailingMoeSelfAttention(SelfAttention):
+
+ def get_query_key_value_tensors(
+ self,
+ hidden_states: Tensor,
+ key_value_states: Optional[Tensor] = None,
+ *args,
+ **kwargs,
+ ):
+ """Override to handle BailingMoE's non-interleaved QKV weight layout.
+
+ BailingMoE stores weights as [Q_all | K_all | V_all] (split by head count),
+ not Megatron's interleaved [q1 q2 k1 v1 | q3 q4 k2 v2 | ...].
+ """
+ # [sq, b, h] --> [sq, b, (num_heads + 2 * num_kv_heads) * head_dim]
+ mixed_qkv, _ = self.linear_qkv(hidden_states)
+
+ # [sq, b, (num_heads + 2 * num_kv_heads) * head_dim]
+ # --> [sq, b, num_heads + 2 * num_kv_heads, head_dim]
+ new_tensor_shape = mixed_qkv.size()[:-1] + (
+ self.num_attention_heads_per_partition + 2 * self.num_query_groups_per_partition,
+ self.hidden_size_per_attention_head,
+ )
+ mixed_qkv = mixed_qkv.view(*new_tensor_shape)
+
+ # Split by head count: [sq, b, num_heads, hn], [sq, b, num_kv_heads, hn], [sq, b, num_kv_heads, hn]
+ query, key, value = torch.split(
+ mixed_qkv,
+ [
+ self.num_attention_heads_per_partition, self.num_query_groups_per_partition,
+ self.num_query_groups_per_partition
+ ],
+ dim=2,
+ )
+
+ if self.q_layernorm is not None:
+ query = self.q_layernorm(query)
+
+ if self.k_layernorm is not None:
+ key = self.k_layernorm(key)
+
+ return query, key, value
+
+
+class BailingMoeLoader(ModelLoader):
+
+ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
+ transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
+ for layer_spec in transformer_layer_spec.layer_specs:
+ layer_spec.submodules.self_attention.module = BailingMoeSelfAttention
+ return transformer_layer_spec
+
+
+class BailingMoeBridge(GPTBridge):
+ hf_embed_key = 'model.word_embeddings.weight'
+ hf_attn_prefix = 'attention'
+ hf_q_norm_key = 'query_layernorm.weight'
+ hf_k_norm_key = 'key_layernorm.weight'
+ hf_expert_bias_key = 'gate.expert_bias'
+ hf_o_proj_key = 'dense'
+
+ def _set_qkv(self, mg_attn, hf_state_dict, to_mcore: bool):
+ self._set_state_dict(mg_attn, 'linear_qkv.weight', hf_state_dict, 'query_key_value.weight', to_mcore)
+ assert not self.config.add_bias_linear
+ return hf_state_dict
+
+
+register_model(
+ ModelMeta(
+ ModelType.bailing_moe,
+ ['bailing_moe'],
+ bridge_cls=BailingMoeBridge,
+ loader=BailingMoeLoader,
+ ))
diff --git a/tests/test_llm.py b/tests/test_llm.py
index 152df42..17f0b52 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -157,6 +157,10 @@ def test_olmoe():
_test_model('allenai/OLMoE-1B-7B-0125-Instruct')
+def test_bailing():
+ _test_model('inclusionAI/Ling-mini-2.0')
+
+
if __name__ == '__main__':
# test_qwen2()
# test_llama2()
@@ -169,7 +173,7 @@ def test_olmoe():
# test_megrez()
# test_llama3_1()
# test_llama3_2()
- test_qwen3()
+ # test_qwen3()
# test_qwen2_moe()
# test_qwen3_moe()
# test_internlm3()
@@ -190,3 +194,4 @@ def test_olmoe():
# test_minimax_m2()
# test_glm4_moe_lite()
# test_olmoe()
+ test_bailing()