fix(trainer-rank): price dense CP1 no-grad MLP memory per row

src/art/trainer_rank/_impl.py

@@ -1190,6 +1190,27 @@ def _moe_layer_count(model: torch.nn.Module) -> int:
     return sum(isinstance(module, BaseMoELayer) for module in model.modules())
 
 
+def _dense_fc1_adapted(model: torch.nn.Module) -> bool:
+    """Whether every decoder layer's MLP FC1 is ART's gated LoRA wrapper, which
+    keeps a 2F adapter output and the sum beside the 2F base output, even for
+    inactive slots."""
+
+    try:
+        from art.megatron.lora import SharedExpertsLinearFC1LoRA
+
+        layers = _language_model(model).decoder.layers
+    except (AttributeError, ImportError, RuntimeError):
+        return False
+    fc1s = [
+        getattr(getattr(layer, "mlp", None), "linear_fc1", None) for layer in layers
+    ]
+    return len(fc1s) > 0 and all(
+        type(fc1) is SharedExpertsLinearFC1LoRA
+        and getattr(fc1, "non_gated", True) is False
+        for fc1 in fc1s
+    )
+
+
 def _expert_parallel_shape(provider: object) -> tuple[int, int]:
     """(EP, ETP) of the initialized runtime, else the provider's configuration."""
 
@@ -1943,6 +1964,7 @@ def memory_field(name: str, default: Any = None) -> Any:
         }
         spec = getattr(runtime, "model_support_spec", None)
         self._moe_layers = _moe_layer_count(runtime.model[0])
