diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml
new file mode 100644
index 0000000..3c3778f
--- /dev/null
+++ b/.github/workflows/lint.yaml
@@ -0,0 +1,22 @@
+name: Lint test
+
+on: [push, pull_request]
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.ref }}
+ cancel-in-progress: true
+
+jobs:
+ lint:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v2
+ - name: Set up Python 3.12
+ uses: actions/setup-python@v2
+ with:
+ python-version: '3.12'
+ - name: Install pre-commit hook
+ run: |
+ pip install pre-commit
+ - name: Linting
+ run: pre-commit run --all-files
diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml
new file mode 100644
index 0000000..7b60e49
--- /dev/null
+++ b/.github/workflows/publish.yaml
@@ -0,0 +1,29 @@
+name: release
+
+on:
+ push:
+ tags:
+ - 'v**'
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.ref }}-publish
+ cancel-in-progress: true
+
+jobs:
+ build-n-publish:
+ runs-on: ubuntu-22.04
+ #if: startsWith(github.event.ref, 'refs/tags')
+ steps:
+ - uses: actions/checkout@v2
+ - name: Set up Python 3.12
+ uses: actions/setup-python@v2
+ with:
+ python-version: '3.12'
+ - name: Install wheel
+ run: pip install wheel packaging setuptools==69.5.1
+ - name: Build Mcore Bridge
+ run: python setup.py sdist bdist_wheel
+ - name: Publish package to PyPI
+ run: |
+ pip install twine
+ twine upload dist/* --skip-existing -u __token__ -p ${{ secrets.PYPI_API_TOKEN }}
diff --git a/README.md b/README.md
index 7dab6a7..85bd378 100644
--- a/README.md
+++ b/README.md
@@ -39,7 +39,7 @@
- [News](#-news)
- [Installation](#%EF%B8%8F-installation)
- [Quick Start](#-quick-Start)
-- [Usage](#-Usage)
+- [Model List](#-Model-List)
- [License](#-License)
@@ -47,14 +47,14 @@
You can contact us and communicate with us by adding our group:
-WeChat Group |
-:-------------------------:
-
+| WeChat Group |
+|:-------------------------:|
+|
|
## 📝 Introduction
## 🎉 News
-- 🎉 2025.03.23: MCore-Bridge is released! Making Megatron training as simple as Transformers, providing Megatron-Core model definitions for state-of-the-art large language models.
+- 🎉 2025.04.01: MCore-Bridge is released! Providing Megatron-Core model definitions for state-of-the-art large language models and making Megatron training as simple as Transformers.
## 🛠️ Installation
To install using pip:
@@ -85,6 +85,7 @@ You need to create the following file (test.py), then run `CUDA_VISIBLE_DEVICES=
The saved model can be used for inference by referring to the [example code in the model card](https://modelscope.cn/models/Qwen/Qwen3.5-35B-A3B).
```python
+# test env: transformers==5.2.0 megatron-core==0.16.1
import os
import torch
import torch.distributed as dist
@@ -93,6 +94,7 @@ from modelscope import snapshot_download
from transformers import AutoConfig, AutoProcessor
from mcore_bridge import ModelConfig, get_mcore_model, hf_to_mcore_config
+is_rank0 = int(os.getenv('RANK')) == 0
torch.cuda.set_device(f"cuda:{os.getenv('LOCAL_RANK')}")
dist.init_process_group(backend='nccl')
TP, PP, EP, ETP = 2, 2, 2, 1
@@ -115,8 +117,6 @@ config = ModelConfig(
expert_tensor_parallel_size=ETP,
sequence_parallel=True,
mtp_num_layers=1,
- processor=processor,
- hf_config=hf_config,
**config_kwargs)
# Create model
@@ -133,12 +133,110 @@ for name, parameter in bridge.export_weights(mg_models):
# Save weights
output_dir = 'Qwen3.5-35B-A3B-HF'
bridge.save_weights(mg_models, output_dir)
-processor.save_pretrained(output_dir)
-hf_config.save_pretrained(output_dir)
+if is_rank0:
+ processor.save_pretrained(output_dir)
+ hf_config.save_pretrained(output_dir)
```
-## ✨ Usage
+### Using Peft
+Mcore-Bridge is fully compatible with [Peft](https://github.com/huggingface/peft) for LoRA training. The following introduces how to use Peft to prepare a PeftModel and save the incremental weights.
+
+You need to create the following file (test.py), then run `CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 test.py`.
+
+```python
+import copy
+import os
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from modelscope import snapshot_download
+from peft import LoraConfig, get_peft_model
+from transformers import AutoConfig, AutoProcessor
+
+from mcore_bridge import ModelConfig, get_mcore_model, hf_to_mcore_config, set_random_seed
+
+is_rank0 = int(os.getenv('RANK')) == 0
+torch.cuda.set_device(f"cuda:{os.getenv('LOCAL_RANK')}")
+dist.init_process_group(backend='nccl')
+TP, PP = 2, 2
+mpu.initialize_model_parallel(
+ tensor_model_parallel_size=TP,
+ pipeline_model_parallel_size=PP,
+)
+# To correctly initialize the model randomly (full parameters/LoRA)
+# you need to set the random seed.
+set_random_seed(42)
+
+model_dir = snapshot_download('Qwen/Qwen3.5-4B')
+hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
+config_kwargs = hf_to_mcore_config(hf_config)
+config = ModelConfig(
+ params_dtype=torch.bfloat16,
+ tensor_model_parallel_size=TP,
+ pipeline_model_parallel_size=PP,
+ sequence_parallel=True,
+ **config_kwargs)
+
+# Create model and load weights
+mg_models = get_mcore_model(config)
+bridge = config.bridge
+bridge.load_weights(mg_models, model_dir)
+
+# Prepare PeftModel and load LoRA weights
+# For multimodal models, it is recommended to use regex to specify target_modules
+target_modules = r'^language_model.*\.(in_proj|out_proj|linear_fc1|linear_fc2|linear_qkv|linear_proj)$'
+# When saving as safetensors, you need to store the corresponding HF target_modules
+hf_target_modules = r'^model.language_model.*\.(in_proj_qkv|in_proj_z|in_proj_b|in_proj_a|out_proj|gate_proj|up_proj|down_proj|q_proj|k_proj|v_proj|o_proj)$'
+lora_config = LoraConfig(task_type='CAUSAL_LM', r=8, lora_alpha=32, lora_dropout=0.05, target_modules=target_modules)
+peft_models = [get_peft_model(model, lora_config) for model in mg_models]
+# Optional
+# bridge.load_weights(peft_models, model_dir, peft_format=True)
+
+# Export LoRA weights
+for name, parameter in bridge.export_weights(mg_models, peft_format=True):
+ pass
+
+# Save LoRA weights
+output_dir = 'Qwen3.5-4B-LoRA'
+bridge.save_weights(mg_models, output_dir, peft_format=True)
+if is_rank0:
+ hf_lora_config = copy.copy(lora_config)
+ hf_lora_config.target_modules = hf_target_modules
+ hf_lora_config.save_pretrained(output_dir)
+```
+
+Using the saved LoRA weights:
+
+```python
+from transformers import Qwen3_5ForConditionalGeneration
+from modelscope import snapshot_download
+from peft import PeftModel
+
+model_dir = snapshot_download('Qwen/Qwen3.5-4B')
+model = Qwen3_5ForConditionalGeneration.from_pretrained(model_dir)
+peft_model = PeftModel.from_pretrained(model, 'Qwen3.5-4B-LoRA')
+```
+
+## ✨ Model List
+
+The following is the list of models supported by MCore-Bridge:
+
+| Series | model_type |
+| -------- | ------------------------------------------------------------ |
+| Qwen | qwen2, qwen2_moe
qwen2_vl, qwen2_5_vl, qwen2_5_omni
qwen3, qwen3_moe
qwen3_vl, qwen3_vl_moe, qwen3_omni_moe
qwen3_next, qwen3_5, qwen3_5_moe |
+| DeepSeek | deepseek_v3, deepseek_v32 |
+| GLM | glm4, glm4_moe, glm4_moe_lite
glm4v, glm4v_moe,
glm_moe_dsa |
+| MiniMax | minimax_m2 |
+| Kimi | kimi_k2, kimi_vl |
+| InternLM | internlm3, internvl_chat, internvl |
+| Ovis | ovis2_5 |
+| Llama | llama, llama4 |
+| GPT-OSS | gpt_oss |
+| ERNIE | ernie4_5, ernie4_5_moe |
+| MiMo | mimo |
+| Dots | dots1 |
+| OLMoE | olmoe |
## 🏛 License
diff --git a/README_zh.md b/README_zh.md
index 4f84d3b..ebb812a 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -39,21 +39,21 @@
- [新闻](#-新闻)
- [安装](#%EF%B8%8F-安装)
- [快速开始](#-快速开始)
-- [如何使用](#-如何使用)
+- [模型列表](#-模型列表)
- [License](#-license)
## ☎ 用户群
请扫描下面的二维码来加入我们的交流群:
-微信群 |
-:-------------------------:
-
+| 微信群 |
+|:-------------------------:|
+|
|
## 📝 简介
## 🎉 新闻
-- 🎁 2026.03.23: Mcore-Bridge发布,让Megatron训练像transformers一样简单,为最先进的大语言模型提供 Megatron-Core 模型定义。
+- 🎉 2026.04.01: MCore-Bridge 正式发布!为最先进的大语言模型提供 Megatron-Core 模型定义,让 Megatron 训练像 Transformers 一样简单。
## 🛠️ 安装
使用pip进行安装:
@@ -84,6 +84,7 @@ uv pip install -e . --torch-backend=auto
保存的模型,可以参考[模型卡片的示例代码](https://modelscope.cn/models/Qwen/Qwen3.5-35B-A3B)进行推理。
```python
+# test env: transformers==5.2.0 megatron-core==0.16.1
import os
import torch
import torch.distributed as dist
@@ -92,6 +93,7 @@ from modelscope import snapshot_download
from transformers import AutoConfig, AutoProcessor
from mcore_bridge import ModelConfig, get_mcore_model, hf_to_mcore_config
+is_rank0 = int(os.getenv('RANK')) == 0
torch.cuda.set_device(f"cuda:{os.getenv('LOCAL_RANK')}")
dist.init_process_group(backend='nccl')
TP, PP, EP, ETP = 2, 2, 2, 1
@@ -114,8 +116,6 @@ config = ModelConfig(
expert_tensor_parallel_size=ETP,
sequence_parallel=True,
mtp_num_layers=1,
- processor=processor,
- hf_config=hf_config,
**config_kwargs)
# 创建模型
@@ -132,11 +132,108 @@ for name, parameter in bridge.export_weights(mg_models):
# 保存权重
output_dir = 'Qwen3.5-35B-A3B-HF'
bridge.save_weights(mg_models, output_dir)
-processor.save_pretrained(output_dir)
-hf_config.save_pretrained(output_dir)
+if is_rank0:
+ processor.save_pretrained(output_dir)
+ hf_config.save_pretrained(output_dir)
```
-## ✨ 如何使用
+### 使用Peft
+
+Mcore-Bridge完全兼容使用[Peft](https://github.com/huggingface/peft)进行LoRA训练。以下介绍如何使用peft准备PeftModel,并保存增量权重。
+
+你需要创建以下文件(test.py),然后运行`CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 test.py`。
+
+```python
+import copy
+import os
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from modelscope import snapshot_download
+from peft import LoraConfig, get_peft_model
+from transformers import AutoConfig, AutoProcessor
+
+from mcore_bridge import ModelConfig, get_mcore_model, hf_to_mcore_config, set_random_seed
+
+is_rank0 = int(os.getenv('RANK')) == 0
+torch.cuda.set_device(f"cuda:{os.getenv('LOCAL_RANK')}")
+dist.init_process_group(backend='nccl')
+TP, PP = 2, 2
+mpu.initialize_model_parallel(
+ tensor_model_parallel_size=TP,
+ pipeline_model_parallel_size=PP,
+)
+# 为了正确随机初始化模型(全参数/LoRA),你需要设置随机种子
+set_random_seed(42)
+
+model_dir = snapshot_download('Qwen/Qwen3.5-4B')
+hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
+config_kwargs = hf_to_mcore_config(hf_config)
+config = ModelConfig(
+ params_dtype=torch.bfloat16,
+ tensor_model_parallel_size=TP,
+ pipeline_model_parallel_size=PP,
+ sequence_parallel=True,
+ **config_kwargs)
+
+# 创建模型并加载权重
+mg_models = get_mcore_model(config)
+bridge = config.bridge
+bridge.load_weights(mg_models, model_dir)
+
+# 准备PeftModel并加载LoRA权重
+# 多模态模型建议使用正则表达式指定target_modules
+target_modules = r'^language_model.*\.(in_proj|out_proj|linear_fc1|linear_fc2|linear_qkv|linear_proj)$'
+# 存储成safetensors时,需要存储hf对应的target_modules
+hf_target_modules = r'^model.language_model.*\.(in_proj_qkv|in_proj_z|in_proj_b|in_proj_a|out_proj|gate_proj|up_proj|down_proj|q_proj|k_proj|v_proj|o_proj)$'
+lora_config = LoraConfig(task_type='CAUSAL_LM', r=8, lora_alpha=32, lora_dropout=0.05, target_modules=target_modules)
+peft_models = [get_peft_model(model, lora_config) for model in mg_models]
+# 可选
+# bridge.load_weights(peft_models, model_dir, peft_format=True)
+
+# 导出LoRA权重
+for name, parameter in bridge.export_weights(mg_models, peft_format=True):
+ pass
+
+# 保存LoRA权重
+output_dir = 'Qwen3.5-4B-LoRA'
+bridge.save_weights(mg_models, output_dir, peft_format=True)
+if is_rank0:
+ hf_lora_config = copy.copy(lora_config)
+ hf_lora_config.target_modules = hf_target_modules
+ hf_lora_config.save_pretrained(output_dir)
+```
+
+使用存储下来的LoRA权重:
+```python
+from transformers import Qwen3_5ForConditionalGeneration
+from modelscope import snapshot_download
+from peft import PeftModel
+
+model_dir = snapshot_download('Qwen/Qwen3.5-4B')
+model = Qwen3_5ForConditionalGeneration.from_pretrained(model_dir)
+peft_model = PeftModel.from_pretrained(model, 'Qwen3.5-4B-LoRA')
+```
+
+## ✨ 模型列表
+
+以下为MCore-Bridge支持的模型列表:
+
+| 系列 | model_type |
+| -------- | ------------------------------------------------------------ |
+| Qwen | qwen2, qwen2_moe
qwen2_vl, qwen2_5_vl, qwen2_5_omni
qwen3, qwen3_moe
qwen3_vl, qwen3_vl_moe, qwen3_omni_moe
qwen3_next, qwen3_5, qwen3_5_moe |
+| DeepSeek | deepseek_v3, deepseek_v32 |
+| GLM | glm4, glm4_moe, glm4_moe_lite
glm4v, glm4v_moe,
glm_moe_dsa |
+| MiniMax | minimax_m2 |
+| Kimi | kimi_k2, kimi_vl |
+| InternLM | internlm3, internvl_chat, internvl |
+| Ovis | ovis2_5 |
+| Llama | llama, llama4 |
+| GPT-OSS | gpt_oss |
+| ERNIE | ernie4_5, ernie4_5_moe |
+| MiMo | mimo |
+| Dots | dots1 |
+| OLMoE | olmoe |
## 🏛 License
diff --git a/setup.py b/setup.py
index 8de9e36..09af602 100644
--- a/setup.py
+++ b/setup.py
@@ -43,7 +43,7 @@ def parse_requirements(path='requirements.txt'):
keywords=['transformers', 'LLM', 'lora', 'megatron', 'peft'],
url='https://github.com/modelscope/mcore-bridge',
package_dir={'': 'src'},
- packages=find_packages(include='src'),
+ packages=find_packages('src'),
python_requires='>=3.8.0',
classifiers=[
'Development Status :: 4 - Beta',
diff --git a/src/mcore_bridge/__init__.py b/src/mcore_bridge/__init__.py
index 8f9cade..5fb7d8c 100644
--- a/src/mcore_bridge/__init__.py
+++ b/src/mcore_bridge/__init__.py
@@ -11,13 +11,16 @@
from .bridge import GPTBridge
from .config import ModelConfig, hf_to_mcore_config
from .model import get_mcore_model
- from .utils import get_logger
+ from .tuners import LoraParallelLinear
+ from .utils import get_logger, set_random_seed
from .version import __release_datetime__, __version__
else:
_import_structure = {
'bridge': ['GPTBridge'],
'config': ['ModelConfig', 'hf_to_mcore_config'],
'model': ['get_mcore_model'],
+ 'tuners': ['LoraParallelLinear'],
+ 'utils': ['get_logger', 'set_random_seed'],
'version': ['__release_datetime__', '__version__'],
}
diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py
index 923afb9..136e647 100644
--- a/src/mcore_bridge/bridge/gpt_bridge.py
+++ b/src/mcore_bridge/bridge/gpt_bridge.py
@@ -40,7 +40,7 @@ class GPTBridge:
hf_k_norm_key = 'k_norm.weight'
hf_mlp_prefix = 'mlp'
hf_gate_key = 'gate.weight'
- hf_shared_expert_key = 'shared_expert'
+ hf_shared_expert_key = None
hf_expert_bias_key = 'gate.e_score_correction_bias'
def __init__(self, config):
@@ -661,8 +661,11 @@ def _set_moe_state(
if config.moe_shared_expert_intermediate_size:
hf_shared_expert_key = self.hf_shared_expert_key
- if self.llm_model_type in {'deepseek', 'deepseek_v2', 'deepseek_v3'}:
- hf_shared_expert_key = 'shared_experts'
+ if hf_shared_expert_key is None:
+ if 'qwen' in self.llm_model_type or self.model_type == 'llama4':
+ hf_shared_expert_key = 'shared_expert'
+ else:
+ hf_shared_expert_key = 'shared_experts'
hf_state_dict.update(
self._set_mlp_state(None if mg_mlp is None else mg_mlp.shared_experts, hf_state_dict,
f'{hf_shared_expert_key}.', layer_idx, to_mcore))
@@ -1568,6 +1571,7 @@ def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqd
else:
hf_state_dict = self._convert_hf_state_dict(hf_state_dict, to_mcore)
yield from list(self._add_prefix(hf_state_dict, hf_prefix).items())
+ prog_bar.close()
def _convert_mtp_extra(self, mtp_layer, hf_state_dict, to_mcore, origin_hf_state_dict):
for key in ['enorm.weight', 'hnorm.weight', 'eh_proj.weight']:
diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py
index ad7dd1f..0cf9851 100644
--- a/src/mcore_bridge/config/model_config.py
+++ b/src/mcore_bridge/config/model_config.py
@@ -5,7 +5,7 @@
from dataclasses import dataclass
from megatron.core import mpu
from megatron.core.transformer import TransformerConfig
-from transformers import PretrainedConfig, PreTrainedTokenizerBase
+from transformers import PretrainedConfig
from transformers.utils import is_torch_npu_available
from transformers.utils.versions import require_version
from typing import List, Literal, Optional, Union
@@ -204,7 +204,6 @@ class ModelConfig(TransformerConfig):
# visual
hf_config: Optional[PretrainedConfig] = None
- processor: Optional[PreTrainedTokenizerBase] = None
vit_gradient_checkpointing: Optional[bool] = None
vit_attn_impl: Optional[str] = None # e.g. 'flash_attention_2'
vit_gradient_checkpointing_kwargs: Optional[Union[dict, str]] = None
@@ -305,22 +304,15 @@ def __post_init__(self):
if self.apply_query_key_layer_scaling:
os.environ['NVTE_APPLY_QK_LAYER_SCALING'] = '1'
# patch rotary_interleaved
- _origin_rotary_interleaved = self.rotary_interleaved
- if self.multi_latent_attention and self.rotary_interleaved:
- self.rotary_interleaved = False
super().__post_init__()
- self.rotary_interleaved = _origin_rotary_interleaved
self._check_npu()
if self.mcore_model_type is None:
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.hf_config is None:
- raise ValueError('Multimodal model must specify hf_config.')
- if self.processor is None:
- raise ValueError('Multimodal model must specify processor.')
+ if self.is_multimodal and self.hf_config is None:
+ raise ValueError('Multimodal model must specify hf_config.')
self.is_moe_model = self.num_moe_experts is not None
self.bridge = self.model_meta.bridge_cls(self)
diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py
index a1a7596..84f2d7a 100644
--- a/src/mcore_bridge/config/parser.py
+++ b/src/mcore_bridge/config/parser.py
@@ -100,6 +100,7 @@ def _convert_config(config, _internal_call=False) -> Dict[str, Any]:
def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
res = _convert_config(hf_config)
+ res['hf_config'] = hf_config
hf_model_type = res.get('hf_model_type')
llm_model_type = res.get('llm_model_type') or hf_model_type
res['llm_model_type'] = llm_model_type
@@ -162,8 +163,6 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
res.pop('num_query_groups', None)
if llm_model_type == 'glm_moe_dsa':
res['experimental_attention_variant'] = 'dsa'
- # https://github.com/modelscope/ms-swift/pull/8085
- # res['rotary_interleaved'] = False
elif llm_model_type == 'qwen3_next' or hf_model_type in {'qwen3_5', 'qwen3_5_moe'}:
use_mcore_gdn = get_env_args('USE_MCORE_GDN', bool, True)
res['layernorm_zero_centered_gamma'] = True
@@ -204,10 +203,6 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
mrope_interleaved = rope_scaling.get('mrope_interleaved', False) or rope_scaling.get('interleaved', False)
res['mrope_interleaved'] = mrope_interleaved
- if res.get('multi_latent_attention') and res.get('position_embedding_type') in {
- 'rope', None
- } and 'rotary_interleaved' not in res:
- res['rotary_interleaved'] = True
if first_k_dense_replace is not None:
res['moe_layer_freq'] = f'[0]*{first_k_dense_replace}+[1]*{res["num_layers"] - first_k_dense_replace}'
if res.get('moe_router_score_function', 'softmax') == 'sigmoid' and 'moe_router_enable_expert_bias' not in res:
diff --git a/src/mcore_bridge/model/gpt_model.py b/src/mcore_bridge/model/gpt_model.py
index f752d92..29db5d0 100644
--- a/src/mcore_bridge/model/gpt_model.py
+++ b/src/mcore_bridge/model/gpt_model.py
@@ -175,6 +175,10 @@ def _apply_rotary_pos_emb_bshd(
# ideally t_pass is empty so rotary pos embedding is applied to all tensor t
t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+ if multi_latent_attention:
+ x1 = t[..., 0::2]
+ x2 = t[..., 1::2]
+ t = torch.cat((x1, x2), dim=-1)
# first part is cosine component
# second part is sine component, need to change signs with _rotate_half method
diff --git a/src/mcore_bridge/model/mm_gpts/glm.py b/src/mcore_bridge/model/mm_gpts/glm.py
index 114f005..e39f496 100644
--- a/src/mcore_bridge/model/mm_gpts/glm.py
+++ b/src/mcore_bridge/model/mm_gpts/glm.py
@@ -19,7 +19,7 @@ def prepare_model(self, hf_config: PretrainedConfig):
self.visual = Glm4vMoeVisionModel._from_config(hf_config.vision_config)
def get_inputs_embeds(self, inputs_embeds, **kwargs):
- return self._hf_get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.processor, self.hf_config)
+ return self._hf_get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.hf_config)
register_model(ModelMeta(
diff --git a/src/mcore_bridge/model/mm_gpts/internvl.py b/src/mcore_bridge/model/mm_gpts/internvl.py
index 828bb91..aa142c0 100644
--- a/src/mcore_bridge/model/mm_gpts/internvl.py
+++ b/src/mcore_bridge/model/mm_gpts/internvl.py
@@ -33,6 +33,7 @@ def prepare_attn_impl(self):
self.hf_config.vision_config.use_flash_attn = use_flash_attn
def prepare_model(self, hf_config: PretrainedConfig):
+ from transformers import AutoProcessor
llm_model_type = self.config.llm_model_type
if llm_model_type not in ['qwen2', 'qwen3', 'qwen3_moe', 'gpt_oss']:
raise ValueError(f'{llm_model_type} is not supported for internvl_chat model')
@@ -51,6 +52,7 @@ def prepare_model(self, hf_config: PretrainedConfig):
self.select_layer = hf_config.select_layer
self.downsample_ratio = hf_config.downsample_ratio
self.ps_version = hf_config.ps_version
+ self.processor = AutoProcessor.from_pretrained(hf_config.name_or_path, trust_remote_code=True)
def get_inputs_embeds(self, inputs_embeds, **kwargs):
input_ids = kwargs['input_ids']
diff --git a/src/mcore_bridge/model/mm_gpts/kimi_vl.py b/src/mcore_bridge/model/mm_gpts/kimi_vl.py
index c4bdc58..aedc0f7 100644
--- a/src/mcore_bridge/model/mm_gpts/kimi_vl.py
+++ b/src/mcore_bridge/model/mm_gpts/kimi_vl.py
@@ -1,6 +1,5 @@
# Copyright (c) ModelScope Contributors. All rights reserved.
import torch
-from PIL import Image
from transformers import PretrainedConfig
from transformers.dynamic_module_utils import get_class_from_dynamic_module
@@ -37,17 +36,18 @@ def prepare_model(self, hf_config: PretrainedConfig):
def get_inputs_embeds(self, inputs_embeds, **kwargs):
input_ids = kwargs['input_ids']
pixel_values = kwargs.get('pixel_values')
+ vision_config = self.hf_config.vision_config
if pixel_values is not None and pixel_values.size(0) > 0:
pixel_values = pixel_values.to(self.vision_tower.dtype)
image_features: torch.Tensor = self._extract_image_features(pixel_values, kwargs['image_grid_hws'])
inputs_embeds = inputs_embeds.to(image_features[0].dtype).clone()
inputs_embeds = self._merge_with_image_features(inputs_embeds, input_ids, image_features)
else:
- image_processor = self.processor.image_processor
- dummy_image = Image.new('RGB', (32, 32), (0, 0, 0))
- image_inputs = image_processor([dummy_image], return_tensors='pt')
- pixel_values = image_inputs['pixel_values'].to(self.vision_tower.dtype)
- image_features: torch.Tensor = self._extract_image_features(pixel_values, image_inputs['image_grid_hws'])
+ pixel_values = torch.zeros((16, 3, vision_config.patch_size, vision_config.patch_size),
+ dtype=self.vision_tower.dtype,
+ device=input_ids.device)
+ image_grid_hws = input_ids.new_tensor([[4, 4]])
+ image_features: torch.Tensor = self._extract_image_features(pixel_values, image_grid_hws)
inputs_embeds = inputs_embeds + image_features.mean() * 0.
return inputs_embeds
diff --git a/src/mcore_bridge/model/mm_gpts/qwen.py b/src/mcore_bridge/model/mm_gpts/qwen.py
index bd6c0c1..c959269 100644
--- a/src/mcore_bridge/model/mm_gpts/qwen.py
+++ b/src/mcore_bridge/model/mm_gpts/qwen.py
@@ -28,7 +28,7 @@ def prepare_model(self, hf_config: PretrainedConfig):
self.visual = VisionModel._from_config(hf_config.vision_config)
def get_inputs_embeds(self, inputs_embeds, **kwargs):
- return self._hf_get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.processor, self.hf_config)
+ return self._hf_get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.hf_config)
class Qwen2_5VLBridge(MultimodalGPTBridge):
@@ -85,7 +85,7 @@ def prepare_model(self, hf_config: PretrainedConfig):
def get_inputs_embeds(self, inputs_embeds, **kwargs):
thinker_config = self.hf_config.thinker_config
- inputs_embeds = self._hf_get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.processor, thinker_config)
+ inputs_embeds = self._hf_get_inputs_embeds(inputs_embeds, kwargs, self.visual, thinker_config)
input_ids = kwargs['input_ids']
input_features = kwargs.get('input_features')
feature_attention_mask = kwargs.get('feature_attention_mask')
diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_5.py b/src/mcore_bridge/model/mm_gpts/qwen3_5.py
index a52f571..91472c3 100644
--- a/src/mcore_bridge/model/mm_gpts/qwen3_5.py
+++ b/src/mcore_bridge/model/mm_gpts/qwen3_5.py
@@ -77,7 +77,7 @@ def prepare_model(self, hf_config):
self.visual = Qwen3_5VisionModel._from_config(hf_config.vision_config)
def get_inputs_embeds(self, inputs_embeds, **kwargs):
- return self._hf_get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.processor, self.hf_config)
+ return self._hf_get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.hf_config)
class Qwen3_5Bridge(Qwen3NextBridge):
diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_omni.py b/src/mcore_bridge/model/mm_gpts/qwen3_omni.py
index 94d366b..315340f 100644
--- a/src/mcore_bridge/model/mm_gpts/qwen3_omni.py
+++ b/src/mcore_bridge/model/mm_gpts/qwen3_omni.py
@@ -34,7 +34,7 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs):
input_ids = kwargs['input_ids']
visual = self.visual
hf_config = self.hf_config.thinker_config
- res = Qwen3VL_Vit._get_inputs_embeds(self, inputs_embeds, kwargs, visual, self.processor, hf_config)
+ res = Qwen3VL_Vit._get_inputs_embeds(self, inputs_embeds, kwargs, visual, hf_config)
inputs_embeds = res['inputs_embeds']
input_features = kwargs.get('input_features')
feature_attention_mask = kwargs.get('feature_attention_mask')
diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_vl.py b/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
index 737ecf2..92a90d8 100644
--- a/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
+++ b/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
@@ -8,11 +8,10 @@
from megatron.core.models.gpt import gpt_model
from megatron.core.packed_seq_params import PackedSeqParams
from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor
-from PIL import Image
from typing import List, Optional, Union
from mcore_bridge.bridge import MultimodalGPTBridge
-from mcore_bridge.utils import split_cp_inputs, to_device
+from mcore_bridge.utils import split_cp_inputs
from ..constant import ModelType
from ..register import ModelLoader, ModelMeta, register_model
@@ -321,9 +320,9 @@ def prepare_model(self, hf_config):
self.visual = VisionModel._from_config(hf_config.vision_config)
def get_inputs_embeds(self, inputs_embeds, **kwargs):
- return self._get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.processor, self.hf_config)
+ return self._get_inputs_embeds(inputs_embeds, kwargs, self.visual, self.hf_config)
- def _get_inputs_embeds(self, inputs_embeds, inputs, visual, processor, hf_config):
+ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config):
input_ids = inputs['input_ids']
packed_seq_params = inputs.get('packed_seq_params')
pixel_values = inputs.get('pixel_values')
@@ -331,12 +330,12 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, processor, hf_config
image_grid_thw = inputs.get('image_grid_thw')
video_grid_thw = inputs.get('video_grid_thw')
dtype = visual.dtype
+ vision_config = HuggingFaceVit._get_vision_config(hf_config)
if pixel_values is None and pixel_values_videos is None: # plain-text
- images = [Image.new('RGB', (32, 32), (0, 0, 0))]
- media_inputs = processor.image_processor(images=images, return_tensors='pt')
- media_inputs = to_device(media_inputs, input_ids.device)
- pixel_values = media_inputs['pixel_values'].type(dtype)
- visual_res = visual(pixel_values, grid_thw=media_inputs['image_grid_thw'])
+ hidden_size = vision_config.in_channels * vision_config.temporal_patch_size * vision_config.patch_size**2
+ pixel_values = torch.zeros(16 * 16, hidden_size, dtype=dtype, device=input_ids.device)
+ image_grid_thw = input_ids.new_tensor([[1, 16, 16]])
+ visual_res = visual(pixel_values, grid_thw=image_grid_thw)
if hasattr(visual_res, 'pooler_output'):
image_embeds = visual_res.pooler_output
deepstack_visual_embeds = visual_res.deepstack_features
@@ -369,7 +368,7 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, processor, hf_config
image_embeds = mixed_embeds
video_embeds = None
else:
- merge_length = processor.image_processor.merge_size**2
+ merge_length = vision_config.spatial_merge_size**2
image_tokens = (image_grid_thw.prod(dim=-1) // merge_length).sum()
image_embeds = mixed_embeds[:image_tokens]
video_embeds = mixed_embeds[image_tokens:]
diff --git a/src/mcore_bridge/model/mm_gpts/utils.py b/src/mcore_bridge/model/mm_gpts/utils.py
index dd2f6c4..4521b5f 100644
--- a/src/mcore_bridge/model/mm_gpts/utils.py
+++ b/src/mcore_bridge/model/mm_gpts/utils.py
@@ -3,11 +3,10 @@
from abc import ABC, abstractmethod
from contextlib import contextmanager
from megatron.core.models.huggingface import HuggingFaceModule as _HuggingFaceModule
-from PIL import Image
from transformers import PretrainedConfig, dynamic_module_utils
from mcore_bridge.config import ModelConfig
-from mcore_bridge.utils import safe_ddp_context, to_device
+from mcore_bridge.utils import safe_ddp_context
@contextmanager
@@ -45,7 +44,6 @@ def __init__(self, config: ModelConfig, ignore_init_model_cls=None):
self.prepare_attn_impl()
with patch_get_dynamic_module():
self.prepare_model(hf_config)
- self.processor = config.processor
self.to(device='cuda')
@abstractmethod
@@ -62,19 +60,25 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs):
pass
@staticmethod
- def _hf_get_inputs_embeds(inputs_embeds, inputs, visual, processor, hf_config):
+ def _get_vision_config(hf_config):
+ for k in ['vision_config', 'vit_config']:
+ if hasattr(hf_config, k):
+ return getattr(hf_config, k)
+
+ @staticmethod
+ def _hf_get_inputs_embeds(inputs_embeds, inputs, visual, hf_config):
input_ids = inputs['input_ids']
pixel_values = inputs.get('pixel_values')
pixel_values_videos = inputs.get('pixel_values_videos')
image_grid_thw = inputs.get('image_grid_thw')
video_grid_thw = inputs.get('video_grid_thw')
dtype = visual.dtype
+ vision_config = HuggingFaceVit._get_vision_config(hf_config)
if pixel_values is None and pixel_values_videos is None: # plain-text
- images = [Image.new('RGB', (32, 32), (0, 0, 0))]
- media_inputs = processor.image_processor(images=images, return_tensors='pt')
- media_inputs = to_device(media_inputs, input_ids.device)
- pixel_values = media_inputs['pixel_values'].type(dtype)
- image_embeds = visual(pixel_values, grid_thw=media_inputs['image_grid_thw'])
+ hidden_size = vision_config.in_channels * vision_config.temporal_patch_size * vision_config.patch_size**2
+ pixel_values = torch.zeros(16 * 16, hidden_size, dtype=dtype, device=input_ids.device)
+ image_grid_thw = input_ids.new_tensor([[1, 16, 16]])
+ image_embeds = visual(pixel_values, grid_thw=image_grid_thw)
if hasattr(image_embeds, 'pooler_output'):
image_embeds = image_embeds.pooler_output
inputs_embeds = inputs_embeds + image_embeds.mean().to(device=inputs_embeds.device) * 0.
@@ -99,7 +103,7 @@ def _hf_get_inputs_embeds(inputs_embeds, inputs, visual, processor, hf_config):
image_embeds = mixed_embeds
video_embeds = None
else:
- merge_length = processor.image_processor.merge_size**2
+ merge_length = vision_config.spatial_merge_size**2
image_tokens = (image_grid_thw.prod(dim=-1) // merge_length).sum()
image_embeds = mixed_embeds[:image_tokens]
video_embeds = mixed_embeds[image_tokens:]
diff --git a/src/mcore_bridge/patcher.py b/src/mcore_bridge/patcher.py
index 7b7cd7d..2de8e51 100644
--- a/src/mcore_bridge/patcher.py
+++ b/src/mcore_bridge/patcher.py
@@ -724,9 +724,9 @@ def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor):
# x_pe [seqlen, batch, *, qk_pos_emb_head_dim]
x_pe, x_nope = torch.split(
x, [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], dim=-1)
- origin_rotary_interleaved = self.config.rotary_interleaved
+ origin_multi_latent_attention = self.config.multi_latent_attention
try:
- self.config.rotary_interleaved = self.config.dsa_indexer_rotary_interleaved
+ self.config.multi_latent_attention = self.config.dsa_indexer_rotary_interleaved
x_pe = apply_rotary_pos_emb(
x_pe,
rotary_pos_emb,
@@ -735,7 +735,7 @@ def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor):
cp_group=self.pg_collection.cp,
)
finally:
- self.config.rotary_interleaved = origin_rotary_interleaved
+ self.config.multi_latent_attention = origin_multi_latent_attention
# [seqlen, batch, *, index_head_dim]
x = torch.cat([x_pe, x_nope], dim=-1)
return x
diff --git a/src/mcore_bridge/tuners/lora.py b/src/mcore_bridge/tuners/lora.py
index 05b5a01..6557178 100644
--- a/src/mcore_bridge/tuners/lora.py
+++ b/src/mcore_bridge/tuners/lora.py
@@ -221,6 +221,8 @@ def update_layer(self, adapter_name, r, *, lora_alpha, lora_dropout, init_lora_w
self.set_adapter(self.active_adapters)
def _get_rng_context(self, lora):
+ if not get_cuda_rng_tracker().is_initialized():
+ return nullcontext()
if self.is_expert:
rng_context = get_cuda_rng_tracker().fork(get_expert_parallel_rng_tracker_name())
elif getattr(lora, 'parallel_mode', None) is None:
diff --git a/src/mcore_bridge/utils/__init__.py b/src/mcore_bridge/utils/__init__.py
index be95fec..34cdbd1 100644
--- a/src/mcore_bridge/utils/__init__.py
+++ b/src/mcore_bridge/utils/__init__.py
@@ -3,7 +3,7 @@
from .env import get_dist_setting, get_node_setting, is_dist, is_last_rank, is_local_master, is_master
from .import_utils import _LazyModule, is_flash_attn_3_available
from .logger import get_logger
-from .megatron_utils import get_local_layer_specs, split_cp_inputs, unwrap_model
+from .megatron_utils import get_local_layer_specs, set_random_seed, split_cp_inputs, unwrap_model
from .safetensors import SafetensorLazyLoader, StreamingSafetensorSaver
from .torch_utils import gc_collect, get_current_device, safe_ddp_context, to_device
from .utils import deep_getattr, get_env_args, json_parse_to_dict
diff --git a/src/mcore_bridge/utils/megatron_utils.py b/src/mcore_bridge/utils/megatron_utils.py
index 102d41a..031b4e5 100644
--- a/src/mcore_bridge/utils/megatron_utils.py
+++ b/src/mcore_bridge/utils/megatron_utils.py
@@ -2,12 +2,13 @@
# code borrowed from modelscope/ms-swift
import megatron.core
import torch
-from megatron.core import mpu
+from megatron.core import mpu, tensor_parallel
from megatron.core.distributed import DistributedDataParallel as DDP
from megatron.core.transformer.module import Float16Module
from megatron.core.transformer.transformer_block import get_num_layers_to_build
from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
from packaging import version
+from transformers import set_seed
from typing import Optional
from .logger import get_logger
@@ -79,3 +80,25 @@ def get_local_layer_specs(config, layer_specs, vp_stage=None):
offset = get_transformer_layer_offset(config, **kwargs)
local_layer_specs = layer_specs[offset:offset + num_layers_to_build]
return local_layer_specs
+
+
+def set_random_seed(
+ seed_: int,
+ data_parallel_random_init: bool = False,
+ te_rng_tracker: bool = False,
+ inference_rng_tracker: bool = False,
+ use_cudagraphable_rng: bool = False,
+):
+ """Set random seed for reproducability."""
+ if seed_ is not None and seed_ > 0:
+ # Ensure that different pipeline MP stages get different seeds.
+ seed = seed_ + (1009 * mpu.get_pipeline_model_parallel_rank())
+ # Ensure different data parallel ranks get different seeds
+ if data_parallel_random_init:
+ seed = seed + (11 * mpu.get_data_parallel_rank())
+ set_seed(seed)
+ if torch.cuda.device_count() > 0:
+ tensor_parallel.model_parallel_cuda_manual_seed(seed, te_rng_tracker, inference_rng_tracker,
+ use_cudagraphable_rng)
+ else:
+ raise ValueError(f'Seed ({seed_}) should be a positive integer.')
diff --git a/src/mcore_bridge/version.py b/src/mcore_bridge/version.py
index 7582cd1..67830af 100644
--- a/src/mcore_bridge/version.py
+++ b/src/mcore_bridge/version.py
@@ -1,5 +1,5 @@
# Make sure to modify __release_datetime__ to release time when making official release.
-__version__ = '0.0.1.dev0'
+__version__ = '1.0.0.dev0'
# default release datetime for branches under active development is set
# to be a time far-far-away-into-the-future
__release_datetime__ = '2099-12-31 23:59:59'