(self, spec, module, quant_type=common_spec.Quantization.CT2)
| 2743 | spec.gamma = layer_norm.weight |
| 2744 | |
| 2745 | def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2): |
| 2746 | spec.scale_embeddings = False |
| 2747 | self.set_embeddings(spec.embeddings, module.embed_tokens) |
| 2748 | self.set_layer_norm(spec.layer_norm, module.norm) |
| 2749 | |
| 2750 | for layer_idx, (layer_spec, layer) in enumerate(zip(spec.layer, module.layers)): |
| 2751 | self.set_layer_norm( |
| 2752 | layer_spec.self_attention.layer_norm, layer.input_layernorm |
| 2753 | ) |
| 2754 | self.set_layer_norm( |
| 2755 | layer_spec.ffn.layer_norm, layer.post_attention_layernorm |
| 2756 | ) |
| 2757 | |
| 2758 | self.set_layer_norm( |
| 2759 | layer_spec.self_attention.q_norm, layer.self_attn.q_norm |
| 2760 | ) |
| 2761 | self.set_layer_norm( |
| 2762 | layer_spec.self_attention.k_norm, layer.self_attn.k_norm |
| 2763 | ) |
| 2764 | |
| 2765 | split_layers = [common_spec.LinearSpec() for _ in range(3)] |
| 2766 | self.set_linear( |
| 2767 | split_layers[0], layer.self_attn.q_proj, quant_type=quant_type |
| 2768 | ) |
| 2769 | self.set_linear( |
| 2770 | split_layers[1], layer.self_attn.k_proj, quant_type=quant_type |
| 2771 | ) |
| 2772 | self.set_linear( |
| 2773 | split_layers[2], layer.self_attn.v_proj, quant_type=quant_type |
| 2774 | ) |
| 2775 | |
| 2776 | if quant_type == common_spec.Quantization.CT2: |
| 2777 | utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) |
| 2778 | else: |
| 2779 | cc_dim = 1 if quant_type == common_spec.Quantization.AWQ_GEMM else 0 |
| 2780 | utils.fuse_linear_prequant( |
| 2781 | layer_spec.self_attention.linear[0], split_layers, cc_dim |
| 2782 | ) |
| 2783 | |
| 2784 | self.set_linear( |
| 2785 | layer_spec.self_attention.linear[1], |
| 2786 | layer.self_attn.o_proj, |
| 2787 | quant_type=quant_type, |
| 2788 | ) |
| 2789 | |
| 2790 | self.set_linear( |
| 2791 | layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type |
| 2792 | ) |
| 2793 | self.set_linear( |
| 2794 | layer_spec.ffn.linear_0_noact, layer.mlp.up_proj, quant_type=quant_type |
| 2795 | ) |
| 2796 | self.set_linear( |
| 2797 | layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type |
| 2798 | ) |
| 2799 | |
| 2800 | delattr(layer, "self_attn") |
| 2801 | delattr(layer, "mlp") |
| 2802 | gc.collect() |
no test coverage detected