Views
No views yet
1import os
2
3import torch
4import transformers
5from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
6 GenerationConfig, pipeline, set_seed)
7
8model_id = "microsoft/Phi-3.5-MoE-instruct"
9repo_id = "yujiepan/phi-3.5-moe-tiny-random"
10save_path = f"/tmp/{repo_id}"
11
12config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
13config.hidden_size = 16
14config.intermediate_size = 32
15config.num_attention_heads = 4
16config.num_hidden_layers = 2
17config.num_key_value_heads = 4
18config.rope_scaling['long_factor'] = [1.0299, 1.0499]
19config.rope_scaling['short_factor'] = [1.05, 1.05]
20
21tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
22tokenizer.save_pretrained(save_path)
23
24model = AutoModelForCausalLM.from_config(
25 config, torch_dtype=torch.bfloat16,
26 # attn_implementation="sdpa",
27 trust_remote_code=True,
28)
29model.generation_config = GenerationConfig.from_pretrained(
30 model_id, trust_remote_code=True
31)
32
33set_seed(42)
34with torch.no_grad():
35 for _, p in sorted(model.named_parameters()):
36 torch.nn.init.uniform_(p, -0.3, 0.3)
37
38model.save_pretrained(save_path)
39
40pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device="cuda",
41 trust_remote_code=True, max_new_tokens=20)
42print(pipe('Hello'))