Views
No views yet
1import json
2import os
3
4import torch
5from huggingface_hub import hf_hub_download
6
7from transformers import (
8 AutoConfig,
9 AutoProcessor,
10 Gemma4ForConditionalGeneration,
11 set_seed,
12)
13
14source_model_id = "google/gemma-4-E4B"
15save_folder = "/tmp/peft/tiny-random-gemma4"
16
17processor = AutoProcessor.from_pretrained(source_model_id)
18
19
20with open(
21 hf_hub_download(source_model_id, filename="config.json", repo_type="model"), "r", encoding="utf-8",
22) as f:
23 config_json = json.load(f)
24
25config_json["audio_config"].update(
26 {
27 "num_attention_heads": 2,
28 "num_hidden_layers": 2,
29 "hidden_size": 64,
30 "output_proj_dims": 32,
31 }
32)
33config_json["text_config"].update(
34 {
35 "global_head_dim": 64,
36 "head_dim": 32,
37 "hidden_size": 8,
38 "hidden_size_per_layer_input": 2,
39 "intermediate_size": 64,
40 "layer_types": [
41 "sliding_attention",
42 "full_attention",
43 "sliding_attention",
44 "full_attention",
45 ],
46 "num_attention_heads": 8,
47 "num_hidden_layers": 4,
48 "num_key_value_heads": 4,
49 "num_kv_shared_layers": 2,
50 }
51)
52config_json["vision_config"].update(
53 {
54 "num_hidden_layers": 2,
55 "hidden_size": 8,
56 "intermediate_size": 64,
57 "head_dim": 32,
58 "global_head_dim": 32,
59 "num_attention_heads": 4,
60 "num_key_value_heads": 4,
61 }
62)
63
64with open(f"{save_folder}/config.json", "w", encoding="utf-8") as f:
65 json.dump(config_json, f, indent=2)
66config = AutoConfig.from_pretrained(save_folder)
67
68torch.set_default_dtype(torch.bfloat16)
69model = Gemma4ForConditionalGeneration(config)
70torch.set_default_dtype(torch.float32)
71set_seed(42)
72model = model.cpu()
73
74all_numels = 0
75for name, p in sorted(model.named_parameters()):
76 all_numels += p.numel()
77with torch.no_grad():
78 for name, p in sorted(model.named_parameters()):
79 torch.nn.init.normal_(p, 0, 0.2)
80 print(name, p.shape, f"{p.numel() / all_numels * 100: .4f}%")
81
82
83token = os.environ.get("HF_TOKEN")
84processor.push_to_hub("peft-internal-testing/tiny-random-gemma4-E2B", token=token)
85model.push_to_hub("peft-internal-testing/tiny-random-gemma4-E2B", token=token)