+        self._dense_fc1_adapted = _dense_fc1_adapted(runtime.model[0])
         self._checkpointed_moe_layers = sum(
             getattr(module, "moe_layer_recompute", False) is True
             for module in runtime.model[0].modules()

src/art/trainer_rank/_memory.py

@@ -1078,6 +1078,32 @@ def _estimate_required_memory_bytes_from_values(
     static_compute = (
         packed_tokens * self._hidden_size * self._param_dtype_size * activation_factor
     )
+    no_grad_stage = 0
+    if (
+        not self._geometry.moe_experts
+        and self._geometry.ffn_hidden_size
+        and self._dense_fc1_adapted
+        and signature.topology[1:3] == (1, 1)
+    ):
+        # A dense no-grad layer peaks at its FC1 stage (the base GEMM output,
+        # the adapter output and their sum: 6F, or the SwiGLU live set if
+        # wider) beside the residual pair, embedding and norm output (4H), per
+        # row: Qwen3.8-27B TP1/CP1 traces at 7k-174k rows. At CP1 every packed
+        # row is local, which the per-packed-token floor above prices at only
+        # 16H. Groups run one after another, so the largest no-grad group
+        # bounds it.
+        rows = max(
+            (n for n, grad in group_rows if not grad),
+            default=0 if signature.grad_enabled else packed_tokens,
+        )
+        no_grad_stage = (
+            rows
+            * self._param_dtype_size
+            * (
+                max(6, self._mlp_activation_factor) * self._geometry.ffn_hidden_size
+                + 4 * self._hidden_size
+            )
+        )
     if signature.grad_enabled and self._recompute_granularity != "full":
         geometry = self._geometry
         hidden = self._hidden_size
@@ -1174,6 +1200,9 @@ def _estimate_required_memory_bytes_from_values(
         max(retained, checkpoint_floor[0])
         + max(workspace, head_workspace_bytes, checkpoint_floor[1])
         + (retained if include_checkpoint_input_gradient else 0),
+        # A no-grad group's forward stage, beside any gradient groups'
+        # boundaries but not the backward's input gradient.
+        max(retained, checkpoint_floor[0]) + no_grad_stage,
     )
     if signature.topology[2] > 1:
         # Local head results coexist with full CP outputs during gathering.

src/art/trainer_rank/_micro_batch_planner.py

@@ -1444,6 +1444,7 @@ def _fill_planner_snapshot(
                 "sequence_parallel",
                 "attention_output_gate",
                 "mlp_activation_factor",
+                "dense_fc1_adapted",
                 "gdn_layers",
                 "checkpointed_moe_layers",
                 "moe_output_bytes_per_token",

src/art/trainer_rank/_planner_misses.py

@@ -531,7 +531,7 @@ def report(
 _RANK_FIELDS = frozenset(
     "num_layers hidden_size param_dtype_size recompute_granularity "
     "one_layer_recompute sequence_parallel attention_output_gate "
-    "mlp_activation_factor gdn_layers "
+    "mlp_activation_factor dense_fc1_adapted gdn_layers "
     "checkpointed_moe_layers recompute_modules moe_output_bytes_per_token "
     "moe_forward_stages".split()
 )
@@ -595,11 +595,15 @@ def replay(
     state = payload["memory_replay"]
     if not state["estimates"]:
         raise ValueError("memory replay has no candidate estimates")
-    values = state["rank"]
+    # Reports from before the dense no-grad FC1 floor omit this fact; their
+    # estimator did not price that floor.
+    values: dict[str, Any] = {"dense_fc1_adapted": False, **state["rank"]}
     if set(values) != _RANK_FIELDS | {"geometry", "topology"}:
         raise ValueError(
             "incomplete replay: immutable rank fields differ (including MoE stages)"
         )
+    if type(values["dense_fc1_adapted"]) is not bool:
+        raise ValueError("incomplete replay: recorded FC1 adapter fact is invalid")
     rank = _planner_replay.ReplayRank.__new__(_planner_replay.ReplayRank)
     for name in _RANK_FIELDS - {"one_layer_recompute"}:
         setattr(rank, "_" + name, values[name])

tests/unit/test_trainer_rank_checkpoint_memory.py

@@ -497,3 +497,177 @@ def test_no_grad_enclosure_config_guard(field, value):
     r = rank()
     setattr(r.runtime.model[0].decoder.config, field, value)
     assert r._checkpoint_memory_floor(((11, False),)) == (0, 0)
+
+
+def dense_rank():
+    """Qwen3.8-27B's dense shape: H 5120, F 17408, 64 layers, fused SwiGLU."""
+    r = rank()
+    config = r.runtime.model[0].config
+    config.hidden_size, config.num_layers = 5120, 64
+    config.ffn_hidden_size, config.bias_activation_fusion = 17408, True
+    decoder = r.runtime.model[0].decoder
+    decoder.layers = torch.nn.ModuleList(
+        [torch.nn.Linear(1, 1).bfloat16() for _ in range(64)]
+    )
+    decoder.num_layers_per_pipeline_rank = 64
+    r.runtime.provider = SimpleNamespace(
+        hidden_size=5120, num_layers=64, ffn_hidden_size=17408
+    )
+    dense = TrainerRank(r.runtime)
+    assert dense._geometry.ffn_hidden_size == 17408
+    assert not dense._geometry.moe_experts and dense._mlp_activation_factor == 3
+    # These stand-in layers have no MLP; the traced model wraps every FC1.
+    assert dense._dense_fc1_adapted is False
+    dense._dense_fc1_adapted = True
+    return dense
+
+
+# Qwen3.8-27B TP1/CP1 no-grad target-only waves (random init, ART a68fa500,
+# 2026-09-29): a cold first wave peaked at 260,789 B per row, warm waves at
+# 250,181-250,507; the floor they share is (6F + 4H) x 2 = 249,856.
+_DENSE_NO_GRAD_ROW = (6 * 17408 + 4 * 5120) * 2
+_COLD_NO_GRAD_ROW = 260_789
+
+
+def no_grad_wave(r, n):
+    return r._estimate_flat_forward(
+        [
+            ForwardInput(
+                input_tokens=torch.arange(n),
+                target_tokens=torch.arange(n),
+                no_grad=True,
+            )
+        ],
+        checkpoint=None,
+    )
+
+
+def estimate(r, signature, groups, packed, logical=None):
+    return r._estimate_required_memory_bytes_from_values(
+        packed_tokens=packed,
+        output_bytes=0,
+        signature=signature,
+        logical_tokens=logical,
+        group_rows=groups,
+    )
+
+
+@pytest.mark.parametrize("n", [7268, 144228])
+def test_cp1_dense_no_grad_floor_covers_each_row(n):
+    r = dense_rank()
+    values = no_grad_wave(r, n)
+    _, _, signature, groups, _ = values
+    assert not signature.grad_enabled and signature.topology == (1, 1, 1, 1)
+    assert groups == ((n, False),)
+    # The 16H per packed token floor admitted 180 KB per row, below the cold
+    # peak; the per-row floor, itself below the warm peaks, admits it with the
+    # usual 1.1 factor.
+    assert estimate(r, signature, groups, n) == int(n * _DENSE_NO_GRAD_ROW * 1.1)
+    assert price(r, values).required >= n * _COLD_NO_GRAD_ROW
+
+
+def test_cp1_dense_no_grad_floor_bounds_profiles():
+    r = dense_rank()
+    _, _, signature, groups, _ = no_grad_wave(r, 7268)
+    floor = estimate(r, signature, groups, 7268)
+    r._memory_profiles[signature] = _MemoryProfile(100_000, 7268)
+    assert estimate(r, signature, groups, 7268) == floor
+    r._memory_profiles[signature] = _MemoryProfile(480_869, 7268)
+    assert estimate(r, signature, groups, 7268) > floor
+
+
+def test_cp1_dense_no_grad_floor_prices_physical_rows_of_the_largest_group():
+    r = dense_rank()
+    _, _, signature, _, _ = no_grad_wave(r, 64)
+    # Shared prefixes add logical rows, not physical ones: beside a low
+    # profile, which does grow with logical rows, the floor stays put.
+    r._memory_profiles[signature] = _MemoryProfile(1_000, 144228)
+    single = estimate(r, signature, ((144228, False),), 144228)
+    assert single == int(144228 * _DENSE_NO_GRAD_ROW * 1.1)
+    assert estimate(r, signature, ((144228, False),), 144228, 1_500_000) == single
+    del r._memory_profiles[signature]
+    # Groups run one after another: the largest group's rows bound the stage,
+    # and the per-packed-token floor still applies to their sum.
+    groups = ((100_000, False), (44_228, False))
+    assert estimate(r, signature, groups, 144228) == int(
+        max(144228 * 5120 * 2 * 16, 100_000 * _DENSE_NO_GRAD_ROW) * 1.1
+    )
+
+
+@pytest.mark.parametrize("case", ["cp2", "tp2", "gradient", "moe", "unwrapped_fc1"])
+def test_dense_no_grad_floor_leaves_other_pricing(case):
+    r = dense_rank()
+    _, _, signature, _, _ = no_grad_wave(r, 64)
+    groups = ((144228, False),)
+    if case == "cp2":
+        signature = replace(signature, topology=(1, 1, 2, 1))
+    elif case == "tp2":
+        signature = replace(signature, topology=(1, 2, 1, 1))
+    elif case == "gradient":
+        signature = replace(signature, grad_enabled=True, grad_modes=(True,))
+        groups = ((144228, True),)
+    elif case == "moe":
+        r._geometry = replace(r._geometry, moe_experts=8)
+    else:
+        # Attention-only adapters leave FC1 unwrapped: no adapter output or sum.
+        r._dense_fc1_adapted = False
+    priced = estimate(r, signature, groups, 144228)
+    r._geometry = replace(r._geometry, ffn_hidden_size=0)
+    assert estimate(r, signature, groups, 144228) == priced
+    if case in ("cp2", "unwrapped_fc1"):
+        # CP2 ranks hold part of the packed rows and are left to the existing
+        # floors (and #986's traced CP2 stage), as is an unwrapped FC1.
+        assert priced == int(144228 * 5120 * 2 * 16 * 1.1)
+
+
+@pytest.mark.parametrize("factor,stage", [(3, 6), (5, 6), (7, 7)])
+def test_cp1_dense_no_grad_stage_takes_the_wider_swiglu_live_set(factor, stage):
+    r = dense_rank()
+    r._mlp_activation_factor = factor
+    _, _, signature, groups, _ = no_grad_wave(r, 4096)
+    assert estimate(r, signature, groups, 4096) == int(
+        4096 * (stage * 17408 + 4 * 5120) * 2 * 1.1
+    )
+
+
+@pytest.mark.parametrize("gradient_first", [False, True])
+def test_dense_mixed_plan_keeps_the_no_grad_stage(gradient_first):
+    r = dense_rank()
+    gradient, reference = requests(1, 10_000)
+    req = [gradient, reference] if gradient_first else [reference, gradient]
+    reference_cost = r._plan_cost(r._plan_flat_forward([reference])).required
+    assert reference_cost >= int(10_000 * _DENSE_NO_GRAD_ROW * 1.1)
+    mixed = r._plan_flat_forward(req)
+    mixed_cost = r._plan_cost(mixed).required
+    # One gradient row cannot make the no-grad group's stage cheaper.
+    assert mixed_cost >= reference_cost
+    assert r._memory_check(mixed).estimated_required_bytes == mixed_cost
+    assert (
+        r._split_chunk_lower_cost(
+            req, tuple(x.input_tokens for x in req), checkpoint=Unset
+        ).required
+        == mixed_cost
+    )
+
+
+def test_dense_fc1_adapted_needs_every_layer_wrapped():
+    from art.megatron.lora import SharedExpertsLinearFC1LoRA
+    from art.trainer_rank._impl import _dense_fc1_adapted
+
+    wrapped = SharedExpertsLinearFC1LoRA.__new__(SharedExpertsLinearFC1LoRA)
+    wrapped.__dict__["non_gated"] = False
+    non_gated = SharedExpertsLinearFC1LoRA.__new__(SharedExpertsLinearFC1LoRA)
+    non_gated.__dict__["non_gated"] = True
+
+    def model(*fc1s: object) -> Any:
+        layers = [SimpleNamespace(mlp=SimpleNamespace(linear_fc1=f)) for f in fc1s]
+        return SimpleNamespace(
+            _preprocess=lambda: None, decoder=SimpleNamespace(layers=layers)
+        )
+
+    assert _dense_fc1_adapted(model(wrapped, wrapped)) is True
+    assert _dense_fc1_adapted(model(wrapped, torch.nn.Linear(1, 1))) is False
+    # A non-gated FC1 outputs F, not 2F: not the traced stage.
+    assert _dense_fc1_adapted(model(non_gated, non_gated)) is False
+    assert _dense_fc1_adapted(model()) is False
+    assert _dense_fc1_adapted(cast(Any, SimpleNamespace())) is False

tests/unit/test_trainer_rank_planner_reports.py

@@ -236,6 +236,99 @@ def test_spool_symlink_refuses(tmp_path):
     assert not list(target.iterdir())
 
 
+@pytest.mark.parametrize("adapted", [True, False, None])
+def test_replay_prices_the_recorded_dense_fc1_floor(adapted, tmp_path):
+    # A dense CP1 no-grad estimate: rows x (6F + 4H) x 2 with F 32, H 8 is
+    # 1,792 B for 4 rows, above 5H x 2 per packed token (320 B). Reports from
+    # before the floor (no recorded fact, None) replay without it.
+    rank_fields: dict[str, Any] = {
+        "num_layers": 4,
+        "hidden_size": 8,
+        "param_dtype_size": 2,
+        "recompute_granularity": "full",
+        "one_layer_recompute": True,
+        "sequence_parallel": False,
+        "attention_output_gate": False,
+        "mlp_activation_factor": 3,
+        "gdn_layers": 0,
+        "checkpointed_moe_layers": 0,
+        "recompute_modules": [],
+        "moe_output_bytes_per_token": 0,
+        "moe_forward_stages": [],
+        "geometry": {
+            "hidden_size": 8,
+            "ffn_hidden_size": 32,
+            "num_attention_heads": 2,
+            "num_query_groups": 2,
+            "kv_channels": 4,
+        },
+        "topology": [1, 1, 1, 1],
+    }
+    if adapted is not None:
+        rank_fields["dense_fc1_adapted"] = adapted
+    required = int((1792 if adapted else 320) * 1.1)
+    estimate = {
+        "signature": {
+            "topology": [1, 1, 1, 1],
+            "planner_coefficients": [2, None],
+            "slot_group_count": 1,
+            "request_mix": ["target:single"],
+            "grad_enabled": False,
+            "grad_modes": [False],
+            "slot_shapes": [],
+        },
+        "profile": None,
+        "arguments": {
+            "packed_tokens": 4,
+            "output_bytes": 0,
+            "logical_tokens": 4,
+            "gdn_segments": 0,
+            "retained_tokens": 4,
+            "group_rows": [],
+        },
+        "expected_required_bytes": required,
+        "retained_bytes": required,
+        "cost_components": {
+            "required": required,
+            "retained": required,
+            "checkpoint_retained": 0,
+            "checkpoint_workspace": 0,
+            "checkpoint_input_gradient": 0,
+            "checkpoint_peak_increment": 0,
+            "hybridep_growth": 0,
+        },
+    }
+    payload = {
+        "memory_replay": {"rank": rank_fields, "estimates": [estimate]},
+        "split_memory_floor_bytes": 0,
+        "local_admission_peak_bytes": required,
+        "reduced_admission_peak_bytes": required,
+        "safety_factor": 1.1,
+        "layouts": [],
+    }
+    path = report(
+        tmp_path,
+        predicted_peak_bytes=required,
+        observed_peak_bytes=2 * required,
+        admission_peak_bytes=required,
+        replay_factory=lambda: payload,
+    )
+    result = reports.replay(reports.validate_report(path.read_bytes()))
+    assert result["estimates"] == [
+        {"required_bytes": required, "retained_bytes": required, "matches": True}
+    ]
+    rank_fields["dense_fc1_adapted"] = 1
+    path = report(
+        tmp_path,
+        predicted_peak_bytes=required,
+        observed_peak_bytes=2 * required,
+        admission_peak_bytes=required,
+        replay_factory=lambda: payload,
+    )
+    with pytest.raises(ValueError, match="FC1 adapter fact"):
+        reports.replay(reports.validate_report(path.read_bytes()))
+
+
 def test_replay_reruns_real_memory_estimator_and_prefix_layout(tmp_path):
     from art.trainer_rank._prefix_tree_planner import (
         build_canonical_prefix_tree,
@@ -469,6 +562,7 @@ def test_actual_emitted_split_recomputes_frozen_runtime_facts(
     assert original["incomplete_reasons"] == []
     snapshot = original["replay"]
     assert snapshot["memory_replay"]["rank"]["moe_forward_stages"] == []
+    assert snapshot["memory_replay"]["rank"]["dense_fc1_adapted"] is False
     assert len(snapshot["layouts"]) == 2
     assert snapshot["local_admission_peak_bytes"] == local
     assert snapshot["reduced_admission_peak_bytes"] == local + 123