diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index a183cc117..88879ff73 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 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( diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index a81905fe1..46bab4762 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/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 + ) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index e37d9d6a7..03f3b6d3c 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/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 diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index a0ab095c1..cd0c3bdef 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/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