From f16b1269ff9c717d2dde8e6eb3a270ac2b1876ef Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 04:20:23 +0000
Subject: [PATCH 1/5] Raise EP routed-row pricing to each checkpoint's observed
 routing

At EP>1 TrainerRank prices the rows each rank's experts receive at a fixed
allowance over balanced routing (1.4x at EP2); skew beyond it went unseen.
Record each checkpoint's worst-layer routed share at the micro-batch
handoff, from HybridEP's retained tokens_per_expert, through the existing
handoff MAX reduce (EP>1 only: share, epoch, -epoch). A share is recorded
only when every rank observed the same checkpoint route epoch. Epochs are
assigned after a load or snapshot commits on every rank, never reused, and
forgotten on reload or discard.

When a checkpoint's share plus a 0.10 margin exceeds the allowance, its
balanced routed rows are scaled up to match. Charges never go down. The
share is reported in last_forward_telemetry() and planner-miss reports.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .github/workflows/prek.yml                   |   4 +-
 src/art/trainer_rank/_checkpoint.py          |   3 +
 src/art/trainer_rank/_impl.py                | 150 ++++++++-
 tests/unit/test_trainer_rank_routed_share.py | 303 +++++++++++++++++++
 4 files changed, 457 insertions(+), 3 deletions(-)
 create mode 100644 tests/unit/test_trainer_rank_routed_share.py

diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml
index b66ac2c0e..375ffccdc 100644
--- a/.github/workflows/prek.yml
+++ b/.github/workflows/prek.yml
@@ -241,6 +241,7 @@ jobs:
             tests/unit/test_trainer_rank_converted_memory.py \
             tests/unit/test_trainer_rank_layout_memory.py \
             tests/unit/test_context_parallel_retained_bytes.py \
+            tests/unit/test_trainer_rank_routed_share.py \
             tests/unit/test_trainer_rank_split.py \
             tests/unit/test_megatron_compile_garbage.py \
             tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \
@@ -287,4 +288,5 @@ jobs:
             --ignore=tests/unit/test_megatron_compile_garbage.py \
             --ignore=tests/unit/test_trainer_rank_converted_memory.py \
             --ignore=tests/unit/test_trainer_rank_layout_memory.py \
-            --ignore=tests/unit/test_context_parallel_retained_bytes.py
+            --ignore=tests/unit/test_context_parallel_retained_bytes.py \
+            --ignore=tests/unit/test_trainer_rank_routed_share.py
diff --git a/src/art/trainer_rank/_checkpoint.py b/src/art/trainer_rank/_checkpoint.py
index 530a7f35e..bee473c9c 100644
--- a/src/art/trainer_rank/_checkpoint.py
+++ b/src/art/trainer_rank/_checkpoint.py
@@ -1643,6 +1643,7 @@ def snapshot_checkpoint(trainer: TrainerRank, source: str, destination: str) ->
         _restore_slots(model_snapshot)
         trainer._checkpoint_slots.pop(destination, None)
         raise
+    trainer._commit_route_epoch(destination)
     trainer._snapshot_checkpoint_names.add(destination)
     return True
 
@@ -1688,6 +1689,7 @@ def discard_snapshot_checkpoint(trainer: TrainerRank, checkpoint: str) -> None:
         _restore_slots(model_snapshot)
         trainer._checkpoint_slots[checkpoint] = slot
         raise
+    trainer._forget_route_epoch(slot)
 
 
 def _commit_slot(trainer: TrainerRank, source: str, destination: str) -> None:
@@ -1959,6 +1961,7 @@ def commit() -> None:
     except BaseException:
         _rollback_load(trainer, snapshot, temporary, name, previous, group)
         raise
