prain: tiered forward with CPU path
This commit is contained in:
+152
@@ -0,0 +1,152 @@
|
||||
"""Run one MoE layer when only some experts fit in VRAM. Needs torch, not vLLM."""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
HEAT_INDEX = {"count": 0, "weight": 1, "recency": 2}
|
||||
|
||||
|
||||
def route_stats(topk_weights, topk_ids, num_experts, recency_half_life):
|
||||
"""{expert: (tokens, router weight mass, recency-weighted mass)} for every routed expert,
|
||||
with one host sync. Later tokens weigh more for recency. Route id -1 counts for nothing."""
|
||||
tokens = topk_ids.shape[0]
|
||||
valid = (topk_ids >= 0).float()
|
||||
ids = topk_ids.clamp_min(0).long().flatten()
|
||||
position = torch.arange(tokens, device=topk_ids.device, dtype=torch.float32)
|
||||
decay = torch.pow(0.5, (tokens - 1 - position) / recency_half_life)[:, None]
|
||||
weight = topk_weights.float() * valid
|
||||
stats = torch.stack([
|
||||
torch.bincount(ids, weights=valid.flatten(), minlength=num_experts),
|
||||
torch.bincount(ids, weights=weight.flatten(), minlength=num_experts),
|
||||
torch.bincount(ids, weights=(weight * decay).flatten(), minlength=num_experts),
|
||||
]).tolist()
|
||||
return {e: (round(stats[0][e]), stats[1][e], stats[2][e])
|
||||
for e in range(num_experts) if stats[0][e] > 0}
|
||||
|
||||
|
||||
def draft_rows(spans, rows, device):
|
||||
"""Bool [rows]: True for the token rows of draft requests (spans from prain.drafts)."""
|
||||
mask = torch.zeros(rows, dtype=torch.bool)
|
||||
for first, end in spans:
|
||||
mask[first:min(end, rows)] = True
|
||||
return mask.to(device)
|
||||
|
||||
|
||||
def available_experts(topk_ids, draft, resident, num_experts):
|
||||
"""Bool [num_experts]: experts in VRAM, plus every expert a full (non-draft) token routes
|
||||
to, because those are copied in this forward anyway."""
|
||||
ids = topk_ids.long()
|
||||
full_routes = ((ids >= 0) & ~draft[:, None]).flatten().float()
|
||||
available = torch.bincount(ids.clamp_min(0).flatten(), weights=full_routes,
|
||||
minlength=num_experts) > 0
|
||||
if resident:
|
||||
available[torch.tensor(sorted(resident), device=ids.device)] = True
|
||||
return available
|
||||
|
||||
|
||||
def degrade(topk_weights, topk_ids, draft, available, min_weight):
|
||||
"""Let draft tokens skip experts that are not available.
|
||||
|
||||
A draft token keeps its heaviest unavailable routes only while less than min_weight of its
|
||||
router weight would be computed; its kept routes are scaled up to its full router weight.
|
||||
Dropped routes get id -1 and weight 0, which route_stats and remap ignore.
|
||||
draft: bool [tokens]; available: bool [num_experts]. Full tokens are left unchanged."""
|
||||
ids = topk_ids.long()
|
||||
valid = ids >= 0
|
||||
w = topk_weights.to(torch.promote_types(topk_weights.dtype, torch.float32)) * valid
|
||||
missing = valid & draft[:, None] & ~available[ids.clamp_min(0)]
|
||||
total = w.sum(1, keepdim=True)
|
||||
short = (min_weight * total - (w * ~missing).sum(1, keepdim=True)).clamp_min(0)
|
||||
order = (w * missing).argsort(1, descending=True) # heaviest missing route first
|
||||
heavy = torch.gather(w * missing, 1, order)
|
||||
keep_sorted = (heavy > 0) & (heavy.cumsum(1) - heavy < short) # still short before it
|
||||
keep = torch.zeros_like(keep_sorted).scatter(1, order, keep_sorted)
|
||||
drop = missing & ~keep
|
||||
kept = (w * ~drop).sum(1, keepdim=True)
|
||||
scale = torch.where(kept > 0, total / kept, torch.ones_like(kept))
|
||||
weights = torch.where(drop, torch.zeros_like(w), w * scale).to(topk_weights.dtype)
|
||||
return weights, torch.where(drop, torch.full_like(topk_ids, -1), topk_ids)
|
||||
|
||||
|
||||
def remap(topk_weights, topk_ids, slot_of, num_experts):
|
||||
"""Point routes at VRAM slots for one round.
|
||||
|
||||
Routes to experts outside this round (and padding routes with id -1) get slot 0
|
||||
and weight 0, so any fused-MoE kernel can run them unchanged."""
|
||||
table = torch.full((num_experts,), -1, dtype=torch.int64)
|
||||
if slot_of:
|
||||
table[torch.tensor(list(slot_of), dtype=torch.int64)] = torch.tensor(
|
||||
list(slot_of.values()), dtype=torch.int64
|
||||
)
|
||||
table = table.to(topk_ids.device)
|
||||
ids = topk_ids.long()
|
||||
slots = torch.where(ids >= 0, table[ids.clamp_min(0)], -1)
|
||||
weights = torch.where(slots >= 0, topk_weights, torch.zeros_like(topk_weights))
|
||||
return weights, slots.clamp_min(0).to(topk_ids.dtype)
|
||||
|
||||
|
||||
def cpu_experts(x, topk_weights, topk_ids, experts, weights_of):
|
||||
"""The routes to `experts`, computed on the CPU like vLLM's fused kernel with activation
|
||||
"gelu_tanh": w13 = [gate; up], y = (gelu_tanh(x @ gate.T) * (x @ up.T)) @ w2.T times the
|
||||
route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]). Returns [tokens, hidden] fp32."""
|
||||
out = torch.zeros(x.shape, dtype=torch.float32)
|
||||
for e in sorted(experts):
|
||||
rows, cols = (topk_ids == e).nonzero(as_tuple=True)
|
||||
if rows.numel() == 0:
|
||||
continue
|
||||
w13, w2 = weights_of(e)
|
||||
h = x[rows].to(w13.dtype) @ w13.t()
|
||||
half = w13.shape[0] // 2
|
||||
act = F.gelu(h[:, :half].float(), approximate="tanh") * h[:, half:].float()
|
||||
y = (act.to(w2.dtype) @ w2.t()).float()
|
||||
out.index_add_(0, rows, y * topk_weights[rows, cols].float()[:, None])
|
||||
return out
|
||||
|
||||
|
||||
def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight",
|
||||
recency_half_life=32.0, miss=None, cpu_weights=None, on_stats=None,
|
||||
after_enqueue=None, draft=None, min_weight=0.5):
|
||||
"""One MoE forward over VRAM slots.
|
||||
|
||||
cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights,
|
||||
slot_ids, used_slots) runs the stock fused-MoE kernel on the slot tensors. With a
|
||||
[miss] config other than "copy", the experts it picks run on the CPU (cpu_weights(e)
|
||||
gives their weights) while the GPU rounds run. on_stats(stats, tokens) sees the routing
|
||||
statistics; after_enqueue(cpu, stats) runs once all GPU work of this forward is queued.
|
||||
draft: bool [tokens] for the rows of draft requests, or None; see degrade().
|
||||
"""
|
||||
raw = None
|
||||
if draft is not None:
|
||||
if on_stats is not None: # the trace records the routing, not what drafts skipped
|
||||
raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
|
||||
available = available_experts(topk_ids, draft, set(cache.pinned) | set(cache.slots),
|
||||
cache.num_experts)
|
||||
topk_weights, topk_ids = degrade(topk_weights, topk_ids, draft, available, min_weight)
|
||||
stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) # host sync
|
||||
tokens = int(topk_ids.shape[0])
|
||||
if on_stats is not None:
|
||||
on_stats(raw if raw is not None else stats, tokens)
|
||||
counts = {e: s[0] for e, s in stats.items()}
|
||||
cpu = cache.choose_cpu(counts, miss) if miss is not None else set()
|
||||
if cpu: # copy the inputs now: later the copy would wait for the queued GPU rounds
|
||||
host = x.cpu(), topk_weights.cpu(), topk_ids.cpu()
|
||||
index = HEAT_INDEX[heat]
|
||||
outs = []
|
||||
for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu):
|
||||
for expert, slot in loads:
|
||||
load(expert, slot)
|
||||
if slot_of:
|
||||
weights, slots = remap(topk_weights, topk_ids, slot_of, cache.num_experts)
|
||||
outs.append(kernel(x, weights, slots, list(slot_of.values())))
|
||||
if after_enqueue is not None:
|
||||
after_enqueue(cpu, stats)
|
||||
if cpu: # runs while the GPU works through its rounds
|
||||
outs.append(cpu_experts(*host, cpu, cpu_weights).to(x.device))
|
||||
if not outs:
|
||||
return torch.zeros_like(x)
|
||||
if len(outs) == 1:
|
||||
return outs[0].to(x.dtype)
|
||||
# ponytail: every extra round re-runs all tokens with weight 0 for experts it does not
|
||||
# hold; only big batches need many rounds. Upgrade: compact tokens per round.
|
||||
acc = torch.promote_types(x.dtype, torch.float32)
|
||||
return torch.stack([o.to(acc) for o in outs]).sum(0).to(x.dtype)
|
||||
@@ -0,0 +1,128 @@
|
||||
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, [])
|
||||
Reference in New Issue
Block a user