tiny-random/lfm2-moe

1---2library_name: transformers3pipeline_tag: text-generation4inference: true5widget:6  - text: Hello!7    example_title: Hello world8    group: Python9base_model:10- LiquidAI/LFM2-8B-A1B11---12 13This tiny model is intended for debugging. It is randomly initialized using the configuration adapted from [LiquidAI/LFM2-8B-A1B](https://huggingface.co/LiquidAI/LFM2-8B-A1B).14 15### Example usage:16 17```python18from transformers import AutoModelForCausalLM, AutoTokenizer19 20# Load model and tokenizer21model_id = "tiny-random/lfm2-moe"22model = AutoModelForCausalLM.from_pretrained(23    model_id,24    device_map="cuda",25    dtype="bfloat16",26    trust_remote_code=True,27    attn_implementation="flash_attention_2",28)29tokenizer = AutoTokenizer.from_pretrained(model_id)30 31# Generate answer32prompt="What is AI?"33input_ids=tokenizer.apply_chat_template(34    [{"role": "user", "content": prompt}],35    add_generation_prompt=True,36    return_tensors="pt",37    tokenize=True,38).to(model.device)39 40output=model.generate(41    input_ids,42    do_sample=True,43    temperature=0.3,44    min_p=0.15,45    repetition_penalty=1.05,46    max_new_tokens=32,47)48 49print(tokenizer.decode(output[0], skip_special_tokens=False))50```51 52### Codes to create this repo:53 54```python55import json56from pathlib import Path57 58import accelerate59import torch60from huggingface_hub import file_exists, hf_hub_download61from transformers import (62    AutoConfig,63    AutoModelForCausalLM,64    AutoProcessor,65    GenerationConfig,66    set_seed,67)68 69source_model_id = "LiquidAI/LFM2-8B-A1B"70save_folder = "/tmp/tiny-random/lfm2-moe"71 72processor = AutoProcessor.from_pretrained(source_model_id, trust_remote_code=True)73processor.save_pretrained(save_folder)74 75with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r', encoding='utf-8') as f:76    config_json = json.load(f)77config_json['hidden_size'] = 6478config_json['intermediate_size'] = 12879config_json['layer_types'] = ['conv', 'conv', 'full_attention']80config_json['moe_intermediate_size'] = 12881config_json['num_dense_layers'] = 282config_json['num_attention_heads'] = 283config_json['num_hidden_layers'] = 384config_json['num_key_value_heads'] = 185config_json['use_cache'] = True86# config_json['tie_word_embeddings'] = True87with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:88    json.dump(config_json, f, indent=2)89 90config = AutoConfig.from_pretrained(91    save_folder,92    trust_remote_code=True,93)94print(config)95torch.set_default_dtype(torch.bfloat16)96model = AutoModelForCausalLM.from_config(config)97torch.set_default_dtype(torch.float32)98if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):99    model.generation_config = GenerationConfig.from_pretrained(100        source_model_id, trust_remote_code=True,101    )102set_seed(42)103model = model.cpu()  # cpu is more stable for random initialization across machines104with torch.no_grad():105    for name, p in sorted(model.named_parameters()):106        torch.nn.init.normal_(p, 0, 0.1)107        print(name, p.shape)108model.save_pretrained(save_folder)109print(model)110```111 112### Printing the model:113 114```text115Lfm2MoeForCausalLM(116  (model): Lfm2MoeModel(117    (embed_tokens): Embedding(65536, 64, padding_idx=0)118    (layers): ModuleList(119      (0-1): 2 x Lfm2MoeDecoderLayer(120        (conv): Lfm2MoeShortConv(121          (conv): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(2,), groups=64, bias=False)122          (in_proj): Linear(in_features=64, out_features=192, bias=False)123          (out_proj): Linear(in_features=64, out_features=64, bias=False)124        )125        (feed_forward): Lfm2MoeMLP(126          (w1): Linear(in_features=64, out_features=128, bias=False)127          (w3): Linear(in_features=64, out_features=128, bias=False)128          (w2): Linear(in_features=128, out_features=64, bias=False)129        )130        (operator_norm): Lfm2MoeRMSNorm((64,), eps=1e-05)131        (ffn_norm): Lfm2MoeRMSNorm((64,), eps=1e-05)132      )133      (2): Lfm2MoeDecoderLayer(134        (self_attn): Lfm2MoeAttention(135          (q_proj): Linear(in_features=64, out_features=64, bias=False)136          (k_proj): Linear(in_features=64, out_features=32, bias=False)137          (v_proj): Linear(in_features=64, out_features=32, bias=False)138          (out_proj): Linear(in_features=64, out_features=64, bias=False)139          (q_layernorm): Lfm2MoeRMSNorm((32,), eps=1e-05)140          (k_layernorm): Lfm2MoeRMSNorm((32,), eps=1e-05)141        )142        (feed_forward): Lfm2MoeSparseMoeBlock(143          (gate): Linear(in_features=64, out_features=32, bias=False)144          (experts): Lfm2MoeExperts(145            (0-31): 32 x Lfm2MoeMLP(146              (w1): Linear(in_features=64, out_features=128, bias=False)147              (w3): Linear(in_features=64, out_features=128, bias=False)148              (w2): Linear(in_features=128, out_features=64, bias=False)149            )150          )151        )152        (operator_norm): Lfm2MoeRMSNorm((64,), eps=1e-05)153        (ffn_norm): Lfm2MoeRMSNorm((64,), eps=1e-05)154      )155    )156    (pos_emb): Lfm2MoeRotaryEmbedding()157    (embedding_norm): Lfm2MoeRMSNorm((64,), eps=1e-05)158  )159  (lm_head): Linear(in_features=64, out_features=65536, bias=False)160)161```