wehub-resource-sync
194 lines
8.4 KiB
Python
194 lines
8.4 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import json
|
|
import os
|
|
from transformers import AutoProcessor, PretrainedConfig, PreTrainedModel
|
|
from transformers.dynamic_module_utils import get_class_from_dynamic_module
|
|
|
|
from swift.template import TemplateType
|
|
from swift.utils import Processor, get_device, get_device_count, get_dist_setting, get_logger
|
|
from ..constant import LLMModelType, MLLMModelType
|
|
from ..model_arch import ModelArch
|
|
from ..model_meta import Model, ModelGroup, ModelMeta
|
|
from ..patcher import patch_ignore_check_imports
|
|
from ..register import ModelLoader, register_model
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class MiniMaxVLLoader(ModelLoader):
|
|
|
|
def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel:
|
|
logger.warn('NOTE: minimax-vl-01 model does not support training.')
|
|
n_gpu = get_device_count()
|
|
_, local_rank, _, local_world_size = get_dist_setting()
|
|
device_ids = list(range(max(local_rank, 0), n_gpu, local_world_size))
|
|
if 'quantization_config' in model_kwargs:
|
|
quantization_config = model_kwargs['quantization_config']
|
|
from transformers import QuantoConfig
|
|
if isinstance(quantization_config, QuantoConfig):
|
|
quantization_config.modules_to_not_convert = (
|
|
[
|
|
'vision_tower',
|
|
'image_newline',
|
|
'multi_modal_projector',
|
|
'lm_head',
|
|
'embed_tokens',
|
|
] + [f'model.layers.{i}.coefficient' for i in range(config.text_config.num_hidden_layers)]
|
|
+ [f'model.layers.{i}.block_sparse_moe.gate' for i in range(config.text_config.num_hidden_layers)])
|
|
|
|
if len(device_ids) > 1:
|
|
model_safetensors_index_path = os.path.join(model_dir, 'model.safetensors.index.json')
|
|
with open(model_safetensors_index_path, 'r') as f:
|
|
model_safetensors_index = json.load(f)
|
|
weight_map = model_safetensors_index['weight_map']
|
|
vision_map = {}
|
|
for key, value in weight_map.items():
|
|
if 'vision_tower' in key or 'image_newline' in key or 'multi_modal_projector' in key:
|
|
new_key = key.replace('.weight', '').replace('.bias', '')
|
|
if new_key not in vision_map:
|
|
vision_map[new_key] = value
|
|
|
|
device_map = {
|
|
'language_model.model.embed_tokens': get_device(device_ids[0]),
|
|
'language_model.model.norm': get_device(device_ids[len(device_ids) - 1]),
|
|
'language_model.lm_head': get_device(device_ids[len(device_ids) - 1])
|
|
}
|
|
for key, value in vision_map.items():
|
|
device_map[key] = get_device(device_ids[0])
|
|
device_map['vision_tower.vision_model.post_layernorm'] = get_device(device_ids[0])
|
|
layers_per_device = config.text_config.num_hidden_layers // len(device_ids)
|
|
for i in range(len(device_ids)):
|
|
for j in range(layers_per_device):
|
|
device_map[f'language_model.model.layers.{i * layers_per_device + j}'] = get_device(device_ids[i])
|
|
model_kwargs['device_map'] = device_map
|
|
with patch_ignore_check_imports():
|
|
return super().get_model(model_dir, config, processor, model_kwargs)
|
|
|
|
def get_processor(self, model_dir: str, config: PretrainedConfig) -> Processor:
|
|
MiniMaxVL01ProcessorKwargs = get_class_from_dynamic_module(
|
|
'processing_minimax_vl_01.MiniMaxVL01ProcessorKwargs', model_dir)
|
|
get_hw_multiple_of = get_class_from_dynamic_module('processing_minimax_vl_01.get_hw_multiple_of', model_dir)
|
|
get_num_token = get_class_from_dynamic_module('processing_minimax_vl_01.get_num_token', model_dir)
|
|
|
|
processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True)
|
|
processor.MiniMaxVL01ProcessorKwargs = MiniMaxVL01ProcessorKwargs
|
|
processor.get_hw_multiple_of = get_hw_multiple_of
|
|
processor.get_num_token = get_num_token
|
|
return processor
|
|
|
|
|
|
register_model(
|
|
ModelMeta(
|
|
MLLMModelType.minimax_vl, [
|
|
ModelGroup([
|
|
Model('MiniMax/MiniMax-VL-01', 'MiniMaxAI/MiniMax-VL-01'),
|
|
]),
|
|
],
|
|
MiniMaxVLLoader,
|
|
template=TemplateType.minimax_vl,
|
|
architectures=['MiniMaxVL01ForConditionalGeneration'],
|
|
tags=['vision']))
|
|
|
|
|
|
class MinimaxTextLoader(ModelLoader):
|
|
|
|
def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel:
|
|
logger.warn('NOTE: minimax-text-01 model does not support training.')
|
|
n_gpu = get_device_count()
|
|
_, local_rank, _, local_world_size = get_dist_setting()
|
|
device_ids = list(range(max(local_rank, 0), n_gpu, local_world_size))
|
|
if 'quantization_config' in model_kwargs:
|
|
quantization_config = model_kwargs['quantization_config']
|
|
from transformers import QuantoConfig
|
|
if isinstance(quantization_config, QuantoConfig):
|
|
quantization_config.modules_to_not_convert = (
|
|
[
|
|
'lm_head',
|
|
'embed_tokens',
|
|
] + [f'model.layers.{i}.coefficient' for i in range(config.num_hidden_layers)]
|
|
+ [f'model.layers.{i}.block_sparse_moe.gate' for i in range(config.num_hidden_layers)])
|
|
|
|
if len(device_ids) > 1:
|
|
layers_per_device = config.num_hidden_layers // len(device_ids)
|
|
# set device map
|
|
device_map = {
|
|
'model.embed_tokens': get_device(0),
|
|
'model.norm': get_device(len(device_ids) - 1),
|
|
'lm_head': get_device(len(device_ids) - 1)
|
|
}
|
|
for i in range(len(device_ids)):
|
|
for j in range(layers_per_device):
|
|
device_map[f'model.layers.{i * layers_per_device + j}'] = get_device(i)
|
|
model_kwargs['device_map'] = device_map
|
|
with patch_ignore_check_imports():
|
|
return super().get_model(model_dir, config, processor, model_kwargs)
|
|
|
|
|
|
register_model(
|
|
ModelMeta(
|
|
LLMModelType.minimax, [
|
|
ModelGroup([
|
|
Model('MiniMax/MiniMax-Text-01', 'MiniMaxAI/MiniMax-Text-01'),
|
|
]),
|
|
],
|
|
MinimaxTextLoader,
|
|
template=TemplateType.minimax,
|
|
architectures=['MiniMaxText01ForCausalLM']))
|
|
|
|
register_model(
|
|
ModelMeta(
|
|
LLMModelType.minimax_m1, [
|
|
ModelGroup([
|
|
Model('MiniMax/MiniMax-M1-40k', 'MiniMaxAI/MiniMax-M1-40k'),
|
|
Model('MiniMax/MiniMax-M1-80k', 'MiniMaxAI/MiniMax-M1-80k'),
|
|
]),
|
|
],
|
|
MinimaxTextLoader,
|
|
template=TemplateType.minimax_m1,
|
|
architectures=['MiniMaxM1ForCausalLM']))
|
|
|
|
register_model(
|
|
ModelMeta(
|
|
LLMModelType.minimax_m2, [
|
|
ModelGroup([
|
|
Model('MiniMax/MiniMax-M2', 'MiniMaxAI/MiniMax-M2'),
|
|
], TemplateType.minimax_m2),
|
|
ModelGroup([
|
|
Model('MiniMax/MiniMax-M2.1', 'MiniMaxAI/MiniMax-M2.1'),
|
|
], TemplateType.minimax_m2_1),
|
|
ModelGroup([
|
|
Model('MiniMax/MiniMax-M2.5', 'MiniMaxAI/MiniMax-M2.5'),
|
|
], TemplateType.minimax_m2_5),
|
|
ModelGroup([
|
|
Model('MiniMax/MiniMax-M2.7', 'MiniMaxAI/MiniMax-M2.7'),
|
|
], TemplateType.minimax_m2_7),
|
|
],
|
|
requires=['transformers==4.57.1'],
|
|
architectures=['MiniMaxM2ForCausalLM']))
|
|
|
|
|
|
class MinimaxM3VLLoader(ModelLoader):
|
|
default_trust_remote_code = False
|
|
|
|
def get_processor(self, model_dir: str, config: PretrainedConfig) -> Processor:
|
|
return AutoProcessor.from_pretrained(model_dir, trust_remote_code=True)
|
|
|
|
def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel:
|
|
from transformers import AutoModelForImageTextToText
|
|
self.auto_model_cls = self.auto_model_cls or AutoModelForImageTextToText
|
|
return super().get_model(model_dir, config, processor, model_kwargs)
|
|
|
|
|
|
register_model(
|
|
ModelMeta(
|
|
MLLMModelType.minimax_m3_vl, [
|
|
ModelGroup([
|
|
Model('MiniMax/MiniMax-M3', 'MiniMaxAI/MiniMax-M3'),
|
|
]),
|
|
],
|
|
MinimaxM3VLLoader,
|
|
template=TemplateType.minimax_m3_vl,
|
|
model_arch=ModelArch.minimax_m3_vl,
|
|
architectures=['MiniMaxM3SparseForConditionalGeneration'],
|
|
tags=['vision', 'video']))
|