92 lines
4.9 KiB
Python
92 lines
4.9 KiB
Python
"""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()
|