From 45544f4822ab7524d030d54d0402d77180f99bc8 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 11:45:50 +0000 Subject: [PATCH 1/5] Price TP4 sequence-parallel boundary shards in the checkpoint floor At TP>1 the full-recompute checkpoint floor declined, so a cold TP4 sequence-parallel wave was priced only by the static per-token floor. That floor covers the saved boundaries but not the recomputed layer's workspace. In two dense Qwen3.8-27B runs the first 25,728-row wave admitted 4.64 GB and peaked at 7.06 GB on every rank. A trace of that exact wave reproduces it (7.13 GB) and puts the rest of the peak in the recomputed layer's transient workspace over the gathered rows. Apply the floor at TP4 with sequence parallelism: each rank saves ceil(rows/TP) rows per boundary, and the repeated boundaries (the input-gradient term) cover that workspace. The floor is limited to dense models at CP1/PP1 with full/uniform/1 recompute that are deep enough for this cover at their widths. Other shapes keep today's pricing. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/prek.yml | 2 + src/art/trainer_rank/_impl.py | 53 ++++++- tests/unit/test_trainer_rank_tp_floor.py | 171 +++++++++++++++++++++++ 3 files changed, 223 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_trainer_rank_tp_floor.py diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 4ff185190..d7991b079 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -230,6 +230,7 @@ jobs: tests/unit/test_trainer_rank_weird_shapes.py \ tests/unit/test_trainer_rank_admission_inputs.py \ tests/unit/test_trainer_rank_checkpoint_memory.py \ + tests/unit/test_trainer_rank_tp_floor.py \ tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ tests/unit/test_trainer_rank_slot_memory.py \ tests/unit/test_trainer_rank_moe_memory.py \ @@ -274,6 +275,7 @@ jobs: --ignore=tests/unit/test_trainer_rank_weird_shapes.py \ --ignore=tests/unit/test_trainer_rank_admission_inputs.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \ + --ignore=tests/unit/test_trainer_rank_tp_floor.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ --ignore=tests/unit/test_trainer_rank_slot_memory.py \ --ignore=tests/unit/test_trainer_rank_moe_memory.py \ diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 572ad7b67..a7034ce07 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4019,6 +4019,42 @@ def _moe_workspace_bytes( else rows * coefficient ) + def _sequence_parallel_floor_covered(self, layers: int, tp: int, cp: int) -> bool: + """Whether the checkpoint floor covers a dense TP x SP recompute peak. + + Traced once: dense Qwen3.8-27B (64 layers) at TP4 with sequence + parallelism and CP1. Over the gathered rows, the recomputed layer's peak + held its SP-gathered norm input (2H per row), the MLP FC1 stage (6F/TP), + the recomputed mixer (within its projection widths / TP) and norm + outputs (under H), plus one input gradient per sharded row. The floor + repeats the sharded boundaries as the input-gradient term, so that + repeat must cover this workspace. Other TP sizes, CP, MoE and models too + shallow or wide for the bound keep today's pricing. + """ + geometry = self._geometry + if tp != 4 or cp != 1 or self._moe_layers or geometry.moe_experts: + return False + hidden = self._hidden_size + ffn = geometry.ffn_hidden_size or 4 * hidden + attention = ( + (7 if self._attention_output_gate else 5) + * geometry.num_attention_heads + * geometry.kv_channels + + 3 * geometry.num_query_groups * geometry.kv_channels + if self._num_layers > self._gdn_layers + else 0 + ) + gdn = ( + 4 * geometry.gdn_key_heads * geometry.gdn_key_head_dim + + 8 * geometry.gdn_value_heads * geometry.gdn_value_head_dim + if self._gdn_layers + else 0 + ) + # Per gathered row, times TP: the repeat is layers x H; the workspace is + # 2H + 6F/TP + mixer/TP + H, and the gradient H/TP. + workspace = 2 * hidden * tp + 6 * ffn + max(attention, gdn) + hidden * tp + return layers * hidden >= workspace + hidden + def _checkpoint_memory_floor( self, group_rows: tuple[tuple[int, bool], ...], @@ -4033,6 +4069,8 @@ def _checkpoint_memory_floor( residual and norm output across the MoE stage. Count these four row tensors separately from returned outputs, allowing storage aliases. This is not a bound for custom preprocessing, attention, or all backward. + With sequence parallelism a rank saves only its shard of each boundary; + that is priced only where ``_sequence_parallel_floor_covered`` holds. """ gradient_rows = sum(rows for rows, grad in group_rows if grad) if not group_rows or len(self.runtime.model) != 1: @@ -4052,12 +4090,13 @@ def _checkpoint_memory_floor( return 0, 0 config = decoder.config layers = len(decoder.layers) + _, tp, cp, pp = self._topology_key() expected = { "recompute_granularity": "full", "recompute_method": "uniform", "recompute_num_layers": 1, "distribute_saved_activations": False, - "sequence_parallel": False, + "sequence_parallel": tp > 1, "fp32_residual_connection": False, "cpu_offloading": False, "cuda_graph_impl": "none", @@ -4071,7 +4110,9 @@ def _checkpoint_memory_floor( or config.params_dtype is not torch.bfloat16 or self._param_dtype_size != 2 or next(self.runtime.model[0].parameters()).dtype is not torch.bfloat16 - or self._topology_key()[1::2] != (1, 1) + or pp != 1 + or tp < 1 + or (tp > 1 and not self._sequence_parallel_floor_covered(layers, tp, cp)) or any( type(getattr(config, name, None)) is not type(value) or getattr(config, name) != value @@ -4087,7 +4128,13 @@ def _checkpoint_memory_floor( or getattr(decoder, "_forward_pre_hooks", None) ): return 0, 0 - retained = gradient_rows * layers * self._hidden_size * 2 + # Physical rows are padded to a multiple of TP; each rank saves its shard. + retained = ( + sum(-(-rows // tp) for rows, grad in group_rows if grad) + * layers + * self._hidden_size + * 2 + ) if gradient_rows: self._checkpoint_moe_bytes_per_token() refs = (None,) * len(group_rows) if slot_refs is None else slot_refs diff --git a/tests/unit/test_trainer_rank_tp_floor.py b/tests/unit/test_trainer_rank_tp_floor.py new file mode 100644 index 000000000..0a6266c09 --- /dev/null +++ b/tests/unit/test_trainer_rank_tp_floor.py @@ -0,0 +1,171 @@ +"""The full-recompute checkpoint floor under TP4 sequence parallelism. + +CPU admission math, not a bound. One four-H200 trace of dense Qwen3.8-27B +(64 layers, TP4, sequence parallel) put the cold peak at each rank's boundary +shards plus a recompute workspace that the floor's repeated shards cover; +other shapes keep today's pricing. +""" + +from dataclasses import replace +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from art.trainer_rank import TrainerRank +from art.trainer_rank._impl import _MemorySignature + +H, F, LAYERS = 5120, 17408, 64 +TP4 = (1, 4, 1, 1) +# The traced 062 wave: one request, 25,727 tokens padded to 25,728 rows. +ROWS, OUTPUT = 25_728, 102_908 + + +def tp_rank(layers=LAYERS, *, ffn=F, topology=TP4, sequence_parallel=True, **config): + from megatron.core.transformer.transformer_block import TransformerBlock + + block = TransformerBlock.__new__(TransformerBlock) + torch.nn.Module.__init__(block) + block.config = SimpleNamespace( + hidden_size=H, + num_layers=layers, + padded_vocab_size=32, + params_dtype=torch.bfloat16, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + distribute_saved_activations=False, + sequence_parallel=sequence_parallel, + fp32_residual_connection=False, + cpu_offloading=False, + cuda_graph_impl="none", + fp8=None, + fp4=None, + **config, + ) + block.layers = torch.nn.ModuleList( + [torch.nn.Linear(1, 1).bfloat16() for _ in range(layers)] + ) + block.num_layers_per_pipeline_rank = layers + model: Any = torch.nn.Module() + model.config = block.config + model.decoder = block + model._preprocess = lambda: None + r = TrainerRank( + cast( + Any, + SimpleNamespace( + model=[model], + optimizer=None, + provider=SimpleNamespace(hidden_size=H, num_layers=layers), + model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), + ), + ) + ) + # Qwen3.8-27B: gated attention every fourth layer, GDN otherwise. + r._geometry = replace( + r._geometry, + hidden_size=H, + ffn_hidden_size=ffn, + num_attention_heads=24, + num_query_groups=4, + kv_channels=256, + gdn_key_heads=16, + gdn_key_head_dim=128, + gdn_value_heads=48, + gdn_value_head_dim=128, + ) + r._attention_output_gate = True + r._gdn_layers = layers * 3 // 4 + r._topology_key = lambda: topology + return r + + +def _required(r, group_rows=((ROWS, True),), topology=TP4): + signature = _MemorySignature( + topology, + (1, None), + len(group_rows), + (), + any(grad for _, grad in group_rows), + tuple(grad for _, grad in group_rows), + ) + return r._subforward_cost( + packed_tokens=sum(rows for rows, _ in group_rows), + output_bytes=OUTPUT, + signature=signature, + logical_tokens=sum(rows for rows, _ in group_rows) - 1, + group_rows=group_rows, + ) + + +def test_the_traced_tp4_wave_prices_its_boundary_shards_and_their_repeat(): + r = tp_rank() + retained, workspace = r._checkpoint_memory_floor(((ROWS, True),)) + # Each rank saves a quarter of every boundary: the traced 4.215 GB. + assert retained == ROWS // 4 * LAYERS * H * 2 == 4_215_275_520 + assert workspace == 0 + cost = _required(r) + assert cost.checkpoint_input_gradient == retained + assert cost.required == int((OUTPUT + 2 * retained) * 1.1) + # Measured cold on all four ranks: 7.130 GB (7.060 GB in production), all + # but the boundaries a transient recompute workspace; this raw floor + # (8.43 GB) covers it. Today's cold admission was 4.637 GB. + assert cost.required / 1.1 > 7.130e9 + + +def test_rows_are_sharded_with_ceiling_and_only_gradient_groups_save_them(): + r = tp_rank() + # Physical rows are padded to TP, but a stray remainder still rounds up. + assert r._checkpoint_memory_floor(((ROWS - 1, True),))[0] == ( + -(-(ROWS - 1) // 4) * LAYERS * H * 2 + ) + mixed = r._checkpoint_memory_floor(((1024, True), (4096, False))) + assert mixed[0] == 256 * LAYERS * H * 2 + # No-grad groups keep today's four full rows; the static floor covers them. + assert mixed[1] == 4 * 4096 * H * 2 + + +@pytest.mark.parametrize( + "case", + [ + "tp2", + "tp8", + "cp2", + "pp2", + "no_sequence_parallel", + "sequence_parallel_at_tp1", + "selective_recompute", + "moe", + "shallow", + "wide_ffn", + ], +) +def test_unproven_shapes_keep_todays_pricing(case): + shapes = { + "tp2": dict(topology=(1, 2, 1, 1)), + "tp8": dict(topology=(1, 8, 1, 1)), + "cp2": dict(topology=(1, 4, 2, 1)), + "pp2": dict(topology=(1, 4, 1, 2)), + "no_sequence_parallel": dict(sequence_parallel=False), + "sequence_parallel_at_tp1": dict(topology=(1, 1, 1, 1)), + "selective_recompute": dict(), + "moe": dict(), + "shallow": dict(layers=44), + "wide_ffn": dict(ffn=4 * F), + } + r = tp_rank(**shapes[case]) + if case == "selective_recompute": + r.runtime.model[0].decoder.config.recompute_granularity = "selective" + if case == "moe": + r._moe_layers = 1 + assert r._checkpoint_memory_floor(((ROWS, True),)) == (0, 0) + + +def test_the_depth_bound_is_the_traced_workspace_at_these_widths(): + # Per gathered row, the repeated shards (layers x H / 4) must cover the + # SP gather (2H), the FC1 stage (6F/4), the wider mixer (GDN here), the + # norm (H) and one gradient shard (H/4): about 45 layers at these widths. + assert tp_rank(layers=45)._checkpoint_memory_floor(((ROWS, True),))[0] > 0 + assert tp_rank(layers=44)._checkpoint_memory_floor(((ROWS, True),)) == (0, 0) From 37179a3235076c8bdcd2c91bdba5767469b5a6f6 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 12:22:21 +0000 Subject: [PATCH 2/5] Bound the TP4 floor's cover and price GDN segment states Round-1 review closures: - The recomputed GDN layer's recurrent states grow with segments, not rows. The floor now prices them from the plan's gradient segments. - Fewer KV groups than TP keep a replicated global QKV on every rank, so those configurations keep today's pricing. So do models with unreadable attention or GDN widths. - The bound uses the SwiGLU live set when it is wider than the FC1 stage, and prices other workspace (the trace's TE and residual allocations) at H per gathered row instead of relying on slack. - The test fixture's rank is typed Any so it can stub the topology. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 52 ++++++++++++++++----- tests/unit/test_trainer_rank_tp_floor.py | 59 +++++++++++++++++++++--- 2 files changed, 92 insertions(+), 19 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index a7034ce07..981542c48 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4025,23 +4025,42 @@ def _sequence_parallel_floor_covered(self, layers: int, tp: int, cp: int) -> boo Traced once: dense Qwen3.8-27B (64 layers) at TP4 with sequence parallelism and CP1. Over the gathered rows, the recomputed layer's peak held its SP-gathered norm input (2H per row), the MLP FC1 stage (6F/TP), - the recomputed mixer (within its projection widths / TP) and norm - outputs (under H), plus one input gradient per sharded row. The floor - repeats the sharded boundaries as the input-gradient term, so that - repeat must cover this workspace. Other TP sizes, CP, MoE and models too - shallow or wide for the bound keep today's pricing. + the recomputed mixer (within its projection widths / TP), norm outputs + and other workspace (each under H), plus one input gradient per + sharded row. The floor repeats the sharded boundaries as the + input-gradient term, so that repeat must cover this workspace; GDN + segment states grow with segments instead and are priced separately. + Other TP sizes, CP, MoE, replicated QKV (KV groups below TP), missing + geometry and models too shallow or wide for the bound keep today's + pricing. """ geometry = self._geometry if tp != 4 or cp != 1 or self._moe_layers or geometry.moe_experts: return False hidden = self._hidden_size ffn = geometry.ffn_hidden_size or 4 * hidden + attention_layers = self._num_layers > self._gdn_layers + if attention_layers and ( + geometry.num_attention_heads <= 0 + or geometry.kv_channels <= 0 + # Replicated QKV keeps a global QKV output on every rank. + or not tp <= geometry.num_query_groups + ): + return False + gdn_widths = ( + geometry.gdn_key_heads, + geometry.gdn_key_head_dim, + geometry.gdn_value_heads, + geometry.gdn_value_head_dim, + ) + if self._gdn_layers and min(gdn_widths) <= 0: + return False attention = ( (7 if self._attention_output_gate else 5) * geometry.num_attention_heads * geometry.kv_channels + 3 * geometry.num_query_groups * geometry.kv_channels - if self._num_layers > self._gdn_layers + if attention_layers else 0 ) gdn = ( @@ -4051,14 +4070,17 @@ def _sequence_parallel_floor_covered(self, layers: int, tp: int, cp: int) -> boo else 0 ) # Per gathered row, times TP: the repeat is layers x H; the workspace is - # 2H + 6F/TP + mixer/TP + H, and the gradient H/TP. - workspace = 2 * hidden * tp + 6 * ffn + max(attention, gdn) + hidden * tp + # 2H + the FC1 stage (6F/TP, or the SwiGLU live set if wider) + + # mixer/TP + H of norms + H of other workspace, and the gradient H/TP. + stage = max(6, self._mlp_activation_factor) * ffn + workspace = 2 * hidden * tp + stage + max(attention, gdn) + 2 * hidden * tp return layers * hidden >= workspace + hidden def _checkpoint_memory_floor( self, group_rows: tuple[tuple[int, bool], ...], slot_refs: tuple["LoRASlotRef | None", ...] | None = None, + gdn_segments: int = 0, ) -> tuple[int, int]: """Conservative saved-boundary charge and one disjoint MoE workspace. @@ -4070,7 +4092,8 @@ def _checkpoint_memory_floor( tensors separately from returned outputs, allowing storage aliases. This is not a bound for custom preprocessing, attention, or all backward. With sequence parallelism a rank saves only its shard of each boundary; - that is priced only where ``_sequence_parallel_floor_covered`` holds. + that is priced only where ``_sequence_parallel_floor_covered`` holds, + and there the recomputed GDN layer's ``gdn_segments`` recurrent states. """ gradient_rows = sum(rows for rows, grad in group_rows if grad) if not group_rows or len(self.runtime.model) != 1: @@ -4111,7 +4134,6 @@ def _checkpoint_memory_floor( or self._param_dtype_size != 2 or next(self.runtime.model[0].parameters()).dtype is not torch.bfloat16 or pp != 1 - or tp < 1 or (tp > 1 and not self._sequence_parallel_floor_covered(layers, tp, cp)) or any( type(getattr(config, name, None)) is not type(value) @@ -4143,6 +4165,10 @@ def _checkpoint_memory_floor( + (0 if grad else 4 * rows * self._hidden_size * 2) for (rows, grad), ref in zip(group_rows, refs, strict=True) ) + if tp > 1 and self._gdn_layers: + # Recurrent states grow with segments, not rows; backward recomputes + # one layer at a time. + workspace += math.ceil(gdn_segments * self._gdn_segment_layer_bytes()) return retained, workspace def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: @@ -4189,7 +4215,7 @@ def _subforward_cost( include_checkpoint_input_gradient=False, ) checkpoint_retained, checkpoint_workspace = self._checkpoint_memory_floor( - group_rows, slot_refs + group_rows, slot_refs, gdn_segments ) retained = self._retained_memory_bytes( signature, @@ -7967,7 +7993,9 @@ def _estimate_required_memory_bytes_from_values( for ref in (slot_refs or (None,)) ), ) - retained, workspace = self._checkpoint_memory_floor(group_rows, slot_refs) + retained, workspace = self._checkpoint_memory_floor( + group_rows, slot_refs, gdn_segments + ) static_compute = max( static_compute, max(retained, checkpoint_floor[0]) diff --git a/tests/unit/test_trainer_rank_tp_floor.py b/tests/unit/test_trainer_rank_tp_floor.py index 0a6266c09..d41c0db25 100644 --- a/tests/unit/test_trainer_rank_tp_floor.py +++ b/tests/unit/test_trainer_rank_tp_floor.py @@ -52,7 +52,7 @@ def tp_rank(layers=LAYERS, *, ffn=F, topology=TP4, sequence_parallel=True, **con model.config = block.config model.decoder = block model._preprocess = lambda: None - r = TrainerRank( + r: Any = TrainerRank( cast( Any, SimpleNamespace( @@ -82,7 +82,7 @@ def tp_rank(layers=LAYERS, *, ffn=F, topology=TP4, sequence_parallel=True, **con return r -def _required(r, group_rows=((ROWS, True),), topology=TP4): +def _required(r, group_rows=((ROWS, True),), topology=TP4, gdn_segments=1): signature = _MemorySignature( topology, (1, None), @@ -96,6 +96,7 @@ def _required(r, group_rows=((ROWS, True),), topology=TP4): output_bytes=OUTPUT, signature=signature, logical_tokens=sum(rows for rows, _ in group_rows) - 1, + gdn_segments=gdn_segments, group_rows=group_rows, ) @@ -108,7 +109,10 @@ def test_the_traced_tp4_wave_prices_its_boundary_shards_and_their_repeat(): assert workspace == 0 cost = _required(r) assert cost.checkpoint_input_gradient == retained - assert cost.required == int((OUTPUT + 2 * retained) * 1.1) + # One segment's recurrent states in the recomputed GDN layer. + state = cost.checkpoint_workspace + assert state == -(-r._gdn_segment_layer_bytes() // 1) < 2 * 2**20 + assert cost.required == int((OUTPUT + 2 * retained + state) * 1.1) # Measured cold on all four ranks: 7.130 GB (7.060 GB in production), all # but the boundaries a transient recompute workspace; this raw floor # (8.43 GB) covers it. Today's cold admission was 4.637 GB. @@ -138,6 +142,9 @@ def test_rows_are_sharded_with_ceiling_and_only_gradient_groups_save_them(): "sequence_parallel_at_tp1", "selective_recompute", "moe", + "moe_geometry", + "replicated_qkv", + "missing_attention_geometry", "shallow", "wide_ffn", ], @@ -152,7 +159,10 @@ def test_unproven_shapes_keep_todays_pricing(case): "sequence_parallel_at_tp1": dict(topology=(1, 1, 1, 1)), "selective_recompute": dict(), "moe": dict(), - "shallow": dict(layers=44), + "moe_geometry": dict(), + "replicated_qkv": dict(), + "missing_attention_geometry": dict(), + "shallow": dict(layers=48), "wide_ffn": dict(ffn=4 * F), } r = tp_rank(**shapes[case]) @@ -160,12 +170,47 @@ def test_unproven_shapes_keep_todays_pricing(case): r.runtime.model[0].decoder.config.recompute_granularity = "selective" if case == "moe": r._moe_layers = 1 + # Geometry-only MoE, fewer KV groups than TP (a replicated global QKV) and + # unreadable attention widths are all outside the traced shape. + edits = { + "moe_geometry": dict(moe_experts=8), + "replicated_qkv": dict(num_query_groups=2), + "missing_attention_geometry": dict(kv_channels=0), + } + if case in edits: + r._geometry = replace(r._geometry, **edits[case]) assert r._checkpoint_memory_floor(((ROWS, True),)) == (0, 0) def test_the_depth_bound_is_the_traced_workspace_at_these_widths(): # Per gathered row, the repeated shards (layers x H / 4) must cover the # SP gather (2H), the FC1 stage (6F/4), the wider mixer (GDN here), the - # norm (H) and one gradient shard (H/4): about 45 layers at these widths. - assert tp_rank(layers=45)._checkpoint_memory_floor(((ROWS, True),))[0] > 0 - assert tp_rank(layers=44)._checkpoint_memory_floor(((ROWS, True),)) == (0, 0) + # norms (H), other workspace (H) and one gradient shard (H/4): about 49 + # layers at these widths. + assert tp_rank(layers=49)._checkpoint_memory_floor(((ROWS, True),))[0] > 0 + assert tp_rank(layers=48)._checkpoint_memory_floor(((ROWS, True),)) == (0, 0) + + +def test_an_ungated_attention_only_model_is_bounded_by_its_attention_width(): + r = tp_rank() + r._gdn_layers = 0 + r._attention_output_gate = False + assert r._checkpoint_memory_floor(((ROWS, True),)) == ( + ROWS // 4 * LAYERS * H * 2, + 0, + ) + + +def test_gdn_segment_states_are_priced_with_the_segments(): + """4,096 two-token requests: states grow with segments, not rows.""" + r = tp_rank() + rows = 8192 + cost = _required(r, group_rows=((rows, True),), gdn_segments=4096) + states = -(-4096 * r._gdn_segment_layer_bytes() // 1) + assert cost.checkpoint_workspace == states + # The recomputed layer's initial states alone: 4,096 x 12 local value + # heads x 128 x 128 x fp32. + assert states > 4096 * 12 * 128 * 128 * 4 + assert cost.required >= int( + (OUTPUT + 2 * rows // 4 * LAYERS * H * 2 + states) * 1.1 + ) From 46d79b12f54d1378955811c418228596a4d9ec8a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 12:32:21 +0000 Subject: [PATCH 3/5] Let the MoE floor stub take the new segment count Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/unit/test_trainer_rank_moe_memory.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index aa30b54d8..a6b1cbf51 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -589,7 +589,9 @@ def held(rows): # Growth enters the forward and checkpoint peaks, not forward retention. rank._update_memory_profile(plan, 10**9, retained_bytes=10**8) monkeypatch.setattr( - rank, "_checkpoint_memory_floor", lambda rows, refs=None: (10**7, 10**6) + rank, + "_checkpoint_memory_floor", + lambda rows, refs=None, segments=0: (10**7, 10**6), ) grown = rank._plan_cost(plan) monkeypatch.setattr(rank, "_plan_hybridep_growth_bytes", lambda plan: 0) From 68e7eff9d243cd96d0c105e16301db9ed3dbc67e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 12:57:57 +0000 Subject: [PATCH 4/5] Count gradient and TP-padding segments for the TP4 state charge Round-2 review closures: - TP padding gives each pad token its own GDN root. The floor now adds up to TP - 1 padding roots per gradient group to the segment count. - Segment states are charged only for gradient waves, so no-grad pricing is unchanged. - Width probing counted every active request, no-grad included, and priced exact layouts with that bound. _estimate_flat_forward now reports gradient groups' segments: twice their requests in cheap mode (still an upper bound for accepting a width) and the selected layouts' actual counts in exact mode, matching the materialized plan. - Tests exercise conv history and pin the per-segment charge exactly. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 45 ++++++++++++++------ tests/unit/test_trainer_rank_tp_floor.py | 54 ++++++++++++++++++------ 2 files changed, 74 insertions(+), 25 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 981542c48..7aa568cb5 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4093,7 +4093,8 @@ def _checkpoint_memory_floor( This is not a bound for custom preprocessing, attention, or all backward. With sequence parallelism a rank saves only its shard of each boundary; that is priced only where ``_sequence_parallel_floor_covered`` holds, - and there the recomputed GDN layer's ``gdn_segments`` recurrent states. + and there, for gradient waves, the recomputed GDN layer's recurrent + states for ``gdn_segments`` (gradient groups' segments) plus padding. """ gradient_rows = sum(rows for rows, grad in group_rows if grad) if not group_rows or len(self.runtime.model) != 1: @@ -4165,10 +4166,12 @@ def _checkpoint_memory_floor( + (0 if grad else 4 * rows * self._hidden_size * 2) for (rows, grad), ref in zip(group_rows, refs, strict=True) ) - if tp > 1 and self._gdn_layers: + if tp > 1 and self._gdn_layers and gradient_rows: # Recurrent states grow with segments, not rows; backward recomputes - # one layer at a time. - workspace += math.ceil(gdn_segments * self._gdn_segment_layer_bytes()) + # one layer at a time. Padding to TP adds up to TP - 1 one-token + # roots per group. Kernel-internal chunk states are not bounded here. + roots = gdn_segments + (tp - 1) * sum(grad for _, grad in group_rows) + workspace += math.ceil(roots * self._gdn_segment_layer_bytes()) return retained, workspace def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: @@ -5243,8 +5246,12 @@ def estimate(width: int) -> tuple[_MemoryCheck, bool, bool] | None: return estimates[width] indices, local_inputs = local_slice(width) local_requests = list(_flatten(local_inputs)) + cheap_segments: list[int] = [] values = self._estimate_flat_forward( - local_requests, checkpoint=checkpoint, sync_planning_errors=True + local_requests, + checkpoint=checkpoint, + sync_planning_errors=True, + gdn_segments=cheap_segments, ) if not self._all_ranks_true(values is not None): estimates[width] = None @@ -5258,6 +5265,8 @@ def priced( signature: _MemorySignature, group_rows: tuple[tuple[int, bool], ...], head_workspace_bytes: int, + *, + gdn_segments: int, ) -> tuple[_MemoryCheck, int, int, _MemorySignature]: with self._planning_status(True): required = self._estimate_required_memory_bytes_from_values( @@ -5265,12 +5274,9 @@ def priced( output_bytes=output_bytes, signature=signature, logical_tokens=logical_tokens, - # A radix tree has fewer than twice as many segments as - # active requests; the exact plan uses its actual count. - gdn_segments=2 - * sum( - _request_mix_key(r) != "inactive" for r in local_requests - ), + # Gradient groups' segments: exact layouts' counts, else + # an upper bound (see _estimate_flat_forward). + gdn_segments=gdn_segments, group_rows=group_rows, head_workspace_bytes=head_workspace_bytes, ) @@ -5284,14 +5290,20 @@ def priced( def priced_estimate( *, exact: bool, memory_minimal: bool ) -> tuple[_MemoryCheck, int, int, _MemorySignature] | None: + segments: list[int] = [] estimated = self._estimate_flat_forward( local_requests, checkpoint=checkpoint, exact=exact, memory_minimal=memory_minimal, sync_planning_errors=True, + gdn_segments=segments, + ) + return ( + None + if estimated is None + else priced(*estimated, gdn_segments=sum(segments)) ) - return None if estimated is None else priced(*estimated) def trusted(packed_tokens: int, signature: _MemorySignature) -> bool: return self._all_ranks_have_memory_profile( @@ -5304,7 +5316,7 @@ def trusted(packed_tokens: int, signature: _MemorySignature) -> bool: # reject on memory, or when it would reject on profile trust while # a profile exists — the selected layout may be far smaller than # the bound and squarely inside the profiled regime. - selected = priced(*values) + selected = priced(*values, gdn_segments=sum(cheap_segments)) profiled = self._all_ranks_true(selected[3] in self._memory_profiles) needs_exact = not selected[0].fits or ( profiled and not trusted(selected[1], selected[3]) @@ -5977,6 +5989,7 @@ def _estimate_flat_forward( exact: bool = False, memory_minimal: bool = False, sync_planning_errors: bool = False, + gdn_segments: list[int] | None = None, ) -> tuple[int, int, _MemorySignature, tuple[tuple[int, bool], ...], int] | None: """Estimate packed tokens for width probing. @@ -5989,6 +6002,8 @@ def _estimate_flat_forward( ``exact=True`` prices the planner's actual layouts (memoized by content) and is used only inside the band where those bounds disagree. Under CP it returns None: per-rank floors need materialized layouts. + ``gdn_segments`` receives each gradient group's segment count: exact + layouts' actual counts, else twice its requests (a radix tree has fewer). """ if sync_planning_errors: @@ -6038,6 +6053,8 @@ def _estimate_flat_forward( physical_rows = self._physical_tokens(layout.packed_tokens) packed_tokens += physical_rows group_rows.append((physical_rows, grad_enabled)) + if grad_enabled and gdn_segments is not None: + gdn_segments.append(len(layout.segments)) projected = upper positions = None mixed_targets = ( @@ -6096,6 +6113,8 @@ def _estimate_flat_forward( physical_rows = self._physical_tokens(group_packed_tokens) packed_tokens += physical_rows group_rows.append((physical_rows, grad_enabled)) + if grad_enabled and gdn_segments is not None: + gdn_segments.append(2 * len(group_indices)) head_workspace_bytes = max( head_workspace_bytes, self._group_head_workspace_bytes( diff --git a/tests/unit/test_trainer_rank_tp_floor.py b/tests/unit/test_trainer_rank_tp_floor.py index d41c0db25..19ad0ec10 100644 --- a/tests/unit/test_trainer_rank_tp_floor.py +++ b/tests/unit/test_trainer_rank_tp_floor.py @@ -20,6 +20,9 @@ TP4 = (1, 4, 1, 1) # The traced 062 wave: one request, 25,727 tokens padded to 25,728 rows. ROWS, OUTPUT = 25_728, 102_908 +# One segment's initial and final fp32 states over 12 local value heads, plus +# conv history over 3 taps, in the recomputed GDN layer. +SEGMENT = (4 * 12 * 128 * 128 + 2 * (2 * 4 * 128 + 12 * 128) * 3) * 2 def tp_rank(layers=LAYERS, *, ffn=F, topology=TP4, sequence_parallel=True, **config): @@ -75,6 +78,7 @@ def tp_rank(layers=LAYERS, *, ffn=F, topology=TP4, sequence_parallel=True, **con gdn_key_head_dim=128, gdn_value_heads=48, gdn_value_head_dim=128, + gdn_conv_kernel=4, ) r._attention_output_gate = True r._gdn_layers = layers * 3 // 4 @@ -106,12 +110,13 @@ def test_the_traced_tp4_wave_prices_its_boundary_shards_and_their_repeat(): retained, workspace = r._checkpoint_memory_floor(((ROWS, True),)) # Each rank saves a quarter of every boundary: the traced 4.215 GB. assert retained == ROWS // 4 * LAYERS * H * 2 == 4_215_275_520 - assert workspace == 0 + # Without a segment count, only the TP-padding roots' states. + assert workspace == 3 * SEGMENT cost = _required(r) assert cost.checkpoint_input_gradient == retained - # One segment's recurrent states in the recomputed GDN layer. + # One segment plus up to three TP-padding roots, each with its states. state = cost.checkpoint_workspace - assert state == -(-r._gdn_segment_layer_bytes() // 1) < 2 * 2**20 + assert state == 4 * SEGMENT assert cost.required == int((OUTPUT + 2 * retained + state) * 1.1) # Measured cold on all four ranks: 7.130 GB (7.060 GB in production), all # but the boundaries a transient recompute workspace; this raw floor @@ -127,8 +132,9 @@ def test_rows_are_sharded_with_ceiling_and_only_gradient_groups_save_them(): ) mixed = r._checkpoint_memory_floor(((1024, True), (4096, False))) assert mixed[0] == 256 * LAYERS * H * 2 - # No-grad groups keep today's four full rows; the static floor covers them. - assert mixed[1] == 4 * 4096 * H * 2 + # No-grad groups keep today's four full rows (the static floor covers + # them); the gradient group adds its TP-padding roots' states. + assert mixed[1] == 4 * 4096 * H * 2 + 3 * SEGMENT @pytest.mark.parametrize( @@ -206,11 +212,35 @@ def test_gdn_segment_states_are_priced_with_the_segments(): r = tp_rank() rows = 8192 cost = _required(r, group_rows=((rows, True),), gdn_segments=4096) - states = -(-4096 * r._gdn_segment_layer_bytes() // 1) - assert cost.checkpoint_workspace == states - # The recomputed layer's initial states alone: 4,096 x 12 local value - # heads x 128 x 128 x fp32. - assert states > 4096 * 12 * 128 * 128 * 4 - assert cost.required >= int( - (OUTPUT + 2 * rows // 4 * LAYERS * H * 2 + states) * 1.1 + assert cost.checkpoint_workspace == (4096 + 3) * SEGMENT + assert cost.required == int( + (OUTPUT + 2 * rows // 4 * LAYERS * H * 2 + (4096 + 3) * SEGMENT) * 1.1 ) + + +def test_tp_padding_roots_carry_their_own_states(): + """One one-token request: padding makes four roots, each with its states.""" + r = tp_rank() + cost = _required(r, group_rows=((4, True),), gdn_segments=1) + # Four roots' initial states alone: 4 x 12 value heads x 128 x 128 x fp32. + assert cost.checkpoint_workspace == 4 * SEGMENT > 4 * 12 * 128 * 128 * 4 + assert cost.required > 4 * SEGMENT + + +def test_no_grad_waves_keep_their_pricing_whatever_the_segments(): + r = tp_rank() + no_grad = r._checkpoint_memory_floor(((8192, False),), None, 8192) + assert no_grad == (0, 4 * 8192 * H * 2) + + +def test_width_probes_count_only_gradient_segments(): + from test_trainer_rank_checkpoint_memory import rank, requests + + r = rank() + # One gradient request and one no-grad reference. + cheap: list[int] = [] + assert r._estimate_flat_forward(requests(), gdn_segments=cheap) is not None + assert cheap == [2] # A radix tree has fewer than twice its requests. + exact: list[int] = [] + assert r._estimate_flat_forward(requests(), exact=True, gdn_segments=exact) + assert exact == [1] # The selected layout's actual segments. From c827edaed1ad779a53083f13f63f0dab0253d86a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 13:19:07 +0000 Subject: [PATCH 5/5] Give the width search's rejection bound a segment lower bound The memory-minimal cheap estimate decides rejection, but it carried the no-sharing upper bound on GDN segments (twice the requests), so exact pricing never ran for many-segment widths. It now reports at least one segment per gradient group, like its full-sharing token count. The no-sharing estimate keeps the upper bound for acceptance, and exact pricing keeps actual counts. The TP4 floor also requires the GDN conv kernel width, which prices each segment's conv history. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 11 ++++++++--- tests/unit/test_trainer_rank_tp_floor.py | 10 ++++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 7aa568cb5..e38a0939f 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4052,6 +4052,7 @@ def _sequence_parallel_floor_covered(self, layers: int, tp: int, cp: int) -> boo geometry.gdn_key_head_dim, geometry.gdn_value_heads, geometry.gdn_value_head_dim, + geometry.gdn_conv_kernel, # Prices each segment's conv history. ) if self._gdn_layers and min(gdn_widths) <= 0: return False @@ -5275,7 +5276,7 @@ def priced( signature=signature, logical_tokens=logical_tokens, # Gradient groups' segments: exact layouts' counts, else - # an upper bound (see _estimate_flat_forward). + # a bound matching the estimate's (_estimate_flat_forward). gdn_segments=gdn_segments, group_rows=group_rows, head_workspace_bytes=head_workspace_bytes, @@ -6003,7 +6004,8 @@ def _estimate_flat_forward( content) and is used only inside the band where those bounds disagree. Under CP it returns None: per-rank floors need materialized layouts. ``gdn_segments`` receives each gradient group's segment count: exact - layouts' actual counts, else twice its requests (a radix tree has fewer). + layouts' actual counts; in cheap mode, the same kind of bound as the + token count (twice the requests, as a radix tree has fewer, or one). """ if sync_planning_errors: @@ -6114,7 +6116,10 @@ def _estimate_flat_forward( packed_tokens += physical_rows group_rows.append((physical_rows, grad_enabled)) if grad_enabled and gdn_segments is not None: - gdn_segments.append(2 * len(group_indices)) + # Bounds like the token counts: at most twice the requests + # without sharing (acceptance), at least one with full + # sharing (rejection); exact pricing counts the rest. + gdn_segments.append(1 if memory_minimal else 2 * len(group_indices)) head_workspace_bytes = max( head_workspace_bytes, self._group_head_workspace_bytes( diff --git a/tests/unit/test_trainer_rank_tp_floor.py b/tests/unit/test_trainer_rank_tp_floor.py index 19ad0ec10..6d5f4d00d 100644 --- a/tests/unit/test_trainer_rank_tp_floor.py +++ b/tests/unit/test_trainer_rank_tp_floor.py @@ -151,6 +151,7 @@ def test_rows_are_sharded_with_ceiling_and_only_gradient_groups_save_them(): "moe_geometry", "replicated_qkv", "missing_attention_geometry", + "missing_conv_kernel", "shallow", "wide_ffn", ], @@ -168,6 +169,7 @@ def test_unproven_shapes_keep_todays_pricing(case): "moe_geometry": dict(), "replicated_qkv": dict(), "missing_attention_geometry": dict(), + "missing_conv_kernel": dict(), "shallow": dict(layers=48), "wide_ffn": dict(ffn=4 * F), } @@ -182,6 +184,8 @@ def test_unproven_shapes_keep_todays_pricing(case): "moe_geometry": dict(moe_experts=8), "replicated_qkv": dict(num_query_groups=2), "missing_attention_geometry": dict(kv_channels=0), + # Unread conv history would leave each segment's price short. + "missing_conv_kernel": dict(gdn_conv_kernel=0), } if case in edits: r._geometry = replace(r._geometry, **edits[case]) @@ -241,6 +245,12 @@ def test_width_probes_count_only_gradient_segments(): cheap: list[int] = [] assert r._estimate_flat_forward(requests(), gdn_segments=cheap) is not None assert cheap == [2] # A radix tree has fewer than twice its requests. + # The full-sharing estimate rejects widths, so it takes a lower bound. + minimal: list[int] = [] + assert r._estimate_flat_forward( + requests(), memory_minimal=True, gdn_segments=minimal + ) + assert minimal == [1] exact: list[int] = [] assert r._estimate_flat_forward(requests(), exact=True, gdn_segments=exact) assert exact == [1] # The selected layout's actual segments.