diff --git a/docs/source_en/Components/Training Middleware/Sequence-Parallel.md b/docs/source_en/Components/Training Middleware/Sequence-Parallel.md index d08b01d67..86b2ebc29 100644 --- a/docs/source_en/Components/Training Middleware/Sequence-Parallel.md +++ b/docs/source_en/Components/Training Middleware/Sequence-Parallel.md @@ -18,7 +18,6 @@ from twinkle.model.transformers.strategy.sequence_parallel import SequenceParall config = SequenceParallelConfig( enabled=True, # Enable sequence parallel ulysses_size=None, # Ulysses SP degree (auto-derived from DeviceMesh if None) - gather_logits=True, # Gather logits after forward for loss computation ) ``` @@ -41,16 +40,16 @@ device_mesh = DeviceMesh.from_sizes( 1. **Pad** — input sequences are padded to a length divisible by SP world size 2. **Split** — padded inputs are evenly split across SP ranks along the sequence dimension -3. **Distributed Attention** — FlashAttention2 is patched to perform Ulysses all-to-all communication before/after attention computation -4. **Gather** — after forward, logits are gathered back to full sequence length for loss computation +3. **Distributed Attention** — SDPA and FlashAttention 2/3 are patched to perform Ulysses all-to-all communication before/after attention computation +4. **Gather** — collect the present logps, entropies, values and requested logits before the unchanged full-sequence loss ## Supported Attention Backends | Backend | Status | |---------|--------| -| FlashAttention2 | Fully supported (including packed/padding-free sequences) | +| FlashAttention 2 / 3 | Ulysses and packed/padding-free paths; requires a compatible installed kernel | | SDPA | Supported (non-packed batches only) | -| Derived Ring Attention | Supported with FlashAttention2 only (`rp_world_size > 1`) | +| Derived Ring Attention | CUDA FlashAttention 2/3 only (`rp_world_size > 1`); incompatible platforms fail during initialization | ## Qwen3.5 Linear Attention @@ -63,6 +62,10 @@ For MoE models, SP automatically installs a forward hook that gathers router log ## Key Constraints - `num_key_value_heads` must be divisible by `ulysses_size` (for Ulysses) or use ring attention fallback -- Packed/padding-free batches require FlashAttention2 +- Packed/padding-free batches require FlashAttention 2 or 3 - Derived ring attention requires `batch_size == 1` (packed format) - `torch.distributed` must be initialized + +## Isolating per-model SP contexts + +Each SequenceParallelStrategy owns its own SequenceParallel context, which stores that model's communication group, ranks, and current batch state, so algorithms that use multiple models, such as DPO, can enable SP training at the same time. diff --git "a/docs/source_zh/\347\273\204\344\273\266/\350\256\255\347\273\203\344\270\255\351\227\264\344\273\266/\345\272\217\345\210\227\345\271\266\350\241\214.md" "b/docs/source_zh/\347\273\204\344\273\266/\350\256\255\347\273\203\344\270\255\351\227\264\344\273\266/\345\272\217\345\210\227\345\271\266\350\241\214.md" index e1d997188..dd0defe73 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\350\256\255\347\273\203\344\270\255\351\227\264\344\273\266/\345\272\217\345\210\227\345\271\266\350\241\214.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/\350\256\255\347\273\203\344\270\255\351\227\264\344\273\266/\345\272\217\345\210\227\345\271\266\350\241\214.md" @@ -18,7 +18,6 @@ from twinkle.model.transformers.strategy.sequence_parallel import SequenceParall config = SequenceParallelConfig( enabled=True, # 启用序列并行 ulysses_size=None, # Ulysses SP 并行度（若为 None 则从 DeviceMesh 自动推导） - gather_logits=True, # 前向后聚合 logits 用于损失计算 ) ``` @@ -41,16 +40,16 @@ device_mesh = DeviceMesh.from_sizes( 1. **填充** — 输入序列被填充到可被 SP 并行度整除的长度 2. **分割** — 填充后的输入沿序列维度均匀分配到各 SP rank -3. **分布式注意力** — FlashAttention2 被 patch 为在注意力计算前后执行 Ulysses all-to-all 通信 -4. **聚合** — 前向传播后，logits 被聚合回完整序列长度用于损失计算 +3. **分布式注意力** — SDPA 和 FlashAttention 2/3 被 patch 为在注意力计算前后执行 Ulysses all-to-all 通信 +4. **聚合** — 收集实际存在的 logps、entropies、values 及请求的 logits，再交给原有完整序列 loss ## 支持的注意力后端 | 后端 | 状态 | |------|------| -| FlashAttention2 | 完全支持（包括打包/padding-free 序列）| +| FlashAttention 2 / 3 | 支持 Ulysses 和 packed/padding-free 路径，需要兼容的 kernel | | SDPA | 支持（仅非打包批次）| -| 派生环形注意力 | 仅支持 FlashAttention2（`rp_world_size > 1`）| +| 派生环形注意力 | 需要 CUDA FlashAttention 2/3；不兼容平台在初始化时明确报错 | ## Qwen3.5 线性注意力 @@ -63,6 +62,10 @@ SP 自动检测 Qwen3.5 GatedDeltaNet 线性注意力层，并应用 `Qwen3_5Gat ## 关键约束 - `num_key_value_heads` 必须能被 `ulysses_size` 整除（Ulysses 模式），否则回退到环形注意力 -- 打包/padding-free 批次需要 FlashAttention2 +- 打包/padding-free 批次需要 FlashAttention 2 或 3 - 派生环形注意力要求 `batch_size == 1`（打包格式） - `torch.distributed` 必须已初始化 + +## 隔离每个模型的 SP 上下文 + +每个 SequenceParallelStrategy 持有自己的 SequenceParallel 上下文，保存该模型的通信组、rank 和当前 batch 状态，以支持 dpo 这种多个模型的算法同时开启 sp 训练。 diff --git a/src/twinkle/model/transformers/moe/router_replay.py b/src/twinkle/model/transformers/moe/router_replay.py index 9c9377d70..b31a569f5 100644 --- a/src/twinkle/model/transformers/moe/router_replay.py +++ b/src/twinkle/model/transformers/moe/router_replay.py @@ -185,8 +185,9 @@ def set_router_replay_data( # SP: slice full-sequence routed_experts (and the per-token replay_mask) to local SP rank tokens. # Both travel in one pad_and_split_inputs call so the mask is split identically to the indices and # stays token-aligned with them. - from ..strategy.sequence_parallel import sequence_parallel as sp - sp_world_size = getattr(sp, 'sp_world_size', None) or 1 + from ..strategy.sequence_parallel import _MODULE_TO_CTX + sp = _MODULE_TO_CTX.get(model) + sp_world_size = getattr(sp, 'world_size', None) or 1 if sp_world_size > 1: extra_split_values = [(routed_experts, 0, 1)] if replay_mask is not None: @@ -255,8 +256,9 @@ def get_router_replay_data(model: nn.Module, batch_size=1) -> torch.Tensor | Non routed_experts = routed_experts.reshape(batch_size, -1, num_layers, topk) # SP: all-gather local routing data across SP ranks along seq_len dim. - from ..strategy.sequence_parallel import sequence_parallel as sp - sp_world_size = getattr(sp, 'sp_world_size', None) or 1 + from ..strategy.sequence_parallel import _MODULE_TO_CTX + sp = _MODULE_TO_CTX.get(model) + sp_world_size = getattr(sp, 'world_size', None) or 1 if sp_world_size > 1: routed_experts = sp.gather(routed_experts, dim=1, position_ids=sp.real_position_ids) diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py index 0112d1aec..2426b49be 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py @@ -1,10 +1,14 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import math +import threading import torch import torch.distributed as dist from copy import copy from dataclasses import asdict, dataclass, is_dataclass -from functools import partial +from contextlib import contextmanager, nullcontext +from functools import wraps +from types import MethodType +from weakref import WeakKeyDictionary, WeakMethod, WeakSet from transformers import PreTrainedTokenizer from types import SimpleNamespace from typing import Any, Dict, List, Optional, Tuple, Union @@ -22,14 +26,61 @@ logger = get_logger() -def is_qwen3_vl(model): - mt = getattr(getattr(model, 'config', None), 'model_type', '') - return 'qwen3_vl' in mt +_MODULE_TO_CTX = WeakKeyDictionary() +_CTX_LOCAL = threading.local() + + +def _current_ctx(): + stack = getattr(_CTX_LOCAL, 'stack', ()) + return stack[-1] if stack else None + + +@contextmanager +def _use_context(sp): + stack = getattr(_CTX_LOCAL, 'stack', None) + if stack is None: + stack = _CTX_LOCAL.stack = [] + depth = len(stack) + stack.append(sp) + try: + yield + finally: + del stack[depth:] + + +def _attention_context(module): + sp = _MODULE_TO_CTX.get(module) + active = _current_ctx() + if sp is None: + if active is not None: + raise RuntimeError(f'SP model {active.model_name}: unregistered attention module ' + f'{type(module).__module__}.{type(module).__name__}') + return None + if module not in sp._text_modules: + return None # Explicitly registered local (e.g. vision) computation. + if active is not sp: + raise RuntimeError(f'SP model {sp.model_name}: attention {sp._module_names[module]} ' + 'called outside its owning context') + return sp + + +def _call_flash_attention(origin, *args, **kwargs): + position_ids = kwargs.get('position_ids') + if torch.is_tensor(position_ids): + # These are boundary metadata, not the positions used for RoPE. HF may + # infer sequence starts from position_ids.min(); SP's -1 padding would + # then hide every real sequence. Continue the preceding positions over + # padding while retaining all real positions and packed restarts. + indices = torch.arange(position_ids.shape[-1], device=position_ids.device) + last_valid = torch.where(position_ids != -1, indices, 0).cummax(dim=-1).values + continued = position_ids.gather(-1, last_valid).clamp_min(0) + indices - last_valid + kwargs['position_ids'] = torch.where(position_ids == -1, continued, position_ids) + return origin(*args, **kwargs) -def is_qwen3_omni(model): +def is_qwen3_vl(model): mt = getattr(getattr(model, 'config', None), 'model_type', '') - return 'qwen3_omni' in mt + return 'qwen3_vl' in mt def _call_create_causal_mask(fn, config, input_embeds, attention_mask, cache_position_or_past_key_values, *args, @@ -80,6 +131,11 @@ def __init__(self): self.num_heads = None self.causal_mask_func = None self.extra_kwargs = {} + self._dist_attns = {} + self._text_modules = WeakSet() + self._module_names = WeakKeyDictionary() + self._defer_inputs = False + self._packed_linear_supported = True @property def real_position_ids(self) -> torch.Tensor: @@ -102,7 +158,217 @@ def sp_rank(self) -> int: def rp_rank(self) -> int: return self._rp_rank - def _prepare_flash_attn(self, base_model: torch.nn.Module): + def _get_dist_attn(self, key: str, local_attention=None) -> DistributedAttention: + """Per-context DistributedAttention cache (one per attention entry point).""" + dist_attn = self._dist_attns.get(key) + if dist_attn is None: + dist_attn = DistributedAttention(local_attention, self) + self._dist_attns[key] = dist_attn + return dist_attn + + @staticmethod + def _sdpa_mask_impl(sp, origin_sdpa, origin_uses_cache_position, batch_size, q_length=None, kv_length=None, *args, + **kwargs): + q_length = q_length if q_length is not None else kwargs.pop('cache_position', None) + device = q_length.device if torch.is_tensor(q_length) else kwargs.pop('device', None) + if device is None and sp is not None: + device = sp.real_position_ids.device + + cache_position = None + if sp is not None and sp.world_size > 1: + padded_position_ids = sp.pad( + sp.real_position_ids[0], + padding_value=-1, + position_ids=sp.real_position_ids, + dim=0, + ) + global_length = padded_position_ids.shape[0] + if origin_uses_cache_position: + cache_position = torch.arange(0, global_length, device=device) + kv_length = global_length + else: + # Newer Transformers passes q_length/q_offset instead of cache_position. In SP training, + # create_causal_mask may still see the local shard length, so restore the global query length + # only for the no-cache prefill path; cache/sliding paths keep their upstream offsets. + q_offset = kwargs.get('q_offset', 0) + kv_offset = kwargs.get('kv_offset', 0) + no_cache_offsets = ((not torch.is_tensor(q_offset) and q_offset == 0) + and (not torch.is_tensor(kv_offset) and kv_offset == 0)) + if no_cache_offsets: + q_length = global_length + attention_mask = kwargs.get('attention_mask') + if attention_mask is not None and torch.is_tensor(attention_mask): + kv_length = attention_mask.shape[-1] + else: + kv_length = global_length + + # padding_mask may be an AttentionMask wrapper instead of Tensor + _pm = kwargs.get('padding_mask') + if _pm is not None and not isinstance(_pm, torch.Tensor): + if hasattr(_pm, 'to_tensor'): + kwargs['padding_mask'] = _pm.to_tensor() + elif hasattr(_pm, 'mask'): + kwargs['padding_mask'] = _pm.mask + + if origin_uses_cache_position: + if cache_position is None: + cache_position = q_length if torch.is_tensor(q_length) else torch.arange(q_length, device=device) + return origin_sdpa(batch_size, cache_position, kv_length, *args, **kwargs) + + return origin_sdpa(batch_size, q_length, kv_length, *args, device=device, **kwargs) + + @staticmethod + def _create_causal_mask_impl(sp, + config, + input_embeds, + attention_mask, + cache_position_or_past_key_values=None, + *args, + **kwargs): + from transformers import masking_utils + + if sp is None or sp.world_size == 1: + return _call_create_causal_mask( + masking_utils.origin_create_causal_mask, + config, + input_embeds, + attention_mask, + cache_position_or_past_key_values, + *args, + **kwargs, + ) + input_embeds = torch.ones( + (input_embeds.shape[0], input_embeds.shape[1] * sp.sp_world_size, input_embeds.shape[2]), + dtype=input_embeds.dtype, + device=input_embeds.device) + if has_signature_parameter(masking_utils.origin_create_causal_mask, 'cache_position'): + cache_position_or_past_key_values = torch.arange( + 0, + input_embeds.shape[1], + device=input_embeds.device, + ) + return _call_create_causal_mask( + masking_utils.origin_create_causal_mask, + config, + input_embeds, + attention_mask, + cache_position_or_past_key_values, + *args, + **kwargs, + ) + + def _local_flash_attn_impl(self, module, query_states, key_states, value_states, attention_mask, *args, **kwargs): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + dist_attn = self._get_dist_attn('flash_attention_2') + if dist_attn.local_attn is None: + + # Pass the current layer at call time. Capturing it in this cached + # closure would keep a weak-registry key alive through its context. + def _attention(query, key, value, *args, _sp_module, **kwargs): + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + if self.rp_world_size > 1: + from .zigzag_ring_attn import zigzag_ring_flash_attn_varlen_func + + position_ids = kwargs.get('position_ids') + if position_ids is None: + position_ids = self.real_position_ids + position_ids = self._extract_real_position_ids(position_ids) + cu_seqlens = get_cu_seqlens_from_position_ids(position_ids).to(torch.int32) + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + position_ids = self._split_packed(position_ids, cu_seqlens, dim=-1) + mask = position_ids != -1 + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + query, key, value = self._mask_qkv(query, key, value, mask) + return zigzag_ring_flash_attn_varlen_func( + query, + key, + value, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + dropout_p=kwargs.get('dropout', 0.0), + softmax_scale=kwargs.get('scaling'), + causal=_sp_module.is_causal, + window_size=kwargs.get('sliding_window') or (-1, -1), + group=self._rp_group, + backend=self.attn_implementation, + ) + elif self.extra_kwargs.get('padding_free', False) or 'cu_seq_lens_q' in kwargs: + position_ids = kwargs.get('position_ids') + if position_ids is None: + position_ids = self.real_position_ids + if position_ids is None: + raise ValueError('Packed/varlen flash_attention_2 requires position_ids to derive ' + 'cu_seq_lens_q.') + position_ids = self._extract_real_position_ids(position_ids) + position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids) + cu_seqlens = get_cu_seqlens_from_position_ids(position_ids).to( + dtype=torch.int32, + device=query.device, + ) + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + total_tokens = int(cu_seqlens[-1].item()) + if query.shape[2] != total_tokens: + raise ValueError('Packed/varlen flash_attention_2 expects query sequence length to match ' + f'cu_seqlens total tokens, got query_seq_len={query.shape[2]} ' + f'and cu_seqlens_total={total_tokens}.') + kwargs['cu_seq_lens_q'] = cu_seqlens + kwargs['cu_seq_lens_k'] = cu_seqlens + kwargs['max_length_q'] = max_seqlen + kwargs['max_length_k'] = max_seqlen + if self.extra_kwargs.get('padding_free', False) and len(args) > 0: + args = (None, *args[1:]) + backend = getattr(getattr(_sp_module, 'config', None), '_attn_implementation', None) + origin = ALL_ATTENTION_FUNCTIONS.get(f'{backend}_origin', + ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin']) + return _call_flash_attention(origin, _sp_module, query, key, value, *args, **kwargs)[0] + + dist_attn.local_attn = _attention + + return dist_attn( + query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask, + *args, _sp_module=module, **kwargs), None + + def _local_sdpa_attn_impl(self, module, query_states, key_states, value_states, attention_mask, *args, **kwargs): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + # Policy: packed (PackingDataset/padding-free) batches require FlashAttention2 varlen/packed semantics. + # SDPA does not have a native packed/varlen interface; supporting packed batches would require building a + # large block-diagonal causal mask (slow / memory heavy). + if self.extra_kwargs.get('padding_free', False): + raise RuntimeError('SequenceParallel: detected padding_free/packed batch. ' + 'SDPA backend is not supported for padding_free/packed batches; please use ' + 'flash_attention_2.') + dist_attn = self._get_dist_attn('sdpa') + if dist_attn.local_attn is None: + + # The cache is per context, not per layer; do not bind the first + # layer or retain it through the global weak-key registry. + def _attention(query, key, value, *args, _sp_module, **kwargs): + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + if self.rp_world_size > 1: + raise NotImplementedError('SDPA does not support derived ring attention.') + return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](_sp_module, query, key, value, *args, **kwargs)[0] + + dist_attn.local_attn = _attention + return dist_attn( + query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask, + *args, _sp_module=module, **kwargs), None + + def _prepare_flash_attn(self): + """Install the global SP-aware attention/mask wrappers (once per process). + + The installed wrappers carry no per-model state; they resolve the owning + SequenceParallel context at call time via ``_MODULE_TO_CTX`` (attention + functions) or the thread-local ctx stack (mask helpers), so multiple + SP-enabled models (and non-SP models) can coexist in one process. + """ try: from transformers import masking_utils @@ -110,54 +376,8 @@ def _prepare_flash_attn(self, base_model: torch.nn.Module): origin_uses_cache_position = has_signature_parameter(origin_sdpa, 'cache_position') def sdpa_mask(batch_size, q_length=None, kv_length=None, *args, **kwargs): - q_length = q_length if q_length is not None else kwargs.pop('cache_position', None) - device = q_length.device if torch.is_tensor(q_length) else kwargs.pop('device', None) - if device is None: - device = self.real_position_ids.device - - cache_position = None - if self.world_size > 1: - padded_position_ids = self.pad( - self.real_position_ids[0], - padding_value=-1, - position_ids=self.real_position_ids, - dim=0, - ) - global_length = padded_position_ids.shape[0] - if origin_uses_cache_position: - cache_position = torch.arange(0, global_length, device=device) - kv_length = global_length - else: - # Newer Transformers passes q_length/q_offset instead of cache_position. In SP training, - # create_causal_mask may still see the local shard length, so restore the global query length - # only for the no-cache prefill path; cache/sliding paths keep their upstream offsets. - q_offset = kwargs.get('q_offset', 0) - kv_offset = kwargs.get('kv_offset', 0) - no_cache_offsets = ((not torch.is_tensor(q_offset) and q_offset == 0) - and (not torch.is_tensor(kv_offset) and kv_offset == 0)) - if no_cache_offsets: - q_length = global_length - attention_mask = kwargs.get('attention_mask') - if attention_mask is not None and torch.is_tensor(attention_mask): - kv_length = attention_mask.shape[-1] - else: - kv_length = global_length - - # padding_mask may be an AttentionMask wrapper instead of Tensor - _pm = kwargs.get('padding_mask') - if _pm is not None and not isinstance(_pm, torch.Tensor): - if hasattr(_pm, 'to_tensor'): - kwargs['padding_mask'] = _pm.to_tensor() - elif hasattr(_pm, 'mask'): - kwargs['padding_mask'] = _pm.mask - - if origin_uses_cache_position: - if cache_position is None: - cache_position = q_length if torch.is_tensor(q_length) else torch.arange( - q_length, device=device) - return origin_sdpa(batch_size, cache_position, kv_length, *args, **kwargs) - - return origin_sdpa(batch_size, q_length, kv_length, *args, device=device, **kwargs) + return SequenceParallel._sdpa_mask_impl(_current_ctx(), origin_sdpa, origin_uses_cache_position, + batch_size, q_length, kv_length, *args, **kwargs) masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'] = origin_sdpa masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask @@ -168,48 +388,14 @@ def create_causal_mask(config, cache_position_or_past_key_values=None, *args, **kwargs): - if self.world_size == 1: - return _call_create_causal_mask( - masking_utils.origin_create_causal_mask, - config, - input_embeds, - attention_mask, - cache_position_or_past_key_values, - *args, - **kwargs, - ) - input_embeds = torch.ones( - (input_embeds.shape[0], input_embeds.shape[1] * self.sp_world_size, input_embeds.shape[2]), - dtype=input_embeds.dtype, - device=input_embeds.device) - if has_signature_parameter(masking_utils.origin_create_causal_mask, 'cache_position'): - cache_position_or_past_key_values = torch.arange( - 0, - input_embeds.shape[1], - device=input_embeds.device, - ) - return _call_create_causal_mask( - masking_utils.origin_create_causal_mask, - config, - input_embeds, - attention_mask, - cache_position_or_past_key_values, - *args, - **kwargs, - ) + return SequenceParallel._create_causal_mask_impl(_current_ctx(), config, input_embeds, attention_mask, + cache_position_or_past_key_values, *args, **kwargs) masking_utils.origin_create_causal_mask = masking_utils.create_causal_mask masking_utils.create_causal_mask = create_causal_mask except ImportError: pass - if hasattr(base_model, 'language_model'): - text_model = base_model.language_model - else: - text_model = base_model - # Precomputed once: the membership test runs per attention layer per forward step. - text_model_classes = {module.__class__ for module in text_model.modules()} - from transformers.modeling_flash_attention_utils import is_flash_attn_available if is_flash_attn_available(): # TODO this works for multi-modal models like qwen2.5-vl @@ -217,124 +403,46 @@ def create_causal_mask(config, # more work for maintaining. from transformers import modeling_flash_attention_utils from transformers.modeling_flash_attention_utils import _flash_attention_forward - _distributed_flash_attention = DistributedAttention(_flash_attention_forward, self) modeling_flash_attention_utils._flash_attention_forward_origin = _flash_attention_forward + def local_flash_attention_forward(*args, **kwargs): + return _call_flash_attention(_flash_attention_forward, *args, **kwargs) + def flash_attention_forward(query_states: torch.Tensor, key_states: torch.Tensor, value_states: torch.Tensor, attention_mask: Optional[torch.Tensor], q_len, *args, **kwargs): - if self.world_size == 1: + sp = _current_ctx() + if sp is None or sp.world_size == 1: return _flash_attention_forward(query_states, key_states, value_states, attention_mask, q_len, *args, **kwargs) - return _distributed_flash_attention(query_states, key_states, value_states, attention_mask, - q_len * self.sp_world_size, *args, **kwargs) + dist_attn = sp._get_dist_attn('flash_attention_utils', local_flash_attention_forward) + return dist_attn(query_states, key_states, value_states, attention_mask, q_len * sp.sp_world_size, + *args, **kwargs) modeling_flash_attention_utils._flash_attention_forward = flash_attention_forward from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS def local_flash_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args, - dist_attn, **kwargs): - if self.world_size == 1 or module.__class__ not in text_model_classes: - return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query_states, key_states, - value_states, attention_mask, *args, - **kwargs) - if dist_attn.local_attn is None: - - def _attention(query, key, value, *args, **kwargs): - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - if self.rp_world_size > 1: - from .zigzag_ring_attn import zigzag_ring_flash_attn_varlen_func - - position_ids = kwargs.get('position_ids') - if position_ids is None: - position_ids = self.real_position_ids - position_ids = self._extract_real_position_ids(position_ids) - cu_seqlens = get_cu_seqlens_from_position_ids(position_ids).to(torch.int32) - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - position_ids = self._split_packed(position_ids, cu_seqlens, dim=-1) - mask = position_ids != -1 - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - query, key, value = self._mask_qkv(query, key, value, mask) - return zigzag_ring_flash_attn_varlen_func( - query, - key, - value, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - dropout_p=kwargs.get('dropout', 0.0), - softmax_scale=kwargs.get('scaling'), - causal=module.is_causal, - window_size=kwargs.get('sliding_window') or (-1, -1), - group=self._rp_group, - backend=self.attn_implementation, - ) - elif self.extra_kwargs.get('padding_free', False) or 'cu_seq_lens_q' in kwargs: - position_ids = kwargs.get('position_ids') - if position_ids is None: - position_ids = self.real_position_ids - if position_ids is None: - raise ValueError('Packed/varlen flash_attention_2 requires position_ids to derive ' - 'cu_seq_lens_q.') - position_ids = self._extract_real_position_ids(position_ids) - position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids) - cu_seqlens = get_cu_seqlens_from_position_ids(position_ids).to( - dtype=torch.int32, - device=query.device, - ) - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - total_tokens = int(cu_seqlens[-1].item()) - if query.shape[2] != total_tokens: - raise ValueError('Packed/varlen flash_attention_2 expects query sequence length to match ' - f'cu_seqlens total tokens, got query_seq_len={query.shape[2]} ' - f'and cu_seqlens_total={total_tokens}.') - kwargs['cu_seq_lens_q'] = cu_seqlens - kwargs['cu_seq_lens_k'] = cu_seqlens - kwargs['max_length_q'] = max_seqlen - kwargs['max_length_k'] = max_seqlen - if self.extra_kwargs.get('padding_free', False) and len(args) > 0: - args = (None, *args[1:]) - return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query, key, value, *args, - **kwargs)[0] - - dist_attn.local_attn = _attention - - return dist_attn( - query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask, - *args, **kwargs), None + **kwargs): + sp = _attention_context(module) + if sp is None or sp.world_size == 1: + backend = getattr(getattr(module, 'config', None), '_attn_implementation', None) + origin = ALL_ATTENTION_FUNCTIONS.get(f'{backend}_origin', + ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin']) + return origin(module, query_states, key_states, value_states, attention_mask, *args, **kwargs) + return sp._local_flash_attn_impl(module, query_states, key_states, value_states, attention_mask, *args, + **kwargs) def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args, - dist_attn, **kwargs): - # Bypass SP logic when world_size == 1 (SP disabled) or module not in text_model - if self.world_size == 1 or module.__class__ not in text_model_classes: + **kwargs): + sp = _attention_context(module) + if sp is None or sp.world_size == 1: return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query_states, key_states, value_states, attention_mask, *args, **kwargs) - # Policy: packed (PackingDataset/padding-free) batches require FlashAttention2 varlen/packed semantics. - # SDPA does not have a native packed/varlen interface; supporting packed batches would require building a - # large block-diagonal causal mask (slow / memory heavy). - if self.extra_kwargs.get('padding_free', False): - raise RuntimeError( - 'SequenceParallel: detected padding_free/packed batch. ' - 'SDPA backend is not supported for padding_free/packed batches; please use flash_attention_2.') - if dist_attn.local_attn is None: - - def _attention(query, key, value, *args, **kwargs): - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - if self.rp_world_size > 1: - raise NotImplementedError('SDPA does not support derived ring attention.') - return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query, key, value, *args, **kwargs)[0] - - dist_attn.local_attn = _attention - return dist_attn( - query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask, - *args, **kwargs), None + return sp._local_sdpa_attn_impl(module, query_states, key_states, value_states, attention_mask, *args, + **kwargs) # `flash_attention_2/3` all hold the same `flash_attention_forward` object in transformers; the # concrete kernel is picked at call time from `module.config._attn_implementation`. So the SP @@ -345,45 +453,111 @@ def _attention(query, key, value, *args, **kwargs): ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'] = ALL_ATTENTION_FUNCTIONS['flash_attention_2'] ALL_ATTENTION_FUNCTIONS['sdpa_origin'] = ALL_ATTENTION_FUNCTIONS['sdpa'] for _impl in FLASH_ATTENTION_IMPLS: - ALL_ATTENTION_FUNCTIONS[_impl] = partial( - local_flash_attn, dist_attn=DistributedAttention(None, self)) - ALL_ATTENTION_FUNCTIONS['sdpa'] = partial(local_sdpa_attn, dist_attn=DistributedAttention(None, self)) + if _impl not in ALL_ATTENTION_FUNCTIONS: + continue + ALL_ATTENTION_FUNCTIONS[f'{_impl}_origin'] = ALL_ATTENTION_FUNCTIONS[_impl] + ALL_ATTENTION_FUNCTIONS[_impl] = local_flash_attn + ALL_ATTENTION_FUNCTIONS['sdpa'] = local_sdpa_attn logger.info_once(f'[SequenceParallel] attention wrappers registered: {list(FLASH_ATTENTION_IMPLS)} -> ' 'local_flash_attn, sdpa -> local_sdpa_attn (SP-aware)') - def _prepare_forward_hook(self, base_model: torch.nn.Module): - - def pre_forward_split_hook(_self, args, kwargs): - if self.world_size == 1: - return args, kwargs - # Pad to multiple of SP size and split inputs per SP rank before forward. - input_ids = kwargs.get('input_ids', None) - inputs_embeds = kwargs.get('inputs_embeds', None) - position_ids = kwargs['position_ids'] - real_position_ids = self._extract_real_position_ids(position_ids) - attention_mask = kwargs.get('attention_mask', None) - cache_position = kwargs.get('cache_position', None) - if hasattr(_self, 'language_model'): - embed_tokens = getattr(_self.language_model, 'embed_tokens', None) - else: - embed_tokens = getattr(_self, 'embed_tokens', None) - input_ids, inputs_embeds, _, position_ids, attention_mask, _, _ = self.pad_and_split_inputs( - input_ids, - inputs_embeds, - None, - position_ids, - attention_mask, - None, - embed_tokens=embed_tokens, - real_position_ids=real_position_ids, - cache_position=cache_position) - kwargs['input_ids'] = input_ids - kwargs['inputs_embeds'] = inputs_embeds - kwargs['position_ids'] = position_ids - kwargs['attention_mask'] = attention_mask - return args, kwargs - - base_model.register_forward_pre_hook(pre_forward_split_hook, with_kwargs=True) + @property + def causal_mask_func(self): + func = self._causal_mask_func + return func() if isinstance(func, WeakMethod) else func + + @causal_mask_func.setter + def causal_mask_func(self, func): + # The global weak-key registry owns this context. A bound method must + # not keep the owning model (and therefore its weak key) alive. + self._causal_mask_func = WeakMethod(func) if getattr(func, '__self__', None) is not None else func + + @contextmanager + def context(self, state=None): + previous = self.extra_kwargs + if state is not None: + self.extra_kwargs = state + try: + with _use_context(self): + yield + finally: + self.extra_kwargs = previous + + def _prepare_forward_boundaries(self, model, text_model): + self.model_name = type(model).__name__ + text_modules = tuple(text_model.modules()) + named_modules = tuple(model.named_modules()) + self._defer_inputs = any(getattr(getattr(module, 'config', None), 'vision_config', None) is not None + for _, module in named_modules) + if not any('Attention' in type(module).__name__ or 'GatedDeltaNet' in type(module).__name__ + for module in text_modules): + raise NotImplementedError(f'SP {self.model_name}: cannot identify the text attention modules') + for module in text_modules: + if 'GatedDeltaNet' in type(module).__name__ and any( + getattr(module, name) % self.sp_world_size for name in ('num_k_heads', 'num_v_heads')): + raise NotImplementedError(f'SP {self.model_name}: incompatible GDN head counts') + vision = [module for name, module in named_modules + if name.rsplit('.', 1)[-1] in ('visual', 'vision_tower', 'vision_model')] + if self._defer_inputs and (not vision or any(module in text_modules for module in vision)): + raise NotImplementedError(f'SP {self.model_name}: cannot isolate local vision from text') + self._text_modules.update(text_modules) + for name, module in named_modules: + _MODULE_TO_CTX[module] = self + self._module_names[module] = name or '' + + def wrap_forward(module, local=False): + origin = WeakMethod(module.forward) + @wraps(module.forward.__func__) + def forward(_module, *args, **kwargs): + if local: + with _use_context(None): + return origin()(*args, **kwargs) + if args: + raise NotImplementedError(f'{self.model_name}: SP text entry requires keyword inputs') + kwargs = self._split_model_inputs(dict(kwargs)) + return origin()(**kwargs) + module.forward = MethodType(forward, module) + + if self._defer_inputs: + wrap_forward(text_model) + for module in vision: + wrap_forward(module, local=True) + for module in text_model.modules(): + original = getattr(module, '_gradient_checkpointing_func', None) + if original is None: + continue + def checkpoint(function, *args, _checkpoint=original, **kwargs): + state = copy(self.extra_kwargs) + def recompute(*args, **kwargs): + with self.context(state): + return function(*args, **kwargs) + return _checkpoint(recompute, *args, **kwargs) + module._gradient_checkpointing_func = checkpoint + + def _split_model_inputs(self, inputs): + positions = self._extract_real_position_ids(inputs.get('position_ids')) + model_positions = inputs.get('position_ids') + if self._defer_inputs and torch.is_tensor(model_positions) and model_positions.dim() == 3: + positions = self.real_position_ids # mRoPE axes are not the packed token layout. + if positions is not None: + self.extra_kwargs['position_ids'] = positions.clone() + if inputs.get('attention_mask') is None and not self.extra_kwargs.get('padding_free', False): + inputs['attention_mask'] = (self.real_position_ids >= 0).long() + mask = inputs.get('attention_mask') + if torch.is_tensor(mask) and mask.dim() == 4: + length = self.pad(self.real_position_ids, padding_value=-1).shape[-1] + if self.attn_implementation != 'sdpa' or mask.shape[1] != 1 or mask.shape[-2:] != (length, length): + raise NotImplementedError('SP needs a 2D padding mask or an aligned broadcast SDPA 4D mask') + if torch.is_tensor(mask) and mask.dim() == 2: + inputs['attention_mask'] = self.pad(mask, padding_value=0, position_ids=self.real_position_ids) + ids, embeds, _, positions, mask, _, _ = self.pad_and_split_inputs( + inputs.get('input_ids'), inputs.get('inputs_embeds'), None, + inputs.get('position_ids'), inputs.get('attention_mask'), None, + real_position_ids=self.real_position_ids, cache_position=inputs.get('cache_position')) + self.extra_kwargs['attention_mask'] = inputs.get('attention_mask') + inputs.update(input_ids=ids, inputs_embeds=embeds, position_ids=positions, attention_mask=mask) + return inputs + def _prepare_multimodal_deepstack(self, base_model: torch.nn.Module): if not is_qwen3_vl(base_model): @@ -417,6 +591,10 @@ def _prepare_qwen35_linear_attention(self, model: torch.nn.Module): raise NotImplementedError( 'SequenceParallel: Qwen3.5 linear attention sequence parallel does not support rp_world_size > 1 ' '(derived ring attention).') + from .linear_attention_sp import _FLA_CAUSAL_CONV1D_FN, _FLA_CHUNK_GATED_DELTA_RULE + self._packed_linear_supported = bool(_FLA_CAUSAL_CONV1D_FN and _FLA_CHUNK_GATED_DELTA_RULE) or all( + getattr(module, '_twinkle_npu_patched', False) for module in model.modules() + if 'GatedDeltaNet' in type(module).__name__) apply_patch(None, Qwen3_5GatedDeltaNetUlyssesPatch, sequence_parallel=self) def _prepare_moe_aux_loss(self, base_model: torch.nn.Module): @@ -441,7 +619,7 @@ def moe_aux_loss_hook(module, args, kwargs, output): _bs_logits = [logit[_slice] for logit in router_logits] compute_device = _bs_logits[0].device _bs_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in _bs_logits], dim=0) - _bs_logits, _ = GatherLoss.apply(_bs_logits, None, 1, self.real_position_ids) + _bs_logits, _ = GatherLoss.apply(self, _bs_logits, None, 1, self.real_position_ids) _gathered_logits.append(_bs_logits) router_logits = torch.stack(_gathered_logits, dim=0) if self.real_position_ids is not None: @@ -459,12 +637,18 @@ def prepare( tokenizer: PreTrainedTokenizer, device_mesh: Optional[DeviceMesh] = None, ): - llm_model = get_llm_model(model) - config_candidates = [getattr(model, 'config', None)] + root_model = model + while isinstance(getattr(root_model, 'module', None), torch.nn.Module): + root_model = root_model.module + llm_model = get_llm_model(root_model) + if not (getattr(llm_model, '_supports_attention_backend', False) or self._is_qwen35_model(llm_model)): + raise NotImplementedError(f'SP {type(llm_model).__name__}: unsupported text attention entry point') + llm_model.config.use_cache = False + config_candidates = [getattr(root_model, 'config', None)] llm_config = getattr(llm_model, 'config', None) if llm_config is not None and llm_config not in config_candidates: config_candidates.append(llm_config) - text_config = getattr(getattr(model, 'config', None), 'text_config', None) + text_config = getattr(getattr(root_model, 'config', None), 'text_config', None) if text_config is not None and text_config not in config_candidates: config_candidates.append(text_config) @@ -477,7 +661,8 @@ def prepare( self.num_heads = get_config_attr(config, 'num_attention_heads') if self.num_heads is not None: break - assert self.num_heads is not None, 'Cannot find num_attention_heads/num_key_value_heads in model config' + if not isinstance(self.num_heads, int) or self.num_heads < 1: + raise NotImplementedError(f'SP {type(model).__name__}: cannot determine positive attention head count') self.seq_world_size = sp_size self.sp_world_size, self.rp_world_size = _derive_sequence_parallel_sizes(self.num_heads, self.seq_world_size) self.world_size = self.seq_world_size @@ -497,30 +682,47 @@ def prepare( if hasattr(llm_model, '_update_causal_mask'): self.causal_mask_func = llm_model._update_causal_mask + if self.attn_implementation not in ('sdpa', *FLASH_ATTENTION_IMPLS): + raise NotImplementedError(f'SP {type(model).__name__}: unsupported attention {self.attn_implementation!r}') + if self.attn_implementation in FLASH_ATTENTION_IMPLS: + from transformers.modeling_flash_attention_utils import is_flash_attn_available + from transformers.utils import is_flash_attn_2_available, is_flash_attn_3_available + device = next(model.parameters()).device.type + available = device in ('cuda', 'npu') and ( + is_flash_attn_2_available() if device == 'cuda' else is_flash_attn_available()) + if self.attn_implementation == 'flash_attention_3': + available = device == 'cuda' and is_flash_attn_3_available() + if not available: + raise NotImplementedError(f'SP {self.attn_implementation} kernel unavailable on {device}') + if self.rp_world_size > 1 and self.attn_implementation == 'sdpa': + raise NotImplementedError('SDPA does not support derived ring attention') + if self.rp_world_size > 1 and next(model.parameters()).device.type != 'cuda': + raise NotImplementedError('Derived ring attention requires CUDA FlashAttention kernels') self._init_device_mesh(device_mesh) if not SequenceParallel._global_inited: # these operations are global initializations and patches - self._prepare_flash_attn(llm_model) + self._prepare_flash_attn() SequenceParallel._global_inited = True self._prepare_qwen35_linear_attention(llm_model) - self._prepare_forward_hook(llm_model) + text_model = getattr(llm_model, 'language_model', llm_model) + self._prepare_forward_boundaries(model, text_model) self._prepare_multimodal_deepstack(llm_model) - if is_moe_config(getattr(model, 'config', None)): + if is_moe_config(getattr(llm_model, 'config', None)): self._prepare_moe_aux_loss(llm_model) self.model_dtype = next(model.parameters()).dtype self.tokenizer = tokenizer if self.rp_world_size > 1: - attn_impl = getattr(model.config, '_attn_implementation', None) + attn_impl = self.attn_implementation if attn_impl not in FLASH_ATTENTION_IMPLS: raise NotImplementedError( f'Derived ring attention only supports {"/".join(FLASH_ATTENTION_IMPLS)} backends, ' f'got {attn_impl!r}.') logger.info_once(f'[SequenceParallel] enabled: world_size={self.world_size} ' f'(ulysses={self.sp_world_size} x ring={self.rp_world_size}), ' - f'attn_implementation={getattr(model.config, "_attn_implementation", None)!r}') + f'attn_implementation={getattr(root_model.config, "_attn_implementation", None)!r}') def _mask_qkv(self, query, key, value, mask): mask = mask.unsqueeze(2).unsqueeze(3) @@ -737,8 +939,7 @@ def pad_and_split_inputs(self, input_ids = self.pad(input_ids, padding_value=tokenizer.pad_token_id, position_ids=real_position_ids) self.extra_kwargs['input_ids'] = input_ids.clone() if input_embeds is not None: - pad_emb = torch.zeros( - (1, embed_tokens.weight.shape[-1])).to(embed_tokens.weight.device).to(embed_tokens.weight.dtype) + pad_emb = input_embeds.new_zeros((1, input_embeds.shape[-1])) input_embeds = self.pad(input_embeds, padding_value=pad_emb, position_ids=real_position_ids) batch_size = input_ids.shape[ 0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else 1 @@ -831,56 +1032,46 @@ def _is_packed_position_ids(position_ids: Optional[torch.Tensor]) -> bool: return True return False - def prepare_inputs(self, inputs): - """Prepare inputs - - 1. set extra_kwargs['position_ids'] - 2. split labels, and completion_mask when present - """ - input_ids = inputs.get('input_ids') - position_ids = inputs.get('position_ids') - padding_free = bool(inputs.pop('padding_free', False)) + def prepare_inputs(self, inputs, token_labels=True): + """Cache the global layout and shard the fixed token-level loss inputs once.""" + positions = self._extract_real_position_ids(inputs.get('position_ids')) + if self.rp_world_size > 1 and (positions[:, 0] != 0).any(): + raise NotImplementedError( + 'Ring SP requires sequences beginning at position zero; left padding is unsupported') + padding_free = bool(inputs.pop('padding_free', False)) or self._is_packed_position_ids(positions) + if padding_free and not self._packed_linear_supported: + raise NotImplementedError('Packed Qwen3.5 SP requires FLA or a supported native GDN kernel') if padding_free and not is_flash_attention_implementation(self.attn_implementation): - raise RuntimeError('Transformers SequenceParallel does not support padding_free/packed inputs with ' - f'attn_implementation={self.attn_implementation!r}. ' - 'Use a FlashAttention backend, or disable padding_free/packing. ' - 'SDPA/eager attention cannot safely preserve packed sequence boundaries in this path.') - real_position_ids = self._extract_real_position_ids(position_ids) - if real_position_ids is not None and input_ids is not None and real_position_ids.shape[0] == input_ids.shape[0]: - self.extra_kwargs['position_ids'] = real_position_ids.clone() - self.extra_kwargs['padding_free'] = padding_free - if input_ids is not None: - self.extra_kwargs['input_ids'] = input_ids.clone() - if 'labels' in inputs: - labels = inputs.get('labels') - # completion_mask sits on the labels' index space, so it is padded and - # split identically -- unlike loss_scale, which is rolled beforehand. - completion_mask = inputs.get('completion_mask') - extra_split_values = None if completion_mask is None else [(completion_mask, 0, -1)] - _, _, labels, _, _, _, extra_values = self.pad_and_split_inputs( - None, - None, - labels, - None, - None, - None, - real_position_ids=real_position_ids, - extra_split_values=extra_split_values, - ) - inputs['labels'] = labels - if extra_values: - inputs['completion_mask'] = extra_values[0] + raise NotImplementedError('Packed SP inputs require FlashAttention 2 or 3') + self.extra_kwargs = {'position_ids': positions.clone(), 'padding_free': padding_free} + if inputs.get('input_ids') is not None: + self.extra_kwargs['input_ids'] = inputs['input_ids'].clone() + padded_positions = self.pad(positions, padding_value=-1, position_ids=positions) + sharded = [] + for name, pad_value in (('labels', -100), ('completion_mask', 0), ('loss_scale', 0.0)): + value = inputs.get(name) + if value is None or not token_labels: + continue + if value.dim() != 2 or value.shape[-1] != positions.shape[-1]: + raise ValueError(f'SP {name} must match the token layout, got {tuple(value.shape)}') + value = self.pad(value, padding_value=pad_value, position_ids=positions) + inputs[name] = self.split(value, dim=-1, position_ids=padded_positions) + sharded.append(name) + indices = torch.arange(positions.shape[-1], device=positions.device).expand_as(positions) + indices = self.pad(indices, padding_value=-1, position_ids=positions) + indices = self.split(indices, dim=-1, position_ids=padded_positions) + valid_positions = positions + mask = inputs.get('attention_mask') + if torch.is_tensor(mask) and mask.dim() == 2: + valid_positions = positions.masked_fill(~mask.bool(), -1) + inputs['_sequence_layout'] = {'position_ids': valid_positions, 'indices': indices, 'sharded': sharded} return inputs -sequence_parallel = SequenceParallel() - - @dataclass(frozen=True) class SequenceParallelConfig: enabled: bool = True ulysses_size: Optional[int] = None - gather_logits: bool = True class SequenceParallelStrategy: @@ -906,6 +1097,7 @@ def __init__( self._tokenizer_id = tokenizer_id self._tokenizer = None self._initialized = False + self._ctx = SequenceParallel() def _get_tokenizer(self) -> Optional[PreTrainedTokenizer]: if self._tokenizer is not None: @@ -923,6 +1115,8 @@ def _get_tokenizer(self) -> Optional[PreTrainedTokenizer]: def initialize(self) -> bool: if not self.enabled or self.ulysses_size <= 1: return False + if self._initialized: + return True if not dist.is_initialized(): raise RuntimeError('torch.distributed must be initialized before enabling sequence parallel.') if not isinstance(self.device_mesh, DeviceMesh): @@ -932,7 +1126,7 @@ def initialize(self) -> bool: tokenizer = self._get_tokenizer() if tokenizer is None: raise RuntimeError('SequenceParallelStrategy requires a tokenizer to initialize.') - sequence_parallel.prepare( + self._ctx.prepare( self.ulysses_size, self._model_ref, tokenizer, @@ -941,40 +1135,46 @@ def initialize(self) -> bool: self._initialized = True return True - def preprocess_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + def prepare_model_inputs(self, inputs, token_labels=True, padding_free=False): if not self.enabled or self.ulysses_size <= 1: return inputs - return sequence_parallel.prepare_inputs(inputs) - - def postprocess_outputs(self, outputs: Any) -> Any: - if (not self.enabled or self.ulysses_size <= 1 or not self.sp_config.get('gather_logits', True)): - return outputs - # Twinkle expects dict-like ModelOutput containers in the main training path - # (uses `.get(...)` and `outputs[...] = ...`). Keep SP postprocess consistent. - if outputs is None or not hasattr(outputs, 'get') or not hasattr(outputs, '__setitem__'): - raise TypeError('SequenceParallelStrategy.postprocess_outputs expects a dict-like ModelOutput. ' - f'Got type={type(outputs)}') - logits = outputs.get('logits', None) - if logits is None or not torch.is_tensor(logits) or logits.dim() < 2: - return outputs - gathered = sequence_parallel.gather(logits, dim=1, position_ids=sequence_parallel.real_position_ids) - # Scheme A: SP pads to make seq_len divisible by sp_size. Trim back to the original - # (unpadded) length using the cached real_position_ids. - real_pos = sequence_parallel.real_position_ids - if real_pos is not None and torch.is_tensor(real_pos) and real_pos.dim() >= 2: - gathered = gathered[:, :real_pos.shape[1]].contiguous() - outputs['logits'] = gathered - return outputs + inputs = dict(inputs) + if inputs.get('past_key_values') is not None or inputs.get('use_cache') is True: + raise NotImplementedError('SP supports training/prefill without KV cache only') + ids = inputs.get('input_ids') + embeds = inputs.get('inputs_embeds') + tensor = ids if ids is not None else embeds + infer_model_positions = inputs.get('position_ids') is None and self._ctx._defer_inputs + if inputs.get('position_ids') is None: + if tensor is None: + raise ValueError('SP requires input_ids, inputs_embeds or position_ids') + positions = torch.arange(tensor.shape[1], device=tensor.device).expand(tensor.shape[:2]) + mask = inputs.get('attention_mask') + inputs['position_ids'] = (positions.masked_fill(~mask.bool(), -1) + if torch.is_tensor(mask) and mask.dim() == 2 else positions) + inputs['padding_free'] = bool(inputs.get('padding_free', False) or padding_free) + self._ctx.prepare_inputs(inputs, token_labels=token_labels) + if not self._ctx._defer_inputs: + inputs = self._ctx._split_model_inputs(inputs) + elif infer_model_positions: + inputs.pop('position_ids') # Let HF compute multimodal RoPE before the text boundary. + return inputs - @staticmethod - def _trim_gathered_sequence_padding(tensor: torch.Tensor, real_position_ids: torch.Tensor) -> torch.Tensor: + def context(self): + if not self.enabled or self.ulysses_size <= 1: + return nullcontext() + if not self._initialized: + raise RuntimeError('Initialize sequence parallel before entering the model context') + return self._ctx.context() + + def _trim_gathered_sequence_padding(self, tensor: torch.Tensor, real_position_ids: torch.Tensor) -> torch.Tensor: if real_position_ids is None or not torch.is_tensor(real_position_ids) or real_position_ids.dim() < 2: return tensor - if sequence_parallel.rp_world_size > 1: + if self._ctx.rp_world_size > 1: cu_seqlens = get_cu_seqlens_from_position_ids(real_position_ids) pieces = [] padded_offset = 0 - divisor = sequence_parallel.world_size * 2 + divisor = self._ctx.world_size * 2 for i in range(len(cu_seqlens) - 1): real_len = int((cu_seqlens[i + 1] - cu_seqlens[i]).item()) padded_len = math.ceil(real_len / divisor) * divisor @@ -983,72 +1183,37 @@ def _trim_gathered_sequence_padding(tensor: torch.Tensor, real_position_ids: tor return torch.cat(pieces, dim=1).contiguous() if pieces else tensor[:, :0].contiguous() return tensor[:, :real_position_ids.shape[-1]].contiguous() - def _gather_completion_mask(self, inputs: Dict[str, Any], real_position_ids) -> None: - """Gather ``completion_mask`` in place, mirroring the labels gather. - - Deliberately not routed through :class:`GatherLoss`: the mask carries no - gradient, and reusing that autograd Function would attach a second backward - path to whichever tensor were passed alongside it, double-scaling its grad. - """ - mask = inputs.get('completion_mask') - if mask is None or not torch.is_tensor(mask) or mask.dim() < 2: - return - gathered = sequence_parallel.gather(mask, dim=1, position_ids=real_position_ids) - inputs['completion_mask'] = self._trim_gathered_sequence_padding(gathered, real_position_ids) - - def gather_loss_tensors( - self, - inputs: Dict[str, Any], - outputs: Dict[str, Any], - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: - if inputs is None or outputs is None: - return inputs, outputs - if not self.enabled or self.ulysses_size <= 1: - return inputs, outputs - labels = inputs.get('labels') - logps = outputs.get('logps') - # Fused-linear-CE path: the model forward skipped the lm_head GEMM and - # stashed the lm_head module under ``outputs['lm_head']`` (see - # ``TransformersFusedCEPatch``), so ``outputs['logps']`` is absent and - # ``outputs['logits']`` holds the last hidden state (a per-rank seq - # shard under SP). Gather hidden states + labels across the sequence - # group so every rank computes the same full loss — mirrors the logps - # gather below. The hidden tensor ([B, T_local, H]) is tiny vs logits, - # and the fused loss still chunks the GEMM internally, so the memory - # win is preserved. - if (logps is None and labels is not None and 'lm_head' in outputs and torch.is_tensor(outputs.get('logits')) - and outputs['logits'].dim() >= 2): - hidden = outputs['logits'] - inputs = copy(inputs) - outputs = copy(outputs) - real_position_ids = sequence_parallel.real_position_ids - gathered_hidden, gathered_labels = GatherLoss.apply(hidden, labels, 1, real_position_ids) - gathered_hidden = self._trim_gathered_sequence_padding(gathered_hidden, real_position_ids) - gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids) - outputs['logits'] = gathered_hidden - inputs['labels'] = gathered_labels - self._gather_completion_mask(inputs, real_position_ids) - return inputs, outputs - if labels is None or logps is None: + def gather_loss_tensors(self, inputs, outputs): + """Restore complete sequence tensors; values do not require labels or logps.""" + if inputs is None or outputs is None or not self.enabled or self.ulysses_size <= 1: return inputs, outputs - if not torch.is_tensor(logps) or logps.dim() < 2: - raise TypeError('SequenceParallelStrategy.gather_loss_inputs expects outputs[\"logps\"] to be a ' - f'sequence tensor, got type={type(logps)} shape={getattr(logps, "shape", None)}') - inputs = copy(inputs) - outputs = copy(outputs) - real_position_ids = sequence_parallel.real_position_ids - gathered_logps, gathered_labels = GatherLoss.apply(logps, labels, 1, real_position_ids) - gathered_logps = self._trim_gathered_sequence_padding(gathered_logps, real_position_ids) - gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids) - outputs['logps'] = gathered_logps - inputs['labels'] = gathered_labels - self._gather_completion_mask(inputs, real_position_ids) - entropies = outputs.get('entropies') - if entropies is not None and torch.is_tensor(entropies) and entropies.dim() >= 2: - gathered_entropies, _ = GatherLoss.apply(entropies, labels, 1, real_position_ids) - outputs['entropies'] = self._trim_gathered_sequence_padding(gathered_entropies, real_position_ids) + inputs, outputs = copy(inputs), copy(outputs) + positions = self._ctx.real_position_ids + for name in ('logps', 'entropies', 'values', 'logits'): + value = outputs.get(name) + if value is None: + continue + if not torch.is_tensor(value) or value.dim() < 2: + raise TypeError(f'SP output {name} must be a batched sequence tensor') + value, _ = GatherLoss.apply(self._ctx, value, None, 1, positions) + outputs[name] = self._trim_gathered_sequence_padding(value, positions) + layout = inputs.pop('_sequence_layout', None) + if layout is None and any(inputs.get(name) is not None for name in ('labels', 'completion_mask', 'loss_scale')): + raise ValueError('SP loss inputs must come from prepare_model_inputs') + for name in (layout or {}).get('sharded', ()): + value = self._ctx.gather(inputs[name], dim=1, position_ids=positions) + inputs[name] = self._trim_gathered_sequence_padding(value, positions) + inputs['position_ids'] = positions return inputs, outputs + def all_reduce(self, tensor, op=dist.ReduceOp.SUM): + from torch.distributed.nn.functional import all_reduce + for group, size in ((self._ctx._sp_group, self._ctx.sp_world_size), + (self._ctx._rp_group, self._ctx.rp_world_size)): + if size and size > 1: + tensor = all_reduce(tensor, op=op, group=group) + return tensor + def wrap_model(self, model, optimizer=None): self.initialize() return model, optimizer diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py b/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py index ab8db95ed..4cc1929d4 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py @@ -184,9 +184,15 @@ def _run_forward( apply_mask_to_padding_states = getattr(modeling_module, 'apply_mask_to_padding_states') local_attention_mask = attention_mask - if torch.is_tensor(attention_mask) and attention_mask.dim() == 2: + if _sp_is_enabled(sequence_parallel_context): + # New HF recurrent-mask helpers crop to the tail on every rank. + # Recover the owning batch's global mask before selecting this rank. + global_mask = sequence_parallel_context.extra_kwargs.get('attention_mask') + if torch.is_tensor(global_mask) and global_mask.dim() == 2: + local_attention_mask = global_mask + if torch.is_tensor(local_attention_mask) and local_attention_mask.dim() == 2: local_attention_mask = _get_local_padding_mask( - attention_mask, + local_attention_mask, hidden_states.shape[1], sequence_parallel_context, ) @@ -340,7 +346,8 @@ def sp_linear_forward( _origin_forward=origin_forward, **extra_kwargs, ): - sequence_parallel_context = extra_kwargs.pop('sequence_parallel_context', sequence_parallel) + from . import _attention_context + sequence_parallel_context = extra_kwargs.pop('sequence_parallel_context', _attention_context(mod)) if not _sp_is_enabled(sequence_parallel_context): return _origin_forward( mod, diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/utils.py b/src/twinkle/model/transformers/strategy/sequence_parallel/utils.py index 6bc8677dc..8ed17e09c 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/utils.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/utils.py @@ -237,33 +237,39 @@ def head_to_seq_shard(tensor: torch.Tensor, sequence_parallel) -> torch.Tensor: class GatherLoss(torch.autograd.Function): - """Gather loss from sequence group.""" + """Gather loss from sequence group. + + The owning ``SequenceParallel`` context is passed explicitly as the first + argument: backward() runs after other models may have forwarded, so reading + the process-global singleton here could pick up the wrong context in + multi-model training (DPO/PPO/GKD). + """ @staticmethod - def forward(ctx, loss, labels, gather_idx=None, position_ids=None): - from . import sequence_parallel + def forward(ctx, sp, loss, labels, gather_idx=None, position_ids=None): + ctx.sp = sp ctx.scatter_shape = loss.shape[gather_idx or 0] ctx.gather_idx = gather_idx or 0 if position_ids is not None: - position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids) + position_ids = sp.pad(position_ids, padding_value=-1, position_ids=position_ids) ctx.position_ids = position_ids - output = sequence_parallel.gather(loss, dim=ctx.gather_idx, position_ids=position_ids) + output = sp.gather(loss, dim=ctx.gather_idx, position_ids=position_ids) if labels is not None: - labels_output = sequence_parallel.gather(labels, dim=ctx.gather_idx, position_ids=position_ids) + labels_output = sp.gather(labels, dim=ctx.gather_idx, position_ids=position_ids) else: labels_output = None return output, labels_output @staticmethod def backward(ctx, *grad_output): - from . import sequence_parallel - _grad = grad_output[0] * sequence_parallel.world_size - if sequence_parallel.rp_world_size > 1: - _grad = sequence_parallel.split(_grad, dim=ctx.gather_idx, position_ids=ctx.position_ids).contiguous() + sp = ctx.sp + _grad = grad_output[0] * sp.world_size + if sp.rp_world_size > 1: + _grad = sp.split(_grad, dim=ctx.gather_idx, position_ids=ctx.position_ids).contiguous() else: _grad = _grad.split( - ctx.scatter_shape, dim=ctx.gather_idx)[dist.get_rank(group=sequence_parallel._sp_group)].contiguous() - return _grad, None, None, None + ctx.scatter_shape, dim=ctx.gather_idx)[dist.get_rank(group=sp._sp_group)].contiguous() + return None, _grad, None, None, None def _generate_layout_params(scatter_idx, seq_world_size, input): diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 2270a472e..1a9f4f140 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -111,11 +111,11 @@ def _resolve_task_context(model, task, template=None): 'causal_lm' (default) keeps the model untouched (returns ``nullcontext``). 'embedding' swaps lm_head for identity + installs a feature-extraction hook so downstream pooling can run inside - ``InputProcessor.postprocess_tensor_sp(task='embedding', ...)``. + ``InputProcessor.postprocess_outputs(task='embedding', ...)``. 'generative_reranker' wraps lm_head to emit a per-token yes-minus-no score; the - last-valid-token reduction is deferred to ``postprocess_tensor_sp`` too. + last-valid-token reduction is deferred to ``postprocess_outputs`` too. 'seq_cls' / 'reranker' capture per-token scores BEFORE HF's internal pooling so - ``postprocess_tensor_sp`` can do packed-aware last-token reduction (same as embedding). + ``postprocess_outputs`` can do packed-aware last-token reduction (same as embedding). """ if task in (None, 'causal_lm'): return contextlib.nullcontext() @@ -870,7 +870,7 @@ def _postprocess_forward_outputs(self, outputs, *, labels, task, temperature, sa del logits if getattr(loss_instance, 'require_values', False): values = outputs['logits'] - outputs['values'] = values.squeeze(-1) if values.shape[-1] == 1 else values + outputs['values'] = values.squeeze(-1) if values.dim() == 3 and values.shape[-1] == 1 else values outputs['past_key_values'] = None if not (return_logits or loss_require_logits): outputs['logits'] = None @@ -878,8 +878,10 @@ def _postprocess_forward_outputs(self, outputs, *, labels, task, temperature, sa def _finalize_forward_outputs(self, inputs, outputs, *, processor, task, optimizer_config, is_training, forward_kwargs, loss_instance, return_logits, recorded_routing): """Shared forward/forward_only tail: SP postprocess, unpack, status recording, and the return copy.""" - inputs, outputs = processor.postprocess_tensor_sp(inputs, outputs, sp_strategy=self.sp_strategy, task=task) - inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task) + inputs, outputs = processor.postprocess_outputs( + inputs, outputs, task=task, + gather_fn=self.sp_strategy.gather_loss_tensors if self.sp_strategy is not None else None, + reduce_fn=self.sp_strategy.all_reduce if self.sp_strategy is not None else None) status = optimizer_config.train_status if is_training else optimizer_config.eval_status status.inputs = inputs status.outputs = outputs @@ -936,11 +938,14 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' inputs: Dict[str, Any] = processor( inputs, - sp_strategy=self.sp_strategy, hf_config=self.hf_config, enable_sp=getattr(self, '_enable_sp', False), gdn_padding_free_enabled=self._gdn_padding_free_enabled, ) + if self.sp_strategy is not None: + inputs = self.sp_strategy.prepare_model_inputs( + inputs, token_labels=processor.token_level_labels(task), padding_free=processor.padding_free) + sequence_layout = inputs.pop('_sequence_layout', None) labels: torch.Tensor = inputs.pop('labels', None) loss_scale = inputs.pop('loss_scale', None) channel = inputs.pop('channel', None) @@ -979,11 +984,14 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec unwrapped_model = self.strategy.unwrap_model(self.model) with _resolve_task_context(unwrapped_model, task, template=optimizer_config.template), \ _resolve_mm_graph_context(unwrapped_model, optimizer_config.template): - outputs = self.model(**inputs) + with self.sp_strategy.context() if self.sp_strategy is not None else contextlib.nullcontext(): + outputs = self.model(**inputs) recorded_routing = rr.finish() self._restore_loss_inputs(inputs, labels, loss_scale, channel, completion_mask) + if sequence_layout is not None: + inputs['_sequence_layout'] = sequence_layout self._postprocess_forward_outputs( outputs, labels=labels, @@ -1103,11 +1111,14 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs): raise ValueError('sampling replay does not support sequence or context parallelism') inputs: Dict[str, Any] = processor( inputs, - sp_strategy=self.sp_strategy, hf_config=self.hf_config, enable_sp=getattr(self, '_enable_sp', False), gdn_padding_free_enabled=self._gdn_padding_free_enabled, ) + if self.sp_strategy is not None: + inputs = self.sp_strategy.prepare_model_inputs( + inputs, token_labels=processor.token_level_labels(task), padding_free=processor.padding_free) + sequence_layout = inputs.pop('_sequence_layout', None) labels = inputs.pop('labels', None) loss_scale = inputs.pop('loss_scale', None) channel = inputs.pop('channel', None) @@ -1145,11 +1156,14 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs): if disable_lora and isinstance(unwrapped_model, PeftModel) else contextlib.nullcontext()) # unwrapped on purpose: task patches probe head/module attributes DDP does not forward. with _resolve_task_context(unwrapped_model, task, template=optimizer_config.template), lora_ctx: - outputs = self.model(**inputs) + with self.sp_strategy.context() if self.sp_strategy is not None else contextlib.nullcontext(): + outputs = self.model(**inputs) recorded_routing = rr.finish() self._restore_loss_inputs(inputs, labels, loss_scale, channel, completion_mask) + if sequence_layout is not None: + inputs['_sequence_layout'] = sequence_layout self._postprocess_forward_outputs( outputs, labels=labels, diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index cbea6cd3b..37a260abc 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -84,7 +84,6 @@ def __init__( self.prepare_gdn_padding_free_args, self.drop_causal_4d_mask, self.split_cp, - self.apply_transformers_sp, self.prepare_outputs, ] @@ -148,49 +147,25 @@ def to_tensor(_input): return [to_tensor(_input) for _input in inputs] - def apply_transformers_sp(self, inputs: list[InputFeature], **kwargs) -> list[InputFeature]: - sp_strategy = kwargs.get("sp_strategy") - if self.framework != "transformers" or sp_strategy is None: - return inputs - padding_free = bool(self.padding_free or self._any_packing(inputs)) - results = [] - for _input in inputs: - payload = dict(_input) - payload["padding_free"] = padding_free - results.append(InputFeature(**sp_strategy.preprocess_inputs(payload))) - return results + @staticmethod + def token_level_labels(task): + return task not in ('embedding', 'generative_reranker', 'seq_cls', 'reranker') - def postprocess_tensor_sp( - self, inputs: dict[str, Any], outputs: dict[str, Any], **kwargs - ) -> tuple[dict[str, Any], dict[str, Any]]: - """Adjust SP tensors after forward and before loss computation. - - Pipeline: SP gather → packed-sequence unpack. - After this call, logps and labels are in per-sequence batch format - ``[num_sequences, max_seq_len]`` when the input was packed, or left - unchanged for normal (non-packed) batches. - - For ``task='embedding'`` this also performs the last-valid-token - pooling (with padding-free / SP gather awareness) and writes the - pooled ``[n_seqs, H]`` tensor to ``outputs['embeddings']``; the raw - per-token ``outputs['features']`` is consumed and removed. - - For ``task='seq_cls'``/``'reranker'`` the model outputs per-token scores - ``[B, T, num_labels]`` (TransformersValuePatch skips HF's internal pooling); - this method pools them to ``[n_seqs, num_labels]`` via the same packed-aware - last-valid-token logic as embedding. + def postprocess_outputs(self, inputs, outputs, *, task='causal_lm', gather_fn=None, reduce_fn=None): + """Task reduction, loss-tensor collection and packed layout restoration. + + Callbacks contain only communication; token selection and task semantics + stay here. Without callbacks, both training backends keep their local path. """ - sp_strategy = kwargs.get("sp_strategy") - task = kwargs.get("task", "causal_lm") - if task == "embedding": - return self._postprocess_embedding(inputs, outputs, sp_strategy=sp_strategy) - if task == "generative_reranker": - return self._postprocess_generative_reranker(inputs, outputs, sp_strategy=sp_strategy) - if task in ("seq_cls", "reranker"): - return self._postprocess_seq_cls(inputs, outputs, sp_strategy=sp_strategy) - if self.framework == "transformers" and sp_strategy is not None: - return sp_strategy.gather_loss_tensors(inputs, outputs) - return inputs, outputs + if task == 'embedding': + inputs, outputs = self._postprocess_embedding(inputs, outputs, reduce_fn=reduce_fn) + elif task == 'generative_reranker': + inputs, outputs = self._postprocess_generative_reranker(inputs, outputs, reduce_fn=reduce_fn) + elif task in ('seq_cls', 'reranker'): + inputs, outputs = self._postprocess_seq_cls(inputs, outputs, reduce_fn=reduce_fn) + elif gather_fn is not None: + inputs, outputs = gather_fn(inputs, outputs) + return self.unpack_packed_sequences(inputs, outputs, task=task) @staticmethod def _packed_last_indices(position_ids: torch.Tensor, total_len: int) -> torch.Tensor: @@ -202,7 +177,7 @@ def _packed_last_indices(position_ids: torch.Tensor, total_len: int) -> torch.Te return (boundaries[1:] - 1).long() def _postprocess_embedding( - self, inputs: dict[str, Any], outputs: dict[str, Any], sp_strategy=None + self, inputs: dict[str, Any], outputs: dict[str, Any], reduce_fn=None ) -> tuple[dict[str, Any], dict[str, Any]]: """Pool per-token features to per-sequence embeddings (last-valid-token). @@ -217,7 +192,7 @@ def _postprocess_embedding( "outputs['features'] missing; ensure TransformersEmbeddingPatch is applied for task='embedding'." ) - pooled, is_packed = self._pool_last_valid_token(features, inputs, sp_strategy=sp_strategy) + pooled, is_packed = self._pool_last_valid_token(features, inputs, reduce_fn=reduce_fn) # Normalize after pool+reduce so swapping pooling strategy (last/mean/CLS) stays mathematically correct. embeddings = F.normalize(pooled, p=2, dim=-1) @@ -228,7 +203,7 @@ def _postprocess_embedding( return inputs, outputs def _postprocess_generative_reranker( - self, inputs: dict[str, Any], outputs: dict[str, Any], sp_strategy=None + self, inputs: dict[str, Any], outputs: dict[str, Any], reduce_fn=None ) -> tuple[dict[str, Any], dict[str, Any]]: """Reduce per-token generative-reranker scores ``[B, T, 1]`` to ``[n_seqs, 1]``. @@ -242,13 +217,13 @@ def _postprocess_generative_reranker( "outputs['logits'] missing; ensure TransformersGenerativeRerankerPatch is applied and the " "reranker loss sets require_logits=True." ) - pooled, _ = self._pool_last_valid_token(logits, inputs, sp_strategy=sp_strategy) + pooled, _ = self._pool_last_valid_token(logits, inputs, reduce_fn=reduce_fn) outputs = copy(outputs) outputs["logits"] = pooled.contiguous() return inputs, outputs def _postprocess_seq_cls( - self, inputs: dict[str, Any], outputs: dict[str, Any], sp_strategy=None + self, inputs: dict[str, Any], outputs: dict[str, Any], reduce_fn=None ) -> tuple[dict[str, Any], dict[str, Any]]: """Pool per-token seq_cls/reranker scores ``[B, T, num_labels]`` to ``[n_seqs, num_labels]``. @@ -261,35 +236,30 @@ def _postprocess_seq_cls( "outputs['logits'] missing; ensure TransformersValuePatch is applied for task='seq_cls'/'reranker'." ) if logits.dim() == 2: + if reduce_fn is not None and inputs.get('_sequence_layout') is not None: + raise NotImplementedError('Sequence-parallel scoring requires a per-token score head') # A pre-pooled head (XLM-R-style ``*ForSequenceClassification`` pools the last valid # token INSIDE forward, before the score head) already answered one score per sequence; # the patch captured ``[n_seqs, num_labels]`` and there is nothing left to reduce. return inputs, outputs - pooled, _ = self._pool_last_valid_token(logits, inputs, sp_strategy=sp_strategy) + pooled, _ = self._pool_last_valid_token(logits, inputs, reduce_fn=reduce_fn) outputs = copy(outputs) outputs["logits"] = pooled.contiguous() return inputs, outputs def _pool_last_valid_token( - self, features: torch.Tensor, inputs: dict[str, Any], sp_strategy=None + self, features: torch.Tensor, inputs: dict[str, Any], reduce_fn=None ) -> tuple[torch.Tensor, bool]: """Gather the last valid token per sequence from ``[B, T, C]`` -> ``[n_seqs, C]``. - Build a one-hot end-token mask in the un-padded global frame, route it - through the same pad+split as ``input_ids`` so it aligns with local - features, pool locally, then ``all_reduce`` only the ``[n_seqs, C]`` - tensor across SP × RP. No feature gather; uniform across - DP / Ulysses / zigzag-ring / padding-free. Returns ``(pooled, is_packed)``. + Select endpoints in the unpadded global frame, then use the batch-local + token index mapping to align the mask with local features. The optional + callback reduces only the pooled tensor; this method does not own any + process groups or sequence-splitting logic. """ - sp_enabled = ( - self.framework == "transformers" - and sp_strategy is not None - and getattr(sp_strategy, "enabled", False) - and getattr(sp_strategy, "world_size", 1) > 1 - ) - - if sp_enabled: - ref_pos = sp_strategy.real_position_ids + layout = inputs.get('_sequence_layout') + if layout is not None: + ref_pos = layout['position_ids'] elif inputs.get("position_ids") is not None: ref_pos = inputs["position_ids"] else: @@ -308,6 +278,9 @@ def _pool_last_valid_token( ref_pos = torch.where(am.bool(), arange, arange.new_tensor(-1)) if ref_pos.dim() == 3: ref_pos = ref_pos[0] + attention_mask = inputs.get('attention_mask') + if torch.is_tensor(attention_mask) and attention_mask.shape == ref_pos.shape: + ref_pos = ref_pos.masked_fill(~attention_mask.bool(), -1) # Accept both flash-attn (cu_seq_lens_q) and Megatron (cu_seqlens_q) conventions. cu_seq_lens_q = inputs.get("cu_seq_lens_q", inputs.get("cu_seqlens_q")) @@ -328,7 +301,8 @@ def _pool_last_valid_token( else: n_seqs = ref_pos.shape[0] # clamp_min(0) prevents negative indexing when a row is fully padded (sum == 0). - end_idx = ((ref_pos >= 0).long().sum(-1) - 1).clamp_min(0) + indices = torch.arange(T_real, device=ref_pos.device).expand_as(ref_pos) + end_idx = indices.masked_fill(ref_pos < 0, -1).amax(-1).clamp_min(0) batch_idx = torch.arange(n_seqs, device=device) target_idx = torch.zeros(n_seqs, dtype=torch.long, device=device) B_mask, n_targets = n_seqs, 1 @@ -336,26 +310,21 @@ def _pool_last_valid_token( mask = torch.zeros(B_mask, T_real, n_targets, dtype=dtype, device=device) mask[batch_idx, end_idx, target_idx] = 1.0 - if sp_enabled: - # Route mask through the same pad+split as input_ids to align with local features. - rp = sp_strategy.real_position_ids - rp_padded = sp_strategy.pad(rp, padding_value=-1, position_ids=rp, dim=-1) - mask = sp_strategy.pad(mask, padding_value=0, position_ids=rp, dim=1) - mask = sp_strategy.split(mask, dim=1, position_ids=rp_padded) + if layout is not None: + indices = layout['indices'] + mask = mask.gather(1, indices.clamp_min(0).unsqueeze(-1).expand(-1, -1, n_targets)) + mask = mask * (indices >= 0).unsqueeze(-1) assert mask.shape[1] == features.shape[1], ( f"mask seq-dim {mask.shape[1]} != features seq-dim {features.shape[1]}; " - "sp_strategy.split contract broken or T_real mismatched." + "sequence layout does not match local features." ) pooled = torch.einsum("bth,btn->bnh", features, mask) pooled = pooled.squeeze(0) if is_packed else pooled.squeeze(1) - if sp_enabled and dist.is_available() and dist.is_initialized(): - for grp_attr, size_attr in (("_sp_group", "sp_world_size"), ("_rp_group", "rp_world_size")): - grp = getattr(sp_strategy, grp_attr, None) - if grp is not None and getattr(sp_strategy, size_attr, 1) > 1: - dist.all_reduce(pooled, op=dist.ReduceOp.SUM, group=grp) + if reduce_fn is not None: + pooled = reduce_fn(pooled, torch.distributed.ReduceOp.SUM) return pooled, is_packed @@ -712,7 +681,7 @@ def unpack_packed_sequences( skipped. For ``task='embedding'`` the outputs are already pooled to ``[n_seqs, H]`` - by ``postprocess_tensor_sp``, so this is a no-op. Same for + by ``postprocess_outputs``, so this is a no-op. Same for ``task='generative_reranker'``, whose logits are already reduced to ``[n_seqs, 1]``, and ``task='seq_cls'``/``'reranker'``, whose per-token scores are pooled to ``[n_seqs, num_labels]`` by ``_postprocess_seq_cls``. diff --git a/tests/loss/test_channel_loss.py b/tests/loss/test_channel_loss.py index 2d6c8f008..35410fe9e 100644 --- a/tests/loss/test_channel_loss.py +++ b/tests/loss/test_channel_loss.py @@ -125,7 +125,7 @@ def __init__(self): def __call__(self, inputs, **kwargs): return dict(inputs) - def postprocess_tensor_sp(self, inputs, outputs, **kwargs): + def postprocess_outputs(self, inputs, outputs, **kwargs): return inputs, outputs def unpack_packed_sequences(self, inputs, outputs=None, **kwargs): diff --git a/tests/processor/test_processor.py b/tests/processor/test_processor.py index 6840a9364..6add14a6b 100644 --- a/tests/processor/test_processor.py +++ b/tests/processor/test_processor.py @@ -206,3 +206,14 @@ def test_grpo_padding_free(self): out = proc.collate_fn(batch) assert out[0]['input_ids'].shape == (1, 6) assert out[0]['labels'].shape == (1, 6) + + +def test_prepooled_score_head_keeps_local_path_but_rejects_sequence_shards(): + import pytest + processor = InputProcessor() + outputs = {'logits': torch.ones(2, 3)} + inputs, actual = processor.postprocess_outputs({}, outputs, task='seq_cls') + assert actual['logits'] is outputs['logits'] + with pytest.raises(NotImplementedError, match='per-token score head'): + processor.postprocess_outputs({'_sequence_layout': {}}, outputs, task='seq_cls', + reduce_fn=lambda value, op: value) diff --git a/tests/transformers/test_qwen35_linear_attention_sp.py b/tests/transformers/test_qwen35_linear_attention_sp.py index d46c423cf..77e1347e1 100644 --- a/tests/transformers/test_qwen35_linear_attention_sp.py +++ b/tests/transformers/test_qwen35_linear_attention_sp.py @@ -12,7 +12,7 @@ from types import MethodType, SimpleNamespace from twinkle.loss import CrossEntropyLoss -from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy, sequence_parallel +from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy from twinkle.model.transformers.strategy.sequence_parallel.linear_attention_sp import Qwen3_5GatedDeltaNetUlyssesPatch from twinkle.utils import DeviceMesh, selective_log_softmax @@ -41,8 +41,7 @@ _FLA_CAUSAL_CONV1D_FN = None _FLA_CHUNK_GATED_DELTA_RULE = None -pytestmark = pytest.mark.skip( - reason='Heavy Qwen3.5 SP linear-attention test (FLA + multi-GPU); not viable on dual-V100 CI, run manually.') + WORLD_SIZE = 2 LOGITS_RTOL = 5e-3 @@ -277,6 +276,7 @@ def _compute_training_path_loss( loss_inputs = {'labels': labels} loss_outputs = {'logps': selective_log_softmax(logits, masked_labels)} if strategy is not None: + loss_inputs['_sequence_layout'] = {'sharded': ('labels',)} loss_inputs, loss_outputs = strategy.gather_loss_tensors(loss_inputs, loss_outputs) result = CrossEntropyLoss(reduction='sum')(loss_inputs, loss_outputs) num_tokens = result['num_tokens'] @@ -337,20 +337,20 @@ def _run_prefill_alignment_worker(rank: int, baseline_qkv_grad = _get_qkv_weight(baseline_model).grad.detach().float().cpu() / max(baseline_num_tokens, 1) strategy = _make_strategy(sp_model, world_size) - processed_inputs = strategy.preprocess_inputs({ + processed_inputs = strategy.prepare_model_inputs({ 'input_ids': input_ids, 'position_ids': position_ids, 'labels': labels, 'padding_free': packed, + 'attention_mask': attention_mask, }) local_labels = processed_inputs['labels'] - sp_outputs = sp_model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - use_cache=False, - ) - gathered_outputs = strategy.postprocess_outputs(copy.copy(sp_outputs)) + model_inputs = dict(processed_inputs) + model_inputs.pop('labels', None) + model_inputs.pop('_sequence_layout', None) + with strategy.context(): + sp_outputs = sp_model(**model_inputs, use_cache=False) + _, gathered_outputs = strategy.gather_loss_tensors({}, copy.copy(sp_outputs)) gathered_logits = gathered_outputs.logits.float() if not torch.allclose(gathered_logits, baseline_logits, rtol=LOGITS_RTOL, atol=LOGITS_ATOL): max_diff = (gathered_logits - baseline_logits).abs().max().item() @@ -367,7 +367,7 @@ def _run_prefill_alignment_worker(rank: int, if sp_num_tokens != baseline_num_tokens: raise AssertionError( f'prefill num_tokens mismatch on rank {rank}: baseline={baseline_num_tokens} sp={sp_num_tokens}') - sp_qkv_grad = _average_qkv_grad_over_group(sp_model, sequence_parallel._data_rank_group) / max(sp_num_tokens, 1) + sp_qkv_grad = _average_qkv_grad_over_group(sp_model, strategy._ctx._data_rank_group) / max(sp_num_tokens, 1) if not torch.allclose(sp_qkv_grad, baseline_qkv_grad, rtol=GRAD_RTOL, atol=GRAD_ATOL): max_diff = (sp_qkv_grad - baseline_qkv_grad).abs().max().item() raise AssertionError(f'qkv grad mismatch on rank {rank}: max_diff={max_diff}') @@ -382,6 +382,8 @@ def _run_prefill_alignment_worker(rank: int, not (torch.cuda.is_available() and torch.cuda.device_count() >= WORLD_SIZE), reason='requires 2 CUDA devices') @pytest.mark.skipif( not _HAS_FLA_PREFILL, reason='requires flash-linear-attention kernels for Qwen3.5 SP linear attention tests') +@pytest.mark.slow +@pytest.mark.accel(2) class TestQwen35LinearAttentionSP: def test_qwen35_linear_attention_prefill_logits_and_qkv_grad_alignment(self): @@ -455,3 +457,60 @@ def test_qwen35_linear_attention_packed_prefill_logits_and_qkv_grad_alignment(se if __name__ == '__main__': pytest.main([__file__]) + + +def _qwen35_cpu_worker(rank, port, layer_types, padded=False): + from types import SimpleNamespace + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + torch.set_num_threads(1) + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, world_size=2) + try: + torch.manual_seed(92) + config = Qwen3_5TextConfig(vocab_size=32, hidden_size=32, intermediate_size=64, + num_hidden_layers=len(layer_types), num_attention_heads=2, num_key_value_heads=2, head_dim=16, + linear_conv_kernel_dim=4, linear_key_head_dim=8, linear_value_head_dim=8, + linear_num_key_heads=2, linear_num_value_heads=4, + layer_types=list(layer_types), use_cache=False) + config._attn_implementation = 'sdpa' + baseline = Qwen3_5ForCausalLM(config).train() + model = copy.deepcopy(baseline) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) + mesh = DeviceMesh.from_sizes(fsdp_size=2, dp_size=1, ulysses_size=2, device_type='cpu') + strategy = SequenceParallelStrategy(mesh, {'ulysses_size': 2}, model=model) + strategy._tokenizer = SimpleNamespace(pad_token_id=0) + strategy.initialize() + ids = torch.tensor([[1, 3, 5, 7, 9, 11, 13]]) + if padded: + # The HF recurrent mask takes the tail on both ranks. Rank zero + # needs the prefix instead, including all four valid short tokens. + ids = torch.cat((ids, torch.tensor([[2, 4, 6, 8, 0, 0, 0]]))) + mask = ids.ne(0).long() + expected = baseline(input_ids=ids, attention_mask=mask, use_cache=False).logits + inputs = strategy.prepare_model_inputs({'input_ids': ids, 'attention_mask': mask}) + inputs.pop('_sequence_layout') + with strategy.context(): + output = model(**inputs, use_cache=False) + _, gathered = strategy.gather_loss_tensors({}, output) + # Padding outputs are not supervised; compare every valid token and + # backpropagate the same full-sequence objective on both paths. + torch.testing.assert_close(gathered.logits[mask.bool()], expected[mask.bool()], rtol=3e-4, atol=3e-6) + expected[mask.bool()].square().mean().backward() + gathered.logits[mask.bool()].square().mean().backward() + for (name, parameter), (_, original) in zip(model.named_parameters(), baseline.named_parameters()): + if original.grad is None: + assert parameter.grad is None + continue + assert parameter.grad is not None, name + dist.all_reduce(parameter.grad) + parameter.grad.div_(2) + torch.testing.assert_close(parameter.grad, original.grad, rtol=5e-4, atol=3e-6, msg=name) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not _HAS_QWEN35, reason='transformers Qwen3.5 is unavailable') +@pytest.mark.parametrize('layer_types', [('linear_attention',), ('full_attention',), ('linear_attention', 'full_attention')]) +@pytest.mark.parametrize('padded', [False, True]) +def test_qwen35_cpu_torch_fallback_sp_and_checkpoint_match_baseline(layer_types, padded): + mp.spawn(_qwen35_cpu_worker, args=(_find_free_port(), layer_types, padded), nprocs=2, join=True) diff --git a/tests/transformers/test_sequence_parallel_and_cp.py b/tests/transformers/test_sequence_parallel_and_cp.py index a2ce49e61..d427c114f 100644 --- a/tests/transformers/test_sequence_parallel_and_cp.py +++ b/tests/transformers/test_sequence_parallel_and_cp.py @@ -13,11 +13,10 @@ from types import SimpleNamespace from twinkle.loss import CrossEntropyLoss -from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy, sequence_parallel +from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy from twinkle.utils import DeviceMesh, Platform, ensure_hccl_socket_env, selective_log_softmax, torch_util -pytestmark = pytest.mark.skip( - reason='Heavy SP/CP test (mp.spawn, world_size up to 4); not viable on dual-V100 CI, run manually.') + LOGITS_RTOL = 1e-2 LOGITS_ATOL = 5e-3 @@ -263,7 +262,7 @@ def _make_strategy(model: LlamaForCausalLM, device_mesh: DeviceMesh, ulysses_siz def _prepare_label_inputs(strategy: SequenceParallelStrategy, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: labels = _make_labels(input_ids) - processed = strategy.preprocess_inputs({ + processed = strategy.prepare_model_inputs({ 'input_ids': input_ids, 'position_ids': position_ids, 'labels': labels, @@ -306,6 +305,7 @@ def _compute_training_path_loss( loss_inputs = {'labels': labels} loss_outputs = {'logps': selective_log_softmax(logits, masked_labels)} if strategy is not None: + loss_inputs['_sequence_layout'] = {'sharded': ('labels',)} loss_inputs, loss_outputs = strategy.gather_loss_tensors(loss_inputs, loss_outputs) result = CrossEntropyLoss(reduction='sum')(loss_inputs, loss_outputs) num_tokens = result['num_tokens'] @@ -421,12 +421,10 @@ def _measure_peak_memory( device_module = _get_device_module(device_type) device_module.empty_cache() device_module.reset_peak_memory_stats(device) - outputs = model( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=None, - use_cache=False, - ) + prepared = strategy.prepare_model_inputs({'input_ids': input_ids, 'position_ids': position_ids}) + prepared.pop('_sequence_layout') + with strategy.context(): + outputs = model(**prepared, use_cache=False) local_logits = outputs.logits loss_sum, _ = _compute_training_path_loss(local_logits, local_labels, strategy) loss_sum.backward() @@ -519,28 +517,25 @@ def _run_precision_worker(rank: int, device_type=backend_config['device_type'], ) strategy = _make_strategy(sp_model, device_mesh, int(case['ulysses_size'])) - processed_inputs = strategy.preprocess_inputs({ - 'input_ids': input_ids, - 'position_ids': position_ids, - 'labels': labels, + processed_inputs = strategy.prepare_model_inputs({ + 'input_ids': input_ids, 'position_ids': position_ids, 'labels': labels, + 'attention_mask': attention_mask, }) local_labels = processed_inputs['labels'] sp_model.zero_grad(set_to_none=True) - sp_outputs = sp_model( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=attention_mask, - use_cache=False, - ) + processed_inputs.pop('labels') + processed_inputs.pop('_sequence_layout') + with strategy.context(): + sp_outputs = sp_model(**processed_inputs, use_cache=False) local_logits = sp_outputs.logits - gathered_outputs = strategy.postprocess_outputs(sp_outputs) + _, gathered_outputs = strategy.gather_loss_tensors({}, sp_outputs) sp_logits = gathered_outputs.logits.detach().float() sp_loss_sum, sp_num_tokens = _compute_training_path_loss(local_logits, local_labels, strategy) sp_loss_sum.backward() global_loss = sp_loss_sum / max(sp_num_tokens, 1) - _average_model_grads_over_group(sp_model, sequence_parallel._data_rank_group) + _average_model_grads_over_group(sp_model, strategy._ctx._data_rank_group) sp_attention_grads = _normalize_grad_dict(_collect_attention_param_grads(sp_model), sp_num_tokens) if sp_num_tokens != base_num_tokens: @@ -653,6 +648,8 @@ def _run_memory_worker(rank: int, world_size: int, port: int, case_name: str, ba dist.destroy_process_group() +@pytest.mark.slow +@pytest.mark.accel(2) class TestDerivedRingPrecision: def _get_backend_or_skip(self, world_size: int = 4) -> dict: @@ -720,6 +717,8 @@ def test_cp_sp_precision_alignment(self): join=True) +@pytest.mark.slow +@pytest.mark.accel(2) class TestDerivedRingMemoryProfile: def _get_backend_or_skip(self, world_size: int = 4) -> dict: @@ -760,3 +759,333 @@ def test_cp_sp_memory_profile_grid(self): backend = self._get_backend_or_skip(world_size) port = _find_free_port() mp.spawn(_run_memory_worker, args=(world_size, port, 'cp_sp_memory', backend), nprocs=world_size, join=True) + + +def test_explicit_sp_context_restores_nested_and_exception_paths(): + from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallel, _current_ctx, _attention_context + first, second = SequenceParallel(), SequenceParallel() + first.model_name = 'policy' + assert _current_ctx() is None + with first.context(): + with pytest.raises(RuntimeError, match='unregistered attention'): + _attention_context(torch.nn.Linear(2, 2)) + with pytest.raises(ValueError, match='forward failed'): + with second.context(): + assert _current_ctx() is second + raise ValueError('forward failed') + assert _current_ctx() is first + assert _current_ctx() is None + + +def test_flash_padding_metadata_preserves_real_packed_boundaries(): + from twinkle.model.transformers.strategy.sequence_parallel import _call_flash_attention + positions = torch.tensor([[0, 1, 2, -1, 0, 1, -1, -1]]) + original = positions.clone() + normalized = _call_flash_attention(lambda **kw: kw['position_ids'], position_ids=positions) + assert normalized.tolist() == [[0, 1, 2, 3, 0, 1, 2, 3]] + # Equivalent boundary rule used by TF 5.15.1; the tail must not become a new sequence. + assert (normalized == normalized.min()).nonzero()[:, -1].tolist() == [0, 4] + torch.testing.assert_close(positions, original) + + +def _slim_cpu_worker(rank, port, checkpoint): + import torch.nn.functional as F + from torch.utils.checkpoint import DefaultDeviceType + from twinkle.processor import InputProcessor + from twinkle.model.transformers.strategy.sequence_parallel import _current_ctx, SequenceParallel + DefaultDeviceType.set_device_type('cpu') + torch.set_num_threads(1) + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, world_size=2) + try: + torch.manual_seed(42) + config = LlamaConfig(vocab_size=32, hidden_size=16, intermediate_size=32, + num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=4, + attention_dropout=0., use_cache=False) + config._attn_implementation = 'sdpa' + base_p, base_r = LlamaForCausalLM(config).train(), LlamaForCausalLM(config).train() + policy, reference = copy.deepcopy(base_p), copy.deepcopy(base_r) + if checkpoint is not None: + for model in (base_p, base_r, policy, reference): + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': checkpoint}) + mesh = DeviceMesh.from_sizes(fsdp_size=2, dp_size=1, ulysses_size=2, device_type='cpu') + policy_sp, ref_sp = _make_strategy(policy, mesh, 2), _make_strategy(reference, mesh, 2) + inputs_p = torch.tensor([[1, 3, 5, 7, 9, 11, 13]]) + inputs_r = torch.tensor([[1, 4, 6, 8, 10, 12, 14, 16, 18]]) + + def baseline(model, ids): + logits = model(input_ids=ids, use_cache=False).logits + labels = _make_labels(ids) + loss = F.cross_entropy(logits.reshape(-1, 32), labels.reshape(-1), ignore_index=-100) + return logits, loss + + def parallel(strategy, model, ids): + inputs = strategy.prepare_model_inputs({'input_ids': ids, 'labels': _make_labels(ids)}) + labels, layout = inputs.pop('labels'), inputs.pop('_sequence_layout') + with strategy.context(): + output = model(**inputs, use_cache=False) + inputs.update(labels=labels, _sequence_layout=layout) + inputs, output = strategy.gather_loss_tensors(inputs, output) + loss = F.cross_entropy(output.logits.reshape(-1, 32), inputs['labels'].reshape(-1), ignore_index=-100) + return output.logits, loss + + baseline_p, loss_p = baseline(base_p, inputs_p) + baseline_r, loss_r = baseline(base_r, inputs_r) + logits_p, parallel_p = parallel(policy_sp, policy, inputs_p) + logits_r, parallel_r = parallel(ref_sp, reference, inputs_r) + # Reuse a context with another batch before the first backward: recompute must restore its own positions. + parallel(policy_sp, policy, inputs_r) + for actual, expected in ((logits_p, baseline_p), (logits_r, baseline_r), + (parallel_p, loss_p), (parallel_r, loss_r)): + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-6) + parallel_r.backward() + parallel_p.backward() + loss_r.backward() + loss_p.backward() + for model, base in ((policy, base_p), (reference, base_r)): + for (_, parameter), (_, original) in zip(model.named_parameters(), base.named_parameters()): + assert parameter.grad is not None and original.grad is not None + dist.all_reduce(parameter.grad) + parameter.grad.div_(2) + torch.testing.assert_close(parameter.grad, original.grad, rtol=3e-4, atol=3e-6) + assert _current_ctx() is None + # The non-SP baseline still uses the origin after all global wrappers have been installed. + again, _ = baseline(base_p, inputs_p) + torch.testing.assert_close(again, baseline_p, rtol=0, atol=0) + + processor = InputProcessor() + # No logps and no labels: critic-style outputs must still gather with autograd. + projection = torch.nn.Linear(3, 1, bias=False) + full = torch.randn(1, 7, 3) + original = projection(full) + expected_loss = original.square().mean() + expected_loss.backward() + expected_grad = projection.weight.grad.clone() + projection.zero_grad() + prepared = policy_sp.prepare_model_inputs({'input_ids': inputs_p}) + indices = prepared['_sequence_layout']['indices'] + padded = F.pad(full, (0, 0, 0, 1)) + local = padded.gather(1, indices.clamp_min(0).unsqueeze(-1).expand(-1, -1, 3)) + _, gathered = policy_sp.gather_loss_tensors(prepared, {'values': projection(local)}) + torch.testing.assert_close(gathered['values'], original, rtol=1e-6, atol=1e-7) + gathered['values'].square().mean().backward() + dist.all_reduce(projection.weight.grad) + projection.weight.grad.div_(2) + torch.testing.assert_close(projection.weight.grad, expected_grad, rtol=1e-6, atol=1e-7) + + # Both Ulysses and zigzag Ring layouts; pure communication, no Flash kernel substitution. + for ring in (False, True): + strategy = SequenceParallelStrategy(mesh, {'ulysses_size': 2}) + ctx = strategy._ctx + ctx.seq_world_size = ctx.world_size = 2 + ctx.sp_world_size, ctx.rp_world_size = (1, 2) if ring else (2, 1) + ctx._init_device_mesh(mesh) + ctx.tokenizer = SimpleNamespace(pad_token_id=0) + ctx.attn_implementation = 'flash_attention_2' # packed layout, not an attention-kernel test + for positions in (torch.tensor([[-1, 0, 1, -1, 3, 4, -1]]), + torch.tensor([[0, 1, 0, 1, 2, 0, 1]])): + if ring and positions[0, 0] < 0: + with pytest.raises(NotImplementedError, match='left padding'): + strategy.prepare_model_inputs({'input_ids': inputs_p, 'position_ids': positions}) + continue + inputs = {'input_ids': inputs_p, 'position_ids': positions, + 'labels': torch.arange(7).unsqueeze(0), + 'completion_mask': torch.tensor([[0, 1, 0, 1, 1, 1, 1]]), + 'loss_scale': torch.linspace(0., 1., 7).unsqueeze(0)} + expected_input = {'position_ids': positions} + full = torch.randn(1, 7, 3) + projection = torch.nn.Linear(3, 2, bias=False) + expected, _ = processor._pool_last_valid_token(projection(full), expected_input) + expected.square().sum().backward() + expected_grad = projection.weight.grad.clone() + projection.zero_grad() + prepared = strategy.prepare_model_inputs(inputs) + padded = ctx.pad(full, padding_value=0., position_ids=positions) + padded_positions = ctx.pad(positions, padding_value=-1, position_ids=positions) + local = ctx.split(padded, dim=1, position_ids=padded_positions) + pooled, _ = processor._pool_last_valid_token(projection(local), prepared, + reduce_fn=strategy.all_reduce) + torch.testing.assert_close(pooled, expected, rtol=1e-6, atol=1e-7) + restored, _ = strategy.gather_loss_tensors(prepared, {}) + for name in ('labels', 'completion_mask', 'loss_scale'): + torch.testing.assert_close(restored[name], inputs[name]) + pooled.square().sum().backward() + dist.all_reduce(projection.weight.grad) + projection.weight.grad.div_(2) + torch.testing.assert_close(projection.weight.grad, expected_grad, rtol=1e-5, atol=1e-6) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize('checkpoint', [None, False, True]) +def test_slim_sp_cpu_logits_gradients_values_pooling_and_recompute(checkpoint): + mp.spawn(_slim_cpu_worker, args=(_find_free_port(), checkpoint), nprocs=2, join=True) + + +def _slim_qwen_vl_worker(rank, port): + from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + torch.set_num_threads(1) + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, world_size=2) + try: + torch.manual_seed(81) + import transformers + text_config = dict(vocab_size=64, hidden_size=32, intermediate_size=64, num_hidden_layers=1, + num_attention_heads=2, num_key_value_heads=2, head_dim=16) + # 4.x names this test configuration rope_scaling; 5.x uses rope_parameters. + rope_key = 'rope_scaling' if int(transformers.__version__.split('.')[0]) < 5 else 'rope_parameters' + text_config[rope_key] = {'rope_type': 'default', 'rope_theta': 10000., 'mrope_section': [2, 3, 3]} + config = Qwen3VLConfig( + text_config=text_config, + vision_config=dict(depth=1, hidden_size=32, intermediate_size=64, num_heads=4, + patch_size=2, spatial_merge_size=2, temporal_patch_size=1, + out_hidden_size=32, num_position_embeddings=16, deepstack_visual_indexes=[0]), + image_token_id=60, video_token_id=61, vision_start_token_id=59, vision_end_token_id=58) + config._attn_implementation = 'sdpa' + baseline = Qwen3VLForConditionalGeneration(config).train() + model = torch.nn.parallel.DistributedDataParallel(copy.deepcopy(baseline)) + mesh = DeviceMesh.from_sizes(fsdp_size=2, dp_size=1, ulysses_size=2, device_type='cpu') + strategy = _make_strategy(model, mesh, 2) + ids = torch.tensor([[1, 59, 60, 60, 60, 60, 58, 5, 6]]) + inputs = dict(input_ids=ids, pixel_values=torch.randn(16, 12), image_grid_thw=torch.tensor([[1, 4, 4]]), + mm_token_type_ids=torch.tensor([[0, 0, 1, 1, 1, 1, 0, 0, 0]]), use_cache=False) + expected = baseline(**inputs).logits + prepared = strategy.prepare_model_inputs(inputs) + prepared.pop('_sequence_layout') + assert 'position_ids' not in prepared + with strategy.context(): + output = model(**prepared) + _, gathered = strategy.gather_loss_tensors({}, output) + torch.testing.assert_close(gathered.logits, expected, rtol=3e-5, atol=3e-6) + expected.square().mean().backward() + gathered.logits.square().mean().backward() + for (name, parameter), (_, original) in zip(model.named_parameters(), baseline.named_parameters()): + if original.grad is None: + assert parameter.grad is None + continue + assert parameter.grad is not None, name + dist.all_reduce(parameter.grad) + parameter.grad.div_(2) + torch.testing.assert_close(parameter.grad, original.grad, rtol=3e-4, atol=3e-6, msg=name) + finally: + dist.destroy_process_group() + + +def test_qwen3_vl_vision_stays_local_and_text_is_split_after_visual_merge(): + from transformers.utils import is_torch_available + assert is_torch_available() + mp.spawn(_slim_qwen_vl_worker, args=(_find_free_port(),), nprocs=2, join=True) + + +def test_sp_weak_registry_does_not_retain_the_model(): + import gc + import weakref + from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallel + model = LlamaForCausalLM(LlamaConfig(vocab_size=16, hidden_size=8, intermediate_size=16, + num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2)) + ref = weakref.ref(model) + ctx = SequenceParallel() + ctx._prepare_forward_boundaries(model, model.model) + ctx.causal_mask_func = model.forward + del model + gc.collect() + assert ref() is None + assert ctx.causal_mask_func is None + + +@pytest.mark.parametrize('backend', ['eager', 'flash_attention_2', 'flash_attention_3']) +def test_sp_rejects_unsupported_backend_before_building_groups(backend): + from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallel + model = LlamaForCausalLM(LlamaConfig(vocab_size=16, hidden_size=8, intermediate_size=16, + num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2)) + model.config._attn_implementation = backend + with pytest.raises(NotImplementedError, match='unsupported attention|kernel unavailable'): + SequenceParallel().prepare(2, model, SimpleNamespace(pad_token_id=0)) + + +def test_single_local_token_value_keeps_its_sequence_axis(): + from twinkle.model.transformers.transformers import TransformersModel + outputs = {'logits': torch.ones(2, 1)} + TransformersModel._postprocess_forward_outputs( + None, outputs, labels=None, task='value', temperature=1., sampling_masks=None, + return_logits=False, loss_instance=SimpleNamespace(require_values=True)) + assert outputs['values'].shape == (2, 1) + + +def test_sentence_loss_metadata_stays_global_while_model_tokens_split(): + strategy = SequenceParallelStrategy(sp_config={'ulysses_size': 2}) + ctx = strategy._ctx + ctx.world_size = ctx.sp_world_size = 2 + ctx.rp_world_size = 1 + ctx.attn_implementation = 'sdpa' + ctx.tokenizer = SimpleNamespace(pad_token_id=0) + labels, weights = torch.tensor([[1]]), torch.tensor([[0.3]]) + inputs = strategy.prepare_model_inputs( + {'input_ids': torch.arange(7).unsqueeze(0), 'labels': labels, 'loss_scale': weights}, + token_labels=False) + assert inputs['input_ids'].shape == (1, 4) + assert inputs['labels'] is labels and inputs['loss_scale'] is weights + assert inputs['_sequence_layout']['sharded'] == [] + + +def _heterogeneous_context_cpu_worker(rank, port): + import torch.nn.functional as F + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + torch.set_num_threads(1) + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, world_size=4) + try: + torch.manual_seed(73) + cases = [] + for sp_size, heads, length in ((4, 4, 7), (2, 2, 9)): + config = LlamaConfig(vocab_size=32, hidden_size=16, intermediate_size=32, + num_hidden_layers=1, num_attention_heads=heads, num_key_value_heads=heads, + attention_dropout=0., use_cache=False) + config._attn_implementation = 'sdpa' + base = LlamaForCausalLM(config) + model = copy.deepcopy(base) + ids = (torch.arange(length) + 1).unsqueeze(0) + labels = _make_labels(ids) + baseline = base(input_ids=ids, use_cache=False).logits + base_loss = F.cross_entropy(baseline.flatten(0, 1), labels.flatten(), ignore_index=-100) + mesh = DeviceMesh.from_sizes(fsdp_size=4, dp_size=1, ulysses_size=sp_size, device_type='cpu') + strategy = _make_strategy(model, mesh, sp_size) + prepared = strategy.prepare_model_inputs({'input_ids': ids, 'labels': labels}) + local_labels, layout = prepared.pop('labels'), prepared.pop('_sequence_layout') + with strategy.context(): + output = model(**prepared, use_cache=False) + prepared.update(labels=local_labels, _sequence_layout=layout) + prepared, output = strategy.gather_loss_tensors(prepared, output) + torch.testing.assert_close(output.logits, baseline, rtol=2e-5, atol=2e-6) + loss = F.cross_entropy(output.logits.flatten(0, 1), prepared['labels'].flatten(), ignore_index=-100) + torch.testing.assert_close(loss, base_loss, rtol=2e-5, atol=2e-6) + cases.append((strategy, model, base, loss, base_loss)) + assert cases[0][0]._ctx is not cases[1][0]._ctx + for strategy, model, base, loss, base_loss in reversed(cases): + loss.backward() + base_loss.backward() + for parameter, expected in zip(model.parameters(), base.parameters()): + dist.all_reduce(parameter.grad) + parameter.grad.div_(4) + torch.testing.assert_close(parameter.grad, expected.grad, rtol=3e-4, atol=3e-6) + finally: + dist.destroy_process_group() + + +@pytest.mark.slow +def test_heterogeneous_sp4_sp2_contexts_logits_loss_and_parameter_gradients_cpu(): + mp.spawn(_heterogeneous_context_cpu_worker, args=(_find_free_port(),), nprocs=4, join=True) + + +def test_disabled_strategy_context_stays_local_and_failed_registration_is_atomic(): + from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallel, _current_ctx, _MODULE_TO_CTX + strategy = SequenceParallelStrategy(sp_config={'enabled': False, 'ulysses_size': 2}) + with strategy.context(): + assert _current_ctx() is None + model = torch.nn.Linear(2, 2) + ctx = SequenceParallel() + with pytest.raises(NotImplementedError, match='text attention modules'): + ctx._prepare_forward_boundaries(model, model) + assert model not in _MODULE_TO_CTX + with pytest.raises(RuntimeError, match='Initialize sequence parallel'): + SequenceParallelStrategy(sp_config={'ulysses_size': 2}).context()