| --- |
| 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 = "tiny-random/olmo-3" |
| |
| 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/tiny-random/olmo-3" |
| |
| 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) |
| ) |
| ``` |