From 8780416ac9e877e8e8aafbe67844467c8b6190de Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Tue, 29 Sep 2026 18:24:41 +0000 Subject: [PATCH 1/3] Price CP1 dense no-grad waves per row at their FC1 stage A dense no-grad layer peaks at its FC1 stage (base GEMM output, adapter output and their sum) beside the residual pair, embedding and norm output: (6F + 4H) x 2 bytes per row on Qwen3.8-27B. The static floor charged 16H x 2 per packed token, so cold CP1 waves (every packed row local) were admitted at 180 KB per row against a 261 KB peak. Price the largest group's rows at that stage for dense no-grad waves at TP1/CP1; CP>1, TP>1, MoE and gradient pricing are unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_memory.py | 22 ++++ .../test_trainer_rank_checkpoint_memory.py | 112 ++++++++++++++++++ 2 files changed, 134 insertions(+) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 70e3c9250..bd1ea6137 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -1078,6 +1078,28 @@ def _estimate_required_memory_bytes_from_values( static_compute = ( packed_tokens * self._hidden_size * self._param_dtype_size * activation_factor ) + if ( + not signature.grad_enabled + and not self._geometry.moe_experts + and self._geometry.ffn_hidden_size + 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 does not + # cover. Groups run one after another, so the largest bounds it. + rows = max((n for n, _ in group_rows), default=packed_tokens) + static_compute = max( + static_compute, + 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 diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index c6e33ee9d..66585d052 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -497,3 +497,115 @@ 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 + 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),) + assert _DENSE_NO_GRAD_ROW <= 250_181 + # The 16H per packed token floor admitted 180 KB per row, below the cold + # peak; the per-row floor admits it with the usual 1.1 factor. + assert int(n * 5120 * 2 * 16 * 1.1) < n * _COLD_NO_GRAD_ROW + 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. + single = estimate(r, signature, ((144228, False),), 144228) + assert estimate(r, signature, ((144228, False),), 144228, 149832) == single + # 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"]) +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),) + else: + r._geometry = replace(r._geometry, moe_experts=8) + priced = estimate(r, signature, groups, 144228) + r._geometry = replace(r._geometry, ffn_hidden_size=0) + assert estimate(r, signature, groups, 144228) == priced + if case == "cp2": + # At CP2 each rank holds part of the packed rows; left to the + # existing floors (and #986's traced CP2 stage). + assert priced == int(144228 * 5120 * 2 * 16 * 1.1) From e8cef39cbb92d65e052d9268c34b96ee0ef84616 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Tue, 29 Sep 2026 18:34:19 +0000 Subject: [PATCH 2/3] Apply the dense no-grad FC1 floor only where ART wraps FC1 The traced 6F stage is the base output, the adapter output and their sum, which ART's FC1 LoRA wrapper keeps even for an inactive slot; with attention-only adapters FC1 is unwrapped and that stage is not there. Record whether every decoder layer's FC1 is wrapped, gate the floor on it, and replay it from the report (older reports, which priced without the floor, replay without it). Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 18 ++++ src/art/trainer_rank/_memory.py | 1 + src/art/trainer_rank/_micro_batch_planner.py | 1 + src/art/trainer_rank/_planner_misses.py | 8 +- .../test_trainer_rank_checkpoint_memory.py | 34 ++++++- .../unit/test_trainer_rank_planner_reports.py | 94 +++++++++++++++++++ 6 files changed, 149 insertions(+), 7 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index a25505c53..cfa205569 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1190,6 +1190,23 @@ 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 LoRA wrapper, which keeps + an adapter output and the sum beside the 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 + return len(layers) > 0 and all( + type(getattr(getattr(layer, "mlp", None), "linear_fc1", None)) + is SharedExpertsLinearFC1LoRA + for layer in layers + ) + + def _expert_parallel_shape(provider: object) -> tuple[int, int]: """(EP, ETP) of the initialized runtime, else the provider's configuration.""" @@ -1943,6 +1960,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() diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index bd1ea6137..32a9fbcce 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -1082,6 +1082,7 @@ def _estimate_required_memory_bytes_from_values( not signature.grad_enabled and 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, diff --git a/src/art/trainer_rank/_micro_batch_planner.py b/src/art/trainer_rank/_micro_batch_planner.py index 2a077b57d..0c753727b 100644 --- a/src/art/trainer_rank/_micro_batch_planner.py +++ b/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", diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index bb7e8297c..5e267e2e9 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/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]) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 66585d052..8f4fb7870 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -516,6 +516,9 @@ def dense_rank(): 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 @@ -588,7 +591,7 @@ def test_cp1_dense_no_grad_floor_prices_physical_rows_of_the_largest_group(): ) -@pytest.mark.parametrize("case", ["cp2", "tp2", "gradient", "moe"]) +@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) @@ -600,12 +603,33 @@ def test_dense_no_grad_floor_leaves_other_pricing(case): elif case == "gradient": signature = replace(signature, grad_enabled=True, grad_modes=(True,)) groups = ((144228, True),) - else: + 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 == "cp2": - # At CP2 each rank holds part of the packed rows; left to the - # existing floors (and #986's traced CP2 stage). + 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) + + +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) + + 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 + assert _dense_fc1_adapted(model()) is False + assert _dense_fc1_adapted(cast(Any, SimpleNamespace())) is False diff --git a/tests/unit/test_trainer_rank_planner_reports.py b/tests/unit/test_trainer_rank_planner_reports.py index a58adfbc1..c34a42698 100644 --- a/tests/unit/test_trainer_rank_planner_reports.py +++ b/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 From b87a8f02a7bfd25c455fb59ea9729f79ea56ffde Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Tue, 29 Sep 2026 18:45:45 +0000 Subject: [PATCH 3/3] Keep the dense no-grad stage in mixed plans; require a gated FC1 A gradient row beside a no-grad group left the no-grad group's FC1 stage unpriced, so a mixed plan priced below its no-grad part alone. Price the largest no-grad group's stage in any plan, beside the gradient groups' boundaries but not the backward's input gradient, so the exact plan, admission check and split bound agree. A non-gated FC1 outputs F, not 2F, so only the gated wrapper qualifies. Tests cover the SwiGLU live set above 6F, mixed plans through the planner, and physical rows beside a profile. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 16 ++++--- src/art/trainer_rank/_memory.py | 22 +++++---- .../test_trainer_rank_checkpoint_memory.py | 48 +++++++++++++++++-- 3 files changed, 67 insertions(+), 19 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index cfa205569..f3e9703ae 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1191,8 +1191,9 @@ def _moe_layer_count(model: torch.nn.Module) -> int: def _dense_fc1_adapted(model: torch.nn.Module) -> bool: - """Whether every decoder layer's MLP FC1 is ART's LoRA wrapper, which keeps - an adapter output and the sum beside the base output even for inactive slots.""" + """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 @@ -1200,10 +1201,13 @@ def _dense_fc1_adapted(model: torch.nn.Module) -> bool: layers = _language_model(model).decoder.layers except (AttributeError, ImportError, RuntimeError): return False - return len(layers) > 0 and all( - type(getattr(getattr(layer, "mlp", None), "linear_fc1", None)) - is SharedExpertsLinearFC1LoRA - for layer in layers + 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 ) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 32a9fbcce..58da620d4 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -1078,9 +1078,9 @@ 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 signature.grad_enabled - and not self._geometry.moe_experts + not self._geometry.moe_experts and self._geometry.ffn_hidden_size and self._dense_fc1_adapted and signature.topology[1:3] == (1, 1) @@ -1089,17 +1089,20 @@ def _estimate_required_memory_bytes_from_values( # 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 does not - # cover. Groups run one after another, so the largest bounds it. - rows = max((n for n, _ in group_rows), default=packed_tokens) - static_compute = max( - static_compute, + # 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 @@ -1197,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. diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 8f4fb7870..22ad97f42 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -559,10 +559,9 @@ def test_cp1_dense_no_grad_floor_covers_each_row(n): _, _, signature, groups, _ = values assert not signature.grad_enabled and signature.topology == (1, 1, 1, 1) assert groups == ((n, False),) - assert _DENSE_NO_GRAD_ROW <= 250_181 # The 16H per packed token floor admitted 180 KB per row, below the cold - # peak; the per-row floor admits it with the usual 1.1 factor. - assert int(n * 5120 * 2 * 16 * 1.1) < n * _COLD_NO_GRAD_ROW + # 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 @@ -580,9 +579,13 @@ def test_cp1_dense_no_grad_floor_bounds_profiles(): 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. + # 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 estimate(r, signature, ((144228, False),), 144228, 149832) == single + 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)) @@ -617,11 +620,44 @@ def test_dense_no_grad_floor_leaves_other_pricing(case): 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] @@ -631,5 +667,7 @@ def model(*fc1s: object) -> Any: 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