Update README.md with weight comparison and hardware info
Browse files
README.md
CHANGED
@@ -1,199 +1,269 @@
|
|
1 |
---
|
2 |
-
|
3 |
-
tags:
|
|
|
|
|
|
|
|
|
|
|
4 |
---
|
5 |
|
6 |
-
#
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
-
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
|
80 |
-
|
81 |
-
|
82 |
-
|
83 |
-
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
|
117 |
-
|
118 |
-
|
119 |
-
|
120 |
-
|
121 |
-
|
122 |
-
|
123 |
-
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
|
128 |
-
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
-
|
173 |
-
|
174 |
-
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
|
187 |
-
|
188 |
-
|
189 |
-
|
190 |
-
|
191 |
-
|
192 |
-
|
193 |
-
|
194 |
-
|
195 |
-
|
196 |
-
|
197 |
-
|
198 |
-
|
199 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
---
|
2 |
+
language: en
|
3 |
+
tags:
|
4 |
+
- jax
|
5 |
+
- flax
|
6 |
+
- text-generation
|
7 |
+
- transformers
|
8 |
+
- google/gemma-2b # Add the specific model name as a tag
|
9 |
---
|
10 |
|
11 |
+
# google/gemma-2b - JAX/Flax
|
12 |
+
|
13 |
+
This repository contains the JAX/Flax version of the google/gemma-2b model, originally a PyTorch model from google. This conversion enables efficient inference and training on TPUs and GPUs using the JAX/Flax framework.
|
14 |
+
|
15 |
+
## Model Description
|
16 |
+
|
17 |
+
google/gemma-2b is a transformer-based language model developed by google.
|
18 |
+
|
19 |
+
## Conversion Details
|
20 |
+
|
21 |
+
This model was converted from the original PyTorch implementation to JAX/Flax. The conversion process involved the following steps:
|
22 |
+
|
23 |
+
1. **Loading the PyTorch model and configuration:** The pretrained PyTorch model and its configuration were loaded using the Hugging Face Transformers library.
|
24 |
+
2. **Creating an equivalent Flax model architecture:** A Flax model with the same architecture as the original PyTorch model was created.
|
25 |
+
3. **Converting the PyTorch weights to Flax format:** The weights from the PyTorch model were converted to the Flax format using the `convert_pytorch_state_dict_to_flax` utility function provided by Hugging Face.
|
26 |
+
4. **Verifying the converted weights:** The converted Flax weights were compared against the original PyTorch weights to ensure that the conversion process was performed accurately.
|
27 |
+
|
28 |
+
### Important Note about `max_position_embeddings`
|
29 |
+
|
30 |
+
During the conversion process, it was necessary to modify the `max_position_embeddings` parameter in the model's configuration. The original value of {original_max_pos_embed} led to out-of-memory (OOM) errors on the hardware used for conversion. To resolve this, `max_position_embeddings` was adjusted to {new_max_pos_embed}.
|
31 |
+
|
32 |
+
**Implications of this change:**
|
33 |
+
|
34 |
+
* The model may not be able to handle sequences longer than 8192 tokens without truncation or other modifications.
|
35 |
+
* If you fine-tune this model, keep in mind the revised `max_position_embeddings` when preparing your training data.
|
36 |
+
|
37 |
+
## Weight Comparison
|
38 |
+
|
39 |
+
The following table summarizes the comparison between the weights of the original PyTorch model and the converted JAX/Flax model. This detailed verification confirms that the conversion was accurate and that both models should produce (approximately) the same outputs given the same inputs.
|
40 |
+
|
41 |
+
| Layer | PyTorch Shape | Flax Shape | Allclose | Max Diff | Mean Diff | Std Diff |
|
42 |
+
| :---- | :------------ | :--------- | :------- | :------- | :-------- | :------- |
|
43 |
+
| model.embed_tokens.weight | (256000, 2048) | (256000, 2048) | True | 0 | 0 | 0 |
|
44 |
+
| model.layers.0.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
45 |
+
| model.layers.0.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
46 |
+
| model.layers.0.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
47 |
+
| model.layers.0.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
48 |
+
| model.layers.0.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
49 |
+
| model.layers.0.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
50 |
+
| model.layers.0.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
51 |
+
| model.layers.0.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
52 |
+
| model.layers.0.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
53 |
+
| model.layers.1.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
54 |
+
| model.layers.1.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
55 |
+
| model.layers.1.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
56 |
+
| model.layers.1.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
57 |
+
| model.layers.1.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
58 |
+
| model.layers.1.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
59 |
+
| model.layers.1.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
60 |
+
| model.layers.1.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
61 |
+
| model.layers.1.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
62 |
+
| model.layers.2.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
63 |
+
| model.layers.2.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
64 |
+
| model.layers.2.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
65 |
+
| model.layers.2.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
66 |
+
| model.layers.2.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
67 |
+
| model.layers.2.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
68 |
+
| model.layers.2.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
69 |
+
| model.layers.2.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
70 |
+
| model.layers.2.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
71 |
+
| model.layers.3.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
72 |
+
| model.layers.3.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
73 |
+
| model.layers.3.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
74 |
+
| model.layers.3.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
75 |
+
| model.layers.3.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
76 |
+
| model.layers.3.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
77 |
+
| model.layers.3.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
78 |
+
| model.layers.3.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
79 |
+
| model.layers.3.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
80 |
+
| model.layers.4.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
81 |
+
| model.layers.4.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
82 |
+
| model.layers.4.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
83 |
+
| model.layers.4.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
84 |
+
| model.layers.4.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
85 |
+
| model.layers.4.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
86 |
+
| model.layers.4.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
87 |
+
| model.layers.4.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
88 |
+
| model.layers.4.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
89 |
+
| model.layers.5.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
90 |
+
| model.layers.5.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
91 |
+
| model.layers.5.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
92 |
+
| model.layers.5.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
93 |
+
| model.layers.5.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
94 |
+
| model.layers.5.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
95 |
+
| model.layers.5.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
96 |
+
| model.layers.5.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
97 |
+
| model.layers.5.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
98 |
+
| model.layers.6.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
99 |
+
| model.layers.6.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
100 |
+
| model.layers.6.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
101 |
+
| model.layers.6.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
102 |
+
| model.layers.6.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
103 |
+
| model.layers.6.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
104 |
+
| model.layers.6.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
105 |
+
| model.layers.6.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
106 |
+
| model.layers.6.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
107 |
+
| model.layers.7.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
108 |
+
| model.layers.7.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
109 |
+
| model.layers.7.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
110 |
+
| model.layers.7.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
111 |
+
| model.layers.7.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
112 |
+
| model.layers.7.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
113 |
+
| model.layers.7.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
114 |
+
| model.layers.7.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
115 |
+
| model.layers.7.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
116 |
+
| model.layers.8.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
117 |
+
| model.layers.8.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
118 |
+
| model.layers.8.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
119 |
+
| model.layers.8.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
120 |
+
| model.layers.8.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
121 |
+
| model.layers.8.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
122 |
+
| model.layers.8.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
123 |
+
| model.layers.8.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
124 |
+
| model.layers.8.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
125 |
+
| model.layers.9.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
126 |
+
| model.layers.9.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
127 |
+
| model.layers.9.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
128 |
+
| model.layers.9.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
129 |
+
| model.layers.9.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
130 |
+
| model.layers.9.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
131 |
+
| model.layers.9.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
132 |
+
| model.layers.9.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
133 |
+
| model.layers.9.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
134 |
+
| model.layers.10.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
135 |
+
| model.layers.10.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
136 |
+
| model.layers.10.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
137 |
+
| model.layers.10.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
138 |
+
| model.layers.10.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
139 |
+
| model.layers.10.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
140 |
+
| model.layers.10.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
141 |
+
| model.layers.10.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
142 |
+
| model.layers.10.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
143 |
+
| model.layers.11.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
144 |
+
| model.layers.11.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
145 |
+
| model.layers.11.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
146 |
+
| model.layers.11.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
147 |
+
| model.layers.11.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
148 |
+
| model.layers.11.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
149 |
+
| model.layers.11.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
150 |
+
| model.layers.11.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
151 |
+
| model.layers.11.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
152 |
+
| model.layers.12.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
153 |
+
| model.layers.12.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
154 |
+
| model.layers.12.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
155 |
+
| model.layers.12.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
156 |
+
| model.layers.12.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
157 |
+
| model.layers.12.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
158 |
+
| model.layers.12.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
159 |
+
| model.layers.12.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
160 |
+
| model.layers.12.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
161 |
+
| model.layers.13.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
162 |
+
| model.layers.13.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
163 |
+
| model.layers.13.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
164 |
+
| model.layers.13.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
165 |
+
| model.layers.13.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
166 |
+
| model.layers.13.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
167 |
+
| model.layers.13.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
168 |
+
| model.layers.13.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
169 |
+
| model.layers.13.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
170 |
+
| model.layers.14.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
171 |
+
| model.layers.14.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
172 |
+
| model.layers.14.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
173 |
+
| model.layers.14.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
174 |
+
| model.layers.14.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
175 |
+
| model.layers.14.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
176 |
+
| model.layers.14.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
177 |
+
| model.layers.14.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
178 |
+
| model.layers.14.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
179 |
+
| model.layers.15.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
180 |
+
| model.layers.15.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
181 |
+
| model.layers.15.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
182 |
+
| model.layers.15.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
183 |
+
| model.layers.15.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
184 |
+
| model.layers.15.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
185 |
+
| model.layers.15.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
186 |
+
| model.layers.15.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
187 |
+
| model.layers.15.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
188 |
+
| model.layers.16.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
189 |
+
| model.layers.16.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
190 |
+
| model.layers.16.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
191 |
+
| model.layers.16.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
192 |
+
| model.layers.16.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
193 |
+
| model.layers.16.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
194 |
+
| model.layers.16.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
195 |
+
| model.layers.16.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
196 |
+
| model.layers.16.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
197 |
+
| model.layers.17.self_attn.q_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
198 |
+
| model.layers.17.self_attn.k_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
199 |
+
| model.layers.17.self_attn.v_proj.weight | (2048, 256) | (2048, 256) | True | 0 | 0 | 0 |
|
200 |
+
| model.layers.17.self_attn.o_proj.weight | (2048, 2048) | (2048, 2048) | True | 0 | 0 | 0 |
|
201 |
+
| model.layers.17.mlp.gate_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
202 |
+
| model.layers.17.mlp.up_proj.weight | (2048, 16384) | (2048, 16384) | True | 0 | 0 | 0 |
|
203 |
+
| model.layers.17.mlp.down_proj.weight | (16384, 2048) | (16384, 2048) | True | 0 | 0 | 0 |
|
204 |
+
| model.layers.17.input_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
205 |
+
| model.layers.17.post_attention_layernorm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
206 |
+
| model.norm.weight | (2048,) | (2048,) | True | 0 | 0 | 0 |
|
207 |
+
| lm_head.weight | (2048, 256000) | (2048, 256000) | True | 0 | 0 | 0 |
|
208 |
+
|
209 |
+
**Note:**
|
210 |
+
|
211 |
+
* `Allclose` indicates whether the weights are approximately equal within the specified relative (`rtol=1e-5`) and absolute (`atol=1e-3`) tolerances using `jnp.allclose()`.
|
212 |
+
* `Max Diff`, `Mean Diff`, and `Std Diff` provide further details on the differences between the weights if `Allclose` is `False`, which might be expected for some layers due to numerical precision differences between frameworks.
|
213 |
+
|
214 |
+
## Hardware Used for Conversion
|
215 |
+
|
216 |
+
The conversion process was performed on the following hardware configuration:
|
217 |
+
|
218 |
+
* **CPU:**
|
219 |
+
* **RAM:** 251.67 GB
|
220 |
+
* **OS:** Linux-5.15.0-107-generic-x86_64-with-glibc2.36
|
221 |
+
* **JAX version:** 0.3.22
|
222 |
+
* **Flax version:** 0.6.2
|
223 |
+
* **Transformers version:** 4.47.0
|
224 |
+
* **GPU:** NVIDIA A100-SXM4-40GB
|
225 |
+
|
226 |
+
This conversion took approximately 184.13 seconds to complete.
|
227 |
+
|
228 |
+
## Usage
|
229 |
+
|
230 |
+
Here's how you can use the converted model in JAX/Flax for text generation:
|
231 |
+
|
232 |
+
```python
|
233 |
+
import jax
|
234 |
+
import jax.numpy as jnp
|
235 |
+
from transformers import FlaxAutoModelForCausalLM, AutoTokenizer
|
236 |
+
|
237 |
+
model_name = "Erland/gemma-2b-JAX" # Replace with your repository name
|
238 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
239 |
+
model = FlaxAutoModelForCausalLM.from_pretrained(model_name, from_pt=False) # from_pt should be False since it's already flax
|
240 |
+
|
241 |
+
# Example prompt
|
242 |
+
prompt = "The quick brown fox"
|
243 |
+
|
244 |
+
# Tokenize the prompt
|
245 |
+
tokenized_prompt = tokenizer(prompt, return_tensors="np")
|
246 |
+
|
247 |
+
# Generate text
|
248 |
+
output_ids = model.generate(tokenized_prompt.input_ids, max_length=50)
|
249 |
+
|
250 |
+
# Decode the generated text
|
251 |
+
generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
|
252 |
+
```
|
253 |
+
## Limitations
|
254 |
+
|
255 |
+
Sequence Length: As mentioned earlier, the max_position_embeddings has been modified to 8192. Be mindful of this limitation when working with long sequences.
|
256 |
+
|
257 |
+
Numerical Precision: Minor differences in outputs compared to the original PyTorch model might be observed due to numerical precision variations between PyTorch and JAX/Flax, particularly on different hardware.
|
258 |
+
|
259 |
+
## Acknowledgements
|
260 |
+
|
261 |
+
We thank the original authors of google/gemma-2b at `google` for their groundbreaking work in developing this powerful language model.
|
262 |
+
|
263 |
+
We acknowledge the Hugging Face Transformers library for providing the essential tools and infrastructure that made this conversion possible.
|
264 |
+
|
265 |
+
Thanks to the JAX and Flax teams for developing such performant and flexible frameworks for numerical computation and deep learning.
|
266 |
+
|
267 |
+
## License
|
268 |
+
|
269 |
+
This JAX/Flax model is released under the original model license.
|