(
*,
input_path: str = "",
)
| 13 | |
| 14 | @torch.inference_mode() |
| 15 | def convert_ts_checkpoint( |
| 16 | *, |
| 17 | input_path: str = "", |
| 18 | ) -> None: |
| 19 | |
| 20 | config = model.ModelArgs() |
| 21 | print(f"Model config {config.__dict__}") |
| 22 | |
| 23 | def quant_weight_int8(weight): |
| 24 | s = 1.0 / weight.abs().mean().clamp_(min=1e-5) |
| 25 | new_weight = (weight * s).round().clamp(-1, 1).to(torch.int8) |
| 26 | new_scale = (1.0 / s).to(torch.bfloat16) |
| 27 | return new_weight, new_scale.reshape(1) |
| 28 | |
| 29 | def quant_weight_fp16(weight): |
| 30 | s = 1.0 / weight.abs().mean().clamp_(min=1e-5) |
| 31 | new_weight = (weight * s).round().clamp(-1, 1) / s |
| 32 | return new_weight |
| 33 | |
| 34 | def convert_int8_to_int2(weight): |
| 35 | return convert_weight_int8_to_int2(weight) |
| 36 | |
| 37 | merged_result = torch.load(input_path, map_location="cpu", mmap=True, weights_only=True) |
| 38 | int2_result = {} |
| 39 | fp16_result = {} |
| 40 | zero = torch.zeros(1).to(torch.bfloat16) |
| 41 | for key, value in merged_result.items(): |
| 42 | if 'wqkv' in key: |
| 43 | wq = value[:config.dim] |
| 44 | wk = value[config.dim:config.dim // config.n_heads * config.n_kv_heads + config.dim] |
| 45 | wv = value[config.dim // config.n_heads * config.n_kv_heads + config.dim:] |
| 46 | wq_weight, wa_scale = quant_weight_int8(wq) |
| 47 | wk_weight, wb_scale = quant_weight_int8(wk) |
| 48 | wv_weight, wc_scale = quant_weight_int8(wv) |
| 49 | wqkv_weight = torch.cat([wq_weight, wk_weight, wv_weight], dim=0) |
| 50 | wqkv_scale = torch.cat([wa_scale, wb_scale, wc_scale, zero], dim=0) |
| 51 | int2_result[key] = convert_int8_to_int2(wqkv_weight) |
| 52 | int2_result[key.replace('weight', 'weight_scale')] = wqkv_scale |
| 53 | |
| 54 | wq_weight = quant_weight_fp16(wq) |
| 55 | wk_weight = quant_weight_fp16(wk) |
| 56 | wv_weight = quant_weight_fp16(wv) |
| 57 | wqkv_weight = torch.cat([wq_weight, wk_weight, wv_weight], dim=0) |
| 58 | fp16_result[key] = wqkv_weight |
| 59 | elif 'w13' in key: |
| 60 | w1 = value[:config.ffn_dim] |
| 61 | w3 = value[config.ffn_dim:] |
| 62 | w1_weight, w1_scale = quant_weight_int8(w1) |
| 63 | w3_weight, w3_scale = quant_weight_int8(w3) |
| 64 | w13_weight = torch.cat([w1_weight, w3_weight], dim=0) |
| 65 | w13_scale = torch.cat([w1_scale, w3_scale, zero, zero], dim=0) |
| 66 | int2_result[key] = convert_int8_to_int2(w13_weight) |
| 67 | int2_result[key.replace('weight', 'weight_scale')] = w13_scale |
| 68 | |
| 69 | w1_weight = quant_weight_fp16(w1) |
| 70 | w3_weight = quant_weight_fp16(w3) |
| 71 | w13_weight = torch.cat([w1_weight, w3_weight], dim=0) |
| 72 | fp16_result[key] = w13_weight |
no test coverage detected