Views
No views yet
| Model ID | Notes |
|---|---|
| tiny-random/gemma-3n | hidden size is 32 |
| tiny-random/gemma-3n-dim4 | hidden size is 4; potentially not supported in paged attention kernels |
1import torch
2
3from transformers import pipeline
4
5model_id = "tiny-random/gemma-3n-dim4"
6pipe = pipeline(
7 task="image-text-to-text",
8 model=model_id,
9 device=0,
10 torch_dtype=torch.bfloat16
11)
12
13# temporary patch for audio tower
14from accelerate.hooks import ModelHook, add_hook_to_module
15
16class EnsureDtype(ModelHook):
17 def pre_forward(self, module, *args, **kwargs):
18 args = list(args)
19 args[0] = args[0].to(module.dtype)
20 return super().pre_forward(module, *args, **kwargs)
21add_hook_to_module(pipe.model.audio_tower, EnsureDtype())
22
23messages = [
24 {
25 "role": "system",
26 "content": [
27 {"type": "text", "text": "You are a helpful assistant."}
28 ]
29 },
30 {
31 "role": "user",
32 "content": [
33 {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"},
34 # audio is buggy for now: bf16 x fp32
35 {"type": "audio", "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Audio/glass-breaking-151256.mp3"},
36 {"type": "text", "text": "Which image is cuter?"},
37 ]
38 },
39]
40result = pipe(messages, min_new_tokens=512, max_new_tokens=512, do_sample=True)
41print(result)1import json
2from pathlib import Path
3
4import torch
5
6import accelerate
7from huggingface_hub import file_exists, hf_hub_download
8from timm.models.mobilenetv5 import decode_arch_def
9from transformers import (
10 AutoConfig,
11 AutoModelForCausalLM,
12 AutoProcessor,
13 AutoTokenizer,
14 Gemma3nForConditionalGeneration,
15 GenerationConfig,
16 set_seed,
17)
18
19source_model_id = "google/gemma-3n-E4B-it"
20save_folder = "/tmp/tiny-random/gemma-3n-dim4"
21
22processor = AutoProcessor.from_pretrained(source_model_id)
23processor.save_pretrained(save_folder)
24
25with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r', encoding='utf-8') as f:
26 config_json = json.load(f)
27
28config_json['audio_config'].update({
29 "conf_num_attention_heads": 2,
30 "conf_num_hidden_layers": 2,
31 "hidden_size": 4,
32})
33config_json['text_config'].update({
34 "activation_sparsity_pattern": [0.95, 0.95, 0.0, 0.0],
35 "head_dim": 2,
36 "hidden_size": 4,
37 "hidden_size_per_layer_input": 1,
38 "intermediate_size": 8,
39 "laurel_rank": 1,
40 "layer_types": ['sliding_attention', 'full_attention', 'sliding_attention', 'full_attention'],
41 "num_attention_heads": 2,
42 "num_hidden_layers": 4,
43 "num_key_value_heads": 1,
44 "num_kv_shared_layers": 2,
45 "sliding_window": 512,
46})
47block_args = decode_arch_def(
48 [
49 # Stage 0: 128x128 in
50 [
51 'er_r1_k3_s2_e4_c4',
52 'er_r1_k3_s1_e4_c4',
53 ],
54 # Stage 1: 256x256 in
55 [
56 'uir_r1_a3_k5_s2_e6_c4',
57 'uir_r1_a5_k0_s1_e4_c4',
58 'uir_r1_a3_k0_s1_e4_c4',
59 ],
60 # Stage 2: 640x640 in
61 [
62 "uir_r1_a5_k5_s2_e6_c4",
63 "uir_r1_a0_k0_s1_e1_c4",
64 "mqa_r1_k3_h2_v2_s1_d8_c4",
65 "uir_r1_a0_k0_s1_e2_c4",
66 ],
67 # Stage 3: 1280x1280 in
68 [
69 "uir_r1_a5_k5_s2_e6_c4",
70 "mqa_r1_k3_h2_s1_d8_c4",
71 "uir_r1_a0_k0_s1_e2_c4",
72 ],
73 ]
74)
75config_json['vision_config'].update({
76 "hidden_size": 2048, # hard-coded in timm
77 "model_args": {
78 "block_args": block_args,
79 }
80})
81config_json['tie_word_embeddings'] = True
82
83with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
84 json.dump(config_json, f, indent=2)
85
86config = AutoConfig.from_pretrained(
87 save_folder,
88 trust_remote_code=True,
89)
90print(config)
91
92torch.set_default_dtype(torch.bfloat16)
93model = Gemma3nForConditionalGeneration(config)
94torch.set_default_dtype(torch.float32)
95if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):
96 model.generation_config = GenerationConfig.from_pretrained(
97 source_model_id, trust_remote_code=True,
98 )
99set_seed(42)
100model = model.cpu()
101all_numels = 0
102for name, p in sorted(model.named_parameters()):
103 all_numels += p.numel()
104with torch.no_grad():
105 for name, p in sorted(model.named_parameters()):
106 torch.nn.init.normal_(p, 0, 0.2)
107 print(name, p.shape, f'{p.numel() / all_numels * 100: .4f}%')
108model.save_pretrained(save_folder)1Gemma3nForConditionalGeneration(
2 (model): Gemma3nModel(
3 (vision_tower): TimmWrapperModel(
4 (timm_model): MobileNetV5Encoder(
5 (conv_stem): ConvNormAct(
6 (conv): Conv2dSame(3, 64, kernel_size=(3, 3), stride=(2, 2), bias=False)
7 (bn): RmsNormAct2d(
8 (drop): Identity()
9 (act): GELU(approximate='none')
10 )
11 )
12 (blocks): Sequential(
13 (0): Sequential(
14 (0): EdgeResidual(
15 (conv_exp): Conv2dSame(64, 256, kernel_size=(3, 3), stride=(2, 2), bias=False)
16 (bn1): RmsNormAct2d(
17 (drop): Identity()
18 (act): GELU(approximate='none')
19 )
20 (aa): Identity()
21 (se): Identity()
22 (conv_pwl): Conv2d(256, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
23 (bn2): RmsNormAct2d(
24 (drop): Identity()
25 (act): Identity()
26 )
27 (drop_path): Identity()
28 )
29 (1): EdgeResidual(
30 (conv_exp): Conv2d(8, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
31 (bn1): RmsNormAct2d(
32 (drop): Identity()
33 (act): GELU(approximate='none')
34 )
35 (aa): Identity()
36 (se): Identity()
37 (conv_pwl): Conv2d(32, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
38 (bn2): RmsNormAct2d(
39 (drop): Identity()
40 (act): Identity()
41 )
42 (drop_path): Identity()
43 )
44 )
45 (1): Sequential(
46 (0): UniversalInvertedResidual(
47 (dw_start): ConvNormAct(
48 (conv): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=8, bias=False)
49 (bn): RmsNormAct2d(
50 (drop): Identity()
51 (act): Identity()
52 )
53 )
54 (pw_exp): ConvNormAct(
55 (conv): Conv2d(8, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)
56 (bn): RmsNormAct2d(
57 (drop): Identity()
58 (act): GELU(approximate='none')
59 )
60 )
61 (dw_mid): ConvNormAct(
62 (conv): Conv2dSame(48, 48, kernel_size=(5, 5), stride=(2, 2), groups=48, bias=False)
63 (bn): RmsNormAct2d(
64 (drop): Identity()
65 (act): GELU(approximate='none')
66 )
67 )
68 (se): Identity()
69 (pw_proj): ConvNormAct(
70 (conv): Conv2d(48, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
71 (bn): RmsNormAct2d(
72 (drop): Identity()
73 (act): Identity()
74 )
75 )
76 (dw_end): Identity()
77 (layer_scale): LayerScale2d()
78 (drop_path): Identity()
79 )
80 (1): UniversalInvertedResidual(
81 (dw_start): ConvNormAct(
82 (conv): Conv2d(8, 8, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=8, bias=False)
83 (bn): RmsNormAct2d(
84 (drop): Identity()
85 (act): Identity()
86 )
87 )
88 (pw_exp): ConvNormAct(
89 (conv): Conv2d(8, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)
90 (bn): RmsNormAct2d(
91 (drop): Identity()
92 (act): GELU(approximate='none')
93 )
94 )
95 (dw_mid): Identity()
96 (se): Identity()
97 (pw_proj): ConvNormAct(
98 (conv): Conv2d(32, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
99 (bn): RmsNormAct2d(
100 (drop): Identity()
101 (act): Identity()
102 )
103 )
104 (dw_end): Identity()
105 (layer_scale): LayerScale2d()
106 (drop_path): Identity()
107 )
108 (2): UniversalInvertedResidual(
109 (dw_start): ConvNormAct(
110 (conv): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=8, bias=False)
111 (bn): RmsNormAct2d(
112 (drop): Identity()
113 (act): Identity()
114 )
115 )
116 (pw_exp): ConvNormAct(
117 (conv): Conv2d(8, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)
118 (bn): RmsNormAct2d(
119 (drop): Identity()
120 (act): GELU(approximate='none')
121 )
122 )
123 (dw_mid): Identity()
124 (se): Identity()
125 (pw_proj): ConvNormAct(
126 (conv): Conv2d(32, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
127 (bn): RmsNormAct2d(
128 (drop): Identity()
129 (act): Identity()
130 )
131 )
132 (dw_end): Identity()
133 (layer_scale): LayerScale2d()
134 (drop_path): Identity()
135 )
136 )
137 (2): Sequential(
138 (0): UniversalInvertedResidual(
139 (dw_start): ConvNormAct(
140 (conv): Conv2d(8, 8, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=8, bias=False)
141 (bn): RmsNormAct2d(
142 (drop): Identity()
143 (act): Identity()
144 )
145 )
146 (pw_exp): ConvNormAct(
147 (conv): Conv2d(8, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)
148 (bn): RmsNormAct2d(
149 (drop): Identity()
150 (act): GELU(approximate='none')
151 )
152 )
153 (dw_mid): ConvNormAct(
154 (conv): Conv2dSame(48, 48, kernel_size=(5, 5), stride=(2, 2), groups=48, bias=False)
155 (bn): RmsNormAct2d(
156 (drop): Identity()
157 (act): GELU(approximate='none')
158 )
159 )
160 (se): Identity()
161 (pw_proj): ConvNormAct(
162 (conv): Conv2d(48, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
163 (bn): RmsNormAct2d(
164 (drop): Identity()
165 (act): Identity()
166 )
167 )
168 (dw_end): Identity()
169 (layer_scale): LayerScale2d()
170 (drop_path): Identity()
171 )
172 (1): UniversalInvertedResidual(
173 (dw_start): Identity()
174 (pw_exp): ConvNormAct(
175 (conv): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
176 (bn): RmsNormAct2d(
177 (drop): Identity()
178 (act): GELU(approximate='none')
179 )
180 )
181 (dw_mid): Identity()
182 (se): Identity()
183 (pw_proj): ConvNormAct(
184 (conv): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
185 (bn): RmsNormAct2d(
186 (drop): Identity()
187 (act): Identity()
188 )
189 )
190 (dw_end): Identity()
191 (layer_scale): LayerScale2d()
192 (drop_path): Identity()
193 )
194 (2): MobileAttention(
195 (norm): RmsNormAct2d(
196 (drop): Identity()
197 (act): Identity()
198 )
199 (attn): MultiQueryAttention2d(
200 (query): Sequential(
201 (proj): Conv2d(8, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
202 )
203 (key): Sequential(
204 (down_conv): Conv2dSame(8, 8, kernel_size=(3, 3), stride=(2, 2), groups=8, bias=False)
205 (norm): RmsNorm2d()
206 (proj): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
207 )
208 (value): Sequential(
209 (down_conv): Conv2dSame(8, 8, kernel_size=(3, 3), stride=(2, 2), groups=8, bias=False)
210 (norm): RmsNorm2d()
211 (proj): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
212 )
213 (attn_drop): Dropout(p=0.0, inplace=False)
214 (output): Sequential(
215 (proj): Conv2d(16, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
216 (drop): Dropout(p=0.0, inplace=False)
217 )
218 )
219 (layer_scale): LayerScale2d()
220 (drop_path): Identity()
221 )
222 (3): UniversalInvertedResidual(
223 (dw_start): Identity()
224 (pw_exp): ConvNormAct(
225 (conv): Conv2d(8, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
226 (bn): RmsNormAct2d(
227 (drop): Identity()
228 (act): GELU(approximate='none')
229 )
230 )
231 (dw_mid): Identity()
232 (se): Identity()
233 (pw_proj): ConvNormAct(
234 (conv): Conv2d(16, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
235 (bn): RmsNormAct2d(
236 (drop): Identity()
237 (act): Identity()
238 )
239 )
240 (dw_end): Identity()
241 (layer_scale): LayerScale2d()
242 (drop_path): Identity()
243 )
244 )
245 (3): Sequential(
246 (0): UniversalInvertedResidual(
247 (dw_start): ConvNormAct(
248 (conv): Conv2d(8, 8, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=8, bias=False)
249 (bn): RmsNormAct2d(
250 (drop): Identity()
251 (act): Identity()
252 )
253 )
254 (pw_exp): ConvNormAct(
255 (conv): Conv2d(8, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)
256 (bn): RmsNormAct2d(
257 (drop): Identity()
258 (act): GELU(approximate='none')
259 )
260 )
261 (dw_mid): ConvNormAct(
262 (conv): Conv2dSame(48, 48, kernel_size=(5, 5), stride=(2, 2), groups=48, bias=False)
263 (bn): RmsNormAct2d(
264 (drop): Identity()
265 (act): GELU(approximate='none')
266 )
267 )
268 (se): Identity()
269 (pw_proj): ConvNormAct(
270 (conv): Conv2d(48, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
271 (bn): RmsNormAct2d(
272 (drop): Identity()
273 (act): Identity()
274 )
275 )
276 (dw_end): Identity()
277 (layer_scale): LayerScale2d()
278 (drop_path): Identity()
279 )
280 (1): MobileAttention(
281 (norm): RmsNormAct2d(
282 (drop): Identity()
283 (act): Identity()
284 )
285 (attn): MultiQueryAttention2d(
286 (query): Sequential(
287 (proj): Conv2d(8, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
288 )
289 (key): Sequential(
290 (proj): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
291 )
292 (value): Sequential(
293 (proj): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
294 )
295 (attn_drop): Dropout(p=0.0, inplace=False)
296 (output): Sequential(
297 (proj): Conv2d(16, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
298 (drop): Dropout(p=0.0, inplace=False)
299 )
300 )
301 (layer_scale): LayerScale2d()
302 (drop_path): Identity()
303 )
304 (2): UniversalInvertedResidual(
305 (dw_start): Identity()
306 (pw_exp): ConvNormAct(
307 (conv): Conv2d(8, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
308 (bn): RmsNormAct2d(
309 (drop): Identity()
310 (act): GELU(approximate='none')
311 )
312 )
313 (dw_mid): Identity()
314 (se): Identity()
315 (pw_proj): ConvNormAct(
316 (conv): Conv2d(16, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
317 (bn): RmsNormAct2d(
318 (drop): Identity()
319 (act): Identity()
320 )
321 )
322 (dw_end): Identity()
323 (layer_scale): LayerScale2d()
324 (drop_path): Identity()
325 )
326 )
327 )
328 (msfa): MobileNetV5MultiScaleFusionAdapter(
329 (ffn): UniversalInvertedResidual(
330 (dw_start): Identity()
331 (pw_exp): ConvNormAct(
332 (conv): Conv2d(16, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)
333 (bn): RmsNormAct2d(
334 (drop): Identity()
335 (act): GELU(approximate='none')
336 )
337 )
338 (dw_mid): Identity()
339 (se): Identity()
340 (pw_proj): ConvNormAct(
341 (conv): Conv2d(32, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
342 (bn): RmsNormAct2d(
343 (drop): Identity()
344 (act): Identity()
345 )
346 )
347 (dw_end): Identity()
348 (layer_scale): Identity()
349 (drop_path): Identity()
350 )
351 (norm): RmsNorm2d()
352 )
353 )
354 )
355 (language_model): Gemma3nTextModel(
356 (embed_tokens): Gemma3nTextScaledWordEmbedding(262400, 4, padding_idx=0)
357 (layers): ModuleList(
358 (0-3): 4 x Gemma3nTextDecoderLayer(
359 (self_attn): Gemma3nTextAttention(
360 (q_proj): Linear(in_features=4, out_features=4, bias=False)
361 (k_proj): Linear(in_features=4, out_features=2, bias=False)
362 (v_proj): Linear(in_features=4, out_features=2, bias=False)
363 (o_proj): Linear(in_features=4, out_features=4, bias=False)
364 (q_norm): Gemma3nRMSNorm((2,), eps=1e-06)
365 (k_norm): Gemma3nRMSNorm((2,), eps=1e-06)
366 (v_norm): Gemma3nRMSNorm((), eps=1e-06)
367 )
368 (mlp): Gemma3nTextMLP(
369 (gate_proj): Linear(in_features=4, out_features=8, bias=False)
370 (up_proj): Linear(in_features=4, out_features=8, bias=False)
371 (down_proj): Linear(in_features=8, out_features=4, bias=False)
372 (act_fn): PytorchGELUTanh()
373 )
374 (input_layernorm): Gemma3nRMSNorm((4,), eps=1e-06)
375 (post_attention_layernorm): Gemma3nRMSNorm((4,), eps=1e-06)
376 (pre_feedforward_layernorm): Gemma3nRMSNorm((4,), eps=1e-06)
377 (post_feedforward_layernorm): Gemma3nRMSNorm((4,), eps=1e-06)
378 (act_fn): PytorchGELUTanh()
379 (altup): Gemma3nTextAltUp(
380 (correction_coefs): Linear(in_features=4, out_features=4, bias=False)
381 (prediction_coefs): Linear(in_features=4, out_features=16, bias=False)
382 (modality_router): Linear(in_features=4, out_features=4, bias=False)
383 (router_norm): Gemma3nRMSNorm((4,), eps=1e-06)
384 )
385 (laurel): Gemma3nTextLaurelBlock(
386 (linear_left): Linear(in_features=4, out_features=1, bias=False)
387 (linear_right): Linear(in_features=1, out_features=4, bias=False)
388 (post_laurel_norm): Gemma3nRMSNorm((4,), eps=1e-06)
389 )
390 (per_layer_input_gate): Linear(in_features=4, out_features=1, bias=False)
391 (per_layer_projection): Linear(in_features=1, out_features=4, bias=False)
392 (post_per_layer_input_norm): Gemma3nRMSNorm((4,), eps=1e-06)
393 )
394 )
395 (norm): Gemma3nRMSNorm((4,), eps=1e-06)
396 (rotary_emb): Gemma3nTextRotaryEmbedding()
397 (rotary_emb_local): Gemma3nTextRotaryEmbedding()
398 (embed_tokens_per_layer): Gemma3nTextScaledWordEmbedding(262144, 4, padding_idx=0)
399 (per_layer_model_projection): Linear(in_features=4, out_features=4, bias=False)
400 (per_layer_projection_norm): Gemma3nRMSNorm((1,), eps=1e-06)
401 (altup_projections): ModuleList(
402 (0-2): 3 x Linear(in_features=4, out_features=4, bias=False)
403 )
404 (altup_unembed_projections): ModuleList(
405 (0-2): 3 x Linear(in_features=4, out_features=4, bias=False)
406 )
407 )
408 (audio_tower): Gemma3nAudioEncoder(
409 (subsample_conv_projection): Gemma3nAudioSubSampleConvProjection(
410 (conv_0): Gemma3nAudioSSCPConvBlock(
411 (conv): Conv2d(1, 128, kernel_size=(3, 3), stride=(2, 2), bias=False)
412 (norm): Gemma3nAudioCumulativeGroupNorm()
413 (activation): ReLU()
414 )
415 (conv_1): Gemma3nAudioSSCPConvBlock(
416 (conv): Conv2d(128, 32, kernel_size=(3, 3), stride=(2, 2), bias=False)
417 (norm): Gemma3nAudioCumulativeGroupNorm()
418 (activation): ReLU()
419 )
420 (input_proj_linear): Linear(in_features=1024, out_features=4, bias=False)
421 )
422 (conformer): ModuleList(
423 (0-1): 2 x Gemma3nAudioConformerBlock(
424 (ffw_layer_start): Gemma3nAudioConformerFeedForward(
425 (pre_layer_norm): Gemma3nRMSNorm((4,), eps=1e-06)
426 (ffw_layer_1): Linear(in_features=4, out_features=16, bias=False)
427 (ffw_layer_2): Linear(in_features=16, out_features=4, bias=False)
428 (post_layer_norm): Gemma3nRMSNorm((4,), eps=1e-06)
429 )
430 (attention): Gemma3nAudioConformerAttention(
431 (pre_attn_norm): Gemma3nRMSNorm((4,), eps=1e-06)
432 (attn): Gemma3nAudioAttention(
433 (relative_position_embedding): Gemma3nAudioRelativePositionEmbedding(
434 (pos_proj): Linear(in_features=4, out_features=4, bias=False)
435 )
436 (q_proj): Linear(in_features=4, out_features=4, bias=False)
437 (k_proj): Linear(in_features=4, out_features=4, bias=False)
438 (v_proj): Linear(in_features=4, out_features=4, bias=False)
439 )
440 (post): Linear(in_features=4, out_features=4, bias=False)
441 (post_norm): Gemma3nRMSNorm((4,), eps=1e-06)
442 )
443 (lconv1d): Gemma3nAudioConformerLightConv1d(
444 (pre_layer_norm): Gemma3nRMSNorm((4,), eps=1e-06)
445 (linear_start): Linear(in_features=4, out_features=8, bias=False)
446 (depthwise_conv1d): Conv1d(4, 4, kernel_size=(5,), stride=(1,), groups=4, bias=False)
447 (conv_norm): Gemma3nRMSNorm((4,), eps=1e-06)
448 (linear_end): Linear(in_features=4, out_features=4, bias=False)
449 )
450 (ffw_layer_end): Gemma3nAudioConformerFeedForward(
451 (pre_layer_norm): Gemma3nRMSNorm((4,), eps=1e-06)
452 (ffw_layer_1): Linear(in_features=4, out_features=16, bias=False)
453 (ffw_layer_2): Linear(in_features=16, out_features=4, bias=False)
454 (post_layer_norm): Gemma3nRMSNorm((4,), eps=1e-06)
455 )
456 (norm): Gemma3nRMSNorm((4,), eps=1e-06)
457 )
458 )
459 )
460 (embed_vision): Gemma3nMultimodalEmbedder(
461 (embedding): Embedding(128, 2048)
462 (hard_embedding_norm): Gemma3nRMSNorm((2048,), eps=1e-06)
463 (soft_embedding_norm): Gemma3nRMSNorm((2048,), eps=1e-06)
464 (embedding_projection): Linear(in_features=2048, out_features=4, bias=False)
465 (embedding_post_projection_norm): Gemma3nRMSNorm((), eps=1e-06)
466 )
467 (embed_audio): Gemma3nMultimodalEmbedder(
468 (embedding): Embedding(128, 4)
469 (hard_embedding_norm): Gemma3nRMSNorm((4,), eps=1e-06)
470 (soft_embedding_norm): Gemma3nRMSNorm((4,), eps=1e-06)
471 (embedding_projection): Linear(in_features=4, out_features=4, bias=False)
472 (embedding_post_projection_norm): Gemma3nRMSNorm((), eps=1e-06)
473 )
474 )
475 (lm_head): Linear(in_features=4, out_features=262400, bias=False)
476)