Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / expert_parallel / static_splits.py: 100%
30 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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# ============================================================================
15"""Cache one *static* MoE routing plan so the dispatch can stop syncing every layer.
17Why this exists
18---------------
19A MoE dispatch has to give ``all_to_all_single`` Python split lists, which costs a
20device-to-host drain of the token counts: an extra counts all-to-all, its wait, and two
21``.tolist()`` calls, per layer, in *both* dispatch implementations
22(``core/expert_parallel/expert_parallel.py`` and
23``distributed/expert_parallel/experts.py``). A ``.tolist()`` waits for every op already
24enqueued on the stream -- including the collective it just issued -- so the host cannot
25enqueue the token all-to-all until the counts have landed.
27That serialisation is visible in the host API profile of the 2-SN step: **8.8 s blocked in
28``aclrtSynchronizeStream``** over a 25 s window (470 calls, average 18.7 ms, maximum
29237 ms).
31When the plan is *static* the answer never changes, so it can be computed once:
33* the benchmark's ``fix_router: true`` adapter assigns slots deterministically, so every
34 rank sends and receives exactly the same number of tokens to every peer, in every layer;
35* a real router reshuffles tokens every step, and reusing a stale plan there would silently
36 mis-route tokens.
38Hence the explicit opt-in (``HP_EP_STATIC_SPLITS=1``), a plan-shape-keyed cache, and a log
39line the first time a plan is reused. Benchmark-only, exactly like ``fix_router`` itself.
41BLOCKED -- measured, not usable yet
42-----------------------------------
43On the 2-SN GBS-256 config (``fix_router: true``, so the plan really is uniform: the log
44reports ``tokens/peer=512`` = 8192 * 8 / 128 exactly) the fast path is engaged and then the
45run dies in the **first forward**, in the loss-normalization all-reduce
46(``trainer/runtime/metrics.py:75``), with
48 Memory_Allocation_Failure(EL0004): Failed to allocate memory requested by HCCL module
50Measured: **4 of 4 full fast-path runs OOM** (r32_static_r1/r2, r33_static_r1/r2) while
51**0 of 2 control runs** did (min step 20.6625 s / 20.4525 s).
53The bisect (r34) names the responsible half -- and it is *not* the collective:
55=========================== ================================== ==================
56probe removes result
57=========================== ================================== ==================
58``HP_EP_STATIC_SPLITS_PROBE=a2a`` the two D2H drains **OOM** (24 hits)
59``HP_EP_STATIC_SPLITS_PROBE=drains`` the counts all-to-all passes, min 20.78 s
60=========================== ================================== ==================
62So the ``.tolist()`` drains are not merely a cost, they are the **host-side pacing**: they
63stall the enqueue once per layer, and removing them lets the host run far enough ahead that
64the live-tensor peak rises until the first forward's loss all-reduce fails to get its HCCL
65buffer. A drain is a throttle as much as an expense.
67Consequences:
69* removing only the collective (keeping the drains) is safe but **slower** (20.78 s vs the
70 20.45-20.66 s control band), so that half is not worth shipping on its own;
71* making the full fast path safe needs an explicit run-ahead bound, not just the removal --
72 e.g. a cheaper pacing point, a lower ``forward_prefetch_depth``, or computing the splits
73 one layer ahead so the drain is hidden instead of deleted.
75**Do not enable this switch on a real run until a pacing solution is measured**; it defaults
76to off precisely so a blocked lever cannot leak into the champion.
77"""
78from __future__ import annotations
80import logging
81import os
82from typing import Any, Hashable, Optional
84logger = logging.getLogger(__name__)
86__all__ = [
87 "clear_static_plans",
88 "probe_mode",
89 "get_static_plan",
90 "static_plan_key",
91 "static_splits_enabled",
92 "store_static_plan",
93]
95_ENABLED = os.environ.get("HP_EP_STATIC_SPLITS", "0") == "1"
96_PLANS: dict[Hashable, Any] = {}
97_LOGGED: set[Hashable] = set()
100def static_splits_enabled() -> bool:
101 """Return whether ``HP_EP_STATIC_SPLITS`` asked for the cached-plan fast path."""
102 return _ENABLED
105def static_plan_key(*parts: Any, device: Optional[Any] = None) -> Hashable:
106 """Build a cache key from the quantities that define a routing plan.
108 Args:
109 *parts: Shape/config values that pin the plan (token count, top-k, ep size, ...).
110 device: Device the plan's tensors live on, so one process cannot mix devices.
112 Returns:
113 A hashable key.
114 """
115 return (*parts, str(device))
118def get_static_plan(key: Hashable) -> Any:
119 """Return the cached plan for ``key``, or ``None`` when there is none.
121 Args:
122 key: Key from :func:`static_plan_key`.
124 Returns:
125 The stored value, or ``None``.
126 """
127 if not _ENABLED:
128 return None
129 return _PLANS.get(key)
132def store_static_plan(key: Hashable, value: Any, description: str = "") -> None:
133 """Cache a plan and log it once, so the fast path is never invisible.
135 Args:
136 key: Key from :func:`static_plan_key`.
137 value: The plan to reuse (split lists, counts tensor, ...).
138 description: Short human-readable note for the log line.
139 """
140 if not _ENABLED:
141 return
142 _PLANS[key] = value
143 if key not in _LOGGED:
144 _LOGGED.add(key)
145 logger.info(
146 "EP static splits: reusing one routing plan for every layer %s -- "
147 "benchmark-only, requires a static router such as fix_router=True",
148 description or f"(key={key})",
149 )
152def clear_static_plans() -> None:
153 """Drop every cached plan (used by tests and by shape changes)."""
154 _PLANS.clear()
155 _LOGGED.clear()
158def probe_mode() -> str:
159 """Return the bisect probe selected by ``HP_EP_STATIC_SPLITS_PROBE``.
161 The fast path removes two things at once -- the counts all-to-all and the two
162 ``.tolist()`` drains -- and a 4/4 OOM on the 2-SN config means one of them (or the
163 combination) raises the live-tensor peak. These modes remove exactly one half, so the
164 culprit can be named:
166 * ``"a2a"`` -- keep the collective, drop only the host drains;
167 * ``"drains"``-- keep the host drains, drop only the collective;
168 * ``""`` -- normal fast path (both removed).
170 Returns:
171 ``"a2a"``, ``"drains"`` or ``""``.
172 """
173 mode = os.environ.get("HP_EP_STATIC_SPLITS_PROBE", "").strip().lower()
174 return mode if mode in ("a2a", "drains") else ""