|
|
--- |
|
|
library_name: transformers |
|
|
base_model: |
|
|
- allenai/Olmo-3-32B-Think |
|
|
--- |
|
|
|
|
|
This tiny model is intended for debugging. It is randomly initialized using the configuration adapted from [allenai/Olmo-3-32B-Think](https://huggingface.co/allenai/Olmo-3-32B-Think). |
|
|
|
|
|
### Example usage: |
|
|
|
|
|
```python |
|
|
import os |
|
|
import re |
|
|
|
|
|
import torch |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
|
model_id = "yujiepan/olmo-3-tiny-random" |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
model_id, device_map="auto", torch_dtype=torch.bfloat16) |
|
|
messages = [ |
|
|
{"role": "user", "content": "How to make pasta?" * 1500}, |
|
|
] |
|
|
inputs = tokenizer.apply_chat_template( |
|
|
messages, |
|
|
tokenize=True, |
|
|
add_generation_prompt=True, |
|
|
return_tensors="pt", |
|
|
)['input_ids'] |
|
|
print(inputs.shape) |
|
|
outputs = model.generate(inputs.to( |
|
|
model.device), max_new_tokens=32) |
|
|
print(outputs) |
|
|
``` |
|
|
|
|
|
### Codes to create this repo: |
|
|
|
|
|
```python |
|
|
import json |
|
|
from pathlib import Path |
|
|
|
|
|
import accelerate |
|
|
import torch |
|
|
from huggingface_hub import file_exists, hf_hub_download |
|
|
from transformers import ( |
|
|
AutoConfig, |
|
|
AutoModelForCausalLM, |
|
|
AutoProcessor, |
|
|
GenerationConfig, |
|
|
set_seed, |
|
|
) |
|
|
|
|
|
source_model_id = "allenai/Olmo-3-32B-Think" |
|
|
save_folder = "/tmp/yujiepan/olmo-3-tiny-random" |
|
|
|
|
|
processor = AutoProcessor.from_pretrained( |
|
|
source_model_id, trust_remote_code=True) |
|
|
processor.save_pretrained(save_folder) |
|
|
|
|
|
with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r', encoding='utf-8') as f: |
|
|
config_json = json.load(f) |
|
|
config_json['hidden_size'] = 8 |
|
|
config_json['head_dim'] = 32 # vllm requirement |
|
|
config_json['intermediate_size'] = 32 |
|
|
config_json['num_attention_heads'] = 8 |
|
|
config_json['num_hidden_layers'] = 2 |
|
|
config_json['num_key_value_heads'] = 4 # better support tensor parallel |
|
|
config_json['tie_word_embeddings'] = False |
|
|
config_json['layer_types'] = ['sliding_attention', 'full_attention'] |
|
|
with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f: |
|
|
json.dump(config_json, f, indent=2) |
|
|
|
|
|
config = AutoConfig.from_pretrained( |
|
|
save_folder, |
|
|
trust_remote_code=True, |
|
|
) |
|
|
print(config) |
|
|
torch.set_default_dtype(torch.bfloat16) |
|
|
model = AutoModelForCausalLM.from_config(config, trust_remote_code=True) |
|
|
torch.set_default_dtype(torch.float32) |
|
|
if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'): |
|
|
model.generation_config = GenerationConfig.from_pretrained( |
|
|
source_model_id, trust_remote_code=True, |
|
|
) |
|
|
model.generation_config.do_sample = True |
|
|
set_seed(42) |
|
|
model = model.cpu() # cpu is more stable for random initialization across machines |
|
|
with torch.no_grad(): |
|
|
for name, p in sorted(model.named_parameters()): |
|
|
torch.nn.init.normal_(p, 0, 0.1) |
|
|
print(name, p.shape) |
|
|
model.save_pretrained(save_folder) |
|
|
``` |
|
|
|
|
|
### Printing the model: |
|
|
|
|
|
```text |
|
|
Olmo3ForCausalLM( |
|
|
(model): Olmo3Model( |
|
|
(embed_tokens): Embedding(100278, 8, padding_idx=100277) |
|
|
(layers): ModuleList( |
|
|
(0-1): 2 x Olmo3DecoderLayer( |
|
|
(self_attn): Olmo3Attention( |
|
|
(q_proj): Linear(in_features=8, out_features=256, bias=False) |
|
|
(k_proj): Linear(in_features=8, out_features=128, bias=False) |
|
|
(v_proj): Linear(in_features=8, out_features=128, bias=False) |
|
|
(o_proj): Linear(in_features=256, out_features=8, bias=False) |
|
|
(q_norm): Olmo3RMSNorm((256,), eps=1e-06) |
|
|
(k_norm): Olmo3RMSNorm((128,), eps=1e-06) |
|
|
) |
|
|
(mlp): Olmo3MLP( |
|
|
(gate_proj): Linear(in_features=8, out_features=32, bias=False) |
|
|
(up_proj): Linear(in_features=8, out_features=32, bias=False) |
|
|
(down_proj): Linear(in_features=32, out_features=8, bias=False) |
|
|
(act_fn): SiLUActivation() |
|
|
) |
|
|
(post_attention_layernorm): Olmo3RMSNorm((8,), eps=1e-06) |
|
|
(post_feedforward_layernorm): Olmo3RMSNorm((8,), eps=1e-06) |
|
|
) |
|
|
) |
|
|
(norm): Olmo3RMSNorm((8,), eps=1e-06) |
|
|
(rotary_emb): Olmo3RotaryEmbedding() |
|
|
) |
|
|
(lm_head): Linear(in_features=8, out_features=100278, bias=False) |
|
|
) |
|
|
``` |