124 lines
6.7 KiB
Python
124 lines
6.7 KiB
Python
|
"""CPU tests for numerical conversion, MLA, shared experts, MTP and resume."""
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from safetensors import safe_open
|
||
|
|
from safetensors.torch import save_file
|
||
|
|
|
||
|
|
import convert_deepseek_fp8_to_int8 as convert
|
||
|
|
|
||
|
|
|
||
|
|
def fixture(source):
|
||
|
|
source.mkdir()
|
||
|
|
config = {
|
||
|
|
"model_type": "deepseek_v3", "architectures": ["DeepseekV3ForCausalLM"],
|
||
|
|
"num_hidden_layers": 2, "num_nextn_predict_layers": 1, "hidden_size": 128,
|
||
|
|
"q_lora_rank": 128, "kv_lora_rank": 128, "num_attention_heads": 8,
|
||
|
|
"qk_nope_head_dim": 16, "qk_rope_head_dim": 64, "v_head_dim": 16,
|
||
|
|
"first_k_dense_replace": 1, "intermediate_size": 256,
|
||
|
|
"moe_intermediate_size": 128, "n_routed_experts": 2, "n_shared_experts": 1,
|
||
|
|
"moe_layer_freq": 1, "torch_dtype": "bfloat16",
|
||
|
|
"quantization_config": {"quant_method": "fp8", "fmt": "e4m3",
|
||
|
|
"activation_scheme": "dynamic", "weight_block_size": [128, 128]},
|
||
|
|
}
|
||
|
|
(source / "config.json").write_text(json.dumps(config))
|
||
|
|
(source / "tokenizer_config.json").write_text('{}')
|
||
|
|
weights, scales = {}, {}
|
||
|
|
torch.manual_seed(12)
|
||
|
|
for name, shape in convert.expected_projection_shapes(config).items():
|
||
|
|
weights[name] = torch.randn(*shape).clamp(-5, 5).to(torch.float8_e4m3fn)
|
||
|
|
scales[name.removesuffix("weight") + "weight_scale_inv"] = (
|
||
|
|
torch.rand((shape[0] + 127) // 128, (shape[1] + 127) // 128) * 0.1)
|
||
|
|
weights["model.embed_tokens.weight"] = torch.ones(8, 128, dtype=torch.bfloat16)
|
||
|
|
weights["lm_head.weight"] = torch.ones(8, 128, dtype=torch.bfloat16)
|
||
|
|
weights["model.layers.1.mlp.gate.weight"] = torch.ones(2, 128, dtype=torch.bfloat16)
|
||
|
|
weights["model.layers.1.mlp.gate.e_score_correction_bias"] = torch.ones(2)
|
||
|
|
weights["model.norm.weight"] = torch.ones(128, dtype=torch.bfloat16)
|
||
|
|
# MTP has no destination in the main model and must not leak into output.
|
||
|
|
weights["model.layers.2.eh_proj.weight"] = torch.ones(128, 256, dtype=torch.bfloat16)
|
||
|
|
weights["model.layers.2.self_attn.q_a_proj.weight"] = torch.ones(128, 128).to(torch.float8_e4m3fn)
|
||
|
|
scales["model.layers.2.self_attn.q_a_proj.weight_scale_inv"] = torch.ones(1, 1)
|
||
|
|
save_file(weights, str(source / "weights.safetensors"))
|
||
|
|
save_file(scales, str(source / "scales.safetensors"))
|
||
|
|
mapping = {key: "weights.safetensors" for key in weights}
|
||
|
|
mapping.update({key: "scales.safetensors" for key in scales})
|
||
|
|
(source / "model.safetensors.index.json").write_text(json.dumps({
|
||
|
|
"metadata": {"total_size": 999999999999999}, "weight_map": mapping}))
|
||
|
|
return config, weights, scales
|
||
|
|
|
||
|
|
|
||
|
|
class ConversionTests(unittest.TestCase):
|
||
|
|
def test_numeric(self):
|
||
|
|
torch.manual_seed(42)
|
||
|
|
weight = torch.randn(129, 257).to(torch.float8_e4m3fn)
|
||
|
|
weight[0] = torch.zeros(257).to(torch.float8_e4m3fn)
|
||
|
|
block_scale = torch.tensor([[0.1, 0.2, 0.3], [0.5, 0.6, 0.7]])
|
||
|
|
reference = weight.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:129, :257]
|
||
|
|
q, scales = convert.quantize_weight(weight, block_scale, chunk_rows=37)
|
||
|
|
self.assertTrue(torch.all((q.float() * scales - reference).abs() <= scales / 2 + 1e-5))
|
||
|
|
self.assertEqual(q[0].count_nonzero().item(), 0)
|
||
|
|
self.assertEqual(scales[0].item(), 1)
|
||
|
|
self.assertTrue(torch.equal(convert.dequantize_bf16(weight, block_scale, 37), reference.bfloat16()))
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
convert.quantize_weight(weight, block_scale * float('nan'))
|
||
|
|
|
||
|
|
def test_checkpoint(self):
|
||
|
|
with tempfile.TemporaryDirectory() as temp:
|
||
|
|
source, output = Path(temp) / "source", Path(temp) / "output"
|
||
|
|
config, weights, scales = fixture(source)
|
||
|
|
original = {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in source.iterdir()}
|
||
|
|
command = [sys.executable, str(Path(convert.__file__)), "--source", str(source),
|
||
|
|
"--output", str(output), "--threads", "1"]
|
||
|
|
subprocess.run(command + ["--dry-run"], check=True)
|
||
|
|
self.assertFalse(output.exists())
|
||
|
|
subprocess.run(command, check=True)
|
||
|
|
mapping = json.loads((output / "model.safetensors.index.json").read_text())["weight_map"]
|
||
|
|
self.assertFalse(any(k.startswith('model.layers.2.') for k in mapping))
|
||
|
|
self.assertFalse(any(k.endswith('weight_scale_inv') for k in mapping))
|
||
|
|
for key, filename in mapping.items():
|
||
|
|
with safe_open(str(output / filename), framework="pt") as reader:
|
||
|
|
value = reader.get_tensor(key)
|
||
|
|
if key.endswith("weight_scale"):
|
||
|
|
self.assertEqual(value.dtype, torch.float32)
|
||
|
|
self.assertEqual(value.shape[-1], 1)
|
||
|
|
elif weights[key].dtype != torch.float8_e4m3fn:
|
||
|
|
self.assertTrue(torch.equal(value, weights[key]))
|
||
|
|
elif convert.is_bf16_projection(key):
|
||
|
|
self.assertEqual(value.dtype, torch.bfloat16)
|
||
|
|
self.assertNotIn(key.removesuffix('weight') + 'weight_scale', mapping)
|
||
|
|
self.assertTrue(torch.equal(value, convert.dequantize_bf16(
|
||
|
|
weights[key], scales[key.removesuffix('weight') + 'weight_scale_inv'])))
|
||
|
|
else:
|
||
|
|
self.assertEqual(value.dtype, torch.int8)
|
||
|
|
final_config = json.loads((output / 'config.json').read_text())
|
||
|
|
self.assertEqual(final_config['num_nextn_predict_layers'], 0)
|
||
|
|
self.assertEqual(final_config['quantization_config'], convert.quantization_config())
|
||
|
|
for name in ('config.json', 'model.safetensors.index.json', 'conversion_report.json'):
|
||
|
|
(output / name).unlink()
|
||
|
|
subprocess.run(command + ['--resume'], check=True)
|
||
|
|
self.assertTrue((output / 'conversion_report.json').exists())
|
||
|
|
self.assertEqual(original, {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in source.iterdir()})
|
||
|
|
self.assertNotEqual(subprocess.run(command, capture_output=True).returncode, 0)
|
||
|
|
|
||
|
|
def test_missing_projection_rejected(self):
|
||
|
|
with tempfile.TemporaryDirectory() as temp:
|
||
|
|
source = Path(temp) / 'source'
|
||
|
|
fixture(source)
|
||
|
|
path = source / 'model.safetensors.index.json'
|
||
|
|
index = json.loads(path.read_text())
|
||
|
|
del index['weight_map']['model.layers.1.mlp.experts.0.up_proj.weight']
|
||
|
|
path.write_text(json.dumps(index))
|
||
|
|
with self.assertRaisesRegex(ValueError, 'Projection coverage'):
|
||
|
|
convert.plan_source(source)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
torch.set_num_threads(1)
|
||
|
|
unittest.main()
|