trainer_rank: preserve memory estimates and account for overlapping expert buffers

src/art/trainer_rank/_impl.py

@@ -1205,7 +1205,7 @@ def _configure_moe_dispatcher_caches(model: Sequence[torch.nn.Module]) -> None:
 def _moe_output_bytes_per_token(
     model: Sequence[torch.nn.Module], shape: ParallelShape
 ) -> int:
-    """Known eager FC2 input/outputs, not the complete or compiled working set."""
+    """Known routed-expert working set, not a complete model/compiled bound."""
     if shape != ParallelShape(tp=1, cp=1):
         return 0
     from megatron.core.extensions.transformer_engine import TERowParallelGroupedLinear
@@ -1216,7 +1216,7 @@ def _moe_output_bytes_per_token(
         MoEAlltoAllTokenDispatcher,
     )
 
-    from art.megatron.lora import LoRA, MLPExpertsLinearFC2LoRA
+    from art.megatron.lora import LoRA, MLPExpertsLinearFC1LoRA, MLPExpertsLinearFC2LoRA
 
     coefficient = 0
     for chunk in model:
@@ -1288,6 +1288,27 @@ def _moe_output_bytes_per_token(
                 # their sum. Compilers may reuse storage; other workspace is
                 # not covered. Unknown input metadata keeps the prior pair.
                 features = inputs.shape[-2] + 3 * fc2.out_features
+                fc1 = getattr(experts, "linear_fc1", None)
+                if (
+                    type(fc1) is MLPExpertsLinearFC1LoRA
+                    and "forward" not in vars(fc1)
+                    and not getattr(fc1, "_forward_hooks", None)
+                    and not getattr(fc1, "_forward_pre_hooks", None)
+                    and fc1.fused_gate_up
+                    and not fc1.non_gated
+                    and fc1.out_features == 2 * inputs.shape[-2]
+                    and getattr(dispatcher, "ep_size", None) == 1
+                    and getattr(dispatcher, "tp_size", None) == 1
+                    and getattr(dispatcher, "num_local_experts", 0) > 1
+                    and getattr(config, "moe_permute_fusion", False)
+                    and getattr(experts, "offload_expert_fc1", None) is False
+                    and getattr(experts, "offload_moe_act", None) is False
+                    and getattr(experts, "activation_recompute", None) is False
+                ):
+                    # The two dispatched H-wide inputs and FC1 gate/up sum
+                    # remain live at the FC2 sum, including in the observed
+                    # compiled path. This is one stage, not a backward bound.
+                    features += 2 * fc2.out_features + fc1.out_features
             coefficient = max(
                 coefficient,
                 config.moe_router_topk * features * weights.element_size(),
@@ -2649,23 +2670,6 @@ def _split_chunk_lower_cost(
             logical_tokens=logical_tokens,
         )
         profile = self._memory_profiles.get(signature)
-        if profile is not None:
-            cap = profile.packed_tokens * _MEMORY_PROFILE_TRUST_GROWTH
-            if packed_tokens <= cap < unshared_packed_tokens:
-                # Required cost can drop when a larger layout leaves the
-                # profile window. Cold cost grows with packed tokens, so its
-                # first integer count bounds every possible post-cap layout,
-                # even when TP padding makes that count itself unattainable.
-                cold = self._subforward_cost(
-                    packed_tokens=cap + 1,
-                    output_bytes=output_bytes,
-                    signature=signature,
-                    logical_tokens=logical_tokens,
-                )
-                cost = _SubforwardCost(
-                    required=min(cost.required, cold.required),
-                    retained=min(cost.retained, cold.retained),
-                )
         if (
             profile is not None
             and profile.retained_compute_bytes_per_token is not None
@@ -4839,10 +4843,10 @@ def _estimate_required_memory_bytes_from_values(
             profiled_tokens = max(
                 packed_tokens, logical_tokens / profiled.logical_per_packed
             )
-        if (
-            profiled is None
-            or profiled.packed_tokens * _MEMORY_PROFILE_TRUST_GROWTH < packed_tokens
-        ):
+        # The trust window limits calibration growth, not the empirical floor.
+        # Dropping that floor beyond the window can admit a larger request that
+        # was refused just inside it, even below a previously observed peak.
+        if profiled is None:
             compute = static_compute
         else:
             compute = max(

tests/unit/test_trainer_rank_active_memory.py

@@ -7,7 +7,13 @@
 import pytest
 import torch
 
-from art.trainer_rank import ForwardInput, ForwardOutput, TrainerRank, Unset
+from art.trainer_rank import (
+    ForwardInput,
+    ForwardOutput,
+    TrainerRank,
+    TrainerRankMemoryError,
+    Unset,
+)
 
 
 class _Model(torch.nn.Module):
@@ -199,3 +205,58 @@ def test_distinct_output_and_pure_grad_modes_require_their_own_profile():
         assert not rank._all_ranks_have_memory_profile(
             packed_tokens=plan.packed_tokens, signature=plan.signature
         )
+
+
+@pytest.mark.parametrize("no_grad", [False, True])
+def test_direct_forward_does_not_drop_observed_peak_outside_trust(monkeypatch, no_grad):
+    rank = _rank()
+
+    def request(length):
+        tokens = torch.arange(length)
+        return ForwardInput(input_tokens=tokens, target_tokens=tokens, no_grad=no_grad)
+
+    observed = rank._plan_flat_forward([request(10)])
+    rank._update_memory_profile(observed, 10_000, retained_bytes=1000)
+    monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 8000)
+
+    def unexpected_execution(*args, **kwargs):
+        raise AssertionError("An unsafe single request reached model execution")
+
+    monkeypatch.setattr(
+        rank, "_run_flat_plan_with_memory_tracking", unexpected_execution
+    )
+    for length in (80, 81):
+        candidate = rank._plan_flat_forward([request(length)])
+        assert rank._all_ranks_have_memory_profile(
+            packed_tokens=length, signature=candidate.signature
+        ) == (length == 80)
+        # Calibration trust still ends at 8x. Crossing it must not discard a
+        # peak already observed at a smaller size and admit the larger request.
+        with pytest.raises(
+            TrainerRankMemoryError, match="single request cannot be split"
+        ):
+            rank.dp_rank_forward([request(length)])
+        assert rank.last_forward_telemetry()["predicted_peak_bytes"] >= 10_000
+
+
+@pytest.mark.parametrize("logical_ratio", [1, 2, 10])
+def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio):
+    rank = _rank()
+    observed = rank._plan_flat_forward(_requests("target_tokens"))
+    rank._update_memory_profile(observed, 10_000, retained_bytes=1000)
+    estimate = rank._estimate_required_memory_bytes_from_values
+    values = [
+        estimate(
+            packed_tokens=count,
+            logical_tokens=count * logical_ratio,
+            output_bytes=count * 4,
+            signature=observed.signature,
+        )
+        for count in (8, 63, 64, 65, 800)
+    ]
+    assert values == sorted(values)
+    rate = rank._memory_profiles[observed.signature].bytes_per_token
+    assert values[-1] == int((800 * 4 + rate * 800 * logical_ratio) * 1.1)
+    assert not rank._all_ranks_have_memory_profile(
+        packed_tokens=800, signature=observed.signature
+    )

tests/unit/test_trainer_rank_moe_memory.py

@@ -233,7 +233,7 @@ def test_topk_not_expert_count_controls_envelope(layer):
     assert _moe_output_bytes_per_token([layer], shape) == old // 2
 
 
-def test_profiles_outputs_trust_and_empty_plan_unchanged():
+def test_profiles_outputs_and_empty_plan_preserve_empirical_floor():
     rank = _rank()
     signature = _signature()
     assert rank._moe_output_bytes_per_token == 0
@@ -248,12 +248,12 @@ def test_profiles_outputs_trust_and_empty_plan_unchanged():
         packed_tokens=100,
         logical_per_packed=1,
     )
-    for tokens in (100, 800):
+    for tokens in (100, 800, 801):
         assert estimate(
             packed_tokens=tokens, output_bytes=123, signature=signature
         ) == int((tokens * 100000 + 123) * 1.1)
-    assert estimate(packed_tokens=801, output_bytes=123, signature=signature) == int(
-        (801 * 65536 + 123) * 1.1
+    assert not rank._all_ranks_have_memory_profile(
+        packed_tokens=801, signature=signature
     )
     assert estimate(
         packed_tokens=100, logical_tokens=200, output_bytes=123, signature=signature
@@ -382,3 +382,77 @@ def reduce(value, op, group):
         (float(estimate), impl.dist.ReduceOp.MAX, group),
         (6_800_000_000.0, impl.dist.ReduceOp.MIN, group),
     ]
+
+
+def _enclosing_moe(layer):
+    from art.megatron.lora import MLPExpertsLinearFC1LoRA
+
+    fc1 = MLPExpertsLinearFC1LoRA.__new__(MLPExpertsLinearFC1LoRA)
+    torch.nn.Module.__init__(fc1)
+    fc1.fused_gate_up = True
+    fc1.non_gated = False
+    fc1.out_features = 1024
+    layer.experts.linear_fc1 = fc1
+    layer.experts.offload_expert_fc1 = False
+    layer.experts.offload_moe_act = False
+    layer.experts.activation_recompute = False
+    layer.token_dispatcher.ep_size = 1
+    layer.token_dispatcher.tp_size = 1
+    layer.token_dispatcher.num_local_experts = 256
+    layer.config.moe_permute_fusion = True
+    return layer
+
+
+def test_compiled_moe_retained_inputs_cold_floor(layer):
+    # Three complete intervals in a retained native trace have these seven
+    # distinct storages live together. Their sum is not the whole-model peak.
+    rank = _rank(_enclosing_moe(layer))
+    rows, hidden, intermediate = 45_981 * 8, 2048, 512
+    widths = (hidden, hidden, 2 * intermediate, intermediate, hidden, hidden, hidden)
+    storages = [
+        torch.empty(rows, width, dtype=torch.bfloat16, device="meta")
+        for width in widths
+    ]
+    assert len({tensor.untyped_storage()._cdata for tensor in storages}) == 7
+    observed_component_bytes = sum(
+        tensor.untyped_storage().nbytes() for tensor in storages
+    )
+    assert observed_component_bytes == 8_663_556_096
+    assert rank._moe_output_bytes_per_token * 45_981 == observed_component_bytes
+    for no_grad in (False, True):
+        signature = replace(
+            _signature(), grad_enabled=not no_grad, grad_modes=(not no_grad,)
+        )
+        assert rank._estimate_required_memory_bytes_from_values(
+            packed_tokens=45_981, output_bytes=0, signature=signature
+        ) == int(observed_component_bytes * 1.1)
+
+
+@pytest.mark.parametrize(
+    "site,attribute,value",
+    [
+        ("dispatcher", "ep_size", 2),
+        ("dispatcher", "tp_size", 2),
+        ("dispatcher", "num_local_experts", 1),
+        ("config", "moe_permute_fusion", False),
+        ("experts", "offload_expert_fc1", True),
+        ("experts", "offload_moe_act", True),
+        ("experts", "activation_recompute", True),
+        ("fc1", "fused_gate_up", False),
+        ("fc1", "non_gated", True),
+        ("fc1", "out_features", 2048),
+        ("fc1", "forward", lambda *args: None),
+    ],
+)
+def test_unknown_enclosing_lifetimes_keep_previous_fc2_floor(
+    layer, site, attribute, value
+):
+    _enclosing_moe(layer)
+    sites = {
+        "dispatcher": layer.token_dispatcher,
+        "config": layer.config,
+        "experts": layer.experts,
+        "fc1": layer.experts.linear_fc1,
+    }
+    setattr(sites[site], attribute, value)
+    assert _moe_output_bytes_per_token([layer], ParallelShape(tp=1, cp=1)) == 106_496

tests/unit/test_trainer_rank_split.py

@@ -604,11 +604,35 @@ def test_retained_ratio_lower_bound_reaches_exact_split_admission(
         plan, check = rank._plan_admissible_forward(
             requests, checkpoint=Unset, context="test"
         )
-        assert isinstance(plan, _SplitForwardPlan) and plan.subforward_count == 2
-        assert check == exact
+        assert isinstance(plan, _SplitForwardPlan)
+        # The smaller profile cannot discount larger layouts after its trust
+        # window; the original 20 GiB budget now requires four subforwards.
+        assert plan.subforward_count == (2 if profile_packed == 8000 else 4)
+        assert check == rank._split_rung_check(
+            [rank._plan_cost(child) for child in plan.subforwards]
+        )
         assert sorted(
             index for chunk in plan.request_indices for index in chunk
         ) == list(range(32))
+        retained = 0
+        leaf_ids = {}
+        for child, indices in zip(plan.subforwards, plan.request_indices, strict=True):
+            cost = rank._plan_cost(child)
+            assert retained + cost.required <= check.estimated_required_bytes <= budget
+            retained += cost.retained
+            assert child.request_count == len(indices)
+            leaf_ids[id(child)] = indices
+
+        def run(child, **kwargs):
+            assert kwargs["check"] == check
+            return [
+                ForwardOutput(None, None, None, None, checkpoint=str(index))
+                for index in leaf_ids[id(child)]
+            ], None
+
+        monkeypatch.setattr(rank, "_run_flat_plan_with_memory_tracking", run)
+        outputs = rank._execute_admitted_plan(plan, check=check, context="test")
+        assert [output.checkpoint for output in outputs] == list(map(str, range(32)))
     lower = rank._split_rung_check(
         [
             rank._split_chunk_lower_cost(
@@ -621,10 +645,11 @@ def test_retained_ratio_lower_bound_reaches_exact_split_admission(
     )
     assert [child.packed_tokens for child in children] == [64_000, 64_000]
     assert lower.estimated_required_bytes <= exact.estimated_required_bytes
-    assert lower.fits and exact.fits == admit
-    if not admit:
-        # The optimistic bound survives, but exact pricing must reject this
-        # rung. A later rung with smaller chunks may still fit this budget.
+    assert lower.fits == (not admit or profile_packed == 8000)
+    assert exact.fits == (admit and profile_packed == 8000)
+    if not exact.fits:
+        # Reject this rung, at the lower bound or exact pricing; a later rung
+        # with smaller chunks may still fit this same budget.
         split, rejected = rank._admit_split_rung(
             [tuple(range(16)), tuple(range(16, 32))],
             requests,
@@ -730,8 +755,8 @@ def test_packed_profile_bound_handles_unattainable_tp_boundary_count(
     assert (minimal.packed_tokens, minimal.logical_tokens) == (8, 112)
     rank._memory_profiles[minimal.signature] = _MemoryProfile(profile_rate, 4)
     lower = rank._split_chunk_lower_cost(requests, [tokens] * 16, checkpoint=Unset)
-    # The cutoff is32. A hypothetical cold count33 is cheaper than the first
-    # physically possible post-cutoff count36, so it remains a valid bound.
+    # The calibration cutoff is32, with the first physically possible larger
+    # count36. Required cost retains its empirical floor across that boundary.
     for packed in range(8, 113, 4):
         exact = rank._subforward_cost(
             packed_tokens=packed,
@@ -767,7 +792,7 @@ def test_packed_profile_bound_counts_each_groups_tp_padding(
         requests, [request.input_tokens for request in requests], checkpoint=Unset
     )
     exact = rank._plan_cost(maximum)
-    assert lower.required <= exact.required < rank._plan_cost(minimal).required
+    assert lower.required == rank._plan_cost(minimal).required < exact.required
     assert lower.retained <= exact.retained