diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 18a49a302..9005b367c 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -241,6 +241,8 @@ jobs: tests/unit/test_trainer_rank_pending_memory.py \ tests/unit/test_trainer_rank_shared_memory.py \ tests/unit/test_trainer_rank_converted_memory.py \ + tests/unit/test_trainer_rank_layout_memory.py \ + tests/unit/test_context_parallel_retained_bytes.py \ tests/unit/test_trainer_rank_split.py \ tests/unit/test_megatron_compile_garbage.py \ tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \ @@ -287,4 +289,6 @@ jobs: --ignore=tests/unit/test_trainer_rank_pending_memory.py \ --ignore=tests/unit/test_trainer_rank_shared_memory.py \ --ignore=tests/unit/test_megatron_compile_garbage.py \ - --ignore=tests/unit/test_trainer_rank_converted_memory.py + --ignore=tests/unit/test_trainer_rank_converted_memory.py \ + --ignore=tests/unit/test_trainer_rank_layout_memory.py \ + --ignore=tests/unit/test_context_parallel_retained_bytes.py diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index 4015011fe..c59751966 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -35,6 +35,7 @@ DkvReducePlan, ExactMaskMetadata, FlexMaskSpec, + RankRuntimePlan, StageExecutionSpec, StagePlan, TokenRange, @@ -1550,6 +1551,112 @@ def _merge_stage_output_grads_from_tape( return stage_out_grads, stage_lse_grads +def minimum_retained_bytes_per_row( + *, + q_heads: int, + kv_heads: int, + head_dim: int, + value_head_dim: int, + element_size: int, +) -> int: + """The least ``retained_stage_record_bytes`` keeps per own row. + + Every rank with rows runs a local stage over all of them (each row attends + to itself); an aligned one keeps at least the contiguous copies of + multi-head views, flex's output and its two LSEs. + """ + copies = (q_heads > 1) * q_heads * head_dim + (kv_heads > 1) * kv_heads * ( + head_dim + value_head_dim + ) + return (copies + q_heads * value_head_dim) * element_size + 2 * q_heads * 4 + + +def retained_stage_record_bytes( + rank_plan: RankRuntimePlan, + *, + q_heads: int, + kv_heads: int, + head_dim: int, + value_head_dim: int, + element_size: int, + block_size: SparseBlockSize, +) -> int: + """Bytes a recomputed attention keeps for backward beyond its own-row tensors. + + A size-only mirror of ``_forward_stage_records`` with ``record_for_backward`` + and of ``_run_stage_attention``; keep the three in step. Every stage that + runs keeps the Q/K/V flex consumed: a padded copy when the execution length + differs, else a contiguous copy of the permuted ``q_flat``/``k_flat`` view; + partial-range gathers and remote fetch buffers are kept as the stage's + inputs. It keeps flex's output and LSE at the execution length, and + logical-length copies of them when padded, plus flex's own LSE beside the + normalized one the FLASH backend returns (counted on every backend). Every + producing stage after the first keeps a merge-tape clone of the + accumulators. Accumulators themselves are transient. + """ + own = int(rank_plan.local_valid_lengths[0]) if rank_plan.local_valid_lengths else 0 + accum_size = 4 if element_size < 4 else element_size + q_row = q_heads * head_dim * element_size + k_row = kv_heads * head_dim * element_size + v_row = kv_heads * value_head_dim * element_size + out_row = q_heads * value_head_dim * element_size + lse_row = q_heads * 4 + tape_row = q_heads * (value_head_dim + 1) * accum_size + total = 0 + local_produced = False + tapes: list[int] = [] + for stage in _ordered_stage_plans(rank_plan.stage_plans): + if not (stage.q_len > 0 and stage.k_len > 0 and stage.slices): + continue + q_len = _logical_stage_q_len(stage) + k_len = _logical_stage_k_len(stage) + q_pad, k_pad, _family = select_sparse_execution_family( + is_local_stage=bool(stage.is_local_stage), + q_len=int(stage.q_len), + k_len=int(stage.k_len), + block_size=block_size, + ) + q_full = _ranges_cover_full_length(stage.owner_local_q_ranges, length=own) + # Queries: the full range is a permuted view of q_flat; a partial range + # is gathered into a contiguous tensor kept as the stage's input. + if not q_full: + total += q_row * q_len + if q_pad != q_len: + total += q_row * q_pad + elif q_full: + # A view of the projection output: flex copies it unless it is + # already contiguous, which one head alone does not guarantee + # (a fused QKV split keeps the projection's token stride). + total += q_row * q_len + # Keys and values: local ranges as for queries; remote ones land in + # contiguous head-major fetch buffers kept as the stage's inputs. + k_full = bool(stage.is_local_stage) and _ranges_cover_full_length( + stage.owner_local_k_ranges, length=own + ) + if not k_full: + total += (k_row + v_row) * k_len + if k_pad != k_len: + total += (k_row + v_row) * k_pad + elif k_full: + total += (k_row + v_row) * k_len + total += (out_row + 2 * lse_row) * q_pad + if q_pad != q_len: + total += (out_row + lse_row) * q_len + tape = tape_row * (own if q_full else q_len) + if not q_full: + # Its int64 row index is kept even by the first producing stage. + total += 8 * q_len + if stage.is_local_stage: + local_produced = True + else: + tapes.append(tape) + # The first producing stage records no tape: the local stage when it ran, + # else whichever remote stage is ready first, so drop the smallest tape. + if tapes and not local_produced: + tapes.remove(min(tapes)) + return total + sum(tapes) + + def _forward_stage_records( *, q_flat: torch.Tensor, diff --git a/src/art/megatron/context_parallel/runtime.py b/src/art/megatron/context_parallel/runtime.py index 531cb8829..2f7ac2dad 100644 --- a/src/art/megatron/context_parallel/runtime.py +++ b/src/art/megatron/context_parallel/runtime.py @@ -433,6 +433,55 @@ def context_parallel_rank_model_token_counts( ) +def context_parallel_rank_layouts( + *, + group_ids: torch.Tensor, + parent_ids: torch.Tensor, + topology: ParallelTopology, + config: ContextParallelConfig, + original_seq_len: int, + build_gdn_execution_spec: bool, + gdn_planner_config: Any | None = None, +) -> tuple[tuple[int, ...], tuple[int, ...] | None, tuple[RankRuntimePlan, ...]]: + """Each CP rank's attention rows, GDN rows and attention stage plan. + + Uses the cached planning bundle and per-rank runtime plans that execution + builds, so a memory estimate sees the layouts the ranks will run. + """ + planning_key, bundle, _group_ids_cpu, _parent_ids_cpu = ( + _get_or_build_planning_bundle( + group_ids=group_ids, + parent_ids=parent_ids, + topology=topology, + config=config, + original_seq_len=original_seq_len, + build_gdn_execution_spec=build_gdn_execution_spec, + ) + ) + attention = tuple(bundle.token_layout_index.token_counts_by_rank) + gdn = None + if build_gdn_execution_spec: + gdn = tuple( + _plan_gdn_global_execution( + planning_key=planning_key, + bundle=bundle, + topology=topology, + gdn_planner_config=gdn_planner_config, + ).gdn_token_counts_by_rank + ) + plans = tuple( + _get_or_build_bundle_rank_plan( + planning_key=planning_key, + bundle=bundle, + original_seq_len=original_seq_len, + target_rank=rank, + block_size=config.block_size, + ) + for rank in range(len(attention)) + ) + return attention, gdn, plans + + def context_parallel_model_token_total( *, group_ids: torch.Tensor, diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 0378c8c24..af1b5bdfd 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1041,6 +1041,17 @@ def signature(self) -> _MemorySignature: _AnyForwardPlan = _FlatForwardPlan | _SplitForwardPlan +@dataclass(frozen=True) +class _GroupLayout: + """One packed group's CP layouts on every rank, for layout-aware pricing.""" + + attention_rows: tuple[int, ...] + gdn_rows: tuple[int, ...] | None + # What each rank's recomputed CP attention keeps for backward beyond its + # own-row activations (``retained_stage_record_bytes``). + attention_retained: tuple[int, ...] + + @dataclass(frozen=True) class _SubforwardCost: """Memory terms of one candidate subforward while all graphs stay live. @@ -3693,6 +3704,7 @@ def _split_chunk_lower_cost( unshared_packed_tokens = 0 head_workspace_bytes = 0 group_rows: list[tuple[int, bool]] = [] + group_physical_rows: list[int] = [] for (_slot, grad_enabled), group_indices in groups: estimated = estimate_prefix_tree_packed_tokens( (rows[index] for index in group_indices), @@ -3700,6 +3712,7 @@ def _split_chunk_lower_cost( ) assert estimated is not None # rows are CPU copies physical_rows = self._physical_tokens(estimated) + group_physical_rows.append(physical_rows) packed_tokens += physical_rows # The most loaded CP rank holds at least an even share. cp = max(1, self._topology_key()[2]) @@ -3725,12 +3738,23 @@ def _split_chunk_lower_cost( slot_groups=tuple(key for key, _ in groups), ) logical_tokens = _active_logical_tokens(requests) + # Where exact plans price each rank's layouts, bound them from below + # the same way rather than with the busiest-rank widths. + layouts = ( + self._minimum_layouts(group_physical_rows, signature.topology[2]) + if self._layout_pricing_supported( + signature.topology, + gradient_groups=all(grad for _, grad in group_rows), + ) + else None + ) cost = self._subforward_cost( packed_tokens=packed_tokens, output_bytes=output_bytes, signature=signature, logical_tokens=logical_tokens, group_rows=tuple(group_rows), + group_layouts=layouts, slot_refs=tuple(ref for (ref, _), _ in groups), head_workspace_bytes=head_workspace_bytes, # The average CP load is an optimistic bound, not an admission cost. @@ -4181,8 +4205,10 @@ def _moe_workspace_bytes( raise ValueError("Invalid constructor converted-weight stages") if type(shared) is not int or not 0 <= shared <= coefficient: raise ValueError("Invalid constructor shared-expert coefficient") - routed = rows if routed_rows is None else max(0, min(rows, routed_rows)) - return (rows - routed) * shared + ( + # A rank can receive more routed rows than it holds (an uneven CP + # split); the shared part then stays on the routed rows' count. + routed = rows if routed_rows is None else max(0, routed_rows) + return max(0, rows - routed) * shared + ( max( routed * coefficient, *(routed * per_row + fixed for per_row, fixed in stages), @@ -4255,11 +4281,14 @@ def _checkpoint_memory_floor( slot_refs: tuple["LoRASlotRef | None", ...] | None = None, gdn_segments: int = 0, routed_rows: tuple[int, ...] | None = None, + layouts: tuple[_GroupLayout, ...] | None = None, ) -> tuple[int, int]: """Conservative saved-boundary charge and one recomputed layer's workspace. ``routed_rows`` are each group's balanced dispatched rows per rank - (``_plan_group_routed_rows``); by default, its local rows. + (``_plan_group_routed_rows``); by default, its local rows. With + ``layouts`` (``_plan_group_layouts``), price every rank on its own CP + layouts instead of the busiest rank's rows (``_layout_checkpoint_floor``). Count actual local full/uniform/1 boundaries, including aliases, rather than claiming measured distinct storage. Only this call's new groups enter the term; already-live graphs remain in the availability baseline. @@ -4336,6 +4365,10 @@ def _checkpoint_memory_floor( retained, workspace = self._sequence_parallel_checkpoint_floor( group_rows, refs, layers, tp, gdn_segments ) + elif layouts is not None and all(grad for _, grad in group_rows): + retained, workspace = self._layout_checkpoint_floor( + decoder.layers, refs, routed, layouts + ) else: retained, workspace = self._generic_checkpoint_floor( group_rows, refs, routed, layers @@ -4429,6 +4462,221 @@ def _generic_checkpoint_floor( workspace += self._te_workspace_growth_bytes() return retained, workspace + def _layout_checkpoint_floor( + self, + layers: Sequence[torch.nn.Module], + refs: tuple["LoRASlotRef | None", ...], + routed: tuple[int | None, ...], + layouts: tuple[_GroupLayout, ...], + ) -> tuple[int, int]: + """Each CP rank's boundaries and recomputed layer on its own layouts. + + A saved layer input arrives in the GDN layout when the layer follows a + GDN layer in its island (``_art_gdn_island_boundary``), and in the + attention layout otherwise. A recomputed attention layer keeps its + activations on its attention rows plus what the CP executor keeps for + backward; a GDN layer keeps the GDN width on its GDN rows. Either one's + residual, norm, routing state and MoE stage use that layer's rows, and + routed rows the EP share. Traced CP2 ranks: boundaries match exactly and + attention within 0.3%. Returns the largest rank's boundaries and the + rest of the largest rank total, so the two sum to that total. + """ + hidden = self._hidden_size * 2 + gdn_inputs = sum( + getattr( + getattr(layer, "_art_gdn_island_boundary", None), "input_layout", "" + ) + == "gdn" + for layer in layers + ) + attention_inputs = len(layers) - gdn_inputs + widths = self._recomputed_mixer_widths(stage_buffers=False) + moe = self._checkpoint_moe_bytes_per_token() + beside = 2 * hidden + self._moe_checkpoint_state_bytes_per_token() if moe else 0 + retained_by_rank: list[int] = [] + totals: list[int] = [] + for rank in range(len(layouts[0].attention_rows)): + retained = workspace = 0 + for ref, dispatched, layout in zip(refs, routed, layouts, strict=True): + attention = max(1, layout.attention_rows[rank]) + gdn = ( + attention + if layout.gdn_rows is None + else max(1, layout.gdn_rows[rank]) + ) + retained += hidden * (attention_inputs * attention + gdn_inputs * gdn) + for kind, rows, extra in ( + ("attention", attention, layout.attention_retained[rank]), + ("gdn", gdn, 0), + ): + if kind not in widths: + continue + stage = ( + rows * (widths[kind] + beside) + + extra + + self._moe_workspace_bytes( + rows, + routed_rows=dispatched, + checkpoint_grad=True, + slot_ref=ref, + ) + ) + workspace = max(workspace, stage) + retained_by_rank.append(retained) + totals.append(retained + workspace) + retained = max(retained_by_rank) + workspace = max(totals) - retained + if moe: + workspace += self._te_workspace_growth_bytes() + return retained, workspace + + def _layout_pricing_supported( + self, topology: tuple[int, int, int, int], *, gradient_groups: bool + ) -> bool: + """Whether layout-aware pricing models this runtime and plan shape.""" + _dp, tp, cp, pp = topology + if (tp, cp, pp) != (1, 2, 1) or not gradient_groups: + return False + geometry = self._geometry + if not geometry.num_attention_heads or not geometry.kv_channels: + return False + if len(self.runtime.model) != 1: + return False + try: + decoder = _language_model(self.runtime.model[0]).decoder + from art.megatron.context_parallel.core_attention import ( + ArtContextParallelCoreAttention, + ) + except (AttributeError, RuntimeError, ModuleNotFoundError): + return False + layers = getattr(decoder, "layers", None) + if layers is None: + return False + for layer in layers: + boundary = getattr(layer, "_art_gdn_island_boundary", None) + if boundary is not None and boundary.is_gdn: + continue + if self._gdn_layers and boundary is None: + return False + core = getattr( + getattr(layer, "self_attention", None), "core_attention", None + ) + if ( + type(core) is not ArtContextParallelCoreAttention + or getattr(core, "softmax_offset", None) is not None + ): + return False + return True + + def _minimum_layouts( + self, physical_rows: Sequence[int], cp: int + ) -> tuple[_GroupLayout, ...]: + """Even-share layouts keeping the least attention state: a lower bound. + + Every rank's total grows with its own rows, and every rank keeps at + least an aligned local stage's state per row (each row attends to + itself), so the largest rank's total is at least the total at the mean + rows. The mean is at least the floor of an even share, which is why + this rounds down; rounding up can exceed a split's exact cost. + """ + from art.megatron.context_parallel.executor import ( + minimum_retained_bytes_per_row, + ) + + geometry = self._geometry + per_row = minimum_retained_bytes_per_row( + q_heads=int(geometry.num_attention_heads), + kv_heads=int(geometry.num_query_groups), + head_dim=int(geometry.kv_channels), + value_head_dim=int(geometry.kv_channels), + element_size=self._param_dtype_size, + ) + return tuple( + _GroupLayout( + attention_rows=(rows // cp,) * cp, + gdn_rows=(rows // cp,) * cp if self._gdn_layers else None, + attention_retained=(rows // cp * per_row,) * cp, + ) + for rows in physical_rows + ) + + def _plan_group_layouts( + self, plan: _FlatForwardPlan + ) -> tuple[_GroupLayout, ...] | None: + """Every rank's CP layouts per group, where layout pricing is modeled. + + Only CP2 at TP1/PP1 with gradient groups, ART's CP core attention with + no softmax offset, and GDN layers marked with island boundaries; the + executor's retained set is validated there. Elsewhere ``None`` keeps + the busiest-rank pricing. + """ + if not plan.groups or not self._layout_pricing_supported( + plan.signature.topology, + gradient_groups=all(group.grad_enabled for group in plan.groups), + ): + return None + started = time.perf_counter() + try: + return self._compute_group_layouts(plan) + finally: + # Planning work: every rank's CP plan, cached by planning key. + self._planning_seconds_accum += time.perf_counter() - started + + def _compute_group_layouts( + self, plan: _FlatForwardPlan + ) -> tuple[_GroupLayout, ...]: + geometry = self._geometry + from art.megatron.context_parallel.executor import retained_stage_record_bytes + from art.megatron.context_parallel.runtime import context_parallel_rank_layouts + from art.megatron.flex_attn.compiled import flash_sparse_block_size_for_head_dim + from art.megatron.training.microbatches import ( + _context_parallel_config_for_provider, + _gdn_planner_config_for_provider, + ) + + topology = self._topology() + handler = self.runtime.model_support_handler + config = _context_parallel_config_for_provider( + self.runtime.provider, self.device, handler + ) + head = int(geometry.kv_channels) + block = flash_sparse_block_size_for_head_dim( + head_dim=head, head_dim_v=head, device=self.device + ) + layouts = [] + for group in plan.groups: + batch = _pad_packed_batch(group.packed, multiple=int(topology.tp)) + attention, gdn, rank_plans = context_parallel_rank_layouts( + group_ids=batch.group_ids, + parent_ids=batch.parent_ids, + topology=topology, + config=config, + original_seq_len=int(batch.tokens.shape[1]), + build_gdn_execution_spec=handler.build_gdn_execution_spec, + gdn_planner_config=_gdn_planner_config_for_provider( + self.runtime.provider, handler + ), + ) + layouts.append( + _GroupLayout( + attention_rows=attention, + gdn_rows=gdn, + attention_retained=tuple( + retained_stage_record_bytes( + rank_plan, + q_heads=int(geometry.num_attention_heads), + kv_heads=int(geometry.num_query_groups), + head_dim=head, + value_head_dim=head, + element_size=self._param_dtype_size, + block_size=block, + ) + for rank_plan in rank_plans + ), + ) + ) + return tuple(layouts) + def _te_workspace_growth_bytes(self) -> int: """Transformer Engine's cuBLAS workspaces, until its GEMMs allocate them. @@ -4525,6 +4773,7 @@ def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: gdn_segments=plan.grad_segment_count, group_rows=self._plan_group_rows(plan), group_routed_rows=self._plan_group_routed_rows(plan), + group_layouts=self._plan_group_layouts(plan), slot_refs=tuple(g.slot_ref for g in plan.groups), head_workspace_bytes=self._plan_head_workspace_bytes(plan), checkpoint_floor=_gdn_memory.plan_floor(self, plan), @@ -4542,6 +4791,7 @@ def _subforward_cost( gdn_segments: int = 0, group_rows: tuple[tuple[int, bool], ...] = (), group_routed_rows: tuple[int, ...] | None = None, + group_layouts: tuple[_GroupLayout, ...] | None = None, slot_refs: tuple["LoRASlotRef | None", ...] | None = None, head_workspace_bytes: int = 0, checkpoint_floor: tuple[int, int] = (0, 0), @@ -4556,6 +4806,7 @@ def _subforward_cost( gdn_segments=gdn_segments, group_rows=group_rows, group_routed_rows=group_routed_rows, + group_layouts=group_layouts, slot_refs=slot_refs, head_workspace_bytes=head_workspace_bytes, checkpoint_floor=checkpoint_floor, @@ -4567,6 +4818,7 @@ def _subforward_cost( slot_refs, gdn_segments, routed_rows=group_routed_rows, + layouts=group_layouts, ) retained = self._retained_memory_bytes( signature, @@ -6696,6 +6948,11 @@ def _fill_planner_snapshot( "retained_tokens": self._plan_retained_tokens(child), "group_rows": self._plan_group_rows(child), "group_routed_rows": self._plan_group_routed_rows(child), + "group_layouts": ( + None + if (layouts := self._plan_group_layouts(child)) is None + else [asdict(layout) for layout in layouts] + ), "hybridep_growth_bytes": ( self._plan_hybridep_growth_bytes(child) ), @@ -7421,6 +7678,7 @@ def _memory_check( gdn_segments=forward.grad_segment_count, group_rows=self._plan_group_rows(forward), group_routed_rows=self._plan_group_routed_rows(forward), + group_layouts=self._plan_group_layouts(forward), slot_refs=tuple(g.slot_ref for g in forward.groups), head_workspace_bytes=self._plan_head_workspace_bytes(forward), checkpoint_floor=_gdn_memory.plan_floor(self, forward), @@ -8268,6 +8526,7 @@ def _estimate_required_memory_bytes_from_values( gdn_segments: int = 0, group_rows: tuple[tuple[int, bool], ...] = (), group_routed_rows: tuple[int, ...] | None = None, + group_layouts: tuple[_GroupLayout, ...] | None = None, slot_refs: tuple["LoRASlotRef | None", ...] | None = None, head_workspace_bytes: int = 0, checkpoint_floor: tuple[int, int] = (0, 0), @@ -8364,6 +8623,7 @@ def _estimate_required_memory_bytes_from_values( slot_refs, gdn_segments, routed_rows=group_routed_rows, + layouts=group_layouts, ) static_compute = max( static_compute, @@ -8527,18 +8787,26 @@ def _recomputed_mixer_bytes_per_token(self) -> int: more rows; its hidden-width input exchange and value-width output allowance price that (88 KB measured at CP2, 94 KB priced). """ + return max(self._recomputed_mixer_widths().values(), default=0) + + def _recomputed_mixer_widths(self, *, stage_buffers: bool = True) -> dict[str, int]: + """Recomputed attention and GDN mixer bytes per row, by layer type. + + ``stage_buffers=False`` leaves out the CP attention stage allowance, for + callers that price the executor's stage buffers from its stage plan. + """ geometry = self._geometry hidden = self._hidden_size tp = max(1, self._topology_key()[1]) cp = self._topology_key()[2] > 1 - widths = [] + widths: dict[str, float] = {} if self._gdn_layers < self._num_layers: attention, _gdn = self._mixer_activation_widths() - if cp: + if cp and stage_buffers: q = geometry.num_attention_heads * geometry.kv_channels or hidden kv = geometry.num_query_groups * geometry.kv_channels or hidden attention += (3 * q + 2 * kv) / tp - widths.append(attention) + widths["attention"] = attention if self._gdn_layers: key = geometry.gdn_key_heads * geometry.gdn_key_head_dim value = geometry.gdn_value_heads * geometry.gdn_value_head_dim @@ -8547,8 +8815,10 @@ def _recomputed_mixer_bytes_per_token(self) -> int: gdn = hidden + (2 * key + normalized + 6 * value + chunk) / tp if cp: gdn += hidden + value / tp - widths.append(gdn) - return int(max(widths, default=0) * self._param_dtype_size) + widths["gdn"] = gdn + return { + kind: int(width * self._param_dtype_size) for kind, width in widths.items() + } def _gdn_segment_layer_bytes(self) -> float: """Initial and final fp32 recurrent states plus convolution history.""" diff --git a/tests/unit/test_context_parallel_retained_bytes.py b/tests/unit/test_context_parallel_retained_bytes.py new file mode 100644 index 000000000..18f0f4b48 --- /dev/null +++ b/tests/unit/test_context_parallel_retained_bytes.py @@ -0,0 +1,157 @@ +"""Size-only retained-set mirror of the CP attention executor's recompute path.""" + +import pytest + +pytest.importorskip("triton") + +from art.megatron.context_parallel.executor import ( # noqa: E402 + minimum_retained_bytes_per_row, + retained_stage_record_bytes, +) +from art.megatron.context_parallel.types import ( # noqa: E402 + RankRuntimePlan, + StagePlan, + TokenRange, +) + + +# Qwen3.6-35B-A3B attention: 16 query heads, 2 KV heads of 256, BF16, and the +# H200 flash block for a 256-wide head (128 query, 64 key rows). +def retained(runtime_plan, *, q_heads=16, kv_heads=2): + return retained_stage_record_bytes( + runtime_plan, + q_heads=q_heads, + kv_heads=kv_heads, + head_dim=256, + value_head_dim=256, + element_size=2, + block_size=(128, 64), + ) + + +def minimum(*, q_heads=16, kv_heads=2): + return minimum_retained_bytes_per_row( + q_heads=q_heads, + kv_heads=kv_heads, + head_dim=256, + value_head_dim=256, + element_size=2, + ) + + +# Flex's output keeps its own LSE beside the normalized copy it returns; +# logical-length copies keep one. +Q, KV, FLEX, OUT, TAPE = 8192, 2048, 8192 + 2 * 64, 8192 + 64, 16 * 257 * 4 + + +def stage(index, *, local, q, k, q_len=None, k_len=None, own=None, source=0): + q_ranges = ( + (TokenRange(0, q),) if own is None or q == own else (TokenRange(1, 1 + q),) + ) + return StagePlan( + stage_index=index, + source_rank=source, + is_local_stage=local, + slices=("slice",), # ty: ignore[invalid-argument-type] + owner_local_q_ranges=q_ranges if q else (), + owner_local_k_ranges=(TokenRange(0, k),) if k else (), + q_len=q if q_len is None else q_len, + k_len=k if k_len is None else k_len, + ) + + +def plan(own, *stages): + return RankRuntimePlan( + rank=0, + original_seq_len=own, + token_layout_index=None, # ty: ignore[invalid-argument-type] + local_valid_lengths=(own,), + local_row_ranges=(TokenRange(0, own),), + stage_plans=stages, + remote_dkv_reduce_plan=None, # ty: ignore[invalid-argument-type] + ) + + +def test_aligned_single_stage_copies_views_without_output_copies(): + # Real-data rank 0: one aligned local stage. Contiguous copies of the + # permuted Q/K/V views, flex output and LSE; no padding, copies or tape. + rows = 52480 + kept = retained(plan(rows, stage(0, local=True, q=rows, k=rows))) + assert kept == Q * rows + KV * rows + FLEX * rows # 0.971 GB traced + assert kept == rows * minimum() + + +def test_unaligned_single_stage_pads_and_copies_the_logical_output(): + # The tail chunk on the single-stage rank: the planner rounds both lengths + # up to its 128-row block, so both pad. + rows = 52481 + stage_len = 411 * 128 + kept = retained( + plan( + rows, stage(0, local=True, q=rows, k=rows, q_len=stage_len, k_len=stage_len) + ) + ) + assert kept == Q * stage_len + KV * stage_len + FLEX * stage_len + OUT * rows + + +def test_tiny_stage_pads_to_two_blocks(): + kept = retained(plan(5, stage(0, local=True, q=5, k=5))) + assert kept == Q * 256 + KV * 128 + FLEX * 256 + OUT * 5 + + +def test_single_head_views_may_still_be_copied(): + # One head does not make a view of a fused QKV split contiguous, so the + # mirror still charges the copies; only the lower bound leaves them out. + rows = 1024 + kept = retained( + plan(rows, stage(0, local=True, q=rows, k=rows)), q_heads=1, kv_heads=1 + ) + assert kept == (512 + 1024 + 512 + 2 * 4) * rows + assert minimum(q_heads=1, kv_heads=1) == 512 + 2 * 4 + + +def test_full_query_remote_stage_keeps_fetch_buffers_and_a_merge_tape(): + # Real-data rank 1: a local stage and a remote stage over all of its + # queries. Planner lengths round up, so both stages pad. + own, remote_k = 44314, 16504 + local = stage(0, local=True, q=own, k=own, q_len=44352, k_len=44352) + remote = stage(1, local=False, q=own, k=remote_k, q_len=44352, k_len=16512) + kept = retained(plan(own, local, remote)) + q_pad = 347 * 128 # 44,416, as flex's traced output size shows + local_bytes = Q * q_pad + KV * 44352 + FLEX * q_pad + OUT * own + remote_bytes = ( + Q * q_pad + KV * remote_k + KV * 16512 + FLEX * q_pad + OUT * own + TAPE * own + ) + assert kept == local_bytes + remote_bytes + assert 3.05e9 < kept < 3.10e9 # 3.07 GB traced + + +def test_partial_query_remote_stage_keeps_its_gather_and_a_partial_tape(): + # Random-data rank 1: a large local stage and a 768-row remote stage. + own = 105153 + local = stage(0, local=True, q=own, k=own, q_len=105216, k_len=105216) + remote = stage(1, local=False, q=768, k=20608, own=own) + kept = retained(plan(own, local, remote)) + local_bytes = Q * 105216 + KV * 105216 + FLEX * 105216 + OUT * own + # Aligned: the query gather and fetch buffers feed flex without copies; the + # partial tape also keeps its int64 row index. + remote_bytes = Q * 768 + KV * 20608 + FLEX * 768 + (TAPE + 8) * 768 + assert kept == local_bytes + remote_bytes + + +def test_empty_remote_stage_and_missing_local_stage(): + rows = 1024 + empty = stage(1, local=False, q=0, k=0) + alone = retained(plan(rows, stage(0, local=True, q=rows, k=rows), empty)) + assert alone == Q * rows + KV * rows + FLEX * rows + # Without a local stage, the first ready remote stage records no tape; the + # mirror drops the smallest so it never under-counts the order. + small = stage(1, local=False, q=256, k=256, own=rows) + full = stage(2, local=False, q=rows, k=512) + both = retained(plan(rows, small, full)) + without_small_tape = retained(plan(rows, full)) + # The dropped tape's int64 index stays: the executor keeps it regardless. + assert both - without_small_tape == ( + Q * 256 + KV * 256 + FLEX * 256 + 8 * 256 + TAPE * rows + ) + assert retained(plan(rows)) == 0 diff --git a/tests/unit/test_trainer_rank_admission_inputs.py b/tests/unit/test_trainer_rank_admission_inputs.py index 3481bfd6d..66b4bf812 100644 --- a/tests/unit/test_trainer_rank_admission_inputs.py +++ b/tests/unit/test_trainer_rank_admission_inputs.py @@ -33,6 +33,7 @@ def assert_plan_values(rank, plan, values): assert values["gdn_segments"] == plan.grad_segment_count assert values["group_rows"] == rank._plan_group_rows(plan) assert values["group_routed_rows"] == rank._plan_group_routed_rows(plan) + assert values["group_layouts"] == rank._plan_group_layouts(plan) assert values["head_workspace_bytes"] == rank._plan_head_workspace_bytes(plan) assert values["checkpoint_floor"] == _gdn_memory.plan_floor(rank, plan) assert values["retained_tokens"] == rank._plan_retained_tokens(plan) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 78fcb7f73..223c2dd54 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -534,10 +534,11 @@ def test_routed_rows_move_only_the_routed_moe_part(): fewer = r._checkpoint_memory_floor(((local, True),), None, routed_rows=(routed,)) assert fewer[0] == retained assert workspace - fewer[1] == (local - routed) * (188416 - 8192) - # Never more routed rows than local ones. - assert r._moe_workspace_bytes( - 10, routed_rows=20, checkpoint_grad=True - ) == r._moe_workspace_bytes(10, checkpoint_grad=True) + # A rank can receive more routed rows than it holds: all are priced, and + # the shared part stays on the routed count. + assert r._moe_workspace_bytes(10, routed_rows=20, checkpoint_grad=True) == ( + 20 * 188416 + ) r._moe_gradient_shared_bytes = 188417 with pytest.raises(ValueError, match="shared-expert"): r._moe_workspace_bytes(10, checkpoint_grad=True) diff --git a/tests/unit/test_trainer_rank_layout_memory.py b/tests/unit/test_trainer_rank_layout_memory.py new file mode 100644 index 000000000..12942e7a3 --- /dev/null +++ b/tests/unit/test_trainer_rank_layout_memory.py @@ -0,0 +1,304 @@ +"""CP2 layout-aware recompute pricing: per-rank ledgers and gates, CPU only.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +from test_trainer_rank_checkpoint_memory import rank +import torch + +from art.trainer_rank import ForwardInput +from art.trainer_rank._impl import _TE_CUBLAS_WORKSPACE_BYTES, Unset, _GroupLayout + +H = 2048 * 2 + + +def qwen36(r): + """Qwen3.6-35B-A3B's 3:1 GDN/attention pattern and mixer geometry at CP2.""" + r._topology_key = lambda: (1, 1, 2, 1) + r._geometry = replace( + r._geometry, + num_attention_heads=16, + num_query_groups=2, + kv_channels=256, + gdn_key_heads=16, + gdn_key_head_dim=128, + gdn_value_heads=32, + gdn_value_head_dim=128, + ) + r._attention_output_gate = True + r._gdn_layers = 30 + for index, layer in enumerate(r.runtime.model[0].decoder.layers): + is_gdn = index % 4 != 3 + after_gdn = index > 0 and (index - 1) % 4 != 3 + layer._art_gdn_island_boundary = SimpleNamespace( + is_gdn=is_gdn, input_layout="gdn" if is_gdn and after_gdn else "attention" + ) + return r + + +def test_boundaries_follow_each_layer_inputs_layout(): + # Traced real-data CP2 ranks: 20 inputs in each layout, and boundaries of + # 8.248 and 7.612 GB. The busiest attention rank is not the busiest overall. + r = qwen36(rank()) + layers = r.runtime.model[0].decoder.layers + layout = _GroupLayout((52480, 44314), (48194, 48600), (0, 0)) + retained, _ = r._layout_checkpoint_floor(layers, (None,), (48397,), (layout,)) + assert retained == H * (20 * 52480 + 20 * 48194) + assert H * (20 * 44314 + 20 * 48600) < retained < 40 * 52480 * H + + +def test_largest_rank_total_not_a_sum_of_rank_maxima(): + r = qwen36(rank()) + layers = r.runtime.model[0].decoder.layers + widths = r._recomputed_mixer_widths(stage_buffers=False) + assert widths == {"attention": 2 * (2 * 2048 + 7 * 4096 + 3 * 512), "gdn": 94208} + # Rank 1's two full-query stages keep 3.08 GB against rank 0's 0.97 GB. + layout = _GroupLayout((52480, 44314), (48194, 48600), (970_720_000, 3_079_000_000)) + retained, workspace = r._layout_checkpoint_floor( + layers, (None,), (48397,), (layout,) + ) + + def total(rank_index): + attention = layout.attention_rows[rank_index] + assert layout.gdn_rows is not None + gdn = layout.gdn_rows[rank_index] + ledger = H * (20 * attention + 20 * gdn) + stages = ( + attention * (widths["attention"] + 2 * H) + + layout.attention_retained[rank_index] + + r._moe_workspace_bytes( + attention, routed_rows=48397, checkpoint_grad=True + ), + gdn * (widths["gdn"] + 2 * H) + + r._moe_workspace_bytes(gdn, routed_rows=48397, checkpoint_grad=True), + ) + return ledger + max(stages) + + assert retained + workspace == max(total(0), total(1)) + _TE_CUBLAS_WORKSPACE_BYTES + # Rank 1 has fewer rows but is the largest total: its attention runs two + # full-query stages. Pricing it on its own rows is about neutral against + # the busiest-rank floor, which over-counts boundaries and GDN instead. + assert total(1) > total(0) + busiest = sum( + r._checkpoint_memory_floor(((52480, True),), None, routed_rows=(48397,)) + ) + assert abs(retained + workspace - busiest) < 0.01 * busiest + + +def test_routed_rows_above_a_ranks_own_rows_are_all_priced(): + r = qwen36(rank()) + layers = r.runtime.model[0].decoder.layers + few = _GroupLayout((100, 100), (100, 100), (0, 0)) + many = _GroupLayout((100, 100), (100, 100), (0, 0)) + low = sum(r._layout_checkpoint_floor(layers, (None,), (100,), (few,))) + high = sum(r._layout_checkpoint_floor(layers, (None,), (400,), (many,))) + assert high - low == 300 * 188416 + + +def test_no_grad_groups_keep_busiest_rank_pricing(): + r = qwen36(rank()) + layout = _GroupLayout((10, 8), (9, 9), (0, 0)) + groups = ((10, True), (12, False)) + assert r._checkpoint_memory_floor( + groups, None, routed_rows=(9, 12), layouts=(layout, layout) + ) == r._checkpoint_memory_floor(groups, None, routed_rows=(9, 12)) + + +def _plan(r, *, no_grad=False): + plan = r._plan_flat_forward( + [ + ForwardInput( + input_tokens=torch.arange(64), hidden_states=True, no_grad=no_grad + ) + ] + ) + return replace(plan, signature=replace(plan.signature, topology=(1, 1, 2, 1))) + + +@pytest.mark.parametrize( + "topology", [(1, 1, 1, 1), (1, 2, 2, 1), (1, 1, 4, 1), (1, 1, 2, 2)] +) +def test_layout_pricing_is_cp2_tp1_pp1_only(topology): + r = qwen36(rank()) + plan = _plan(r) + plan = replace(plan, signature=replace(plan.signature, topology=topology)) + assert r._plan_group_layouts(plan) is None + + +def test_layout_pricing_needs_gradient_groups_and_art_cp_attention(): + r = qwen36(rank()) + assert r._plan_group_layouts(_plan(r, no_grad=True)) is None + # The stub decoder's attention layers carry no ART CP core attention. + assert r._plan_group_layouts(_plan(r)) is None + # A GDN model without island boundaries is not modeled either. + for layer in r.runtime.model[0].decoder.layers: + del layer._art_gdn_island_boundary + assert r._plan_group_layouts(_plan(r)) is None + + +def test_plan_cost_and_admission_use_the_same_layouts(monkeypatch): + r = qwen36(rank()) + plan = _plan(r) + layout = _GroupLayout((40, 24), (32, 32), (1_000_000, 3_000_000)) + monkeypatch.setattr(r, "_plan_group_layouts", lambda plan: (layout,)) + monkeypatch.setattr(r, "_plan_group_rows", lambda plan: ((40, True),)) + monkeypatch.setattr(r, "_plan_group_routed_rows", lambda plan: (32,)) + monkeypatch.setattr(r, "_plan_hybridep_growth_bytes", lambda plan: 0) + monkeypatch.setattr(r, "_plan_retained_tokens", lambda plan: 32) + cost = r._plan_cost(plan) + # Admission adds CP output coexistence the plan cost leaves out (as it did + # before layouts); the layouts themselves enter both the same way. + gap = r._memory_check(plan).estimated_required_bytes - cost.required + monkeypatch.setattr(r, "_plan_group_layouts", lambda plan: None) + assert gap == r._memory_check(plan).estimated_required_bytes - ( + r._plan_cost(plan).required + ) + monkeypatch.setattr(r, "_plan_group_layouts", lambda plan: (layout,)) + retained, _ = r._layout_checkpoint_floor( + r.runtime.model[0].decoder.layers, + (None,), + r._plan_group_routed_rows(plan), + (layout,), + ) + assert cost.checkpoint_retained == plan.output_bytes + retained + + +def art_cp(r, monkeypatch): + """Qwen3.6 at CP2 with ART's CP core attention and a CPU planning config.""" + from art.megatron.context_parallel.core_attention import ( + ArtContextParallelCoreAttention, + ) + from art.megatron.context_parallel.types import ParallelTopology + + r = qwen36(r) + for index, layer in enumerate(r.runtime.model[0].decoder.layers): + if index % 4 == 3: + core = ArtContextParallelCoreAttention.__new__( + ArtContextParallelCoreAttention + ) + torch.nn.Module.__init__(core) + core.softmax_offset = None + layer.self_attention = torch.nn.Module() + layer.self_attention.core_attention = core + provider = r.runtime.provider + provider.kv_channels = 256 + provider.ffn_hidden_size = 8192 + provider.linear_num_key_heads = 16 + provider.linear_num_value_heads = 32 + provider.linear_key_head_dim = 128 + provider.linear_value_head_dim = 128 + provider.params_dtype = torch.bfloat16 + r.runtime.model_support_handler = SimpleNamespace( + build_gdn_execution_spec=True, + context_parallel_workload_profile=lambda provider: None, + ) + monkeypatch.setattr(r, "_topology", lambda: ParallelTopology(tp=1, cp=2)) + return r + + +def _requests(lengths=(900, 700, 500)): + start, requests = 0, [] + for n in lengths: + requests.append( + ForwardInput( + input_tokens=torch.arange(start, start + n), hidden_states=True + ) + ) + start += n + return requests + + +def test_rank_layouts_are_the_executors_plans(monkeypatch): + from art.megatron.context_parallel.runtime import context_parallel_rank_layouts + from art.megatron.context_parallel.types import ( + ContextParallelConfig, + ParallelTopology, + ) + from art.megatron.prefix_tree_packing import prefix_tree_pack + + packed = prefix_tree_pack([r.input_tokens for r in _requests()], max_depth=1) + attention, gdn, plans = context_parallel_rank_layouts( + group_ids=packed.group_ids, + parent_ids=packed.parent_ids, + topology=ParallelTopology(tp=1, cp=2), + config=ContextParallelConfig(), + original_seq_len=int(packed.tokens.shape[1]), + build_gdn_execution_spec=True, + ) + total = int(packed.tokens.numel()) + assert sum(attention) == total and gdn is not None and sum(gdn) == total + # The ledger's attention rows and the mirror's own rows are the same count. + assert attention == tuple(plan.local_valid_lengths[0] for plan in plans) + + +def test_gated_plans_price_every_rank_from_its_stage_plan(monkeypatch): + from art.megatron.context_parallel.executor import retained_stage_record_bytes + from art.megatron.context_parallel.runtime import context_parallel_rank_layouts + from art.megatron.context_parallel.types import ParallelTopology + from art.megatron.training.microbatches import ( + _context_parallel_config_for_provider, + ) + + r = art_cp(rank(), monkeypatch) + plan = _plan_with(r, _requests()) + (layout,) = r._plan_group_layouts(plan) + assert sum(layout.attention_rows) == plan.packed_tokens + assert layout.gdn_rows is not None and sum(layout.gdn_rows) == plan.packed_tokens + # Each rank's retention is the executor mirror over that rank's own plan. + (group,) = plan.groups + _, _, rank_plans = context_parallel_rank_layouts( + group_ids=group.packed.group_ids, + parent_ids=group.packed.parent_ids, + topology=ParallelTopology(tp=1, cp=2), + config=_context_parallel_config_for_provider( + r.runtime.provider, r.device, r.runtime.model_support_handler + ), + original_seq_len=int(group.packed.tokens.shape[1]), + build_gdn_execution_spec=True, + ) + assert layout.attention_retained == tuple( + retained_stage_record_bytes( + rank_plan, + q_heads=16, + kv_heads=2, + head_dim=256, + value_head_dim=256, + element_size=2, + block_size=(128, 128), # the CPU device's flex block + ) + for rank_plan in rank_plans + ) + assert all(retained > 0 for retained in layout.attention_retained) + + +def test_softmax_offset_leaves_the_busiest_rank_floor(monkeypatch): + r = art_cp(rank(), monkeypatch) + plan = _plan_with(r, _requests()) + assert r._plan_group_layouts(plan) is not None + for layer in r.runtime.model[0].decoder.layers: + core = getattr(getattr(layer, "self_attention", None), "core_attention", None) + if core is not None: + core.softmax_offset = torch.zeros(16) + assert r._plan_group_layouts(plan) is None + + +@pytest.mark.parametrize( + "lengths", + [(2048, 1536, 1024, 512), (4099, 3, 5, 7), (1, 2, 3, 4, 5, 6, 7), (8191,)], +) +def test_split_lower_bound_stays_below_the_layout_cost(monkeypatch, lengths): + # Even and skewed CP splits, odd row counts, and a single long sequence. + r = art_cp(rank(), monkeypatch) + requests = _requests(lengths) + plan = _plan_with(r, requests) + lower = r._split_chunk_lower_cost( + requests, tuple(item.input_tokens for item in requests), checkpoint=Unset + ) + assert lower.required <= r._plan_cost(plan).required + # It prices even-share layouts at the least attention state. + assert r._layout_pricing_supported((1, 1, 2, 1), gradient_groups=True) + + +def _plan_with(r, requests): + return r._plan_flat_forward(requests) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 57b1d1aca..1cc118bc3 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -606,7 +606,10 @@ def held(rows): monkeypatch.setattr( rank, "_checkpoint_memory_floor", - lambda rows, refs=None, segments=0, routed_rows=None: (10**7, 10**6), + lambda rows, refs=None, segments=0, routed_rows=None, layouts=None: ( + 10**7, + 10**6, + ), ) grown = rank._plan_cost(plan) monkeypatch.setattr(rank, "_plan_hybridep_growth_bytes", lambda plan: 0) @@ -864,6 +867,48 @@ def test_hybridep_recompute_prices_fresh_dense_output_without_buffer_growth( assert not torch.cuda.is_initialized() +def test_hybridep_combine_extent_floors_the_layout_path( + hybrid_checkpoint_rank, monkeypatch +): + rank = hybrid_checkpoint_rank + groups = ((2, True),) + calls = [] + + def layout_floor(layers, refs, routed, layouts): + calls.append(layouts) + return 7, 11 + + def generic_floor(*args): + raise AssertionError("per-rank layouts must take the layout path") + + monkeypatch.setattr(rank, "_layout_checkpoint_floor", layout_floor) + monkeypatch.setattr(rank, "_generic_checkpoint_floor", generic_floor) + layouts = (object(),) + te = rank._te_workspace_growth_bytes() + # Two rows round up to four; the layout floor's small workspace loses. + assert rank._checkpoint_memory_floor(groups, layouts=layouts) == ( + 7, + 4 * 2048 * 2 + te, + ) + marker = torch.empty(0) + rank._pending_hybridep_graphs.append(weakref.ref(marker)) + rank._hybridep_rows_high_water = 218751 + assert rank._checkpoint_memory_floor(groups, layouts=layouts) == ( + 7, + 218752 * 2048 * 2 + te, + ) + # A larger stage already carries its TE growth; the combine floor adds none. + stage = 218752 * 2048 * 2 + te + 1 + + def larger_layout_floor(layers, refs, routed, layouts): + calls.append(layouts) + return 7, stage + + monkeypatch.setattr(rank, "_layout_checkpoint_floor", larger_layout_floor) + assert rank._checkpoint_memory_floor(groups, layouts=layouts) == (7, stage) + assert calls == [layouts, layouts, layouts] + + @pytest.mark.parametrize("reference", ["absent", "expired", "smaller"]) def test_hybridep_high_water_needs_a_live_larger_graph( hybrid_checkpoint_rank, reference