# Copyright (c) ModelScope Contributors. All rights reserved. import sys import torch from transformers import AutoModel, PretrainedConfig, PreTrainedModel from typing import Any, Dict from swift.template import TemplateType from swift.utils import Processor, get_logger, git_clone_github from ..constant import LLMModelType, MLLMModelType from ..model_arch import ModelArch from ..model_meta import Model, ModelGroup, ModelMeta from ..patcher import patch_output_clone, patch_output_to_input_device from ..register import ModelLoader, register_model from ..utils import use_submodel_func class DeepseekLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: model = super().get_model(model_dir, *args, **kwargs) # fix dtype bug mlp_cls = model.model.layers[-1].mlp.__class__ for module in model.modules(): if isinstance(module, mlp_cls): patch_output_to_input_device(module) return model register_model( ModelMeta( LLMModelType.deepseek, [ ModelGroup([ Model('deepseek-ai/deepseek-moe-16b-chat', 'deepseek-ai/deepseek-moe-16b-chat'), Model('deepseek-ai/deepseek-moe-16b-base', 'deepseek-ai/deepseek-moe-16b-base'), ], ), ], DeepseekLoader, template=TemplateType.deepseek, architectures=['DeepseekForCausalLM'], )) register_model( ModelMeta( LLMModelType.deepseek_v2, [ ModelGroup([ Model('deepseek-ai/DeepSeek-Coder-V2-Instruct', 'deepseek-ai/DeepSeek-Coder-V2-Instruct'), Model('deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct', 'deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct'), Model('deepseek-ai/DeepSeek-Coder-V2-Base', 'deepseek-ai/DeepSeek-Coder-V2-Base'), Model('deepseek-ai/DeepSeek-Coder-V2-Lite-Base', 'deepseek-ai/DeepSeek-Coder-V2-Lite-Base'), Model('deepseek-ai/DeepSeek-V2-Lite', 'deepseek-ai/DeepSeek-V2-Lite'), Model('deepseek-ai/DeepSeek-V2-Lite-Chat', 'deepseek-ai/DeepSeek-V2-Lite-Chat'), Model('deepseek-ai/DeepSeek-V2', 'deepseek-ai/DeepSeek-V2'), Model('deepseek-ai/DeepSeek-V2-Chat', 'deepseek-ai/DeepSeek-V2-Chat'), ], TemplateType.deepseek), ModelGroup([ Model('deepseek-ai/DeepSeek-V2.5', 'deepseek-ai/DeepSeek-V2.5'), Model('deepseek-ai/DeepSeek-V2.5-1210', 'deepseek-ai/DeepSeek-V2.5-1210') ], TemplateType.deepseek_v2_5) ], DeepseekLoader, model_arch=ModelArch.deepseek_v2, architectures=['DeepseekV2ForCausalLM'], requires=['transformers>=4.39.3'], )) register_model( ModelMeta( LLMModelType.deepseek_v3, [ ModelGroup([ Model('deepseek-ai/DeepSeek-V3-Base', 'deepseek-ai/DeepSeek-V3-Base'), Model('deepseek-ai/DeepSeek-V3', 'deepseek-ai/DeepSeek-V3'), Model('deepseek-ai/DeepSeek-V3-0324', 'deepseek-ai/DeepSeek-V3-0324'), ], TemplateType.deepseek_v2_5), ModelGroup([ Model('cognitivecomputations/DeepSeek-V3-awq', 'cognitivecomputations/DeepSeek-V3-AWQ'), Model('cognitivecomputations/DeepSeek-V3-0324-AWQ', 'cognitivecomputations/DeepSeek-V3-0324-AWQ') ], TemplateType.deepseek_v2_5), ModelGroup([ Model('deepseek-ai/DeepSeek-Prover-V2-7B', 'deepseek-ai/DeepSeek-Prover-V2-7B'), Model('deepseek-ai/DeepSeek-Prover-V2-671B', 'deepseek-ai/DeepSeek-Prover-V2-671B'), ], TemplateType.deepseek_v2_5), ModelGroup([ Model('unsloth/DeepSeek-V3-bf16', 'unsloth/DeepSeek-V3-bf16'), Model('unsloth/DeepSeek-V3-0324-BF16', 'unsloth/DeepSeek-V3-0324-BF16'), Model('unsloth/DeepSeek-Prover-V2-671B-BF16', 'unsloth/DeepSeek-Prover-V2-671B-BF16'), ], TemplateType.deepseek_v2_5), ModelGroup([ Model('deepseek-ai/DeepSeek-R1', 'deepseek-ai/DeepSeek-R1'), Model('deepseek-ai/DeepSeek-R1-Zero', 'deepseek-ai/DeepSeek-R1-Zero'), Model('deepseek-ai/DeepSeek-R1-0528', 'deepseek-ai/DeepSeek-R1-0528'), ], TemplateType.deepseek_r1), ModelGroup([ Model('cognitivecomputations/DeepSeek-R1-awq', 'cognitivecomputations/DeepSeek-R1-AWQ'), Model('cognitivecomputations/DeepSeek-R1-0528-AWQ', 'cognitivecomputations/DeepSeek-R1-0528-AWQ'), ], TemplateType.deepseek_r1), ModelGroup([ Model('unsloth/DeepSeek-R1-BF16', 'unsloth/DeepSeek-R1-BF16'), Model('unsloth/DeepSeek-R1-Zero-BF16', 'unsloth/DeepSeek-R1-Zero-BF16'), Model('unsloth/DeepSeek-R1-0528-BF16', 'unsloth/DeepSeek-R1-0528-BF16'), ], TemplateType.deepseek_r1), ModelGroup([ Model('moonshotai/Moonlight-16B-A3B', 'moonshotai/Moonlight-16B-A3B'), Model('moonshotai/Moonlight-16B-A3B-Instruct', 'moonshotai/Moonlight-16B-A3B-Instruct'), ], TemplateType.moonlight, requires=['transformers<4.49']), ModelGroup([ Model('moonshotai/Kimi-K2-Base', 'moonshotai/Kimi-K2-Base'), Model('moonshotai/Kimi-K2-Instruct', 'moonshotai/Kimi-K2-Instruct'), Model('moonshotai/Kimi-K2-Instruct-0905', 'moonshotai/Kimi-K2-Instruct-0905'), Model('moonshotai/Kimi-K2-Thinking', 'moonshotai/Kimi-K2-Thinking'), ], TemplateType.kimi_k2), ModelGroup([ Model('deepseek-ai/DeepSeek-V3.1-Base', 'deepseek-ai/DeepSeek-V3.1-Base'), Model('deepseek-ai/DeepSeek-V3.1', 'deepseek-ai/DeepSeek-V3.1'), Model('deepseek-ai/DeepSeek-V3.1-Terminus', 'deepseek-ai/DeepSeek-V3.1-Terminus'), ], TemplateType.deepseek_v3_1), ], DeepseekLoader, model_arch=ModelArch.deepseek_v2, architectures=['DeepseekV3ForCausalLM'], requires=['transformers>=4.39.3'], )) class DeepseekV32Loader(ModelLoader): def get_config(self, model_dir: str): try: from transformers.models.deepseek_v32 import DeepseekV32Config except ImportError: from transformers.models.deepseek_v3 import DeepseekV3Config as DeepseekV32Config return DeepseekV32Config.from_pretrained(model_dir) def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: try: from transformers.models.deepseek_v32 import DeepseekV32ForCausalLM except ImportError: # It’s only for compatibility with Megatron training or vllm/sglang infer, # while we wait for Transformers to support deepseek_v32. from transformers.models.deepseek_v3 import DeepseekV3ForCausalLM as DeepseekV32ForCausalLM if not self.return_dummy_model: raise ValueError('DeepSeek-V3.2 is not supported in transformers.') self.auto_model_cls = DeepseekV32ForCausalLM return super().get_model(model_dir, *args, **kwargs) register_model( ModelMeta( LLMModelType.deepseek_v32, [ ModelGroup([ Model('deepseek-ai/DeepSeek-V3.2', 'deepseek-ai/DeepSeek-V3.2'), Model('deepseek-ai/DeepSeek-V3.2-Speciale', 'deepseek-ai/DeepSeek-V3.2-Speciale'), Model('deepseek-ai/DeepSeek-V3.2-Exp', 'deepseek-ai/DeepSeek-V3.2-Exp'), Model('deepseek-ai/DeepSeek-V3.2-Exp-Base', 'deepseek-ai/DeepSeek-V3.2-Exp-Base'), Model('deepseek-ai/DeepSeek-Math-V2', 'deepseek-ai/DeepSeek-Math-V2'), ]), ], DeepseekV32Loader, template=TemplateType.deepseek_v3_1, architectures=['DeepseekV32ForCausalLM'], )) register_model( ModelMeta( LLMModelType.deepseek_v4, [ ModelGroup([ Model('deepseek-ai/DeepSeek-V4-Flash', 'deepseek-ai/DeepSeek-V4-Flash'), Model('deepseek-ai/DeepSeek-V4-Flash-Base', 'deepseek-ai/DeepSeek-V4-Flash-Base'), ]), ModelGroup([ Model('deepseek-ai/DeepSeek-V4-Pro', 'deepseek-ai/DeepSeek-V4-Pro'), Model('deepseek-ai/DeepSeek-V4-Pro-Base', 'deepseek-ai/DeepSeek-V4-Pro-Base'), ]), ], template=TemplateType.deepseek_v4, architectures=['DeepseekV4ForCausalLM'], )) class DeepseekVLLoader(ModelLoader): def get_config(self, model_dir: str): # compat with python==3.10 if sys.version_info.minor >= 10: import collections import collections.abc for type_name in collections.abc.__all__: setattr(collections, type_name, getattr(collections.abc, type_name)) local_repo_path = self.local_repo_path if not local_repo_path: local_repo_path = git_clone_github('https://github.com/deepseek-ai/DeepSeek-VL') sys.path.append(local_repo_path) from deepseek_vl.models import VLChatProcessor self.auto_tokenizer_cls = VLChatProcessor return super().get_config(model_dir) def _get_model(self, model_dir: str, llm_prefix, *args, **kwargs) -> PreTrainedModel: model = super().get_model(model_dir, *args, **kwargs) llm = getattr(model, llm_prefix) patch_output_clone(llm.model.embed_tokens) patch_output_to_input_device(llm.model.embed_tokens) use_submodel_func(model, llm_prefix) model.generation_config = llm.generation_config return model def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: return self._get_model(model_dir, 'language_model', *args, **kwargs) register_model( ModelMeta( MLLMModelType.deepseek_vl, [ ModelGroup([ Model('deepseek-ai/deepseek-vl-1.3b-chat', 'deepseek-ai/deepseek-vl-1.3b-chat'), Model('deepseek-ai/deepseek-vl-7b-chat', 'deepseek-ai/deepseek-vl-7b-chat'), ], ), ], DeepseekVLLoader, template=TemplateType.deepseek_vl, architectures=['MultiModalityCausalLM'], model_arch=ModelArch.deepseek_vl, tags=['vision'], )) class DeepseekJanusLoader(DeepseekVLLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: return self._get_model(model_dir, 'language_model', *args, **kwargs) def get_config(self, model_dir: str): local_repo_path = self.local_repo_path if not local_repo_path: local_repo_path = git_clone_github('https://github.com/deepseek-ai/Janus') sys.path.append(local_repo_path) from janus.models import VLChatProcessor self.auto_tokenizer_cls = VLChatProcessor return super(DeepseekVLLoader, self).get_config(model_dir) register_model( ModelMeta( MLLMModelType.deepseek_janus, [ ModelGroup([ Model('deepseek-ai/Janus-1.3B', 'deepseek-ai/Janus-1.3B'), ]), ], DeepseekJanusLoader, template=TemplateType.deepseek_janus, model_arch=ModelArch.deepseek_janus, tags=['vision'], )) register_model( ModelMeta( MLLMModelType.deepseek_janus_pro, [ ModelGroup([ Model('deepseek-ai/Janus-Pro-1B', 'deepseek-ai/Janus-Pro-1B'), Model('deepseek-ai/Janus-Pro-7B', 'deepseek-ai/Janus-Pro-7B'), ]), ], DeepseekJanusLoader, template=TemplateType.deepseek_janus_pro, model_arch=ModelArch.deepseek_janus, tags=['vision'], )) class DeepseekVL2Loader(DeepseekVLLoader): def get_config(self, model_dir: str): local_repo_path = self.local_repo_path if not local_repo_path: local_repo_path = git_clone_github('https://github.com/deepseek-ai/DeepSeek-VL2') sys.path.append(local_repo_path) try: from deepseek_vl2.models import DeepseekVLV2Processor except ImportError: # compat transformers>=4.42 import transformers transformers.models.llama.modeling_llama.LlamaFlashAttention2 = None from deepseek_vl2.models import DeepseekVLV2Processor self.auto_tokenizer_cls = DeepseekVLV2Processor return super(DeepseekVLLoader, self).get_config(model_dir) def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: return super()._get_model(model_dir, 'language', *args, **kwargs) register_model( ModelMeta( MLLMModelType.deepseek_vl2, [ ModelGroup([ Model('deepseek-ai/deepseek-vl2-tiny', 'deepseek-ai/deepseek-vl2-tiny'), Model('deepseek-ai/deepseek-vl2-small', 'deepseek-ai/deepseek-vl2-small'), Model('deepseek-ai/deepseek-vl2', 'deepseek-ai/deepseek-vl2'), ]), ], DeepseekVL2Loader, template=TemplateType.deepseek_vl2, model_arch=ModelArch.deepseek_vl2, requires=['transformers<4.42'], tags=['vision'], )) class DeepseekOCRLoader(ModelLoader): visual_name = 'vision_model' def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: self.auto_model_cls = self.auto_model_cls or AutoModel model = super().get_model(model_dir, *args, **kwargs) patch_output_clone(model.model.embed_tokens) patch_output_to_input_device(model.model.sam_model) patch_output_to_input_device(getattr(model.model, self.visual_name)) patch_output_to_input_device(model.model.projector) return model def get_processor(self, model_dir: str, config: PretrainedConfig) -> Processor: from transformers import AutoProcessor, AutoTokenizer # When not loading model (e.g., vllm backend), avoid triggering AutoConfig which would execute # trust_remote_code and cause transformers version compatibility issues # For vllm backend, we only need the processor/tokenizer try: processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True) except Exception: # Fallback to AutoTokenizer if AutoProcessor is not available processor = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) return processor class DeepseekOCR2Loader(DeepseekOCRLoader): visual_name = 'qwen2_model' register_model( ModelMeta( MLLMModelType.deepseek_ocr, [ ModelGroup([ Model('deepseek-ai/DeepSeek-OCR', 'deepseek-ai/DeepSeek-OCR'), ]), ], DeepseekOCRLoader, template=TemplateType.deepseek_ocr, model_arch=ModelArch.deepseek_ocr, architectures=['DeepseekOCRForCausalLM'], requires=['transformers==4.46.3', 'easydict'], tags=['vision'], )) register_model( ModelMeta( MLLMModelType.deepseek_ocr2, [ ModelGroup([ Model('deepseek-ai/DeepSeek-OCR-2', 'deepseek-ai/DeepSeek-OCR-2'), ]), ], DeepseekOCR2Loader, template=TemplateType.deepseek_ocr2, model_arch=ModelArch.deepseek_ocr2, architectures=['DeepseekOCR2ForCausalLM'], requires=['transformers==4.46.3', 'easydict'], tags=['vision'], )) class UnlimitedOCRLoader(DeepseekOCRLoader): visual_name = 'vision_model' @staticmethod def _apply_multi_gpu_patch(): """ Fixed two bugs affecting `UnlimitedOCRModel` in multi-GPU scenarios using `device_map='auto'`: Bug 1 - Device mismatch in `torch.cat`: `image_newline` and `view_seperator` are `nn.Parameter`s; under `device_map='auto'`, their device placement might not align with the image features. Bug 2 - Device mismatch in `masked_scatter_`: Hard-coded `.cuda()` usage caused a conflict where `images_in_this_batch` resided on the projector's device (e.g., `cuda:7`), while `inputs_embeds` resided on the device hosting `embed_tokens` (e.g., `cuda:0`). Fix strategy: Temporarily replace `torch.cat` and `torch.Tensor.masked_scatter_` during the forward pass to handle device placement automatically, then restore the original methods after execution. """ modeling_module = None for mod_name, mod in sys.modules.items(): if 'modeling_unlimitedocr' in mod_name: modeling_module = mod break if modeling_module is None: return False UnlimitedOCRModel = getattr(modeling_module, 'UnlimitedOCRModel', None) if UnlimitedOCRModel is None: return False # Avoid redundant patching if getattr(UnlimitedOCRModel, '_swift_multi_gpu_patched', False): return True _original_forward = UnlimitedOCRModel.forward def _patched_forward(self, *args, **kwargs): _orig_cat = torch.cat _orig_masked_scatter_ = torch.Tensor.masked_scatter_ def _safe_cat(tensors, dim=0, **cat_kwargs): # Using the device of the first tensor as the reference, the others are aligned to it. ref_device = None for t in tensors: if isinstance(t, torch.Tensor): ref_device = t.device break if ref_device is None: return _orig_cat(tensors, dim, **cat_kwargs) aligned = [ t.to(ref_device) if isinstance(t, torch.Tensor) and t.device != ref_device else t for t in tensors ] return _orig_cat(aligned, dim, **cat_kwargs) def _safe_masked_scatter_(tensor_self, mask, source): # Use the device of tensor_self (inputs_embeds[idx]) as the reference. dev = tensor_self.device if mask.device != dev: mask = mask.to(dev) if source.device != dev: source = source.to(dev) return _orig_masked_scatter_(tensor_self, mask, source) # Simultaneously replace the module namespace and the global scope (double insurance). modeling_module.torch.cat = _safe_cat torch.cat = _safe_cat torch.Tensor.masked_scatter_ = _safe_masked_scatter_ try: return _original_forward(self, *args, **kwargs) finally: # Restore the state to avoid contaminating other modules. modeling_module.torch.cat = _orig_cat torch.cat = _orig_cat torch.Tensor.masked_scatter_ = _orig_masked_scatter_ UnlimitedOCRModel.forward = _patched_forward UnlimitedOCRModel._swift_multi_gpu_patched = True return True def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: logger = get_logger() self.auto_model_cls = self.auto_model_cls or AutoModel model = super(DeepseekOCRLoader, self).get_model(model_dir, *args, **kwargs) patch_output_clone(model.model.embed_tokens) patch_output_to_input_device(model.model.sam_model) patch_output_to_input_device(getattr(model.model, self.visual_name)) patch_output_to_input_device(model.model.projector) patch_output_to_input_device(model.model) _orig_sw = (getattr(model.config, 'sliding_window_size', None) or getattr(model.config, 'sliding_window', None)) if _orig_sw is not None: model.config._ring_window = _orig_sw model.config.sliding_window = None logger.info('[UnlimitedOCR] R-SWA enabled: ring_window=%d', _orig_sw) else: logger.warning('[UnlimitedOCR] sliding_window config not found, R-SWA may not work.') n_devices = len(set(str(p.device) for p in model.parameters() if p.device.type == 'cuda')) if n_devices > 1: if self._apply_multi_gpu_patch(): logger.info('[UnlimitedOCR] Multi-GPU patch applied (%d GPUs).', n_devices) else: logger.warning('[UnlimitedOCR] Multi-GPU deployment failed to apply patch.' 'If an inference error occurs, please check whether' ' `modeling_unlimitedocr` has been loaded correctly.') return model register_model( ModelMeta( MLLMModelType.unlimited_ocr, [ ModelGroup([ Model('PaddlePaddle/Unlimited-OCR', 'PaddlePaddle/Unlimited-OCR'), ]), ], UnlimitedOCRLoader, template=TemplateType.unlimited_ocr, model_arch=ModelArch.unlimited_ocr, architectures=['UnlimitedOCRForCausalLM'], requires=['transformers==4.46.3', 'easydict'], tags=['vision'], ))