523 lines
20 KiB
Python
523 lines
20 KiB
Python
"""Deterministic memory, KV, and network planner for the GLM-5.2 Max alpha route.
|
|
|
|
Everything here is arithmetic over the exact pinned artifact bytes and the exact
|
|
pinned architecture. There is no measurement, no probing, and no heuristic tuned
|
|
to a result — the planner is written *before* the target runs so that a later
|
|
story cannot discover a topology that "works" and then rationalise it.
|
|
|
|
Three ideas do the real work.
|
|
|
|
**Unified memory is one pool.** On an integrated-GPU machine the "VRAM" the driver
|
|
reports is carved out of the same physical DRAM the OS is already counting. Adding
|
|
them produces a node that appears to hold twice what it holds, and the failure mode
|
|
is not a clean admission rejection — it is an OOM or a swap-thrash halfway through
|
|
a 200 GiB load. :class:`NodeMemory` therefore refuses to be constructed from an
|
|
additive claim about one shared pool.
|
|
|
|
**The reserve is not optional headroom.** Weights plus KV are not the whole
|
|
resident cost: backend workspaces, quantization scratch, the graph plan, the
|
|
process, and the OS all live outside them, and the largest of those scale with the
|
|
backend rather than with the shard. Alpha reserves ``max(20% of physically usable
|
|
memory, 8 GiB)`` per node, and the *remainder* is the placement budget.
|
|
|
|
**Equal layer counts are not equal bytes.** Embeddings and the output head are
|
|
endpoint-only; three layers are dense and 75 are MoE; shared experts, indexer
|
|
tensors, and quant block alignment all skew the per-node share. Until DGR-018/019
|
|
report measured per-tensor placement, the planner carries an explicit
|
|
:data:`PLACEMENT_IMBALANCE_FACTOR` and reports the arithmetic minimum and the
|
|
recommended count as two separate numbers. The arithmetic minimum is a fit probe;
|
|
it is admissible only with exact measured placement evidence behind it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import Literal
|
|
|
|
from .manifest import GIB, ArchitectureSnapshot, TargetManifest
|
|
|
|
# Q8_0 stores 32 int8 quants plus one fp16 scale per block: 34 bytes / 32 values.
|
|
Q8_0_BYTES_PER_VALUE = 34 / 32
|
|
F16_BYTES_PER_VALUE = 2.0
|
|
|
|
KV_DTYPES: dict[str, float] = {
|
|
"Q8_0": Q8_0_BYTES_PER_VALUE,
|
|
"F16": F16_BYTES_PER_VALUE,
|
|
}
|
|
|
|
# The alpha KV configuration, locked by the roadmap.
|
|
ALPHA_KV_DTYPE = "Q8_0"
|
|
ALPHA_CONTEXT_TOKENS = 16384
|
|
ALPHA_CONCURRENCY = 1
|
|
|
|
# The reserve every node holds outside its weight-plus-KV placement budget.
|
|
RESERVE_FRACTION = 0.20
|
|
RESERVE_FLOOR_GIB = 8.0
|
|
|
|
# The aggregate runtime-accessible memory at which the artifact *just* fits.
|
|
# This is an experimental hard-fit floor, not an operational envelope: it has no
|
|
# room for a backend that allocates more scratch than another, and none for the
|
|
# imbalance below.
|
|
AGGREGATE_HARD_FIT_FLOOR_GIB = 224.0
|
|
|
|
# How much more than an equal share the worst-placed node is expected to hold.
|
|
# 1.10 is the roadmap's recommended-topology column expressed as arithmetic: it
|
|
# reproduces 10 / 6 / 5 / 3 / 3 nodes for the 32 / 48 / 64 / 96 / 128 GiB tiers.
|
|
# DGR-019 must replace it with measured per-tensor placement.
|
|
PLACEMENT_IMBALANCE_FACTOR = 1.10
|
|
|
|
# Alpha network floor. A link rate is a bandwidth claim, never a speed claim.
|
|
MIN_LINK_RATE_GBPS = 2.5
|
|
RECOMMENDED_LINK_RATE_GBPS = 10.0
|
|
|
|
BF16_BYTES = 2
|
|
DSA_SIDEBAND_INT32_BYTES = 4
|
|
|
|
IndexerLayout = Literal["optimized", "conservative"]
|
|
|
|
|
|
class ResourcePlanError(ValueError):
|
|
"""Raised when a node or route cannot be accounted for honestly."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class NodeMemory:
|
|
"""One node's physically usable memory, counted once.
|
|
|
|
``physical_usable_gib`` is what the node can actually place bytes into after
|
|
firmware and fixed carve-outs — not the marketing capacity, and not a sum of
|
|
two views of the same DRAM.
|
|
"""
|
|
|
|
name: str
|
|
physical_usable_gib: float
|
|
unified: bool
|
|
|
|
def __post_init__(self) -> None:
|
|
if not isinstance(self.name, str) or not self.name.strip():
|
|
raise ResourcePlanError("node name must be a non-empty physical-host identity")
|
|
if (
|
|
isinstance(self.physical_usable_gib, bool)
|
|
or not isinstance(self.physical_usable_gib, (int, float))
|
|
or not math.isfinite(self.physical_usable_gib)
|
|
or self.physical_usable_gib <= 0
|
|
):
|
|
raise ResourcePlanError(
|
|
f"node {self.name!r} must declare finite positive usable memory"
|
|
)
|
|
|
|
@classmethod
|
|
def from_host(
|
|
cls,
|
|
name: str,
|
|
*,
|
|
system_ram_gib: float,
|
|
gpu_memory_gib: float = 0.0,
|
|
unified: bool,
|
|
) -> "NodeMemory":
|
|
"""Build a node from a host's reported RAM and GPU memory.
|
|
|
|
On a unified machine the GPU memory *is* system RAM, so it is counted once
|
|
and never added. Passing a non-zero ``gpu_memory_gib`` alongside
|
|
``unified=True`` is the double-count this project has already decided is a
|
|
bug (RALPH-CONTEXT runtime decision 16), so it is rejected rather than
|
|
silently discarded: a caller who believes an integrated GPU adds memory has
|
|
a wrong model of the machine, and quietly ignoring the argument would let
|
|
that belief survive.
|
|
"""
|
|
if not isinstance(unified, bool):
|
|
raise ResourcePlanError(f"node {name!r} unified flag must be boolean")
|
|
for value, label, allow_zero in (
|
|
(system_ram_gib, "system RAM", False),
|
|
(gpu_memory_gib, "GPU memory", True),
|
|
):
|
|
if (
|
|
isinstance(value, bool)
|
|
or not isinstance(value, (int, float))
|
|
or not math.isfinite(value)
|
|
or value < 0
|
|
or (not allow_zero and value == 0)
|
|
):
|
|
qualifier = "finite non-negative" if allow_zero else "finite positive"
|
|
raise ResourcePlanError(f"node {name!r} must declare {qualifier} {label}")
|
|
if unified:
|
|
if gpu_memory_gib:
|
|
raise ResourcePlanError(
|
|
f"node {name!r} declares unified memory and {gpu_memory_gib} GiB of "
|
|
"separate GPU memory. Integrated-GPU memory is carved out of the same "
|
|
"physical DRAM as system RAM; adding them double-counts one pool. "
|
|
"Pass unified=True with system_ram_gib only."
|
|
)
|
|
usable = system_ram_gib
|
|
else:
|
|
usable = system_ram_gib + gpu_memory_gib
|
|
return cls(name=name, physical_usable_gib=usable, unified=unified)
|
|
|
|
@property
|
|
def reserve_gib(self) -> float:
|
|
"""``max(20% of physically usable memory, 8 GiB)``."""
|
|
return max(RESERVE_FRACTION * self.physical_usable_gib, RESERVE_FLOOR_GIB)
|
|
|
|
@property
|
|
def placement_budget_gib(self) -> float:
|
|
"""What remains for weights plus KV after the reserve."""
|
|
return self.physical_usable_gib - self.reserve_gib
|
|
|
|
|
|
def kv_bytes(
|
|
snapshot: ArchitectureSnapshot,
|
|
*,
|
|
context_tokens: int = ALPHA_CONTEXT_TOKENS,
|
|
concurrency: int = ALPHA_CONCURRENCY,
|
|
dtype: str = ALPHA_KV_DTYPE,
|
|
indexer_layout: IndexerLayout = "conservative",
|
|
include_indexer: bool = True,
|
|
) -> int:
|
|
"""Bytes of MLA (and DSA indexer) KV cache for the whole model.
|
|
|
|
``indexer_layout`` is the honest part. Correct DSA only needs indexer keys for
|
|
the Full producer layers, but the current experimental implementation may
|
|
allocate them across every backbone layer. Alpha budgets ``conservative``
|
|
(all 78) so that a route admitted by this planner cannot be surprised by the
|
|
implementation it actually gets.
|
|
"""
|
|
if (
|
|
not isinstance(context_tokens, int)
|
|
or isinstance(context_tokens, bool)
|
|
or context_tokens <= 0
|
|
or not isinstance(concurrency, int)
|
|
or isinstance(concurrency, bool)
|
|
or concurrency <= 0
|
|
):
|
|
raise ResourcePlanError("context_tokens and concurrency must be positive integers")
|
|
if dtype not in KV_DTYPES:
|
|
raise ResourcePlanError(
|
|
f"unsupported KV dtype {dtype!r}; alpha locks {ALPHA_KV_DTYPE} "
|
|
f"(known: {', '.join(sorted(KV_DTYPES))})"
|
|
)
|
|
bytes_per_value = KV_DTYPES[dtype]
|
|
|
|
layers = int(snapshot["num_hidden_layers"])
|
|
mla_values = int(snapshot["mla_cached_values_per_token_per_layer"])
|
|
total_values = mla_values * layers
|
|
|
|
if include_indexer:
|
|
if indexer_layout == "optimized":
|
|
indexer_layers = int(snapshot["indexer_full_layers"])
|
|
elif indexer_layout == "conservative":
|
|
indexer_layers = layers
|
|
else: # pragma: no cover - Literal keeps this unreachable from typed callers
|
|
raise ResourcePlanError(f"unknown indexer_layout {indexer_layout!r}")
|
|
total_values += int(snapshot["index_head_dim"]) * indexer_layers
|
|
|
|
return int(total_values * context_tokens * concurrency * bytes_per_value)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TopologyPlan:
|
|
"""The node count a homogeneous tier needs, and how it was reached."""
|
|
|
|
physical_usable_gib: float
|
|
reserve_gib: float
|
|
placement_budget_gib: float
|
|
weight_gib: float
|
|
kv_gib: float
|
|
total_placement_gib: float
|
|
arithmetic_minimum_nodes: int
|
|
recommended_nodes: int
|
|
imbalance_factor: float
|
|
|
|
@property
|
|
def is_arithmetic_minimum_topology(self) -> bool:
|
|
"""True when the recommendation offers no imbalance headroom at all."""
|
|
return self.recommended_nodes == self.arithmetic_minimum_nodes
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"physical_usable_gib": round(self.physical_usable_gib, 3),
|
|
"reserve_gib": round(self.reserve_gib, 3),
|
|
"placement_budget_gib": round(self.placement_budget_gib, 3),
|
|
"weight_gib": round(self.weight_gib, 3),
|
|
"kv_gib": round(self.kv_gib, 3),
|
|
"total_placement_gib": round(self.total_placement_gib, 3),
|
|
"arithmetic_minimum_nodes": self.arithmetic_minimum_nodes,
|
|
"recommended_nodes": self.recommended_nodes,
|
|
"imbalance_factor": self.imbalance_factor,
|
|
}
|
|
|
|
|
|
def plan_topology(
|
|
manifest: TargetManifest,
|
|
snapshot: ArchitectureSnapshot,
|
|
*,
|
|
physical_usable_gib: float,
|
|
context_tokens: int = ALPHA_CONTEXT_TOKENS,
|
|
concurrency: int = ALPHA_CONCURRENCY,
|
|
kv_dtype: str = ALPHA_KV_DTYPE,
|
|
indexer_layout: IndexerLayout = "conservative",
|
|
imbalance_factor: float = PLACEMENT_IMBALANCE_FACTOR,
|
|
) -> TopologyPlan:
|
|
"""Minimum and recommended node count for a homogeneous tier of this size."""
|
|
if (
|
|
isinstance(imbalance_factor, bool)
|
|
or not isinstance(imbalance_factor, (int, float))
|
|
or not math.isfinite(imbalance_factor)
|
|
or imbalance_factor < 1.0
|
|
):
|
|
raise ResourcePlanError(
|
|
"imbalance_factor must be finite and at least 1.0; a lower value would "
|
|
"assume the worst-placed node holds less than an equal share"
|
|
)
|
|
|
|
node = NodeMemory(
|
|
name=f"{physical_usable_gib:g}GiB-tier",
|
|
physical_usable_gib=physical_usable_gib,
|
|
unified=False,
|
|
)
|
|
budget = node.placement_budget_gib
|
|
if budget <= 0:
|
|
raise ResourcePlanError(
|
|
f"a {physical_usable_gib:g} GiB node has no placement budget after its "
|
|
f"{node.reserve_gib:.1f} GiB reserve"
|
|
)
|
|
|
|
weight_gib = manifest.total_bytes / GIB
|
|
kv_gib = (
|
|
kv_bytes(
|
|
snapshot,
|
|
context_tokens=context_tokens,
|
|
concurrency=concurrency,
|
|
dtype=kv_dtype,
|
|
indexer_layout=indexer_layout,
|
|
)
|
|
/ GIB
|
|
)
|
|
total = weight_gib + kv_gib
|
|
|
|
return TopologyPlan(
|
|
physical_usable_gib=physical_usable_gib,
|
|
reserve_gib=node.reserve_gib,
|
|
placement_budget_gib=budget,
|
|
weight_gib=weight_gib,
|
|
kv_gib=kv_gib,
|
|
total_placement_gib=total,
|
|
arithmetic_minimum_nodes=math.ceil(total / budget),
|
|
recommended_nodes=math.ceil(total * imbalance_factor / budget),
|
|
imbalance_factor=imbalance_factor,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RouteFit:
|
|
"""Whether a concrete, possibly heterogeneous set of nodes can hold the target."""
|
|
|
|
node_count: int
|
|
aggregate_usable_gib: float
|
|
aggregate_placement_budget_gib: float
|
|
required_placement_gib: float
|
|
fits: bool
|
|
meets_hard_fit_floor: bool
|
|
no_single_node_can_admit_target: bool
|
|
headroom_gib: float
|
|
reasons: tuple[str, ...]
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"node_count": self.node_count,
|
|
"aggregate_usable_gib": round(self.aggregate_usable_gib, 3),
|
|
"aggregate_placement_budget_gib": round(self.aggregate_placement_budget_gib, 3),
|
|
"required_placement_gib": round(self.required_placement_gib, 3),
|
|
"fits": self.fits,
|
|
"meets_hard_fit_floor": self.meets_hard_fit_floor,
|
|
"no_single_node_can_admit_target": self.no_single_node_can_admit_target,
|
|
"headroom_gib": round(self.headroom_gib, 3),
|
|
"reasons": list(self.reasons),
|
|
}
|
|
|
|
|
|
def plan_route(
|
|
manifest: TargetManifest,
|
|
snapshot: ArchitectureSnapshot,
|
|
nodes: list[NodeMemory],
|
|
*,
|
|
context_tokens: int = ALPHA_CONTEXT_TOKENS,
|
|
concurrency: int = ALPHA_CONCURRENCY,
|
|
kv_dtype: str = ALPHA_KV_DTYPE,
|
|
indexer_layout: IndexerLayout = "conservative",
|
|
) -> RouteFit:
|
|
"""Evaluate a concrete route. Every node's memory is already counted once."""
|
|
if len(nodes) < 2:
|
|
raise ResourcePlanError(
|
|
"the alpha target is distributed by definition; a route needs at least two "
|
|
"physical nodes"
|
|
)
|
|
names = [node.name for node in nodes]
|
|
if len(set(names)) != len(names):
|
|
raise ResourcePlanError(
|
|
"duplicate node names in the route; one physical machine counted twice is "
|
|
"the same double-count as adding integrated-GPU memory to system RAM"
|
|
)
|
|
|
|
weight_gib = manifest.total_bytes / GIB
|
|
kv_gib = (
|
|
kv_bytes(
|
|
snapshot,
|
|
context_tokens=context_tokens,
|
|
concurrency=concurrency,
|
|
dtype=kv_dtype,
|
|
indexer_layout=indexer_layout,
|
|
)
|
|
/ GIB
|
|
)
|
|
required = weight_gib + kv_gib
|
|
|
|
aggregate_usable = sum(node.physical_usable_gib for node in nodes)
|
|
aggregate_budget = sum(node.placement_budget_gib for node in nodes)
|
|
fits = aggregate_budget >= required
|
|
largest_budget = max(node.placement_budget_gib for node in nodes)
|
|
no_single_node = largest_budget < required
|
|
|
|
reasons: list[str] = []
|
|
if not fits:
|
|
reasons.append(
|
|
f"aggregate placement budget {aggregate_budget:.1f} GiB is below the "
|
|
f"{required:.1f} GiB the target needs after each node's reserve"
|
|
)
|
|
if not no_single_node:
|
|
reasons.append(
|
|
"at least one node could admit the complete target alone; that is a "
|
|
"single-host run, not distributed alpha"
|
|
)
|
|
if aggregate_usable < AGGREGATE_HARD_FIT_FLOOR_GIB:
|
|
reasons.append(
|
|
f"aggregate usable memory {aggregate_usable:.1f} GiB is below the "
|
|
f"{AGGREGATE_HARD_FIT_FLOOR_GIB:g} GiB experimental hard-fit floor"
|
|
)
|
|
|
|
return RouteFit(
|
|
node_count=len(nodes),
|
|
aggregate_usable_gib=aggregate_usable,
|
|
aggregate_placement_budget_gib=aggregate_budget,
|
|
required_placement_gib=required,
|
|
fits=fits,
|
|
meets_hard_fit_floor=aggregate_usable >= AGGREGATE_HARD_FIT_FLOOR_GIB,
|
|
no_single_node_can_admit_target=no_single_node,
|
|
headroom_gib=aggregate_budget - required,
|
|
reasons=tuple(reasons),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SeamPlan:
|
|
"""Bytes and latency across the activation seams of a route.
|
|
|
|
Bandwidth and latency are reported separately on purpose. Decode moves almost
|
|
nothing — 12 KiB per token per seam — so a faster link barely helps it. What
|
|
decode pays is *serial*: every generated token crosses every seam in order, so
|
|
the cost that matters is ``seams x per-hop latency``. A route that claims to be
|
|
fast because it is on 10 GbE has confused the two.
|
|
"""
|
|
|
|
node_count: int
|
|
seam_count: int
|
|
hidden_size: int
|
|
bytes_per_token_per_seam: int
|
|
prefill_bytes_per_seam: int
|
|
decode_bytes_per_seam_per_token: int
|
|
dsa_sideband_bytes_per_query: int
|
|
link_rate_gbps: float
|
|
meets_alpha_minimum: bool
|
|
is_recommended_link: bool
|
|
decode_serialization_ms_per_token: float
|
|
decode_latency_ms_per_token: float
|
|
decode_bandwidth_share_ms_per_token: float
|
|
prefill_serialization_ms: float
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"node_count": self.node_count,
|
|
"seam_count": self.seam_count,
|
|
"hidden_size": self.hidden_size,
|
|
"bytes_per_token_per_seam": self.bytes_per_token_per_seam,
|
|
"prefill_bytes_per_seam": self.prefill_bytes_per_seam,
|
|
"decode_bytes_per_seam_per_token": self.decode_bytes_per_seam_per_token,
|
|
"dsa_sideband_bytes_per_query": self.dsa_sideband_bytes_per_query,
|
|
"link_rate_gbps": self.link_rate_gbps,
|
|
"meets_alpha_minimum": self.meets_alpha_minimum,
|
|
"is_recommended_link": self.is_recommended_link,
|
|
"decode_serialization_ms_per_token": round(self.decode_serialization_ms_per_token, 4),
|
|
"decode_latency_ms_per_token": round(self.decode_latency_ms_per_token, 4),
|
|
"decode_bandwidth_share_ms_per_token": round(
|
|
self.decode_bandwidth_share_ms_per_token, 4
|
|
),
|
|
"prefill_serialization_ms": round(self.prefill_serialization_ms, 3),
|
|
}
|
|
|
|
|
|
def plan_seams(
|
|
snapshot: ArchitectureSnapshot,
|
|
*,
|
|
node_count: int,
|
|
context_tokens: int = ALPHA_CONTEXT_TOKENS,
|
|
link_rate_gbps: float = MIN_LINK_RATE_GBPS,
|
|
per_hop_latency_ms: float = 0.5,
|
|
) -> SeamPlan:
|
|
"""Model seam bytes, wire serialization, and serial per-hop latency separately."""
|
|
if not isinstance(node_count, int) or isinstance(node_count, bool) or node_count < 2:
|
|
raise ResourcePlanError("a seam exists only between two nodes")
|
|
if not isinstance(context_tokens, int) or isinstance(context_tokens, bool) or context_tokens <= 0:
|
|
raise ResourcePlanError("context_tokens must be a positive integer")
|
|
if (
|
|
isinstance(link_rate_gbps, bool)
|
|
or not isinstance(link_rate_gbps, (int, float))
|
|
or not math.isfinite(link_rate_gbps)
|
|
or link_rate_gbps <= 0
|
|
):
|
|
raise ResourcePlanError("link_rate_gbps must be finite and positive")
|
|
if (
|
|
isinstance(per_hop_latency_ms, bool)
|
|
or not isinstance(per_hop_latency_ms, (int, float))
|
|
or not math.isfinite(per_hop_latency_ms)
|
|
or per_hop_latency_ms < 0
|
|
):
|
|
raise ResourcePlanError("per_hop_latency_ms must be finite and non-negative")
|
|
|
|
hidden = int(snapshot["hidden_size"])
|
|
bytes_per_token = hidden * BF16_BYTES
|
|
seams = node_count - 1
|
|
|
|
bits_per_ms = link_rate_gbps * 1e9 / 1e3
|
|
decode_serialization_ms = (bytes_per_token * 8) / bits_per_ms
|
|
prefill_serialization_ms = (bytes_per_token * context_tokens * 8) / bits_per_ms
|
|
|
|
return SeamPlan(
|
|
node_count=node_count,
|
|
seam_count=seams,
|
|
hidden_size=hidden,
|
|
bytes_per_token_per_seam=bytes_per_token,
|
|
prefill_bytes_per_seam=bytes_per_token * context_tokens,
|
|
decode_bytes_per_seam_per_token=bytes_per_token,
|
|
dsa_sideband_bytes_per_query=int(snapshot["index_topk"]) * DSA_SIDEBAND_INT32_BYTES,
|
|
link_rate_gbps=link_rate_gbps,
|
|
meets_alpha_minimum=link_rate_gbps >= MIN_LINK_RATE_GBPS,
|
|
is_recommended_link=link_rate_gbps >= RECOMMENDED_LINK_RATE_GBPS,
|
|
decode_serialization_ms_per_token=decode_serialization_ms * seams,
|
|
decode_latency_ms_per_token=per_hop_latency_ms * seams,
|
|
decode_bandwidth_share_ms_per_token=decode_serialization_ms * seams,
|
|
prefill_serialization_ms=prefill_serialization_ms * seams,
|
|
)
|
|
|
|
|
|
ALPHA_TIERS_GIB: tuple[float, ...] = (32.0, 48.0, 64.0, 96.0, 128.0)
|
|
|
|
|
|
def plan_all_tiers(
|
|
manifest: TargetManifest, snapshot: ArchitectureSnapshot
|
|
) -> dict[str, TopologyPlan]:
|
|
"""The alpha tier table, recomputed from the pinned artifact and architecture."""
|
|
return {
|
|
f"{tier:g}": plan_topology(manifest, snapshot, physical_usable_gib=tier)
|
|
for tier in ALPHA_TIERS_GIB
|
|
}
|