first commit

validate_sample_weights.py

@@ -0,0 +1,51 @@
+"""Check real DeepSeek layer-0 MLA/Dense and layer-3 shared/routed experts."""
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from convert_deepseek_fp8_to_int8 import quantize_weight, dequantize_bf16, is_bf16_projection
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--source', required=True, type=Path)
+    parser.add_argument('--report', required=True, type=Path)
+    args = parser.parse_args()
+    torch.set_num_threads(8)
+    mapping = json.loads((args.source / 'model.safetensors.index.json').read_text())['weight_map']
+    names = [f'model.layers.0.self_attn.{name}.weight' for name in
+             ['q_a_proj', 'q_b_proj', 'kv_a_proj_with_mqa', 'kv_b_proj', 'o_proj']]
+    names += ['model.layers.0.mlp.gate_proj.weight',
+              'model.layers.3.mlp.shared_experts.gate_proj.weight',
+              'model.layers.3.mlp.experts.0.gate_proj.weight',
+              'model.layers.3.mlp.experts.0.down_proj.weight']
+    report = []
+    for name in names:
+        scale_name = name.removesuffix('weight') + 'weight_scale_inv'
+        with safe_open(str(args.source / mapping[name]), framework='pt') as reader:
+            weight = reader.get_tensor(name)
+        with safe_open(str(args.source / mapping[scale_name]), framework='pt') as reader:
+            block_scale = reader.get_tensor(scale_name)
+        with torch.inference_mode():
+            reference = weight.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:weight.shape[0], :weight.shape[1]]
+            if is_bf16_projection(name):
+                value = dequantize_bf16(weight, block_scale)
+                assert torch.equal(value, reference.bfloat16())
+                error = value.float() - reference
+                check = 'exact BF16 rounding'
+            else:
+                q, scales = quantize_weight(weight, block_scale)
+                error = q.float() * scales - reference
+                assert torch.all(error.abs() <= scales / 2 + scales * 1e-4 + 1e-7), name
+                check = 'half-scale rounding bound'
+            item = {'name': name, 'shape': list(weight.shape), 'check': check,
+                    'relative_rmse': (error.square().mean() / reference.square().mean()).sqrt().item()}
+            report.append(item)
+            print(json.dumps(item), flush=True)
+    args.report.write_text(json.dumps(report, indent=2) + '\n')
+
+
+if __name__ == '__main__':
+    main()