first commit

test_conversion.py

@@ -0,0 +1,91 @@
+"""CPU numerical and cross-shard checkpoint integration tests."""
+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
+
+from convert_fp8_to_int8 import quantize_weight, quantization_config
+
+SCRIPT = Path(__file__).with_name("convert_fp8_to_int8.py")
+
+
+class ConversionTests(unittest.TestCase):
+    def test_quantization_error_and_partial_blocks(self):
+        torch.manual_seed(42)
+        w = (torch.randn(129, 257) * 5).to(torch.float8_e4m3fn)
+        w[0] = torch.zeros(257).to(torch.float8_e4m3fn)
+        block_scale = torch.tensor([[0.1, 0.2, 0.3], [0.5, 0.6, 0.7]])
+        q, scale = quantize_weight(w, block_scale, chunk_rows=37)
+        ref = w.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:129, :257]
+        self.assertEqual(q.dtype, torch.int8)
+        self.assertEqual(tuple(scale.shape), (129, 1))
+        self.assertTrue(torch.all((q.float() * scale - ref).abs() <= scale / 2 + 1e-5))
+        self.assertEqual(scale[0].item(), 1.0)
+        self.assertEqual(q[0].count_nonzero().item(), 0)
+        self.assertTrue(torch.allclose(scale[1:], ref.abs().amax(1, keepdim=True)[1:] / 127))
+        with self.assertRaises(ValueError):
+            quantize_weight(w, torch.full((2, 3), float("nan")))
+
+    def test_cross_shard_conversion_and_resume(self):
+        with tempfile.TemporaryDirectory() as temp:
+            root = Path(temp)
+            src, dst = root / "source", root / "out"
+            src.mkdir()
+            config = {"model_type": "minimax_m2", "num_hidden_layers": 1,
+                      "num_local_experts": 1, "quantization_config": {
+                          "quant_method": "fp8", "fmt": "float8_e4m3fn",
+                          "activation_scheme": "dynamic", "weight_block_size": [128, 128]}}
+            (src / "config.json").write_text(json.dumps(config))
+            (src / "tokenizer_config.json").write_text('{}')
+            modules = [f"self_attn.{p}_proj" for p in "qkvo"]
+            modules += [f"block_sparse_moe.experts.0.w{i}" for i in (1, 2, 3)]
+            tensors, scales = {}, {}
+            for module in modules:
+                key = f"model.layers.0.{module}.weight"
+                tensors[key] = torch.arange(24).reshape(4, 6).float().to(torch.float8_e4m3fn)
+                scales[key.removesuffix("weight") + "weight_scale_inv"] = torch.tensor([[0.25]])
+            tensors["lm_head.weight"] = torch.ones(4, 6, dtype=torch.bfloat16)
+            tensors["model.layers.0.block_sparse_moe.gate.weight"] = torch.ones(1, 6)
+            save_file(tensors, str(src / "weights.safetensors"))
+            save_file(scales, str(src / "scales.safetensors"))
+            mapping = {k: "weights.safetensors" for k in tensors}
+            mapping.update({k: "scales.safetensors" for k in scales})
+            (src / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping}))
+            before = {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in src.iterdir()}
+            command = [sys.executable, str(SCRIPT), "--source", str(src), "--output", str(dst), "--threads", "1"]
+            subprocess.run(command + ["--dry-run"], check=True)
+            self.assertFalse(dst.exists())
+            subprocess.run(command, check=True)
+            out_map = json.loads((dst / "model.safetensors.index.json").read_text())["weight_map"]
+            self.assertFalse(any(k.endswith("weight_scale_inv") for k in out_map))
+            self.assertEqual(len(out_map), len(tensors) + 7)
+            for key, filename in out_map.items():
+                with safe_open(str(dst / filename), framework="pt") as reader:
+                    value = reader.get_tensor(key)
+                if key.endswith("weight_scale"):
+                    self.assertEqual(tuple(value.shape), (4, 1))
+                    self.assertEqual(value.dtype, torch.float32)
+                elif key in tensors and "proj.weight" not in key and ".experts." not in key:
+                    self.assertTrue(torch.equal(value, tensors[key]))
+                else:
+                    self.assertEqual(value.dtype, torch.int8)
+            self.assertEqual(json.loads((dst / "config.json").read_text())["quantization_config"], quantization_config())
+            self.assertTrue((dst / "tokenizer_config.json").exists())
+            for filename in ("config.json", "model.safetensors.index.json", "conversion_report.json"):
+                (dst / filename).unlink()
+            subprocess.run(command + ["--resume"], check=True)
+            self.assertTrue((dst / "model.safetensors.index.json").exists())
+            self.assertEqual(before, {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in src.iterdir()})
+            self.assertNotEqual(subprocess.run(command, capture_output=True).returncode, 0)
+
+
+if __name__ == "__main__":
+    torch.set_num_threads(1)
+    unittest.main()