| """Sequence-packing support for the sparse-attention path. |
| |
| Ported from Megatron-Bridge (``deepseek_ocr2/modeling_deepseek_ocr_mcore.py``): |
| the two-stage warmup/sparse pipeline there runs with packing, and does so by |
| making every causal / local-window / block mask *segment aware* instead of |
| assuming a single contiguous sequence per row. |
| |
| Why this is needed |
| ------------------- |
| With packing, ms-swift concatenates several samples into one row and emits |
| per-sample-RESET ``position_ids`` (``[0,1,..,L1-1, 0,1,..,L2-1, ...]``; see |
| ``swift/template/base.py::packing_row``). The sparse code compares the query |
| ``position_ids`` against ``k_positions = arange(k_len)`` (the GLOBAL flat key |
| index) to build causal / local masks. With reset positions that comparison is |
| wrong (a sample-2 query at local pos 0 would be "allowed" to see only key 0, |
| which belongs to sample 1), so packing silently corrupts attention. |
| |
| Fix (mirrors Megatron-Bridge) |
| ----------------------------- |
| * Derive a per-token SEGMENT id from the ``position_ids`` resets. |
| * Evaluate causal / local windows in GLOBAL key space using the flat query |
| index (for a packed row with batch==1 the flat index *is* the global |
| position, and ``k_positions`` is already global). |
| * AND every keep-mask with a ``same_segment`` constraint so no query attends |
| across a sample boundary. |
| * Bypass the fused Triton kernels when packed (they assume ``position_ids == |
| arange`` and a single segment); the eager PyTorch paths below are made |
| segment aware instead. |
| |
| batch==1 is assumed for packed rows (``padding_free``/packing collapses the |
| batch dimension into the sequence). |
| |
| Block-grid alignment (sparse stage) |
| ----------------------------------- |
| Key blocks are cut on the FLAT (global) token index, so a sample that does not |
| start on a ``block_size`` boundary gets a block grid shifted by |
| ``start % block_size`` relative to single-sample inference. That is safe (the |
| segment masks above still prevent any cross-sample attention) but it is a |
| train/inference discrepancy in *how the keys are grouped*, plus the first block |
| of each sample straddles the previous one and gets trimmed. Run packing with |
| ``PACKING_ALIGN_SIZE=<block_size>`` (ms-swift ``Template._align_packing_row``) |
| to make every sample start on a block boundary; ``check_block_alignment`` warns |
| once when it is not the case. |
| """ |
| import torch |
|
|
|
|
| def is_packed(position_ids): |
| """A packed row restarts ``position_ids`` at 0 at each sample boundary, so a |
| 0 appears at some column >= 1. A plain (non-packed) sequence only has its |
| single 0 at column 0. Robust for training (no left padding under |
| padding_free).""" |
| if position_ids is None or position_ids.dim() != 2 or position_ids.shape[1] < 2: |
| return False |
| return bool((position_ids[:, 1:] == 0).any()) |
|
|
|
|
| def segment_ids(position_ids): |
| """[B, S] long: 0-based segment index of every token, from position_ids |
| resets. ``(position_ids == 0)`` marks each segment start; cumsum-1 numbers |
| them. For a non-packed row this is all zeros (single segment).""" |
| return (position_ids == 0).to(torch.long).cumsum(dim=1) - 1 |
|
|
|
|
| def packed_masks(position_ids, k_positions, q_index, local_window): |
| """All masks are built in GLOBAL key space and are [B, 1, C, Sk] bool. |
| |
| Args: |
| position_ids: [B, S] reset positions (training prefill: S == k_len). |
| k_positions: [Sk] global key indices == ``arange(k_len)``. |
| q_index: [C] global (flat) indices of the query rows in this |
| chunk (e.g. ``arange(c0, c1)`` or the strided |
| ``sel`` used by the KL loss). These ARE the global |
| query positions for a packed batch==1 row. |
| local_window: int inclusive local window length in tokens. |
| |
| Returns dict with: |
| token_future [B,1,C,Sk] key is strictly after the query (to mask out) |
| local_token [B,1,C,Sk] key in [gpos-local_window+1, gpos] |
| same_seg [B,1,C,Sk] key and query share a segment |
| gpos [B,C] global query positions |
| """ |
| device = position_ids.device |
| B = position_ids.shape[0] |
| Sk = k_positions.shape[0] |
| seg = segment_ids(position_ids) |
| seg_k = seg[:, :Sk] |
| q_index = q_index.to(device=device, dtype=torch.long) |
| seg_q = seg.index_select(1, q_index) |
| gpos = q_index.view(1, -1).expand(B, -1) |
|
|
| kp = k_positions.view(1, 1, 1, -1) |
| gpos_e = gpos[:, None, :, None] |
| token_future = kp > gpos_e |
| ls = (gpos - (local_window - 1)).clamp(min=0)[:, None, :, None] |
| local_token = (kp >= ls) & (kp <= gpos_e) |
| same_seg = seg_k[:, None, None, :] == seg_q[:, None, :, None] |
| return { |
| 'token_future': token_future, |
| 'local_token': local_token, |
| 'same_seg': same_seg, |
| 'gpos': gpos, |
| } |
|
|
|
|
| def segment_of(position_ids, token_indices): |
| """Gather the segment id for arbitrary GLOBAL token indices. |
| |
| Args: |
| position_ids: [B, S] |
| token_indices: [B, ...] long, global key indices (clamped by caller). |
| Returns tensor shaped like ``token_indices`` with the segment id of each. |
| """ |
| seg = segment_ids(position_ids) |
| flat = token_indices.reshape(token_indices.shape[0], -1) |
| out = seg.gather(1, flat.clamp(min=0, max=seg.shape[1] - 1)) |
| return out.reshape(token_indices.shape) |
|
|
|
|
| def segment_starts(position_ids): |
| """[N] long: GLOBAL start index of every segment in a packed row (batch 0). |
| ``start = flat_index - position_ids`` at the reset points.""" |
| return (position_ids[0] == 0).nonzero(as_tuple=True)[0] |
|
|
|
|
| _ALIGN_WARNED = False |
|
|
|
|
| def check_block_alignment(position_ids, block_size): |
| """one-shot warning when a packed row's samples do NOT start on a |
| ``block_size`` boundary. |
| |
| The sparse stage cuts its key blocks on the FLAT (global) token index, so a sample |
| starting at ``s`` gets a block grid shifted by ``s % block_size`` relative to |
| single-sample inference (which starts at 0), and its first block straddles the previous |
| sample (the segment mask then trims it, wasting part of the remote top-k budget). |
| Correctness is unaffected (no cross-sample leakage), but the train/inference block |
| partitioning differs. Set ``PACKING_ALIGN_SIZE=<block_size>`` so ms-swift's |
| ``packing_row`` pads every sample up to that multiple and the two grids coincide. |
| """ |
| global _ALIGN_WARNED |
| if _ALIGN_WARNED or position_ids is None or block_size <= 1: |
| return |
| starts = segment_starts(position_ids) |
| n_bad = int((starts % block_size != 0).sum()) |
| if n_bad: |
| _ALIGN_WARNED = True |
| print(f'[PackingAlign] {n_bad}/{starts.numel()} packed samples do not start on a ' |
| f'{block_size}-token block boundary -> block grid is shifted w.r.t. ' |
| f'single-sample inference. Set PACKING_ALIGN_SIZE={block_size} to align.', |
| flush=True) |
|
|
|
|
| def dense_segment_keep(position_ids, device): |
| """Boolean [B, 1, S, S] keep-mask for the DENSE (warmup) main branch: |
| True where a query may attend (same segment AND causal). Passed straight to |
| ``scaled_dot_product_attention(attn_mask=...)`` (bool: True == attend). |
| Mirrors the ``same_seg & causal`` SDPA mask of Megatron-Bridge's dense path. |
| |
| NOTE: this materialises a [B, 1, S, S] bool tensor (~1 GiB at S=32k, B=1). |
| A chunked-query SDPA (as in MB's ``modeling_deepseek_ocr_mcore`` cp>1 loop) |
| would avoid it and is a follow-up memory optimisation.""" |
| B, S = position_ids.shape |
| seg = segment_ids(position_ids) |
| pos = torch.arange(S, device=device) |
| causal = pos[None, :] <= pos[:, None] |
| same_seg = seg[:, :, None] == seg[:, None, :] |
| return (same_seg & causal[None])[:, None] |
|
|
