From 9221156c32aa58e82f1441b2493d1b7f8f196e8b Mon Sep 17 00:00:00 2001 From: AlanX-Yan Date: Thu, 15 Oct 2026 00:22:13 +0800 Subject: [PATCH 1/4] [DSv4][P3-T08] Add versioned placement map supporting both physical-id semantics T08 needs a versioned bijection between the router's canonical logical_expert_id and the physical_expert_id recorded in the Envelope. Without one, nothing verifies that a recorded physical id corresponds to the logical id on the same row: a dispatch-table off-by-one would route a token to the wrong expert while every semantic comparison still passed. Two readings of physical_expert_id are in circulation and they disagree. assembler.assemble and checker.validate_plan treat it as a global permutation of range(E). The task contract defines it as the rank-local id under the current EP placement, which makes the bijection logical <-> (ep_rank, physical). With EP=8 and E=256 the rank-local ids are 0..31 on each rank, so sorted(mapping) is [0]*8 + [1]*8 + ... and the current check rejects a correct placement. This module does not pick a winner. A PlacementMap declares its semantics and validation applies the matching rule, so a map recorded under one reading can never be silently judged under the other, the fingerprint separates them, and whichever way the question is resolved the implementation already exists. Construction paths all return the same type: identity, adopt an existing flat list, explicit (ep_rank, physical) pairs, or generate block/interleave from an EP configuration. A global physical id is the expert's position in the rank-major concatenation of per-rank tables, the only numbering that agrees with the rank-local one about where an expert sits inside its own rank. plan_fragment extends the plan's placement_map additively: version and logical_to_physical keep their meaning, so consumers reading only those two are unaffected and no contract delta is needed. The Envelope rank is deliberately not compared against ep_rank. EP is not a token dimension, so a token recorded on one DP/CP/TP/PP rank routinely routes to experts hosted elsewhere; a control test pins this. Co-Authored-By: Claude Opus 5 --- docs/design/p3-placement-map.md | 240 +++++++++++++++ rl_engine/p3/placement.py | 502 ++++++++++++++++++++++++++++++++ tests/p3/test_placement.py | 418 ++++++++++++++++++++++++++ 3 files changed, 1160 insertions(+) create mode 100644 docs/design/p3-placement-map.md create mode 100644 rl_engine/p3/placement.py create mode 100644 tests/p3/test_placement.py diff --git a/docs/design/p3-placement-map.md b/docs/design/p3-placement-map.md new file mode 100644 index 000000000..023430447 --- /dev/null +++ b/docs/design/p3-placement-map.md @@ -0,0 +1,240 @@ +# P3 Placement Map: Design and Implementation Record + +| Field | Value | +|---|---| +| Document status | Delivered; one contract question open (Section 2) | +| Document date | 2026-10-15 | +| Work package | T08 (CP/DP/PP/EP ownership, placement, P7 `cross_config`) | +| Target | DSV4-Flash MoE router (E=256, K=6) | +| Test baseline | 53 passed, 2 skipped (native bitmath unavailable on the dev host) | +| Scope | The `logical <-> physical` expert placement map and its gates | +| Depends on | #498 (`P3Verdict`, `require`, `assembler`, `checker`) | + +> **What this is.** The design record for the expert placement map: the +> versioned bijection relating the router's canonical expert id to the +> physical id recorded in the Envelope, under *both* readings of +> "physical id" currently in circulation. Section 2 is the open question and +> is the part reviewers should read first. + +--- + +## 1. Why A Placement Map Exists + +`logical_expert_id` is the canonical route key: the router's semantic output, +which must not change when expert parallelism moves an expert to another +device. `physical_expert_id` is Envelope metadata describing where that expert +currently lives; it is audited by the artifact hash and excluded from the +semantic hash. + +Keeping them apart is the point. Without an explicit object relating the two, +nothing verifies that a recorded `physical_expert_id` actually corresponds to +the `logical_expert_id` on the same row: a dispatch-table off-by-one would +send a token to the wrong expert while every semantic comparison still passed, +because both sides would be wrong in the same way. + +The start kit already carries that relation as a flat list plus a version, and +`checker.validate_plan` already gates it. This module adds the EP-rank +dimension, generation from an EP configuration, a provenance fingerprint, +granular diagnostics, and — the substantive part — an explicit declaration of +which numbering a map was built under. + +## 2. Open Question: What Is `physical_expert_id`? + +Two readings are in circulation, and they disagree. + +**Global permutation.** `physical_expert_id` is unique across all `E` experts, +so the map is a permutation of `range(E)`. This is what `assembler.assemble` +builds by default and what `checker.validate_plan` enforces today: + +```python +require(len(mapping) == E and sorted(mapping) == list(range(E)), + P3Verdict.INVALID_PLACEMENT_MAP, "archived placement not bijective") +``` + +**Rank-local.** The task contract (§1.2) defines `physical_expert_id` as *"the +rank-local id under the current EP placement"*. Under that reading ids index +the hosting rank's own expert table, so they repeat across EP ranks and the +bijection is `logical <-> (ep_rank, physical)`. + +The two are not reconcilable. With `EP=8` and `E=256`, rank-local ids are +`0..31` on each of the eight ranks, so `sorted(mapping)` is +`[0]*8 + [1]*8 + ...`, and the current check rejects a correct placement. + +This module does not pick a winner. A map declares its `semantics` and +validation applies the matching rule, so: + +- a map recorded under one reading can never be silently judged under the + other (`test_rank_local_table_is_rejected_when_declared_global`); +- the fingerprint separates the two, so the readings cannot collide in + provenance; +- whichever way the question is resolved, the implementation already exists. + +P4 needs the answer too: it decides what a dispatch target id means. + +## 3. The Object + +``` +PlacementMap + version placement identity, e.g. "ep8-block-v1" (compared literally) + semantics global_permutation | rank_local + ep_size number of EP ranks + ep_rank_of indexed by logical_expert_id -> ep_rank + physical_of indexed by logical_expert_id -> physical_expert_id + strategy explicit | identity | block | interleave (provenance) + num_experts E + schema_version p3-placement.v1 +``` + +Construction paths, all returning the same type: + +| Constructor | Use | +|---|---| +| `identity()` | the single-rank default the assembler already builds | +| `from_logical_to_physical()` | adopt a plan already recorded under the flat list | +| `from_pairs()` | explicit `logical -> (ep_rank, physical)`, irregular placements | +| `from_ep_config()` | generate `block` / `interleave` from an EP configuration | +| `from_plan_fragment()` | rebuild from a RoutePlan's `placement_map` | + +A global physical id is defined as the expert's position in the rank-major +concatenation of all per-rank tables. That is the only numbering that agrees +with the rank-local one about where an expert sits *inside* its own rank +(`test_both_semantics_agree_on_position_within_a_rank`). For `block` it +collapses to the identity permutation; for `interleave` it is a non-trivial +one. + +## 4. Design Decisions + +### D1: Both semantics in one type, declared not inferred + +Inferring the reading from the data is possible for most tables and wrong for +the rest: a single-expert-per-rank layout satisfies both rules, and an +EP=1 placement is identical under either. An inferred reading would therefore +be stable in testing and wrong in production. The map carries its semantics; +validation never guesses. + +### D2: Construction never validates + +A map is a claim; the gates test it. `from_pairs` turns a missing logical id +into a `(-1, -1)` hole rather than raising, because a map that cannot be +constructed also cannot be *rejected with a verdict* — and the project's +fail-closed discipline requires defects to surface as verdict codes rather +than as exceptions in a caller's stack. It also lets each negative fixture +build exactly one broken map, which is how the gates are proved. + +`from_ep_config` is the exception: a non-positive shape or an ungeneratable +strategy raises `INVALID_PLACEMENT_MAP` immediately, because there is no map +to report on. + +### D3: Two verdicts, because the repair paths differ + +`PLACEMENT_MAP_VERSION_MISMATCH(65)` means the evidence was produced under a +*different* placement, so validating the table in hand proves nothing about +it and the fix is to fetch the right map. `INVALID_PLACEMENT_MAP(64)` means +this table is broken and the fix is to repair it. Collapsing them would route +both to the same wrong repair. + +`validate_plan_placement` therefore gates `version -> structure -> rows`, a +dependency order rather than a preference. + +### D4: Envelope `rank` is deliberately not compared with `ep_rank` + +This is the decision most likely to look like a missing check, so it is +recorded explicitly. + +The contract classifies EP as `not_a_token_dimension`. An Envelope `rank` is a +*token-dimension* coordinate — which DP/CP/TP/PP rank recorded this token — +while `ep_rank` is an *expert-dimension* coordinate. Under any combined +configuration a token recorded on one rank routinely routes to experts hosted +on another; that is normal operation. Comparing them would fail every +realistic configuration. Pinned as a control test, not left implicit. + +Token ownership is a separate concern and already has hooks in `mocks.py`. + +### D5: `plan_fragment` extends the plan additively + +`version` and `logical_to_physical` keep their meaning and position, so a +consumer reading only those two is unaffected. The extra keys +(`semantics`, `ep_size`, `logical_to_ep_rank`, `strategy`, +`placement_fingerprint`) are additive, so this needs no contract delta and no +schema major bump. + +A fragment *without* the extended keys is read as the global reading at +`ep_size=1` — that is what such a plan was recorded under, and assuming +anything else would reinterpret old evidence. + +### D6: Defects are collected, not short-circuited + +`placement_problems` returns every structural defect. The gate only needs to +know whether the list is empty, but a caller repairing a table wants the whole +list. Checks run in the order defects corrupt the table — shape, then range, +then the semantics-specific rule — and shape failures suppress per-entry +findings, which would otherwise be noise generated by the shape error itself. + +## 5. Test Matrix + +Every negative case injects exactly one defect. + +| Injected defect | Verdict | +|---|---| +| unmapped logical expert (gap) | 64 | +| two logical experts on one `(ep_rank, physical)` slot | 64 | +| repeated physical id under `global_permutation` | 64 | +| `ep_rank` outside `[0, ep_size)` | 64 | +| rank-local slots not `0..n-1` | 64 | +| table length disagrees with `num_experts` | 64 | +| rank table and physical table of different length | 64 | +| empty placement version | 64 | +| unknown semantics / unknown strategy | 64 | +| lookup of a logical id outside the table | 64 | +| declared version != map version | 65 | +| Envelope physical disagrees with the map | 64 | +| padding row carrying a physical id other than `-1` | 64 | +| Envelope version drift inside one plan | 65 | +| core and envelope of different length | 64 | +| plan recorded under `interleave`, judged against `block` | 64 | +| rank-local table declared as `global_permutation` | 64 | +| version mismatch **and** broken map | **65** (version precedes structure) | +| broken map **and** bad row | **64 structural** (structure precedes rows) | + +Controls (must **not** be flagged): Envelope `rank` differing from every +`ep_rank`; padding rows with `-1` on both sides; an irregular but bijective +explicit table; a plan fragment round-tripping without an external map. + +Interop: a generated global placement (`block` and `interleave`) is driven +through the real `assemble` and accepted by both `validate_plan_placement` and +the start kit's own `checker.validate_plan` — the compatibility claim is that +under the global reading nothing existing has to change. + +## 6. Boundary and Pending Work + +- Token ownership is not decided here. `mocks.py` holds the `OWNERSHIP` table + and the `materialize_tokens` / `validate_ownership` hooks; turning that + table into per-mode dispatch (including `pp: replica_by_layer` and + `ep: not_a_token_dimension`) is the next T08 change. +- `expected_token_set` generation from a parallel configuration is still T08. +- No P7 `cross_config` materializer is wired yet. +- If the rank-local reading wins, `assembler.assemble` and + `checker.validate_plan` need their bijection check routed through + `placement_problems` instead of `sorted(mapping) == list(range(E))`. That is + a one-line change in each, deliberately not made here while the question is + open. + +### Platform note + +Two integration tests skip where the native bitmath library cannot be loaded. +On this Windows development host `rl_engine/p3/provider.py` (`import fcntl`) +and `rl_engine/p3/bitmath.py` (ctypes load of a Unix build product) are not +importable at all, so the whole of `tests/p3` except this module fails +collection. The unit layer here builds plans synthetically and runs anywhere; +the interop test runs wherever the kit itself runs. + +## 7. Reproduction + +```bash +cd RL-Kernel # repo root +python -m pytest tests/p3/test_placement.py -q +# -> 53 passed, 2 skipped (skips are the native-bitmath interop cases) + +python -m ruff check rl_engine/p3/placement.py tests/p3/test_placement.py +python -m ruff format --check rl_engine/p3/placement.py tests/p3/test_placement.py +``` diff --git a/rl_engine/p3/placement.py b/rl_engine/p3/placement.py new file mode 100644 index 000000000..be156ae4d --- /dev/null +++ b/rl_engine/p3/placement.py @@ -0,0 +1,502 @@ +"""Versioned logical/physical expert placement map (T08). + +``logical_expert_id`` is the canonical route key: the router's semantic output, +which must not change when expert parallelism moves an expert to another +device. ``physical_expert_id`` is Envelope metadata describing where that +expert currently lives. This module owns the versioned bijection between them. + +Two readings of ``physical_expert_id`` are in circulation and this module +supports both explicitly, because they are not interchangeable and the choice +has not yet been settled: + +``global_permutation`` + ``physical_expert_id`` is unique across all ``E`` experts, so the map is a + permutation of ``range(E)``. This is what ``assembler.assemble`` builds by + default and what ``checker.validate_plan`` enforces today + (``sorted(mapping) == list(range(E))``). + +``rank_local`` + ``physical_expert_id`` indexes the hosting rank's own expert table, so ids + repeat across EP ranks and the bijection is + ``logical <-> (ep_rank, physical)``. This is the reading in the task + contract, which defines ``physical_expert_id`` as "the rank-local id under + the current EP placement". + +Under ``EP=8`` with ``E=256`` the two disagree: rank-local ids are ``0..31`` on +each of the eight ranks, so ``sorted(mapping)`` is ``[0]*8 + [1]*8 + ...`` and +the current bijection check would reject a correct placement. A map carries its +``semantics`` so validation can apply the matching rule instead of guessing, +and so a map recorded under one reading can never be silently judged under the +other. + +``plan_fragment`` extends the existing ``placement_map`` object additively: +``version`` and ``logical_to_physical`` keep their meaning and position, so +consumers that read only those two keep working unchanged. +""" + +from dataclasses import dataclass + +from .contract import E, P3Verdict +from .oracle import require +from .serialization import fingerprint + +PLACEMENT_SCHEMA = "p3-placement.v1" + +#: How ``physical_expert_id`` is numbered. Recorded, never inferred. +SEMANTICS_GLOBAL = "global_permutation" +SEMANTICS_RANK_LOCAL = "rank_local" +SEMANTICS = (SEMANTICS_GLOBAL, SEMANTICS_RANK_LOCAL) + +#: How the table was produced (provenance only; validation never branches here). +STRATEGY_EXPLICIT = "explicit" +STRATEGY_IDENTITY = "identity" +STRATEGY_BLOCK = "block" # contiguous chunk per rank +STRATEGY_INTERLEAVE = "interleave" # round robin: rank = logical % ep_size +STRATEGIES = (STRATEGY_EXPLICIT, STRATEGY_IDENTITY, STRATEGY_BLOCK, STRATEGY_INTERLEAVE) + +#: Padding rows carry expert -1 on both sides and are never looked up. +PADDING_EXPERT_ID = -1 + + +@dataclass(frozen=True) +class PlacementMap: + """A versioned expert placement under one declared numbering. + + ``ep_rank_of`` and ``physical_of`` are both indexed by + ``logical_expert_id``. Tuples keep the instance hashable and the + fingerprint order fixed. + + Construction never validates. A map is a claim; ``placement_problems`` and + the gates below are where the claim is tested. A map that cannot be built + also cannot be rejected with a verdict, and the negative fixtures need to + build broken maps in order to prove the gates reject them. + """ + + version: str + semantics: str + ep_size: int + ep_rank_of: tuple[int, ...] + physical_of: tuple[int, ...] + strategy: str = STRATEGY_EXPLICIT + num_experts: int = E + schema_version: str = PLACEMENT_SCHEMA + + # --- lookups --------------------------------------------------------- + + def locate(self, logical_expert_id): + """Return ``(ep_rank, physical_expert_id)`` for a logical expert.""" + require( + 0 <= logical_expert_id < len(self.physical_of), + P3Verdict.INVALID_PLACEMENT_MAP, + f"logical expert {logical_expert_id} outside [0, {len(self.physical_of)})", + ) + return self.ep_rank_of[logical_expert_id], self.physical_of[logical_expert_id] + + def ep_rank(self, logical_expert_id): + """Return the EP rank hosting a logical expert.""" + return self.locate(logical_expert_id)[0] + + def physical_expert_id(self, logical_expert_id): + """Return the physical id of a logical expert under this numbering.""" + return self.locate(logical_expert_id)[1] + + def experts_on_rank(self, ep_rank): + """Logical experts hosted by one EP rank, ordered by physical id.""" + owned = sorted( + (self.physical_of[logical], logical) + for logical, rank in enumerate(self.ep_rank_of) + if rank == ep_rank + ) + return tuple(logical for _, logical in owned) + + # --- interop --------------------------------------------------------- + + def logical_to_physical(self): + """The flat ``logical -> physical`` list used by the existing plan. + + Under ``rank_local`` this list contains repeats by construction, which + the current ``assembler``/``checker`` bijection check rejects. That is + the open question this module makes explicit rather than papering over. + """ + return list(self.physical_of) + + def plan_fragment(self): + """The ``placement_map`` object for a RoutePlan. + + Additive: ``version`` and ``logical_to_physical`` are unchanged, so a + consumer reading only those two is unaffected by the extra keys. + """ + return { + "version": self.version, + "logical_to_physical": self.logical_to_physical(), + "placement_schema_version": self.schema_version, + "semantics": self.semantics, + "ep_size": self.ep_size, + "logical_to_ep_rank": list(self.ep_rank_of), + "strategy": self.strategy, + "placement_fingerprint": self.fingerprint(), + } + + def fingerprint(self): + """Canonical hash over schema, identity, numbering, shape and table. + + ``semantics`` and ``strategy`` are included: two tables that happen to + coincide but were declared under different readings are different + provenance, and must not compare equal. + """ + return fingerprint( + { + "schema_version": self.schema_version, + "version": self.version, + "semantics": self.semantics, + "strategy": self.strategy, + "num_experts": self.num_experts, + "ep_size": self.ep_size, + "ep_rank_of": list(self.ep_rank_of), + "physical_of": list(self.physical_of), + } + ) + + # --- construction ---------------------------------------------------- + + @classmethod + def identity(cls, version="identity.v1", num_experts=E): + """The single-rank identity placement the assembler defaults to.""" + return cls( + version=version, + semantics=SEMANTICS_GLOBAL, + ep_size=1, + ep_rank_of=(0,) * num_experts, + physical_of=tuple(range(num_experts)), + strategy=STRATEGY_IDENTITY, + num_experts=num_experts, + ) + + @classmethod + def from_logical_to_physical( + cls, + version, + mapping, + ep_size=1, + ep_rank_of=None, + semantics=SEMANTICS_GLOBAL, + strategy=STRATEGY_EXPLICIT, + ): + """Adopt an existing flat ``logical -> physical`` list. + + This is the bridge from plans already recorded under the global + reading; ``ep_rank_of`` defaults to rank 0 for every expert, which is + accurate for the single-rank case and informational otherwise. + """ + physical = tuple(int(p) for p in mapping) + ranks = ( + tuple(int(r) for r in ep_rank_of) if ep_rank_of is not None else (0,) * len(physical) + ) + return cls( + version=version, + semantics=semantics, + ep_size=ep_size, + ep_rank_of=ranks, + physical_of=physical, + strategy=strategy, + num_experts=len(physical), + ) + + @classmethod + def from_pairs( + cls, + version, + pairs, + ep_size, + semantics=SEMANTICS_RANK_LOCAL, + strategy=STRATEGY_EXPLICIT, + num_experts=E, + ): + """Build from an explicit ``logical -> (ep_rank, physical)`` table. + + Logical ids missing from ``pairs`` become ``(-1, -1)`` holes so + ``placement_problems`` reports them as gaps. + """ + table = {int(k): (int(v[0]), int(v[1])) for k, v in dict(pairs).items()} + entries = [table.get(i, (-1, -1)) for i in range(num_experts)] + return cls( + version=version, + semantics=semantics, + ep_size=ep_size, + ep_rank_of=tuple(r for r, _ in entries), + physical_of=tuple(p for _, p in entries), + strategy=strategy, + num_experts=num_experts, + ) + + @classmethod + def from_ep_config( + cls, + version, + ep_size, + strategy=STRATEGY_BLOCK, + semantics=SEMANTICS_RANK_LOCAL, + num_experts=E, + ): + """Generate a regular placement; bijective under either semantics. + + ``block`` gives each rank a contiguous chunk, with the first + ``num_experts % ep_size`` ranks taking one extra expert; ``interleave`` + assigns round robin. + + A global physical id is the expert's position in the rank-major + concatenation of all per-rank tables, which is the only numbering that + agrees with the rank-local one on where an expert sits inside its own + rank. For ``block`` that collapses to the identity permutation. + """ + require( + num_experts > 0 and ep_size > 0, + P3Verdict.INVALID_PLACEMENT_MAP, + f"num_experts and ep_size must be positive, got {num_experts}/{ep_size}", + ) + require( + strategy in (STRATEGY_BLOCK, STRATEGY_INTERLEAVE), + P3Verdict.INVALID_PLACEMENT_MAP, + f"cannot generate strategy {strategy!r}", + ) + require( + semantics in SEMANTICS, + P3Verdict.INVALID_PLACEMENT_MAP, + f"unknown placement semantics {semantics!r}", + ) + + base, extra = divmod(num_experts, ep_size) + widths = [base + (1 if r < extra else 0) for r in range(ep_size)] + starts, offset = [], 0 + for width in widths: + starts.append(offset) + offset += width + + ranks, locals_ = [0] * num_experts, [0] * num_experts + for rank in range(ep_size): + if strategy == STRATEGY_BLOCK: + owned = range(starts[rank], starts[rank] + widths[rank]) + else: + owned = range(rank, num_experts, ep_size) + for local, logical in enumerate(owned): + ranks[logical] = rank + locals_[logical] = local + + if semantics == SEMANTICS_RANK_LOCAL: + physical = locals_ + else: + physical = [starts[ranks[i]] + locals_[i] for i in range(num_experts)] + + return cls( + version=version, + semantics=semantics, + ep_size=ep_size, + ep_rank_of=tuple(ranks), + physical_of=tuple(physical), + strategy=strategy, + num_experts=num_experts, + ) + + @classmethod + def from_plan_fragment(cls, fragment): + """Rebuild a map from a plan's ``placement_map`` object. + + A fragment without the extended keys is read as the global reading at + ``ep_size=1``: that is what such a plan was recorded under, and + assuming anything else would reinterpret old evidence. + """ + mapping = fragment["logical_to_physical"] + ranks = fragment.get("logical_to_ep_rank") + return cls.from_logical_to_physical( + fragment["version"], + mapping, + ep_size=int(fragment.get("ep_size", 1)), + ep_rank_of=ranks, + semantics=fragment.get("semantics", SEMANTICS_GLOBAL), + strategy=fragment.get("strategy", STRATEGY_EXPLICIT), + ) + + +# --- structural validation --------------------------------------------------- + + +def placement_problems(placement): + """Return every structural defect in a map (empty list == valid). + + Defects are collected rather than short-circuited: a caller repairing a + table wants the whole list, while the gate only needs to know whether the + list is empty. Ordered by how early a defect corrupts the table, so the + first entry is usually the root cause. + """ + problems = [] + + if not placement.version: + problems.append("placement version is empty") + if placement.semantics not in SEMANTICS: + problems.append(f"unknown semantics {placement.semantics!r}") + if placement.strategy not in STRATEGIES: + problems.append(f"unknown strategy {placement.strategy!r}") + if placement.num_experts <= 0: + problems.append(f"num_experts must be positive, got {placement.num_experts}") + if placement.ep_size <= 0: + problems.append(f"ep_size must be positive, got {placement.ep_size}") + if len(placement.physical_of) != placement.num_experts: + problems.append( + f"table holds {len(placement.physical_of)} physical entries but " + f"num_experts is {placement.num_experts}" + ) + if len(placement.ep_rank_of) != len(placement.physical_of): + problems.append( + f"{len(placement.ep_rank_of)} rank entries against " + f"{len(placement.physical_of)} physical entries" + ) + if problems: + return problems # shape is broken; per-entry findings would be noise + + gaps, out_of_range = [], [] + for logical, (rank, physical) in enumerate( + zip(placement.ep_rank_of, placement.physical_of, strict=True) + ): + if rank < 0 or physical < 0: + gaps.append(logical) + elif rank >= placement.ep_size: + out_of_range.append(logical) + if gaps: + problems.append(f"unmapped logical experts (gaps): {gaps[:16]}") + if out_of_range: + problems.append( + f"logical experts on an ep_rank outside [0, {placement.ep_size}): {out_of_range[:16]}" + ) + if problems: + return problems + + if placement.semantics == SEMANTICS_GLOBAL: + problems.extend(_global_problems(placement)) + else: + problems.extend(_rank_local_problems(placement)) + return problems + + +def _global_problems(placement): + """Under the global reading the physical ids are a permutation of range(E).""" + physical = sorted(placement.physical_of) + if physical == list(range(placement.num_experts)): + return [] + seen, duplicates = set(), [] + for value in placement.physical_of: + if value in seen: + duplicates.append(value) + seen.add(value) + missing = [v for v in range(placement.num_experts) if v not in seen] + detail = [] + if duplicates: + detail.append(f"repeated physical ids {sorted(set(duplicates))[:16]}") + if missing: + detail.append(f"physical ids never used {missing[:16]}") + outside = [v for v in placement.physical_of if v >= placement.num_experts] + if outside: + detail.append(f"physical ids >= num_experts {sorted(set(outside))[:16]}") + return [ + "global_permutation requires physical ids to be a permutation of " + f"[0, {placement.num_experts}): " + "; ".join(detail) + ] + + +def _rank_local_problems(placement): + """Under the rank-local reading each rank's ids are exactly ``0..n-1``.""" + problems = [] + claimed, slots_by_rank = {}, {} + for logical, (rank, physical) in enumerate( + zip(placement.ep_rank_of, placement.physical_of, strict=True) + ): + claimed.setdefault((rank, physical), []).append(logical) + slots_by_rank.setdefault(rank, []).append(physical) + + overlaps = {slot: ids for slot, ids in claimed.items() if len(ids) > 1} + if overlaps: + shown = sorted(overlaps.items())[:8] + problems.append( + "physical slots claimed by more than one logical expert: " + + ", ".join(f"(ep_rank={r}, physical={p})<-{ids}" for (r, p), ids in shown) + ) + + for rank in sorted(slots_by_rank): + slots = sorted(slots_by_rank[rank]) + if slots != list(range(len(slots))): + problems.append( + f"ep_rank {rank} physical slots are not 0..{len(slots) - 1} (got {slots[:16]}): " + "a rank-local id must index that rank's own expert table" + ) + return problems + + +# --- gates ------------------------------------------------------------------- + + +def validate_placement_map(placement): + """Fail closed on a structurally invalid map -> ``INVALID_PLACEMENT_MAP``.""" + problems = placement_problems(placement) + require(not problems, P3Verdict.INVALID_PLACEMENT_MAP, "; ".join(problems)) + + +def check_placement_version(declared_version, placement): + """Fail closed when evidence was produced under another placement. + + Separate from the structural gate because the repair paths differ: a + version mismatch means validating the table in hand proves nothing about + the evidence, and the fix is to fetch the right map. + """ + require( + declared_version == placement.version, + P3Verdict.PLACEMENT_MAP_VERSION_MISMATCH, + f"evidence declares placement {declared_version!r} but the map is {placement.version!r}", + ) + + +def validate_plan_placement(plan, placement=None): + """Check every Envelope physical id in a RoutePlan against the map. + + Gate order is version -> structure -> rows, a dependency order: rows cannot + be judged against a map from another placement, nor against a table that is + not a bijection under its own declared semantics. + + The Envelope ``rank`` is deliberately not compared with ``ep_rank``. Expert + parallelism is not a token dimension: the rank that recorded a token is a + token-dimension coordinate (DP/CP/TP/PP) while ``ep_rank`` is an + expert-dimension one, so under any combined configuration a token recorded + on one rank routinely routes to experts hosted on another. Token ownership + is a separate concern with its own hooks. + """ + fragment = plan["placement_map"] + placement = PlacementMap.from_plan_fragment(fragment) if placement is None else placement + check_placement_version(fragment["version"], placement) + validate_placement_map(placement) + + rows, envelopes = plan["core"], plan["envelope"] + require( + len(rows) == len(envelopes), + P3Verdict.INVALID_PLACEMENT_MAP, + f"{len(rows)} core rows against {len(envelopes)} envelopes", + ) + for row, (core, env) in enumerate(zip(rows, envelopes, strict=True)): + require( + env["placement_map_version"] == placement.version, + P3Verdict.PLACEMENT_MAP_VERSION_MISMATCH, + f"row {row}: envelope placement version mismatch", + ) + expert = core["logical_expert_id"] + if expert == PADDING_EXPERT_ID: + require( + env["physical_expert_id"] == PADDING_EXPERT_ID, + P3Verdict.INVALID_PLACEMENT_MAP, + f"row {row}: padding carries physical {env['physical_expert_id']}, expected -1", + ) + continue + expected = placement.physical_expert_id(expert) + require( + env["physical_expert_id"] == expected, + P3Verdict.INVALID_PLACEMENT_MAP, + f"row {row}: logical expert {expert} maps to physical {expected} under " + f"{placement.version!r} ({placement.semantics}), " + f"envelope records {env['physical_expert_id']}", + ) diff --git a/tests/p3/test_placement.py b/tests/p3/test_placement.py new file mode 100644 index 000000000..6344a921c --- /dev/null +++ b/tests/p3/test_placement.py @@ -0,0 +1,418 @@ +"""Placement map tests (T08): both physical-id semantics, gates, and interop. + +The unit layer builds plans synthetically so it runs anywhere. One integration +test drives the real ``assemble`` and the start-kit ``checker`` to prove the +global-permutation path stays byte-compatible with what the kit already +produces; it is skipped where the native bitmath library cannot be loaded. + +Every negative case injects exactly one defect, so a correct verdict cannot be +a coincidence of two interacting faults. +""" + +from dataclasses import replace + +import pytest + +from rl_engine.p3.contract import E, P3Error, P3Verdict +from rl_engine.p3.placement import ( + PADDING_EXPERT_ID, + SEMANTICS_GLOBAL, + SEMANTICS_RANK_LOCAL, + STRATEGY_BLOCK, + STRATEGY_EXPLICIT, + STRATEGY_IDENTITY, + STRATEGY_INTERLEAVE, + PlacementMap, + check_placement_version, + placement_problems, + validate_placement_map, + validate_plan_placement, +) + +VERSION = "ep8-block-v1" +SEMANTIC_IDS = [SEMANTICS_GLOBAL, SEMANTICS_RANK_LOCAL] + + +def _map(ep_size=8, strategy=STRATEGY_BLOCK, semantics=SEMANTICS_RANK_LOCAL, version=VERSION): + return PlacementMap.from_ep_config(version, ep_size, strategy, semantics) + + +def _plan(placement, *, tokens=3, padding=1): + """A minimal RoutePlan carrying only what the placement gates read.""" + core, envelopes, row = [], [], 0 + for token in range(tokens): + for slot in range(6): + logical = (token * 6 + slot) % placement.num_experts + core.append({"logical_expert_id": logical, "global_token_id": token}) + envelopes.append( + { + "source_row": row, + "physical_expert_id": placement.physical_expert_id(logical), + "placement_map_version": placement.version, + "rank": 3, # a token-dimension rank, unrelated to ep_rank + } + ) + row += 1 + for _ in range(padding): + for _slot in range(6): + core.append({"logical_expert_id": PADDING_EXPERT_ID, "global_token_id": -1}) + envelopes.append( + { + "source_row": row, + "physical_expert_id": PADDING_EXPERT_ID, + "placement_map_version": placement.version, + "rank": 3, + } + ) + row += 1 + return {"core": core, "envelope": envelopes, "placement_map": placement.plan_fragment()} + + +# --- generation -------------------------------------------------------------- + + +@pytest.mark.parametrize("semantics", SEMANTIC_IDS) +@pytest.mark.parametrize("strategy", [STRATEGY_BLOCK, STRATEGY_INTERLEAVE]) +@pytest.mark.parametrize("ep_size", [1, 2, 4, 8]) +def test_generated_placements_are_valid(semantics, strategy, ep_size): + placement = _map(ep_size, strategy, semantics) + assert placement_problems(placement) == [] + assert sum(len(placement.experts_on_rank(r)) for r in range(ep_size)) == E + assert len(placement.experts_on_rank(0)) == E // ep_size + + +@pytest.mark.parametrize("semantics", SEMANTIC_IDS) +def test_uneven_division_spreads_the_remainder(semantics): + placement = PlacementMap.from_ep_config("v", 4, STRATEGY_BLOCK, semantics, num_experts=10) + assert placement_problems(placement) == [] + assert [len(placement.experts_on_rank(r)) for r in range(4)] == [3, 3, 2, 2] + assert placement.experts_on_rank(0) == (0, 1, 2) + assert placement.experts_on_rank(3) == (8, 9) + + +def test_rank_local_ids_repeat_across_ranks(): + placement = _map(8, STRATEGY_BLOCK, SEMANTICS_RANK_LOCAL) + assert placement.locate(0) == (0, 0) + assert placement.locate(32) == (1, 0) # same physical id, different rank + assert max(placement.physical_of) == E // 8 - 1 + + +def test_global_ids_are_unique_across_ranks(): + placement = _map(8, STRATEGY_BLOCK, SEMANTICS_GLOBAL) + assert sorted(placement.physical_of) == list(range(E)) + # Rank-major concatenation: a contiguous block layout is the identity. + assert list(placement.physical_of) == list(range(E)) + + +def test_interleave_global_is_a_nontrivial_permutation(): + placement = PlacementMap.from_ep_config("v", 4, STRATEGY_INTERLEAVE, SEMANTICS_GLOBAL, 10) + assert list(placement.physical_of) == [0, 3, 6, 8, 1, 4, 7, 9, 2, 5] + assert sorted(placement.physical_of) == list(range(10)) + + +def test_both_semantics_agree_on_position_within_a_rank(): + local = _map(8, STRATEGY_INTERLEAVE, SEMANTICS_RANK_LOCAL) + glob = _map(8, STRATEGY_INTERLEAVE, SEMANTICS_GLOBAL) + for rank in range(8): + assert local.experts_on_rank(rank) == glob.experts_on_rank(rank) + + +def test_generation_rejects_bad_arguments(): + for kwargs in ( + {"ep_size": 0}, + {"num_experts": 0}, + {"strategy": STRATEGY_EXPLICIT}, + {"semantics": "guesswork"}, + ): + with pytest.raises(P3Error) as excinfo: + PlacementMap.from_ep_config("v", **{"ep_size": 8, **kwargs}) + assert excinfo.value.verdict == P3Verdict.INVALID_PLACEMENT_MAP + + +# --- explicit tables --------------------------------------------------------- + + +def test_identity_matches_the_assembler_default(): + placement = PlacementMap.identity() + assert placement_problems(placement) == [] + assert placement.strategy == STRATEGY_IDENTITY + assert placement.semantics == SEMANTICS_GLOBAL + assert placement.logical_to_physical() == list(range(E)) + + +def test_explicit_pairs_accept_an_irregular_but_bijective_table(): + pairs = {0: (1, 0), 1: (0, 0), 2: (1, 1), 3: (0, 1)} + placement = PlacementMap.from_pairs("irregular-v1", pairs, 2, num_experts=4) + assert placement_problems(placement) == [] + assert placement.experts_on_rank(0) == (1, 3) + assert placement.experts_on_rank(1) == (0, 2) + + +def test_lookup_outside_the_table_fails_closed(): + with pytest.raises(P3Error) as excinfo: + _map().locate(E) + assert excinfo.value.verdict == P3Verdict.INVALID_PLACEMENT_MAP + + +# --- the ambiguity this module makes explicit -------------------------------- + + +def test_rank_local_table_is_rejected_when_declared_global(): + """The open contract question, pinned as a test. + + A rank-local table judged under the global reading fails, and vice versa. + Neither reading is assumed: the map declares which one it was built under. + """ + rank_local = _map(8, STRATEGY_BLOCK, SEMANTICS_RANK_LOCAL) + assert placement_problems(rank_local) == [] + + mislabeled = replace(rank_local, semantics=SEMANTICS_GLOBAL) + problems = placement_problems(mislabeled) + assert problems and "permutation" in problems[0] + + glob = _map(8, STRATEGY_BLOCK, SEMANTICS_GLOBAL) + assert placement_problems(replace(glob, semantics=SEMANTICS_RANK_LOCAL)) + + +def test_fingerprint_separates_the_two_readings(): + local = _map(8, STRATEGY_BLOCK, SEMANTICS_RANK_LOCAL) + glob = _map(8, STRATEGY_BLOCK, SEMANTICS_GLOBAL) + assert local.fingerprint() != glob.fingerprint() + + +def test_fingerprint_separates_version_strategy_and_table(): + base = _map() + assert base.fingerprint() == _map().fingerprint() + assert base.fingerprint() != _map(version="ep8-block-v2").fingerprint() + assert base.fingerprint() != _map(strategy=STRATEGY_INTERLEAVE).fingerprint() + assert base.fingerprint() != _map(ep_size=4).fingerprint() + + +# --- structural gate --------------------------------------------------------- + + +def _assert_invalid(placement, needle): + problems = placement_problems(placement) + assert problems, "expected a structural defect" + assert any(needle in p for p in problems), problems + with pytest.raises(P3Error) as excinfo: + validate_placement_map(placement) + assert excinfo.value.verdict == P3Verdict.INVALID_PLACEMENT_MAP + + +def test_gap(): + pairs = {i: (0, i) for i in range(4)} + del pairs[2] # one hole, nothing else touched + _assert_invalid(PlacementMap.from_pairs("v", pairs, 1, num_experts=4), "gaps") + + +def test_overlap_rank_local(): + pairs = {0: (0, 0), 1: (0, 1), 2: (0, 1), 3: (0, 2)} # 1 and 2 collide + _assert_invalid( + PlacementMap.from_pairs("v", pairs, 1, num_experts=4), "more than one logical expert" + ) + + +def test_duplicate_physical_under_global(): + pairs = {0: (0, 0), 1: (0, 1), 2: (0, 1), 3: (0, 2)} + placement = PlacementMap.from_pairs("v", pairs, 1, semantics=SEMANTICS_GLOBAL, num_experts=4) + _assert_invalid(placement, "repeated physical ids") + + +def test_ep_rank_out_of_range(): + pairs = {0: (0, 0), 1: (0, 1), 2: (0, 2), 3: (9, 0)} # ep_size is 1 + _assert_invalid(PlacementMap.from_pairs("v", pairs, 1, num_experts=4), "outside [0, 1)") + + +def test_noncontiguous_rank_local_slots(): + pairs = {0: (0, 0), 1: (0, 1), 2: (0, 2), 3: (0, 7)} + _assert_invalid(PlacementMap.from_pairs("v", pairs, 1, num_experts=4), "not 0..3") + + +def test_table_length_disagrees_with_num_experts(): + _assert_invalid(replace(_map(), num_experts=E + 1), "num_experts is 257") + + +def test_rank_and_physical_tables_must_align(): + placement = _map() + _assert_invalid(replace(placement, ep_rank_of=placement.ep_rank_of[:-1]), "rank entries") + + +def test_empty_version(): + _assert_invalid(_map(version=""), "version is empty") + + +def test_unknown_semantics(): + _assert_invalid(replace(_map(), semantics="guesswork"), "unknown semantics") + + +def test_unknown_strategy(): + _assert_invalid(replace(_map(), strategy="mystery"), "unknown strategy") + + +# --- version gate ------------------------------------------------------------ + + +def test_version_agreement_passes(): + check_placement_version(VERSION, _map()) + + +def test_version_mismatch_fails_closed(): + with pytest.raises(P3Error) as excinfo: + check_placement_version("ep8-block-v2", _map()) + assert excinfo.value.verdict == P3Verdict.PLACEMENT_MAP_VERSION_MISMATCH + + +# --- plan gate --------------------------------------------------------------- + + +@pytest.mark.parametrize("semantics", SEMANTIC_IDS) +def test_wellformed_plan_passes(semantics): + placement = _map(8, STRATEGY_BLOCK, semantics) + validate_plan_placement(_plan(placement), placement) + + +@pytest.mark.parametrize("semantics", SEMANTIC_IDS) +def test_plan_fragment_round_trips_without_an_external_map(semantics): + placement = _map(8, STRATEGY_INTERLEAVE, semantics) + plan = _plan(placement) + validate_plan_placement(plan) # map rebuilt from the fragment + rebuilt = PlacementMap.from_plan_fragment(plan["placement_map"]) + assert rebuilt.fingerprint() == placement.fingerprint() + + +def test_envelope_rank_may_differ_from_ep_rank(): + """Control: EP is not a token dimension, so this must not be flagged.""" + placement = _map(8, STRATEGY_BLOCK, SEMANTICS_RANK_LOCAL) + plan = _plan(placement) + assert all(env["rank"] == 3 for env in plan["envelope"]) + active = [c for c in plan["core"] if c["logical_expert_id"] != PADDING_EXPERT_ID] + assert {placement.ep_rank(c["logical_expert_id"]) for c in active} != {3} + validate_plan_placement(plan, placement) + + +def test_row_physical_disagreeing_with_the_map(): + placement = _map() + plan = _plan(placement) + plan["envelope"][7]["physical_expert_id"] += 1 + with pytest.raises(P3Error) as excinfo: + validate_plan_placement(plan, placement) + assert excinfo.value.verdict == P3Verdict.INVALID_PLACEMENT_MAP + assert "row 7" in str(excinfo.value) + + +def test_padding_row_must_carry_minus_one(): + placement = _map() + plan = _plan(placement) + plan["envelope"][-1]["physical_expert_id"] = 0 + with pytest.raises(P3Error) as excinfo: + validate_plan_placement(plan, placement) + assert excinfo.value.verdict == P3Verdict.INVALID_PLACEMENT_MAP + assert "padding" in str(excinfo.value) + + +def test_envelope_version_drift_inside_a_plan(): + placement = _map() + plan = _plan(placement) + plan["envelope"][2]["placement_map_version"] = "ep8-block-v2" + with pytest.raises(P3Error) as excinfo: + validate_plan_placement(plan, placement) + assert excinfo.value.verdict == P3Verdict.PLACEMENT_MAP_VERSION_MISMATCH + + +def test_core_and_envelope_must_align(): + placement = _map() + plan = _plan(placement) + plan["envelope"].pop() + with pytest.raises(P3Error) as excinfo: + validate_plan_placement(plan, placement) + assert excinfo.value.verdict == P3Verdict.INVALID_PLACEMENT_MAP + + +def test_version_gate_precedes_the_structural_gate(): + """Both defects present: the verdict must name the version. + + A map from another placement cannot be judged structurally relevant. + """ + placement = _map() + plan = _plan(placement) + broken = replace(placement, version="other-v9", strategy="mystery") + with pytest.raises(P3Error) as excinfo: + validate_plan_placement(plan, broken) + assert excinfo.value.verdict == P3Verdict.PLACEMENT_MAP_VERSION_MISMATCH + + +def test_structural_gate_precedes_row_comparison(): + placement = _map() + plan = _plan(placement) + plan["envelope"][0]["physical_expert_id"] = 99 + with pytest.raises(P3Error) as excinfo: + validate_plan_placement(plan, replace(placement, strategy="mystery")) + assert excinfo.value.verdict == P3Verdict.INVALID_PLACEMENT_MAP + assert "unknown strategy" in str(excinfo.value) + + +def test_plan_recorded_under_interleave_fails_against_a_block_map(): + interleave = _map(8, STRATEGY_INTERLEAVE, SEMANTICS_RANK_LOCAL) + plan = _plan(interleave) + with pytest.raises(P3Error) as excinfo: + validate_plan_placement(plan, _map(8, STRATEGY_BLOCK, SEMANTICS_RANK_LOCAL)) + assert excinfo.value.verdict == P3Verdict.INVALID_PLACEMENT_MAP + + +# --- interop with the start kit ---------------------------------------------- + + +def _assembled_plan(case, placement): + from rl_engine.p3 import oracle + from rl_engine.p3.assembler import assemble + from rl_engine.p3.contract import P3OpResult + + score = oracle.router_sqrt_softplus_fwd(case.z, case.round_policy) + s = score["s"] + route = ( + oracle.hash_route_fwd(case.input_token_id, s, case.table) + if case.router_mode == "hash" + else oracle.learned_route_fwd(s, case.bias) + ) + bundle = assemble( + case, + P3OpResult(P3Verdict.PASS, score), + P3OpResult(P3Verdict.PASS, route), + placement=placement.logical_to_physical(), + placement_version=placement.version, + ) + return bundle["RoutePlan"] + + +def _native_math_available(): + from rl_engine.p3 import bitmath + + try: + import numpy as np + + bitmath.apply(np.zeros(1, dtype=np.float32), "sqrt") + except OSError: + return False + return True + + +@pytest.mark.skipif( + not _native_math_available(), reason="native bitmath library unavailable on this platform" +) +@pytest.mark.parametrize("strategy", [STRATEGY_BLOCK, STRATEGY_INTERLEAVE]) +def test_global_placement_round_trips_through_the_real_assembler(strategy): + """A generated global placement must satisfy both gates and the kit's own. + + This is the compatibility claim: under the global reading nothing about + the existing assembler or checker has to change. + """ + from rl_engine.p3.checker import validate_plan + from rl_engine.p3.fixtures import catalog + + placement = PlacementMap.from_ep_config(f"ep8-{strategy}-v1", 8, strategy, SEMANTICS_GLOBAL) + case = next(c for c in catalog() if c.row_active.any()) + plan = _assembled_plan(case, placement) + + validate_plan_placement(plan, placement) + validate_plan({"RoutePlan": plan}) From caf5e86c0d8d1708633cc23491a443d238867b1b Mon Sep 17 00:00:00 2001 From: AlanX-Yan Date: Thu, 15 Oct 2026 01:58:53 +0800 Subject: [PATCH 2/4] [DSv4][P3-T08] Implement token ownership under combined parallelism mocks.OWNERSHIP records what each parallel dimension does to token ownership. Each mode is clear in isolation and none is checkable in isolation: a rank under tp=2 cp=2 dp=2 pp=2 ep=2 is at once a sequence-replica, a sequence-partition, a batch-partition, a layer-replica and token-irrelevant, and no row of the table says whether rank 19 should hold token 5 at layer 2. The unit of ownership is a routing decision (absolute_layer, global_token_id), not a token. A token-keyed model cannot express pp: replica_by_layer at all, since the same token at two layers lives on two stages. Everything goes through carriers(): the set of global ranks that must hold a decision. Each dimension either shrinks that set or leaves it alone, so composition is an intersection and three rules cover every configuration -- a rank holding what it does not carry is AMBIGUOUS_GLOBAL_TOKEN_MAPPING, a carrier missing what it owes is INCOMPLETE_ARTIFACT, carriers disagreeing is ROUTE_SEMANTIC_FINGERPRINT_MISMATCH. carriers is built by inverting expected_decisions over the world rather than recomposing per-dimension rules, so the two directions cannot drift apart. rank_order and the CP-before-SP sequence split are declared, not inferred: guessing either yields a validator that agrees with the truth on symmetric configurations and silently disagrees on the rest. sp is a flag rather than a size because sequence parallelism subdivides the existing TP group, and the resulting tp mode is read from mocks.OWNERSHIP rather than restated. ep is kept separate from tp although the carrier arithmetic is identical: a fingerprint difference tracking the EP coordinate means expert placement leaked into the routing decision, a different bug with a different owner, so the failure message names the implicated dimension. Gate order is rank set -> per-rank ownership -> replica agreement. check_cross_config is token-centric, never rank-against-rank, since changing the layout is precisely what makes rank N hold different tokens. Co-Authored-By: Claude Opus 5 --- docs/design/p3-token-ownership.md | 207 ++++++++++++++ rl_engine/p3/ownership.py | 438 ++++++++++++++++++++++++++++++ tests/p3/test_ownership.py | 383 ++++++++++++++++++++++++++ 3 files changed, 1028 insertions(+) create mode 100644 docs/design/p3-token-ownership.md create mode 100644 rl_engine/p3/ownership.py create mode 100644 tests/p3/test_ownership.py diff --git a/docs/design/p3-token-ownership.md b/docs/design/p3-token-ownership.md new file mode 100644 index 000000000..9a720b07e --- /dev/null +++ b/docs/design/p3-token-ownership.md @@ -0,0 +1,207 @@ +# P3 Token Ownership Under Combined Parallelism: Design Record + +| Field | Value | +|---|---| +| Document status | Delivered | +| Document date | 2026-10-15 | +| Work package | T08 (CP/DP/PP/EP ownership, placement, P7 `cross_config`) | +| Target | DSV4-Flash MoE router | +| Test baseline | 40 passed (ownership); 93 passed / 2 skipped with the placement map | +| Scope | Token ownership and cross-config comparison; no collectives, no second P7 runner | +| Depends on | #498 (`P3Verdict`, `require`, `mocks.OWNERSHIP`) | + +> **What this is.** `mocks.py` records *what* each parallel dimension does to +> token ownership. This is the module that turns that table into behaviour. +> Section 2 is the modelling decision everything else follows from. + +--- + +## 1. Why The Table Is Not Enough + +The contract assigns each dimension a mode: + +``` +tp: replica sp: partition_by_sequence cp: partition +dp: partition pp: replica_by_layer ep: not_a_token_dimension +``` + +Each is clear in isolation and none of them is checkable in isolation. A rank +under `tp=2, cp=2, dp=2, pp=2, ep=2` is simultaneously a sequence-replica, a +sequence-partition, a batch-partition, a layer-replica and token-irrelevant. +There are 32 ranks and the question "should rank 19 hold token 5 at layer 2?" +has one right answer that no single row of the table gives. + +Composition is the work. + +## 2. The Unit Is A Decision, And The Model Is Carriers + +### 2.1 `(absolute_layer, global_token_id)`, not a token + +A token-keyed model cannot express `pp: replica_by_layer` at all. The same +token routed at layer 1 and layer 2 is two routing decisions, and under +pipeline parallelism they live on different stages. Keying by token forces +either "this stage is missing most of its tokens" or "layers are not checked", +and both are wrong. + +This also matches the semantic hash, which is per routing decision. + +### 2.2 One function: `carriers` + +`carriers(config, layout)` maps every decision to the set of global ranks that +must hold it. Each dimension either shrinks that set (partitioning) or leaves +it alone (replication), so composition is an intersection and the six modes +need no special cases downstream: + +| Finding | Verdict | +|---|---| +| a rank holds a decision it does not carry | `AMBIGUOUS_GLOBAL_TOKEN_MAPPING(63)` | +| a carrier is missing a decision it owes | `INCOMPLETE_ARTIFACT(13)` | +| two carriers of one decision disagree | `ROUTE_SEMANTIC_FINGERPRINT_MISMATCH(59)` | + +Three rules cover every configuration. Without the carrier set each +combination would need its own predicate, and the predicates would disagree +with each other at the edges. + +`carriers` is built by **inverting** `expected_decisions` over the world +rather than by composing per-dimension rules a second time. One derivation +means the two directions cannot drift apart — a class of bug that is invisible +until a configuration exercises exactly the case where they differ. + +## 3. Declared, Not Inferred + +Two things could be guessed from the data. Both are declared instead, because +a validator that guesses them agrees with the truth on symmetric +configurations and silently disagrees on the rest — stable in testing, wrong +in production. + +**`rank_order`** — how a global rank decomposes into coordinates. Default +`("tp", "cp", "ep", "dp", "pp")`, fastest-varying first. Under `tp=2, pp=2`, +rank 1 is `tp=1` in this order and `pp=1` under a different one; nothing in +the artifact distinguishes them. + +**Sequence split order** — CP cuts the sequence first, then SP subdivides each +CP chunk inside the TP group. Reversing it produces a different, equally +self-consistent assignment. `test_cp_then_sp_split_the_sequence_in_that_order` +pins the chosen one. + +`sp` is a flag rather than a size, because sequence parallelism subdivides +along the existing TP group and has no axis of its own. Setting it changes TP +from a replica dimension into a partitioning one, which is exactly what the +contract means by "`tp` is `replica` or `partition_by_sequence`". The mode is +read through `mocks.OWNERSHIP` rather than restated, so that table stays the +single source. + +## 4. Design Decisions + +### D1: `ep` is listed apart from `tp` although the arithmetic is identical + +Both replicate, so carrier computation treats them the same. They are kept +separate because the *defect* they expose is different: a fingerprint +difference that tracks the EP coordinate means expert placement leaked into +the routing decision, which is a different bug with a different owner from a +TP replica drifting. `replica_dimensions` names the responsible dimensions in +the failure message, so the report points at EP when EP is implicated. + +### D2: Gate order is rank set -> per-rank ownership -> replica agreement + +A dependency order, not a preference. Ownership computed over a contaminated +rank set is meaningless. Comparing carriers before knowing each rank holds the +right decisions would report a partition error as a numeric disagreement, and +send it to the wrong owner. + +Within the rank-set gate, duplicate/unexpected ranks and mixed groups outrank +a simultaneous missing rank: contaminated evidence makes the missing-rank +finding an artefact of the contamination rather than a fact about the run. + +Within per-rank ownership, unauthorized decisions outrank gaps: holding +another rank's decision means the partition itself is wrong, and the gaps +elsewhere are its consequence. + +### D3: Ownership is judged per rank, over all its layers at once + +A rank reports one evidence record per layer it owns. Judging a single record +against the rank's full expected set reads every other layer of that rank as a +gap — this was a real defect during development, caught by the `pp=1` control +where one rank legitimately owns four layers. `check_rank_ownership` takes the +rank and all of its records. + +### D4: Replica agreement compares every carrier, never a deduplicated copy + +Folding carriers to one representative before comparing would hide exactly the +failure replication exists to rule out: a single rank drifting while its peers +agree. Every carrier's fingerprint is collected and any split is reported with +the rank groups on each side. + +### D5: Cross-config comparison is token-centric + +Never rank-against-rank. The whole point of changing the parallel layout is +that rank N holds different tokens afterwards; `rank 0 vs rank 0` compares +unrelated work. Each side is folded to `(layer, token) -> fingerprint` first, +and only then compared. Both sides are validated for ownership before any +comparison, so a broken side is reported as broken rather than as drift. + +## 5. Test Matrix + +Every negative case injects exactly one defect. + +| Injected defect | Verdict | +|---|---| +| a rank absent from the evidence | 20 `MISSING_RANK` | +| duplicate `(rank, layer)` record | 22 `STALE_RUN_METADATA` | +| rank outside the world | 22 | +| mixed group labels | 22 | +| rank holds a decision it does not carry | 63 | +| rank missing a decision it owes | 13 | +| one TP replica drifting | 59 (message names `tp`) | +| one EP replica drifting | 59 (message names `ep`) | +| cross-config: side missing decisions | 13 | +| cross-config: ghost decision | 63 | +| cross-config: decision fingerprint differs | 59 | +| layout repeating a `global_token_id` | 63 | +| empty sequence or layer space | 11 `SCHEMA_MISMATCH` | +| non-positive dimension size, bad `rank_order` | 11 | +| rank outside `[0, world_size)` | 20 | +| stale evidence **and** a missing rank | **22** (integrity precedes) | +| unauthorized decision **and** a gap | **63** (partition precedes) | +| broken rank set **and** a ghost decision | **20** (rank set precedes) | + +Controls (must **not** be flagged): + +| Case | Why it is legal | +|---|---| +| pipeline stages holding disjoint layers | `pp: replica_by_layer` | +| two TP ranks holding the same token | `tp: replica` | +| partitioned configuration with nothing to agree on | no carrier has a peer | +| two configurations where rank 0 holds different tokens | D5 | + +Positive coverage pins each dimension's effect on the carrier set +(`per_rank` size and whether anything is replicated), the combined +`tp2 cp2 dp2 pp2 ep2` case (world 32, four carriers per decision = `tp x ep`), +and `coordinates`/`global_rank` round-tripping across the whole world. + +## 6. Boundary and Pending Work + +- No collectives are implemented or wrapped here, and no second P7 runner is + introduced. This module answers "who should hold what", nothing else. +- `TokenLayout` describes the decision space explicitly (sequences plus + layers). Deriving it from a live job's batch and model configuration is not + done; the validator takes it as input. +- The P7 `cross_config` materializer is still unwired. `check_cross_config` + provides the comparison; connecting it to P7's existing planner/runner and + `mocks.materialize_tokens` is the remaining T08 item. +- Capacity overflow does not interact with ownership here: an overflowed slot + is still a decision the owning rank must report. + +## 7. Reproduction + +```bash +cd RL-Kernel # repo root +python -m pytest tests/p3/test_ownership.py -q +# -> 40 passed + +python -m pytest tests/p3/test_ownership.py tests/p3/test_placement.py -q +# -> 93 passed, 2 skipped (skips are the native-bitmath interop cases) + +python -m ruff check rl_engine/p3/ownership.py tests/p3/test_ownership.py +python -m ruff format --check rl_engine/p3/ownership.py tests/p3/test_ownership.py +``` diff --git a/rl_engine/p3/ownership.py b/rl_engine/p3/ownership.py new file mode 100644 index 000000000..4a94ae2b8 --- /dev/null +++ b/rl_engine/p3/ownership.py @@ -0,0 +1,438 @@ +"""Token ownership under combined parallelism (T08). + +`mocks.py` records *what* each parallel dimension does to token ownership: + + tp: replica sp: partition_by_sequence cp: partition + dp: partition pp: replica_by_layer ep: not_a_token_dimension + +That table is the contract. This module turns it into behaviour, because the +six modes only become checkable once they are composed: a rank under +``tp=2, cp=2, dp=2, pp=2, ep=2`` is simultaneously a sequence-replica, a +sequence-partition, a batch-partition, a layer-replica and token-irrelevant, +and no single rule covers it. + +The unit of ownership is a **routing decision** ``(absolute_layer, +global_token_id)``, not a token. The same token routed at two layers is two +decisions and they live on different pipeline stages, so a token-keyed model +cannot express ``pp: replica_by_layer`` at all. + +Everything is expressed through one function: ``carriers`` — the set of global +ranks that must hold a given decision. Each dimension either shrinks that set +(partitioning) or leaves it alone (replication), so composition is an +intersection and the six modes need no special cases downstream: + +- a rank holding a decision it does not carry is ``AMBIGUOUS_GLOBAL_TOKEN_MAPPING`` +- a carrier missing a decision it owes is ``INCOMPLETE_ARTIFACT`` +- two carriers of one decision disagreeing is ``ROUTE_SEMANTIC_FINGERPRINT_MISMATCH`` + +Two things are declared rather than inferred, because guessing either one +produces a validator that is stable in testing and wrong in production: +``rank_order`` (how a global rank decomposes into coordinates) and the +sequence split order (CP before SP within a TP group). +""" + +from dataclasses import dataclass, field + +from .contract import P3Verdict +from .mocks import OWNERSHIP +from .oracle import require + +OWNERSHIP_SCHEMA = "p3-ownership.v1" + +#: Dimensions that subdivide the decision space, and what they cut. +PARTITION_BY_BATCH = ("dp",) +PARTITION_BY_SEQUENCE = ("cp", "sp") +PARTITION_BY_LAYER = ("pp",) +#: Dimensions that replicate it. ``ep`` is listed separately from ``tp`` +#: because an EP-induced token difference is a different defect, worth its own +#: diagnostic, even though the carrier arithmetic is the same. +REPLICATE = ("tp", "ep") + +#: Default decomposition of a global rank, fastest-varying dimension first. +#: Declared, never inferred: a validator that guesses this agrees with the +#: truth on symmetric configurations and silently disagrees on the rest. +DEFAULT_RANK_ORDER = ("tp", "cp", "ep", "dp", "pp") + + +def partition_slices(count, parts): + """Contiguous ranges over ``count`` items split across ``parts`` holders. + + The first ``count % parts`` holders take one extra item. This is the same + rule the placement map uses for ``block``; keeping one convention means a + reader never has to ask which split a given remainder followed. + """ + require(parts > 0, P3Verdict.SCHEMA_MISMATCH, f"parts must be positive, got {parts}") + base, extra = divmod(count, parts) + slices, start = [], 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + slices.append(range(start, start + width)) + start += width + return slices + + +@dataclass(frozen=True) +class ParallelConfig: + """One parallel configuration and how its global ranks decompose. + + ``sp`` is a flag rather than a size: sequence parallelism subdivides along + the existing TP group, so it has no axis of its own. With ``sp`` set, TP + stops being a pure replica dimension and starts partitioning the sequence + (contract: ``tp`` is ``replica`` or ``partition_by_sequence``). + """ + + tp: int = 1 + sp: bool = False + cp: int = 1 + dp: int = 1 + pp: int = 1 + ep: int = 1 + rank_order: tuple = DEFAULT_RANK_ORDER + schema_version: str = OWNERSHIP_SCHEMA + + def __post_init__(self): + for name in ("tp", "cp", "dp", "pp", "ep"): + size = getattr(self, name) + require( + isinstance(size, int) and size > 0, + P3Verdict.SCHEMA_MISMATCH, + f"{name} must be a positive int, got {size!r}", + ) + require( + tuple(sorted(self.rank_order)) == tuple(sorted(DEFAULT_RANK_ORDER)), + P3Verdict.SCHEMA_MISMATCH, + f"rank_order must permute {DEFAULT_RANK_ORDER}, got {self.rank_order}", + ) + + @property + def world_size(self): + return self.tp * self.cp * self.ep * self.dp * self.pp + + def size_of(self, dimension): + """Axis length; ``sp`` has none of its own and reports the TP size.""" + return self.tp if dimension == "sp" else getattr(self, dimension) + + def ownership_of(self, dimension): + """The contract's ownership mode for one dimension. + + ``tp`` is the only dimension whose mode depends on configuration, and + it is read through the shared ``OWNERSHIP`` table rather than restated + here so that the table stays the single source. + """ + if dimension == "tp": + return OWNERSHIP["sp"] if self.sp else OWNERSHIP["tp"] + return OWNERSHIP[dimension] + + def coordinates(self, global_rank): + """Decompose a global rank into per-dimension coordinates.""" + require( + 0 <= global_rank < self.world_size, + P3Verdict.MISSING_RANK, + f"rank {global_rank} outside [0, {self.world_size})", + ) + coords, remainder = {}, global_rank + for dimension in self.rank_order: + size = getattr(self, dimension) + coords[dimension] = remainder % size + remainder //= size + return coords + + def global_rank(self, **coords): + """Compose per-dimension coordinates back into a global rank.""" + rank, stride = 0, 1 + for dimension in self.rank_order: + size = getattr(self, dimension) + value = int(coords.get(dimension, 0)) + require( + 0 <= value < size, + P3Verdict.MISSING_RANK, + f"{dimension} coordinate {value} outside [0, {size})", + ) + rank += value * stride + stride *= size + return rank + + def describe(self): + """Human-readable configuration label used in reports and groups.""" + parts = [f"tp{self.tp}" + ("sp" if self.sp else "")] + parts += [f"{d}{getattr(self, d)}" for d in ("cp", "dp", "pp", "ep")] + return "-".join(parts) + + +@dataclass(frozen=True) +class TokenLayout: + """The global decision space a configuration is sliced out of. + + ``sequences`` holds one tuple of ``global_token_id`` per batch item, in + order; ``layers`` holds the model's absolute layer ids. Together they + enumerate every ``(layer, token)`` decision the whole job must produce. + """ + + sequences: tuple + layers: tuple + + def __post_init__(self): + require(bool(self.sequences), P3Verdict.SCHEMA_MISMATCH, "no sequences in layout") + require(bool(self.layers), P3Verdict.SCHEMA_MISMATCH, "no layers in layout") + seen = [token for sequence in self.sequences for token in sequence] + require( + len(set(seen)) == len(seen), + P3Verdict.AMBIGUOUS_GLOBAL_TOKEN_MAPPING, + "global_token_id repeated across the layout; it must be unique job-wide", + ) + + @classmethod + def flat(cls, tokens, layers): + """A single-sequence layout; convenient for single-layer fixtures.""" + return cls(sequences=(tuple(tokens),), layers=tuple(layers)) + + def decisions(self): + """Every ``(layer, token)`` decision in the job, in canonical order.""" + return tuple( + (layer, token) + for layer in self.layers + for sequence in self.sequences + for token in sequence + ) + + +# --- carriers ---------------------------------------------------------------- + + +def _sequence_owner(config, sequence, cp_rank, tp_rank): + """Tokens of one sequence held by a (cp, tp) pair. + + CP splits the sequence first, then SP subdivides each CP chunk inside the + TP group. That order is the declared composition; reversing it produces a + different, equally self-consistent assignment, which is exactly why it has + to be written down rather than discovered. + """ + chunk = partition_slices(len(sequence), config.cp)[cp_rank] + if not config.sp: + return tuple(sequence[i] for i in chunk) + sub = partition_slices(len(chunk), config.tp)[tp_rank] + return tuple(sequence[chunk.start + i] for i in sub) + + +def expected_decisions(config, layout, global_rank): + """Decisions one global rank must hold, and only those.""" + coords = config.coordinates(global_rank) + layer_slice = partition_slices(len(layout.layers), config.pp)[coords["pp"]] + layers = [layout.layers[i] for i in layer_slice] + sequence_slice = partition_slices(len(layout.sequences), config.dp)[coords["dp"]] + + tokens = [] + for index in sequence_slice: + tokens.extend(_sequence_owner(config, layout.sequences[index], coords["cp"], coords["tp"])) + return frozenset((layer, token) for layer in layers for token in tokens) + + +def carriers(config, layout): + """Map every decision to the set of global ranks that must hold it. + + Built by inverting ``expected_decisions`` over the world rather than by + composing per-dimension rules a second time: one derivation, so the two + directions cannot drift apart. + """ + out = {} + for rank in range(config.world_size): + for decision in expected_decisions(config, layout, rank): + out.setdefault(decision, set()).add(rank) + return {decision: frozenset(ranks) for decision, ranks in out.items()} + + +def replica_dimensions(config): + """Dimensions that replicate decisions under this configuration.""" + dims = [d for d in REPLICATE if config.size_of(d) > 1] + if config.sp: + dims = [d for d in dims if d != "tp"] + return tuple(dims) + + +# --- evidence ---------------------------------------------------------------- + + +@dataclass(frozen=True) +class RankEvidence: + """What one rank reported for one layer.""" + + rank: int + group: str + layer: int + fingerprints: dict = field(default_factory=dict) # global_token_id -> semantic hash + + @classmethod + def from_plan(cls, rank, group, layer, plan): + """Read a RoutePlan's per-token semantic fingerprints.""" + return cls( + rank=rank, + group=group, + layer=layer, + fingerprints={int(k): v for k, v in plan["per_token_fingerprints"].items()}, + ) + + def decisions(self): + return frozenset((self.layer, token) for token in self.fingerprints) + + +# --- gates ------------------------------------------------------------------- + + +def check_rank_set(config, evidence): + """The reported rank set must be exactly one configuration's world. + + Duplicate ranks, unexpected ranks and mixed group labels are stale or + mixed evidence and take precedence over a simultaneous missing-rank + symptom: a contaminated rank set makes the missing-rank finding an + artefact of the contamination rather than a fact about the run. + """ + seen, groups = {}, {item.group for item in evidence} + for item in evidence: + seen.setdefault((item.rank, item.layer), []).append(item) + + duplicated = sorted(key for key, items in seen.items() if len(items) > 1) + ranks = {rank for rank, _ in seen} + unexpected = sorted(r for r in ranks if not 0 <= r < config.world_size) + if duplicated or unexpected or len(groups) > 1: + problems = [] + if duplicated: + problems.append(f"duplicated (rank, layer) evidence: {duplicated[:8]}") + if unexpected: + problems.append(f"ranks outside [0, {config.world_size}): {unexpected[:8]}") + if len(groups) > 1: + problems.append(f"mixed groups: {sorted(groups)}") + require(False, P3Verdict.STALE_RUN_METADATA, "; ".join(problems)) + + missing = sorted(set(range(config.world_size)) - ranks) + require( + not missing, + P3Verdict.MISSING_RANK, + f"missing ranks for {config.describe()}: {missing[:16]}", + ) + + +def check_rank_ownership(config, layout, rank, items): + """One rank must hold exactly the decisions it carries. + + ``items`` is every evidence record for this rank. A rank reports one + record per layer it owns, so ownership can only be judged over the union: + checking a single record would read every other layer of the same rank as + a gap. + + Unauthorized decisions are reported before missing ones: holding a + decision that belongs to another rank means the partition itself is wrong, + and the gaps elsewhere are its consequence. + """ + expected = expected_decisions(config, layout, rank) + actual = frozenset().union(*(item.decisions() for item in items)) if items else frozenset() + + unauthorized = sorted(actual - expected) + if unauthorized: + owners = carriers(config, layout) + detail = ", ".join( + f"(layer={layer}, token={token})->ranks{sorted(owners.get((layer, token), ()))[:4]}" + for layer, token in unauthorized[:6] + ) + require( + False, + P3Verdict.AMBIGUOUS_GLOBAL_TOKEN_MAPPING, + f"rank {rank} reports decisions it does not carry: {detail}", + ) + + missing = sorted(expected - actual) + require( + not missing, + P3Verdict.INCOMPLETE_ARTIFACT, + f"rank {rank} owes {len(missing)} decisions, first missing {missing[:6]}", + ) + + +def check_replica_agreement(config, layout, evidence): + """Every carrier of a decision must report the same semantic fingerprint. + + Comparing a deduplicated single copy would miss a single-rank drift, which + is the whole failure mode replication is supposed to rule out. Under + ``ep`` the message names the dimension explicitly: a token-set or + fingerprint difference that tracks the EP coordinate means expert + placement leaked into the routing decision, which is a different bug from + a TP replica drifting. + """ + by_rank = {(item.rank, item.layer): item for item in evidence} + for decision, ranks in carriers(config, layout).items(): + if len(ranks) < 2: + continue + layer, token = decision + reported = {} + for rank in sorted(ranks): + item = by_rank.get((rank, layer)) + if item is not None and token in item.fingerprints: + reported.setdefault(item.fingerprints[token], []).append(rank) + if len(reported) > 1: + dims = replica_dimensions(config) or ("replica",) + spread = "; ".join(f"{h[:12]}<-ranks{r}" for h, r in sorted(reported.items())) + require( + False, + P3Verdict.ROUTE_SEMANTIC_FINGERPRINT_MISMATCH, + f"carriers of (layer={layer}, token={token}) disagree across " + f"{'/'.join(dims)}: {spread}", + ) + + +def validate_ownership(config, layout, evidence): + """Full ownership gate for one configuration's evidence. + + Order is rank set -> per-rank ownership -> replica agreement, a dependency + order: ownership of a contaminated rank set is meaningless, and comparing + carriers before knowing each rank holds the right decisions would attribute + a partition error to a numeric disagreement. + """ + check_rank_set(config, evidence) + by_rank = {} + for item in evidence: + by_rank.setdefault(item.rank, []).append(item) + for rank in sorted(by_rank): + check_rank_ownership(config, layout, rank, by_rank[rank]) + check_replica_agreement(config, layout, evidence) + + +def check_cross_config(layout, base, other): + """Two configurations must produce identical decisions for every token. + + Token-centric, never rank-against-rank: the whole point of changing the + parallel layout is that rank N holds different tokens afterwards. Each + side is folded to ``(layer, token) -> fingerprint`` first, and only then + compared. + """ + base_config, base_evidence = base + other_config, other_evidence = other + validate_ownership(base_config, layout, base_evidence) + validate_ownership(other_config, layout, other_evidence) + + folded = [] + for evidence in (base_evidence, other_evidence): + merged = {} + for item in evidence: + for token, value in item.fingerprints.items(): + merged[(item.layer, token)] = value + folded.append(merged) + base_map, other_map = folded + + missing = sorted(set(base_map) - set(other_map)) + require( + not missing, + P3Verdict.INCOMPLETE_ARTIFACT, + f"{other_config.describe()} is missing {len(missing)} decisions, first {missing[:6]}", + ) + ghost = sorted(set(other_map) - set(base_map)) + require( + not ghost, + P3Verdict.AMBIGUOUS_GLOBAL_TOKEN_MAPPING, + f"{other_config.describe()} reports decisions no base rank owns: {ghost[:6]}", + ) + for decision in sorted(base_map): + require( + base_map[decision] == other_map[decision], + P3Verdict.ROUTE_SEMANTIC_FINGERPRINT_MISMATCH, + f"(layer={decision[0]}, token={decision[1]}) differs between " + f"{base_config.describe()} and {other_config.describe()}", + ) diff --git a/tests/p3/test_ownership.py b/tests/p3/test_ownership.py new file mode 100644 index 000000000..90c0c5fa6 --- /dev/null +++ b/tests/p3/test_ownership.py @@ -0,0 +1,383 @@ +"""Token ownership tests (T08): six parallel modes, composed and gated. + +The positive layer pins what each dimension does to the carrier set, alone and +combined. Every negative case injects exactly one defect, so a correct verdict +cannot be a coincidence of two interacting faults. Controls pin the inverse: +pipeline stages holding different layers, and TP replicas holding the same +token twice, must not be flagged. +""" + +from dataclasses import replace + +import pytest + +from rl_engine.p3.contract import P3Error, P3Verdict +from rl_engine.p3.mocks import OWNERSHIP +from rl_engine.p3.ownership import ( + DEFAULT_RANK_ORDER, + ParallelConfig, + RankEvidence, + TokenLayout, + carriers, + check_cross_config, + check_rank_ownership, + check_rank_set, + check_replica_agreement, + expected_decisions, + partition_slices, + replica_dimensions, + validate_ownership, +) + +LAYOUT = TokenLayout(sequences=((0, 1, 2, 3), (4, 5, 6, 7)), layers=(0, 1, 2, 3)) +TOTAL = len(LAYOUT.decisions()) # 4 layers x 8 tokens + + +def _fp(layer, token): + """A deterministic stand-in for the per-decision semantic fingerprint.""" + return f"{layer:04d}-{token:04d}" + + +def _evidence(config, layout=LAYOUT, fp=_fp, group=None): + """Perfect evidence: every rank reports exactly what it carries.""" + group = group or config.describe() + out = [] + for rank in range(config.world_size): + decisions = expected_decisions(config, layout, rank) + for layer in sorted({layer for layer, _ in decisions}): + tokens = {token: fp(layer, token) for lay, token in decisions if lay == layer} + out.append(RankEvidence(rank, group, layer, tokens)) + return out + + +def _for_rank(evidence, rank): + """Every evidence record one rank produced, across all layers it owns.""" + return [item for item in evidence if item.rank == rank] + + +def _verdict(excinfo): + return excinfo.value.verdict + + +# --- configuration arithmetic ------------------------------------------------ + + +def test_partition_slices_gives_the_remainder_to_the_first_holders(): + assert [len(s) for s in partition_slices(10, 4)] == [3, 3, 2, 2] + assert [len(s) for s in partition_slices(8, 4)] == [2, 2, 2, 2] + assert [len(s) for s in partition_slices(2, 4)] == [1, 1, 0, 0] + + +def test_rank_coordinates_round_trip(): + config = ParallelConfig(tp=2, cp=2, dp=2, pp=2, ep=2) + assert config.world_size == 32 + for rank in range(config.world_size): + assert config.global_rank(**config.coordinates(rank)) == rank + + +def test_rank_order_is_declared_and_changes_the_decomposition(): + layout = ParallelConfig(tp=2, pp=2) + reordered = ParallelConfig(tp=2, pp=2, rank_order=("pp", "cp", "ep", "dp", "tp")) + assert layout.coordinates(1) == {"tp": 1, "cp": 0, "ep": 0, "dp": 0, "pp": 0} + assert reordered.coordinates(1) == {"pp": 1, "cp": 0, "ep": 0, "dp": 0, "tp": 0} + + +def test_invalid_configuration_fails_closed(): + for kwargs in ({"tp": 0}, {"dp": -1}, {"rank_order": ("tp", "cp")}): + with pytest.raises(P3Error) as excinfo: + ParallelConfig(**kwargs) + assert _verdict(excinfo) == P3Verdict.SCHEMA_MISMATCH + + +def test_rank_outside_the_world_fails_closed(): + with pytest.raises(P3Error) as excinfo: + ParallelConfig(tp=2).coordinates(2) + assert _verdict(excinfo) == P3Verdict.MISSING_RANK + + +def test_ownership_modes_come_from_the_shared_table(): + assert ParallelConfig(tp=2).ownership_of("tp") == OWNERSHIP["tp"] == "replica" + assert ParallelConfig(tp=2, sp=True).ownership_of("tp") == OWNERSHIP["sp"] + assert ParallelConfig().ownership_of("pp") == "replica_by_layer" + assert ParallelConfig().ownership_of("ep") == "not_a_token_dimension" + + +# --- what each dimension does ------------------------------------------------ + + +@pytest.mark.parametrize( + ("config", "per_rank", "replicated"), + [ + (ParallelConfig(), TOTAL, False), + (ParallelConfig(dp=2), TOTAL // 2, False), + (ParallelConfig(cp=2), TOTAL // 2, False), + (ParallelConfig(pp=2), TOTAL // 2, False), + (ParallelConfig(tp=2, sp=True), TOTAL // 2, False), + (ParallelConfig(tp=2), TOTAL, True), + (ParallelConfig(ep=2), TOTAL, True), + ], +) +def test_single_dimension_ownership(config, per_rank, replicated): + owners = carriers(config, LAYOUT) + assert len(owners) == TOTAL # the job's decision space is always covered + for rank in range(config.world_size): + assert len(expected_decisions(config, LAYOUT, rank)) == per_rank + multi = [d for d, ranks in owners.items() if len(ranks) > 1] + assert bool(multi) is replicated + + +def test_pipeline_stages_own_disjoint_layers(): + config = ParallelConfig(pp=2) + stage0 = {layer for layer, _ in expected_decisions(config, LAYOUT, 0)} + stage1 = {layer for layer, _ in expected_decisions(config, LAYOUT, 1)} + assert stage0 == {0, 1} + assert stage1 == {2, 3} + assert not stage0 & stage1 + + +def test_sequence_parallel_subdivides_inside_the_tp_group(): + plain = ParallelConfig(tp=2) + sharded = ParallelConfig(tp=2, sp=True) + assert expected_decisions(plain, LAYOUT, 0) == expected_decisions(plain, LAYOUT, 1) + assert not expected_decisions(sharded, LAYOUT, 0) & expected_decisions(sharded, LAYOUT, 1) + + +def test_cp_then_sp_split_the_sequence_in_that_order(): + layout = TokenLayout.flat(range(8), layers=(0,)) + config = ParallelConfig(tp=2, sp=True, cp=2) + held = {r: sorted(t for _, t in expected_decisions(config, layout, r)) for r in range(4)} + # CP halves the sequence first; SP halves each CP chunk inside the TP group. + assert sorted(held.values()) == [[0, 1], [2, 3], [4, 5], [6, 7]] + + +def test_combined_configuration_replicates_across_tp_and_ep_only(): + config = ParallelConfig(tp=2, cp=2, dp=2, pp=2, ep=2) + owners = carriers(config, LAYOUT) + assert len(owners) == TOTAL + assert {len(ranks) for ranks in owners.values()} == {4} # tp x ep + assert replica_dimensions(config) == ("tp", "ep") + assert replica_dimensions(ParallelConfig(tp=2, sp=True, ep=2)) == ("ep",) + + +def test_perfect_evidence_passes_for_every_configuration(): + for config in ( + ParallelConfig(), + ParallelConfig(dp=2), + ParallelConfig(cp=2), + ParallelConfig(pp=2), + ParallelConfig(tp=2), + ParallelConfig(tp=2, sp=True), + ParallelConfig(ep=2), + ParallelConfig(tp=2, cp=2, dp=2, pp=2, ep=2), + ): + validate_ownership(config, LAYOUT, _evidence(config)) + + +# --- layout ------------------------------------------------------------------ + + +def test_layout_rejects_a_repeated_global_token_id(): + with pytest.raises(P3Error) as excinfo: + TokenLayout(sequences=((0, 1), (1, 2)), layers=(0,)) + assert _verdict(excinfo) == P3Verdict.AMBIGUOUS_GLOBAL_TOKEN_MAPPING + + +def test_layout_rejects_an_empty_space(): + for kwargs in ({"sequences": ()}, {"layers": ()}): + with pytest.raises(P3Error) as excinfo: + TokenLayout(**{"sequences": ((0,),), "layers": (0,), **kwargs}) + assert _verdict(excinfo) == P3Verdict.SCHEMA_MISMATCH + + +# --- rank-set gate ----------------------------------------------------------- + + +def test_missing_rank(): + config = ParallelConfig(dp=2) + evidence = [e for e in _evidence(config) if e.rank != 1] + with pytest.raises(P3Error) as excinfo: + check_rank_set(config, evidence) + assert _verdict(excinfo) == P3Verdict.MISSING_RANK + + +def test_duplicate_rank_layer_evidence(): + config = ParallelConfig(dp=2) + evidence = _evidence(config) + with pytest.raises(P3Error) as excinfo: + check_rank_set(config, evidence + [evidence[0]]) + assert _verdict(excinfo) == P3Verdict.STALE_RUN_METADATA + + +def test_unexpected_rank(): + config = ParallelConfig(dp=2) + evidence = _evidence(config) + with pytest.raises(P3Error) as excinfo: + check_rank_set(config, evidence + [replace(evidence[0], rank=9)]) + assert _verdict(excinfo) == P3Verdict.STALE_RUN_METADATA + + +def test_mixed_groups(): + config = ParallelConfig(dp=2) + evidence = _evidence(config) + evidence[0] = replace(evidence[0], group="another-run") + with pytest.raises(P3Error) as excinfo: + check_rank_set(config, evidence) + assert _verdict(excinfo) == P3Verdict.STALE_RUN_METADATA + + +def test_stale_evidence_outranks_a_simultaneous_missing_rank(): + """A contaminated rank set makes the missing-rank finding an artefact.""" + config = ParallelConfig(dp=4) + evidence = [e for e in _evidence(config) if e.rank != 3] + evidence[0] = replace(evidence[0], group="another-run") + with pytest.raises(P3Error) as excinfo: + check_rank_set(config, evidence) + assert _verdict(excinfo) == P3Verdict.STALE_RUN_METADATA + + +# --- per-rank ownership gate ------------------------------------------------- + + +def test_rank_holding_a_decision_it_does_not_carry(): + config = ParallelConfig(dp=2) + evidence = _evidence(config) + item = evidence[0] + stolen = next(t for t in range(8) if t not in item.fingerprints) + item.fingerprints[stolen] = _fp(item.layer, stolen) + with pytest.raises(P3Error) as excinfo: + check_rank_ownership(config, LAYOUT, 0, _for_rank(evidence, 0)) + assert _verdict(excinfo) == P3Verdict.AMBIGUOUS_GLOBAL_TOKEN_MAPPING + + +def test_rank_missing_a_decision_it_owes(): + config = ParallelConfig(dp=2) + evidence = _evidence(config) + evidence[0].fingerprints.pop(sorted(evidence[0].fingerprints)[0]) + with pytest.raises(P3Error) as excinfo: + check_rank_ownership(config, LAYOUT, 0, _for_rank(evidence, 0)) + assert _verdict(excinfo) == P3Verdict.INCOMPLETE_ARTIFACT + + +def test_unauthorized_outranks_a_simultaneous_gap(): + """Holding another rank's decision means the partition itself is wrong.""" + config = ParallelConfig(dp=2) + evidence = _evidence(config) + item = evidence[0] + owned = {token for _, token in expected_decisions(config, LAYOUT, 0)} + item.fingerprints.pop(sorted(item.fingerprints)[0]) + stolen = next(t for t in range(8) if t not in owned) # belongs to rank 1 + item.fingerprints[stolen] = _fp(item.layer, stolen) + with pytest.raises(P3Error) as excinfo: + check_rank_ownership(config, LAYOUT, 0, _for_rank(evidence, 0)) + assert _verdict(excinfo) == P3Verdict.AMBIGUOUS_GLOBAL_TOKEN_MAPPING + + +def test_pipeline_stage_is_not_flagged_for_missing_other_layers(): + """Control: a stage owes only its own layers.""" + config = ParallelConfig(pp=2) + evidence = _evidence(config) + for rank in range(config.world_size): + check_rank_ownership(config, LAYOUT, rank, _for_rank(evidence, rank)) + + +def test_tp_replicas_holding_the_same_token_are_not_duplicates(): + """Control: replication is not ambiguous ownership.""" + config = ParallelConfig(tp=2) + evidence = _evidence(config) + rank0 = next(e for e in evidence if e.rank == 0 and e.layer == 0) + rank1 = next(e for e in evidence if e.rank == 1 and e.layer == 0) + assert rank0.decisions() == rank1.decisions() + validate_ownership(config, LAYOUT, evidence) + + +# --- replica agreement ------------------------------------------------------- + + +def test_single_tp_replica_drifting_is_caught(): + config = ParallelConfig(tp=2) + evidence = _evidence(config) + token = sorted(evidence[1].fingerprints)[0] + evidence[1].fingerprints[token] = "drifted" + with pytest.raises(P3Error) as excinfo: + check_replica_agreement(config, LAYOUT, evidence) + assert _verdict(excinfo) == P3Verdict.ROUTE_SEMANTIC_FINGERPRINT_MISMATCH + assert "tp" in str(excinfo.value) + + +def test_ep_induced_drift_names_the_ep_dimension(): + """An EP-tracking difference means placement leaked into the decision.""" + config = ParallelConfig(ep=2) + evidence = _evidence(config) + token = sorted(evidence[1].fingerprints)[0] + evidence[1].fingerprints[token] = "drifted" + with pytest.raises(P3Error) as excinfo: + check_replica_agreement(config, LAYOUT, evidence) + assert _verdict(excinfo) == P3Verdict.ROUTE_SEMANTIC_FINGERPRINT_MISMATCH + assert "ep" in str(excinfo.value) + + +def test_partitioned_configurations_have_nothing_to_agree_on(): + config = ParallelConfig(dp=2, cp=2) + check_replica_agreement(config, LAYOUT, _evidence(config)) + + +def test_rank_set_gate_precedes_ownership(): + config = ParallelConfig(dp=2) + evidence = _evidence(config) + evidence[0].fingerprints[99] = "ghost" # would be AMBIGUOUS on its own + with pytest.raises(P3Error) as excinfo: + validate_ownership(config, LAYOUT, [e for e in evidence if e.rank != 1]) + assert _verdict(excinfo) == P3Verdict.MISSING_RANK + + +# --- cross-config ------------------------------------------------------------ + + +BASE = ParallelConfig(dp=2) +OTHER = ParallelConfig(tp=2, cp=2, pp=2) + + +def test_cross_config_agreement_passes(): + check_cross_config(LAYOUT, (BASE, _evidence(BASE)), (OTHER, _evidence(OTHER))) + + +def test_cross_config_is_token_centric_not_rank_against_rank(): + """Rank 0 holds different tokens under the two configurations.""" + base_rank0 = {t for _, t in expected_decisions(BASE, LAYOUT, 0)} + other_rank0 = {t for _, t in expected_decisions(OTHER, LAYOUT, 0)} + assert base_rank0 != other_rank0 + check_cross_config(LAYOUT, (BASE, _evidence(BASE)), (OTHER, _evidence(OTHER))) + + +def test_cross_config_decision_drift(): + other = _evidence(OTHER) + token = sorted(other[0].fingerprints)[0] + for item in other: # change it on every carrier: a real drift, not a replica split + if token in item.fingerprints and item.layer == other[0].layer: + item.fingerprints[token] = "drifted" + with pytest.raises(P3Error) as excinfo: + check_cross_config(LAYOUT, (BASE, _evidence(BASE)), (OTHER, other)) + assert _verdict(excinfo) == P3Verdict.ROUTE_SEMANTIC_FINGERPRINT_MISMATCH + + +def test_cross_config_rejects_a_broken_side_before_comparing(): + other = [e for e in _evidence(OTHER) if e.rank != 1] + with pytest.raises(P3Error) as excinfo: + check_cross_config(LAYOUT, (BASE, _evidence(BASE)), (OTHER, other)) + assert _verdict(excinfo) == P3Verdict.MISSING_RANK + + +# --- evidence construction --------------------------------------------------- + + +def test_rank_evidence_reads_a_route_plan(): + plan = {"per_token_fingerprints": {"0": "aa", "1": "bb"}} + item = RankEvidence.from_plan(3, "tp1-cp1-dp1-pp1-ep1", 7, plan) + assert item.fingerprints == {0: "aa", 1: "bb"} + assert item.decisions() == frozenset({(7, 0), (7, 1)}) + + +def test_describe_is_stable_and_names_sequence_parallelism(): + assert ParallelConfig().describe() == "tp1-cp1-dp1-pp1-ep1" + assert ParallelConfig(tp=2, sp=True).describe().startswith("tp2sp-") + assert DEFAULT_RANK_ORDER == ("tp", "cp", "ep", "dp", "pp") From 1def611d434d2319494051e2556f770092768276 Mon Sep 17 00:00:00 2001 From: AlanX-Yan Date: Thu, 15 Oct 2026 22:26:06 +0800 Subject: [PATCH 3/4] [DSv4][P3-T08] Bind the router to P7 cross_config P7 already owns the control plane: Planner builds the grid, RuntimeTools materializes, PairedRunner executes, _resume restarts. The contract says T08 supplies a materializer hook and RoutePlan evidence and does not write a second planner, runner or tolerance, so this is an adapter and nothing else. P3_ROUTER_KNOBS extends V1_KNOBS additively, following the precedent in adapters/knobs.py. V1_KNOBS has no training-side pipeline, data or expert parallel size, so the configurations the router must be compared across cannot be expressed in it; nine paths are added and nothing is edited. The parallel sizes are PROCESS-scoped because a process cannot change its EP group mid-run, and router.placement_semantics is carried as a knob so P7's own identity machinery keeps the two readings of physical_expert_id from ever being spanned by one comparison. P3RouterMaterializer builds and validates the router's view -- a ParallelConfig and a PlacementMap -- before reporting anything as applied. A knob the router does not model is reported UNOBSERVABLE rather than waved through: APPLIED means the value was observed in the running system, and claiming it would make a router-only materialization look like a whole-system one. UNOBSERVABLE is in require_executable's rejected set, so a case mixing engine knobs is refused rather than half-run. An unbuildable configuration is ERROR, never FALLBACK, because there is no "something else" that ran instead. rank_token_manifest and materialize_rank_rows connect P7's rank ownership manifest to mocks.materialize_tokens, the hook the start kit left for this. Row slicing is delegated rather than reimplemented so row selection keeps one definition. materialize_rank_rows separates a rank working on a different pipeline stage (MISSING_RANK) from an over-provisioned rank that owns nothing at all (ZERO_ACTIVE_TOKENS); both are legal states and an empty slice would make either look like completed work. Tests drive the real RuntimeTools, not a stand-in: the claim under test is that P7's existing control plane accepts this adapter. Co-Authored-By: Claude Opus 5 --- docs/design/p3-p7-binding.md | 198 ++++++++++++++++++ rl_engine/p3/p7_binding.py | 379 +++++++++++++++++++++++++++++++++++ tests/p3/test_p7_binding.py | 331 ++++++++++++++++++++++++++++++ 3 files changed, 908 insertions(+) create mode 100644 docs/design/p3-p7-binding.md create mode 100644 rl_engine/p3/p7_binding.py create mode 100644 tests/p3/test_p7_binding.py diff --git a/docs/design/p3-p7-binding.md b/docs/design/p3-p7-binding.md new file mode 100644 index 000000000..5ba865632 --- /dev/null +++ b/docs/design/p3-p7-binding.md @@ -0,0 +1,198 @@ +# P3 Router P7 `cross_config` Binding: Design Record + +| Field | Value | +|---|---| +| Document status | Delivered | +| Document date | 2026-10-15 | +| Work package | T08 (CP/DP/PP/EP ownership, placement, P7 `cross_config`) | +| Test baseline | 24 passed (binding); 117 passed / 2 skipped across placement + ownership + binding | +| Scope | A `RuntimeMaterializer` and consumer hooks; no planner, runner, tolerance or collective | +| Depends on | #498, plus `rl_engine/alignment/cross_config` already on the branch | + +> **What this is.** The adapter that lets P7's existing control plane drive +> the router's ownership and placement contracts. The contract says T08 +> supplies a materializer hook and RoutePlan evidence and explicitly does not +> write a second planner, runner or tolerance; this module is bounded by that. + +--- + +## 1. What P7 Already Owns + +`Planner` builds the case grid, `RuntimeTools` materializes a case, +`PairedRunner` executes, `_resume` handles restart, and `ResultCube` collects. +None of that is rebuilt here. The only extension point P7 exposes is the +`RuntimeMaterializer` protocol: + +```python +class RuntimeMaterializer(Protocol): + runtime_kind: str + @property + def implementation_fingerprint(self) -> str: ... + def materialize(self, normalized, descriptors) -> AdapterMaterialization: ... +``` + +`RuntimeTools` enforces a strict contract on what comes back: exactly one +`KnobApplication` per flattened path, no duplicates, no unknown paths, every +path backed by a descriptor, lifecycle and `critical` matching the descriptor, +and — for anything claimed `APPLIED` — `actual == materialized` and +`materialized == requested` unless the descriptor is `derived`. + +The tests drive the real `RuntimeTools` rather than a stand-in. Mocking it +would prove the adapter satisfies a mock, which is precisely the thing not +worth knowing. + +## 2. Knobs: Additive, Following The WS2 Precedent + +`V1_KNOBS` was written against a HuggingFace/FSDP pair. It has `batch.size`, +rollout-side TP and CP, dtypes and a few engine switches. It has **no** +training-side pipeline, data or expert parallel size, so the configurations +the router must be compared across cannot be expressed in it at all. + +`adapters/knobs.py` already set the precedent for this situation: extend +`V1_KNOBS`, never edit it, because deleting or redefining an entry breaks +existing cross-config tests. `P3_ROUTER_KNOBS` follows that exactly — +`V1_KNOBS` entries are reused by identity, and nine paths are added: + +| Path | Lifecycle | Why | +|---|---|---| +| `training.{tensor,context,pipeline,data,expert}_parallel_size` | `PROCESS` | the five axes ownership composes over | +| `training.sequence_parallel` | `PROCESS` | flips TP from `replica` to `partition_by_sequence` | +| `router.placement_version` | `ENGINE_CONSTRUCTION` | placement identity, compared literally | +| `router.placement_strategy` | `ENGINE_CONSTRUCTION` | `block` / `interleave` | +| `router.placement_semantics` | `ENGINE_CONSTRUCTION` | which reading of `physical_expert_id` | + +The parallel sizes are `PROCESS`-scoped: a process cannot change its EP group +mid-run, and P7 uses the lifecycle to decide what it may reuse. A test pins +that a materialized router case reports `process` isolation, so a future +loosening of that scope cannot pass silently. + +`router.placement_semantics` is carried as a knob for a specific reason: the +two readings of `physical_expert_id` (see `p3-placement-map.md`) must never be +spanned by one comparison. Making it a knob means P7's own identity machinery +separates them instead of the router having to notice. + +## 3. Design Decisions + +### D1: A knob the router does not model is `UNOBSERVABLE`, not `APPLIED` + +The materializer owns the parallel sizes, the placement knobs and +`batch.size`. Everything else in a case — `rollout.dtype`, +`rollout.enable_prefix_caching`, `training.attention_backend` — belongs to the +Megatron or vLLM adapter. + +Reporting those as `APPLIED` would be a lie with consequences: `APPLIED` means +the value was observed in the running system, and `RuntimeTools` would then +treat a router-only materialization as a fully executable whole-system one. +`UNOBSERVABLE` is both honest and load-bearing — it is in `require_executable`'s +rejected set, so a case mixing engine knobs with the router adapter is refused +rather than half-run. The reason string names the owning adapter. + +The practical consequence is that the router materializer expects +*router-scoped* cases. A test pins the mixed case degrading correctly. + +### D2: Problems are collected, then reported per knob + +`router_view` returns `(config, placement, problems)` instead of raising. A +raise inside `materialize` would collapse the whole case into one opaque +error with no per-path attribution, and `RuntimeTools` requires one +application per path regardless. Each offending knob therefore carries +`ERROR` with the same reason text, and the aggregate status follows P7's own +precedence rules. + +### D3: An unbuildable configuration is `ERROR`, never `FALLBACK` + +`FALLBACK` means "we ran something else instead". There is no "something +else" for a parallel configuration: if `expert_parallel_size=0` or the +placement strategy has no generator, there is no router view at all. Marking +it `FALLBACK` would let a non-strict run proceed against a configuration that +was never built. + +### D4: No placement is declared rather than invented + +A case that names no `router.placement_version` yields `placement=None`, and +the binding records `None`. Defaulting to the identity placement would make +two configurations look comparable on a dimension neither declared. + +### D5: Two kinds of empty ownership are kept apart + +`materialize_rank_rows` distinguishes a rank that owns *other* layers but not +this one (a different pipeline stage — `MISSING_RANK`) from a rank that owns +nothing at all (more DP ranks than sequences — `ZERO_ACTIVE_TOKENS`, the +kit's own name for it). Both are legal states of a legal configuration; +returning an empty slice for either would make "this rank has no work" look +like "this rank finished its work". + +This surfaced from a test: with one sequence and `dp=3`, ranks 1 and 2 own +nothing, which is degenerate but not a defect. + +### D6: Row slicing delegates to `mocks.materialize_tokens` + +The start kit left that hook for exactly this, and it already defines what a +sliced case is. Re-implementing the slice here would be a second definition +of row selection, which the DoD forbids. This module supplies the token set; +the kit supplies the slice. + +## 4. Consumer Hooks + +| Function | Role | +|---|---| +| `rank_token_manifest(config, layout)` | `{rank: {layer: (token, ...)}}` — P7's rank ownership manifest, one entry per layer because a rank under `pp=1` owns several | +| `materialize_rank_rows(case, config, layout, rank)` | the manifest applied to a recorded case, via `mocks.materialize_tokens` | +| `config_from_binding(binding)` | recover the `ParallelConfig` from what the adapter *reported it built*, not by re-reading the case | +| `layout_from_case(case, layers)` | a single-sequence `TokenLayout` over a case's active tokens; padding never enters the decision space | + +`config_from_binding` reads the binding rather than the case on purpose: a +cross-config comparison should run against what was materialized, so a +divergence between requested and materialized shows up as a comparison input +rather than being silently normalized away. It fails closed with +`UPSTREAM_CONTRACT_MISMATCH` on a binding from another adapter. + +## 5. Test Matrix + +| Case | Expected | +|---|---| +| valid five-axis case | `APPLIED`, executable in strict mode | +| topology reported | `world_size == 32`, all five sizes, `sequence_parallel` | +| ownership table surfaced | the six contract modes, verbatim | +| `sequence_parallel=True` | TP mode flips to `partition_by_sequence` | +| config recovered from the binding | round-trips to the same `describe()` | +| foreign binding | `UPSTREAM_CONTRACT_MISMATCH` | +| no placement declared | `APPLIED`, placement `None` | +| ungeneratable placement strategy | `ERROR`, `require_executable` refuses | +| `expert_parallel_size=0` | `ERROR` | +| engine knob present | `UNOBSERVABLE`, reason names the adapter, not executable | +| path without a descriptor | `SCHEMA_MISMATCH` | +| identical cases / different EP | reusable / not reusable | +| isolation scope | `process` | +| absent dimensions | default to 1 | +| rank order override | honoured, changes the decomposition | +| manifest vs `expected_decisions` | identical for every rank | +| rank rows under `cp=N` | each rank gets exactly its token, union is the whole case | +| rank asked for a layer it does not own | `MISSING_RANK` | +| over-provisioned DP rank | `ZERO_ACTIVE_TOKENS` | +| `layout_from_case` | active tokens only, no padding sentinel | + +## 6. Boundary + +- No collective is implemented or wrapped. +- No second planner, runner, tolerance or ResultCube. +- The materializer builds the router's *view* of a configuration. Engine + construction stays with the Megatron and vLLM adapters; a case needing both + needs both adapters, which is P7's composition question, not the router's. +- `TokenLayout` is still supplied by the caller. Deriving it from a live job's + batch and model configuration remains open. + +## 7. Reproduction + +```bash +cd RL-Kernel # repo root +python -m pytest tests/p3/test_p7_binding.py -q +# -> 24 passed + +python -m pytest tests/p3/test_placement.py tests/p3/test_ownership.py \ + tests/p3/test_p7_binding.py -q +# -> 117 passed, 2 skipped + +python -m ruff check rl_engine/p3/p7_binding.py tests/p3/test_p7_binding.py +python -m ruff format --check rl_engine/p3/ tests/p3/ +``` diff --git a/rl_engine/p3/p7_binding.py b/rl_engine/p3/p7_binding.py new file mode 100644 index 000000000..a3bb63d9f --- /dev/null +++ b/rl_engine/p3/p7_binding.py @@ -0,0 +1,379 @@ +"""P7 ``cross_config`` binding for the MoE router (T08). + +P7 already owns the control plane: `Planner` builds the case grid, +`RuntimeTools` materializes a case, `PairedRunner` executes, and `_resume` +handles restart. The contract says T08 supplies a materializer hook and +RoutePlan evidence and does **not** write a second planner, runner or +tolerance, so this module is an adapter and nothing else. + +Three pieces: + +``P3_ROUTER_KNOBS`` + An *additive* extension of ``V1_KNOBS``, following the precedent set by + the WS2 attention knobs. ``V1_KNOBS`` has no training-side parallel sizes + beyond TP/CP and no expert parallelism at all, so the configurations the + router has to be compared across cannot even be expressed in it. Nothing + in ``V1_KNOBS`` is edited or deleted. + +``P3RouterMaterializer`` + A ``RuntimeMaterializer`` that turns a normalized case into the router's + view of it: a ``ParallelConfig`` and a ``PlacementMap``, both validated + before anything is reported as applied. A knob the router does not model + is reported ``UNOBSERVABLE`` rather than waved through as applied — + claiming to have applied a rollout-engine knob would make the router + materializer look like a whole-system one. + +Consumer hooks + ``rank_token_manifest`` and ``materialize_rank_rows`` connect P7's rank + ownership manifest to ``mocks.materialize_tokens``, which is the hook the + start kit left for exactly this. + +No collective is implemented or wrapped here. +""" + +from collections.abc import Mapping + +from rl_engine.alignment.cross_config.adapters._common import application, flatten +from rl_engine.alignment.cross_config.planner import ( + _NORMALIZERS, + V1_KNOBS, + _positive_int, + _strict_bool, +) +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import ( + IsolationScope, + KnobDescriptor, + MaterializationStatus, +) + +from .contract import P3Error, P3Verdict +from .mocks import materialize_tokens +from .ownership import ParallelConfig, TokenLayout, expected_decisions +from .placement import ( + SEMANTICS, + SEMANTICS_RANK_LOCAL, + STRATEGY_BLOCK, + STRATEGY_INTERLEAVE, + PlacementMap, + placement_problems, +) + +BINDING_SCHEMA = "p3-p7-binding.v1" +RUNTIME_KIND = "p3_router" + +#: Parallel sizes and router placement, added to the V1 allowlist rather than +#: edited into it. TP and CP already exist as training-side knobs in the WS2 +#: attention set; PP, DP and EP do not exist anywhere, and the router is the +#: first consumer that needs them. +P3_ROUTER_KNOB_DESCRIPTORS = ( + KnobDescriptor("training.tensor_parallel_size", IsolationScope.PROCESS, ("training",)), + KnobDescriptor("training.context_parallel_size", IsolationScope.PROCESS, ("training",)), + KnobDescriptor("training.pipeline_parallel_size", IsolationScope.PROCESS, ("training",)), + KnobDescriptor("training.data_parallel_size", IsolationScope.PROCESS, ("training",)), + KnobDescriptor("training.expert_parallel_size", IsolationScope.PROCESS, ("training",)), + KnobDescriptor("training.sequence_parallel", IsolationScope.PROCESS, ("training",)), + KnobDescriptor( + "router.placement_version", + IsolationScope.ENGINE_CONSTRUCTION, + ("training", "rollout"), + ), + KnobDescriptor( + "router.placement_strategy", + IsolationScope.ENGINE_CONSTRUCTION, + ("training", "rollout"), + allowed_values=(STRATEGY_BLOCK, STRATEGY_INTERLEAVE), + ), + KnobDescriptor( + # Which reading of physical_expert_id the recorded evidence uses; see + # docs/design/p3-placement-map.md. Carried as a knob so a comparison + # can never silently span the two readings. + "router.placement_semantics", + IsolationScope.ENGINE_CONSTRUCTION, + ("training", "rollout"), + allowed_values=SEMANTICS, + ), +) + +P3_ROUTER_KNOBS = { + **V1_KNOBS, + **{descriptor.path: descriptor for descriptor in P3_ROUTER_KNOB_DESCRIPTORS}, +} + +P3_ROUTER_NORMALIZERS = { + **_NORMALIZERS, + "training.tensor_parallel_size": _positive_int, + "training.context_parallel_size": _positive_int, + "training.pipeline_parallel_size": _positive_int, + "training.data_parallel_size": _positive_int, + "training.expert_parallel_size": _positive_int, + "training.sequence_parallel": _strict_bool, +} + +#: Paths this materializer actually acts on. Everything else in a case is +#: another adapter's and is reported unobservable. +ROUTER_OWNED_PATHS = frozenset(descriptor.path for descriptor in P3_ROUTER_KNOB_DESCRIPTORS) | { + "batch.size" +} + +_PARALLEL_PATHS = { + "tp": "training.tensor_parallel_size", + "cp": "training.context_parallel_size", + "pp": "training.pipeline_parallel_size", + "dp": "training.data_parallel_size", + "ep": "training.expert_parallel_size", +} + + +def parallel_config_from_knobs(flat, rank_order=None): + """Build a ``ParallelConfig`` from a flattened case. + + Absent dimensions default to 1 rather than raising: a case that varies + only EP should not have to restate the four sizes it leaves alone, and the + planner emits exactly the paths the experiment declares. + """ + sizes = {name: int(flat.get(path, 1)) for name, path in _PARALLEL_PATHS.items()} + kwargs = dict(sizes, sp=bool(flat.get("training.sequence_parallel", False))) + if rank_order is not None: + kwargs["rank_order"] = tuple(rank_order) + return ParallelConfig(**kwargs) + + +def placement_from_knobs(flat, num_experts=None): + """Build the ``PlacementMap`` a case declares, or ``None`` if it declares none. + + A case that names no placement is not given a default one. Inventing a + placement would make two configurations look comparable on a dimension + neither of them declared. + """ + version = flat.get("router.placement_version") + if not version: + return None + ep_size = int(flat.get(_PARALLEL_PATHS["ep"], 1)) + kwargs = {} + if num_experts is not None: + kwargs["num_experts"] = num_experts + return PlacementMap.from_ep_config( + version, + ep_size, + flat.get("router.placement_strategy", STRATEGY_BLOCK), + flat.get("router.placement_semantics", SEMANTICS_RANK_LOCAL), + **kwargs, + ) + + +class P3RouterMaterializer: + """Materialize the router's view of one cross-config case. + + The router is not an engine: this builds and validates the objects the + router's ownership and placement contracts are expressed in, and reports + the configuration that results. Engine construction stays with the + Megatron and vLLM adapters. + """ + + runtime_kind = RUNTIME_KIND + + def __init__( + self, *, layout=None, num_experts=None, rank_order=None, backend_id="rl_engine.p3" + ): + self.layout = layout + self.num_experts = num_experts + self.rank_order = rank_order + self.backend_id = backend_id + + @property + def implementation_fingerprint(self): + from rl_engine.kernels.semantic_registry import implementation_fingerprint + + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "router_view"), + ) + + def router_view(self, flat): + """``(config, placement, problems)`` for a flattened case. + + Problems are returned rather than raised so that each offending knob + can be reported with its own status; a raise here would collapse the + whole case into one opaque error. + """ + problems = [] + config = None + try: + config = parallel_config_from_knobs(flat, self.rank_order) + except P3Error as error: + problems.append(str(error)) + + placement = None + if config is not None: + try: + placement = placement_from_knobs(flat, self.num_experts) + except P3Error as error: + problems.append(str(error)) + if placement is not None: + problems.extend(placement_problems(placement)) + return config, placement, problems + + def materialize(self, normalized, descriptors): + flat = flatten(normalized) + config, placement, problems = self.router_view(flat) + applications = [ + self._application(descriptors, path, value, config, placement, problems) + for path, value in sorted(flat.items()) + ] + + world_size = config.world_size if config is not None else 1 + side = { + "schema_version": BINDING_SCHEMA, + "parallel": config.describe() if config is not None else None, + "rank_order": list(config.rank_order) if config is not None else None, + "ownership": ( + {dim: config.ownership_of(dim) for dim in ("tp", "cp", "dp", "pp", "ep")} + if config is not None + else None + ), + "placement": placement.plan_fragment() if placement is not None else None, + "problems": problems, + } + if self.layout is not None: + side["decisions"] = len(self.layout.decisions()) + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"training": side, "rollout": {}}, + topology={ + "training": { + "world_size": world_size, + **( + {name: getattr(config, name) for name in _PARALLEL_PATHS} + if config is not None + else {} + ), + "sequence_parallel": bool(config.sp) if config is not None else False, + }, + "rollout": {"world_size": 1}, + }, + scorer={"mode": "teacher_forcing", "framework": "rl_engine.p3", "unit": "decision"}, + operator_backends={"training": self.backend_id, "rollout": self.backend_id}, + runtime_kind=self.runtime_kind, + ), + ) + + def _application(self, descriptors, path, value, config, placement, problems): + descriptor = descriptors.get(path) + if descriptor is None: + raise P3Error( + P3Verdict.SCHEMA_MISMATCH, + f"case path {path!r} has no descriptor; pass P3_ROUTER_KNOBS to RuntimeTools", + ) + if path not in ROUTER_OWNED_PATHS: + return application( + descriptor, + value, + value, + None, + MaterializationStatus.UNOBSERVABLE, + f"{path} belongs to an engine adapter; the router materializer cannot observe it", + ) + if problems: + return application( + descriptor, + value, + value, + None, + MaterializationStatus.ERROR, + "; ".join(problems), + ) + return application( + descriptor, + value, + value, + value, + MaterializationStatus.APPLIED, + f"router view built: {config.describe()}" + + (f", placement {placement.version}" if placement is not None else ""), + ) + + +# --- consumer hooks ---------------------------------------------------------- + + +def rank_token_manifest(config, layout): + """P7's rank ownership manifest for one configuration. + + ``{global_rank: {absolute_layer: (global_token_id, ...)}}`` — the shape + ``mocks.materialize_tokens`` consumes, one entry per layer because a rank + under ``pp=1`` owns several. + """ + manifest = {} + for rank in range(config.world_size): + by_layer = {} + for layer, token in sorted(expected_decisions(config, layout, rank)): + by_layer.setdefault(layer, []).append(token) + manifest[rank] = {layer: tuple(tokens) for layer, tokens in by_layer.items()} + return manifest + + +def materialize_rank_rows(case, config, layout, rank): + """Slice a recorded case down to the rows one rank owns at the case's layer. + + Delegates the actual slicing to ``mocks.materialize_tokens``, the hook the + start kit left for T07/T08, so the row-selection semantics stay in one + place. + + Two different empty results are kept apart. A rank that owns other layers + but not this one is working on a different pipeline stage + (``MISSING_RANK``). A rank that owns nothing at all is over-provisioned — + more DP ranks than sequences, say — which is a degenerate but legal + configuration the kit already names ``ZERO_ACTIVE_TOKENS``. Returning an + empty slice for either would make "this rank has no work" look like + "this rank completed its work". + """ + expected = rank_token_manifest(config, layout).get(rank, {}) + if not expected: + raise P3Error( + P3Verdict.ZERO_ACTIVE_TOKENS, + f"rank {rank} owns no decisions at all under {config.describe()}; " + "the configuration provisions more ranks than the layout can fill", + ) + tokens = expected.get(case.absolute_layer) + if tokens is None: + raise P3Error( + P3Verdict.MISSING_RANK, + f"rank {rank} owns layers {sorted(expected)} under {config.describe()}, " + f"not layer {case.absolute_layer}", + ) + return materialize_tokens(case, tokens) + + +def config_from_binding(binding): + """Recover the ``ParallelConfig`` a materialization recorded. + + Reads the binding's topology rather than re-reading the case, so the + comparison runs against what the adapter reported it built. + """ + topology = binding.topology["training"] + if not isinstance(topology, Mapping) or "tp" not in topology: + raise P3Error( + P3Verdict.UPSTREAM_CONTRACT_MISMATCH, + "binding topology was not produced by the P3 router materializer", + ) + return ParallelConfig( + tp=int(topology["tp"]), + cp=int(topology["cp"]), + dp=int(topology["dp"]), + pp=int(topology["pp"]), + ep=int(topology["ep"]), + sp=bool(topology.get("sequence_parallel", False)), + ) + + +def layout_from_case(case, layers): + """A single-sequence ``TokenLayout`` over one recorded case's active tokens.""" + active = case.global_token_id[case.row_active].tolist() + return TokenLayout.flat(tuple(dict.fromkeys(int(t) for t in active)), tuple(layers)) diff --git a/tests/p3/test_p7_binding.py b/tests/p3/test_p7_binding.py new file mode 100644 index 000000000..f302c3863 --- /dev/null +++ b/tests/p3/test_p7_binding.py @@ -0,0 +1,331 @@ +"""P7 cross_config binding tests (T08). + +These drive the *real* `RuntimeTools`, not a stand-in: the point of the +binding is that P7's existing control plane accepts it, and a mock of +`RuntimeTools` would prove nothing about that. The adapter contract +`RuntimeTools` enforces (one application per path, lifecycle and critical +flags matching the descriptor, applied values matching the normalized case) +is therefore exercised on every case here. +""" + +import pytest + +from rl_engine.alignment.cross_config.planner import V1_KNOBS +from rl_engine.alignment.cross_config.runtime import ( + RuntimeBinding, + RuntimeMaterializationError, + RuntimeTools, +) +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + MaterializationStatus, + SemanticIdentitySpec, +) +from rl_engine.p3.contract import P3Error, P3Verdict +from rl_engine.p3.fixtures import catalog +from rl_engine.p3.ownership import ParallelConfig, TokenLayout, expected_decisions +from rl_engine.p3.p7_binding import ( + P3_ROUTER_KNOB_DESCRIPTORS, + P3_ROUTER_KNOBS, + P3_ROUTER_NORMALIZERS, + ROUTER_OWNED_PATHS, + P3RouterMaterializer, + config_from_binding, + layout_from_case, + materialize_rank_rows, + parallel_config_from_knobs, + placement_from_knobs, + rank_token_manifest, +) +from rl_engine.p3.placement import SEMANTICS_RANK_LOCAL + +ROW = ((0, 1, 2, 3),) +IDENTITY = SemanticIdentitySpec( + checkpoint_id="ckpt", + model_version="m1", + tokenizer_policy="frozen", + token_ids=ROW, + selected_token_ids=ROW, + active_mask=((True,) * 4,), + attention_mask=((True,) * 4,), + position_ids=ROW, + pre_update_state="locked", +) +LAYOUT = TokenLayout(sequences=((0, 1, 2, 3), (4, 5, 6, 7)), layers=(0, 1, 2, 3)) + + +def _case(training=None, router=None, extra=None): + requested = { + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "pipeline_parallel_size": 2, + "data_parallel_size": 2, + "expert_parallel_size": 2, + "sequence_parallel": False, + **(training or {}), + }, + "router": { + "placement_version": "ep2-block-v1", + "placement_strategy": "block", + "placement_semantics": SEMANTICS_RANK_LOCAL, + **(router or {}), + }, + "batch": {"size": 8}, + **(extra or {}), + } + return ExperimentCase( + case_id="c0", + experiment_id="e0", + scenario_id="s0", + identity=IDENTITY, + requested=requested, + ) + + +def _materialize(case=None, **kwargs): + tools = RuntimeTools(P3_ROUTER_KNOBS) + kwargs.setdefault("num_experts", 8) + return tools.materialize(case or _case(), P3RouterMaterializer(**kwargs)) + + +# --- knob allowlist ---------------------------------------------------------- + + +def test_router_knobs_extend_v1_without_editing_it(): + assert set(V1_KNOBS) < set(P3_ROUTER_KNOBS) + for path, descriptor in V1_KNOBS.items(): + assert P3_ROUTER_KNOBS[path] is descriptor + assert set(P3_ROUTER_NORMALIZERS) >= set(P3_ROUTER_KNOBS) - { + "router.placement_version", + "router.placement_strategy", + "router.placement_semantics", + } + + +def test_router_knobs_cover_every_parallel_dimension(): + paths = {descriptor.path for descriptor in P3_ROUTER_KNOB_DESCRIPTORS} + for dimension in ("tensor", "context", "pipeline", "data", "expert"): + assert f"training.{dimension}_parallel_size" in paths + assert "training.sequence_parallel" in paths + assert paths <= ROUTER_OWNED_PATHS + + +# --- materialization --------------------------------------------------------- + + +def test_valid_case_is_applied_and_executable(): + materialization = _materialize() + assert materialization.materialized_case.status is MaterializationStatus.APPLIED + assert materialization.executable_in_strict_mode + RuntimeTools.require_executable(materialization, strict=True) + + +def test_binding_reports_the_composed_topology(): + binding = _materialize().binding + topology = binding.topology["training"] + assert topology["world_size"] == 32 # 2*2*2*2*2 + assert [topology[d] for d in ("tp", "cp", "pp", "dp", "ep")] == [2, 2, 2, 2, 2] + assert topology["sequence_parallel"] is False + assert binding.runtime_kind == "p3_router" + + +def test_binding_surfaces_the_ownership_table(): + side = _materialize().binding.side_configs["training"] + assert dict(side["ownership"]) == { + "tp": "replica", + "cp": "partition", + "dp": "partition", + "pp": "replica_by_layer", + "ep": "not_a_token_dimension", + } + assert side["parallel"] == "tp2-cp2-dp2-pp2-ep2" + assert side["placement"]["version"] == "ep2-block-v1" + + +def test_sequence_parallel_flips_the_tp_ownership_mode(): + side = _materialize(_case(training={"sequence_parallel": True})).binding.side_configs + assert side["training"]["ownership"]["tp"] == "partition_by_sequence" + + +def test_config_round_trips_through_the_binding(): + binding = _materialize().binding + assert config_from_binding(binding).describe() == "tp2-cp2-dp2-pp2-ep2" + + +def test_config_from_a_foreign_binding_fails_closed(): + binding = RuntimeBinding( + batch_size=1, + side_configs={"rollout": {}, "training": {}}, + topology={"rollout": {"world_size": 1}, "training": {"world_size": 1}}, + scorer={}, + operator_backends={"rollout": "x", "training": "x"}, + runtime_kind="something_else", + ) + with pytest.raises(P3Error) as excinfo: + config_from_binding(binding) + assert excinfo.value.verdict == P3Verdict.UPSTREAM_CONTRACT_MISMATCH + + +def test_a_case_declaring_no_placement_still_materializes(): + case = _case(router={"placement_version": ""}) + materialization = _materialize(case) + assert materialization.materialized_case.status is MaterializationStatus.APPLIED + assert materialization.binding.side_configs["training"]["placement"] is None + + +# --- failure reporting ------------------------------------------------------- + + +def test_ungeneratable_placement_strategy_is_an_error_not_a_fallback(): + materialization = _materialize(_case(router={"placement_strategy": "nonsense"})) + assert materialization.materialized_case.status is MaterializationStatus.ERROR + assert not materialization.executable_in_strict_mode + reasons = [ + a.evidence["reason"] + for a in materialization.applications + if a.status is MaterializationStatus.ERROR + ] + assert reasons and "nonsense" in reasons[0] + with pytest.raises(RuntimeMaterializationError): + RuntimeTools.require_executable(materialization, strict=False) + + +def test_non_positive_parallel_size_is_an_error(): + materialization = _materialize(_case(training={"expert_parallel_size": 0})) + assert materialization.materialized_case.status is MaterializationStatus.ERROR + + +def test_an_engine_knob_is_unobservable_not_silently_applied(): + """The router materializer must not claim to have applied another adapter's knob.""" + case = _case(extra={"rollout": {"dtype": "bfloat16"}}) + materialization = _materialize(case) + unobservable = { + a.path: a.evidence["reason"] + for a in materialization.applications + if a.status is MaterializationStatus.UNOBSERVABLE + } + assert "rollout.dtype" in unobservable + assert "engine adapter" in unobservable["rollout.dtype"] + assert not materialization.executable_in_strict_mode + + +def test_a_path_without_a_descriptor_fails_closed(): + tools = RuntimeTools({path: P3_ROUTER_KNOBS[path] for path in ("batch.size",)}) + with pytest.raises(P3Error) as excinfo: + tools.materialize(_case(), P3RouterMaterializer(num_experts=8)) + assert excinfo.value.verdict == P3Verdict.SCHEMA_MISMATCH + + +# --- reuse ------------------------------------------------------------------- + + +def test_identical_cases_are_reusable_and_different_topologies_are_not(): + first = _materialize() + second = _materialize() + assert RuntimeTools.can_reuse(first, second) + other = _materialize(_case(training={"expert_parallel_size": 4})) + assert not RuntimeTools.can_reuse(first, other) + + +def test_parallel_sizes_are_process_scoped(): + """Changing EP cannot be done inside a live process, and must say so.""" + assert _materialize().materialized_case.isolation_scope.value == "process" + + +# --- knob readers ------------------------------------------------------------ + + +def test_absent_dimensions_default_to_one(): + config = parallel_config_from_knobs({"training.expert_parallel_size": 4}) + assert (config.tp, config.cp, config.dp, config.pp, config.ep) == (1, 1, 1, 1, 4) + assert config.world_size == 4 + + +def test_no_placement_is_declared_rather_than_invented(): + assert placement_from_knobs({"training.expert_parallel_size": 4}) is None + + +def test_placement_follows_the_declared_semantics(): + flat = { + "router.placement_version": "v1", + "router.placement_strategy": "interleave", + "router.placement_semantics": "global_permutation", + "training.expert_parallel_size": 4, + } + placement = placement_from_knobs(flat, num_experts=8) + assert placement.semantics == "global_permutation" + assert sorted(placement.physical_of) == list(range(8)) + + +def test_rank_order_override_is_honoured(): + order = ("pp", "cp", "ep", "dp", "tp") + config = parallel_config_from_knobs( + {"training.tensor_parallel_size": 2, "training.pipeline_parallel_size": 2}, + rank_order=order, + ) + assert config.rank_order == order + assert config.coordinates(1)["pp"] == 1 + + +# --- consumer hooks ---------------------------------------------------------- + + +def test_rank_token_manifest_covers_the_world_exactly_once_per_carrier(): + config = ParallelConfig(dp=2, pp=2) + manifest = rank_token_manifest(config, LAYOUT) + assert set(manifest) == set(range(config.world_size)) + for rank, by_layer in manifest.items(): + flat = {(layer, token) for layer, tokens in by_layer.items() for token in tokens} + assert flat == expected_decisions(config, LAYOUT, rank) + + +def test_materialize_rank_rows_slices_a_recorded_case(): + case = next(c for c in catalog() if c.row_active.any()) + tokens = [int(t) for t in case.global_token_id[case.row_active]] + layout = layout_from_case(case, (case.absolute_layer,)) + # CP splits within a sequence; DP splits across sequences, and + # layout_from_case builds exactly one sequence. + config = ParallelConfig(cp=len(tokens)) + + seen = [] + for rank in range(config.world_size): + sliced = materialize_rank_rows(case, config, layout, rank) + held = [int(t) for t in sliced.global_token_id[sliced.row_active]] + assert len(held) == 1 + seen.extend(held) + assert sorted(seen) == sorted(tokens) + + +def test_over_provisioned_ranks_are_named_as_such(): + """More DP ranks than sequences is degenerate but legal, and not a gap.""" + case = next(c for c in catalog() if c.row_active.any()) + layout = layout_from_case(case, (case.absolute_layer,)) # one sequence + config = ParallelConfig(dp=3) + assert materialize_rank_rows(case, config, layout, 0) is not None + for rank in (1, 2): + with pytest.raises(P3Error) as excinfo: + materialize_rank_rows(case, config, layout, rank) + assert excinfo.value.verdict == P3Verdict.ZERO_ACTIVE_TOKENS + + +def test_materialize_rank_rows_rejects_a_layer_the_rank_does_not_own(): + case = next(c for c in catalog() if c.row_active.any()) + layout = layout_from_case(case, (case.absolute_layer, case.absolute_layer + 1)) + config = ParallelConfig(pp=2) + # Stage 1 owns the second layer, so the case's own layer is not its work. + with pytest.raises(P3Error) as excinfo: + materialize_rank_rows(case, config, layout, 1) + assert excinfo.value.verdict == P3Verdict.MISSING_RANK + + +def test_layout_from_case_uses_active_tokens_only(): + case = next(c for c in catalog() if not case_fully_active(c)) + layout = layout_from_case(case, (case.absolute_layer,)) + active = {int(t) for t in case.global_token_id[case.row_active]} + assert set(layout.sequences[0]) == active + assert -1 not in layout.sequences[0] # padding never enters the decision space + + +def case_fully_active(case): + return bool(case.row_active.all()) From 00358fa252efd8ceb8d9af9a93baa29bca681205 Mon Sep 17 00:00:00 2001 From: AlanX-Yan Date: Wed, 21 Oct 2026 04:50:53 +0800 Subject: [PATCH 4/4] [DSv4][P3-T08] Record the DP gate-gradient collective as provenance P3 produces dz and stops; P1 owns the gate GEMM backward and the all-reduce that sums dW_gate across DP ranks belongs to the training framework. Nothing here implements or wraps a collective. It is recorded because a reduction has an order. Bucket composition, bucket order, algorithm and accumulation dtype fix the summation order of the gate weight gradient; move any of them and dW_gate shifts in the last bits, which shifts the next step's gate logits and router scores. The router is a discrete decision layer, so a last-bit move near the Top-K boundary flips a selection -- and by then it looks like a router defect. The record lets an unexplained routing change be checked against the collective first. CollectiveProvenance keeps configured_algorithm and algorithm as separate fields: a configured value may not stand in for a readback, and a difference with no fallback_reason is SILENT_FALLBACK. An unsourced record is not an observation, so it is missing evidence rather than weaker evidence. A missing field fails closed with MISSING_PROVENANCE instead of defaulting, since a default would make two unrelated runs compare equal. At DP=1 no all-reduce runs and none() is valid; demanding evidence of a collective that does not run would be theatre. The fingerprint covers everything that can move the summation order and excludes source and notes, so one collective observed two ways does not read as drift. collective_delta describes differences rather than raising, because two DP sizes must have different collectives; require_same_reduction is the stricter gate for a comparison that claims the collective is not its axis. In the P7 binding a missing record and a malformed one take different paths: DP>1 with no record is UNOBSERVABLE on training.data_parallel_size -- nobody failed to build anything, the router was handed no observation -- while a supplied record that does not hold together is ERROR. UNOBSERVABLE is in require_executable's rejected set, so a DP case with no collective evidence is not strictly executable. Co-Authored-By: Claude Opus 5 --- docs/design/p3-gate-grad-collective.md | 155 ++++++++++++ docs/design/p3-p7-binding.md | 11 +- rl_engine/p3/collective_provenance.py | 333 +++++++++++++++++++++++++ rl_engine/p3/p7_binding.py | 59 ++++- tests/p3/test_collective_provenance.py | 278 +++++++++++++++++++++ tests/p3/test_p7_binding.py | 66 +++++ 6 files changed, 897 insertions(+), 5 deletions(-) create mode 100644 docs/design/p3-gate-grad-collective.md create mode 100644 rl_engine/p3/collective_provenance.py create mode 100644 tests/p3/test_collective_provenance.py diff --git a/docs/design/p3-gate-grad-collective.md b/docs/design/p3-gate-grad-collective.md new file mode 100644 index 000000000..ca50cfaf1 --- /dev/null +++ b/docs/design/p3-gate-grad-collective.md @@ -0,0 +1,155 @@ +# Gate Parameter-Gradient Collective Provenance: Design Record + +| Field | Value | +|---|---| +| Document status | Delivered | +| Document date | 2026-10-21 | +| Work package | T08, contract §4 ("记录 DP 下 gate parameter-grad collective 的 bucket 顺序/算法") | +| Test baseline | 45 passed (collective); 166 passed / 2 skipped across the T08 modules | +| Scope | A provenance record and its gates. **No collective is implemented or wrapped.** | +| Depends on | #498 | + +--- + +## 1. Why P3 Records Something It Does Not Own + +P3 produces `dz` and stops. P1 owns the gate GEMM backward, and the all-reduce +that sums `dW_gate` across data-parallel ranks belongs to the training +framework. The contract is explicit in three places that P3 implements neither +(§2.1, §2.7, §4). + +It is recorded because it is a **reduction, and a reduction has an order**. +Bucket composition, bucket order, the algorithm and the accumulation dtype +together fix the floating-point summation order of the gate weight gradient. +Change any of them and `dW_gate` moves in the last bits, which moves the next +step's gate logits, which moves router scores — and the router is a discrete +decision layer, so a last-bit move near the Top-K boundary flips a selection. + +By the time anyone sees it, that looks like a router defect. The record exists +so an unexplained routing change can be checked against the collective +*before* anyone goes looking inside the router. It is an attribution tool, not +a computation. + +## 2. Design Decisions + +### D1: Configured and actual are separate fields + +The contract repeats that a configured value may not stand in for an actual +readback (§2.3, and `fallback_reason` non-empty is fail-closed). The record +carries `configured_algorithm` and `algorithm` as two fields, and a +difference with no `fallback_reason` is `SILENT_FALLBACK(66)` — a silent +substitution, not a rounding of the truth. + +`source` names how the readback was obtained. A record without it is not an +observation, so it is treated as missing evidence rather than as weaker +evidence; there is no third category between "observed" and "not observed". + +### D2: Absent evidence is not blank evidence + +A missing field fails closed with `MISSING_PROVENANCE(67)` instead of +defaulting. A default would make two unrelated runs compare equal, which is +the one outcome a provenance record must never produce. + +### D3: DP=1 needs no record + +No all-reduce runs, so there is nothing to observe. +`CollectiveProvenance.none()` says exactly that and validates. Demanding +evidence of a collective that does not run would be fail-closed theatre, and +theatre trains people to route around the gate. + +### D4: A missing record is a gap; a malformed record is a defect + +In the P7 binding these take different paths: + +| Situation | Status | Why | +|---|---|---| +| DP>1, no record supplied | `UNOBSERVABLE` on `training.data_parallel_size` | nobody failed to build anything; the router was handed no observation of a reduction P1 owns | +| DP>1, record supplied but inconsistent | `ERROR` | someone did hand over evidence and it does not hold together | + +`UNOBSERVABLE` is in `require_executable`'s rejected set, so a DP case with no +collective evidence is not strictly executable. The DP size is the knob that +reports the hole, because it is the knob whose numeric effect the missing +record hides. + +### D5: An unrecognised algorithm is rejected + +`KNOWN_ALGORITHMS` is a closed set. A reduction nobody has named cannot be +reasoned about, and "unattributable" is precisely the state this record +exists to prevent. A new implementation is addable, but through a contract +delta rather than by appearing in a field. + +### D6: The fingerprint covers order, not observation + +`source` and `notes` are excluded from `fingerprint()`. How the same behaviour +was observed is not a numeric difference; including it would make one +collective read back two ways look like drift. Everything that can move the +summation order is included, and a test parametrizes over all nine such +fields. + +### D7: A difference between configurations is attribution, not failure + +`collective_delta` returns descriptions, never raises. Two configurations with +different DP sizes *must* have different collectives. The delta names what +moved — including the first position at which the reduction order diverged — +so a downstream drift can be checked against it. + +`require_same_reduction` is the separate, stricter gate: in a one-at-a-time +experiment where the collective is not the declared axis, any difference is an +uncontrolled variable and the comparison would misattribute its effect to +whatever *was* declared. It validates both sides first, so a broken side is +reported as broken rather than as a difference. + +## 3. Test Matrix + +Every negative case injects exactly one defect. + +| Injected defect | Verdict | +|---|---| +| unsupported schema version | 11 `SCHEMA_MISMATCH` | +| `dp_size < 1`, `gradient_accumulation_steps < 1` | 11 | +| rank list shorter than `dp_size`, or a duplicate rank | 11 | +| unrecognised algorithm / reduce op | 11 | +| bucket with no parameters, or no dtype | 11 | +| one parameter reduced in two buckets | 11 | +| bucket indices not distinct and ascending | 11 | +| any of six required fields empty | 67 `MISSING_PROVENANCE` | +| no buckets | 67 | +| `fallback_reason` set | 66 `SILENT_FALLBACK` | +| configured != actual with no reason | 66 | +| malformed **and** incomplete | **11** (structure precedes completeness) | +| incomplete **and** a declared fallback | **67** (completeness precedes fallback) | +| comparison where the collective is not held constant | 69 `UPSTREAM_CONTRACT_MISMATCH` | + +Controls (must **not** be flagged or alter the fingerprint): + +| Case | Expected | +|---|---| +| DP=1 with an empty record | valid | +| the same collective read back by a different `source` | identical fingerprint | +| differing `notes` | identical fingerprint | +| two valid records with different `dp_size` | a delta, both still valid | + +Binding integration: DP>1 with no record is `UNOBSERVABLE` and not strictly +executable; DP=1 needs none; a supplied record's fingerprint reaches +`side_configs`; a malformed record is `ERROR`. + +## 4. Boundary + +- Nothing here calls, configures or wraps a collective. +- The record is supplied by whoever can observe the framework; this module + defines its shape, validates it, fingerprints it and diffs it. +- Bucket contents are taken as given. Verifying that the declared buckets + match the parameters a framework actually registered needs a live training + job and belongs with P1. + +## 5. Reproduction + +```bash +cd RL-Kernel # repo root +python -m pytest tests/p3/test_collective_provenance.py -q +# -> 45 passed + +python -m pytest tests/p3/test_placement.py tests/p3/test_ownership.py \ + tests/p3/test_p7_binding.py tests/p3/test_collective_provenance.py -q +# -> 166 passed, 2 skipped +``` diff --git a/docs/design/p3-p7-binding.md b/docs/design/p3-p7-binding.md index 5ba865632..34de25afc 100644 --- a/docs/design/p3-p7-binding.md +++ b/docs/design/p3-p7-binding.md @@ -5,7 +5,7 @@ | Document status | Delivered | | Document date | 2026-10-15 | | Work package | T08 (CP/DP/PP/EP ownership, placement, P7 `cross_config`) | -| Test baseline | 24 passed (binding); 117 passed / 2 skipped across placement + ownership + binding | +| Test baseline | 28 passed (binding); 166 passed / 2 skipped across the T08 modules | | Scope | A `RuntimeMaterializer` and consumer hooks; no planner, runner, tolerance or collective | | Depends on | #498, plus `rl_engine/alignment/cross_config` already on the branch | @@ -171,10 +171,17 @@ rather than being silently normalized away. It fails closed with | rank asked for a layer it does not own | `MISSING_RANK` | | over-provisioned DP rank | `ZERO_ACTIVE_TOKENS` | | `layout_from_case` | active tokens only, no padding sentinel | +| DP>1 with no collective record | `UNOBSERVABLE`, not strictly executable | +| DP=1 | no collective record required | +| supplied collective record | fingerprint reaches `side_configs` | +| malformed collective record | `ERROR`, not a gap | ## 6. Boundary -- No collective is implemented or wrapped. +- No collective is implemented or wrapped. The gate parameter-gradient + all-reduce is *recorded* as provenance and carried in `side_configs`; see + `p3-gate-grad-collective.md`. A DP case supplying no such record is not + strictly executable, because its reduction order is unobserved. - No second planner, runner, tolerance or ResultCube. - The materializer builds the router's *view* of a configuration. Engine construction stays with the Megatron and vLLM adapters; a case needing both diff --git a/rl_engine/p3/collective_provenance.py b/rl_engine/p3/collective_provenance.py new file mode 100644 index 000000000..b3548775e --- /dev/null +++ b/rl_engine/p3/collective_provenance.py @@ -0,0 +1,333 @@ +"""Gate parameter-gradient collective provenance under DP (T08). + +P3 produces ``dz`` and stops. P1 owns the gate GEMM backward, and the +all-reduce that sums ``dW_gate`` across data-parallel ranks is the training +framework's. **Nothing here implements or wraps a collective.** + +It is recorded because it is a reduction, and a reduction has an order. Bucket +composition, bucket order, the algorithm and the accumulation dtype together +decide the floating-point summation order of the gate weight gradient. Change +any of them and ``dW_gate`` changes in the last bits, which moves the next +step's gate logits, which moves router scores, which -- the router being a +discrete decision layer -- can flip a Top-K selection. By then the drift looks +like a router defect. + +So the record exists to make that attributable: an unexplained routing change +between two runs can be checked against their collective provenance before +anyone goes looking inside the router. + +Two things the contract insists on and this module enforces: + +- **A configured value may not pose as an actual readback.** The record + carries both and a silent substitution is ``SILENT_FALLBACK``, not a + rounding of the truth. +- **Absent evidence is not blank evidence.** A missing field fails closed with + ``MISSING_PROVENANCE`` rather than defaulting, because a default would make + two unrelated runs compare equal. + +At ``dp_size == 1`` there is no all-reduce, so there is nothing to record. +``CollectiveProvenance.none()`` says exactly that and is valid; demanding +evidence of a collective that does not run would be fail-closed theatre. +""" + +from dataclasses import dataclass, field + +from .contract import P3Verdict +from .oracle import require +from .serialization import fingerprint + +COLLECTIVE_SCHEMA = "p3-gate-grad-collective.v1" + +#: Reduction algorithms whose summation order is determined by the algorithm +#: plus the rank count. An implementation outside this set may still be used, +#: but it has to be named in a contract delta first: an unrecognised algorithm +#: is unattributable, which is the one thing this record exists to prevent. +KNOWN_ALGORITHMS = ( + "ring", + "tree", + "reduce_scatter_all_gather", + "double_binary_tree", +) + +KNOWN_REDUCE_OPS = ("sum", "avg", "premul_sum") + + +@dataclass(frozen=True) +class GradientBucket: + """One all-reduce bucket, in the order its parameters are reduced. + + ``parameters`` is ordered, not a set: the order inside a bucket is part of + the summation order, and two frameworks that bucket the same parameters + differently produce different last bits. + """ + + index: int + parameters: tuple + numel: int + dtype: str + + def to_record(self): + return { + "index": self.index, + "parameters": list(self.parameters), + "numel": self.numel, + "dtype": self.dtype, + } + + +@dataclass(frozen=True) +class CollectiveProvenance: + """What the gate-gradient all-reduce actually did, as evidence. + + ``configured_algorithm`` is what was asked for and ``algorithm`` is what + was read back. They are separate fields on purpose; collapsing them would + let a configured value stand in for an observation. + + ``source`` names how the readback was obtained. An unsourced record is not + an observation, so it is treated as missing evidence rather than as a + weaker kind of evidence. + """ + + dp_size: int + dp_ranks: tuple = () + algorithm: str = "" + configured_algorithm: str = "" + reduce_op: str = "" + accumulate_dtype: str = "" + buckets: tuple = () + overlap_with_backward: bool = False + gradient_accumulation_steps: int = 1 + bucket_cap_mb: float = 0.0 + process_group: str = "" + source: str = "" + fallback_reason: str = "" + schema_version: str = COLLECTIVE_SCHEMA + notes: dict = field(default_factory=dict) + + @classmethod + def none(cls, dp_size=1): + """The record for a configuration that runs no gate-gradient collective.""" + return cls(dp_size=dp_size, source="no_collective") + + @property + def required(self): + """Whether this configuration runs a collective at all.""" + return self.dp_size > 1 + + def reduction_order(self): + """The parameter-level summation order this record pins down. + + Returned as data rather than a string so a caller can diff two orders + and name the first parameter that moved. + """ + return tuple( + (bucket.index, parameter) for bucket in self.buckets for parameter in bucket.parameters + ) + + def to_record(self): + return { + "schema_version": self.schema_version, + "dp_size": self.dp_size, + "dp_ranks": list(self.dp_ranks), + "algorithm": self.algorithm, + "configured_algorithm": self.configured_algorithm, + "reduce_op": self.reduce_op, + "accumulate_dtype": self.accumulate_dtype, + "buckets": [bucket.to_record() for bucket in self.buckets], + "overlap_with_backward": self.overlap_with_backward, + "gradient_accumulation_steps": self.gradient_accumulation_steps, + "bucket_cap_mb": self.bucket_cap_mb, + "process_group": self.process_group, + "source": self.source, + "fallback_reason": self.fallback_reason, + "notes": dict(sorted(self.notes.items())), + } + + def fingerprint(self): + """Hash over everything that can change the summation order. + + ``source`` and ``notes`` are excluded: how the same behaviour was + observed is not itself a numeric difference, and including it would + make two identical collectives read back by different mechanisms look + like a drift. + """ + record = self.to_record() + for excluded in ("source", "notes"): + record.pop(excluded) + return fingerprint(record) + + +# --- validation -------------------------------------------------------------- + + +def collective_problems(record): + """Structural defects in a record (empty list == usable evidence).""" + problems = [] + if record.schema_version != COLLECTIVE_SCHEMA: + problems.append(f"unsupported schema {record.schema_version!r}") + if record.dp_size < 1: + problems.append(f"dp_size must be at least 1, got {record.dp_size}") + if record.gradient_accumulation_steps < 1: + problems.append( + f"gradient_accumulation_steps must be at least 1, " + f"got {record.gradient_accumulation_steps}" + ) + if not record.required: + return problems + + ranks = tuple(record.dp_ranks) + if len(ranks) != record.dp_size or len(set(ranks)) != len(ranks): + problems.append( + f"dp_ranks must list {record.dp_size} distinct ranks, got {list(ranks)[:8]}" + ) + if record.algorithm and record.algorithm not in KNOWN_ALGORITHMS: + problems.append( + f"algorithm {record.algorithm!r} is not one of {KNOWN_ALGORITHMS}; " + "an unrecognised reduction is unattributable" + ) + if record.reduce_op and record.reduce_op not in KNOWN_REDUCE_OPS: + problems.append(f"reduce_op {record.reduce_op!r} is not one of {KNOWN_REDUCE_OPS}") + + seen, indices = {}, [] + for bucket in record.buckets: + indices.append(bucket.index) + if not bucket.parameters: + problems.append(f"bucket {bucket.index} lists no parameters") + if bucket.numel < 0: + problems.append(f"bucket {bucket.index} has negative numel {bucket.numel}") + if not bucket.dtype: + problems.append(f"bucket {bucket.index} declares no dtype") + for parameter in bucket.parameters: + if parameter in seen: + problems.append( + f"parameter {parameter!r} appears in buckets " + f"{seen[parameter]} and {bucket.index}; a gradient is reduced once" + ) + seen[parameter] = bucket.index + if record.buckets and indices != sorted(set(indices)): + problems.append(f"bucket indices must be distinct and ascending, got {indices[:8]}") + return problems + + +def validate_collective_provenance(record): + """Fail closed on unusable gate-gradient collective evidence. + + Order is structure -> completeness -> fallback. Structure first because a + malformed record cannot be judged complete; completeness before fallback + because a fallback reported by a record that is missing its readback is + not an observation of a fallback. + """ + problems = collective_problems(record) + require(not problems, P3Verdict.SCHEMA_MISMATCH, "; ".join(problems)) + + if not record.required: + return + + missing = [ + name + for name in ( + "algorithm", + "configured_algorithm", + "reduce_op", + "accumulate_dtype", + "process_group", + "source", + ) + if not getattr(record, name) + ] + if not record.buckets: + missing.append("buckets") + require( + not missing, + P3Verdict.MISSING_PROVENANCE, + f"gate-gradient collective evidence is incomplete at dp={record.dp_size}: " + f"missing {sorted(missing)}", + ) + + require( + not record.fallback_reason, + P3Verdict.SILENT_FALLBACK, + f"gate-gradient collective fell back: {record.fallback_reason}", + ) + require( + record.algorithm == record.configured_algorithm, + P3Verdict.SILENT_FALLBACK, + f"configured collective {record.configured_algorithm!r} but read back " + f"{record.algorithm!r} with no fallback_reason; a configured value may not " + "pose as an observation", + ) + + +# --- attribution -------------------------------------------------------------- + + +def collective_delta(base, other): + """Everything that differs between two records, as attribution evidence. + + Returns descriptions rather than raising. Two configurations with + different DP sizes *must* have different collectives, so a difference is + not by itself a defect -- it is the thing a downstream drift should be + checked against first. + """ + deltas = [] + for name in ( + "dp_size", + "algorithm", + "reduce_op", + "accumulate_dtype", + "gradient_accumulation_steps", + "overlap_with_backward", + "bucket_cap_mb", + "process_group", + ): + left, right = getattr(base, name), getattr(other, name) + if left != right: + deltas.append(f"{name}: {left!r} -> {right!r}") + + left_order, right_order = base.reduction_order(), other.reduction_order() + if left_order != right_order: + first = next( + ( + index + for index, (a, b) in enumerate(zip(left_order, right_order, strict=False)) + if a != b + ), + min(len(left_order), len(right_order)), + ) + detail = f"reduction order differs from position {first}" + if first < min(len(left_order), len(right_order)): + detail += f": {left_order[first]} -> {right_order[first]}" + elif len(left_order) != len(right_order): + detail += f": {len(left_order)} vs {len(right_order)} reduced parameters" + deltas.append(detail) + return tuple(deltas) + + +def require_same_reduction(base, other): + """Fail closed when a comparison claims DP is not the variable under test. + + In a one-at-a-time experiment where the gate-gradient collective is not + the declared axis, a difference here is an uncontrolled variable: the + comparison would attribute its effect to whatever *was* declared. + """ + validate_collective_provenance(base) + validate_collective_provenance(other) + deltas = collective_delta(base, other) + require( + not deltas, + P3Verdict.UPSTREAM_CONTRACT_MISMATCH, + "gate-gradient collective is not held constant across the comparison: " + "; ".join(deltas), + ) + + +def collective_evidence(record): + """The provenance fragment a RoutePlan or binding carries. + + Additive and self-describing: a consumer that does not know about the + gate-gradient collective sees one more key, and one that does gets the + fingerprint without having to re-derive it. + """ + return { + "gate_grad_collective": record.to_record(), + "gate_grad_collective_fingerprint": record.fingerprint(), + } diff --git a/rl_engine/p3/p7_binding.py b/rl_engine/p3/p7_binding.py index a3bb63d9f..101d1c685 100644 --- a/rl_engine/p3/p7_binding.py +++ b/rl_engine/p3/p7_binding.py @@ -50,6 +50,11 @@ MaterializationStatus, ) +from .collective_provenance import ( + CollectiveProvenance, + collective_evidence, + collective_problems, +) from .contract import P3Error, P3Verdict from .mocks import materialize_tokens from .ownership import ParallelConfig, TokenLayout, expected_decisions @@ -177,11 +182,18 @@ class P3RouterMaterializer: runtime_kind = RUNTIME_KIND def __init__( - self, *, layout=None, num_experts=None, rank_order=None, backend_id="rl_engine.p3" + self, + *, + layout=None, + num_experts=None, + rank_order=None, + collective=None, + backend_id="rl_engine.p3", ): self.layout = layout self.num_experts = num_experts self.rank_order = rank_order + self.collective = collective self.backend_id = backend_id @property @@ -221,8 +233,10 @@ def router_view(self, flat): def materialize(self, normalized, descriptors): flat = flatten(normalized) config, placement, problems = self.router_view(flat) + collective, collective_gap, collective_defects = self._collective_for(config) + problems = problems + collective_defects applications = [ - self._application(descriptors, path, value, config, placement, problems) + self._application(descriptors, path, value, config, placement, problems, collective_gap) for path, value in sorted(flat.items()) ] @@ -239,6 +253,8 @@ def materialize(self, normalized, descriptors): "placement": placement.plan_fragment() if placement is not None else None, "problems": problems, } + if collective is not None: + side.update(collective_evidence(collective)) if self.layout is not None: side["decisions"] = len(self.layout.decisions()) return AdapterMaterialization( @@ -264,7 +280,33 @@ def materialize(self, normalized, descriptors): ), ) - def _application(self, descriptors, path, value, config, placement, problems): + def _collective_for(self, config): + """Resolve the gate-gradient collective record for a configuration. + + Returns ``(record, gap, defects)``. A *missing* record under DP is a + gap, not a defect: the router materializer did not fail to build + something, it was handed no observation of a reduction that P1 and the + training framework own. A *malformed* record is a defect, because + someone did hand over evidence and it does not hold together. + + At DP=1 no all-reduce runs, so the record is legitimately empty. + """ + if config is None: + return None, None, [] + if config.dp == 1: + return CollectiveProvenance.none(dp_size=1), None, [] + if self.collective is None: + return ( + None, + f"dp={config.dp} runs a gate-gradient all-reduce but no collective " + "provenance was supplied; its reduction order is unobserved", + [], + ) + return self.collective, None, collective_problems(self.collective) + + def _application( + self, descriptors, path, value, config, placement, problems, collective_gap=None + ): descriptor = descriptors.get(path) if descriptor is None: raise P3Error( @@ -289,6 +331,17 @@ def _application(self, descriptors, path, value, config, placement, problems): MaterializationStatus.ERROR, "; ".join(problems), ) + if collective_gap is not None and path == _PARALLEL_PATHS["dp"]: + # The DP size is the knob whose numeric effect the missing record + # hides, so it is the knob that reports the hole. + return application( + descriptor, + value, + value, + None, + MaterializationStatus.UNOBSERVABLE, + collective_gap, + ) return application( descriptor, value, diff --git a/tests/p3/test_collective_provenance.py b/tests/p3/test_collective_provenance.py new file mode 100644 index 000000000..b03cbd21f --- /dev/null +++ b/tests/p3/test_collective_provenance.py @@ -0,0 +1,278 @@ +"""Gate parameter-gradient collective provenance tests (T08). + +P3 does not implement the all-reduce; it records what the all-reduce did, so +that a routing change can be checked against the reduction order before +anyone goes looking inside the router. These tests pin that the record is +usable as evidence: complete or fail-closed, fingerprinted over exactly what +changes the summation order, and able to name what moved between two runs. + +Every negative case injects exactly one defect. +""" + +import dataclasses + +import pytest + +from rl_engine.p3.collective_provenance import ( + COLLECTIVE_SCHEMA, + CollectiveProvenance, + GradientBucket, + collective_delta, + collective_evidence, + collective_problems, + require_same_reduction, + validate_collective_provenance, +) +from rl_engine.p3.contract import P3Error, P3Verdict + +WEIGHT = GradientBucket(0, ("gate.weight",), 4096 * 256, "fp32") +BIAS = GradientBucket(1, ("gate.bias",), 256, "fp32") + + +def _record(**overrides): + fields = { + "dp_size": 2, + "dp_ranks": (0, 1), + "algorithm": "ring", + "configured_algorithm": "ring", + "reduce_op": "avg", + "accumulate_dtype": "fp32", + "process_group": "dp:0-1", + "source": "torch.distributed readback", + "buckets": (WEIGHT, BIAS), + **overrides, + } + return CollectiveProvenance(**fields) + + +def _verdict(excinfo): + return excinfo.value.verdict + + +# --- usable evidence --------------------------------------------------------- + + +def test_complete_record_validates(): + record = _record() + assert collective_problems(record) == [] + validate_collective_provenance(record) + assert record.required + + +def test_dp1_runs_no_collective_and_needs_no_evidence(): + record = CollectiveProvenance.none() + assert not record.required + validate_collective_provenance(record) + assert record.to_record()["source"] == "no_collective" + + +def test_reduction_order_is_bucket_then_parameter(): + assert _record().reduction_order() == ((0, "gate.weight"), (1, "gate.bias")) + + +def test_evidence_fragment_carries_record_and_fingerprint(): + record = _record() + fragment = collective_evidence(record) + assert fragment["gate_grad_collective"]["schema_version"] == COLLECTIVE_SCHEMA + assert fragment["gate_grad_collective_fingerprint"] == record.fingerprint() + + +# --- fingerprint ------------------------------------------------------------- + + +@pytest.mark.parametrize( + "overrides", + [ + {"dp_size": 4, "dp_ranks": (0, 1, 2, 3)}, + {"algorithm": "tree", "configured_algorithm": "tree"}, + {"reduce_op": "sum"}, + {"accumulate_dtype": "bf16"}, + {"buckets": (BIAS, WEIGHT)}, + {"gradient_accumulation_steps": 2}, + {"overlap_with_backward": True}, + {"bucket_cap_mb": 25.0}, + {"process_group": "dp:2-3"}, + ], +) +def test_anything_that_moves_the_summation_order_moves_the_fingerprint(overrides): + assert _record().fingerprint() != _record(**overrides).fingerprint() + + +@pytest.mark.parametrize( + "overrides", + [{"source": "nccl env dump"}, {"notes": {"host": "node-7"}}], +) +def test_how_it_was_observed_is_not_a_numeric_difference(overrides): + """Control: the same collective seen two ways must not read as drift.""" + assert _record().fingerprint() == _record(**overrides).fingerprint() + + +# --- structural defects ------------------------------------------------------ + + +def _assert_schema_mismatch(record, needle): + problems = collective_problems(record) + assert any(needle in problem for problem in problems), problems + with pytest.raises(P3Error) as excinfo: + validate_collective_provenance(record) + assert _verdict(excinfo) == P3Verdict.SCHEMA_MISMATCH + + +def test_unsupported_schema_version(): + _assert_schema_mismatch(_record(schema_version="p3-gate-grad-collective.v0"), "unsupported") + + +def test_non_positive_dp_size(): + _assert_schema_mismatch(_record(dp_size=0), "dp_size") + + +def test_non_positive_accumulation_steps(): + _assert_schema_mismatch(_record(gradient_accumulation_steps=0), "gradient_accumulation_steps") + + +def test_rank_list_shorter_than_dp_size(): + _assert_schema_mismatch(_record(dp_ranks=(0,)), "distinct ranks") + + +def test_duplicate_rank(): + _assert_schema_mismatch(_record(dp_ranks=(0, 0)), "distinct ranks") + + +def test_unrecognised_algorithm_is_unattributable(): + _assert_schema_mismatch(_record(algorithm="homegrown"), "unattributable") + + +def test_unrecognised_reduce_op(): + _assert_schema_mismatch(_record(reduce_op="max"), "reduce_op") + + +def test_empty_bucket(): + _assert_schema_mismatch( + _record(buckets=(GradientBucket(0, (), 0, "fp32"),)), "lists no parameters" + ) + + +def test_bucket_without_a_dtype(): + _assert_schema_mismatch( + _record(buckets=(GradientBucket(0, ("gate.weight",), 1, ""),)), "declares no dtype" + ) + + +def test_one_parameter_reduced_in_two_buckets(): + duplicate = GradientBucket(1, ("gate.weight",), 1, "fp32") + _assert_schema_mismatch(_record(buckets=(WEIGHT, duplicate)), "reduced once") + + +def test_bucket_indices_must_ascend(): + _assert_schema_mismatch( + _record(buckets=(BIAS, dataclasses.replace(WEIGHT, index=0))), "ascending" + ) + + +# --- missing evidence -------------------------------------------------------- + + +@pytest.mark.parametrize( + "field", + [ + "algorithm", + "configured_algorithm", + "reduce_op", + "accumulate_dtype", + "process_group", + "source", + ], +) +def test_a_missing_field_is_missing_evidence_not_a_default(field): + with pytest.raises(P3Error) as excinfo: + validate_collective_provenance(_record(**{field: ""})) + assert _verdict(excinfo) == P3Verdict.MISSING_PROVENANCE + assert field in str(excinfo.value) + + +def test_no_buckets_is_missing_evidence(): + with pytest.raises(P3Error) as excinfo: + validate_collective_provenance(_record(buckets=())) + assert _verdict(excinfo) == P3Verdict.MISSING_PROVENANCE + + +def test_an_unsourced_record_is_not_an_observation(): + with pytest.raises(P3Error) as excinfo: + validate_collective_provenance(_record(source="")) + assert _verdict(excinfo) == P3Verdict.MISSING_PROVENANCE + + +# --- fallback ---------------------------------------------------------------- + + +def test_declared_fallback_fails_closed(): + with pytest.raises(P3Error) as excinfo: + validate_collective_provenance(_record(fallback_reason="nccl unavailable, used gloo")) + assert _verdict(excinfo) == P3Verdict.SILENT_FALLBACK + + +def test_configured_value_may_not_pose_as_an_observation(): + """Read back something other than what was configured, with no reason given.""" + with pytest.raises(P3Error) as excinfo: + validate_collective_provenance(_record(configured_algorithm="tree", algorithm="ring")) + assert _verdict(excinfo) == P3Verdict.SILENT_FALLBACK + assert "pose as an observation" in str(excinfo.value) + + +def test_structure_precedes_completeness(): + """A malformed record cannot be judged complete.""" + with pytest.raises(P3Error) as excinfo: + validate_collective_provenance(_record(dp_ranks=(0, 0), source="")) + assert _verdict(excinfo) == P3Verdict.SCHEMA_MISMATCH + + +def test_completeness_precedes_fallback(): + """A fallback reported by a record with no readback is not an observation.""" + with pytest.raises(P3Error) as excinfo: + validate_collective_provenance(_record(source="", fallback_reason="gloo")) + assert _verdict(excinfo) == P3Verdict.MISSING_PROVENANCE + + +# --- attribution ------------------------------------------------------------- + + +def test_identical_records_have_no_delta(): + assert collective_delta(_record(), _record()) == () + + +def test_delta_names_the_first_position_the_order_moved(): + deltas = collective_delta( + _record(), _record(buckets=(BIAS, dataclasses.replace(WEIGHT, index=1))) + ) + assert any("position 0" in d and "gate.weight" in d and "gate.bias" in d for d in deltas) + + +def test_delta_reports_a_shorter_or_longer_order(): + deltas = collective_delta(_record(), _record(buckets=(WEIGHT,))) + assert any("2 vs 1 reduced parameters" in d for d in deltas) + + +def test_a_different_dp_size_is_a_delta_not_a_defect(): + """Two DP sizes must have different collectives; that is attribution, not failure.""" + other = _record(dp_size=4, dp_ranks=(0, 1, 2, 3)) + deltas = collective_delta(_record(), other) + assert any(d.startswith("dp_size:") for d in deltas) + validate_collective_provenance(other) # still perfectly valid evidence + + +def test_require_same_reduction_passes_on_identical_records(): + require_same_reduction(_record(), _record()) + + +def test_require_same_reduction_rejects_an_uncontrolled_variable(): + with pytest.raises(P3Error) as excinfo: + require_same_reduction(_record(), _record(accumulate_dtype="bf16")) + assert _verdict(excinfo) == P3Verdict.UPSTREAM_CONTRACT_MISMATCH + assert "held constant" in str(excinfo.value) + + +def test_require_same_reduction_validates_both_sides_first(): + """A broken side is reported as broken, not as a difference.""" + with pytest.raises(P3Error) as excinfo: + require_same_reduction(_record(), _record(source="")) + assert _verdict(excinfo) == P3Verdict.MISSING_PROVENANCE diff --git a/tests/p3/test_p7_binding.py b/tests/p3/test_p7_binding.py index f302c3863..352543434 100644 --- a/tests/p3/test_p7_binding.py +++ b/tests/p3/test_p7_binding.py @@ -21,6 +21,7 @@ MaterializationStatus, SemanticIdentitySpec, ) +from rl_engine.p3.collective_provenance import CollectiveProvenance, GradientBucket from rl_engine.p3.contract import P3Error, P3Verdict from rl_engine.p3.fixtures import catalog from rl_engine.p3.ownership import ParallelConfig, TokenLayout, expected_decisions @@ -83,9 +84,27 @@ def _case(training=None, router=None, extra=None): ) +def _collective(dp_size=2, **overrides): + """A complete gate-gradient collective record for a DP configuration.""" + fields = { + "dp_size": dp_size, + "dp_ranks": tuple(range(dp_size)), + "algorithm": "ring", + "configured_algorithm": "ring", + "reduce_op": "avg", + "accumulate_dtype": "fp32", + "process_group": f"dp:0-{dp_size - 1}", + "source": "torch.distributed readback", + "buckets": (GradientBucket(0, ("gate.weight",), 4096 * 256, "fp32"),), + **overrides, + } + return CollectiveProvenance(**fields) + + def _materialize(case=None, **kwargs): tools = RuntimeTools(P3_ROUTER_KNOBS) kwargs.setdefault("num_experts", 8) + kwargs.setdefault("collective", _collective()) return tools.materialize(case or _case(), P3RouterMaterializer(**kwargs)) @@ -329,3 +348,50 @@ def test_layout_from_case_uses_active_tokens_only(): def case_fully_active(case): return bool(case.row_active.all()) + + +# --- gate-gradient collective provenance ------------------------------------- + + +def test_dp_without_collective_evidence_is_unobservable(): + """DP runs an all-reduce; its reduction order cannot simply go unrecorded.""" + tools = RuntimeTools(P3_ROUTER_KNOBS) + materialization = tools.materialize(_case(), P3RouterMaterializer(num_experts=8)) + assert materialization.materialized_case.status is MaterializationStatus.UNOBSERVABLE + assert not materialization.executable_in_strict_mode + gaps = { + a.path: a.evidence["reason"] + for a in materialization.applications + if a.status is MaterializationStatus.UNOBSERVABLE + } + assert "training.data_parallel_size" in gaps + assert "all-reduce" in gaps["training.data_parallel_size"] + + +def test_dp1_needs_no_collective_record(): + """No all-reduce runs at DP=1, so demanding evidence of one is theatre.""" + tools = RuntimeTools(P3_ROUTER_KNOBS) + case = _case(training={"data_parallel_size": 1}) + materialization = tools.materialize(case, P3RouterMaterializer(num_experts=8)) + assert materialization.materialized_case.status is MaterializationStatus.APPLIED + side = materialization.binding.side_configs["training"] + assert side["gate_grad_collective"]["source"] == "no_collective" + + +def test_binding_carries_the_collective_fingerprint(): + side = _materialize().binding.side_configs["training"] + assert side["gate_grad_collective"]["algorithm"] == "ring" + assert side["gate_grad_collective_fingerprint"] == _collective().fingerprint() + + +def test_a_malformed_collective_record_is_an_error_not_a_gap(): + """Evidence was handed over and does not hold together: that is a defect.""" + broken = _collective(dp_ranks=(0, 0)) # two ranks, one of them twice + materialization = _materialize(collective=broken) + assert materialization.materialized_case.status is MaterializationStatus.ERROR + reasons = [ + a.evidence["reason"] + for a in materialization.applications + if a.status is MaterializationStatus.ERROR + ] + assert reasons and "distinct ranks" in reasons[0]