Views
No views yet
1import os
2
3import torch
4import transformers
5from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, pipeline, set_seed
6
7model_id = "meta-llama/Llama-3.3-70B-Instruct"
8repo_id = "yujiepan/llama-3.3-tiny-random"
9save_path = f"/tmp/{repo_id}"
10
11config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
12config._name_or_path = model_id
13config.hidden_size = 16
14config.intermediate_size = 32
15config.num_attention_heads = 2
16config.num_key_value_heads = 1
17config.head_dim = 8
18config.num_hidden_layers = 2
19config.torch_dtype = "bfloat16"
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, trust_remote_code=True
26)
27model.generation_config = GenerationConfig.from_pretrained(
28 model_id, trust_remote_code=True)
29
30set_seed(42)
31with torch.no_grad():
32 for _, p in sorted(model.named_parameters()):
33 torch.nn.init.uniform_(p, -0.2, 0.2)
34
35model.save_pretrained(save_path)
36
37pipe = pipeline("text-generation", model=save_path, device="cpu",
38 trust_remote_code=True, max_new_tokens=20)
39print(pipe("Hello World!"))