Views
No views yet
1from transformers import pipeline
2model_id = "tiny-random/zaya1"
3pipe = pipeline('text-generation', model=model_id,
4 device='cuda', dtype="bfloat16")
5print(pipe('Hello World!'))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 AutoTokenizer,
12 GenerationConfig,
13 set_seed,
14)
15
16source_model_id = "Zyphra/ZAYA1-reasoning-base"
17save_folder = "/tmp/tiny-random/zaya1"
18
19processor = AutoTokenizer.from_pretrained(
20 source_model_id, trust_remote_code=True)
21processor.save_pretrained(save_folder)
22
23with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r', encoding='utf-8') as f:
24 config_json = json.load(f)
25config_json['hidden_size'] = 512
26config_json['num_attention_heads'] = 4
27config_json['num_key_value_heads'] = 1
28config_json['num_hidden_layers'] = 2
29# bug. need to first set False and then hack
30config_json['tie_word_embeddings'] = False
31config_json['cca_num_q_heads'] = [2, 0]
32config_json['ffn_hidden_size_list'] = [0, 32]
33config_json['num_query_groups_list'] = [1, 0]
34config_json['zaya_layers'] = ['a', 16]
35config_json['zaya_mlp_expansion'] = [0, 8]
36
37with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
38 json.dump(config_json, f, indent=2)
39
40config = AutoConfig.from_pretrained(
41 save_folder,
42 trust_remote_code=True,
43)
44print(config)
45torch.set_default_dtype(torch.bfloat16)
46model = AutoModelForCausalLM.from_config(config)
47model.lm_head = None
48torch.set_default_dtype(torch.float32)
49if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):
50 model.generation_config = GenerationConfig.from_pretrained(
51 source_model_id, trust_remote_code=True,
52 )
53set_seed(42)
54model = model.cpu()
55with torch.no_grad():
56 for name, p in sorted(model.named_parameters()):
57 torch.nn.init.normal_(p, 0, 0.1)
58 print(name, p.shape)
59model.save_pretrained(save_folder)
60with open(f"{save_folder}/config.json", 'r', encoding='utf-8') as f:
61 config_json = json.load(f)
62 config_json['tie_word_embeddings'] = True
63with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
64 json.dump(config_json, f, indent=2)1ZayaForCausalLM(
2 (model): ZayaModel(
3 (embed_tokens): Embedding(262272, 512, padding_idx=0)
4 (layers): ModuleList(
5 (0): ZayaDecoderATTLayer(
6 (self_attn): ZayaSdpaAttention(
7 (o_proj): Linear(in_features=256, out_features=512, bias=False)
8 (qkv): CCA(
9 (linear_q): Linear(in_features=512, out_features=256, bias=False)
10 (linear_k): Linear(in_features=512, out_features=128, bias=False)
11 (val_proj1): Linear(in_features=512, out_features=64, bias=False)
12 (val_proj2): Linear(in_features=512, out_features=64, bias=False)
13 (conv_qk): Sequential(
14 (0): Conv1d(384, 384, kernel_size=(2,), stride=(1,), groups=384)
15 (1): Conv1d(384, 384, kernel_size=(2,), stride=(1,), groups=3)
16 )
17 )
18 )
19 (input_norm): ZayaRMSNorm((512,), eps=1e-05)
20 (res_scale): ResidualScaling()
21 )
22 (1): ZayaDecoderMLPLayer(
23 (zaya_block): ZayaBlock(
24 (router): ZayaRouter(
25 (down_proj): Linear(in_features=512, out_features=8, bias=True)
26 (rmsnorm_eda): ZayaRMSNorm((8,), eps=1e-06)
27 (non_linearity): GELU(approximate='none')
28 (router_mlp): Sequential(
29 (0): Linear(in_features=8, out_features=8, bias=True)
30 (1): GELU(approximate='none')
31 (2): Linear(in_features=8, out_features=8, bias=True)
32 (3): GELU(approximate='none')
33 (4): Linear(in_features=8, out_features=17, bias=False)
34 )
35 )
36 (experts): SequentialMLP(
37 (local_experts): ModuleList(
38 (0-15): 16 x MLP(
39 (linear_fc1): Linear(in_features=512, out_features=32, bias=False)
40 (linear_fc2): Linear(in_features=16, out_features=512, bias=False)
41 )
42 )
43 )
44 )
45 (input_norm): ZayaRMSNorm((512,), eps=1e-05)
46 (res_scale): ResidualScaling()
47 )
48 )
49 (res_scale): ResidualScaling()
50 (final_norm): ZayaRMSNorm((512,), eps=1e-05)
51 (rotary_emb): ZayaRotaryEmbedding()
52 )
53 (lm_head): None
54)