Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model_id = "tiny-random/smollm3"
5device = "cuda" # for GPU usage or "cpu" for CPU usage
6
7# load the tokenizer and the model
8tokenizer = AutoTokenizer.from_pretrained(model_id)
9model = AutoModelForCausalLM.from_pretrained(
10 model_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=True
11).to(device)
12
13# prepare the model input
14prompt = "Give me a brief explanation of gravity in simple terms."
15messages_think = [
16 {"role": "user", "content": prompt}
17]
18
19text = tokenizer.apply_chat_template(
20 messages_think,
21 tokenize=False,
22 add_generation_prompt=True,
23)
24model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
25
26# Generate the output
27generated_ids = model.generate(**model_inputs, max_new_tokens=200)
28
29# Get and decode the output
30output_ids = generated_ids[0][len(model_inputs.input_ids[0]):]
31print(tokenizer.decode(output_ids, skip_special_tokens=True))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 = "HuggingFaceTB/SmolLM3-3B"
16save_folder = "/tmp/tiny-random/smollm3"
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'] = 64
24config_json['intermediate_size'] = 128
25config_json['num_attention_heads'] = 2
26config_json['num_hidden_layers'] = 2
27config_json['num_key_value_heads'] = 1
28config_json['tie_word_embeddings'] = True
29config_json['layer_types'] = None
30config_json['no_rope_layer_interval'] = 2
31config_json['use_sliding_window'] = True
32config_json['sliding_window'] = 128
33config_json['use_cache'] = True
34config_json['layer_types'] = None
35config_json['no_rope_layers'] = None
36with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
37 json.dump(config_json, f, indent=2)
38
39config = AutoConfig.from_pretrained(
40 save_folder,
41 trust_remote_code=True,
42)
43print(config)
44torch.set_default_dtype(torch.bfloat16)
45model = AutoModelForCausalLM.from_config(config)
46torch.set_default_dtype(torch.float32)
47if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):
48 model.generation_config = GenerationConfig.from_pretrained(
49 source_model_id, trust_remote_code=True,
50 )
51set_seed(42)
52model = model.cpu() # cpu is more stable for random initialization across machines
53with torch.no_grad():
54 for name, p in sorted(model.named_parameters()):
55 torch.nn.init.normal_(p, 0, 0.2)
56 print(name, p.shape)
57model.save_pretrained(save_folder)
58print(model)1SmolLM3ForCausalLM(
2 (model): SmolLM3Model(
3 (embed_tokens): Embedding(128256, 64, padding_idx=128004)
4 (layers): ModuleList(
5 (0-1): 2 x SmolLM3DecoderLayer(
6 (self_attn): SmolLM3Attention(
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): SmolLM3MLP(
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): SmolLM3RMSNorm((64,), eps=1e-06)
19 (post_attention_layernorm): SmolLM3RMSNorm((64,), eps=1e-06)
20 )
21 )
22 (norm): SmolLM3RMSNorm((64,), eps=1e-06)
23 (rotary_emb): SmolLM3RotaryEmbedding()
24 )
25 (lm_head): Linear(in_features=64, out_features=128256, bias=False)
26)