Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / expert_parallel / routing.py: 88%
168 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
16"""expert_parallel.routing: MoE router adapters (05 §6.4.7, D-09).
18Routing semantics are model-specific (softmax/sigmoid, top-k normalization,
19scaling), while the expert MLP structure is uniform (SwiGLU). A factory picks
20an adapter from ``MOE_ROUTER_ADAPTERS`` BY NAME in its own code — the choice
21is explicit, never inferred. Each adapter maps
22``(module, hidden) -> (topk_idx [T,K] int64, topk_w [T,K] float)``.
24Benchmark-only balanced variants are derived from a real adapter with
25:func:`make_balanced_router` and registered as ``<adapter>_fixed``, so an
26archetype opts in with one flag (``fix_router``, see :func:`resolve_moe_router`)
27instead of carrying its own copy of the wrapper.
29Split out of components/distributed/ep_utils.py in stage 4e.
30"""
32import logging
33import math
34import os
35from typing import Optional
37import torch
39logger = logging.getLogger(__name__)
41# float32 represents every integer up to 2**24 exactly, so a smaller expert count
42# can use the AICore float sort (see ``experts._SORT_KEY_FP32_LIMIT``).
43_SORT_KEY_FP32_LIMIT = 1 << 24
45# Which slots inside an over-capacity expert survive (``HP_EP_CAPACITY_ORDER``):
46#
47# * ``score`` (default): the heaviest routing weights -- the contributions whose removal
48# perturbs the routed sum the least. A cap exists to bound memory, and spending the bound
49# on the cheapest slots is what makes "which slots die" a defined policy instead of an
50# accident of the sort order;
51# * ``first``: the first slots of the expert-major order, i.e. what the archived
52# ``HP_EP_CAPACITY_FACTOR`` sweeps (drop rate / peak / step time) were measured with.
53#
54# Both policies keep exactly the same number of slots per expert, which is the invariant the
55# memory bound rests on, so the choice cannot move a peak, a row count or a step time -- it
56# only decides whose contribution is lost.
57_CAPACITY_ORDER_SCORE = "score"
58_CAPACITY_ORDER_FIRST = "first"
59_CAPACITY_ORDERS = (_CAPACITY_ORDER_SCORE, _CAPACITY_ORDER_FIRST)
60_FP32_KEY_WARNED = False
63def _softmax_topk_router(module, hidden_states):
64 """default adapter: softmax -> topk -> normalize by sum (Mixtral/Qwen3 semantics).
66 top_k source: config.num_experts_per_tok / config.top_k / module.top_k (default 2);
67 normalization switch: config.norm_topk_prob (default True).
68 """
69 gate = getattr(module, "gate", None)
70 if gate is None:
71 gate = getattr(module, "router", None)
72 if gate is None:
73 raise AttributeError(
74 f"{type(module).__name__}: router not found (neither gate nor router "
75 "attribute exists); please register a custom MOE_ROUTER_ADAPTERS entry"
76 )
77 cfg = getattr(module, "config", None)
78 logits = gate(hidden_states)
79 logits = logits.view(-1, logits.shape[-1])
80 top_k = (getattr(cfg, "num_experts_per_tok", None)
81 or getattr(cfg, "top_k", None)
82 or getattr(module, "top_k", 2))
83 weights = logits.softmax(-1)
84 topk_w, topk_idx = weights.topk(int(top_k), dim=-1)
85 if getattr(cfg, "norm_topk_prob", True):
86 topk_w = topk_w / topk_w.sum(-1, keepdim=True).clamp_min(1e-20)
87 return topk_idx, topk_w
90def _topk_router_module(module, hidden_states):
91 """Adapt both tuple-returning TopKRouter and legacy linear gate modules."""
92 gate = getattr(module, "gate", None)
93 if gate is None:
94 gate = getattr(module, "router", None)
95 out = gate(hidden_states)
96 if isinstance(out, (tuple, list)) and len(out) == 3:
97 _, scores, indices = out
98 return indices, scores
99 if isinstance(out, torch.Tensor):
100 logits = out.reshape(-1, out.shape[-1])
101 config = getattr(module, "config", None)
102 top_k = (getattr(module, "top_k", None)
103 or getattr(config, "num_experts_per_tok", None)
104 or getattr(config, "top_k", None))
105 if top_k is None:
106 raise ValueError(f"{type(module).__name__}: linear router requires top_k")
107 weights = torch.softmax(logits, dim=-1, dtype=torch.float32)
108 scores, indices = weights.topk(int(top_k), dim=-1)
109 norm_topk_prob = getattr(module, "norm_topk_prob", None)
110 if norm_topk_prob is None:
111 norm_topk_prob = getattr(config, "norm_topk_prob", True)
112 if norm_topk_prob:
113 scores = scores / scores.sum(dim=-1, keepdim=True).clamp_min(1e-20)
114 return indices, scores.to(hidden_states.dtype)
115 raise TypeError(
116 f"{type(module).__name__}: router should return a Tensor or "
117 f"(logits, scores, indices), got {type(out).__name__}"
118 )
121def _mask_scores_by_group(scores, n_group, topk_group):
122 """Keep scores only in the highest-scoring expert groups."""
123 expert_count = scores.shape[-1]
124 group_scores = (
125 scores.view(-1, n_group, expert_count // n_group)
126 .topk(2, dim=-1)[0]
127 .sum(dim=-1)
128 )
129 group_indices = group_scores.topk(topk_group, dim=-1, sorted=False)[1]
130 group_mask = torch.zeros_like(group_scores)
131 group_mask.scatter_(1, group_indices, 1)
132 score_mask = (
133 group_mask.unsqueeze(-1)
134 .expand(-1, n_group, expert_count // n_group)
135 .reshape(-1, expert_count)
136 )
137 return scores.masked_fill(~score_mask.bool(), float("-inf"))
140def _sigmoid_group_router(module, hidden_states):
141 """deepseekv3/glm4moe adapter: sigmoid + e_score_correction_bias +
142 group-limited topk + (optional) normalization + routed_scaling_factor
143 (step-by-step consistent with HF DeepseekV3MoE.route_tokens_to_experts /
144 Glm4MoeMoE).
146 Parameter source: the module's own attributes take precedence
147 (n_group/topk_group/top_k/norm_topk_prob/routed_scaling_factor), falling
148 back to module.config; when n_group is missing or <=1, the group-limited
149 filter is skipped.
150 """
151 gate = getattr(module, "gate", None)
152 if gate is None:
153 gate = getattr(module, "router", None)
154 cfg = getattr(module, "config", None)
156 def _attr(name, default=None):
157 v = getattr(module, name, None)
158 if v is None and cfg is not None:
159 v = getattr(cfg, name, None)
160 return default if v is None else v
162 logits = gate(hidden_states)
163 if isinstance(logits, (tuple, list)):
164 logits = logits[0]
165 logits = logits.view(-1, logits.shape[-1]).float()
166 scores = logits.sigmoid()
167 bias = getattr(gate, "e_score_correction_bias", None)
168 scores_for_choice = scores + bias if bias is not None else scores
170 n_group = int(_attr("n_group", 0) or 0)
171 topk_group = int(_attr("topk_group", 0) or 0)
172 top_k = int(_attr("top_k", None) or _attr("num_experts_per_tok", 2))
173 if n_group > 1 and topk_group > 0:
174 scores_for_choice = _mask_scores_by_group(
175 scores_for_choice, n_group, topk_group
176 )
178 topk_idx = scores_for_choice.topk(top_k, dim=-1, sorted=False)[1]
179 topk_w = scores.gather(1, topk_idx)
180 if _attr("norm_topk_prob", False):
181 topk_w = topk_w / (topk_w.sum(dim=-1, keepdim=True) + 1e-20)
182 topk_w = topk_w * float(_attr("routed_scaling_factor", 1.0))
183 return topk_idx, topk_w
186def apply_capacity_limit(
187 topk_idx: "torch.Tensor",
188 capacity_factor: Optional[float],
189 num_experts: int,
190 weights: Optional["torch.Tensor"] = None,
191) -> tuple["torch.Tensor", "torch.Tensor"]:
192 """Capacity mask for a routed assignment (MoE token dropping).
194 Real routing is skewed: on the 2-SN benchmark the busiest rank has been measured at
195 **6x** the mean token count. A rank's expert buffers are sized by the rows it
196 receives, so that skew -- not the balanced case -- decides whether the step fits in
197 memory: at GBS 256 / seq 8192 the balanced peak is 41.5 GB while real routing reached
198 55.3 GiB and failed to allocate a 5.23 GiB expert output. Every *switch* that does not
199 change the routed row count (GBS, prefetch depth, ``swap_inputs``, ``overlap_shared_expert``,
200 even a 25% shorter sequence) bought at most 1.8 GB, because none of them bound the
201 busiest rank.
203 Capping each expert at ``capacity_factor x (T*K / E)`` slots bounds it directly, which is
204 the standard GShard/Switch/Megatron answer to routing skew. Tokens beyond capacity are
205 *dropped* (their contribution is removed), so this changes the model, exactly as it does
206 in those implementations; the returned drop rate is what makes the trade measurable.
208 The cap is what bounds memory, and it is policy-independent: every expert keeps exactly
209 ``min(count_e, capacity)`` slots. *Which* slots those are is decided by
210 ``HP_EP_CAPACITY_ORDER``:
212 * ``score`` (default): the ``capacity`` heaviest ``weights`` of each expert, so the
213 contributions dropped are the cheapest ones. Cheap and deterministic -- the expert index
214 and the inverted weight share a single fp32 sort key, keeping the AICore float sort (see
215 :func:`_expert_score_order`);
216 * ``first``: the first slots of the expert-major order (see :func:`_expert_major_order`).
217 Deterministic for a given input and backend as well, but not a policy with a meaning:
218 within one expert it is "not necessarily first in token order".
220 Args:
221 topk_idx: Routed expert per slot, shape ``[T, K]``.
222 capacity_factor: Multiplier on the mean load, or ``None`` to disable the cap.
223 num_experts: Total routed experts ``E``.
224 weights: Routing weight of every slot, shaped like ``topk_idx``. Only the ``score``
225 policy reads it, so a caller that passes ``None`` gets the ``first`` policy.
227 Returns:
228 ``(keep, dropped)`` where ``keep`` is a bool tensor shaped like ``topk_idx`` and
229 ``dropped`` is the number of dropped slots as a 0-dim tensor on the plan's device.
230 Draining it costs a host sync, so callers must not read it on the hot path -- the one
231 caller logs it once and then sparsely (see ``ep_routed_dispatch``).
233 Raises:
234 ValueError: If ``num_experts`` is not positive, or ``HP_EP_CAPACITY_ORDER`` is neither
235 ``score`` nor ``first``.
236 """
237 if capacity_factor is None or float(capacity_factor) <= 0:
238 return (
239 torch.ones_like(topk_idx, dtype=torch.bool),
240 torch.zeros((), dtype=torch.int64, device=topk_idx.device),
241 )
242 if num_experts <= 0:
243 raise ValueError(f"num_experts must be positive, got {num_experts}")
245 flat_expert = topk_idx.reshape(-1).to(torch.int64)
246 slots = flat_expert.numel()
247 capacity = int(math.ceil(float(capacity_factor) * slots / float(num_experts)))
248 keep_flat = torch.zeros_like(flat_expert, dtype=torch.bool)
249 if capacity > 0:
250 order = _slot_order(flat_expert, weights, num_experts)
251 counts = torch.bincount(flat_expert, minlength=num_experts)
252 starts = torch.cumsum(counts, dim=0) - counts
253 ranks = torch.arange(slots, device=flat_expert.device) - starts[flat_expert[order]]
254 keep_flat[order] = ranks < capacity
255 # ``dropped`` stays a device tensor on purpose: draining it costs a host sync, and this
256 # runs once per MoE layer, so callers that want the number must ask for it explicitly
257 # (they log it once, see ``ep_routed_dispatch``) instead of paying it every layer.
258 return keep_flat.view_as(topk_idx), (~keep_flat).sum()
261def _resolve_capacity_order() -> str:
262 """Resolve ``HP_EP_CAPACITY_ORDER`` for the current call.
264 Read per call rather than at import time, like the capacity factor itself, so the policy
265 can be switched without reloading the module.
267 Returns:
268 ``"score"`` or ``"first"``.
270 Raises:
271 ValueError: If the knob is neither; a typo must not silently keep a policy the
272 operator did not ask for.
273 """
274 raw = os.environ.get("HP_EP_CAPACITY_ORDER", _CAPACITY_ORDER_SCORE).strip().lower()
275 if raw not in _CAPACITY_ORDERS:
276 raise ValueError(
277 f"HP_EP_CAPACITY_ORDER must be one of {list(_CAPACITY_ORDERS)}, but got {raw!r}"
278 )
279 return raw
282def _slot_order(
283 flat_expert: "torch.Tensor",
284 flat_weights: Optional["torch.Tensor"],
285 num_experts: int,
286) -> "torch.Tensor":
287 """Order slots by expert, and under the ``score`` policy by descending weight as well."""
288 if flat_weights is None or _resolve_capacity_order() == _CAPACITY_ORDER_FIRST:
289 return _expert_major_order(flat_expert, num_experts)
290 return _expert_score_order(flat_expert, flat_weights, num_experts)
293def _expert_score_order(
294 flat_expert: "torch.Tensor",
295 flat_weights: "torch.Tensor",
296 num_experts: int,
297) -> "torch.Tensor":
298 """Expert-major order that puts the heaviest slot of every expert first.
300 Both keys travel in ONE fp32 sort key: the integer part is the expert index and the
301 fractional part carries the inverted, normalized weight, so a heavier slot always sorts
302 earlier and the fraction can never carry into the next expert. One key also means one sort
303 -- the AICore float sort the expert-major order already uses. A lexicographic or int64 key
304 would lose that kernel (the same reason ``HP_EP_SORT_FP32`` exists for the dispatch sort),
305 and a second pass is not needed either.
307 Resolution limit: fp32 spends the integer part's exponent on the expert index, so the
308 fractional step is about ``2**-15`` around ``num_experts = 2**8..2**9`` (and ``2**-11``
309 around ``2**12``). Weights closer together than that, relative to the plan's spread, tie
310 and their relative order is arbitrary -- the *count* kept per expert, which is what bounds
311 the busiest rank, is exact either way. ``HP_EP_CAPACITY_ORDER=first`` stays available for a
312 bit-exact reproduction of the archived sweeps.
314 Args:
315 flat_expert: Expert index of every routed slot, shape ``[T*K]``.
316 flat_weights: Routing weight of every slot, same shape.
317 num_experts: Size of the expert index space.
319 Returns:
320 An int64 permutation of ``range(flat_expert.numel())``.
322 Raises:
323 ValueError: If ``flat_weights`` does not carry exactly one weight per slot.
324 """
325 if flat_weights.numel() != flat_expert.numel():
326 raise ValueError(
327 f"weights must have one entry per slot: got {flat_weights.numel()} for "
328 f"{flat_expert.numel()} slots"
329 )
330 if not 0 < num_experts <= _SORT_KEY_FP32_LIMIT:
331 global _FP32_KEY_WARNED # pylint: disable=global-statement
332 if not _FP32_KEY_WARNED:
333 _FP32_KEY_WARNED = True
334 logger.warning(
335 "HP_EP_CAPACITY_ORDER=score needs the expert index to be exact in a float32 "
336 "key, but num_experts=%d is outside (0, 2**24]; keeping the expert-major order",
337 num_experts,
338 )
339 return _expert_major_order(flat_expert, num_experts)
340 weight = flat_weights.reshape(-1).to(torch.float32)
341 inverted = weight.max() - weight
342 # ``clamp_min`` rather than a host-side ``if span == 0``: reading the span would drain the
343 # device, and this runs once per MoE layer. A zero span leaves every ``inverted`` at zero,
344 # so the fraction stays zero and the expert-major order stands. ``0 * inf`` cannot happen:
345 # a zero span means ``inverted`` is exactly zero.
346 span = inverted.max().clamp_min(torch.finfo(weight.dtype).tiny)
347 # One multiply-add instead of three elementwise passes: ``expert + inverted * 0.5 / span``.
348 key = torch.addcmul(flat_expert.to(torch.float32), inverted, 0.5 / span)
349 return key.argsort()
352def _expert_major_order(flat_expert: "torch.Tensor", num_experts: int) -> "torch.Tensor":
353 """Return the permutation that groups slots by expert.
355 Integer ``argsort`` has no AICore kernel on this stack (the same reason
356 ``HP_EP_SORT_FP32`` exists for the dispatch sort), so keys are cast to float32 when the
357 expert count is exactly representable. Ties inside one expert may then be ordered
358 arbitrarily rather than by token: the *count* kept per expert -- which is what bounds the
359 busiest rank's memory -- is unaffected, and the result is still deterministic for a given
360 input and backend.
362 Args:
363 flat_expert: Expert index of every routed slot, shape ``[T*K]``.
364 num_experts: Size of the expert index space.
366 Returns:
367 An int64 permutation of ``range(flat_expert.numel())`` ordering slots by expert.
368 """
369 if 0 < num_experts <= _SORT_KEY_FP32_LIMIT:
370 return flat_expert.to(torch.float32).argsort()
371 return torch.argsort(flat_expert, stable=True)
374def _global_expert_count(module):
375 """Model-level routed expert count (mirrors ``experts._get_global_expert_count``).
377 Kept local so this module stays a torch-only leaf (it is imported by
378 model adapters that do not want the expert compute machinery).
379 """
380 for owner in (getattr(module, "experts", None), module):
381 count = getattr(owner, "num_experts", None)
382 if count is not None:
383 return int(count)
384 cfg = getattr(module, "config", None)
385 for name in ("num_experts", "n_routed_experts"):
386 count = getattr(cfg, name, None)
387 if count is not None:
388 return int(count)
389 raise ValueError(
390 f"{type(module).__name__}: cannot determine the global routed expert count"
391 )
394def _balanced_expert_indices(
395 token_count: int,
396 experts_per_token: int,
397 expert_count: int,
398 local_count: int,
399 *,
400 device,
401) -> "torch.Tensor":
402 """Round-robin the ``T*K`` slots over the DESTINATION ranks.
404 Local slot ``i`` goes to one of the ``L`` experts owned by destination
405 ``i mod Q`` (``Q = E/L``), so every rank receives exactly ``T*K/Q`` slots
406 instead of a data-dependent count.
408 Spreading by *rank* rather than by expert is what makes this exact: the
409 per-expert token count cannot be made uniform at all when ``E`` does not
410 divide ``T*K``, and a per-expert round robin (``i mod E``) replays the same
411 residue pattern on every rank, so the leftover slots pile onto the same
412 fixed group of ranks (measured: 0.8% above the mean on the busiest rank,
413 2.4% spread between rank groups). Going through the destination first
414 absorbs that leftover per rank instead.
416 Args:
417 token_count: Local tokens ``T``.
418 experts_per_token: Routed experts per token ``K``.
419 expert_count: Global routed experts ``E``.
420 local_count: Experts this rank owns ``L``.
421 device: Device for the returned index tensor.
423 Returns:
424 ``[T, K]`` int64 indices naming the expert each slot is routed to.
426 Raises:
427 ValueError: If ``E`` is not divisible by ``L``.
428 """
429 if local_count <= 0 or expert_count % local_count != 0:
430 raise ValueError(
431 f"num_experts ({expert_count}) must be divisible by the local "
432 f"expert count ({local_count})"
433 )
434 destinations = expert_count // local_count
435 slots = torch.arange(token_count * experts_per_token, device=device)
436 indices = (slots % destinations) * local_count + (slots // destinations) % local_count
437 return indices.view(token_count, experts_per_token)
440def make_balanced_router(real_router):
441 """Wrap ``real_router`` in a benchmark-only variant with a uniform expert load.
443 The real adapter still runs -- so its weights stay differentiable and the
444 router GEMM stays in the compute profile -- but the chosen experts are
445 replaced by a round-robin assignment that spreads the local slots over the
446 **destination ranks** (see :func:`_balanced_expert_indices`). Every rank
447 therefore receives exactly ``T*K/Q`` tokens instead of a data-dependent
448 count, whatever routing semantics ``real_router`` implements.
450 Why: unbalanced routing is the dominant source of step-time jitter on MoE
451 training runs -- the slowest rank in a step is the one whose experts drew
452 the most tokens, so identical configurations can differ by tens of percent
453 step to step. Measured on the 293B/18-layer config, the real router put
454 **6.3x** the mean token count on the busiest rank (leaving some experts
455 empty) while the balanced one stayed at 1.00x. Flattening the load makes a
456 step-to-step comparison converge in 2-3 steps instead of needing a dozen,
457 which is what makes small (1-3%) kernel-level effects measurable at all.
459 **Benchmark-only.** The assignment ignores the gate scores, so:
461 - ``loss`` / ``grad_norm`` / any convergence signal is meaningless here;
462 - expert hit distributions and "effective MFU" are optimistic (the load is
463 artificially flat), so numbers from this mode must not be reported as an
464 efficiency result.
466 Correctness must be validated with the switch OFF (``fix_router: False``),
467 where ``real_router`` runs unchanged. The true ``topk_w`` are kept so the
468 gate still receives gradient.
470 Args:
471 real_router: A ``(module, hidden) -> (topk_idx, topk_w)`` adapter, e.g.
472 any entry of :data:`MOE_ROUTER_ADAPTERS`.
474 Returns:
475 A ``(module, hidden) -> (topk_idx, topk_w)`` adapter with a uniform
476 destination load; its ``__name__`` records the wrapped adapter.
477 """
479 def balanced_router(module, hidden_states):
480 """Run ``real_router`` and replace its destinations round-robin."""
481 topk_idx, topk_w = real_router(module, hidden_states)
482 token_count, experts_per_token = topk_idx.shape
483 balanced_idx = _balanced_expert_indices(
484 int(token_count),
485 int(experts_per_token),
486 _global_expert_count(module),
487 # Set by the EP binder (bind_local_expert_forward) before any forward; a
488 # missing value degrades to a plain per-expert round robin.
489 int(getattr(getattr(module, "experts", None),
490 "local_expert_count", None) or 1),
491 device=topk_idx.device,
492 )
493 return balanced_idx.to(topk_idx.dtype), topk_w
495 # A callable object or a functools.partial has no ``__name__``; its class name is
496 # then the most informative thing the log line can carry.
497 real_name = getattr(real_router, "__name__", None) or type(real_router).__name__
498 balanced_router.__name__ = f"balanced_{real_name}"
499 balanced_router.__doc__ = (
500 f"Benchmark-only variant of ``{real_name}`` with a uniform destination "
501 "load (see :func:`make_balanced_router`)."
502 )
503 return balanced_router
506MOE_ROUTER_ADAPTERS = {
507 "softmax_topk": _softmax_topk_router,
508 "topk_router_module": _topk_router_module,
509 "sigmoid_group": _sigmoid_group_router,
510 "default": _softmax_topk_router,
511 "qwen2moe": _topk_router_module,
512 "qwen2_moe": _topk_router_module,
513 "qwen3moe": _topk_router_module,
514 "qwen3_moe": _topk_router_module,
515 "mixtral": _topk_router_module,
516 "deepseekv3": _sigmoid_group_router,
517 "deepseek_v3": _sigmoid_group_router,
518 "glm4moe": _sigmoid_group_router,
519 "glm4_moe": _sigmoid_group_router,
520}
522# One balanced wrapper per distinct real adapter, published as ``<name>_fixed`` for
523# every alias: the convention ``fix_router`` resolves through.
524_BALANCED_VARIANTS = {
525 adapter: make_balanced_router(adapter)
526 for adapter in dict.fromkeys(MOE_ROUTER_ADAPTERS.values())
527}
528# The archived benchmark router: the balanced wrapper over the sigmoid-group adapter,
529# kept as a module-level name because the recipes' docs and the earlier 2-SN sweeps
530# refer to it.
531_balanced_router = _BALANCED_VARIANTS[_sigmoid_group_router]
532for _alias, _adapter in list(MOE_ROUTER_ADAPTERS.items()):
533 MOE_ROUTER_ADAPTERS[f"{_alias}_fixed"] = _BALANCED_VARIANTS[_adapter]
534del _alias, _adapter
537def resolve_moe_router(router, fix_router: bool = False):
538 """Resolve a router adapter, optionally to its benchmark-only balanced variant.
540 ``fix_router`` is the switch every archetype factory exposes; it has to mean the
541 same thing everywhere, so it is resolved here rather than re-implemented per
542 factory (see :func:`make_balanced_router` for what it does and why its numbers are
543 not efficiency results).
545 Args:
546 router: An adapter callable, or the name of an entry in
547 :data:`MOE_ROUTER_ADAPTERS`.
548 fix_router: When true, return the balanced variant: the ``<name>_fixed``
549 registry entry for a name, or ``make_balanced_router(router)`` for a
550 callable.
552 Returns:
553 A ``(module, hidden) -> (topk_idx, topk_w)`` adapter.
555 Raises:
556 KeyError: If ``router`` is a name that is not registered.
557 TypeError: If ``router`` is neither a registry name nor callable.
558 """
559 if isinstance(router, str):
560 key = f"{router}_fixed" if fix_router else router
561 return MOE_ROUTER_ADAPTERS[key]
562 if not callable(router):
563 raise TypeError(
564 "router must be a registry name or a callable, got "
565 f"{type(router).__name__}"
566 )
567 return make_balanced_router(router) if fix_router else router