MCore-Bridge

🛠️ Installation

To install using pip: ```shell pip install mcore-bridge -U

pip install git+https://github.com/modelscope/mcore-bridge.git

git clone https://github.com/modelscope/mcore-bridge.git cd mcore-bridge pip install -e .

🚀 Quick Start

How to use MCore-Bridge for training can be referred to the ms-swift project. Here we introduce how to use MCore-Bridge programmatically.

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. Below is sample code demonstrating how to use Mcore-Bridge for model creation, weight loading, export, and saving.

The saved model can be used for inference by referring to the example code in the model card.

```python import os import torch import torch.distributed as dist from megatron.core import mpu 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 mpu.initialize_model_parallel( tensor_model_parallel_size=TP, pipeline_model_parallel_size=PP, expert_model_parallel_size=EP, expert_tensor_parallel_size=ETP, )

model_dir = snapshot_download('Qwen/Qwen3.5-35B-A3B') hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) processor = AutoProcessor.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, expert_model_parallel_size=EP, expert_tensor_parallel_size=ETP, sequence_parallel=True, mtp_num_layers=1, **config_kwargs)

Minimal forward example

Mcore-Bridge integrates seamlessly with the ms-swift template for model training. You can also replace the ms-swift template module with a custom data processing pipeline to suit your own workflow.

The following provides a minimal example demonstrating how to perform a forward pass and compute the loss using a model created with Mcore-Bridge, helping users quickly integrate Mcore-Bridge into other projects.

Create the following file (test.py) and run it with: CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 test.py.

import os
import torch
import torch.distributed as dist
from megatron.core import mpu
from modelscope import snapshot_download
from swift import get_processor, get_template
from swift.megatron.utils import get_packed_seq_params, get_padding_to
from swift.utils import to_device

from mcore_bridge import ModelConfig, get_mcore_model, hf_to_mcore_config, set_random_seed

data = {
    'messages': [{
        'role': 'user',
        'content': '<image>describe the image.'
    }, {
        'role':
        'assistant',
        'content':
        'The image depicts a close-up of a kitten with striking features. '
        'The kitten has a white and gray coat with distinct black stripes, '
        'particularly noticeable on its face and ears. Its eyes are large '
        'and expressive, with a captivating blue hue that stands out against '
        "the darker fur around them. The kitten's nose is small and pink, "
        'and it has long, delicate whiskers extending from either side of its mouth. '
        "The background is blurred, drawing attention to the kitten's face and "
        'making it the focal point of the image. The overall impression is '
        'one of cuteness and charm.'
    }],
    'images': ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png']
}


def forward_mg_model(mg_model, template):
    template.use_megatron = True
    template.set_mode('train')
    inputs = template.encode(data, return_length=True)
    mg_inputs = to_device(template.data_collator([inputs], padding_to=get_padding_to(mg_model.config)), 'cuda')
    text_position_ids = mg_inputs.pop('text_position_ids', None)
    if text_position_ids is None:
        text_position_ids = mg_inputs.get('position_ids')
    for key in ['num_samples', 'attention_mask_2d', 'loss_scale']:
        mg_inputs.pop(key, None)
    if template.padding_free:
        mg_inputs['packed_seq_params'] = get_packed_seq_params(text_position_ids)
    mg_inputs['labels'] = torch.roll(mg_inputs['labels'], -1, dims=-1)
    loss = mg_model(**mg_inputs)
    loss_mask = mg_inputs['labels'] != -100
    loss = loss * loss_mask
    return loss.sum() / loss_mask.sum()


torch.cuda.set_device(f"cuda:{os.getenv('LOCAL_RANK')}")
dist.init_process_group(backend='nccl')
TP, PP, EP, ETP = 2, 1, 2, 1
mpu.initialize_model_parallel(
    tensor_model_parallel_size=TP,
    pipeline_model_parallel_size=PP,
    expert_model_parallel_size=EP,
    expert_tensor_parallel_size=ETP,
)
set_random_seed(42)

model_dir = snapshot_download('Qwen/Qwen3.5-35B-A3B')
template = get_template(get_processor(model_dir), padding_free=True)
config_kwargs = hf_to_mcore_config(template.config)
config = ModelConfig(
    params_dtype=torch.bfloat16,
    tensor_model_parallel_size=TP,
    pipeline_model_parallel_size=PP,
    expert_model_parallel_size=EP,
    expert_tensor_parallel_size=ETP,
    sequence_parallel=True,
    mtp_num_layers=1,
    **config_kwargs)

mg_model = get_mcore_model(config)[0]
mg_model.cuda()
config.bridge.load_weights([mg_model], model_dir)
loss = forward_mg_model(mg_model, template)
print(f'loss: {loss}')  # loss: 0.8161308169364929

Optional

✨ Model List

The following is the list of models supported by MCore-Bridge:

text-only large models:

multimodal large models: <div class="rdm-tbl-wrap"><table class="rdm-tbl"><thead><tr><th>Series</th><th>model_type</th></tr></thead><tbody><tr><td>Qwen</td><td>qwen2_vl, qwen2_5_vl, qwen2_5_omni&lt;br /&gt;qwen3_vl, qwen3_vl_moe, qwen3_omni_moe, qwen3_asr&lt;br /&gt;qwen3_5, qwen3_5_moe</td></tr><tr><td>Gemma</td><td>gemma4</td></tr><tr><td>GLM</td><td>glm4v, glm4v_moe</td></tr><tr><td>Kimi</td><td>kimi_vl</td></tr><tr><td>InternVL</td><td>internvl_chat, internvl</td></tr><tr><td>Ovis</td><td>ovis2_5</td></tr><tr><td>Llama</td><td>llama4</td></tr><tr><td>Llava</td><td>llava-onevision</td></tr></tbody></table></div>

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]