Views
No views yet
1import os
2import re
3
4import torch
5from transformers import AutoModelForCausalLM, AutoTokenizer
6
7model_id = "yujiepan/olmo-3-tiny-random"
8
9tokenizer = AutoTokenizer.from_pretrained(model_id)
10model = AutoModelForCausalLM.from_pretrained(
11 model_id, device_map="auto", torch_dtype=torch.bfloat16)
12messages = [
13 {"role": "user", "content": "How to make pasta?" * 1500},
14]
15inputs = tokenizer.apply_chat_template(
16 messages,
17 tokenize=True,
18 add_generation_prompt=True,
19 return_tensors="pt",
20)['input_ids']
21print(inputs.shape)
22outputs = model.generate(inputs.to(
23 model.device), max_new_tokens=32)
24print(outputs)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 = "allenai/Olmo-3-32B-Think"
16save_folder = "/tmp/yujiepan/olmo-3-tiny-random"
17
18processor = AutoProcessor.from_pretrained(
19 source_model_id, trust_remote_code=True)
20processor.save_pretrained(save_folder)
21
22with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r', encoding='utf-8') as f:
23 config_json = json.load(f)
24config_json['hidden_size'] = 8
25config_json['head_dim'] = 32 # vllm requirement
26config_json['intermediate_size'] = 32
27config_json['num_attention_heads'] = 8
28config_json['num_hidden_layers'] = 2
29config_json['num_key_value_heads'] = 4 # better support tensor parallel
30config_json['tie_word_embeddings'] = False
31config_json['layer_types'] = ['sliding_attention', 'full_attention']
32with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
33 json.dump(config_json, f, indent=2)
34
35config = AutoConfig.from_pretrained(
36 save_folder,
37 trust_remote_code=True,
38)
39print(config)
40torch.set_default_dtype(torch.bfloat16)
41model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
42torch.set_default_dtype(torch.float32)
43if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):
44 model.generation_config = GenerationConfig.from_pretrained(
45 source_model_id, trust_remote_code=True,
46 )
47 model.generation_config.do_sample = True
48set_seed(42)
49model = model.cpu() # cpu is more stable for random initialization across machines
50with torch.no_grad():
51 for name, p in sorted(model.named_parameters()):
52 torch.nn.init.normal_(p, 0, 0.1)
53 print(name, p.shape)
54model.save_pretrained(save_folder)1Olmo3ForCausalLM(
2 (model): Olmo3Model(
3 (embed_tokens): Embedding(100278, 8, padding_idx=100277)
4 (layers): ModuleList(
5 (0-1): 2 x Olmo3DecoderLayer(
6 (self_attn): Olmo3Attention(
7 (q_proj): Linear(in_features=8, out_features=256, bias=False)
8 (k_proj): Linear(in_features=8, out_features=128, bias=False)
9 (v_proj): Linear(in_features=8, out_features=128, bias=False)
10 (o_proj): Linear(in_features=256, out_features=8, bias=False)
11 (q_norm): Olmo3RMSNorm((256,), eps=1e-06)
12 (k_norm): Olmo3RMSNorm((128,), eps=1e-06)
13 )
14 (mlp): Olmo3MLP(
15 (gate_proj): Linear(in_features=8, out_features=32, bias=False)
16 (up_proj): Linear(in_features=8, out_features=32, bias=False)
17 (down_proj): Linear(in_features=32, out_features=8, bias=False)
18 (act_fn): SiLUActivation()
19 )
20 (post_attention_layernorm): Olmo3RMSNorm((8,), eps=1e-06)
21 (post_feedforward_layernorm): Olmo3RMSNorm((8,), eps=1e-06)
22 )
23 )
24 (norm): Olmo3RMSNorm((8,), eps=1e-06)
25 (rotary_emb): Olmo3RotaryEmbedding()
26 )
27 (lm_head): Linear(in_features=8, out_features=100278, bias=False)
28)