"""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()