Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3# Load model and tokenizer
4model_id = "tiny-random/ernie-4.5-moe"
5model = AutoModelForCausalLM.from_pretrained(
6 model_id,
7 device_map="auto",
8 torch_dtype="bfloat16",
9 trust_remote_code=True,
10)
11tokenizer = AutoTokenizer.from_pretrained(model_id)
12
13# Generate answer
14prompt = "What is AI?"
15input_ids = tokenizer.apply_chat_template(
16 [{"role": "user", "content": prompt}],
17 add_generation_prompt=True,
18 return_tensors="pt",
19 tokenize=True,
20).to(model.device)
21
22output = model.generate(
23 input_ids,
24 do_sample=True,
25 max_new_tokens=32,
26)
27print(tokenizer.decode(output[0], skip_special_tokens=False))1import json
2from pathlib import Path
3
4import accelerate
5import torch
6from huggingface_hub import file_exists, hf_hub_download
7from transformers import (
8 AutoConfig,
9 AutoModelForCausalLM,
10 AutoProcessor,
11 GenerationConfig,
12 set_seed,
13)
14
15source_model_id = "baidu/ERNIE-4.5-21B-A3B-Thinking"
16save_folder = "/tmp/tiny-random/ernie-4.5-moe"
17
18processor = AutoProcessor.from_pretrained(source_model_id, trust_remote_code=True)
19processor.save_pretrained(save_folder)
20
21with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r', encoding='utf-8') as f:
22 config_json = json.load(f)
23config_json['hidden_size'] = 8
24config_json['intermediate_size'] = 32
25config_json['moe_intermediate_size'] = 32
26# config_json['moe_k'] = 6
27config_json['head_dim'] = 32
28config_json['num_attention_heads'] = 16
29config_json['num_hidden_layers'] = 2
30config_json['num_key_value_heads'] = 8
31config_json['tie_word_embeddings'] = True
32config_json['use_cache'] = True
33with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
34 json.dump(config_json, f, indent=2)
35
36config = AutoConfig.from_pretrained(
37 save_folder,
38 trust_remote_code=True,
39)
40print(config)
41torch.set_default_dtype(torch.bfloat16)
42model = AutoModelForCausalLM.from_config(config)
43torch.set_default_dtype(torch.float32)
44if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):
45 model.generation_config = GenerationConfig.from_pretrained(
46 source_model_id, trust_remote_code=True,
47 )
48 model.generation_config.do_sample = True
49 print(model.generation_config)
50model = model.cpu()
51with torch.no_grad():
52 for name, p in sorted(model.named_parameters()):
53 torch.nn.init.normal_(p, 0, 0.1)
54 print(name, p.shape)
55model.save_pretrained(save_folder)
56print(model)1Ernie4_5_MoeForCausalLM(
2 (model): Ernie4_5_MoeModel(
3 (embed_tokens): Embedding(103424, 8, padding_idx=0)
4 (layers): ModuleList(
5 (0): Ernie4_5_MoeDecoderLayer(
6 (self_attn): Ernie4_5_MoeAttention(
7 (q_proj): Linear(in_features=8, out_features=512, bias=False)
8 (k_proj): Linear(in_features=8, out_features=256, bias=False)
9 (v_proj): Linear(in_features=8, out_features=256, bias=False)
10 (o_proj): Linear(in_features=512, out_features=8, bias=False)
11 )
12 (mlp): Ernie4_5_MoeMLP(
13 (gate_proj): Linear(in_features=8, out_features=32, bias=False)
14 (up_proj): Linear(in_features=8, out_features=32, bias=False)
15 (down_proj): Linear(in_features=32, out_features=8, bias=False)
16 (act_fn): SiLU()
17 )
18 (input_layernorm): Ernie4_5_MoeRMSNorm((8,), eps=1e-05)
19 (post_attention_layernorm): Ernie4_5_MoeRMSNorm((8,), eps=1e-05)
20 )
21 (1): Ernie4_5_MoeDecoderLayer(
22 (self_attn): Ernie4_5_MoeAttention(
23 (q_proj): Linear(in_features=8, out_features=512, bias=False)
24 (k_proj): Linear(in_features=8, out_features=256, bias=False)
25 (v_proj): Linear(in_features=8, out_features=256, bias=False)
26 (o_proj): Linear(in_features=512, out_features=8, bias=False)
27 )
28 (mlp): Ernie4_5_MoeSparseMoeBlock(
29 (moe_statics): Ernie4_5_MoeStatics()
30 (gate): Linear(in_features=8, out_features=64, bias=False)
31 (experts): ModuleList(
32 (0-63): 64 x Ernie4_5_MoeMLP(
33 (gate_proj): Linear(in_features=8, out_features=32, bias=False)
34 (up_proj): Linear(in_features=8, out_features=32, bias=False)
35 (down_proj): Linear(in_features=32, out_features=8, bias=False)
36 (act_fn): SiLU()
37 )
38 )
39 (shared_experts): Ernie4_5_MoeMLP(
40 (gate_proj): Linear(in_features=8, out_features=64, bias=False)
41 (up_proj): Linear(in_features=8, out_features=64, bias=False)
42 (down_proj): Linear(in_features=64, out_features=8, bias=False)
43 (act_fn): SiLU()
44 )
45 )
46 (input_layernorm): Ernie4_5_MoeRMSNorm((8,), eps=1e-05)
47 (post_attention_layernorm): Ernie4_5_MoeRMSNorm((8,), eps=1e-05)
48 )
49 )
50 (norm): Ernie4_5_MoeRMSNorm((8,), eps=1e-05)
51 (rotary_emb): Ernie4_5_MoeRotaryEmbedding()
52 )
53 (lm_head): Linear(in_features=8, out_features=103424, bias=False)
54)