Spaces:
Sleeping
Sleeping
| import io | |
| import json | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| import struct | |
| import tempfile | |
| import time | |
| import unittest | |
| from unittest.mock import patch | |
| from fastapi.testclient import TestClient | |
| import ml_dtypes | |
| import numpy as np | |
| from tensor_pruning import prune_queries | |
| from tensor_pruning.checkpoint import _queries, prune_checkpoint | |
| from tensor_pruning.service import create_app | |
| def _name(layer): | |
| return f"model.layers.{layer}.self_attn.q_proj.weight" | |
| def _shard(path, entries, *, sparse=False): | |
| header, offset = {}, 0 | |
| for name, dtype, shape, length, _ in entries: | |
| header[name] = dict(dtype=dtype, shape=shape, data_offsets=[offset, offset + length]) | |
| offset += length | |
| raw = json.dumps(header).encode() | |
| raw += b" " * (-len(raw) % 8) | |
| with path.open("wb") as f: | |
| f.write(struct.pack("<Q", len(raw))) | |
| f.write(raw) | |
| if sparse: | |
| f.truncate(8 + len(raw) + offset) | |
| else: | |
| for _, _, _, _, payload in entries: | |
| f.write(payload) | |
| class CheckpointTests(unittest.TestCase): | |
| def setUp(self): | |
| self.temp = tempfile.TemporaryDirectory() | |
| self.addCleanup(self.temp.cleanup) | |
| self.root = Path(self.temp.name) | |
| self.inputs = self.root / "inputs" | |
| self.models = self.root / "models" | |
| self.inputs.mkdir() | |
| self.models.mkdir() | |
| self.jobs = self.root / "jobs" | |
| def fixture(self): | |
| dims = dict(num_hidden_layers=2, num_attention_heads=2, hidden_size=4) | |
| (self.models / "config.json").write_text(json.dumps(dims)) | |
| arrays = [np.arange(16, dtype=np.float16).reshape(2, 2, 4) - 4, | |
| np.asarray(np.arange(16) / 3, dtype=ml_dtypes.bfloat16).reshape(2, 2, 4)] | |
| mapping = {} | |
| moments = [np.array([1, 2, 3, 4]), np.array([4, 3, 2, 1])] | |
| moment_names = {} | |
| for i, (array, dtype) in enumerate(zip(arrays, ("F16", "BF16"))): | |
| name = _name(i) | |
| filename = f"shard-{1-i}.safetensors" | |
| mapping[name] = filename | |
| noise = ("unused", "F16", [2], 4, b"\xff" * 4) | |
| query = (name, dtype, [4, 4], array.nbytes, array.tobytes()) | |
| _shard(self.models / filename, [noise, query]) | |
| moment_names[name] = f"moments-{i}.npy" | |
| np.save(self.inputs / moment_names[name], moments[i]) | |
| (self.models / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping})) | |
| (self.inputs / "moments.json").write_text(json.dumps(moment_names)) | |
| return arrays, moments | |
| def test_service_selects_one_global_mask_from_fp16_and_bf16(self): | |
| arrays, moments = self.fixture() | |
| expected = prune_queries(arrays, 11, input_second_moments=moments, seed=47) | |
| app = create_app(self.inputs, self.jobs, checkpoint_root=self.models, | |
| token="example", require_token=True) | |
| headers = {"X-Tensor-Pruning-Token": "example"} | |
| body = {"keep": 11, "seed": 47, "moments_manifest": "moments.json"} | |
| with TestClient(app) as client: | |
| self.assertEqual(client.post("/v1/checkpoint-jobs", json=body).status_code, 401) | |
| self.assertEqual(client.post("/v1/checkpoint-jobs", headers=headers, | |
| json={**body, "checkpoint": "../"}).status_code, 400) | |
| created = client.post("/v1/checkpoint-jobs", headers=headers, json=body) | |
| self.assertEqual(created.status_code, 202, created.text) | |
| url = "/v1/jobs/" + created.json()["id"] | |
| for _ in range(500): | |
| status = client.get(url, headers=headers).json() | |
| if status["state"] in {"completed", "failed"}: | |
| break | |
| time.sleep(.01) | |
| self.assertEqual(status["state"], "completed", status) | |
| self.assertEqual(status["result"]["kept_slots"], 11) | |
| self.assertEqual([p["name"] for p in status["result"]["projections"]], | |
| [_name(0), _name(1)]) | |
| for i, want in enumerate(expected): | |
| got = client.get(f"{url}/masks/{i}", headers=headers) | |
| self.assertEqual(got.status_code, 200, got.text if got.status_code != 200 else "") | |
| np.testing.assert_array_equal(np.load(io.BytesIO(got.content)), want) | |
| self.assertEqual(list(self.jobs.glob("*/query-staging-*")), []) | |
| def test_job_submission_fails_closed_without_token(self): | |
| app = create_app(self.inputs, self.jobs, checkpoint_root=self.models, require_token=True) | |
| with TestClient(app) as client: | |
| self.assertEqual(client.get("/health").status_code, 200) | |
| self.assertEqual(client.post("/v1/checkpoint-jobs", json={"keep": 1}).status_code, 503) | |
| def test_disk_preflight_rejects_before_staging(self): | |
| self.fixture() | |
| usage = SimpleNamespace(total=100, used=100, free=0) | |
| with patch("tensor_pruning.checkpoint.shutil.disk_usage", return_value=usage): | |
| with self.assertRaisesRegex(OSError, "Insufficient local disk"): | |
| prune_checkpoint(self.models, self.root / "out", keep=11, | |
| workspace=self.root, max_weights=100) | |
| self.assertEqual(list(self.root.glob("query-staging-*")), []) | |
| def test_llama_70b_layout_exceeds_32_bit_global_indices(self): | |
| shape = [8192, 8192] | |
| byte_length = 2 * shape[0] * shape[1] | |
| config = dict(num_hidden_layers=80, num_attention_heads=64, hidden_size=8192) | |
| (self.models / "config.json").write_text(json.dumps(config)) | |
| entries = [(_name(i), "BF16", shape, byte_length, None) for i in range(80)] | |
| _shard(self.models / "model.safetensors", entries, sparse=True) | |
| query, total = _queries(self.models, 6_000_000_000) | |
| self.assertEqual(total, 5_368_709_120) | |
| self.assertGreater(total, 1 << 32) | |
| self.assertEqual(query[-1].shape, (64, 128, 8192)) | |
| self.assertGreater(query[-1].offset, 1 << 33) | |
| with self.assertRaisesRegex(ValueError, "eligible weight limit"): | |
| _queries(self.models, 1 << 32) | |