Views
No views yet
1# combine weights of expert befor inference:
2for layer in self.model.layers:
3 if isinstance(layer.mlp, DeepseekV2MoE):
4 moe_layer = layer.mlp
5 # combine experts
6 moe_layer.w1 = nn.Parameter(torch.stack([moe_layer.experts[i].gate_proj.weight.T for i in range(moe_layer.config.n_routed_experts)]), requires_grad=False)
7 moe_layer.w2 = nn.Parameter(torch.stack([moe_layer.experts[i].down_proj.weight.T for i in range(moe_layer.config.n_routed_experts)]), requires_grad=False)
8 moe_layer.w3 = nn.Parameter(torch.stack([moe_layer.experts[i].up_proj.weight.T for i in range(moe_layer.config.n_routed_experts)]), requires_grad=False)
9
10# patch the new forward method of DeepseekV2MoE
11
12def new_forward_for_moe(self, hidden_states):
13 batch_size, sequence_length, hidden_dim = hidden_states.shape
14 selected_experts, routing_weights = self.gate(hidden_states)
15 router_scores = torch.zeros(size=(batch_size * sequence_length, self.config.n_routed_experts), device=hidden_states.device, dtype=hidden_states.dtype)
16 # we cast back to the input dtype
17 routing_weights = routing_weights.to(hidden_states.dtype)
18 router_scores = torch.scatter_add(router_scores, -1, selected_experts, routing_weights)
19 hidden_states = hidden_states.view(-1, hidden_dim)
20 if self.config.n_shared_experts is not None:
21 shared_expert_output = self.shared_experts(hidden_states)
22
23 hidden_w1 = torch.matmul(hidden_states, self.w1)
24 hidden_w3 = torch.matmul(hidden_states, self.w3)
25 hidden_states = self.act(hidden_w1) * hidden_w3
26 hidden_states = torch.bmm(hidden_states, self.w2) * torch.transpose(router_scores, 0, 1).unsqueeze(-1)
27 final_hidden_states = hidden_states.sum(dim=0, dtype=hidden_states.dtype)
28 if self.config.n_shared_experts is not None:
29 hidden_states = final_hidden_states + shared_expert_output
30 return hidden_states.view(batch_size, sequence_length, hidden_dim)
31mindspore==2.7.0
mindnlp==0.5.0rc4
transformers==4.57.1
tokenizers
einops
addict
easydict1import os
2import mindnlp
3import mindspore
4from transformers import AutoModel, AutoTokenizer
5
6model_name = 'lvyufeng/DeepSeek-OCR'
7
8tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
9model = AutoModel.from_pretrained(model_name, dtype=mindspore.float16, _attn_implementation='sdpa', trust_remote_code=True, use_safetensors=True, device_map='auto')
10model = model.eval()
11
12# combine experts
13model.combine_moe()
14
15# prompt = "<image>\nFree OCR. "
16prompt = "<image>\n<|grounding|>Convert the document to markdown. "
17image_file = 'your_image.jpg'
18output_path = 'your/output/dir'
19
20# infer(self, tokenizer, prompt='', image_file='', output_path = ' ', base_size = 1024, image_size = 640, crop_mode = True, test_compress = False, save_results = False):
21
22# Tiny: base_size = 512, image_size = 512, crop_mode = False
23# Small: base_size = 640, image_size = 640, crop_mode = False
24# Base: base_size = 1024, image_size = 1024, crop_mode = False
25# Large: base_size = 1280, image_size = 1280, crop_mode = False
26
27# Gundam: base_size = 1024, image_size = 640, crop_mode = True
28
29res = model.infer(tokenizer, prompt=prompt, image_file=image_file, output_path = output_path, base_size = 1024, image_size = 640, crop_mode=True, save_results = True, test_compress = True)torch
transformers==4.57.1
tokenizers
einops
addict
easydict
pip install flash-attn1from transformers import AutoModel, AutoTokenizer
2import torch
3import os
4os.environ["CUDA_VISIBLE_DEVICES"] = '0'
5model_name = 'lvyufeng/DeepSeek-OCR'
6tokenizer = AutoTokenizer.from_pretrained(model_name, dtype=torch.bfloat16,trust_remote_code=True, device_map='auto')
7model = AutoModel.from_pretrained(model_name, _attn_implementation='flash_attention_2', trust_remote_code=True, use_safetensors=True)
8model = model.eval()
9
10# combine experts
11model.combine_moe()
12
13# prompt = "<image>\nFree OCR. "
14prompt = "<image>\n<|grounding|>Convert the document to markdown. "
15image_file = 'your_image.jpg'
16output_path = 'your/output/dir'
17# infer(self, tokenizer, prompt='', image_file='', output_path = ' ', base_size = 1024, image_size = 640, crop_mode = True, test_compress = False, save_results = False):
18# Tiny: base_size = 512, image_size = 512, crop_mode = False
19# Small: base_size = 640, image_size = 640, crop_mode = False
20# Base: base_size = 1024, image_size = 1024, crop_mode = False
21# Large: base_size = 1280, image_size = 1280, crop_mode = False
22# Gundam: base_size = 1024, image_size = 640, crop_mode = True
23res = model.infer(tokenizer, prompt=prompt, image_file=image_file, output_path = output_path, base_size = 1024, image_size = 640, crop_mode=True, save_results = True, test_compress = True)