(self, spec, module, quant_type=common_spec.Quantization.CT2)
| 1792 | spec.gamma = layer_norm.weight |
| 1793 | |
| 1794 | def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2): |
| 1795 | spec.scale_embeddings = False |
| 1796 | self.set_embeddings(spec.embeddings, module.embed_tokens) |
| 1797 | self.set_layer_norm(spec.layer_norm, module.norm) |
| 1798 | |
| 1799 | for layer_spec, layer in zip(spec.layer, module.layers): |
| 1800 | self.set_layer_norm( |
| 1801 | layer_spec.self_attention.layer_norm, layer.input_layernorm |
| 1802 | ) |
| 1803 | self.set_layer_norm( |
| 1804 | layer_spec.ffn.layer_norm, layer.post_attention_layernorm |
| 1805 | ) |
| 1806 | |
| 1807 | split_layers = [common_spec.LinearSpec() for _ in range(3)] |
| 1808 | self.set_linear( |
| 1809 | split_layers[0], layer.self_attn.q_proj, quant_type=quant_type |
| 1810 | ) |
| 1811 | self.set_linear( |
| 1812 | split_layers[1], layer.self_attn.k_proj, quant_type=quant_type |
| 1813 | ) |
| 1814 | self.set_linear( |
| 1815 | split_layers[2], layer.self_attn.v_proj, quant_type=quant_type |
| 1816 | ) |
| 1817 | |
| 1818 | if quant_type == common_spec.Quantization.CT2: |
| 1819 | utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) |
| 1820 | else: |
| 1821 | cc_dim = 1 if quant_type == common_spec.Quantization.AWQ_GEMM else 0 |
| 1822 | utils.fuse_linear_prequant( |
| 1823 | layer_spec.self_attention.linear[0], split_layers, cc_dim |
| 1824 | ) |
| 1825 | self.set_linear( |
| 1826 | layer_spec.self_attention.linear[1], |
| 1827 | layer.self_attn.o_proj, |
| 1828 | quant_type=quant_type, |
| 1829 | ) |
| 1830 | |
| 1831 | self.set_linear( |
| 1832 | layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type |
| 1833 | ) |
| 1834 | self.set_linear( |
| 1835 | layer_spec.ffn.linear_0_noact, layer.mlp.up_proj, quant_type=quant_type |
| 1836 | ) |
| 1837 | self.set_linear( |
| 1838 | layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type |
| 1839 | ) |
| 1840 | |
| 1841 | delattr(layer, "self_attn") |
| 1842 | delattr(layer, "mlp") |
| 1843 | gc.collect() |
| 1844 | |
| 1845 | |
| 1846 | @register_loader("Gemma3TextConfig") |
no test coverage detected