diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index a25505c53..f3e9703ae 100644 --- a/src/art/trainer_rank/_impl.py +++ b/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() diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 70e3c9250..58da620d4 100644 --- a/src/art/trainer_rank/_memory.py +++ b/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. 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 c6e33ee9d..22ad97f42 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/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 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