Views
No views yet
1from transformers import pipeline
2model_id = "yujiepan/llama-3.3-tiny-random-dim64"
3pipe = pipeline(
4 "text-generation", model=model_id, device="cuda",
5 trust_remote_code=True, max_new_tokens=3,
6)
7print(pipe("Hello World!"))1import torch
2
3from transformers import (
4 AutoConfig,
5 AutoModelForCausalLM,
6 AutoTokenizer,
7 GenerationConfig,
8 pipeline,
9 set_seed,
10)
11
12source_model_id = "meta-llama/Llama-3.3-70B-Instruct"
13save_folder = "/tmp/yujiepan/llama-3.3-tiny-random-dim64"
14
15tokenizer = AutoTokenizer.from_pretrained(
16 source_model_id, trust_remote_code=True,
17)
18tokenizer.save_pretrained(save_folder)
19
20config = AutoConfig.from_pretrained(
21 source_model_id, trust_remote_code=True,
22)
23config.hidden_size = 64
24config.intermediate_size = 128
25config.num_attention_heads = 2
26config.num_key_value_heads = 1
27config.head_dim = 32
28config.num_hidden_layers = 2
29config.tie_word_embeddings = True
30
31model = AutoModelForCausalLM.from_config(
32 config,
33 torch_dtype=torch.bfloat16,
34 trust_remote_code=True,
35)
36model.generation_config = GenerationConfig.from_pretrained(
37 source_model_id, trust_remote_code=True,
38)
39set_seed(42)
40with torch.no_grad():
41 for name, p in sorted(model.named_parameters()):
42 torch.nn.init.normal_(p, 0, 0.2)
43 print(name, p.shape)
44model.save_pretrained(save_folder)1LlamaForCausalLM(
2 (model): LlamaModel(
3 (embed_tokens): Embedding(128256, 64)
4 (layers): ModuleList(
5 (0-1): 2 x LlamaDecoderLayer(
6 (self_attn): LlamaAttention(
7 (q_proj): Linear(in_features=64, out_features=64, bias=False)
8 (k_proj): Linear(in_features=64, out_features=32, bias=False)
9 (v_proj): Linear(in_features=64, out_features=32, bias=False)
10 (o_proj): Linear(in_features=64, out_features=64, bias=False)
11 )
12 (mlp): LlamaMLP(
13 (gate_proj): Linear(in_features=64, out_features=128, bias=False)
14 (up_proj): Linear(in_features=64, out_features=128, bias=False)
15 (down_proj): Linear(in_features=128, out_features=64, bias=False)
16 (act_fn): SiLU()
17 )
18 (input_layernorm): LlamaRMSNorm((64,), eps=1e-05)
19 (post_attention_layernorm): LlamaRMSNorm((64,), eps=1e-05)
20 )
21 )
22 (norm): LlamaRMSNorm((64,), eps=1e-05)
23 (rotary_emb): LlamaRotaryEmbedding()
24 )
25 (lm_head): Linear(in_features=64, out_features=128256, bias=False)
26)