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..e38a0939f 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4019,10 +4019,69 @@ 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), 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, + geometry.gdn_conv_kernel, # Prices each segment's conv history. + ) + 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 attention_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 + 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. @@ -4033,6 +4092,10 @@ 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, + 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: @@ -4052,12 +4115,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 +4135,8 @@ 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 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 +4152,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 @@ -4096,6 +4167,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 and gradient_rows: + # Recurrent states grow with segments, not rows; backward recomputes + # 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: @@ -4142,7 +4219,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, @@ -5170,8 +5247,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 @@ -5185,6 +5266,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( @@ -5192,12 +5275,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 + # a bound matching the estimate's (_estimate_flat_forward). + gdn_segments=gdn_segments, group_rows=group_rows, head_workspace_bytes=head_workspace_bytes, ) @@ -5211,14 +5291,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( @@ -5231,7 +5317,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]) @@ -5904,6 +5990,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. @@ -5916,6 +6003,9 @@ 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; 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: @@ -5965,6 +6055,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 = ( @@ -6023,6 +6115,11 @@ 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: + # 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( @@ -7920,7 +8017,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_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) 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..6d5f4d00d --- /dev/null +++ b/tests/unit/test_trainer_rank_tp_floor.py @@ -0,0 +1,256 @@ +"""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 +# 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): + 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: Any = 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, + gdn_conv_kernel=4, + ) + 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, gdn_segments=1): + 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, + gdn_segments=gdn_segments, + 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 + # 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 plus up to three TP-padding roots, each with its states. + state = cost.checkpoint_workspace + 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 + # (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); the gradient group adds its TP-padding roots' states. + assert mixed[1] == 4 * 4096 * H * 2 + 3 * SEGMENT + + +@pytest.mark.parametrize( + "case", + [ + "tp2", + "tp8", + "cp2", + "pp2", + "no_sequence_parallel", + "sequence_parallel_at_tp1", + "selective_recompute", + "moe", + "moe_geometry", + "replicated_qkv", + "missing_attention_geometry", + "missing_conv_kernel", + "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(), + "moe_geometry": dict(), + "replicated_qkv": dict(), + "missing_attention_geometry": dict(), + "missing_conv_kernel": dict(), + "shallow": dict(layers=48), + "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 + # 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), + # 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]) + 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 + # 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) + 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. + # 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.