From 8df527d88fce043953ecdc0b9e837ed6e7328b33 Mon Sep 17 00:00:00 2001 From: xxyux <1650459510@qq.com> Date: Mon, 17 Aug 2026 16:34:22 +0800 Subject: [PATCH] [Performance] Balance the latent-MQA indexer's CP rows with a dual-chunk swap Under contiguous CP the latent-MQA indexer scores each rank's own row slice, so with a causal mask its cost grows with the rank index: measured at 256k/cp16, 2.2ms on cp0 against 66.8ms on cp15 per layer per pass (a 30x spread), which makes the slowest rank do 1.94x the average while every other rank waits for it at the next collective. Only the indexer has this shape -- attention already runs a fixed ``index_topk + window`` columns per row, and the indexer backward a fixed ``index_topk``, both measured flat across ranks. ``mqa_indexer_cp_mode="dualchunk_p2p"`` splits the global sequence into ``2 * cp_size`` chunks and has rank ``r`` score chunks ``(2r, 2*cp_size-1-2r)`` instead of its own ``(2r, 2r+1)``. The ids sum to ``2*cp_size-1`` everywhere, and a causal row's candidate count grows linearly with its global position, so the work is equal on every rank. Three properties keep the change contained: * Only the indexer's per-row inputs move. ``query`` and the layer output stay put, which at s_local=16384/h=64 is 2.3GB of traffic not incurred. * Rank ``r`` keeps the chunk contiguous CP already gave it and swaps the other with rank ``cp_size-1-r``, so the exchange is one point-to-point sendrecv per tensor rather than an all-to-all, and it is an involution -- the same call undoes it. * The permutation lives entirely inside ``paddle.no_grad()`` on detached inputs. The indexer gradient reaches the weights through the loss scaler applied to the *unpermuted* tensors, so nothing here has to be differentiable and the backward is byte-identical between the two modes. The cuDNN indexer call is split in two because the kernel learns where its rows sit globally from ``q_causal_offsets``, one scalar per batch, so a single affine map cannot describe two disjoint segments. Each chunk is internally contiguous and carries its own ``seq_offset`` -- the same shape as the query tiling the dense path already does. The results are concatenated back before the loss scaler sees them: applying it per chunk would halve ``target.shape[1]``, where its backward reads the row-count denominator, and double the gradient. The layer output is bit-identical either way, which is what the new CP test asserts. That equality is only possible because the per-(row, column) score is itself bit-identical: a CTA covers ``q_stage * q_tokens_per_tile`` rows and both offsets are multiples of that, so a given global row lands in the same aligned tile and the kernel walks the same key blocks in the same order. Gradients are held to the rel ~2e-3 the shared cuDNN DSA backward already carries for ``dkv`` (atomics), with the off-vs-off pair checked against the same bound so the tolerance is visibly the kernel's and not this switch's. --- src/paddlefleet/transformer/cp_utils.py | 89 +++++ .../transformer/mqa_latent_attention.py | 171 +++++++- .../transformer/transformer_config.py | 94 +++++ .../test_mqa_indexer_dualchunk_cp.py | 332 ++++++++++++++++ .../test_cp_balance_mode_dispatch.py | 152 +++++++ .../distributed/test_dualchunk_cp_utils.py | 375 ++++++++++++++++++ tests/test_configs.yaml | 4 + 7 files changed, 1205 insertions(+), 12 deletions(-) create mode 100644 tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py create mode 100644 tests/single_card_tests/ai_edited_test/distributed/test_dualchunk_cp_utils.py diff --git a/src/paddlefleet/transformer/cp_utils.py b/src/paddlefleet/transformer/cp_utils.py index 234106d565..2ddbe7ea4d 100644 --- a/src/paddlefleet/transformer/cp_utils.py +++ b/src/paddlefleet/transformer/cp_utils.py @@ -25,6 +25,7 @@ from __future__ import annotations import paddle +import paddle.distributed as dist from paddle import Tensor # =========================================================================== @@ -164,3 +165,91 @@ def build_causal_mask_cp( paddle.zeros([1], dtype="float32"), ) # [sq_local, n_comp] return mask.unsqueeze(0).expand([batch_size, -1, -1]) + + +# =========================================================================== +# Dual-chunk (zigzag) row swap — CP load balancing for the causal indexer +# =========================================================================== + + +def dualchunk_chunk_ids(cp_rank: int, cp_size: int) -> tuple[int, int]: + """The two ids, out of ``2 * cp_size`` equal chunks, this rank computes. + + Contiguous CP hands rank ``r`` chunks ``(2r, 2r+1)``. This layout keeps + ``2r`` and takes ``2*cp_size-1-2r`` instead, so the two ids sum to + ``2*cp_size-1`` on **every** rank. A causal row's candidate count grows + linearly with its global position, so a constant id sum means constant + indexer work per rank: at cp16 the 31x spread between rank 0 and rank 15 + collapses to 1x. + + Keeping ``2r`` — rather than the more familiar ``(r, 2*cp_size-1-r)`` + pairing, which balances just as well — is what reduces the exchange to a + single pairwise swap; see ``dualchunk_swap``. + """ + lo = 2 * cp_rank + return lo, 2 * cp_size - 1 - lo + + +def dualchunk_partner(cp_rank: int, cp_size: int) -> int: + """Rank holding the chunk this one wants, or ``-1`` when nothing moves. + + Rank ``r`` wants chunk ``2*cp_size-1-2r``, which contiguous CP placed on + rank ``cp_size-1-r``; that rank symmetrically wants ``2r+1`` from here, so + the pairing is the involution ``r <-> cp_size-1-r``. Returns ``-1`` for a + single-rank group and for the self-paired middle rank of an odd group. + """ + if cp_size <= 1: + return -1 + partner = cp_size - 1 - cp_rank + return -1 if partner == cp_rank else partner + + +def dualchunk_swap(x: Tensor, group, axis: int = 1) -> Tensor: + """Exchange the second half of ``x`` along ``axis`` with the partner rank. + + In: ``x`` is this rank's contiguous CP shard, i.e. chunks ``(2r, 2r+1)``. + Out: chunks ``(2r, 2*cp_size-1-2r)`` — see ``dualchunk_chunk_ids``. + + Both ends of a pair give up their odd chunk and want the other's, so this + is one symmetric pairwise exchange rather than a general all-to-all: each + rank moves ``x.shape[axis] // 2`` rows in each direction and talks to + exactly one peer. + + The map is an **involution**, so calling this again undoes it — which is + why one function serves both directions. + + Not differentiable, deliberately. The MQA indexer forward runs wholly + under ``paddle.no_grad()`` on detached inputs, and its gradient reaches the + weights through ``TileLangCSAIndexerLossAutoScaler`` applied to the + *unpermuted* tensors, so there is no gradient to route back through here. + """ + if group is None or group.nranks <= 1: + return x + partner = dualchunk_partner(group.rank, group.nranks) + if partner < 0: + return x + + n = int(x.shape[axis]) + if n % 2 != 0: + raise ValueError( + "dualchunk_swap needs an even extent on the swapped axis, got " + f"{n} on axis {axis} of shape {list(x.shape)}" + ) + + keep, give = paddle.split(x, 2, axis=axis) + give = give.contiguous() + take = paddle.empty(give.shape, dtype=give.dtype) + + # ``peer`` is a global rank: ``group.ranks`` maps the group-local index, + # matching context_parallel_utils.py:1252-1261. + peer = group.ranks[partner] + send_op = dist.P2POp(dist.isend, give, peer, group) + recv_op = dist.P2POp(dist.irecv, take, peer, group) + # Order by rank so a pair never issues two sends before either recv. NCCL's + # grouped p2p does not need this, but it costs one branch and removes a + # hang mode if the ops ever degrade to blocking. + ops = [send_op, recv_op] if group.rank < partner else [recv_op, send_op] + for task in dist.batch_isend_irecv(ops): + task.wait() + + return paddle.concat([keep, take], axis=axis).contiguous() diff --git a/src/paddlefleet/transformer/mqa_latent_attention.py b/src/paddlefleet/transformer/mqa_latent_attention.py index a74239a0b1..3bad6d0fd0 100644 --- a/src/paddlefleet/transformer/mqa_latent_attention.py +++ b/src/paddlefleet/transformer/mqa_latent_attention.py @@ -109,7 +109,11 @@ preprocess_index, ) from paddlefleet.process_groups_config import ProcessGroupCollection -from paddlefleet.transformer.cp_utils import all_gather_cp +from paddlefleet.transformer.cp_utils import ( + all_gather_cp, + dualchunk_chunk_ids, + dualchunk_swap, +) from paddlefleet.transformer.csa_attention import ( TileLangCSAIndexerLossAutoScaler, _build_window_topk_idxs_from_doc_bounds, @@ -755,6 +759,132 @@ def __init__( # Fused Triton epilogue for the analytic sink gradient. Only reachable # when ``softmax_offset`` exists; a sinkless layer has no sink gradient. self.sink_grad_fusion = getattr(config, "dsa_sink_grad_fusion", False) + # Row layout the indexer forward runs on; see ``mqa_indexer_cp_mode`` in + # ``transformer_config`` and ``_indexer_topk_dualchunk``. Gated on a real + # CP group: with ``cp_size == 1`` there is nothing to rebalance. + self.indexer_dualchunk = ( + getattr(config, "mqa_indexer_cp_mode", None) == "dualchunk_p2p" + and self.cp_enabled + and self.cp_size > 1 + ) + + def _chunk_valid_range( + self, meta, s_global, doc_start, doc_len, is_valid, offset, length + ): + """``valid_range [1, length, 2]`` for the rows at ``[offset, offset+length)``. + + Both the cached ``MQADocMeta`` and the eager fallback expose the same + ``(offset, length)`` slice of one global table, which is what lets the + dual-chunk layout ask for its two segments by chunk offset instead of + needing a second table. + """ + if meta is not None: + return meta.indexer_valid_range(self.window_size, offset, length)[0] + return self._indexer_valid_range( + s_global, doc_start, doc_len, is_valid, offset, length + )[0] + + def _dualchunk_valid_range( + self, meta, s_global, doc_start, doc_len, is_valid, s + ): + """``valid_range [1, s, 2]`` for this rank's rows in dual-chunk order. + + The same global table the contiguous path slices once, read instead as + the two chunks ``dualchunk_chunk_ids`` assigns, so it lines up + row-for-row with what ``dualchunk_swap`` produces and each kernel call + can be given its own ``seq_offset``. + + ``row_empty`` deliberately has no counterpart here: it is applied to the + results after they come back, i.e. in contiguous order. + """ + lo, hi = dualchunk_chunk_ids(self.cp_rank, self.cp_size) + m = s // 2 + return paddle.concat( + [ + self._chunk_valid_range( + meta, s_global, doc_start, doc_len, is_valid, c * m, m + ) + for c in (lo, hi) + ], + axis=1, + ) + + def _indexer_topk_dualchunk( + self, q_idx, w_idx, k_idx, topk, doc_lens_arg, vr_zz, need_loss + ): + """Indexer top-k on the balanced dual-chunk row layout. + + Rank ``r`` scores global chunks ``(2r, 2*cp_size-1-2r)`` out of + ``2*cp_size`` instead of its contiguous ``(2r, 2r+1)``. The ids sum to + ``2*cp_size-1`` on every rank, and a causal row's candidate count grows + linearly with its global position, so the two halves' work sums to a + constant: measured 31x between cp0 and cp15 at 256k/cp16 collapses to 1x. + + **Two calls, not one.** The kernel learns where its rows sit globally + from ``q_causal_offsets``, one scalar per batch + (``csa_indexer_fwd_cudnn.py``), so a single affine map cannot describe two + disjoint segments. Each chunk is internally contiguous, so each call + carries its own ``seq_offset``. This is the same shape as the query tiling + the dense path already does (it passes ``seq_offset + start`` per tile); + only the second segment's start jumps. + + ``vr_zz`` must be the ``valid_range`` of exactly these rows, in the same + order. A ``seq_offset`` that disagrees with it is silently wrong in one + direction: too large only wastes work (the extra columns are masked out of + the top-k anyway), too small writes ``-inf`` over legal candidates so they + can never be selected. + + The two results are concatenated and swapped back to the contiguous + layout before returning, so ``row_empty``, ``window_idxs``, attention, the + KL and ``TileLangCSAIndexerLossAutoScaler`` all keep seeing one unpermuted + ``[b, s_local, topk]`` tensor. Applying the loss PyLayer per chunk instead + would halve ``target.shape[1]``, which is where its backward reads the + row-count denominator, and double the gradient. + + The swaps are not differentiable and do not need to be: this whole path + runs under ``paddle.no_grad()`` on detached inputs, and the indexer + gradient reaches the weights through the loss scaler applied to the + *unpermuted* tensors. + """ + from paddlefleet.cudnn_ops.indexer.csa_indexer_fwd_cudnn import ( + cudnn_indexer_topk_fwd, + ) + + m = int(q_idx.shape[1]) // 2 + lo, hi = dualchunk_chunk_ids(self.cp_rank, self.cp_size) + + q_zz = dualchunk_swap(q_idx.detach(), self.cp_group, axis=1) + w_zz = dualchunk_swap(w_idx.detach(), self.cp_group, axis=1) + + def _chunk(sl, chunk_id): + return cudnn_indexer_topk_fwd( + q_zz[:, sl].contiguous(), + k_idx.detach(), + w_zz[:, sl].contiguous(), + ratio=1, + topk_effective=topk, + valid_range=vr_zz[:, sl], + doc_lens=doc_lens_arg, + seq_offset=chunk_id * m, + return_topk_scores=need_loss, + ) + + r_lo = _chunk(slice(0, m), lo) + r_hi = _chunk(slice(m, 2 * m), hi) + + selected = dualchunk_swap( + paddle.concat([r_lo[0], r_hi[0]], axis=1), self.cp_group, axis=1 + ) + scores_out = [] + if need_loss: + scores_out = [ + dualchunk_swap( + paddle.concat([r_lo[2], r_hi[2]], axis=1), + self.cp_group, + axis=1, + ) + ] + return selected, scores_out def _needs_indexer_loss(self) -> bool: """Whether this forward should build and attach the indexer loss. @@ -1889,6 +2019,18 @@ def _forward_sparse( valid_range, row_empty = self._indexer_valid_range( s_global, doc_start, doc_len, is_valid, position_offset, s ) + # ``row_empty`` stays contiguous -- it is applied after the results + # come back. ``vr_zz`` is the same table read in dual-chunk row + # order, for the two kernel calls: both sources slice a cached + # global table by ``(offset, length)``, so asking twice with the + # chunk offsets is exactly the row set the swap produces. + vr_zz = ( + self._dualchunk_valid_range( + meta, s_global, doc_start, doc_len, is_valid, s + ) + if self.indexer_dualchunk + else None + ) q_idx, k_idx, w_idx = self._indexer_projections( x, qr, position_offset, grad_enabled=need_loss @@ -1920,17 +2062,22 @@ def _forward_sparse( ) with paddle.no_grad(): - selected, _, *scores_out = cudnn_indexer_topk_fwd( - q_idx.detach(), - k_idx.detach(), - w_idx.detach(), - ratio=1, - topk_effective=topk, - valid_range=valid_range, - doc_lens=doc_lens_arg, - seq_offset=position_offset, - return_topk_scores=need_loss, - ) + if self.indexer_dualchunk: + selected, scores_out = self._indexer_topk_dualchunk( + q_idx, w_idx, k_idx, topk, doc_lens_arg, vr_zz, need_loss + ) + else: + selected, _, *scores_out = cudnn_indexer_topk_fwd( + q_idx.detach(), + k_idx.detach(), + w_idx.detach(), + ratio=1, + topk_effective=topk, + valid_range=valid_range, + doc_lens=doc_lens_arg, + seq_offset=position_offset, + return_topk_scores=need_loss, + ) topk_indices = paddle.where( row_empty, paddle.full_like(selected, -1), selected ) diff --git a/src/paddlefleet/transformer/transformer_config.py b/src/paddlefleet/transformer/transformer_config.py index 8e78690024..da518f9ed4 100644 --- a/src/paddlefleet/transformer/transformer_config.py +++ b/src/paddlefleet/transformer/transformer_config.py @@ -998,6 +998,36 @@ class TransformerConfig(ModelParallelConfig): layout, which is what makes mixing them safe. """ + mqa_indexer_cp_mode: str | None = None + """Row layout the latent-MQA indexer's forward runs on, under context parallel. + + ``None`` (default) inherits ``cp_balance_mode``: the indexer scores this + rank's own contiguous row slice, so under a causal mask its cost grows with + the rank index. Measured at 256k/cp16: 2.2ms on cp0 vs 66.8ms on cp15 per + layer per pass, i.e. the slowest rank does 1.94x the average and every other + rank waits for it at the next collective. + + ``"dualchunk_p2p"`` splits the global sequence into ``2 * cp_size`` chunks + and has rank ``r`` score chunks ``(2r, 2*cp_size-1-2r)`` instead of its own + ``(2r, 2r+1)``. The ids sum to ``2*cp_size-1`` on every rank, and a causal + row's candidate count grows linearly with its global position, so the work is + equal everywhere. Only the indexer's rows move -- attention is already + balanced at a fixed ``index_topk + window`` columns per row, so ``query`` and + the layer output never travel. Rank ``r`` keeps the chunk it already owns and + swaps the other with rank ``cp_size-1-r``, which reduces the exchange to a + single point-to-point sendrecv rather than an all-to-all. + + The global token layout is untouched: this is a layer-local permutation, + undone before the layer returns, so the HCA layers of the same model are + unaffected and ``cp_balance_mode`` must stay contiguous. + + Only the sparse training phase honours this, i.e. ``hybrid_mla_attention= + "mqa_dsa"`` with ``dsa_indexer_use_sparse_loss=True``. That is the only + phase whose indexer runs a top-k over per-rank rows; the warmup phase scores + the whole causal set through a different code path that does not permute + rows, so the other combinations are rejected rather than accepted-and-inert. + """ + v_head_dim: int | None = None """Dimension of the head in the V projection.""" @@ -2470,6 +2500,70 @@ def __post_init__(self): "-2 in csa_compress_ratios." ) + # only support mqa_indexer_cp_mode == dualchunk_p2p if not None + if self.mqa_indexer_cp_mode is not None: + if self.mqa_indexer_cp_mode != "dualchunk_p2p": + raise ValueError( + f"mqa_indexer_cp_mode={self.mqa_indexer_cp_mode!r} is " + "invalid. Must be None or 'dualchunk_p2p'." + ) + # The permutation is layer-local: the rows are swapped inside the + # MQA layer and swapped back before it returns, which is only sound + # while the *global* layout is the contiguous one the index tables + # are built against ("build over the global sequence, then take this + # rank's rows"). A dualchunk global layout would double-permute. + # + # Exact value rather than ``startswith("contiguous")``: under a real + # CP group ``MQALatentAttention.__init__`` accepts only + # ``contiguous_allgather``, so admitting ``contiguous_a2a`` here + # would pass config validation and then raise NotImplementedError at + # module construction -- a contract split between two layers. + if self.cp_balance_mode != "contiguous_allgather": + raise ValueError( + f"mqa_indexer_cp_mode={self.mqa_indexer_cp_mode!r} needs " + "cp_balance_mode='contiguous_allgather' (the only mode the " + "latent-MQA layer supports under context parallel), got " + f"{self.cp_balance_mode!r}." + ) + # Same membership test as hybrid_mla_cp_mode above: no latent-MQA + # layer means no indexer to rebalance. + if self.experimental_attention_variant != "dsv4_hybrid" or ( + -2 not in (self.csa_compress_ratios or ()) + ): + raise ValueError( + "mqa_indexer_cp_mode can only be set in dsv4_hybrid with " + "-2 in csa_compress_ratios." + ) + # ``MQALatentAttention`` reads the switch in ``_forward_sparse`` + # only. The other two phases reach a different indexer path that + # does not permute rows -- ``hybrid_mla_attention="mha"`` builds no + # latent-MQA layer at all, and ``"mqa_dsa"`` with + # ``dsa_indexer_use_sparse_loss=False`` is the warmup phase, whose + # KL spans the whole per-document causal set with no top-k anywhere. + # Accepting those would start successfully and silently deliver none + # of the rebalance this switch advertises. + if ( + self.hybrid_mla_attention != "mqa_dsa" + or not self.dsa_indexer_use_sparse_loss + ): + raise ValueError( + f"mqa_indexer_cp_mode={self.mqa_indexer_cp_mode!r} only " + "takes effect in the sparse training phase, which is " + "hybrid_mla_attention='mqa_dsa' together with " + "dsa_indexer_use_sparse_loss=True. This config has " + f"hybrid_mla_attention={self.hybrid_mla_attention!r} and " + "dsa_indexer_use_sparse_loss=" + f"{self.dsa_indexer_use_sparse_loss!r}, which runs an " + "indexer path that scores the full causal set and never " + "permutes rows, so the rebalance would be silently inert. " + "Drop mqa_indexer_cp_mode, or move to the sparse phase. " + "Note that train_indexer_only=True is the warmup phase by " + "definition and so cannot carry it either." + ) + # The two-chunks-per-rank split needs an even per-rank row count. + # ``TransformerConfig`` does not carry the sequence length, so that + # is checked at the swap itself (``cp_utils.dualchunk_swap``). + # separate_mtp_headloss validation. if self.separate_mtp_headloss: import warnings as _warnings diff --git a/tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py b/tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py new file mode 100644 index 0000000000..c7a337c159 --- /dev/null +++ b/tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py @@ -0,0 +1,332 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``mqa_indexer_cp_mode="dualchunk_p2p"`` must not change what the layer computes. + +The switch rebalances *which rows* each rank's indexer scores: rank ``r`` takes +global chunks ``(2r, 2*cp_size-1-2r)`` of ``2*cp_size`` instead of its own +``(2r, 2r+1)``, so the causal work is equal across ranks instead of growing 31x +from cp0 to cp15. The permutation is undone before the layer returns, so the +output must be unchanged. + +Two things make that testable as an equality rather than a tolerance: + +* the per-``(row, column)`` indexer score is bit-identical between the layouts. + A CTA covers ``q_stage * q_tokens_per_tile`` rows and both offsets are + multiples of that, so the same global row lands in the same aligned tile and + the kernel visits the same key blocks in the same order; +* the swap itself is pure data movement, and it is an involution, so the + round trip restores the contiguous order exactly. + +What is *not* guaranteed is the top-k tie-break: the two calls see different +``max_k`` on the THD path, so tied scores could in principle resolve +differently. Ties are vanishingly rare (a sum of ``index_n_heads`` fp32 +products), which is why this asserts equality and would flag a systematic +divergence rather than paper over it with a tolerance. +""" + +import types +import unittest + +import numpy as np +import paddle +import paddle.distributed as dist +from paddle.distributed import fleet + +from paddlefleet.transformer.cp_utils import ( + dualchunk_chunk_ids, + dualchunk_partner, + dualchunk_swap, +) +from tests.single_card_tests.transformer import hybrid_mla_utils as U + +CP_SIZE = None +CP_RANK = None +CP_GROUP = None + +S_GLOBAL = 512 + +_FA4_PIN = U._fa4_pin() + + +def setUpModule(): + global CP_SIZE, CP_RANK, CP_GROUP + if dist.get_world_size() < 2: + raise unittest.SkipTest( + "MQA dual-chunk tests require >= 2 GPUs (one peer to swap with)" + ) + _FA4_PIN.__enter__() + world = dist.get_world_size() + strategy = fleet.DistributedStrategy() + strategy.hybrid_configs = { + "dp_degree": 1, + "mp_degree": 1, + "pp_degree": 1, + "sharding_degree": world, + "sep_degree": 1, + "cp_degree": world, + "ep_degree": world, + "moe_sharding_degree": 1, + "order": [ + "sharding", + "moe_sharding", + "pp", + "sep", + "cp", + "dp", + "ep", + "mp", + ], + } + fleet.init(is_collective=True, strategy=strategy) + CP_GROUP = fleet.get_hybrid_communicate_group().get_context_parallel_group() + CP_RANK = CP_GROUP.rank + CP_SIZE = CP_GROUP.nranks + + +def tearDownModule(): + _FA4_PIN.__exit__(None, None, None) + + +def _build(dualchunk, loss_coeff=0.0, seed=7): + """CP layer with the rebalance off/on; identical weights either way.""" + cfg = U._create_mqa_config(mode="mqa_dsa", loss_coeff=loss_coeff) + cfg.dsa_indexer_use_sparse_loss = True + cfg.cp_balance_mode = "contiguous_allgather" + cfg.mqa_indexer_cp_mode = "dualchunk_p2p" if dualchunk else None + cfg.experimental_dataflow = True + cfg.pad_token_id = 0 + cfg.context_parallel_size = CP_SIZE + paddle.seed(seed) + return U._build_module( + cfg, + bf16=True, + pg_collection=types.SimpleNamespace(tp=None, cp=CP_GROUP), + ) + + +def _inputs(s_global, seed=1234): + """Global-length inputs, bitwise identical on every rank.""" + rng = np.random.RandomState(seed) + + def t(shape, scale): + return paddle.to_tensor( + (rng.standard_normal(shape) * scale).astype("float32") + ).cast("bfloat16") + + return { + "query": t([1, s_global, U.H, U.DK], 0.5), + "key": t([1, s_global, 1, U.DK], 0.5), + "w_v": t([U.DV, U.H, U.V_HEAD_DIM], 0.05), + "x": t([1, s_global, U.HIDDEN], 0.5), + "qr": t([1, s_global, U.Q_LORA], 0.5), + } + + +def _leaf(t): + out = t.clone().detach() + out.stop_gradient = False + return out + + +def _max_abs(a, b): + """``max |a - b|`` in fp32, so bf16 tensors compare without overflow.""" + return float((a.astype("float32") - b.astype("float32")).abs().max()) + + +def _rel(a, b): + """``max |a - b| / max |a|``, the shape the dkv tolerance is quoted in.""" + scale = float(a.astype("float32").abs().max()) + return _max_abs(a, b) / max(scale, 1e-12) + + +# The shared cuDNN DSA backward quotes rel ~2e-3 for ``dkv`` between identical +# calls; 5e-3 leaves headroom for the max-abs tail without hiding a real drift. +_GRAD_RTOL = 5e-3 + + +def _run(dualchunk, inp, row_end, loss_coeff): + """One forward+backward of this rank's slice. + + Returns ``(output, input_grads, param_grads)``. The parameter grads are the + only place the indexer loss shows up: ``_indexer_projections`` detaches + ``x``/``qr`` before the indexer touches them, so the KL gradient reaches the + indexer's own weights (``wq_b``, ``wk``, ``weights_proj``) and never the + caller's leaves. Asserting on the output alone would pass even if the + permutation corrupted ``topk_probs``, because the loss scaler forwards + ``output`` unchanged and only injects the gradient in backward. + """ + sl = S_GLOBAL // CP_SIZE + off = CP_RANK * sl + layer = _build(dualchunk, loss_coeff) + if dualchunk: + # Same weights on both sides: rebuilding from the reference state dict is + # unnecessary because ``paddle.seed`` is pinned in ``_build``, but assert + # the switch actually took so a silent no-op cannot pass this test. + assert layer.indexer_dualchunk, "dualchunk_p2p did not take effect" + else: + assert not layer.indexer_dualchunk + + r = { + "query": _leaf(inp["query"][:, off : off + sl]), + "key": _leaf(inp["key"][:, off : off + sl]), + "w_v": _leaf(inp["w_v"]), + "x": _leaf(inp["x"][:, off : off + sl]), + "qr": _leaf(inp["qr"][:, off : off + sl]), + } + out = layer( + r["query"], + r["key"], + None, + None, + attn_mask_startend_row_indices=row_end.clone(), + x=r["x"], + qr=r["qr"], + v_b_proj_weight=r["w_v"], + input_ids=None, + ) + out.sum().backward() + return ( + out, + {k: v.grad for k, v in r.items() if v.grad is not None}, + {n: p.grad for n, p in layer.named_parameters() if p.grad is not None}, + ) + + +class TestMqaIndexerDualChunkCp(unittest.TestCase): + def test_chunk_ids_balance_and_cover(self): + """Ids sum to a constant and every chunk has exactly one owner.""" + owner = {} + for r in range(CP_SIZE): + lo, hi = dualchunk_chunk_ids(r, CP_SIZE) + self.assertEqual(lo + hi, 2 * CP_SIZE - 1) + # The chunk this rank keeps is the one contiguous CP already gave + # it, which is what makes the exchange a single sendrecv. + self.assertEqual(lo // 2, r) + self.assertEqual(hi // 2, dualchunk_partner(r, CP_SIZE)) + for c in (lo, hi): + self.assertNotIn(c, owner) + owner[c] = r + self.assertEqual(sorted(owner), list(range(2 * CP_SIZE))) + + def test_swap_is_an_involution(self): + """Applying the swap twice restores the contiguous rows exactly.""" + sl = S_GLOBAL // CP_SIZE + base = CP_RANK * sl + rows = paddle.arange(base, base + sl, dtype="float32") + x = rows.reshape([1, sl, 1]).expand([1, sl, 4]).clone() + once = dualchunk_swap(x, CP_GROUP, axis=1) + twice = dualchunk_swap(once, CP_GROUP, axis=1) + self.assertEqual(float((twice - x).abs().max()), 0.0) + # And the single swap really did move the expected chunks. + m = sl // 2 + lo, hi = dualchunk_chunk_ids(CP_RANK, CP_SIZE) + want = paddle.concat( + [ + paddle.arange(lo * m, (lo + 1) * m, dtype="float32"), + paddle.arange(hi * m, (hi + 1) * m, dtype="float32"), + ] + ) + self.assertEqual(float((once[0, :, 0] - want).abs().max()), 0.0) + + def test_odd_extent_rejected(self): + """Two chunks per rank needs an even local row count.""" + odd = paddle.zeros([1, 2 * (S_GLOBAL // CP_SIZE) + 1, 2]) + with self.assertRaisesRegex(ValueError, "even extent"): + dualchunk_swap(odd, CP_GROUP, axis=1) + + @U._GPU + def test_output_unchanged(self): + """The rebalance is a pure scheduling change: identical output. + + The output being *exactly* equal is the whole contract, and it is the + strong assertion here: it can only hold if the two layouts selected the + same ``token_indices``, in the same order, from the same scores. + + Gradients then flow from identical inputs through identical code, so they + are only checked against the tolerance the shared cuDNN DSA backward + already carries: it accumulates ``dkv`` with atomics and is not + reproducible even between two identical ``off`` runs (documented as + rel ~2e-3 in ``fusions/mqa_sparse_attn.py``). Both the off-vs-off and the + off-vs-dualchunk pair are held to the same bound, which is what shows the + tolerance belongs to the kernel and not to this switch. + """ + inp = _inputs(S_GLOBAL) + row_end = U._row_end([S_GLOBAL], S_GLOBAL) + ref, gref, pref = _run(False, inp, row_end, loss_coeff=0.0) + ref2, gref2, pref2 = _run(False, inp, row_end, loss_coeff=0.0) + got, ggot, pgot = _run(True, inp, row_end, loss_coeff=0.0) + + self.assertEqual(ref.shape, got.shape) + self.assertEqual( + _max_abs(ref, got), 0.0, "dual-chunk changed the layer output" + ) + self._assert_grads_match( + (gref, gref2, ggot), (pref, pref2, pgot), "no-loss" + ) + + @U._GPU + def test_indexer_loss_gradients_unchanged(self): + """With the KL live, the indexer's own weight grads must also match. + + This is the case the output alone cannot cover: the loss scaler forwards + ``output`` unchanged and only injects the gradient in backward, so a + permutation that corrupted ``topk_probs`` or ``topk_indices`` would show + up nowhere else. ``_indexer_projections`` detaches ``x``/``qr``, so the KL + gradient lands on ``indexer.*`` parameters -- which is what gets compared + here, against the same off-vs-off baseline. + """ + inp = _inputs(S_GLOBAL, seed=99) + row_end = U._row_end([S_GLOBAL // 2, S_GLOBAL // 2], S_GLOBAL) + ref, gref, pref = _run(False, inp, row_end, loss_coeff=0.01) + ref2, gref2, pref2 = _run(False, inp, row_end, loss_coeff=0.01) + got, ggot, pgot = _run(True, inp, row_end, loss_coeff=0.01) + + # The KL has to actually reach the indexer, or every assertion below is + # comparing zeros and the test is vacuous. + indexer_grads = [ + n for n in pref if ".indexer." in n or n.startswith("indexer.") + ] + self.assertTrue(indexer_grads, "no indexer parameter received a grad") + self.assertTrue( + any( + float(pref[n].astype("float32").abs().max()) > 0.0 + for n in indexer_grads + ), + "indexer parameter grads are all zero: the KL never fired, so " + "this test would pass for any permutation", + ) + + self.assertEqual(_max_abs(ref, got), 0.0) + self._assert_grads_match( + (gref, gref2, ggot), (pref, pref2, pgot), "indexer-loss" + ) + + def _assert_grads_match(self, inputs, params, tag): + """Hold off-vs-off and off-vs-dualchunk to the same relative bound.""" + for kind, (a, a2, b) in (("input", inputs), ("param", params)): + self.assertEqual(sorted(a), sorted(b), f"{tag}: {kind} grad keys") + for name in a: + for label, other in (("off-vs-off", a2), ("dualchunk", b)): + rel = _rel(a[name], other[name]) + self.assertLessEqual( + rel, + _GRAD_RTOL, + f"{tag} {kind} grad of {name} ({label}): rel " + f"{rel:.2e} exceeds {_GRAD_RTOL:.0e}", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/single_card_tests/ai_edited_test/distributed/test_cp_balance_mode_dispatch.py b/tests/single_card_tests/ai_edited_test/distributed/test_cp_balance_mode_dispatch.py index 48e36ffc51..79281ca0ab 100644 --- a/tests/single_card_tests/ai_edited_test/distributed/test_cp_balance_mode_dispatch.py +++ b/tests/single_card_tests/ai_edited_test/distributed/test_cp_balance_mode_dispatch.py @@ -880,5 +880,157 @@ def test_requires_dsv4_hybrid_with_mla_layer(self): ) +class TestTransformerConfigMqaIndexerCpMode(unittest.TestCase): + """Tests for the mqa_indexer_cp_mode row-rebalance switch.""" + + # Same context hybrid_mla_cp_mode needs: a dsv4_hybrid whose + # csa_compress_ratios declares a latent-MQA layer (the -2), which is the + # only layer carrying an indexer to rebalance. + MQA_HYBRID = { + "experimental_attention_variant": "dsv4_hybrid", + "num_hidden_layers": 2, + "csa_compress_ratios": [-2, 128], + "hybrid_mla_q_lora_rank": 8, + "hybrid_mla_kv_lora_rank": 8, + "hybrid_mla_qk_nope_head_dim": 8, + "hybrid_mla_qk_rope_head_dim": 8, + "hybrid_mla_v_head_dim": 8, + "hybrid_mla_num_attention_heads": 4, + "hybrid_mla_num_key_value_heads": 4, + # The indexer shapes ``hybrid_mla_attention="mqa_dsa"`` requires in both + # DSA phases. ``index_head_dim`` is pinned to 128 by the cuDNN indexer + # and ``index_topk`` must be a multiple of 128. + "dsa_index_n_heads": 4, + "dsa_index_head_dim": 128, + "dsa_index_topk": 128, + } + + # The only phase that reads the switch: ``MQALatentAttention`` consults + # ``indexer_dualchunk`` in ``_forward_sparse`` and nowhere else, so the + # warmup phase would accept the config and rebalance nothing. + SPARSE_PHASE = { + "hybrid_mla_attention": "mqa_dsa", + "dsa_indexer_use_sparse_loss": True, + } + + @staticmethod + def _config(**kwargs): + from paddlefleet.transformer.transformer_config import ( + TransformerConfig, + ) + + return TransformerConfig(**kwargs) + + def test_default_is_none(self): + """Unset means 'score this rank's own contiguous rows'.""" + self.assertIsNone(self._config().mqa_indexer_cp_mode) + + def test_dualchunk_accepted(self): + """The rebalance is layer-local, so it rides on a contiguous global.""" + config = self._config( + cp_balance_mode="contiguous_allgather", + mqa_indexer_cp_mode="dualchunk_p2p", + **self.MQA_HYBRID, + **self.SPARSE_PHASE, + ) + self.assertEqual(config.mqa_indexer_cp_mode, "dualchunk_p2p") + + def test_invalid_value_rejected(self): + """One mode only; a typo must not silently fall back to the default.""" + for mode in ("nonexistent_mode", "dualchunk_allgather"): + with ( + self.subTest(mode=mode), + self.assertRaisesRegex(ValueError, "mqa_indexer_cp_mode"), + ): + self._config( + mqa_indexer_cp_mode=mode, + **self.MQA_HYBRID, + **self.SPARSE_PHASE, + ) + + def test_non_contiguous_global_layout_rejected(self): + """A dualchunk global layout would permute the same rows twice.""" + with self.assertRaisesRegex(ValueError, "contiguous_allgather"): + self._config( + cp_balance_mode="dualchunk_allgather", + mqa_indexer_cp_mode="dualchunk_p2p", + **self.MQA_HYBRID, + **self.SPARSE_PHASE, + ) + + def test_contiguous_a2a_rejected(self): + """Rejected here rather than later: the latent-MQA layer accepts only + contiguous_allgather under CP, so admitting the sibling contiguous mode + would move the failure to module construction.""" + with self.assertRaisesRegex(ValueError, "contiguous_allgather"): + self._config( + cp_balance_mode="contiguous_a2a", + mqa_indexer_cp_mode="dualchunk_p2p", + **self.MQA_HYBRID, + **self.SPARSE_PHASE, + ) + + def test_requires_dsv4_hybrid_with_mqa_layer(self): + """No latent-MQA layer means no indexer, so the setting is dead.""" + for kwargs in ( + {}, + {**self.MQA_HYBRID, "csa_compress_ratios": [128, 0]}, + ): + with ( + self.subTest(ratios=kwargs.get("csa_compress_ratios")), + self.assertRaisesRegex( + ValueError, "only be set in dsv4_hybrid" + ), + ): + self._config( + cp_balance_mode="contiguous_allgather", + mqa_indexer_cp_mode="dualchunk_p2p", + **kwargs, + ) + + def test_requires_sparse_phase(self): + """Only ``_forward_sparse`` reads the switch, so nothing else may set it. + + This is the accepted-but-inert combination the layer cannot catch: + ``hybrid_mla_attention`` defaulting to ``"mha"`` builds no latent-MQA + layer, ``"mqa_full_causal"`` has no indexer, and ``"mqa_dsa"`` with + ``dsa_indexer_use_sparse_loss=False`` (which is what + ``train_indexer_only=True`` pins) runs the warmup indexer path, whose KL + spans the whole causal set and never permutes rows. All three would + start successfully and rebalance nothing. + """ + for label, phase in ( + ("default mha", {}), + ("no indexer", {"hybrid_mla_attention": "mqa_full_causal"}), + ( + "warmup", + { + "hybrid_mla_attention": "mqa_dsa", + "dsa_indexer_use_sparse_loss": False, + }, + ), + ( + "train_indexer_only warmup", + { + "hybrid_mla_attention": "mqa_dsa", + "dsa_indexer_use_sparse_loss": False, + "train_indexer_only": True, + }, + ), + ): + with ( + self.subTest(phase=label), + self.assertRaisesRegex( + ValueError, "takes effect in the sparse training phase" + ), + ): + self._config( + cp_balance_mode="contiguous_allgather", + mqa_indexer_cp_mode="dualchunk_p2p", + **self.MQA_HYBRID, + **phase, + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/single_card_tests/ai_edited_test/distributed/test_dualchunk_cp_utils.py b/tests/single_card_tests/ai_edited_test/distributed/test_dualchunk_cp_utils.py new file mode 100644 index 0000000000..a8c5cd8c93 --- /dev/null +++ b/tests/single_card_tests/ai_edited_test/distributed/test_dualchunk_cp_utils.py @@ -0,0 +1,375 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-card tests for the dual-chunk CP row helpers in ``cp_utils``. + +``tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py`` covers +the same helpers against real NCCL. This file pins the parts that are pure +arithmetic or pure routing, so the chunk assignment and the peer/op layout are +checked on every single-card run and a regression there does not need two GPUs +to surface. + +The ``MQALatentAttention`` row-plumbing helpers are covered here too, with the +cuDNN indexer stubbed. The CP test exercises them for real but is gated on SM100 +kernels, so on every other box the two ``seq_offset`` values, the ``valid_range`` +slicing and the swap-back would otherwise go unchecked. +""" + +import types +import unittest +from unittest import mock + +import paddle + +from paddlefleet.transformer.cp_utils import ( + dualchunk_chunk_ids, + dualchunk_partner, + dualchunk_swap, +) + + +def _make_mock_group(nranks=4, rank=1): + """Mock process group; ``ranks`` maps group-local index to global rank.""" + group = mock.MagicMock() + group.nranks = nranks + group.rank = rank + # Offset so a test cannot pass by confusing the two rank spaces. + group.ranks = [100 + i for i in range(nranks)] + return group + + +class TestDualChunkChunkIds(unittest.TestCase): + def test_ids_partition_and_balance(self): + """Every chunk has one owner and the per-rank id sum is constant.""" + for cp_size in (1, 2, 4, 8, 16): + with self.subTest(cp_size=cp_size): + owner = {} + for rank in range(cp_size): + lo, hi = dualchunk_chunk_ids(rank, cp_size) + # Constant sum is the whole point: a causal row's candidate + # count grows linearly with its global position, so equal + # id sums mean equal work. + self.assertEqual(lo + hi, 2 * cp_size - 1) + # The kept chunk is the one contiguous CP already placed + # here, which is what makes the exchange one sendrecv. + self.assertEqual(lo, 2 * rank) + for c in (lo, hi): + self.assertNotIn(c, owner) + owner[c] = rank + self.assertEqual(sorted(owner), list(range(2 * cp_size))) + + +class TestDualChunkPartner(unittest.TestCase): + def test_involution_on_even_groups(self): + """``partner(partner(r)) == r``, so one function serves both ways.""" + for cp_size in (2, 4, 16): + for rank in range(cp_size): + p = dualchunk_partner(rank, cp_size) + self.assertEqual(dualchunk_partner(p, cp_size), rank) + self.assertEqual(p, cp_size - 1 - rank) + + def test_no_partner_cases(self): + """``-1`` means 'nothing to swap', not 'rank 0'.""" + self.assertEqual(dualchunk_partner(0, 1), -1) + # Odd group: the middle rank is its own partner, so it keeps both + # chunks rather than sending to itself. + self.assertEqual(dualchunk_partner(1, 3), -1) + self.assertEqual(dualchunk_partner(0, 3), 2) + + +class TestDualChunkSwap(unittest.TestCase): + def test_no_op_without_a_peer(self): + """Degenerate groups return the input untouched, not a copy-with-comm.""" + x = paddle.arange(8).reshape([1, 4, 2]).cast("float32") + self.assertIs(dualchunk_swap(x, None, axis=1), x) + self.assertIs(dualchunk_swap(x, _make_mock_group(1, 0), axis=1), x) + # Odd group's middle rank has no partner. + self.assertIs(dualchunk_swap(x, _make_mock_group(3, 1), axis=1), x) + + def test_odd_extent_rejected(self): + """Two chunks per rank needs an even count on the swapped axis.""" + x = paddle.zeros([1, 5, 2]) + with self.assertRaisesRegex(ValueError, "even extent"): + dualchunk_swap(x, _make_mock_group(4, 1), axis=1) + + def test_routing_and_op_order(self): + """The peer, the halves and the isend/irecv order, without NCCL. + + ``batch_isend_irecv`` is stubbed: this pins the routing contract (send + the odd chunk to ``cp_size-1-rank`` as a *global* rank, keep the even + one, and order the ops by rank) which is what a wrong peer or a + deadlock-prone ordering would break. + """ + cp_size = 4 + x = paddle.arange(16).reshape([1, 8, 2]).cast("float32") + for rank in (1, 2): + with self.subTest(rank=rank): + group = _make_mock_group(cp_size, rank) + captured = [] + + def _p2p_op(op, tensor, peer, grp): + captured.append((op, tensor, peer, grp)) + return (op, tensor, peer) + + with ( + mock.patch("paddle.distributed.P2POp", side_effect=_p2p_op), + mock.patch( + "paddle.distributed.batch_isend_irecv", + return_value=[mock.MagicMock()], + ) as batched, + ): + out = dualchunk_swap(x, group, axis=1) + + batched.assert_called_once() + (ops,) = batched.call_args[0] + self.assertEqual(len(ops), 2) + self.assertEqual(len(captured), 2) + + partner = dualchunk_partner(rank, cp_size) + peers = {c[2] for c in captured} + self.assertEqual( + peers, + {group.ranks[partner]}, + "peer must be the partner's *global* rank", + ) + # Lower rank sends first: harmless for NCCL's grouped p2p, but + # it removes a hang mode if the ops ever degrade to blocking. + # The submitted *list* order is what matters, not the order the + # two P2POp objects happened to be constructed in. + self.assertEqual( + ops[0][0] is paddle.distributed.isend, + rank < partner, + "ops must be submitted send-first only on the lower rank", + ) + + # The sent buffer is the second half; the first half survives + # untouched in the result. + sent = next( + t + for op, t, _, _ in captured + if op is paddle.distributed.isend + ) + self.assertEqual( + float((sent - x[:, 4:]).abs().max()), + 0.0, + "the odd chunk is the one that travels", + ) + self.assertEqual(out.shape, x.shape) + self.assertEqual( + float((out[:, :4] - x[:, :4]).abs().max()), + 0.0, + "the kept chunk must not move", + ) + + +def _half_flip(x, group, axis=1): + """Single-rank stand-in for ``dualchunk_swap``. + + An involution that leaves the first half alone and visibly reorders the + second, which is the only property the row plumbing relies on. A real swap + needs a peer rank; this one makes the swap-out/swap-back symmetry observable + on one card. + """ + m = int(x.shape[axis]) // 2 + keep, give = paddle.split(x, 2, axis=axis) + return paddle.concat([keep, paddle.flip(give, axis)], axis=axis) + + +class TestDualChunkValidRange(unittest.TestCase): + """``MQALatentAttention._dualchunk_valid_range`` reads two chunk offsets.""" + + def test_reads_the_two_chunk_offsets_in_order(self): + """Rows come from ``lo`` then ``hi``, each a length-``m`` global slice. + + The two ``(offset, length)`` pairs are what must line up with the two + ``seq_offset`` values the kernel calls get; asking with the local offset, + or in the wrong order, is silently wrong rather than a crash. + """ + from paddlefleet.transformer.mqa_latent_attention import ( + MQALatentAttention, + ) + + asked = [] + + def _chunk( + meta, s_global, doc_start, doc_len, is_valid, offset, length + ): + asked.append((offset, length)) + return paddle.full([1, length, 2], float(offset)) + + fake = types.SimpleNamespace( + cp_rank=1, cp_size=4, _chunk_valid_range=_chunk + ) + out = MQALatentAttention._dualchunk_valid_range( + fake, "meta", 32, "doc_start", "doc_len", "is_valid", 8 + ) + + # cp_rank=1 of 4 owns global chunks (2, 5); m = s // 2 = 4. + self.assertEqual(asked, [(8, 4), (20, 4)]) + self.assertEqual(out.shape, [1, 8, 2]) + self.assertEqual(float(out[0, 0, 0]), 8.0) + self.assertEqual(float(out[0, 4, 0]), 20.0) + + +class TestIndexerTopkDualChunk(unittest.TestCase): + """``MQALatentAttention._indexer_topk_dualchunk`` with the kernel stubbed. + + The real path is covered by ``test_mqa_indexer_dualchunk_cp.py``, but its + value tests need SM100 kernels. What is checked here is the plumbing that + would be wrong on any box: two calls with the two global ``seq_offset`` + values, each fed its own half of the *already dual-chunk-ordered* + ``valid_range``, and the results swapped back to contiguous rows. + """ + + CP_RANK, CP_SIZE, S, TOPK = 1, 4, 8, 4 + + def _run(self, need_loss): + from paddlefleet.transformer.mqa_latent_attention import ( + MQALatentAttention, + ) + + calls = [] + + def _kernel(q, k, w, **kw): + calls.append((q, w, kw)) + m = int(q.shape[1]) + # Rows tagged by their global position so the concat-then-swap-back + # is checkable rather than symmetric-by-accident. + rows = ( + kw["seq_offset"] + paddle.arange(m).cast("float32") + ).reshape([1, m, 1]) + out = paddle.expand(rows, [1, m, self.TOPK]) + return (out, None, out) if kw["return_topk_scores"] else (out, None) + + q = paddle.arange(self.S).reshape([1, self.S, 1, 1]).cast("float32") + w = paddle.arange(self.S).reshape([1, self.S, 1]).cast("float32") + vr = paddle.arange(2 * self.S).reshape([1, self.S, 2]).cast("float32") + fake = types.SimpleNamespace( + cp_rank=self.CP_RANK, cp_size=self.CP_SIZE, cp_group="grp" + ) + + with ( + mock.patch( + "paddlefleet.transformer.mqa_latent_attention.dualchunk_swap", + side_effect=_half_flip, + ) as swap, + mock.patch( + "paddlefleet.cudnn_ops.indexer.csa_indexer_fwd_cudnn" + ".cudnn_indexer_topk_fwd", + side_effect=_kernel, + ), + ): + selected, scores_out = MQALatentAttention._indexer_topk_dualchunk( + fake, + q, + w, + paddle.zeros([1, self.S, 1]), + self.TOPK, + "doc_lens", + vr, + need_loss, + ) + return calls, swap, selected, scores_out, q, w, vr + + def test_two_calls_carry_the_chunk_offsets_and_their_own_rows(self): + calls, swap, selected, scores_out, q, w, vr = self._run(False) + + self.assertEqual(len(calls), 2) + # chunk ids (2, 5) with m = 4. + self.assertEqual([c[2]["seq_offset"] for c in calls], [8, 20]) + + q_zz, w_zz = _half_flip(q, "grp"), _half_flip(w, "grp") + for i, sl in enumerate((slice(0, 4), slice(4, 8))): + q_seen, w_seen, kw = calls[i] + self.assertEqual(float((q_seen - q_zz[:, sl]).abs().max()), 0.0) + self.assertEqual(float((w_seen - w_zz[:, sl]).abs().max()), 0.0) + # ``vr_zz`` is built in dual-chunk order already, so it is sliced, + # never swapped -- swapping it here would double-permute the rows. + self.assertEqual( + float((kw["valid_range"] - vr[:, sl]).abs().max()), + 0.0, + "valid_range must be sliced, not swapped again", + ) + self.assertFalse(kw["return_topk_scores"]) + + # Swapped back: rows 8..11 stay, 20..23 come back reversed by the stub. + self.assertEqual( + [float(v) for v in selected[0, :, 0]], + [8, 9, 10, 11, 23, 22, 21, 20], + ) + # q, w out; selected back. Nothing else travels without the loss. + self.assertEqual(swap.call_count, 3) + self.assertEqual(scores_out, []) + + def test_scores_are_swapped_back_only_when_the_loss_needs_them(self): + calls, swap, selected, scores_out, *_ = self._run(True) + + self.assertTrue(all(c[2]["return_topk_scores"] for c in calls)) + self.assertEqual(swap.call_count, 4) + (scores,) = scores_out + # Same layout as ``selected``: the KL sees unpermuted contiguous rows. + self.assertEqual( + [float(v) for v in scores[0, :, 0]], + [8, 9, 10, 11, 23, 22, 21, 20], + ) + + +class TestChunkValidRange(unittest.TestCase): + """``MQALatentAttention._chunk_valid_range`` picks one of two row sources. + + Both expose the same ``(offset, length)`` slice of one global table, which is + what lets the dual-chunk layout ask for its two segments by chunk offset. The + delegation is what can silently break -- passing the caller's ``window`` + instead of ``self.window_size``, or forgetting to unwrap the + ``(valid_range, row_empty)`` pair -- so pin both paths here rather than only + exercising the eager one from the CP test. + """ + + @staticmethod + def _call(fake_self, meta, offset, length): + from paddlefleet.transformer.mqa_latent_attention import ( + MQALatentAttention, + ) + + return MQALatentAttention._chunk_valid_range( + fake_self, + meta, + 512, + "doc_start", + "doc_len", + "is_valid", + offset, + length, + ) + + def test_uses_meta_when_present(self): + meta = mock.MagicMock() + meta.indexer_valid_range.return_value = ("vr", "row_empty") + fake = types.SimpleNamespace(window_size=128) + self.assertEqual(self._call(fake, meta, 64, 32), "vr") + meta.indexer_valid_range.assert_called_once_with(128, 64, 32) + + def test_falls_back_to_eager_build(self): + eager = mock.MagicMock(return_value=("vr2", "row_empty")) + fake = types.SimpleNamespace( + window_size=128, _indexer_valid_range=eager + ) + self.assertEqual(self._call(fake, None, 64, 32), "vr2") + eager.assert_called_once_with( + 512, "doc_start", "doc_len", "is_valid", 64, 32 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_configs.yaml b/tests/test_configs.yaml index 0dccc54381..4b77c3bbbd 100644 --- a/tests/test_configs.yaml +++ b/tests/test_configs.yaml @@ -154,6 +154,10 @@ tests: - test_case: [tests/multi_card_tests/transformer/test_flash_mask_cp_a2a.py] products: - num_gpus: 2 + # transformer: latent-MQA indexer dual-chunk row rebalance (CP=2) + - test_case: [tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py] + products: + - num_gpus: 2 # transformer: CSA context-parallel, including cuDNN indexer docmask path (CP=2) - test_case: [tests/multi_card_tests/transformer/test_csa_attention_cp.py] products: