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'