From 75d984ce710aa2e20559dc5e02d2619540553378 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 12 Sep 2026 06:58:30 +0000 Subject: [PATCH] Price the known eager MoE FC2 input/output envelope --- src/art/trainer_rank/_impl.py | 20 +++++++-- tests/unit/test_trainer_rank_moe_memory.py | 47 ++++++++++++++++++++-- 2 files changed, 60 insertions(+), 7 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index cc62b781a..a183cc117 100644 --- a/src/art/trainer_rank/_impl.py +++ b/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 ) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 81259c716..e37d9d6a7 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/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 (