Views
No views yet
1import torch
2
3from transformers import AutoProcessor, Llama4ForConditionalGeneration
4
5model_id = "tiny-random/llama-4"
6processor = AutoProcessor.from_pretrained(model_id)
7model = Llama4ForConditionalGeneration.from_pretrained(
8 model_id,
9 attn_implementation="sdpa", # flex attention / flash_attention_2 do not work, debugging...
10 device_map="auto",
11 torch_dtype=torch.bfloat16,
12)
13
14url1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
15url2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png"
16messages = [
17 {
18 "role": "user",
19 "content": [
20 {"type": "image", "url": url1},
21 {"type": "image", "url": url2},
22 {"type": "text", "text": "Can you describe how these two images are similar, and how they differ?"},
23 ]
24 },
25]
26
27inputs = processor.apply_chat_template(
28 messages,
29 add_generation_prompt=True,
30 tokenize=True,
31 return_dict=True,
32 return_tensors="pt",
33).to(model.device)
34
35outputs = model.generate(
36 **inputs,
37 max_new_tokens=32,
38)
39
40response = processor.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:])[0]
41print(response)
42print(outputs[0])1import json
2
3import torch
4
5from huggingface_hub import hf_hub_download
6from transformers import (
7 AutoConfig,
8 AutoModelForCausalLM,
9 AutoProcessor,
10 AutoTokenizer,
11 GenerationConfig,
12 Llama4ForConditionalGeneration,
13 pipeline,
14 set_seed,
15)
16
17source_model_id = "meta-llama/Llama-4-Maverick-17B-128E-Instruct"
18save_folder = "/tmp/tiny-random/llama-4"
19
20processor = AutoProcessor.from_pretrained(source_model_id)
21processor.save_pretrained(save_folder)
22
23with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r') as f:
24 config_json = json.load(f)
25config_json["text_config"]["num_hidden_layers"] = 4 # ensure to trigger no-rope & moe
26config_json["text_config"]["hidden_size"] = 32
27config_json["text_config"]["head_dim"] = 32 # vllm requires dim >= 32
28config_json["text_config"]["num_attention_heads"] = 1
29config_json["text_config"]["num_key_value_heads"] = 1
30config_json['text_config']["use_qk_norm"] = True
31config_json['text_config']["attention_chunk_size"] = 128 # llama4 uses chunked attention
32config_json["text_config"]["intermediate_size"] = 64
33config_json["text_config"]["intermediate_size_mlp"] = 128
34config_json["text_config"]["num_local_experts"] = 8
35config_json["text_config"]["tie_word_embeddings"] = True
36
37config_json["vision_config"]["num_hidden_layers"] = 2
38config_json["vision_config"]["hidden_size"] = 32
39config_json["vision_config"]["intermediate_size"] = 128
40assert config_json["vision_config"]["intermediate_size"] == int(
41 config_json["vision_config"]["hidden_size"] // config_json["vision_config"]["pixel_shuffle_ratio"] ** 2
42)
43config_json["vision_config"]["num_attention_heads"] = 1
44config_json["vision_config"]["projector_input_dim"] = 32
45config_json["vision_config"]["projector_output_dim"] = 32
46config_json["vision_config"]["vision_output_dim"] = 32
47with open(f"{save_folder}/config.json", "w") as f:
48 json.dump(config_json, f, indent=2)
49
50config = AutoConfig.from_pretrained(
51 save_folder,
52)
53print(config)
54torch.set_default_dtype(torch.bfloat16)
55model = Llama4ForConditionalGeneration(config)
56torch.set_default_dtype(torch.float32)
57model.generation_config = GenerationConfig.from_pretrained(
58 source_model_id, trust_remote_code=True,
59)
60set_seed(42)
61with torch.no_grad():
62 for name, p in sorted(model.named_parameters()):
63 torch.nn.init.normal_(p, 0, 0.5)
64 print(name, p.shape)
65 pass
66model.save_pretrained(save_folder)