diff --git a/prain/experts.py b/prain/experts.py index 6fb7b3d..bc5eb04 100644 --- a/prain/experts.py +++ b/prain/experts.py @@ -225,12 +225,10 @@ class TieredRoutedExperts(RoutedExperts): "from ram to disk") ram = list(self.placement.ram) self._ram_row = {e: row for row, e in enumerate(ram)} - index = torch.tensor(ram, dtype=torch.int64, device="cpu") - self._ram_w13, self._ram_w2 = ( - torch.index_select(pack, 0, index, out=torch.empty( - (len(ram), *pack.shape[1:]), dtype=pack.dtype, device="cpu", pin_memory=True)) - for pack in (self._pack_w13, self._pack_w2) - ) + # One pinned tensor per expert: pinning a single 340 MiB block took 6-23 s on an + # eGPU host (hipHostMalloc), small blocks pin at the normal rate. + self._ram_w13 = [self._pack_w13[e].pin_memory() for e in ram] + self._ram_w2 = [self._pack_w2[e].pin_memory() for e in ram] for slot, expert in enumerate(self.placement.vram): self._load(expert, slot) if self._router and (self.prain.prefetch.mode == "pregate" or self.prain.trace_path): diff --git a/prain/sim.py b/prain/sim.py index 939e5bb..f39bdd1 100644 --- a/prain/sim.py +++ b/prain/sim.py @@ -31,9 +31,13 @@ class Cost: sync_ms: float = 0.05 # host sync per layer cpu_ms: float = 0.5 # CPU work per (token, expert) route + def read_ms(self): + """Reading one expert from the pack (the CPU path pays this for disk experts too).""" + return self.expert_mib * 1.048576 / self.disk_gbps + def copy_ms(self, disk): mb = self.expert_mib * 1.048576 - return mb / self.h2d_gbps + (mb / self.disk_gbps if disk else 0.0) + return mb / self.h2d_gbps + (self.read_ms() if disk else 0.0) def load_cost(path, expert_mib): @@ -175,7 +179,8 @@ def simulate(forwards, placements, policy, prefetch, miss, cost, belady=False): gpu = link land = max((ready.pop((layer, e), 0.0) for e in slot_of), default=0.0) gpu = max(gpu, land) + cost.round_ms + cost.route_ms * sum(counts[e] for e in slot_of) - cpu_end = t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu) + cpu_end = (t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu) + + cost.read_ms() * sum(e in disk[layer] for e in cpu)) later = [] if cpu and miss.admit: hot = sorted((e for e in cpu if e not in disk[layer]), key=lambda e: (-heat[e], e)) diff --git a/scripts/measure.py b/scripts/measure.py index 2fc4e6f..197bc39 100644 --- a/scripts/measure.py +++ b/scripts/measure.py @@ -73,15 +73,22 @@ def main(): other_ms = args.dense_gib * 2**30 / 1e6 / device_gbps / args.layers flag = torch.zeros(1, device="cuda") sync_ms = ms(lambda: flag.add_(1).item()) - w13_cpu = torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16).pin_memory() - w2_cpu = torch.randn(HIDDEN, INTER, dtype=torch.bfloat16).pin_memory() + # Serving computes a different expert each time, read from RAM (not from the CPU caches). + experts = [(torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16), + torch.randn(HIDDEN, INTER, dtype=torch.bfloat16)) for _ in range(23)] x = torch.randn(1, HIDDEN, dtype=torch.bfloat16) one = torch.ones(1, 1) ids = torch.zeros(1, 1, dtype=torch.int32) + + def cpu_route(weights): + cpu_experts(x, one, ids, {0}, lambda e: weights) + + for weights in experts[:3]: # warm up like ms(): thread pool and BLAS setup are one-time costs + cpu_route(weights) start = time.perf_counter() - for _ in range(20): - cpu_experts(x, one, ids, {0}, lambda e: (w13_cpu, w2_cpu)) - cpu_ms = (time.perf_counter() - start) * 1000 / 20 + for weights in experts[3:]: + cpu_route(weights) + cpu_ms = (time.perf_counter() - start) * 1000 / len(experts[3:]) print(json.dumps({"h2d_gbps": round(h2d_gbps, 2), "disk_gbps": round(disk_gbps, 2), "round_ms": round(round_ms, 4), "route_ms": round(route_ms, 6), "other_ms": round(other_ms, 4), "sync_ms": round(sync_ms, 4), diff --git a/tests/test_sim.py b/tests/test_sim.py index 6a6adff..111f89d 100644 --- a/tests/test_sim.py +++ b/tests/test_sim.py @@ -57,6 +57,14 @@ class Simulate(unittest.TestCase): copy = self.run_sim().ms_per_decode self.assertLess(self.run_sim(miss=Miss(mode="auto")).ms_per_decode, copy) + def test_cpu_path_pays_the_disk_read_of_disk_experts(self): + on_disk = {layer: LayerPlacement((), (), tuple(range(32)), 3) for layer in range(4)} + cpu = Miss(mode="cpu", admit=0) + ram = simulate(self.forwards, all_ram(3), Policy(), Prefetch(), cpu, Cost()) + disk = simulate(self.forwards, on_disk, Policy(), Prefetch(), cpu, Cost()) + self.assertAlmostEqual(disk.total_ms - ram.total_ms, ram.counters["cpu"] * Cost().read_ms(), + places=3) + def test_prefetch_helps_when_predictions_are_right(self): none = self.run_sim().hit_rate good = self.run_sim(prefetch=Prefetch(mode="pregate", budget=4))