Spaces:
Sleeping
Sleeping
File size: 6,162 Bytes
a42d9ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | 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)
|