(self, spec, module, quant_type=common_spec.Quantization.CT2)
| 1990 | spec.layer_norm_use_residual = True |
| 1991 | |
| 1992 | def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2): |
| 1993 | spec.scale_embeddings = True |
| 1994 | spec.start_from_zero_embedding = False |
| 1995 | self.set_embeddings(spec.embeddings, module.embed_tokens) # Input |
| 1996 | self.set_layer_norm(spec.layer_norm, module.norm) # Output |
| 1997 | |
| 1998 | for layer_spec, layer in zip(spec.layer, module.layers): |
| 1999 | self.set_layer_norm(layer_spec.input_layer_norm, layer.input_layernorm) |
| 2000 | |
| 2001 | self.set_layer_norm( |
| 2002 | layer_spec.post_attention_layer_norm, layer.post_attention_layernorm |
| 2003 | ) |
| 2004 | |
| 2005 | self.set_layer_norm( |
| 2006 | layer_spec.pre_feedforward_layer_norm, layer.pre_feedforward_layernorm |
| 2007 | ) |
| 2008 | |
| 2009 | self.set_layer_norm( |
| 2010 | layer_spec.post_feedforward_layer_norm, layer.post_feedforward_layernorm |
| 2011 | ) |
| 2012 | |
| 2013 | # Set QK-norm weights (Gemma 3 uses this instead of soft-capping) |
| 2014 | self.set_layer_norm( |
| 2015 | layer_spec.self_attention.q_norm, layer.self_attn.q_norm |
| 2016 | ) |
| 2017 | self.set_layer_norm( |
| 2018 | layer_spec.self_attention.k_norm, layer.self_attn.k_norm |
| 2019 | ) |
| 2020 | |
| 2021 | # Set attention projections |
| 2022 | split_layers = [common_spec.LinearSpec() for _ in range(3)] |
| 2023 | self.set_linear( |
| 2024 | split_layers[0], layer.self_attn.q_proj, quant_type=quant_type |
| 2025 | ) |
| 2026 | self.set_linear( |
| 2027 | split_layers[1], layer.self_attn.k_proj, quant_type=quant_type |
| 2028 | ) |
| 2029 | self.set_linear( |
| 2030 | split_layers[2], layer.self_attn.v_proj, quant_type=quant_type |
| 2031 | ) |
| 2032 | |
| 2033 | if quant_type == common_spec.Quantization.CT2: |
| 2034 | utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) |
| 2035 | else: |
| 2036 | cc_dim = 1 if quant_type == common_spec.Quantization.AWQ_GEMM else 0 |
| 2037 | utils.fuse_linear_prequant( |
| 2038 | layer_spec.self_attention.linear[0], split_layers, cc_dim |
| 2039 | ) |
| 2040 | |
| 2041 | self.set_linear( |
| 2042 | layer_spec.self_attention.linear[1], |
| 2043 | layer.self_attn.o_proj, |
| 2044 | quant_type=quant_type, |
| 2045 | ) |
| 2046 | |
| 2047 | # Set FFN weights |
| 2048 | self.set_linear( |
| 2049 | layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type |
no test coverage detected