129 lines
6.4 KiB
Python
129 lines
6.4 KiB
Python
import unittest
|
|
|
|
try:
|
|
import torch
|
|
except ImportError: # the planning host has no torch; the vLLM container does
|
|
torch = None
|
|
|
|
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,
|
|
tiered_forward)
|
|
|
|
HIDDEN, INTER, EXPERTS = 6, 4, 10
|
|
|
|
|
|
def naive_moe(x, w13, w2, topk_weights, topk_ids):
|
|
"""Gated-GELU experts, one route at a time; id -1 = no route."""
|
|
out = torch.zeros(x.shape, dtype=torch.float64)
|
|
for t in range(x.shape[0]):
|
|
for j in range(topk_ids.shape[1]):
|
|
e = int(topk_ids[t, j])
|
|
if e >= 0:
|
|
h = w13[e] @ x[t]
|
|
act = torch.nn.functional.gelu(h[:INTER], approximate="tanh") * h[INTER:]
|
|
out[t] += topk_weights[t, j] * (w2[e] @ act)
|
|
return out
|
|
|
|
|
|
@unittest.skipIf(torch is None, "needs torch")
|
|
class TieredForward(unittest.TestCase):
|
|
def setUp(self):
|
|
torch.manual_seed(0)
|
|
self.w13 = torch.randn(EXPERTS, 2 * INTER, HIDDEN, dtype=torch.float64)
|
|
self.w2 = torch.randn(EXPERTS, HIDDEN, INTER, dtype=torch.float64)
|
|
self.x = torch.randn(6, HIDDEN, dtype=torch.float64)
|
|
self.topk_ids = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 0, 1],
|
|
[2, 3, -1], [4, 5, 6]], dtype=torch.int32)
|
|
self.topk_weights = torch.rand(6, 3, dtype=torch.float64)
|
|
self.expected = naive_moe(self.x, self.w13, self.w2, self.topk_weights, self.topk_ids)
|
|
|
|
def run_tiered(self, pinned, cache_slots, miss=None, policy=Policy(), draft=None):
|
|
cache = SlotCache(pinned, EXPERTS, cache_slots, policy)
|
|
slots13 = torch.zeros(len(pinned) + cache_slots, 2 * INTER, HIDDEN, dtype=torch.float64)
|
|
slots2 = torch.zeros(len(pinned) + cache_slots, HIDDEN, INTER, dtype=torch.float64)
|
|
for slot, e in enumerate(pinned):
|
|
slots13[slot], slots2[slot] = self.w13[e], self.w2[e]
|
|
loads, seen = [], []
|
|
|
|
def load(expert, slot):
|
|
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)
|
|
|
|
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]),
|
|
on_stats=lambda stats, tokens: seen.append((tokens, sorted(stats))),
|
|
draft=draft, min_weight=0.5)
|
|
self.assertEqual(seen[0][0], 6)
|
|
self.seen_experts = seen[0][1]
|
|
return out, loads, cache
|
|
|
|
def test_matches_dense_moe_when_everything_needs_rounds(self):
|
|
out, loads, _ = self.run_tiered(pinned=[0, 1], cache_slots=3) # 8 others -> 3 rounds
|
|
torch.testing.assert_close(out, self.expected)
|
|
self.assertEqual(sorted(loads), [2, 3, 4, 5, 6, 7, 8, 9])
|
|
|
|
def test_matches_dense_moe_in_one_round(self):
|
|
out, _, _ = self.run_tiered(pinned=[], cache_slots=10)
|
|
torch.testing.assert_close(out, self.expected)
|
|
|
|
def test_cpu_path_matches_dense_moe(self):
|
|
for mode in ("cpu", "auto"):
|
|
out, loads, cache = self.run_tiered([0], 3, Miss(mode=mode, cpu_max_tokens=1))
|
|
torch.testing.assert_close(out, self.expected, rtol=1e-5, atol=1e-5)
|
|
self.assertEqual(cache.cpu + cache.misses + cache.hits, 10)
|
|
self.assertEqual(cpu_experts(self.x, self.topk_weights, self.topk_ids, set(),
|
|
None).abs().sum().item(), 0.0)
|
|
|
|
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)
|
|
self.assertEqual(sorted(stats), [1, 2])
|
|
self.assertEqual(stats[2][0], 2)
|
|
self.assertAlmostEqual(stats[2][1], 1.5)
|
|
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_degrade_drops_unavailable_routes_of_draft_tokens_only(self):
|
|
weights = torch.tensor([[0.5, 0.3, 0.2], [0.5, 0.3, 0.2]])
|
|
ids = torch.tensor([[0, 1, 2], [0, 1, 2]], dtype=torch.int32)
|
|
draft = torch.tensor([False, True])
|
|
available = torch.tensor([False, True, False, False])
|
|
w, i = degrade(weights, ids, draft, available, min_weight=0.5)
|
|
# the draft token computes 0.3 of 1.0: it keeps expert 0 (heaviest) to reach 0.5
|
|
self.assertEqual(i.tolist(), [[0, 1, 2], [0, 1, -1]])
|
|
torch.testing.assert_close(w, torch.tensor([[0.5, 0.3, 0.2], [0.625, 0.375, 0.0]]))
|
|
w, i = degrade(weights, ids, draft, available, min_weight=0.0)
|
|
self.assertEqual(i[1].tolist(), [-1, 1, -1])
|
|
torch.testing.assert_close(w[1], torch.tensor([0.0, 1.0, 0.0]))
|
|
|
|
def test_draft_tokens_use_vram_and_what_full_tokens_load(self):
|
|
draft = torch.tensor([False, True, False, True, True, False])
|
|
out, loads, _ = self.run_tiered(pinned=[0], cache_slots=3, draft=draft)
|
|
# available: pinned expert 0 plus every expert of the full rows 0, 2 and 5
|
|
available = torch.zeros(EXPERTS, dtype=torch.bool)
|
|
available[[0, 1, 2, 4, 5, 6, 7, 8]] = True
|
|
self.assertTrue(torch.equal(
|
|
available_experts(self.topk_ids, draft, {0}, EXPERTS), available))
|
|
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)
|
|
torch.testing.assert_close(out[[0, 2, 5]], self.expected[[0, 2, 5]]) # full rows exact
|
|
self.assertEqual(self.seen_experts, list(range(EXPERTS))) # trace sees the routing
|
|
|
|
def test_drafts_are_exact_when_every_expert_is_in_vram(self):
|
|
out, loads, _ = self.run_tiered(pinned=list(range(EXPERTS)), cache_slots=0,
|
|
draft=torch.ones(6, dtype=torch.bool))
|
|
torch.testing.assert_close(out, self.expected)
|
|
self.assertEqual(loads, [])
|