prain: int8/int4 experts quantized once on the GPU, rounds through vLLM's expert_map, [miss] min_weight skips experts not in VRAM; scripts: tune_moe, quality (answer NLL on identical token ids), bench --context, measure --quant
This commit is contained in:
@@ -76,3 +76,17 @@ class Cli(TmpCase):
|
||||
self.assertEqual((tuned.pack_dir, tuned.trace_path), ("/packs", trace))
|
||||
self.assertLessEqual(sum(tuned.placement(layer, 128).cache_slots for layer in range(3)),
|
||||
int(0.2 * 1024 / 11.34) // 3 * 3)
|
||||
|
||||
def test_quantized_experts_check_and_tune(self):
|
||||
hf = self.write("config.json", json.dumps(GEMMA4))
|
||||
out, _ = self.run_cli("check", self.write("q.toml", 'pack_dir = "/packs"\nquant = "int4"\n'), hf)
|
||||
self.assertIn("expert size 3.01 MiB\n", out)
|
||||
trace = os.path.join(self.dir, "trace.jsonl.gz")
|
||||
self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
|
||||
out, _ = self.run_cli("tune", trace, hf, "--quant", "int4", "--expert-vram-gib", "0.2",
|
||||
"--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
|
||||
tuned = load_config(self.write("tuned.toml", out))
|
||||
self.assertEqual(tuned.quant, "int4")
|
||||
slots = sum(tuned.placement(layer, 128).cache_slots for layer in range(3))
|
||||
self.assertGreater(slots, int(0.2 * 1024 / 11.34)) # int4 experts are ~4x smaller
|
||||
self.assertLessEqual(slots, int(0.2 * 1024 / 3.01))
|
||||
|
||||
@@ -107,6 +107,23 @@ mode = "auto"
|
||||
self.assertEqual(config.stats_path, os.path.expanduser("~/s.json"))
|
||||
self.assertEqual(config.trace_path, "/abs/t.gz")
|
||||
|
||||
def test_miss_min_weight(self):
|
||||
config = mc.load_config(self.write("m.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 0.7\n'))
|
||||
self.assertEqual(config.miss.min_weight, 0.7)
|
||||
self.assertEqual(mc.Miss().min_weight, 1.0) # default: compute every routed expert
|
||||
with self.assertRaisesRegex(ValueError, "min_weight must be a number in"):
|
||||
mc.load_config(self.write("m2.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 1.5\n'))
|
||||
|
||||
def test_quant_key_and_expert_sizes(self):
|
||||
self.assertEqual(mc.load_config(self.write("q.toml", 'pack_dir = "/p"\nquant = "int4"\n')).quant, "int4")
|
||||
self.assertEqual(mc.load_config(self.write("b.toml", 'pack_dir = "/p"\n')).quant, "bf16")
|
||||
with self.assertRaisesRegex(ValueError, "quant must be one of"):
|
||||
mc.load_config(self.write("bad.toml", 'pack_dir = "/p"\nquant = "int3"\n'))
|
||||
h, i = 2816, 704
|
||||
self.assertEqual(mc.expert_bytes(h, i), 3 * h * i * 2)
|
||||
self.assertEqual(mc.expert_bytes(h, i, "int8"), 3 * h * i + 4 * (2 * i + h)) # fp32 row scales
|
||||
self.assertEqual(mc.expert_bytes(h, i, "int4"), 3 * h * i // 2 + 2 * 3 * h * i // 64)
|
||||
|
||||
def test_serve_and_draft_tables(self):
|
||||
config = mc.load_config(self.write("s.toml", 'pack_dir = "/p"\n[serve]\nmax_requests = 4\n'
|
||||
'kv_cache_gib = 1.5\n[draft]\nfull_requests = 2\n'
|
||||
|
||||
@@ -31,3 +31,64 @@ class KernelMatch(unittest.TestCase):
|
||||
global_num_experts=experts).float().cpu()
|
||||
cpu = cpu_experts(x, topk_weights, topk_ids, set(range(experts)), lambda e: (w13[e], w2[e]))
|
||||
torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
|
||||
|
||||
def test_a_round_computes_only_the_experts_its_slot_table_maps(self):
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
||||
|
||||
from prain.tiered import cpu_experts, slot_table
|
||||
|
||||
torch.manual_seed(0)
|
||||
tokens, hidden, inter, experts = 64, 256, 128, 6
|
||||
x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
|
||||
w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
|
||||
w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
|
||||
topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
|
||||
topk_ids[0, 1] = experts # a padding route
|
||||
topk_weights = torch.rand(tokens, 2)
|
||||
table = slot_table({1: 0, 4: 1}, experts, "cuda")
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
gpu = fused_experts(x.cuda(), w13[[1, 4]].cuda(), w2[[1, 4]].cuda(), topk_weights.cuda(),
|
||||
topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
|
||||
global_num_experts=table.numel(), expert_map=table).float().cpu()
|
||||
cpu = cpu_experts(x, topk_weights, topk_ids, {1, 4}, lambda e: (w13[e], w2[e]))
|
||||
torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
|
||||
|
||||
def test_quantized_slots_compute_what_the_cpu_path_computes(self):
|
||||
"""Pack layout -> slot layout on the GPU -> vLLM's int8/int4 kernel for a round that
|
||||
holds experts 1 and 3, against the CPU kernels on the pack layout and against the
|
||||
dequantized weights."""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
||||
|
||||
from prain.quant import (cpu_matmul, dequantize, moe_quant_config, quantize, to_cpu_int4,
|
||||
to_gpu_int4)
|
||||
from prain.tiered import cpu_experts, slot_table
|
||||
|
||||
torch.manual_seed(0)
|
||||
tokens, hidden, inter, experts = 8, 256, 128, 4
|
||||
x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
|
||||
w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
|
||||
w2 = torch.randn(experts, hidden, inter) * 0.1
|
||||
topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
|
||||
topk_weights = torch.rand(tokens, 2)
|
||||
table = slot_table({1: 0, 3: 1}, experts, "cuda")
|
||||
for quant in ("int8", "int4"):
|
||||
parts = [quantize(w, quant) for w in (w13, w2)] # [q, scale], gpu layout
|
||||
pack = [[to_cpu_int4(q) if quant == "int4" else q, s] for q, s in parts]
|
||||
slots = [[to_gpu_int4(q[[1, 3]].cuda()) if quant == "int4" else q[[1, 3]].cuda(),
|
||||
s[[1, 3]].cuda()] for q, s in pack]
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
gpu = fused_experts(x.cuda(), slots[0][0], slots[1][0], topk_weights.cuda(),
|
||||
topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
|
||||
global_num_experts=table.numel(), expert_map=table,
|
||||
quant_config=moe_quant_config(quant, slots[0][1], slots[1][1]))
|
||||
cpu = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
|
||||
cpu_matmul(quant, q[e], s[e]) for q, s in pack))
|
||||
exact = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
|
||||
dequantize(q, s, quant)[e] for q, s in parts))
|
||||
self.assertGreater(exact.norm().item(), 1.0)
|
||||
torch.testing.assert_close(cpu, exact, rtol=3e-2, atol=3e-2, msg=quant)
|
||||
torch.testing.assert_close(gpu.float().cpu(), exact, rtol=3e-2, atol=3e-2, msg=quant)
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
import unittest
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ImportError: # the planning host has no torch
|
||||
torch = None
|
||||
|
||||
if torch is not None:
|
||||
from prain.quant import (GROUP, cpu_matmul, dequantize, layout, quantize, to_cpu_int4,
|
||||
to_gpu_int4)
|
||||
|
||||
|
||||
@unittest.skipIf(torch is None, "needs torch")
|
||||
class Quant(unittest.TestCase):
|
||||
def setUp(self):
|
||||
torch.manual_seed(0)
|
||||
self.w = torch.randn(64, 256) * 0.02
|
||||
|
||||
def test_layouts_match_what_quantize_returns(self):
|
||||
for quant in ("int8", "int4"):
|
||||
got = [(tuple(t.shape), t.dtype) for t in quantize(self.w, quant)]
|
||||
self.assertEqual(got, layout(quant, 64, 256))
|
||||
|
||||
def test_round_trip_error(self):
|
||||
for quant, bound in (("int8", 0.01), ("int4", 0.11)):
|
||||
err = (dequantize(*quantize(self.w, quant), quant, torch.float32) - self.w).norm()
|
||||
self.assertLess(err / self.w.norm(), bound)
|
||||
|
||||
def test_int4_clipping_beats_plain_round_to_nearest(self):
|
||||
g = self.w.unflatten(-1, (-1, GROUP))
|
||||
scale = (g.abs().amax(-1, keepdim=True) / 7).bfloat16().float()
|
||||
plain = ((g / scale).round() * scale - g).norm()
|
||||
clipped = (dequantize(*quantize(self.w, "int4"), "int4", torch.float32) - self.w).norm()
|
||||
self.assertLess(clipped, plain)
|
||||
|
||||
def test_cpu_layout_is_torch_avx2_int4pack_and_inverts(self):
|
||||
q, _ = quantize(self.w, "int4")
|
||||
nibbles = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(torch.int32)
|
||||
self.assertTrue(torch.equal(to_cpu_int4(q), torch._convert_weight_to_int4pack_for_cpu(nibbles, 1)))
|
||||
self.assertTrue(torch.equal(to_gpu_int4(to_cpu_int4(q)), q))
|
||||
|
||||
def test_cpu_matmul_matches_dequantized_weights(self):
|
||||
x = torch.randn(3, 256, dtype=torch.bfloat16)
|
||||
for quant in ("int8", "int4"):
|
||||
q, scale = quantize(self.w, quant)
|
||||
ref = x.float() @ dequantize(q, scale, quant, torch.float32).t()
|
||||
stored = to_cpu_int4(q) if quant == "int4" else q
|
||||
got = cpu_matmul(quant, stored, scale)(x).float()
|
||||
self.assertLess(((got - ref).norm() / ref.norm()).item(), 0.01, quant)
|
||||
+37
-7
@@ -9,7 +9,7 @@ from prain.config import Miss, Policy
|
||||
from prain.policy import SlotCache
|
||||
|
||||
if torch is not None:
|
||||
from prain.tiered import (available_experts, cpu_experts, degrade, remap, route_stats,
|
||||
from prain.tiered import (available_experts, cpu_experts, degrade, route_stats, slot_table,
|
||||
tiered_forward)
|
||||
|
||||
HIDDEN, INTER, EXPERTS = 6, 4, 10
|
||||
@@ -52,8 +52,9 @@ class TieredForward(unittest.TestCase):
|
||||
loads.append(expert)
|
||||
slots13[slot], slots2[slot] = self.w13[expert], self.w2[expert]
|
||||
|
||||
def kernel(x, w, ids, used):
|
||||
return naive_moe(x, slots13, slots2, w, ids)
|
||||
def kernel(x, w, ids, table, used):
|
||||
self.assertTrue(bool((ids >= 0).all())) # padding id -1 arrives as num_experts
|
||||
return naive_moe(x, slots13, slots2, w, table[ids.long()])
|
||||
|
||||
out = tiered_forward(self.x, self.topk_weights, self.topk_ids, cache, load, kernel,
|
||||
miss=miss, cpu_weights=lambda e: (self.w13[e], self.w2[e]),
|
||||
@@ -80,6 +81,33 @@ class TieredForward(unittest.TestCase):
|
||||
self.assertEqual(cpu_experts(self.x, self.topk_weights, self.topk_ids, set(),
|
||||
None).abs().sum().item(), 0.0)
|
||||
|
||||
def test_miss_min_weight_skips_missing_experts_for_every_token(self):
|
||||
"""[miss] min_weight < 1: every token skips experts that are not in VRAM while it keeps
|
||||
that share of its router weight; experts some token copies anyway serve all tokens."""
|
||||
everyone = torch.ones(6, dtype=torch.bool)
|
||||
resident = torch.zeros(EXPERTS, dtype=torch.bool)
|
||||
resident[[0, 1, 2]] = True
|
||||
_, first = degrade(self.topk_weights, self.topk_ids, everyone, resident, 0.5)
|
||||
free = resident.clone()
|
||||
free[first[first >= 0].long()] = True
|
||||
w, i = degrade(self.topk_weights, self.topk_ids, everyone, free, 0.5)
|
||||
out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=0.5))
|
||||
torch.testing.assert_close(out, naive_moe(self.x, self.w13, self.w2, w, i))
|
||||
self.assertEqual(sorted(loads), sorted(set(i[i >= 0].tolist()) - {0, 1, 2}))
|
||||
self.assertLess(len(loads), 7) # computing everything copies all 7 others
|
||||
self.assertEqual(self.seen_experts, list(range(EXPERTS))) # the trace sees the routing
|
||||
out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=1.0))
|
||||
torch.testing.assert_close(out, self.expected)
|
||||
self.assertEqual(len(loads), 7)
|
||||
|
||||
def test_cpu_experts_take_matmul_functions(self):
|
||||
"""Quantized experts hand the CPU path x -> x @ w.T functions instead of weights."""
|
||||
def matmuls(e):
|
||||
return (lambda v: v @ self.w13[e].t()), (lambda v: v @ self.w2[e].t())
|
||||
|
||||
out = cpu_experts(self.x, self.topk_weights, self.topk_ids, set(range(EXPERTS)), matmuls)
|
||||
torch.testing.assert_close(out, self.expected.float())
|
||||
|
||||
def test_route_stats(self):
|
||||
stats = route_stats(torch.tensor([[0.5, 0.25], [1.0, 0.0]]),
|
||||
torch.tensor([[2, 1], [2, -1]]), 4, recency_half_life=1.0)
|
||||
@@ -89,10 +117,10 @@ class TieredForward(unittest.TestCase):
|
||||
self.assertAlmostEqual(stats[2][2], 0.5 * 0.5 + 1.0) # the first token is one half-life old
|
||||
self.assertAlmostEqual(stats[1][2], 0.125)
|
||||
|
||||
def test_remap_zeroes_routes_outside_the_round(self):
|
||||
weights, slots = remap(torch.ones(1, 3), torch.tensor([[2, 7, -1]]), {7: 4}, 10)
|
||||
self.assertEqual(weights.tolist(), [[0.0, 1.0, 0.0]])
|
||||
self.assertEqual(slots.tolist(), [[0, 4, 0]])
|
||||
def test_slot_table_maps_the_round_and_skips_the_rest(self):
|
||||
table = slot_table({7: 4, 2: 0}, 10, "cpu")
|
||||
self.assertEqual(table.tolist(), [-1, -1, 0, -1, -1, -1, -1, 4, -1, -1, -1])
|
||||
self.assertEqual(table.dtype, torch.int32) # one extra row: padding id 10 -> -1
|
||||
|
||||
def test_degrade_drops_unavailable_routes_of_draft_tokens_only(self):
|
||||
weights = torch.tensor([[0.5, 0.3, 0.2], [0.5, 0.3, 0.2]])
|
||||
@@ -115,6 +143,8 @@ class TieredForward(unittest.TestCase):
|
||||
available[[0, 1, 2, 4, 5, 6, 7, 8]] = True
|
||||
self.assertTrue(torch.equal(
|
||||
available_experts(self.topk_ids, draft, {0}, EXPERTS), available))
|
||||
_, kept = degrade(self.topk_weights, self.topk_ids, draft, available, 0.5)
|
||||
available[kept[kept >= 0].long()] = True # experts some row copies anyway serve all rows
|
||||
expected = naive_moe(self.x, self.w13, self.w2,
|
||||
*degrade(self.topk_weights, self.topk_ids, draft, available, 0.5))
|
||||
torch.testing.assert_close(out, expected)
|
||||
|
||||
Reference in New Issue
Block a user