+    trainer._commit_route_epoch(name, previous)
 
 
 def snapshot_prepared_checkpoint(
diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index 6d3d61b98..a205ade62 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -854,6 +854,9 @@ class _CheckpointSlot:
     custom: dict[str, _CustomObject] = dataclass_field(default_factory=dict)
     custom_payload: "PreparedCustomPayload | None" = None
     snapshot: bool = False
+    # Key for this committed content's routing observations; see
+    # ``TrainerRank._commit_route_epoch``.
+    route_epoch: int | None = None
 
 
 @dataclass(frozen=True)
@@ -1553,6 +1556,10 @@ def _expert_lora_weight_storage(
 # Unmeasured EP sizes use the next measured one; above EP8 the allowance grows
 # with log2(EP) up to EP itself (every pair on one rank).
 _EP_ROUTED_ROW_ALLOWANCE = {2: 1.4, 4: 1.6, 8: 2.0}
+# Headroom above a checkpoint's observed routed share when that share exceeds
+# the allowance. For the EP2 policy above, call-to-call change was 0.013 and
+# batch resampling reached about 0.04 above the median.
+_ROUTED_SHARE_MARGIN = 0.10
 
 
 def _ep_routed_row_allowance(ep: int) -> float:
@@ -2118,6 +2125,11 @@ def memory_field(name: str, default: Any = None) -> Any:
         self._slot_stack: list[LoRASlotRef] = []
         self._checkpoint_slots: dict[str, _CheckpointSlot] = {}
         self._snapshot_checkpoint_names: set[str] = set()
+        # Highest routed share observed per checkpoint route epoch, identical
+        # on every rank (``_release_cached_memory_for_backward``).
+        self._route_epochs = 0
+        self._routed_share_max: dict[int, float] = {}
+        self._last_routed_share: float | None = None
         self._prepared_lora_exports: dict[str, tuple[str, _PreparedLoraExport]] = {}
         self._checkpoint_prefetches: dict[str, Future[PreparedCheckpoint]] = {}
         self._checkpoint_prefetch_sources: dict[str, str] = {}
@@ -3090,11 +3102,21 @@ def _release_cached_memory_for_backward(
         # outputs. Forward has already executed: never replan or retry here.
         with self._cache_recovery_episode(error=error) as (owner, started):
             exchange_error: BaseException | None = None
+            # EP is the same on every rank, so the payload shape is too. Every
+            # rank must reach the exchange, whatever its observation does.
+            routing = getattr(getattr(self, "_parallel_shape", None), "ep", 1) > 1
+            share, epoch = -1.0, -1
+            if routing and error is None:
+                try:
+                    share, epoch = self._local_routed_share(plan)
+                except Exception:
+                    share, epoch = -1.0, -1
             try:
-                failed, gradients = self._recovery_reduce(
+                failed, gradients, *observed = self._recovery_reduce(
                     [
                         float(error is not None),
                         float(any(group.grad_enabled for group in plan.groups)),
+                        *((share, float(epoch), -float(epoch)) if routing else ()),
                     ],
                     op="MAX",
                     sync_across_dp=True,
@@ -3107,6 +3129,13 @@ def _release_cached_memory_for_backward(
                 raise self._memory_error_with_reduction_note(error, exchange_error)
             if failed:
                 raise RuntimeError("Forward failed on another rank before handoff")
+            self._last_routed_share = None
+            if observed:
+                # Record only when every rank observed the same checkpoint
+                # epoch; any empty, unsupported or other-slot rank sends -1.
+                share, highest, lowest = observed
+                if highest >= 0 and highest == -lowest:
+                    self._record_routed_share(int(highest), share)
             if not gradients:
                 return
             self._try_cache_recovery(
@@ -3117,6 +3146,60 @@ def _release_cached_memory_for_backward(
                 handoff_grad=any(group.grad_enabled for group in plan.groups),
             )
 
+    def _local_routed_share(self, plan: _AnyForwardPlan) -> tuple[float, int]:
+        """This rank's worst-layer routed share and checkpoint epoch, or -1s.
+
+        HybridEP keeps each MoE layer's received rows per local expert after
+        combine, so at the handoff of a flat plan with one gradient group they
+        are this forward's dispatch. The share is the most loaded layer's
+        received pairs over top-k times the balanced rows that pricing scales
+        (``_plan_group_balanced_rows``), so a share above the allowance means
+        more rows arrived than were priced.
+        """
+        if (
+            not isinstance(plan, _FlatForwardPlan)
+            or len(plan.groups) != 1
+            or not plan.groups[0].grad_enabled
+            or plan.groups[0].packed.tokens.numel() == 0
+            or plan.signature.topology[2] <= 1
+            or not getattr(self, "_ep_group_is_cp_group", False)
+            or not getattr(self, "_moe_memory_supported", False)
+        ):
+            return -1.0, -1
+        epoch = self._route_epoch(plan.groups[0].slot_ref)
+        (balanced,) = self._plan_group_balanced_rows(plan)
+        if epoch is None or balanced <= 0:
+            return -1.0, -1
+        from megatron.core.transformer.moe.token_dispatcher import _HybridEPManager
+
+        managers = [
+            manager
+            for chunk in self.runtime.model
+            for module in chunk.modules()
+            if type(
+                manager := getattr(
+                    getattr(module, "token_dispatcher", None), "_comm_manager", None
+                )
+            )
+            is _HybridEPManager
+        ]
+        if not managers or len(managers) != self._moe_layers:
+            return -1.0, -1
+        received = torch.stack(
+            [manager.tokens_per_expert.detach().sum() for manager in managers]
+        )
+        topk = managers[0].config.moe_router_topk
+        return int(received.max().item()) / (topk * balanced), epoch
+
+    def _record_routed_share(self, epoch: int, share: float) -> None:
+        shares = self._routed_share_max
+        shares[epoch] = max(share, shares.get(epoch, share))
+        self._last_routed_share = share
+        observation = getattr(self, "_planner_observation", None)
+        if observation is not None:
+            replay = observation["replay"]
+            observation["replay"] = lambda: {**replay(), "routed_share": share}
+
     @overload
     def dp_rank_forward(
         self,
@@ -4082,6 +4165,56 @@ def _plan_group_rows(self, plan: _FlatForwardPlan) -> tuple[tuple[int, bool], ..
         )
 
     def _plan_group_routed_rows(self, plan: _FlatForwardPlan) -> tuple[int, ...]:
+        """Balanced rows per group, raised to each slot's observed routing.
+
+        The cold allowance prices top-k x allowance pairs per balanced row. A
+        checkpoint whose worst observed share plus a margin exceeds that is
+        priced at the higher share instead; nothing is ever priced lower.
+        """
+        balanced = self._plan_group_balanced_rows(plan)
+        shares = [self._observed_routed_share(group.slot_ref) for group in plan.groups]
+        if all(share is None for share in shares):
+            return balanced
+        cold = _ep_routed_row_allowance(self._parallel_shape.ep)
+        return tuple(
+            rows
+            if share is None or share + _ROUTED_SHARE_MARGIN <= cold
+            else math.ceil(rows * (share + _ROUTED_SHARE_MARGIN) / cold)
+            for rows, share in zip(balanced, shares, strict=True)
+        )
+
+    def _route_epoch(self, ref: "LoRASlotRef | None") -> int | None:
+        if ref is None or ref.name is None:
+            return None
+        slot = getattr(self, "_checkpoint_slots", {}).get(ref.name)
+        return None if slot is None else slot.route_epoch
+
+    def _observed_routed_share(self, ref: "LoRASlotRef | None") -> float | None:
+        epoch = self._route_epoch(ref)
+        if epoch is None:
+            return None
+        return getattr(self, "_routed_share_max", {}).get(epoch)
+
+    def _commit_route_epoch(
+        self, name: str, previous: _CheckpointSlot | None = None
+    ) -> None:
+        """Key a checkpoint's routing observations after every rank committed it.
+
+        Callers run this once the commit agreed across ranks, in the same order
+        everywhere, so each rank gives the same content the same new epoch.
+        Epochs are never reused. Loading over a name forgets the old content's
+        observations; optimizer steps keep the epoch, so its share only grows.
+        """
+        self._checkpoint_slots[name].route_epoch = self._route_epochs
+        self._route_epochs += 1
+        if previous is not None:
+            self._forget_route_epoch(previous)
+
+    def _forget_route_epoch(self, slot: _CheckpointSlot) -> None:
+        if slot.route_epoch is not None:
+            self._routed_share_max.pop(slot.route_epoch, None)
+
+    def _plan_group_balanced_rows(self, plan: _FlatForwardPlan) -> tuple[int, ...]:
         """Rows one rank's experts receive per group at balanced routing.
 
         HybridEP dispatches the whole EP group's rows. When that group is this
@@ -4762,6 +4895,7 @@ def _split_request_order(
 
     def _reset_planning_telemetry(self) -> None:
         self._planning_seconds_accum = 0.0
+        self._last_routed_share = None
         with self._layout_cache_lock:
             self._speculative_planning_seconds = 0.0
 
@@ -4775,6 +4909,9 @@ def _snapshot_planning_telemetry(
             if isinstance(plan, _SplitForwardPlan)
             else (tuple(range(plan.request_count)),)
         )
+        # Consume the share so a later forward without a handoff reports none.
+        routed_share = getattr(self, "_last_routed_share", None)
+        self._last_routed_share = None
         self._last_forward_telemetry_snapshot = {
             "planning_ms": self._planning_seconds_accum * 1_000.0,
             "speculative_planning_ms": speculative_seconds * 1_000.0,
@@ -4785,6 +4922,7 @@ def _snapshot_planning_telemetry(
             "subforward_request_indices": partition,
             "predicted_peak_bytes": check.estimated_required_bytes,
             "usable_limit_bytes": check.available_bytes,
+            "routed_share": routed_share,
         }
 
     def last_forward_telemetry(self) -> dict[str, Any]:
@@ -4799,7 +4937,9 @@ def last_forward_telemetry(self) -> dict[str, Any]:
         are the admitted plan's memory check (for a split: every retained
         graph plus the largest subforward's ephemeral share). A call refused
         with ``TrainerRankMemoryError`` is still reflected, with the binding
-        check that refused it.
+        check that refused it. ``routed_share`` is the micro-batch's worst
+        observed expert-parallel routed share over its balanced rows, when
+        every rank observed the same checkpoint, and otherwise None.
         """
 
         if self._last_forward_telemetry_snapshot is None:
@@ -6754,6 +6894,12 @@ def _fill_planner_snapshot(
                                 self._plan_hybridep_growth_bytes(child)
                             ),
                         },
+                        # Each group's observed share behind its routed rows;
+                        # None prices the cold allowance.
+                        "routed_share_max": [
+                            self._observed_routed_share(group.slot_ref)
+                            for group in child.groups
+                        ],
                         "expected_required_bytes": cost.required,
                         "retained_bytes": cost.retained,
                         "cost_components": asdict(cost),
diff --git a/tests/unit/test_trainer_rank_routed_share.py b/tests/unit/test_trainer_rank_routed_share.py
new file mode 100644
index 000000000..5ca407f2c
--- /dev/null
+++ b/tests/unit/test_trainer_rank_routed_share.py
@@ -0,0 +1,303 @@
+"""Observed EP routed shares: checkpoint epochs, handoff observation, pricing."""
+
+from contextlib import nullcontext
+from dataclasses import replace
+from datetime import timedelta
+from importlib.util import find_spec
+import math
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from art.trainer_rank import ForwardInput, MaterializedCheckpoint
+from art.trainer_rank._impl import _CheckpointSlot, _MemoryCheck
+from art.trainer_rank._planner_cost import ParallelShape
+from tests.unit.test_trainer_rank_moe_memory import _rank
+
+# Real-data rank 0 of a CP2/EP2 call: its busiest rows and the CP group's
+# total, so a balanced EP rank receives 48,397 rows' pairs.
+ROWS, TOTAL, BALANCED, TOPK = 52480, 96794, 48397, 8
+
+
+def _gated(rows=ROWS):
+    rank = _rank()
+    rank._parallel_shape = ParallelShape(tp=1, cp=2, ep=2)
+    rank._ep_group_is_cp_group = True
+    rank._moe_memory_supported = True
+    plan = rank._plan_flat_forward(
+        [ForwardInput(input_tokens=torch.arange(64), target_tokens=torch.arange(64))]
+    )
+    plan = replace(
+        plan,
+        signature=replace(plan.signature, topology=(1, 1, 2, 1)),
+        groups=(replace(plan.groups[0], slot_ref=rank._slot_ref("policy")),),
+    )
+    rank._topology = lambda: SimpleNamespace(tp=1, cp=2)
+    rank._plan_group_rows = lambda plan: tuple(
+        (rows, group.grad_enabled) for group in plan.groups
+    )
+    rank._cp_group_model_tokens = lambda batch, topology: TOTAL
+    rank._checkpoint_slots["policy"] = _CheckpointSlot()
+    rank._commit_route_epoch("policy")
+    return rank, plan
+
+
+class _Layer(torch.nn.Module):
+    def __init__(self, manager):
+        super().__init__()
+        self.token_dispatcher = SimpleNamespace(_comm_manager=manager)
+
+
+def _managers(rank, *layers, topk=TOPK):
+    """HybridEP managers holding each layer's last received rows per expert."""
+    dispatcher = pytest.importorskip("megatron.core.transformer.moe.token_dispatcher")
+    model = torch.nn.Module()
+    for index, counts in enumerate(layers):
+        manager = object.__new__(dispatcher._HybridEPManager)
+        manager.config = SimpleNamespace(moe_router_topk=topk)
+        if counts is not None:
+            manager.tokens_per_expert = torch.tensor(counts, dtype=torch.int64)
+        model.add_module(f"layer{index}", _Layer(manager))
+    rank.runtime.model = [model]
+    rank._moe_layers = len(layers)
+
+
+def _exchange(rank, plan, *peers, error=None, sent=None):
+    """Run the handoff against simulated peers' payloads; return what was sent."""
+    sent = [] if sent is None else sent
+
+    def reduce(values, *, op, sync_across_dp):
+        assert op == "MAX" and sync_across_dp
+        sent.append(list(values))
+        return [max(column) for column in zip(values, *peers, strict=True)]
+
+    rank._recovery_reduce = reduce
+    rank._cache_recovery_episode = lambda error=None: nullcontext((None, None))
+    rank._try_cache_recovery = lambda *args, **kwargs: None
+    rank._release_cached_memory_for_backward(plan, error=error)
+    return sent
+
+
+def test_observed_share_only_ever_raises_routed_rows():
+    rank, plan = _gated()
+    assert rank._plan_group_routed_rows(plan) == (BALANCED,)
+    # Within the cold 1.4 allowance, margin included.
+    rank._record_routed_share(0, 1.25)
+    assert rank._plan_group_routed_rows(plan) == (BALANCED,)
+    rank._record_routed_share(0, 1.45)
+    raised = math.ceil(BALANCED * 1.55 / 1.4)
+    assert rank._plan_group_routed_rows(plan) == (raised,)
+    # A calmer later wave keeps the high-water.
+    rank._record_routed_share(0, 1.0)
+    assert rank._routed_share_max == {0: 1.45}
+    assert rank._plan_group_routed_rows(plan) == (raised,)
+    unnamed = replace(plan.groups[0], slot_ref=rank._slot_ref(None))
+    assert rank._plan_group_routed_rows(replace(plan, groups=(unnamed,))) == (BALANCED,)
+    both = replace(plan, groups=(plan.groups[0], unnamed))
+    assert rank._plan_group_routed_rows(both) == (raised, BALANCED)
+    # New content under the same name is a new epoch and forgets the old share.
+    previous = rank._checkpoint_slots["policy"]
+    rank._checkpoint_slots["policy"] = _CheckpointSlot()
+    rank._commit_route_epoch("policy", previous)
+    assert rank._checkpoint_slots["policy"].route_epoch == 1
+    assert rank._routed_share_max == {}
+    assert rank._plan_group_routed_rows(plan) == (BALANCED,)
+
+
+def test_share_is_the_worst_layer_over_priced_balanced_pairs():
+    rank, plan = _gated()
+    _managers(rank, [200_000, 187_176], [250_000, 253_329])
+    share, epoch = rank._local_routed_share(plan)
+    assert epoch == 0
+    assert share == 503_329 / (TOPK * BALANCED)
+
+
+@pytest.mark.parametrize(
+    "change",
+    [
+        "groups",
+        "no_grad",
+        "empty",
+        "cp1",
+        "unnamed",
+        "unloaded",
+        "ungated",
+        "unpriced",
+        "layers",
+    ],
+)
+def test_share_needs_one_observed_gradient_group(change):
+    rank, plan = _gated()
+    _managers(rank, [1, 2], [3, 4])
+    group = plan.groups[0]
+    if change == "groups":
+        plan = replace(plan, groups=(group, group))
+    elif change == "no_grad":
+        plan = replace(plan, groups=(replace(group, grad_enabled=False),))
+    elif change == "empty":
+        empty = SimpleNamespace(tokens=torch.empty(0, dtype=torch.long))
+        plan = replace(plan, groups=(replace(group, packed=empty),))
+    elif change == "cp1":
+        plan = replace(plan, signature=replace(plan.signature, topology=(1, 1, 1, 1)))
+    elif change == "unnamed":
+        plan = replace(plan, groups=(replace(group, slot_ref=rank._slot_ref(None)),))
+    elif change == "unloaded":
+        rank._checkpoint_slots["policy"] = _CheckpointSlot()
+    elif change == "ungated":
+        rank._ep_group_is_cp_group = False
+    elif change == "unpriced":
+        rank._moe_memory_supported = False
+    else:
+        rank._moe_layers = 3
+    assert rank._local_routed_share(plan) == (-1.0, -1)
+
+
+def test_every_rank_must_observe_the_same_epoch():
+    rank, plan = _gated()
+    _managers(rank, [250_000, 253_329])
+    local = 503_329 / (TOPK * BALANCED)
+    # A peer on the same checkpoint epoch with a busier layer.
+    sent = _exchange(rank, plan, [0.0, 1.0, 1.5, 0.0, -0.0])
+    assert sent == [[0.0, 1.0, local, 0.0, -0.0]]
+    assert rank._routed_share_max == {0: 1.5}
+    # An empty or unsupported peer, or one on another epoch, records nothing.
+    for peer in ([0.0, 0.0, -1.0, -1.0, 1.0], [0.0, 1.0, 2.0, 3.0, -3.0]):
+        _exchange(rank, plan, peer)
+        assert rank._routed_share_max == {0: 1.5}
+    # So does a wave that failed anywhere.
+    with pytest.raises(RuntimeError, match="another rank"):
+        _exchange(rank, plan, [1.0, 1.0, 2.0, 0.0, -0.0])
+    error = RuntimeError("local forward")
+    with pytest.raises(RuntimeError, match="local forward"):
+        _exchange(rank, plan, [0.0, 1.0, 2.0, 0.0, -0.0], error=error)
+    assert rank._routed_share_max == {0: 1.5}
+
+
+def test_a_failed_read_still_joins_the_exchange():
+    rank, plan = _gated()
+    _managers(rank, None)  # No dispatch yet: nothing to read.
+    sent = _exchange(rank, plan, [0.0, 1.0, 1.5, 0.0, -0.0])
+    assert sent == [[0.0, 1.0, -1.0, -1.0, 1.0]]
+    assert rank._routed_share_max == {}
+    # The local forward error skips the read entirely.
+    _managers(rank, [250_000, 253_329])
+    sent = []
+    with pytest.raises(RuntimeError, match="local"):
+        _exchange(
+            rank,
+            plan,
+            [0.0, 1.0, 1.5, 0.0, -0.0],
+            error=RuntimeError("local"),
+            sent=sent,
+        )
+    assert sent == [[1.0, 1.0, -1.0, -1.0, 1.0]]
+    assert rank._routed_share_max == {}
+
+
+def test_ep1_keeps_the_two_value_exchange():
+    rank, plan = _gated()
+    rank._parallel_shape = ParallelShape(tp=1, cp=2, ep=1)
+    assert _exchange(rank, plan, [0.0, 1.0]) == [[0.0, 1.0]]
+
+
+def test_recorded_share_reaches_telemetry_and_planner_reports():
+    rank, plan = _gated()
+    _managers(rank, [250_000, 253_329])
+    rank._planner_observation = {"replay": lambda: {"arguments": {}}}
+    _exchange(rank, plan, [0.0, 1.0, 1.5, 0.0, -0.0])
+    assert rank._planner_observation["replay"]() == {
+        "arguments": {},
+        "routed_share": 1.5,
+    }
+    check = _MemoryCheck(1, 2, True)
+    rank._snapshot_planning_telemetry(plan, check)
+    assert rank.last_forward_telemetry()["routed_share"] == 1.5
+    # A later forward without a handoff reports none.
+    rank._snapshot_planning_telemetry(plan, check)
+    assert rank.last_forward_telemetry()["routed_share"] is None
+
+
+def test_planner_snapshot_freezes_the_share_behind_routed_rows():
+    rank, plan = _gated()
+    rank._record_routed_share(0, 1.45)
+    observation = {}
+    rank._fill_planner_snapshot(plan, _MemoryCheck(1, 2, True), observation)
+    (estimate,) = observation["replay"]()["memory_replay"]["estimates"]
+    assert estimate["routed_share_max"] == [1.45]
+    # Replay keeps pricing the frozen, raised rows through unchanged arguments.
+    assert estimate["arguments"]["group_routed_rows"] == (
+        math.ceil(BALANCED * 1.55 / 1.4),
+    )
+    assert "routed_share_max" not in estimate["arguments"]
+
+
+@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron")
+def test_route_epochs_follow_agreed_commits(tmp_path: Path):
+    from art.trainer_rank import _checkpoint
+    from tests.unit.test_trainer_rank_custom_tensors import _real_lora_trainer
+
+    trainer, _ = _real_lora_trainer()
+    saved = tmp_path / "saved"
+    trainer.save_checkpoint(str(saved), "student")
+    trainer.load_checkpoint(MaterializedCheckpoint("policy", str(saved)))
+    assert trainer._checkpoint_slots["policy"].route_epoch == 0
+    trainer._record_routed_share(0, 1.5)
+    assert trainer.snapshot_checkpoint("policy", "policy:step0")
+    assert trainer._checkpoint_slots["policy:step0"].route_epoch == 1
+    trainer._record_routed_share(1, 1.6)
+
+    def fail(*args):
+        raise RuntimeError("commit")
+
+    # A failed reload rolls back to the committed content and its epoch.
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(_checkpoint, "_commit_slot", fail)
+        with pytest.raises(RuntimeError, match="commit"):
+            trainer.load_checkpoint(MaterializedCheckpoint("policy", str(saved)))
+    assert trainer._checkpoint_slots["policy"].route_epoch == 0
+    assert trainer._route_epochs == 2
+    assert trainer._routed_share_max == {0: 1.5, 1: 1.6}
+    trainer.load_checkpoint(MaterializedCheckpoint("policy", str(saved)))
+    assert trainer._checkpoint_slots["policy"].route_epoch == 2
+    assert trainer._routed_share_max == {1: 1.6}
+    trainer._discard_snapshot_checkpoint("policy:step0")
+    assert trainer._routed_share_max == {}
+
+
+def _exchange_worker(rank_index: int, init_method: str) -> None:
+    dist.init_process_group(
+        "gloo",
+        init_method=init_method,
+        rank=rank_index,
+        world_size=2,
+        timeout=timedelta(seconds=30),
+    )
+    try:
+        rank, plan = _gated()
+        rank._cache_recovery_episode = lambda error=None: nullcontext((None, None))
+        rank._try_cache_recovery = lambda *args, **kwargs: None
+        waves = (
+            ((1.2, 0), (1.5, 0), {0: 1.5}),
+            ((1.7, 0), (-1.0, -1), {0: 1.5}),
+            ((1.7, 0), (1.8, 1), {0: 1.5}),
+        )
+        for first, second, expected in waves:
+            observed = first if rank_index == 0 else second
+            rank._local_routed_share = lambda plan, observed=observed: observed
+            rank._release_cached_memory_for_backward(plan)
+            assert rank._routed_share_max == expected
+    finally:
+        dist.destroy_process_group()
+
+
+def test_gloo_exchange_records_only_unanimous_epochs(tmp_path: Path):
+    mp.spawn(
+        _exchange_worker,
+        args=(f"file://{tmp_path / 'store'}",),
+        nprocs=2,
+        join=True,
+    )

From 6cca90baea6067cb09757bb493a5ca0855799421 Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 04:45:01 +0000
Subject: [PATCH 2/5] Agree on checkpoint targets; observe no-grad waves; seed
 snapshots

Review follow-ups:
- Loads and snapshot discards now check that every rank targets the same
  name and route epoch, so every rank forgets and assigns the same epochs.
- Cancellation during the observation still reaches the handoff exchange,
  as a failed forward.
- Epochs at or above 2**40 are not observed, so float64 transport stays
  exact.
- No-grad waves are observed: they dispatch the same way, and observations
  only raise charges.
- A snapshot starts from its source's share, since it has the same weights.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/art/trainer_rank/_checkpoint.py          |  22 +++-
 src/art/trainer_rank/_impl.py                |  31 ++++-
 tests/unit/test_trainer_rank_routed_share.py | 128 ++++++++++++++++++-
 3 files changed, 165 insertions(+), 16 deletions(-)

diff --git a/src/art/trainer_rank/_checkpoint.py b/src/art/trainer_rank/_checkpoint.py
index bee473c9c..a77f08f91 100644
--- a/src/art/trainer_rank/_checkpoint.py
+++ b/src/art/trainer_rank/_checkpoint.py
@@ -1577,6 +1577,7 @@ def snapshot_checkpoint(trainer: TrainerRank, source: str, destination: str) ->
         destination,
         dict(source_slot.config),
         source_slot.revision,
+        source_slot.route_epoch,
         destination_slot is not None,
     )
     if any(value != identity for value in _gather(identity, group)):
@@ -1643,7 +1644,7 @@ def snapshot_checkpoint(trainer: TrainerRank, source: str, destination: str) ->
         _restore_slots(model_snapshot)
         trainer._checkpoint_slots.pop(destination, None)
         raise
-    trainer._commit_route_epoch(destination)
+    trainer._commit_route_epoch(destination, source=source_slot)
     trainer._snapshot_checkpoint_names.add(destination)
     return True
 
@@ -1656,7 +1657,13 @@ def discard_snapshot_checkpoint(trainer: TrainerRank, checkpoint: str) -> None:
         trainer._default_slot_ref is not None
         and trainer._default_slot_ref.name == checkpoint
     ) or any(ref.name == checkpoint for ref in trainer._slot_stack)
-    state = (slot is not None, False if slot is None else slot.snapshot, active)
+    state = (
+        checkpoint,
+        slot is not None,
+        False if slot is None else slot.snapshot,
+        None if slot is None else slot.route_epoch,
+        active,
+    )
     if any(value != state for value in _gather(state, group)):
         raise trainer._slot_state_error(
             "Checkpoint snapshot state differs across ranks"
@@ -1848,10 +1855,19 @@ def load_checkpoint(
     forward_only: bool = False,
 ) -> None:
     group = _ensure_group(trainer)
-    if any(value != source.digest for value in _gather(source.digest, group)):
+    # Every rank must replace the same name, holding the same content, so each
+    # gives the new content the same route epoch and forgets the same one.
+    current = trainer._checkpoint_slots.get(name)
+    target = (source.digest, name, None if current is None else current.route_epoch)
+    targets = _gather(target, group)
+    if any(digest != source.digest for digest, _, _ in targets):
         raise trainer._slot_state_error(
             f"Checkpoint {name!r} content differs across ranks"
         )
+    if any(value != target for value in targets):
+        raise trainer._slot_state_error(
+            f"Checkpoint {name!r} load target differs across ranks"
+        )
     config = _phase(
         lambda: trainer._validate_checkpoint_adapter_config(
             name, source.config, alpha=None
diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index a205ade62..d735db4da 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -1560,6 +1560,9 @@ def _expert_lora_weight_storage(
 # the allowance. For the EP2 policy above, call-to-call change was 0.013 and
 # batch resampling reached about 0.04 above the median.
 _ROUTED_SHARE_MARGIN = 0.10
+# Route epochs travel as float64 in the handoff exchange; stop observing well
+# before two could round to the same value.
+_ROUTE_EPOCH_LIMIT = 2**40
 
 
 def _ep_routed_row_allowance(ep: int) -> float:
@@ -3111,6 +3114,9 @@ def _release_cached_memory_for_backward(
                     share, epoch = self._local_routed_share(plan)
                 except Exception:
                     share, epoch = -1.0, -1
+                except BaseException as exc:
+                    # Cancellation still exchanges, as a failed forward.
+                    error = exc
             try:
                 failed, gradients, *observed = self._recovery_reduce(
                     [
@@ -3150,8 +3156,8 @@ def _local_routed_share(self, plan: _AnyForwardPlan) -> tuple[float, int]:
         """This rank's worst-layer routed share and checkpoint epoch, or -1s.
 
         HybridEP keeps each MoE layer's received rows per local expert after
-        combine, so at the handoff of a flat plan with one gradient group they
-        are this forward's dispatch. The share is the most loaded layer's
+        combine, so at the handoff of a flat plan with one group they are this
+        forward's dispatch, with or without gradients. The share is the most loaded layer's
         received pairs over top-k times the balanced rows that pricing scales
         (``_plan_group_balanced_rows``), so a share above the allowance means
         more rows arrived than were priced.
@@ -3159,7 +3165,6 @@ def _local_routed_share(self, plan: _AnyForwardPlan) -> tuple[float, int]:
         if (
             not isinstance(plan, _FlatForwardPlan)
             or len(plan.groups) != 1
-            or not plan.groups[0].grad_enabled
             or plan.groups[0].packed.tokens.numel() == 0
             or plan.signature.topology[2] <= 1
             or not getattr(self, "_ep_group_is_cp_group", False)
@@ -3168,7 +3173,7 @@ def _local_routed_share(self, plan: _AnyForwardPlan) -> tuple[float, int]:
             return -1.0, -1
         epoch = self._route_epoch(plan.groups[0].slot_ref)
         (balanced,) = self._plan_group_balanced_rows(plan)
-        if epoch is None or balanced <= 0:
+        if epoch is None or epoch >= _ROUTE_EPOCH_LIMIT or balanced <= 0:
             return -1.0, -1
         from megatron.core.transformer.moe.token_dispatcher import _HybridEPManager
 
@@ -4196,7 +4201,11 @@ def _observed_routed_share(self, ref: "LoRASlotRef | None") -> float | None:
         return getattr(self, "_routed_share_max", {}).get(epoch)
 
     def _commit_route_epoch(
-        self, name: str, previous: _CheckpointSlot | None = None
+        self,
+        name: str,
+        previous: _CheckpointSlot | None = None,
+        *,
+        source: _CheckpointSlot | None = None,
     ) -> None:
         """Key a checkpoint's routing observations after every rank committed it.
 
@@ -4204,9 +4213,19 @@ def _commit_route_epoch(
         everywhere, so each rank gives the same content the same new epoch.
         Epochs are never reused. Loading over a name forgets the old content's
         observations; optimizer steps keep the epoch, so its share only grows.
+        A snapshot starts from its ``source``'s share: same weights, same
+        routing.
         """
-        self._checkpoint_slots[name].route_epoch = self._route_epochs
+        epoch = self._route_epochs
+        self._checkpoint_slots[name].route_epoch = epoch
         self._route_epochs += 1
+        inherited = (
+            None
+            if source is None or source.route_epoch is None
+            else self._routed_share_max.get(source.route_epoch)
+        )
+        if inherited is not None:
+            self._routed_share_max[epoch] = inherited
         if previous is not None:
             self._forget_route_epoch(previous)
 
diff --git a/tests/unit/test_trainer_rank_routed_share.py b/tests/unit/test_trainer_rank_routed_share.py
index 5ca407f2c..4a839cf7e 100644
--- a/tests/unit/test_trainer_rank_routed_share.py
+++ b/tests/unit/test_trainer_rank_routed_share.py
@@ -13,8 +13,8 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from art.trainer_rank import ForwardInput, MaterializedCheckpoint
-from art.trainer_rank._impl import _CheckpointSlot, _MemoryCheck
+from art.trainer_rank import ForwardInput, MaterializedCheckpoint, _impl
+from art.trainer_rank._impl import _CheckpointSlot, _MemoryCheck, _SplitForwardPlan
 from art.trainer_rank._planner_cost import ParallelShape
 from tests.unit.test_trainer_rank_moe_memory import _rank
 
@@ -114,13 +114,16 @@ def test_share_is_the_worst_layer_over_priced_balanced_pairs():
     share, epoch = rank._local_routed_share(plan)
     assert epoch == 0
     assert share == 503_329 / (TOPK * BALANCED)
+    # No-grad waves dispatch the same way, so they are observed too.
+    no_grad = replace(plan.groups[0], grad_enabled=False)
+    assert rank._local_routed_share(replace(plan, groups=(no_grad,))) == (share, 0)
 
 
 @pytest.mark.parametrize(
     "change",
     [
         "groups",
-        "no_grad",
+        "split",
         "empty",
         "cp1",
         "unnamed",
@@ -130,14 +133,14 @@ def test_share_is_the_worst_layer_over_priced_balanced_pairs():
         "layers",
     ],
 )
-def test_share_needs_one_observed_gradient_group(change):
+def test_share_needs_one_observed_group(change):
     rank, plan = _gated()
     _managers(rank, [1, 2], [3, 4])
     group = plan.groups[0]
     if change == "groups":
         plan = replace(plan, groups=(group, group))
-    elif change == "no_grad":
-        plan = replace(plan, groups=(replace(group, grad_enabled=False),))
+    elif change == "split":
+        plan = _SplitForwardPlan((plan,), ((0,),), plan.request_count)
     elif change == "empty":
         empty = SimpleNamespace(tokens=torch.empty(0, dtype=torch.long))
         plan = replace(plan, groups=(replace(group, packed=empty),))
@@ -198,6 +201,60 @@ def test_a_failed_read_still_joins_the_exchange():
     assert rank._routed_share_max == {}
 
 
+def test_cancellation_during_observation_still_exchanges():
+    rank, plan = _gated()
+
+    def cancel(plan):
+        raise KeyboardInterrupt
+
+    rank._local_routed_share = cancel
+    sent = []
+    with pytest.raises(KeyboardInterrupt):
+        _exchange(rank, plan, [0.0, 1.0, 1.5, 0.0, -0.0], sent=sent)
+    # Peers see a failed forward, not a missing collective.
+    assert sent == [[1.0, 1.0, -1.0, -1.0, 1.0]]
+    assert rank._routed_share_max == {}
+
+
+def test_epochs_beyond_exact_float_transport_are_not_observed():
+    rank, plan = _gated()
+    _managers(rank, [250_000, 253_329])
+    rank._checkpoint_slots["policy"].route_epoch = 2**40
+    assert rank._local_routed_share(plan) == (-1.0, -1)
+
+
+def test_raised_rows_reach_admission_beyond_local_rows(monkeypatch):
+    rank, plan = _gated()
+    coefficient, shared = 64 * 1024, 1024
+
+    def priced(*args, shared_bytes=None, **kwargs):
+        # A named slot reprices its MoE layers: one supported layer.
+        if shared_bytes is not None:
+            shared_bytes.append(shared)
+        return coefficient
+
+    monkeypatch.setattr(_impl, "_moe_output_bytes_per_token", priced)
+    rank._moe_layers = 1
+    rank._memory_check_required = lambda required, **kwargs: required
+    before = rank._memory_check(plan)
+    cost_before = rank._plan_cost(plan).required
+    rank._record_routed_share(0, 1.6)
+    (raised,) = rank._plan_group_routed_rows(plan)
+    # More rows arrive than this rank holds: the shared part moves onto them.
+    assert raised == math.ceil(BALANCED * 1.7 / 1.4) > ROWS
+    ref = plan.groups[0].slot_ref
+    assert (
+        rank._moe_workspace_bytes(ROWS, routed_rows=raised, slot_ref=ref)
+        == raised * coefficient
+    )
+    assert rank._moe_workspace_bytes(ROWS, routed_rows=BALANCED, slot_ref=ref) == (
+        (ROWS - BALANCED) * shared + BALANCED * coefficient
+    )
+    after = rank._memory_check(plan)
+    assert after > before and after >= raised * coefficient
+    assert rank._plan_cost(plan).required > cost_before
+
+
 def test_ep1_keeps_the_two_value_exchange():
     rank, plan = _gated()
     rank._parallel_shape = ParallelShape(tp=1, cp=2, ep=1)
@@ -219,6 +276,11 @@ def test_recorded_share_reaches_telemetry_and_planner_reports():
     # A later forward without a handoff reports none.
     rank._snapshot_planning_telemetry(plan, check)
     assert rank.last_forward_telemetry()["routed_share"] is None
+    # Nor does a new public call after a wave that never reached its snapshot.
+    rank._last_routed_share = 1.5
+    rank._reset_planning_telemetry()
+    rank._snapshot_planning_telemetry(plan, check)
+    assert rank.last_forward_telemetry()["routed_share"] is None
 
 
 def test_planner_snapshot_freezes_the_share_behind_routed_rows():
@@ -246,28 +308,80 @@ def test_route_epochs_follow_agreed_commits(tmp_path: Path):
     trainer.load_checkpoint(MaterializedCheckpoint("policy", str(saved)))
     assert trainer._checkpoint_slots["policy"].route_epoch == 0
     trainer._record_routed_share(0, 1.5)
+    # A snapshot has its source's weights, so it starts from its share.
     assert trainer.snapshot_checkpoint("policy", "policy:step0")
     assert trainer._checkpoint_slots["policy:step0"].route_epoch == 1
+    assert trainer._routed_share_max == {0: 1.5, 1: 1.5}
     trainer._record_routed_share(1, 1.6)
 
     def fail(*args):
         raise RuntimeError("commit")
 
-    # A failed reload rolls back to the committed content and its epoch.
+    # A failed reload rolls back to the committed content and its epoch; a
+    # failed first load leaves no slot. Neither advances the counter.
     with pytest.MonkeyPatch.context() as patch:
         patch.setattr(_checkpoint, "_commit_slot", fail)
         with pytest.raises(RuntimeError, match="commit"):
             trainer.load_checkpoint(MaterializedCheckpoint("policy", str(saved)))
+        with pytest.raises(RuntimeError, match="commit"):
+            trainer.load_checkpoint(MaterializedCheckpoint("fresh", str(saved)))
     assert trainer._checkpoint_slots["policy"].route_epoch == 0
+    assert "fresh" not in trainer._checkpoint_slots
     assert trainer._route_epochs == 2
     assert trainer._routed_share_max == {0: 1.5, 1: 1.6}
     trainer.load_checkpoint(MaterializedCheckpoint("policy", str(saved)))
     assert trainer._checkpoint_slots["policy"].route_epoch == 2
     assert trainer._routed_share_max == {1: 1.6}
+    # Prepared snapshots load from disk: a new epoch with nothing inherited.
+    prepared = _checkpoint.prepare_checkpoint(str(saved))
+    assert _checkpoint.snapshot_prepared_checkpoint(trainer, prepared, "frozen")
+    assert trainer._checkpoint_slots["frozen"].route_epoch == 3
     trainer._discard_snapshot_checkpoint("policy:step0")
     assert trainer._routed_share_max == {}
 
 
+@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron")
+@pytest.mark.parametrize("peer", ["name", "epoch"])
+def test_loads_and_discards_must_agree_on_their_target(tmp_path: Path, peer):
+    from art.trainer_rank import TrainerRankSlotStateError, _checkpoint
+    from tests.unit.test_trainer_rank_custom_tensors import _real_lora_trainer
+
+    trainer, _ = _real_lora_trainer()
+    saved = tmp_path / "saved"
+    trainer.save_checkpoint(str(saved), "student")
+    trainer.load_checkpoint(MaterializedCheckpoint("policy", str(saved)))
+    assert trainer.snapshot_checkpoint("policy", "policy:step0")
+    trainer._record_routed_share(0, 1.5)
+    trainer._record_routed_share(1, 1.6)
+    gather = _checkpoint._gather
+
+    def disagree(value, group=None):
+        # One simulated peer targets another name, or holds another epoch.
+        values = gather(value, group)
+        if isinstance(value, tuple) and len(value) in (3, 5):
+            other = list(value)
+            if peer == "name":
+                other[1 if len(value) == 3 else 0] = "other"
+            else:
+                index = 2 if len(value) == 3 else 3
+                other[index] = 7
+            values = (*values, tuple(other))
+        return values
+
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(_checkpoint, "_gather", disagree)
+        with pytest.raises(TrainerRankSlotStateError, match="load target differs"):
+            _checkpoint.load_checkpoint(
+                trainer, _checkpoint.prepare_checkpoint(str(saved)), "policy"
+            )
+        with pytest.raises(TrainerRankSlotStateError, match="state differs"):
+            trainer._discard_snapshot_checkpoint("policy:step0")
+    assert trainer._checkpoint_slots["policy"].route_epoch == 0
+    assert trainer._checkpoint_slots["policy:step0"].route_epoch == 1
+    assert trainer._route_epochs == 2
+    assert trainer._routed_share_max == {0: 1.5, 1: 1.6}
+
+
 def _exchange_worker(rank_index: int, init_method: str) -> None:
     dist.init_process_group(
         "gloo",

From e34327d92fdfbd6969dff7f6365a6cf7f3b3c9ff Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 04:55:26 +0000
Subject: [PATCH 3/5] Also agree on the route epoch counter at loads; cover
 snapshot targets

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/art/trainer_rank/_checkpoint.py          | 14 ++++++---
 src/art/trainer_rank/_impl.py                |  8 ++---
 tests/unit/test_trainer_rank_routed_share.py | 31 +++++++++++++-------
 3 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/src/art/trainer_rank/_checkpoint.py b/src/art/trainer_rank/_checkpoint.py
index a77f08f91..035d4f90b 100644
--- a/src/art/trainer_rank/_checkpoint.py
+++ b/src/art/trainer_rank/_checkpoint.py
@@ -1855,12 +1855,18 @@ def load_checkpoint(
     forward_only: bool = False,
 ) -> None:
     group = _ensure_group(trainer)
-    # Every rank must replace the same name, holding the same content, so each
-    # gives the new content the same route epoch and forgets the same one.
+    # Every rank must replace the same name, holding the same content, at the
+    # same point in its epoch sequence, so each gives the new content the same
+    # route epoch and forgets the same one.
     current = trainer._checkpoint_slots.get(name)
-    target = (source.digest, name, None if current is None else current.route_epoch)
+    target = (
+        source.digest,
+        name,
+        None if current is None else current.route_epoch,
+        trainer._route_epochs,
+    )
     targets = _gather(target, group)
-    if any(digest != source.digest for digest, _, _ in targets):
+    if any(value[0] != source.digest for value in targets):
         raise trainer._slot_state_error(
             f"Checkpoint {name!r} content differs across ranks"
         )
diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index d735db4da..0a46552ea 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -3157,10 +3157,10 @@ def _local_routed_share(self, plan: _AnyForwardPlan) -> tuple[float, int]:
 
         HybridEP keeps each MoE layer's received rows per local expert after
         combine, so at the handoff of a flat plan with one group they are this
-        forward's dispatch, with or without gradients. The share is the most loaded layer's
-        received pairs over top-k times the balanced rows that pricing scales
-        (``_plan_group_balanced_rows``), so a share above the allowance means
-        more rows arrived than were priced.
+        forward's dispatch, with or without gradients. The share is the most
+        loaded layer's received pairs over top-k times the balanced rows that
+        pricing scales (``_plan_group_balanced_rows``), so a share above the
+        allowance means more rows arrived than were priced.
         """
         if (
             not isinstance(plan, _FlatForwardPlan)
diff --git a/tests/unit/test_trainer_rank_routed_share.py b/tests/unit/test_trainer_rank_routed_share.py
index 4a839cf7e..c07239e54 100644
--- a/tests/unit/test_trainer_rank_routed_share.py
+++ b/tests/unit/test_trainer_rank_routed_share.py
@@ -341,8 +341,8 @@ def fail(*args):
 
 
 @pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron")
-@pytest.mark.parametrize("peer", ["name", "epoch"])
-def test_loads_and_discards_must_agree_on_their_target(tmp_path: Path, peer):
+@pytest.mark.parametrize("peer", ["name", "epoch", "counter"])
+def test_loads_and_snapshots_must_agree_on_their_target(tmp_path: Path, peer):
     from art.trainer_rank import TrainerRankSlotStateError, _checkpoint
     from tests.unit.test_trainer_rank_custom_tensors import _real_lora_trainer
 
@@ -355,16 +355,21 @@ def test_loads_and_discards_must_agree_on_their_target(tmp_path: Path, peer):
     trainer._record_routed_share(1, 1.6)
     gather = _checkpoint._gather
 
+    # Gathered tuples: a load's (digest, name, epoch, counter), a snapshot's
+    # (source, destination, config, revision, source epoch, exists) and a
+    # discard's (name, loaded, snapshot, epoch, active).
+    fields = {
+        "name": {4: 1, 6: 1, 5: 0},
+        "epoch": {4: 2, 6: 4, 5: 3},
+        "counter": {4: 3},
+    }[peer]
+
     def disagree(value, group=None):
-        # One simulated peer targets another name, or holds another epoch.
+        # One simulated peer targets another name or holds another epoch.
         values = gather(value, group)
-        if isinstance(value, tuple) and len(value) in (3, 5):
+        if isinstance(value, tuple) and len(value) in fields:
             other = list(value)
-            if peer == "name":
-                other[1 if len(value) == 3 else 0] = "other"
-            else:
-                index = 2 if len(value) == 3 else 3
-                other[index] = 7
+            other[fields[len(value)]] = "other" if peer == "name" else 7
             values = (*values, tuple(other))
         return values
 
@@ -374,8 +379,12 @@ def disagree(value, group=None):
             _checkpoint.load_checkpoint(
                 trainer, _checkpoint.prepare_checkpoint(str(saved)), "policy"
             )
-        with pytest.raises(TrainerRankSlotStateError, match="state differs"):
-            trainer._discard_snapshot_checkpoint("policy:step0")
+        if peer != "counter":
+            with pytest.raises(TrainerRankSlotStateError, match="state differs"):
+                trainer.snapshot_checkpoint("policy", "policy:step1")
+            with pytest.raises(TrainerRankSlotStateError, match="state differs"):
+                trainer._discard_snapshot_checkpoint("policy:step0")
+    assert "policy:step1" not in trainer._checkpoint_slots
     assert trainer._checkpoint_slots["policy"].route_epoch == 0
     assert trainer._checkpoint_slots["policy:step0"].route_epoch == 1
     assert trainer._route_epochs == 2

From 572be044c3e08572c44f829921dceb48c195c3cf Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 05:02:36 +0000
Subject: [PATCH 4/5] Read the route epoch counter with a default for bare
 trainers

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/art/trainer_rank/_checkpoint.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/art/trainer_rank/_checkpoint.py b/src/art/trainer_rank/_checkpoint.py
index 035d4f90b..e6cd3c2d2 100644
--- a/src/art/trainer_rank/_checkpoint.py
+++ b/src/art/trainer_rank/_checkpoint.py
@@ -1863,7 +1863,7 @@ def load_checkpoint(
         source.digest,
         name,
         None if current is None else current.route_epoch,
-        trainer._route_epochs,
+        getattr(trainer, "_route_epochs", 0),
     )
     targets = _gather(target, group)
     if any(value[0] != source.digest for value in targets):

From 43db31ff6e00ffe4b356d1966981dd3f3eb16795 Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 05:22:57 +0000
Subject: [PATCH 5/5] Collect HybridEP managers into a typed list

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/art/trainer_rank/_impl.py | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index 0a46552ea..a734567e4 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -3177,17 +3177,13 @@ def _local_routed_share(self, plan: _AnyForwardPlan) -> tuple[float, int]:
             return -1.0, -1
         from megatron.core.transformer.moe.token_dispatcher import _HybridEPManager
 
-        managers = [
-            manager
-            for chunk in self.runtime.model
-            for module in chunk.modules()
-            if type(
-                manager := getattr(
-                    getattr(module, "token_dispatcher", None), "_comm_manager", None
-                )
-            )
-            is _HybridEPManager
-        ]
+        managers: list[Any] = []
+        for chunk in self.runtime.model:
+            for module in chunk.modules():
+                dispatcher = getattr(module, "token_dispatcher", None)
+                manager = getattr(dispatcher, "_comm_manager", None)
+                if type(manager) is _HybridEPManager:
+                    managers.append(manager)
         if not managers or len(managers) != self._moe_layers:
             return -1.0, -1
         received = torch.stack(