(self, spec, module, quant_type=common_spec.Quantization.CT2)
| 3016 | ) |
| 3017 | |
| 3018 | def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2): |
| 3019 | spec.scale_embeddings = False |
| 3020 | self.set_embeddings(spec.embeddings, module.embed_tokens) |
| 3021 | self.set_layer_norm(spec.layer_norm, module.norm) |
| 3022 | |
| 3023 | for layer_spec, layer in zip(spec.layer, module.layers): |
| 3024 | self.set_layer_norm( |
| 3025 | layer_spec.self_attention.layer_norm, layer.input_layernorm |
| 3026 | ) |
| 3027 | self.set_layer_norm( |
| 3028 | layer_spec.ffn.layer_norm, layer.post_attention_layernorm |
| 3029 | ) |
| 3030 | |
| 3031 | self.set_linear( |
| 3032 | layer_spec.self_attention.linear[0], |
| 3033 | layer.self_attn.qkv_proj, |
| 3034 | quant_type=quant_type, |
| 3035 | ) |
| 3036 | self.set_linear( |
| 3037 | layer_spec.self_attention.linear[1], |
| 3038 | layer.self_attn.o_proj, |
| 3039 | quant_type=quant_type, |
| 3040 | ) |
| 3041 | if ( |
| 3042 | layer.self_attn.rotary_emb.long_factor is not None |
| 3043 | and layer.self_attn.rotary_emb.short_factor is not None |
| 3044 | ): |
| 3045 | self.set_rotary_embeddings( |
| 3046 | layer_spec.self_attention, |
| 3047 | layer.self_attn.rotary_emb.long_factor, |
| 3048 | layer.self_attn.rotary_emb.short_factor, |
| 3049 | ) |
| 3050 | |
| 3051 | # Handle gate_up_proj differently for AWQ vs regular models |
| 3052 | if quant_type == common_spec.Quantization.CT2: |
| 3053 | gate_proj, up_proj = layer.mlp.gate_up_proj.weight.chunk(2, dim=0) |
| 3054 | layer_spec.ffn.linear_0.weight = gate_proj |
| 3055 | layer_spec.ffn.linear_0_noact.weight = up_proj |
| 3056 | else: |
| 3057 | # AWQ: chunk qweight, scales, and qzeros |
| 3058 | gate_qweight, up_qweight = layer.mlp.gate_up_proj.qweight.chunk( |
| 3059 | 2, dim=1 |
| 3060 | ) |
| 3061 | gate_scales, up_scales = layer.mlp.gate_up_proj.scales.chunk(2, dim=1) |
| 3062 | gate_qzeros, up_qzeros = layer.mlp.gate_up_proj.qzeros.chunk(2, dim=1) |
| 3063 | |
| 3064 | layer_spec.ffn.linear_0.weight = gate_qweight |
| 3065 | layer_spec.ffn.linear_0.weight_scale = gate_scales |
| 3066 | layer_spec.ffn.linear_0.weight_zero = gate_qzeros |
| 3067 | |
| 3068 | layer_spec.ffn.linear_0_noact.weight = up_qweight |
| 3069 | layer_spec.ffn.linear_0_noact.weight_scale = up_scales |
| 3070 | layer_spec.ffn.linear_0_noact.weight_zero = up_qzeros |
| 3071 | |
| 3072 | self.set_linear( |
| 3073 | layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type |
| 3074 | ) |
| 3075 |
no test coverage detected