Account for known eager MoE FC2 input and output storage

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 FC2 output pair, not a bound on the complete live working set."""
+    """Known eager FC2 input/outputs, not the complete or compiled working set."""
     if shape != ParallelShape(tp=1, cp=1):
         return 0
     from megatron.core.extensions.transformer_engine import TERowParallelGroupedLinear
@@ -1274,9 +1274,23 @@ def _moe_output_bytes_per_token(
                 or getattr(layer.router, "topk", None) != config.moe_router_topk
             ):
                 return 0
+            features = 2 * fc2.out_features
+            inputs = getattr(lora, "A_T", None)
+            if (
+                isinstance(inputs, torch.Tensor)
+                and inputs.ndim == weights.ndim == 3
+                and inputs.dtype == weights.dtype
+                and inputs.shape[0] == weights.shape[0]
+                and inputs.shape[-1] == weights.shape[-2]
+                and inputs.shape[-2] > 0
+            ):
+                # Eager FC2 keeps x, base_out and adapter_out while producing
+                # 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
             coefficient = max(
                 coefficient,
-                2 * config.moe_router_topk * fc2.out_features * weights.element_size(),
+                config.moe_router_topk * features * weights.element_size(),
             )
     return coefficient
 
@@ -4811,7 +4825,7 @@ def _estimate_required_memory_bytes_from_values(
             * activation_factor
         )
         # Groups execute sequentially: summed packed rows conservatively bound
-        # this output-pair component, not simultaneous workspace or retained graphs.
+        # this FC2 component, not all workspace or retained graphs.
         static_compute = max(
             static_compute, packed_tokens * self._moe_output_bytes_per_token
         )

tests/unit/test_trainer_rank_moe_memory.py

@@ -58,6 +58,9 @@ def module(cls):
     value.experts.linear_fc2.out_features = 2048
     value.experts.linear_fc2.linear_fc2 = module(lora_module.TERowParallelGroupedLinear)
     value.experts.linear_fc2.lora = module(lora_module.LoRA)
+    value.experts.linear_fc2.lora.A_T = torch.nn.Parameter(
+        torch.empty(256, 512, 8, dtype=torch.bfloat16)
+    )
     value.experts.linear_fc2.lora.B_T = torch.nn.Parameter(
         torch.empty(256, 8, 2048, dtype=torch.bfloat16)
     )
@@ -85,24 +88,60 @@ def _signature():
 
 def test_supported_constructor_and_original_shape(layer):
     rank = _rank(layer)
-    assert rank._moe_output_bytes_per_token == 2 * 8 * 2048 * 2
+    assert rank._moe_output_bytes_per_token == (512 + 3 * 2048) * 8 * 2
     estimate = rank._estimate_required_memory_bytes_from_values(
         packed_tokens=106432,
         output_bytes=0,
         signature=_signature(),
     )
-    # This independent storage arithmetic exceeds the former 6,713,560,268.
+    # Eager FC2 arguments and outputs; compiler reuse can reduce this component.
+    inputs = torch.empty(106432 * 8, 512, dtype=torch.bfloat16, device="meta")
     base = torch.empty(106432 * 8, 2048, dtype=torch.bfloat16, device="meta")
     adapter = torch.empty_like(base)
-    assert base.untyped_storage()._cdata != adapter.untyped_storage()._cdata
+    combined = base + adapter
+    tensors = (inputs, base, adapter, combined)
+    assert len({x.untyped_storage()._cdata for x in tensors}) == len(tensors)
     pair = base.untyped_storage().nbytes() + adapter.untyped_storage().nbytes()
     assert pair == 6_975_127_552
-    assert estimate == int(pair * 1.1)
+    assert estimate == int(sum(x.untyped_storage().nbytes() for x in tensors) * 1.1)
     assert rank._memory_check_required(6_800_000_000).fits  # CPU default budget
     rank._available_memory_bytes = lambda: 6_800_000_000
     assert not rank._memory_check_required(estimate).fits
 
 
+def test_cold_fixed_pressure_budget_and_partial_envelope_limit(layer):
+    rank = _rank(layer)
+    signature = replace(_signature(), grad_enabled=False, grad_modes=(False,))
+    required = rank._estimate_required_memory_bytes_from_values(
+        packed_tokens=45981, output_bytes=707325952, signature=signature
+    )
+    assert required == 6_164_530_380
+    rank._available_memory_bytes = lambda: 4_652_784_333
+    assert rank._memory_check_required(4_092_810_444).fits
+    assert not rank._memory_check_required(required).fits
+    # This source component remains below the retained whole-forward increment.
+    # No general physical bound follows from rejecting this recorded budget.
+    assert required < 11_445_665_280
+
+
+@pytest.mark.parametrize("mode", ["missing", "dtype", "rank", "experts"])
+def test_unknown_fc2_input_keeps_previous_pair(layer, mode):
+    lora = layer.experts.linear_fc2.lora
+    if mode == "missing":
+        lora.A_T = None
+    else:
+        shape = (
+            (128, 512, 8)
+            if mode == "experts"
+            else (256, 512, 4)
+            if mode == "rank"
+            else (256, 512, 8)
+        )
+        dtype = torch.float32 if mode == "dtype" else torch.bfloat16
+        lora.A_T = torch.nn.Parameter(torch.empty(shape, dtype=dtype))
+    assert _moe_output_bytes_per_token([layer], ParallelShape(tp=1, cp=1)) == 65536
+
+
 @pytest.mark.parametrize("field", ["tp", "cp", "ep", "etp"])
 def test_sharded_path_unchanged(layer, field):
     assert (