diff --git a/docs/source_en/Components/Training Middleware/Sequence-Parallel.md b/docs/source_en/Components/Training Middleware/Sequence-Parallel.md index d08b01d67..43e4442b5 100644 --- a/docs/source_en/Components/Training Middleware/Sequence-Parallel.md +++ b/docs/source_en/Components/Training Middleware/Sequence-Parallel.md @@ -1,68 +1,97 @@ # Sequence Parallel (SP) -Sequence Parallel splits long sequences across multiple GPUs along the sequence dimension, enabling training with sequence lengths that exceed single-GPU memory. Twinkle implements Ulysses-style sequence parallel with optional derived ring attention. - -## Overview - -| Concept | Description | -|---------|-------------| -| **SequenceParallelConfig** | Configuration dataclass for SP | -| **SequenceParallelStrategy** | Strategy class that wraps SP lifecycle | -| **SequenceParallel** | Core implementation handling pad/split/gather | +Transformers SP shards token sequences and runs attention through Ulysses and optional Ring communication. Each model owns an independent `SequenceParallel` context. Process-wide wrappers resolve that context at call time. ## Configuration ```python from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelConfig -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 -) +config = SequenceParallelConfig(enabled=True, ulysses_size=None) ``` -## Usage with DeviceMesh - -SP is activated by setting `ulysses_size` in `DeviceMesh.from_sizes()`: +A `None` degree is derived from DeviceMesh. The old `gather_logits` option has been removed: output rules and the current loss determine collection. ```python from twinkle.utils import DeviceMesh -# 8 GPUs: 4-way Ulysses SP × 2-way data parallel -device_mesh = DeviceMesh.from_sizes( - world_size=8, - dp_size=2, - ulysses_size=4, +# Eight devices forming two SP=4 groups. +mesh = DeviceMesh.from_sizes(dp_size=8, fsdp_size=1, ulysses_size=4) +``` + +## Processing responsibilities + +1. Processor converts tensors, collates samples, applies ordinary padding and exposes `is_padding_free()` as a layout query. It neither receives nor calls an SP strategy. +2. The model wrapper partitions model inputs and loss inputs. The strategy shards labels and explicitly local auxiliary fields at the prepare boundary. +3. Forward hooks pad/split model token inputs at their declared boundaries. Attention masks retain global layout, including padding masks for single-sample SDPA batches. +4. Attention wrappers perform communication and invoke the configured backend. +5. Model-side `postprocess_task_outputs()` handles task semantics: embedding pooling/normalization, classification/reranker pooling, or token-loss SP collection. +6. Processor's `unpack_packed_sequences()` restores already collected sequences to per-sample layouts before loss calculation. + +Model postprocessing interprets task outputs and performs SP collectives; processor postprocessing restores layouts. + +## Field ownership + +- `data_format/sequence_fields.py` defines `DataKeySpec` and provider contracts, without a global list of all fields. +- `processor/fields.py` owns supported encoded text/media/replay/control inputs and collation/restoration rules. Templates encode raw media first. Internal `length` metadata is not a forward argument. +- `model/transformers/strategy/sequence_parallel/data_keys.py` derives SP execution defaults from the common processor layouts, adding processing boundaries and autograd settings without repeating axes or padding values. It also owns batch-local collections. Media pixels and control arguments are not blindly sharded as tokens. +- Model adapters and losses declare their own extensions. + +Processor defaults do not schedule SP operations. Custom loss fields remain available with SP disabled. + +Declare extensions on the model, adapter or loss; adding a custom key does not require editing either built-in list. Model-side schema resolution is cached for the current model/strategy/loss instances. Direct input/output registrations and `set_loss()` invalidate the cache. Configure provider declarations before forwarding; if reconfiguring an existing loss instance, call `set_loss()` again. Field rules only pad, split and gather; they do not shift token targets. + +```python +from twinkle.data_format.sequence_fields import DataKeySpec + +# Global loss input, restored per sample after packing. +input_key_specs = (DataKeySpec('token_targets', pad_value=-100),) + +# Explicitly scheduled model-forward input. +model_input_key_specs = ( + DataKeySpec('token_features', split_dim=1, stage='forward', + pad_at='forward', split_at='forward'), ) ``` -## How It Works +`pad_at/split_at` accept `prepare`, `forward`, `model_internal`, `router_replay`, or `None`. `need_pad/need_split` are derived properties. `stage` classifies fields (including `auxiliary` and `output`), rather than naming the execution boundary. + +## Loss declarations + +CE uses the default collected `logps`. Other objectives may need values or custom sequence outputs; SP cannot infer their requirements from arbitrary Python code. -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 +- `input_key_specs/output_key_specs` describe layout, routing and optional parallel operations. Processor consumes layouts even without SP. +- `sp_gather_keys` lists full-sequence outputs required by the current loss call. Missing required outputs raise an error. A field's `gather_at='loss'` instead requests default collection when that output is present. +- `require_values` requests per-token `values` from the Transformers wrapper; PPOValueLoss then collects them. +- Earlier `sp_input_key_specs/sp_output_key_specs` names are still read for compatibility. New code uses the common names. + +PPOValueLoss declares: + +```python +require_logps = False +require_values = True +sp_gather_keys = ('values',) +output_key_specs = ( + DataKeySpec('values', split_dim=1, side='output', stage='output', autograd=True), +) +``` -## Supported Attention Backends +Schema defines tensor handling; required-output declarations specify this call's needs. Full vocabulary logits are not gathered unconditionally. Fused CE handles its hidden-state slot separately. Global completion masks, loss weights and default custom targets do not incur an unnecessary split/gather cycle. -| Backend | Status | -|---------|--------| -| FlashAttention2 | Fully supported (including packed/padding-free sequences) | -| SDPA | Supported (non-packed batches only) | -| Derived Ring Attention | Supported with FlashAttention2 only (`rp_world_size > 1`) | +## Attention backends -## Qwen3.5 Linear Attention +| Backend | Implementation scope | +|---|---| +| FlashAttention2 / FlashAttention3 | Shared SP wrapper and packed/varlen preparation; concrete kernels follow model configuration and require their own dependencies/device capabilities | +| SDPA | Ordinary non-packed batches; padding-free/packed inputs are rejected | +| Derived Ring | FA2/FA3 dispatch is implemented and requires batch_size=1; the FA3 path rejects dropout and ALiBi | -SP automatically detects Qwen3.5 GatedDeltaNet linear attention layers and applies the `Qwen3_5GatedDeltaNetUlyssesPatch` for correct sequence-parallel behavior on hybrid attention architectures. +Implementation support does not establish runtime acceptance on every device. Record FA3 validation separately from FA2. -## MoE Auxiliary Loss +## Model adapters and Megatron -For MoE models, SP automatically installs a forward hook that gathers router logits across SP ranks before auxiliary loss computation, ensuring correct load-balancing signals. +Adapters configure Qwen3-VL DeepStack/mRoPE, Qwen3.5 GatedDeltaNet and MoE router auxiliary loss using the owning model's context. -## Key Constraints +This page describes Transformers SP. Megatron has its own CP/TP-SP path: `MegatronEmbeddingPatch` handles embeddings, Megatron forward_step handles classification/reranking, and `postprocess_tensor_cp()` reconstructs CP shards. Both paths share processor collation and packed restoration. Shared-path regression does not replace full Megatron training validation. -- `num_key_value_heads` must be divisible by `ulysses_size` (for Ulysses) or use ring attention fallback -- Packed/padding-free batches require FlashAttention2 -- Derived ring attention requires `batch_size == 1` (packed format) -- `torch.distributed` must be initialized +Initialize `torch.distributed` before distributed execution. Attention-head count and the configured SP degree determine the Ulysses/Ring topology. 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..5c24832f0 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" @@ -1,68 +1,97 @@ # 序列并行 (SP) -序列并行沿序列维度将长序列分割到多个 GPU 上，使训练能处理超出单卡显存的序列长度。Twinkle 实现了 Ulysses 风格的序列并行，并可选地支持派生环形注意力。 - -## 概览 - -| 概念 | 说明 | -|------|------| -| **SequenceParallelConfig** | SP 配置数据类 | -| **SequenceParallelStrategy** | 封装 SP 生命周期的策略类 | -| **SequenceParallel** | 核心实现，处理填充/分割/聚合 | +Transformers SP 沿 token 序列维度切分输入，通过 Ulysses 和可选的 Ring 通信完成注意力计算。每个模型持有独立的 `SequenceParallel` context；进程级包装器在调用时选择当前模型的 context。 ## 配置 ```python from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelConfig -config = SequenceParallelConfig( - enabled=True, # 启用序列并行 - ulysses_size=None, # Ulysses SP 并行度（若为 None 则从 DeviceMesh 自动推导） - gather_logits=True, # 前向后聚合 logits 用于损失计算 -) +config = SequenceParallelConfig(enabled=True, ulysses_size=None) ``` -## 配合 DeviceMesh 使用 - -在 `DeviceMesh.from_sizes()` 中设置 `ulysses_size` 即可激活 SP： +`ulysses_size=None` 时从 DeviceMesh 推导。旧的 `gather_logits` 配置已删除，是否收集完整输出由输出规则与当前 loss 的声明决定。 ```python from twinkle.utils import DeviceMesh -# 8 卡：4 路 Ulysses SP × 2 路数据并行 -device_mesh = DeviceMesh.from_sizes( - world_size=8, - dp_size=2, - ulysses_size=4, +# 8 个设备，两个 SP=4 组。 +mesh = DeviceMesh.from_sizes(dp_size=8, fsdp_size=1, ulysses_size=4) +``` + +## 处理流程与职责 + +1. Processor 处理张量转换、拼批、普通 padding，并通过 `is_padding_free()` 提供 batch 布局查询；它不接收或调用 SP strategy。 +2. 模型包装层区分模型输入与 loss 输入；strategy 在 prepare 阶段切分 labels 等明确声明的局部字段。 +3. 模型 forward hook 按规则 pad/split input_ids、inputs_embeds、position_ids 及扩展模型输入。attention_mask 保留全局布局；SDPA 对单样本也构建或补齐全局 padding mask。 +4. 注意力包装器进行 Ulysses/Ring 通信，并调用配置选择的注意力后端。 +5. 模型侧 `postprocess_task_outputs()` 完成任务处理：embedding 的末 token 池化与归一化、分类/reranker 池化，或普通 token loss 所需的 SP 汇聚。 +6. Processor 的 `unpack_packed_sequences()` 将已经恢复全局序列的张量拆回逐样本布局，供 loss 使用。 + +这里的两类 postprocess 有不同职责：模型侧解释任务输出和执行 SP 通信；Processor 只恢复数据布局。 + +## 字段声明的归属 + +- `data_format/sequence_fields.py` 定义 `DataKeySpec` 类型和 provider 声明接口，不维护全局字段大全。 +- `processor/fields.py` 就近维护 Processor 已支持的文本、多模态、router replay、控制参数和拼批/恢复规则。raw 图片/视频先由模板编码；内部 `length` 元数据不传给模型。 +- `model/transformers/strategy/sequence_parallel/data_keys.py` 从 Processor 的常用布局声明派生 SP 执行规则，只增加处理边界和 autograd 设置，不重复维护序列轴与 padding 值；同时维护 batch-local 集合，不把图片像素或控制参数当成 token 序列切分。 +- 模型 adapter 和 loss 分别声明自己的扩展字段。 + +Processor 的规则不安排 SP 操作。SP 开关不会改变自定义 loss 字段是否被保留、拼批或分区。 + +新增字段由 model、adapter 或 loss 声明，不需要修改两份内置列表。模型侧按当前 model/strategy/loss 实例缓存字段解析结果；直接注册新输入/输出键或调用 `set_loss()` 后重新解析。provider 声明应在首次 forward 前配置；若修改同一个 loss 实例的声明，应再次调用 `set_loss()`。字段规则只负责 pad、split、gather，不负责 token 目标的移位。 + +```python +from twinkle.data_format.sequence_fields import DataKeySpec + +# loss-only 输入，默认保持全局；packed 时按样本恢复。 +input_key_specs = (DataKeySpec('token_targets', pad_value=-100),) + +# 模型 forward 参数，明确安排 SP 操作。 +model_input_key_specs = ( + DataKeySpec('token_features', split_dim=1, stage='forward', + pad_at='forward', split_at='forward'), ) ``` -## 工作原理 +`pad_at/split_at` 可取 `prepare`、`forward`、`model_internal`、`router_replay` 或 `None`；`need_pad/need_split` 是只读推导属性。`stage` 是字段分类，包含 `auxiliary` 和 `output`，不等同于执行时机。 + +## Loss 为什么需要声明 + +普通 CE 使用默认收集的 `logps`。PPO value loss 或自定义 loss 可能读取其他序列输出，SP 不能从任意 Python loss 代码中推断这些需求。 -1. **填充** — 输入序列被填充到可被 SP 并行度整除的长度 -2. **分割** — 填充后的输入沿序列维度均匀分配到各 SP rank -3. **分布式注意力** — FlashAttention2 被 patch 为在注意力计算前后执行 Ulysses all-to-all 通信 -4. **聚合** — 前向传播后，logits 被聚合回完整序列长度用于损失计算 +- `input_key_specs/output_key_specs` 描述字段布局、模型参数归属，以及可选的并行操作；即使 SP 关闭，Processor 仍消费这些布局规则。 +- `sp_gather_keys` 指定当前 loss 必须拿到完整序列的输出。它是当前调用的必需项，缺失输出会报错；字段规则的 `gather_at='loss'` 则表示该输出存在时执行默认汇聚。 +- `require_values` 告诉 Transformers 包装层生成逐 token 的 `values`，随后 PPOValueLoss 请求收集它。 +- `sp_input_key_specs/sp_output_key_specs` 保留此前声明名称的兼容读取；新代码使用公共名称。 + +PPOValueLoss 使用的声明如下： + +```python +require_logps = False +require_values = True +sp_gather_keys = ('values',) +output_key_specs = ( + DataKeySpec('values', split_dim=1, side='output', stage='output', autograd=True), +) +``` -## 支持的注意力后端 +字段 schema 定义如何处理张量，当前 loss 的必需输出声明定义这次计算需要什么。不会无条件 gather 完整词表 logits；fused CE 的 hidden-state 路径单独处理。只在最终 loss 使用的 completion_mask、loss_scale 及默认自定义 target 保持全局。 -| 后端 | 状态 | -|------|------| -| FlashAttention2 | 完全支持（包括打包/padding-free 序列）| -| SDPA | 支持（仅非打包批次）| -| 派生环形注意力 | 仅支持 FlashAttention2（`rp_world_size > 1`）| +## 注意力后端 -## Qwen3.5 线性注意力 +| 后端 | 实现范围 | +|---|---| +| FlashAttention2 / FlashAttention3 | 共用 SP 包装和 packed/varlen 数据处理，实际内核由模型配置选择；具体依赖和设备能力需分别满足 | +| SDPA | 支持普通非 packed batch；padding-free/packed 输入明确拒绝 | +| 派生 Ring | 已接入 FA2/FA3 后端分发，要求 batch_size=1；FA3 路径不支持 dropout 和 ALiBi | -SP 自动检测 Qwen3.5 GatedDeltaNet 线性注意力层，并应用 `Qwen3_5GatedDeltaNetUlyssesPatch`，确保混合注意力架构下序列并行的正确性。 +代码接入不代表每种设备/依赖组合均通过训练验收。FA3 的运行时验证应单独记录；不能用 FA2 的测试结果代替。 -## MoE 辅助损失 +## 模型特殊处理与 Megatron -对于 MoE 模型，SP 自动安装前向 hook，在计算辅助损失前跨 SP rank 聚合路由 logits，确保负载均衡信号的正确性。 +Qwen3-VL 的 DeepStack/mRoPE、Qwen3.5 的 GatedDeltaNet、MoE router auxiliary loss 由对应 adapter 配置，使用所属模型的 context。 -## 关键约束 +本页描述 Transformers SP。Megatron 的 CP/TP-SP 是另一条执行路径：embedding 仍由 `MegatronEmbeddingPatch` 池化，分类/reranker 仍在 Megatron 的 forward_step 中处理；CP 重建继续使用 `postprocess_tensor_cp()`。两者共用 Processor 的拼批和 packed 恢复，因此共享路径需要单独回归，不能用 Transformers 训练通过代替 Megatron 整网验收。 -- `num_key_value_heads` 必须能被 `ulysses_size` 整除（Ulysses 模式），否则回退到环形注意力 -- 打包/padding-free 批次需要 FlashAttention2 -- 派生环形注意力要求 `batch_size == 1`（打包格式） -- `torch.distributed` 必须已初始化 +分布式运行前需要初始化 `torch.distributed`；SP 大小与注意力头数共同决定 Ulysses/Ring 拓扑。 diff --git a/src/twinkle/data_format/sequence_fields.py b/src/twinkle/data_format/sequence_fields.py new file mode 100644 index 000000000..ae6481aed --- /dev/null +++ b/src/twinkle/data_format/sequence_fields.py @@ -0,0 +1,95 @@ +"""Sequence-field declarations shared by data processors and model/loss providers.""" + +from dataclasses import dataclass + +_BOUNDARIES = ("forward", "auxiliary", "model_internal", "router_replay", "output") +_PROCESSING_BOUNDARIES = ("prepare", "forward", "model_internal", "router_replay") + + +@dataclass(frozen=True) +class DataKeySpec: + """Describe collation, routing and explicitly scheduled parallel operations. + + Loss-only fields stay global unless pad_at/split_at opt into local processing. + Packed restoration is independent of sequence-parallel gather and split. + """ + + name: str + split_dim: int = -1 + pad_value: int | float | str = 0 + side: str = "input" + stage: str = "auxiliary" + autograd: bool = False + pad_at: str | None = None + split_at: str | None = None + gather_at: str | None = None + unpack_packed: bool | None = None + model_input: bool | None = None + + def __post_init__(self): + if not isinstance(self.name, str) or not self.name: + raise ValueError("A sequence field must have a non-empty name") + if not isinstance(self.split_dim, int): + raise TypeError("split_dim must be an integer") + if self.side not in ("input", "output"): + raise ValueError("side must be input or output") + if self.stage not in _BOUNDARIES: + raise ValueError("Invalid SP key stage") + if (self.side == "output") != (self.stage == "output"): + raise ValueError("Output keys must use stage=output") + if isinstance(self.pad_value, str) and (self.name != "input_ids" or self.pad_value != "tokenizer"): + raise ValueError("Only input_ids can use the tokenizer padding sentinel") + for operation, boundary in (("padding", self.pad_at), ("split", self.split_at)): + if boundary is not None and boundary not in _PROCESSING_BOUNDARIES: + raise ValueError(f"Invalid SP {operation} boundary: {boundary!r}") + if self.gather_at not in (None, "loss"): + raise ValueError(f"Invalid SP gather boundary: {self.gather_at!r}") + if self.side == "output" and (self.need_pad or self.need_split): + raise ValueError("Output keys cannot request input padding or splitting") + if self.side == "input" and self.gather_at is not None and not self.need_split: + raise ValueError("Only sharded input keys can request gathering") + if self.stage == "forward" and self.gather_at is not None: + raise ValueError("Model-forward keys cannot request loss gathering") + if self.model_input is None: + object.__setattr__(self, "model_input", self.stage == "forward") + if self.unpack_packed is None: + object.__setattr__( + self, "unpack_packed", self.side == "output" or self.gather_at is not None or not self.need_split + ) + + @property + def need_pad(self): + return self.pad_at is not None + + @property + def need_split(self): + return self.split_at is not None + + def sequence_dim(self, ndim): + if not -ndim <= self.split_dim < ndim: + raise ValueError(f"Invalid sequence axis {self.split_dim} for SP key {self.name!r}") + dim = self.split_dim % ndim + if dim == 0: + raise ValueError(f"SP key {self.name!r} cannot split or gather the batch axis") + return dim + + +def declared_key_specs(provider, side): + """Read common field declarations, retaining legacy loss declaration names.""" + specs = getattr(provider, f"{side}_key_specs", ()) + if callable(specs): + specs = specs() + legacy = getattr(provider, f"sp_{side}_key_specs", ()) + for entries in (specs, legacy): + if isinstance(entries, DataKeySpec): + entries = (entries,) + yield from entries or () + + +def partition_model_inputs(inputs, input_key_specs): + """Separate forward arguments from loss-only tensors without mutating the batch.""" + loss_names = {spec.name for spec in input_key_specs if not spec.model_input} + loss_names.add("channel") # sample-level loss metadata, not a sequence tensor + model_inputs = {key: value for key, value in inputs.items() if key not in loss_names} + loss_inputs = {key: value for key, value in inputs.items() if key in loss_names} + return model_inputs, loss_inputs diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py index cf2c3b282..b9297d0a0 100644 --- a/src/twinkle/loss/base.py +++ b/src/twinkle/loss/base.py @@ -7,6 +7,16 @@ class Loss: require_logits = False require_entropy = False require_logps = True + # Full-sequence outputs required by this loss call; missing keys raise an + # error in Transformers SP. output_key_specs defines their layout, while + # this selector requests collection for the current objective only. + sp_gather_keys = () + # Provider-owned layouts/routing for custom tensors, also consumed by the + # processor when SP is disabled. Concrete losses import DataKeySpec from + # data_format; the base class does not depend on an SP strategy. + input_key_specs = () + output_key_specs = () + # Request per-token values from the Transformers wrapper (e.g. PPO critic). require_values = False def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput: diff --git a/src/twinkle/loss/infonce.py b/src/twinkle/loss/infonce.py index 3dffab93f..888c11f27 100644 --- a/src/twinkle/loss/infonce.py +++ b/src/twinkle/loss/infonce.py @@ -52,7 +52,7 @@ def _resolve_distance_metric(metric) -> Callable[[torch.Tensor, torch.Tensor], t def _extract_sentences(outputs) -> torch.Tensor: - """Return [B, D] sentence embeddings from postprocess_tensor_sp output. + """Return [B, D] sentence embeddings from model task postprocessing output. Prefers the canonical ``embeddings`` key (post-pooling); falls back to ``logits`` (legacy hook-side pooling) and applies CLS pooling for 3-D. diff --git a/src/twinkle/loss/value.py b/src/twinkle/loss/value.py index 37a05f163..8ef329b7c 100644 --- a/src/twinkle/loss/value.py +++ b/src/twinkle/loss/value.py @@ -5,6 +5,7 @@ from twinkle.data_format import LossOutput from twinkle.loss.base import Loss from twinkle.loss.grpo import GRPOLoss +from twinkle.data_format.sequence_fields import DataKeySpec if TYPE_CHECKING: import torch @@ -20,6 +21,16 @@ class PPOValueLoss(Loss): require_logps = False require_values = True + # PPO regresses per-token values; gather the value sequence before loss. + sp_gather_keys = ('values',) + output_key_specs = (DataKeySpec( + 'values', + split_dim=1, + side='output', + stage='output', + autograd=True, + unpack_packed=True, + ),) def __init__(self, cliprange_value: float = 0.2, vf_coef: float = 1.0, ignore_index: int = -100, **kwargs): self.cliprange_value = cliprange_value diff --git a/src/twinkle/model/transformers/moe/router_replay.py b/src/twinkle/model/transformers/moe/router_replay.py index 9c9377d70..a38be1a87 100644 --- a/src/twinkle/model/transformers/moe/router_replay.py +++ b/src/twinkle/model/transformers/moe/router_replay.py @@ -47,7 +47,8 @@ class RouterReplay: _registry: dict[str, _RouterReplayState] = {} def __init__(self, *, action, routed_experts=None, batch_size=1, manual_cleanup=False, - replay_mask=None, model=None, apply_patch=None) -> None: + replay_mask=None, model=None, apply_patch=None, sp_context=None) -> None: + self._sp_context = sp_context self._model = model self._unwrapped = None self._batch_size = batch_size @@ -61,7 +62,7 @@ def __init__(self, *, action, routed_experts=None, batch_size=1, manual_cleanup= if self.action == RouterReplayAction.REPLAY_FORWARD: assert routed_experts is not None, 'routed_experts must be not None' RouterReplay.arm_global(self.action) - RouterReplay.set_router_replay_data(routed_experts, self._get_unwrapped(), replay_mask) + RouterReplay.set_router_replay_data(routed_experts, self._get_unwrapped(), replay_mask, sp_context=self._sp_context) else: RouterReplay.arm_global(self.action) @@ -82,7 +83,7 @@ def finish(self) -> torch.Tensor | None: return None recorded = None if self.action == RouterReplayAction.RECORD: - recorded = RouterReplay.get_router_replay_data(self._get_unwrapped(), self._batch_size) + recorded = RouterReplay.get_router_replay_data(self._get_unwrapped(), self._batch_size, sp_context=self._sp_context) if not self._manual_cleanup: RouterReplay.clear_global_router_replay_action() RouterReplay.clear_global_indices() @@ -160,6 +161,7 @@ def set_router_replay_data( routed_experts: torch.Tensor, model: nn.Module, replay_mask: torch.Tensor = None, + sp_context=None, ) -> None: """Slice *routed_experts* into per-block ``target_indices`` and inject them into the registered MoE blocks of *model*. @@ -182,27 +184,22 @@ def set_router_replay_data( raise ValueError(f'Expected routed_experts with shape [bs, seq_len, layers, topk], ' f'got {tuple(routed_experts.shape)}') - # 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 - if sp_world_size > 1: - extra_split_values = [(routed_experts, 0, 1)] + # Use the owning model's context; replay setup runs before its forward hook. + from ..strategy.sequence_parallel import _MODULE_TO_CTX + sp = sp_context if sp_context is not None else _MODULE_TO_CTX.get(model) + if (getattr(sp, 'world_size', None) or 1) > 1: + from ..strategy.sequence_parallel.data_keys import DataKeySpec, SPKeyCollection, SPKeyValue + + keys = SPKeyCollection() + for name, value in (('routed_experts', routed_experts), ('replay_mask', replay_mask)): + if value is not None: + spec = DataKeySpec(name, split_dim=1, pad_value=0, stage='router_replay', + pad_at='router_replay', split_at='router_replay') + keys.add(SPKeyValue(spec, value)) + sp.pad_and_split_inputs(keys, boundary='router_replay', real_position_ids=sp.real_position_ids) + routed_experts = keys.get_value('routed_experts') if replay_mask is not None: - extra_split_values.append((replay_mask, 0, 1)) - _, _, _, _, _, _, extra_values = sp.pad_and_split_inputs( - None, - None, - None, - None, - None, - None, - real_position_ids=sp.real_position_ids, - extra_split_values=extra_split_values) - routed_experts = extra_values[0] - if replay_mask is not None: - replay_mask = extra_values[1] + replay_mask = keys.get_value('replay_mask') # [bs, seq_len, num_moe_layers, topk] -> [total_seq, num_moe_layers, topk] routed_experts = routed_experts.flatten(0, 1).to(Platform.get_local_device()) @@ -225,7 +222,7 @@ def set_router_replay_data( state.target_mask = local_mask @staticmethod - def get_router_replay_data(model: nn.Module, batch_size=1) -> torch.Tensor | None: + def get_router_replay_data(model: nn.Module, batch_size=1, sp_context=None) -> torch.Tensor | None: """Collect ``recorded_indices`` from all registered MoE blocks in *model*. . note:: @@ -255,9 +252,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 - if sp_world_size > 1: + from ..strategy.sequence_parallel import _MODULE_TO_CTX + sp = sp_context if sp_context is not None else _MODULE_TO_CTX.get(model) + if (getattr(sp, 'world_size', None) or 1) > 1: routed_experts = sp.gather(routed_experts, dim=1, position_ids=sp.real_position_ids) # [bs, seq_len, num_layers, topk] diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py index 0112d1aec..66b0f3d55 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py @@ -1,35 +1,88 @@ # 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 transformers import PreTrainedTokenizer from types import SimpleNamespace +from transformers import PreTrainedTokenizer from typing import Any, Dict, List, Optional, Tuple, Union +from weakref import WeakKeyDictionary, WeakMethod -from twinkle.patch import apply_patch +from twinkle.data_format.sequence_fields import _PROCESSING_BOUNDARIES +from twinkle.processor.fields import _FORWARD_KEYS from twinkle.utils import DeviceMesh, get_logger from twinkle.utils.transformers_utils import get_llm_model, is_flash_attention_implementation from twinkle.utils.utils import call_with_supported_kwargs, has_signature_parameter -from twinkle.patch.transformers_qwen3_vl_deepstack import Qwen3VLDeepstackSPPatch -from .linear_attention_sp import Qwen3_5GatedDeltaNetUlyssesPatch, _iter_qwen35_gated_delta_net_classes +from .adapters import ModelSPAdapter, adapters_for_model from .utils import (DistributedAttention, GatherLoss, _derive_sequence_parallel_sizes, _get_seq_groups_from_device_mesh, _get_ulysses_size, _SeqAllToAll, get_config_attr, get_cu_seqlens_from_position_ids, is_hccl_backend, - is_moe_config, post_all2all) + post_all2all) +from .data_keys import DataKeySpec, InputKeyRegistry, OutputKeyRegistry, SPKeyCollection, SPKeyValue + + +# --------------------------------------------------------------------------- +# Per-model SP context routing +# +# Each ``SequenceParallel`` instance is the SP context of exactly one model. +# The global transformers patch points (``ALL_ATTENTION_FUNCTIONS``, +# ``masking_utils``, ``modeling_flash_attention_utils``) are installed once per +# process and resolve the owning context at call time: +# * attention functions receive ``module`` -> ``_MODULE_TO_CTX`` lookup; +# * mask helpers / ``_flash_attention_forward`` receive no module argument -> +# thread-local ctx stack pushed/popped by each model's forward hooks. +# A missing context means the calling model is not SP-enabled: the wrapper +# falls through to the original transformers implementation untouched, so +# SP-enabled and non-SP models can be interleaved freely. +# --------------------------------------------------------------------------- +_MODULE_TO_CTX: WeakKeyDictionary = WeakKeyDictionary() +_CTX_LOCAL = threading.local() + + +def _ctx_stack() -> list: + stack = getattr(_CTX_LOCAL, 'stack', None) + if stack is None: + stack = _CTX_LOCAL.stack = [] + return stack + + +def _push_ctx(ctx) -> None: + _ctx_stack().append(ctx) + + +def _pop_ctx(ctx) -> None: + stack = _ctx_stack() + for i in range(len(stack) - 1, -1, -1): + if stack[i] is ctx: + stack.pop(i) + return logger = get_logger() -def is_qwen3_vl(model): - mt = getattr(getattr(model, 'config', None), 'model_type', '') - return 'qwen3_vl' in mt +def _current_ctx(): + stack = _ctx_stack() + return stack[-1] if stack else None + +def _register_module_tree(model: torch.nn.Module, ctx) -> None: + for module in model.modules(): + _MODULE_TO_CTX[module] = ctx -def is_qwen3_omni(model): - mt = getattr(getattr(model, 'config', None), 'model_type', '') - return 'qwen3_omni' in mt + +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 _call_create_causal_mask(fn, config, input_embeds, attention_mask, cache_position_or_past_key_values, *args, @@ -80,17 +133,23 @@ def __init__(self): self.num_heads = None self.causal_mask_func = None self.extra_kwargs = {} + self._dist_attns = {} + self.input_key_registry = InputKeyRegistry() + self.output_key_registry = OutputKeyRegistry() + self.model_adapters: Tuple[ModelSPAdapter, ...] = () + self.position_ids_mode = 'auto' @property def real_position_ids(self) -> torch.Tensor: """The real position ids, this is different from the position_ids in mrope""" return self.extra_kwargs.get('position_ids') - @staticmethod - def _extract_real_position_ids(position_ids: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + def _extract_real_position_ids(self, position_ids: Optional[torch.Tensor]) -> Optional[torch.Tensor]: if position_ids is None or not torch.is_tensor(position_ids): return position_ids - if position_ids.dim() == 3: + # Cached real positions are already 2D, so only reduce the original + # three-axis mRoPE representation. + if position_ids.dim() == 3 and self.position_ids_mode in ('mrope', 'auto'): return position_ids[0] return position_ids @@ -102,7 +161,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 +379,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 +391,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 +406,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 = _MODULE_TO_CTX.get(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 = _MODULE_TO_CTX.get(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,112 +456,76 @@ 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)') + @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 + def _prepare_forward_hook(self, base_model: torch.nn.Module): def pre_forward_split_hook(_self, args, kwargs): - if self.world_size == 1: + # Make this model's SP context visible to the global mask/attention + # wrappers for the whole duration of the forward. + _push_ctx(self) + try: + if self.world_size == 1: + return args, kwargs + # Pad to multiple of SP size and split inputs per SP rank before forward. + position_ids = kwargs['position_ids'] + real_position_ids = self._extract_real_position_ids(position_ids) + cache_position = kwargs.get('cache_position', None) + specs = tuple(spec for spec in self.input_key_registry + if spec.pad_at == 'forward' or spec.split_at == 'forward') + keys = SPKeyCollection.from_mapping(kwargs, specs) + self.pad_and_split_inputs(keys, boundary='forward', real_position_ids=real_position_ids, + cache_position=cache_position) + for item in keys.values(): + kwargs[item.name] = item.value 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) - - def _prepare_multimodal_deepstack(self, base_model: torch.nn.Module): - if not is_qwen3_vl(base_model): - return - apply_patch(base_model, Qwen3VLDeepstackSPPatch, sequence_parallel=self) - - @staticmethod - def _is_qwen35_model(model: torch.nn.Module) -> bool: - config = getattr(model, 'config', None) - model_type = str(getattr(config, 'model_type', '') or '') - if model_type in {'qwen3_5', 'qwen3_5_moe'}: - return True - - architectures = getattr(config, 'architectures', None) or [] - if any('Qwen3_5' in str(arch) for arch in architectures): - return True - - model_module = getattr(model.__class__, '__module__', '') or '' - return 'transformers.models.qwen3_5' in model_module or 'transformers.models.qwen3_5_moe' in model_module - - def _prepare_qwen35_linear_attention(self, model: torch.nn.Module): - has_qwen35_linear_attention = self._is_qwen35_model(model) - if not has_qwen35_linear_attention: - gated_delta_net_classes = tuple(_iter_qwen35_gated_delta_net_classes()) - if not gated_delta_net_classes: - return - has_qwen35_linear_attention = any(isinstance(module, gated_delta_net_classes) for module in model.modules()) - if not has_qwen35_linear_attention: - return - if int(self.rp_world_size or 1) > 1: - raise NotImplementedError( - 'SequenceParallel: Qwen3.5 linear attention sequence parallel does not support rp_world_size > 1 ' - '(derived ring attention).') - apply_patch(None, Qwen3_5GatedDeltaNetUlyssesPatch, sequence_parallel=self) - - def _prepare_moe_aux_loss(self, base_model: torch.nn.Module): - - def moe_aux_loss_hook(module, args, kwargs, output): - router_logits = getattr(output, 'router_logits', None) - if router_logits is None: - return output - - attention_mask = kwargs['attention_mask'] - if attention_mask is None: - batch_size = 1 - else: - batch_size = attention_mask.shape[0] - - assert router_logits[0].shape[0] % batch_size == 0 - seq_len = router_logits[0].shape[0] // batch_size - - _gathered_logits = [] - for i in range(batch_size): - _slice = slice(i * seq_len, (i + 1) * seq_len) - _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) - _gathered_logits.append(_bs_logits) - router_logits = torch.stack(_gathered_logits, dim=0) - if self.real_position_ids is not None: - router_logits = router_logits[:, :, :self.real_position_ids.shape[1], :] - output['router_logits'] = tuple( - [logit.reshape(-1, logit.shape[-1]) for logit in router_logits.split(1, dim=1)]) + except BaseException: + # PyTorch does not guarantee forward hooks run when a pre-hook + # itself raises. Restore the previous model context immediately. + _pop_ctx(self) + raise + + def post_forward_pop_hook(_self, args, output): + _pop_ctx(self) return output - base_model.register_forward_hook(moe_aux_loss_hook, with_kwargs=True) + base_model.register_forward_pre_hook(pre_forward_split_hook, with_kwargs=True, prepend=True) + # PyTorch also calls this hook if a pre-hook or the forward raises. + base_model.register_forward_hook(post_forward_pop_hook, always_call=True) + + def _configure_model_adapters(self, model: torch.nn.Module, text_model: torch.nn.Module) -> None: + adapters = adapters_for_model(model) + if text_model is not model: + adapters += tuple(adapter for adapter in adapters_for_model(text_model) if adapter not in adapters) + modes = {adapter.position_ids_mode for adapter in adapters if adapter.position_ids_mode != 'auto'} + if len(modes) > 1: + raise ValueError(f'Conflicting SP position_ids modes: {sorted(modes)}') + self.model_adapters = adapters + self.position_ids_mode = next(iter(modes), 'auto') + for adapter in adapters: + for spec in adapter.input_key_specs(): + self.input_key_registry.register(spec) + for spec in adapter.output_key_specs(): + self.output_key_registry.register(spec) + for adapter in adapters: + adapter.configure(self, text_model) def prepare( self, @@ -500,15 +575,10 @@ def prepare( 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) - self._prepare_multimodal_deepstack(llm_model) - - if is_moe_config(getattr(model, 'config', None)): - self._prepare_moe_aux_loss(llm_model) + self._configure_model_adapters(model, llm_model) self.model_dtype = next(model.parameters()).dtype self.tokenizer = tokenizer @@ -522,6 +592,14 @@ def prepare( f'(ulysses={self.sp_world_size} x ring={self.rp_world_size}), ' f'attn_implementation={getattr(model.config, "_attn_implementation", None)!r}') + # Route this model's text-side attention modules to this context. The + # registry is instance-based (not class-based), so a second model of the + # same architecture gets its own context instead of silently sharing + # this one's extra_kwargs, and modules of other models (e.g. vision + # towers) fall through to the origin attention. + text_model = llm_model.language_model if hasattr(llm_model, 'language_model') else llm_model + _register_module_tree(text_model, self) + def _mask_qkv(self, query, key, value, mask): mask = mask.unsqueeze(2).unsqueeze(3) query = query * mask @@ -568,21 +646,29 @@ def pad_and_split_mm_tokens(self, visual_mask, mm_embeds): empty_embeds = torch.empty( (input_ids.shape[0], input_ids.shape[1], mm_embeds.shape[-1])).to(mm_embeds.device).to(mm_embeds.dtype) empty_embeds[visual_mask] = mm_embeds - - embeds = SimpleNamespace(weight=mm_embeds) - - _, split_input_embeds, _, _, _, _, extra_values = self.pad_and_split_inputs( - None, - empty_embeds, - None, - None, - None, - None, - embeds, - self.real_position_ids, - extra_split_values=[(visual_mask, 0, -1)]) - visual_mask = extra_values[0] - return visual_mask, split_input_embeds[visual_mask] + embeds_spec = DataKeySpec( + 'visual_sequence_embeds', + split_dim=1, + pad_value=0, + stage='model_internal', + pad_at='model_internal', + split_at='model_internal', + ) + mask_spec = DataKeySpec( + 'visual_position_mask', + split_dim=-1, + pad_value=0, + stage='model_internal', + pad_at='model_internal', + split_at='model_internal', + ) + keys = SPKeyCollection() + keys.add(SPKeyValue(embeds_spec, empty_embeds)) + keys.add(SPKeyValue(mask_spec, visual_mask)) + self.pad_and_split_inputs(keys, boundary='model_internal', real_position_ids=self.real_position_ids) + local_mask = keys.get_value('visual_position_mask') + local_embeds = keys.get_value('visual_sequence_embeds') + return local_mask, local_embeds[local_mask] def gather(self, local_output, dim: int, position_ids=None): """Gather tensor for sequence parallel - reverse of split.""" @@ -699,105 +785,100 @@ def split(self, input, dim: int, position_ids=None): tensor_list = torch.split(input, dim_size // self.sp_world_size, dim=dim) return tensor_list[self.sp_rank].contiguous() + def _pad_key(self, value, spec, position_ids=None): + if value is None: + return None + pad_value = self.tokenizer.pad_token_id if spec.pad_value == 'tokenizer' else spec.pad_value + return self.pad(value, padding_value=pad_value, dim=spec.sequence_dim(value.dim()), position_ids=position_ids) + + def _split_key(self, value, spec, position_ids=None): + if value is None or not spec.need_split: + return value + return self.split(value, dim=spec.split_dim, position_ids=position_ids) + def pad_and_split_inputs(self, - input_ids, - input_embeds, - labels, - position_ids, - attention_mask, - loss_scale, - embed_tokens=None, + keys: SPKeyCollection, + *, + boundary: str, real_position_ids=None, - cache_position=None, - extra_split_values=None): - """Common implementation for padding and splitting inputs - - Pad to a length divisible by the sequence-parallel size, then split across SP ranks. - - Args: - input_ids: input_ids - input_embeds: input_embeds - labels: labels - position_ids: position_ids or, position_ids for mrope - attention_mask: attention_mask - loss_scale: loss_scale - embed_tokens: embed_tokens - real_position_ids: the real position_ids to represent the seq length information - extra_split_values: List of Tuples for extra split values, e.g.: (tensor, pad_value, split_dim) + cache_position=None): + """Apply the declared pad/split operations to a batch-local key collection. + + ``real_position_ids`` describes the unpadded sequence layout used to + preserve packed boundaries. ``cache_position`` is model-call context; + it is intentionally not treated as a sequence key or transformed. """ - tokenizer = self.tokenizer + if not isinstance(keys, SPKeyCollection): + raise TypeError('pad_and_split_inputs expects an SPKeyCollection') + if boundary not in _PROCESSING_BOUNDARIES: + raise ValueError(f'Unknown SP processing boundary: {boundary!r}') + + position_ids = keys.get_value('position_ids') real_position_ids = real_position_ids if real_position_ids is not None else position_ids - extra_values = [] - batch_size = input_ids.shape[ - 0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else None + input_ids = keys.get_value('input_ids') + inputs_embeds = keys.get_value('inputs_embeds') + model_inputs = input_ids if input_ids is not None else inputs_embeds + batch_size = model_inputs.shape[0] if model_inputs is not None else None if real_position_ids is not None and batch_size is not None and real_position_ids.shape[0] == batch_size: - # TODO clone everytime, but the position_ids is a small tensor self.extra_kwargs['position_ids'] = real_position_ids.clone() - if input_ids is not None: - 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) - 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 - if self.rp_world_size > 1 and batch_size > 1: + if self.rp_world_size > 1 and batch_size is not None and batch_size > 1: raise NotImplementedError( 'Derived ring attention only supports padding-free / packed batches with batch_size == 1.') - if position_ids is not None: - position_ids = self.pad(position_ids, padding_value=-1, position_ids=real_position_ids, dim=-1) - if labels is not None: - labels = self.pad(labels, padding_value=-100, position_ids=real_position_ids) - if loss_scale is not None: - loss_scale = self.pad(loss_scale, padding_value=0., position_ids=real_position_ids) - if real_position_ids is not None: - real_position_ids = self.pad(real_position_ids, padding_value=-1, position_ids=real_position_ids) - # Build a 2D attention_mask whenever we padded for SP alignment so FlashAttention2 can unpad correctly. - # For packed batches (batch_size==1 with multiple position_id resets), relying on position_ids alone is - # unsafe if we also appended SP-alignment padding (position_ids=-1), because HF's FA2 varlen path will - # include the padded tail in the last segment when attention_mask is None. - if (input_ids is not None or input_embeds is not None) and batch_size > 1: - # not padding_free, so not ring-attention - inputs = input_ids if input_ids is not None else input_embeds - attn_shape = inputs.shape[1] # The sequence length - if attention_mask is None: - # Mask out padded positions introduced by sequence-parallel padding. - # `real_position_ids` is padded with `-1` (see above), so use it to build a valid-token mask. - attention_mask = (real_position_ids != -1).to(dtype=torch.int64) - # no need position_ids here, because padding_free does not need attention_mask, - # so this is not ring-attention - attention_mask = self.pad(attention_mask, padding_value=0) - local_cache_position = torch.arange(0, attn_shape, device=inputs.device) - # FlashAttention expects a 2D padding mask (or None). Converting it to a 4D causal mask here breaks - # the later per-rank sequence split and changes the attention contract relative to the baseline path. - if (cache_position is None and hasattr(self, 'causal_mask_func') and self.causal_mask_func is not None + split_is_due = any(item.spec.split_at == boundary and item.value is not None for item in keys.values()) + padded_positions = (self.pad(real_position_ids, padding_value=-1, position_ids=real_position_ids) + if real_position_ids is not None and split_is_due else None) + + # Pad every declared key at this boundary. This includes both built-in + # keys and registered extra keys. + for item in keys.values(): + spec, value = item.spec, item.value + if value is None: + continue + if not torch.is_tensor(value): + raise TypeError(f'SP input {spec.name!r} must be a tensor') + if spec.pad_at == boundary: + # Dense SDPA also needs a global padding mask for batch_size=1; + # FlashAttention's packed path keeps its existing mask handling. + if (spec.name != 'attention_mask' or (batch_size is not None and batch_size > 1) + or not is_flash_attention_implementation(self.attn_implementation)): + value = self._pad_key(value, spec, real_position_ids) + item.value = value + + padded_ids = keys.get_value('input_ids') + padded_embeds = keys.get_value('inputs_embeds') + padded_model_inputs = padded_ids if padded_ids is not None else padded_embeds + if input_ids is not None and boundary == 'forward': + # Attention wrappers and multimodal patches need the padded global IDs. + self.extra_kwargs['input_ids'] = padded_ids.clone() + + # Keep the legacy generated 2D mask and causal-mask conversion on the + # forward boundary; these depend on model inputs and cache state. + if (boundary == 'forward' and padded_model_inputs is not None + and (batch_size > 1 or not is_flash_attention_implementation(self.attn_implementation))): + attn_item = keys.get('attention_mask') + if attn_item is None or attn_item.value is None: + attention_mask = (padded_positions != -1).to(dtype=torch.int64) + if attn_item is None: + spec = self.input_key_registry.get('attention_mask') + attn_item = SPKeyValue(spec, attention_mask) + keys.add(attn_item) + else: + attn_item.value = attention_mask + else: + attention_mask = attn_item.value + local_cache_position = torch.arange(0, padded_model_inputs.shape[1], device=padded_model_inputs.device) + causal_mask_func = self.causal_mask_func + if (cache_position is None and causal_mask_func is not None and not is_flash_attention_implementation(self.attn_implementation)): - attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype), - local_cache_position, None, None) - if extra_split_values is not None: - for (tensor, pad_value, split_dim) in extra_split_values: - extra_values.append( - self.pad(tensor, padding_value=pad_value, position_ids=real_position_ids, dim=split_dim)) - if input_ids is not None: - input_ids = self.split(input_ids, dim=1, position_ids=real_position_ids) - if input_embeds is not None: - input_embeds = self.split(input_embeds, dim=1, position_ids=real_position_ids) - if labels is not None: - labels = self.split(labels, dim=-1, position_ids=real_position_ids) - if loss_scale is not None: - loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1) - loss_scale = self.split(loss_scale, dim=-1, position_ids=real_position_ids) - - if position_ids is not None: - position_ids = self.split(position_ids, dim=-1, position_ids=real_position_ids) - # if attention_mask is not None and torch.is_tensor(attention_mask) and attention_mask.dim() == 2: - # attention_mask = self.split(attention_mask, dim=1, position_ids=real_position_ids) - if extra_split_values is not None: - for i in range(len(extra_values)): - extra_values[i] = self.split( - extra_values[i], dim=extra_split_values[i][2], position_ids=real_position_ids) - return input_ids, input_embeds, labels, position_ids, attention_mask, loss_scale, extra_values + attention_mask = causal_mask_func(attention_mask, padded_model_inputs.to(self.model_dtype), + local_cache_position, None, None) + attn_item.value = attention_mask + + # Split after padding so every field uses the same padded sequence frame. + for item in keys.values(): + if item.value is not None and item.spec.split_at == boundary: + item.value = self._split_key(item.value, item.spec, padded_positions) + return keys def _init_device_mesh(self, device_mesh: Optional[DeviceMesh] = None): """Initialize process groups for sequence parallel.""" @@ -832,11 +913,7 @@ def _is_packed_position_ids(position_ids: Optional[torch.Tensor]) -> bool: return False def prepare_inputs(self, inputs): - """Prepare inputs - - 1. set extra_kwargs['position_ids'] - 2. split labels, and completion_mask when present - """ + """Cache position context and process all keys declared for prepare.""" input_ids = inputs.get('input_ids') position_ids = inputs.get('position_ids') padding_free = bool(inputs.pop('padding_free', False)) @@ -851,36 +928,19 @@ def prepare_inputs(self, inputs): 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] + specs = tuple(spec for spec in self.input_key_registry + if spec.pad_at == 'prepare' or spec.split_at == 'prepare') + keys = SPKeyCollection.from_mapping(inputs, specs) + self.pad_and_split_inputs(keys, boundary='prepare', real_position_ids=real_position_ids) + for item in keys.values(): + inputs[item.name] = item.value 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 +966,11 @@ def __init__( self._tokenizer_id = tokenizer_id self._tokenizer = None self._initialized = False + # Each model owns its SP context. + self._ctx = SequenceParallel() + # Compatibility declarations registered directly through the public API. + # Keep these separate from rules consumed from model/loss providers. + self.field_declarations = SimpleNamespace(input_key_specs=(), output_key_specs=()) def _get_tokenizer(self) -> Optional[PreTrainedTokenizer]: if self._tokenizer is not None: @@ -932,7 +997,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,112 +1006,172 @@ def initialize(self) -> bool: self._initialized = True return True + # -- SP context delegates ------------------------------------------------- + # Last-token pooling reads topology and calls pad/split through this model's + # strategy; processors do not access the parallel context. + + @property + def world_size(self) -> int: + return self._ctx.world_size or 1 + + @property + def sp_world_size(self) -> int: + return self._ctx.sp_world_size or 1 + + @property + def rp_world_size(self) -> int: + return self._ctx.rp_world_size or 1 + + @property + def _sp_group(self): + return self._ctx._sp_group + + @property + def _rp_group(self): + return self._ctx._rp_group + + @property + def real_position_ids(self): + return self._ctx.real_position_ids + + def pad(self, tensor, *args, **kwargs): + if self._ctx.world_size is None: + return tensor + return self._ctx.pad(tensor, *args, **kwargs) + + def split(self, tensor, *args, **kwargs): + if self._ctx.world_size is None: + return tensor + return self._ctx.split(tensor, *args, **kwargs) + def preprocess_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]: 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 + return self._ctx.prepare_inputs(inputs) - @staticmethod - def _trim_gathered_sequence_padding(tensor: torch.Tensor, real_position_ids: torch.Tensor) -> torch.Tensor: + def pool_last_valid_token(self, features, inputs): + from .pooling import pool_last_valid_token + + return pool_last_valid_token(features, inputs, sp_strategy=self) + + def register_input_key(self, spec: DataKeySpec) -> None: + """Declare a token-aligned input on this model only. + + pad_at/split_at select execution boundaries; gather_at selects + loss-side collection. The model wrapper owns input routing and merges + these explicit declarations into the processor schema. + """ + self._ctx.input_key_registry.register(spec) + if spec not in self.field_declarations.input_key_specs: + self.field_declarations.input_key_specs += (spec,) + + @property + def input_key_specs(self) -> Tuple[DataKeySpec, ...]: + """Read-only snapshot of this context's execution rules.""" + return tuple(self._ctx.input_key_registry) + + def register_output_key(self, spec: DataKeySpec) -> None: + """Register a model-produced sequence output to collect before loss.""" + self._ctx.output_key_registry.register(spec) + if spec not in self.field_declarations.output_key_specs: + self.field_declarations.output_key_specs += (spec,) + + @property + def output_key_specs(self) -> Tuple[DataKeySpec, ...]: + return tuple(self._ctx.output_key_registry) + + def register_key_specs(self, input_key_specs, output_key_specs): + """Consume model-resolved field definitions for SP execution.""" + from twinkle.processor.fields import COMMON_INPUT_SPECS, COMMON_OUTPUT_SPECS + + input_defaults = {spec.name: spec for spec in COMMON_INPUT_SPECS} + output_defaults = {spec.name: spec for spec in COMMON_OUTPUT_SPECS} + for spec in input_key_specs: + if spec != input_defaults.get(spec.name) and spec.name not in _FORWARD_KEYS: + self._ctx.input_key_registry.register(spec) + for spec in output_key_specs: + if spec != output_defaults.get(spec.name): + self._ctx.output_key_registry.register(spec) + + def register_loss_key_specs(self, loss) -> None: + """Compatibility entry point; field ownership remains with the provider.""" + from twinkle.data_format.sequence_fields import declared_key_specs + + self.register_key_specs(tuple(declared_key_specs(loss, 'input')), + tuple(declared_key_specs(loss, 'output'))) + + def _trim_gathered_sequence_padding(self, tensor: torch.Tensor, real_position_ids: torch.Tensor, + dim: int = 1) -> 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: + dim = dim % tensor.dim() + + def take(start, end): + slices = [slice(None)] * tensor.dim() + slices[dim] = slice(start, end) + return tensor[tuple(slices)] + + if (self._ctx.rp_world_size or 1) > 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 - pieces.append(tensor[:, padded_offset:padded_offset + real_len]) + pieces.append(take(padded_offset, padded_offset + real_len)) padded_offset += padded_len - 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) + return torch.cat(pieces, dim=dim).contiguous() if pieces else take(0, 0).contiguous() + return take(0, real_position_ids.shape[-1]).contiguous() + + def _gather_key(self, value, spec, real_position_ids): + if not torch.is_tensor(value) or value.dim() < 2: + raise TypeError(f'SP key {spec.name!r} must be a batched sequence tensor') + dim = spec.sequence_dim(value.dim()) + if spec.autograd: + gathered, _ = GatherLoss.apply(self._ctx, value, None, dim, real_position_ids) + else: + gathered = self._ctx.gather(value, dim=dim, position_ids=real_position_ids) + return self._trim_gathered_sequence_padding(gathered, real_position_ids, dim=dim) def gather_loss_tensors( self, inputs: Dict[str, Any], outputs: Dict[str, Any], + sp_gather_keys=(), ) -> 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: - 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) + real_position_ids = self._ctx.real_position_ids + for spec in self.input_key_specs: + value = inputs.get(spec.name) + if spec.stage == 'auxiliary' and spec.gather_at == 'loss' and value is not None: + inputs[spec.name] = self._gather_key(value, spec, real_position_ids) + + fused_hidden = outputs.get('logps') is None and 'lm_head' in outputs and inputs.get('labels') is not None + if isinstance(sp_gather_keys, str): + sp_gather_keys = (sp_gather_keys,) + else: + sp_gather_keys = tuple(sp_gather_keys) + if any(not isinstance(key, str) or not key for key in sp_gather_keys): + raise TypeError('sp_gather_keys must contain non-empty output key names') + required_keys = set(sp_gather_keys) + output_specs = {spec.name: spec for spec in self.output_key_specs} + unknown_keys = required_keys - output_specs.keys() + if unknown_keys: + raise KeyError(f'Loss requested SP gather for unregistered output key(s) {sorted(unknown_keys)}. ' + 'Declare a sequence-aligned DataKeySpec in the model or loss output_key_specs.') + missing_keys = {key for key in required_keys if outputs.get(key) is None} + if missing_keys: + raise KeyError(f'Loss declared SP gather key(s) {sorted(missing_keys)}, but the model output is missing them.') + for spec in self.output_key_specs: + value = outputs.get(spec.name) + if value is not None and (spec.gather_at == 'loss' or spec.name in required_keys + or (spec.name == 'logits' and fused_hidden)): + outputs[spec.name] = self._gather_key(value, spec, real_position_ids) return inputs, outputs def wrap_model(self, model, optimizer=None): diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py b/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py new file mode 100644 index 000000000..0c6e35818 --- /dev/null +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py @@ -0,0 +1,121 @@ +"""Model-family SP adapters and their small registry. + +Adapters own model-family detection and the model-specific setup steps. The +sequence-parallel context remains responsible for generic padding, sharding and +collectives; adapters only install behavior that depends on model internals. +""" +from typing import Tuple + +import torch + +from twinkle.patch import apply_patch +from twinkle.patch.transformers_qwen3_vl_deepstack import Qwen3VLDeepstackSPPatch +from .data_keys import DataKeySpec +from .linear_attention_sp import Qwen3_5GatedDeltaNetUlyssesPatch, _iter_qwen35_gated_delta_net_classes +from .utils import GatherLoss, is_moe_config + + +class ModelSPAdapter: + """Base contract for model-specific SP behavior.""" + + name = 'generic' + position_ids_mode = 'auto' + + def matches(self, model) -> bool: + return False + + def input_key_specs(self) -> Tuple[DataKeySpec, ...]: + return () + + def output_key_specs(self) -> Tuple[DataKeySpec, ...]: + return () + + def configure(self, context, model) -> None: + """Install hooks on the model's text root. Generic adapter is a no-op.""" + + +class Qwen3VLSPAdapter(ModelSPAdapter): + name = 'qwen3_vl' + position_ids_mode = 'mrope' + + def matches(self, model) -> bool: + config = getattr(model, 'config', None) + model_type = str(getattr(config, 'model_type', '') or '') + return 'qwen3_vl' in model_type + + def configure(self, context, model) -> None: + apply_patch(model, Qwen3VLDeepstackSPPatch, sequence_parallel=context) + + +class Qwen35LinearAttentionSPAdapter(ModelSPAdapter): + name = 'qwen3_5_linear_attention' + + @staticmethod + def _matches_config_or_class(model) -> bool: + config = getattr(model, 'config', None) + model_type = str(getattr(config, 'model_type', '') or '') + if model_type in {'qwen3_5', 'qwen3_5_moe'}: + return True + architectures = getattr(config, 'architectures', None) or [] + if any('Qwen3_5' in str(architecture) for architecture in architectures): + return True + module_name = getattr(model.__class__, '__module__', '') or '' + return ('transformers.models.qwen3_5' in module_name + or 'transformers.models.qwen3_5_moe' in module_name) + + def matches(self, model) -> bool: + if self._matches_config_or_class(model): + return True + gated_delta_net_classes = tuple(_iter_qwen35_gated_delta_net_classes()) + return bool(gated_delta_net_classes) and any( + isinstance(module, gated_delta_net_classes) for module in model.modules()) + + def configure(self, context, model) -> None: + if int(context.rp_world_size or 1) > 1: + raise NotImplementedError( + 'SequenceParallel: Qwen3.5 linear attention sequence parallel does not support ' + 'rp_world_size > 1 (derived ring attention).') + apply_patch(None, Qwen3_5GatedDeltaNetUlyssesPatch, sequence_parallel=context) + + +class MoEAuxSPAdapter(ModelSPAdapter): + name = 'moe_aux' + + def matches(self, model) -> bool: + config = getattr(model, 'config', None) + model_type = str(getattr(config, 'model_type', '') or '') + return model_type.endswith('_moe') or is_moe_config(config) + + def configure(self, context, base_model) -> None: + def gather_router_logits(module, args, kwargs, output): + router_logits = getattr(output, 'router_logits', None) + if router_logits is None: + return output + attention_mask = kwargs.get('attention_mask') + batch_size = 1 if attention_mask is None else attention_mask.shape[0] + assert router_logits[0].shape[0] % batch_size == 0 + local_seq_len = router_logits[0].shape[0] // batch_size + gathered_per_sample = [] + for batch_idx in range(batch_size): + seq_slice = slice(batch_idx * local_seq_len, (batch_idx + 1) * local_seq_len) + sample_logits = [layer_logits[seq_slice] for layer_logits in router_logits] + compute_device = sample_logits[0].device + sample_logits = torch.stack([logits.to(compute_device) for logits in sample_logits], dim=0) + sample_logits, _ = GatherLoss.apply(context, sample_logits, None, 1, context.real_position_ids) + gathered_per_sample.append(sample_logits) + router_logits = torch.stack(gathered_per_sample, dim=0) + if context.real_position_ids is not None: + router_logits = router_logits[:, :, :context.real_position_ids.shape[1], :] + output['router_logits'] = tuple( + logits.reshape(-1, logits.shape[-1]) for logits in router_logits.split(1, dim=1)) + return output + + base_model.register_forward_hook(gather_router_logits, with_kwargs=True) + + +_ADAPTERS = [Qwen3VLSPAdapter(), Qwen35LinearAttentionSPAdapter(), MoEAuxSPAdapter()] + + +def adapters_for_model(model): + """Return every matching adapter; families can compose capabilities.""" + return tuple(adapter for adapter in _ADAPTERS if adapter.matches(model)) diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py b/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py new file mode 100644 index 000000000..25784e24c --- /dev/null +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py @@ -0,0 +1,119 @@ +"""SP-local execution defaults, registries and batch-local key collections.""" +from dataclasses import dataclass, replace +from typing import Any, Dict, Iterable, Iterator, Mapping, Optional + +from twinkle.data_format.sequence_fields import DataKeySpec +from twinkle.processor.fields import ( + COMMON_INPUT_SPECS, COMMON_OUTPUT_SPECS, PROCESSOR_PASSTHROUGH_KEYS, validate_input_key_spec) + +__all__ = ['DataKeySpec', 'InputKeyRegistry', 'OutputKeyRegistry', 'BUILTIN_INPUT_KEYS', + 'BUILTIN_OUTPUT_KEYS', 'PROCESSOR_PASSTHROUGH_KEYS', 'SPKeyValue', 'SPKeyCollection'] + + +# Layouts have one source in processor.fields. SP adds only execution rules; +# loss-only completion_mask/loss_scale retain their global layouts unchanged. +_INPUT_OPERATIONS = { + "input_ids": dict(pad_at="forward", split_at="forward"), + "inputs_embeds": dict(pad_at="forward", split_at="forward"), + "position_ids": dict(pad_at="forward", split_at="forward"), + "attention_mask": dict(pad_at="forward"), + "labels": dict(pad_at="prepare", split_at="prepare", gather_at="loss"), +} +BUILTIN_INPUT_KEYS = tuple(replace(spec, **_INPUT_OPERATIONS.get(spec.name, {})) for spec in COMMON_INPUT_SPECS) + +# Real logits are opt-in; fused CE explicitly gathers this hidden-state slot. +BUILTIN_OUTPUT_KEYS = tuple( + replace(spec, autograd=True, gather_at="loss" if spec.name in ("logps", "entropies") else None) + for spec in COMMON_OUTPUT_SPECS +) + + +class InputKeyRegistry: + """Collect provider declarations; reject conflicting definitions. + + Auxiliary keys are processed at their declared boundary (usually + preprocessing). Custom model-forward keys remain global until the forward + hook, so every key is sharded at most once. + """ + + def __init__(self): + self._specs = {spec.name: spec for spec in BUILTIN_INPUT_KEYS} + + def register(self, spec: DataKeySpec): + validate_input_key_spec(spec) + existing = self._specs.get(spec.name) + if existing is not None and existing != spec: + raise ValueError(f"Conflicting SP input rule for {spec.name!r}") + self._specs[spec.name] = spec + + def __iter__(self): + return iter(self._specs.values()) + + def get(self, name): + return self._specs.get(name) + + +class OutputKeyRegistry: + def __init__(self): + self._specs = {spec.name: spec for spec in BUILTIN_OUTPUT_KEYS} + + def register(self, spec): + if not isinstance(spec, DataKeySpec): + raise TypeError("Expected a DataKeySpec") + if spec.side != "output": + raise ValueError("Expected an output key rule") + existing = self._specs.get(spec.name) + if existing is not None and existing != spec: + raise ValueError(f"Conflicting SP output rule for {spec.name!r}") + self._specs[spec.name] = spec + + def __iter__(self): + return iter(self._specs.values()) + + def get(self, name): + return self._specs.get(name) + + +@dataclass +class SPKeyValue: + """One tensor value paired with its immutable SP processing rule.""" + + spec: DataKeySpec + value: Any + + @property + def name(self): + return self.spec.name + + +class SPKeyCollection: + """Batch-local common and registered SP tensors passed through one executor.""" + + def __init__(self): + self._items: Dict[str, SPKeyValue] = {} + + def add(self, item: SPKeyValue): + if not isinstance(item, SPKeyValue): + raise TypeError('SPKeyCollection accepts SPKeyValue items') + if item.name in self._items: + raise ValueError(f'Duplicate SP key in collection: {item.name!r}') + self._items[item.name] = item + + def values(self) -> Iterator[SPKeyValue]: + return iter(self._items.values()) + + def get(self, name: str) -> Optional[SPKeyValue]: + return self._items.get(name) + + def get_value(self, name: str, default=None): + """Return a key's payload without exposing its SP rule wrapper.""" + item = self.get(name) + return default if item is None else item.value + + @classmethod + def from_mapping(cls, values: Mapping[str, Any], specs: Iterable[DataKeySpec]): + collection = cls() + for spec in specs: + if spec.name in values: + collection.add(SPKeyValue(spec, values[spec.name])) + return collection 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..88c1d776d 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 @@ -340,7 +340,11 @@ def sp_linear_forward( _origin_forward=origin_forward, **extra_kwargs, ): - sequence_parallel_context = extra_kwargs.pop('sequence_parallel_context', sequence_parallel) + # The class patch is shared by every model, including models + # without SP. Resolve the owner instead of capturing the first + # model's context in this process-wide closure. + from . import _MODULE_TO_CTX + sequence_parallel_context = extra_kwargs.pop('sequence_parallel_context', _MODULE_TO_CTX.get(mod)) if not _sp_is_enabled(sequence_parallel_context): return _origin_forward( mod, diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/pooling.py b/src/twinkle/model/transformers/strategy/sequence_parallel/pooling.py new file mode 100644 index 000000000..7889f6034 --- /dev/null +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/pooling.py @@ -0,0 +1,102 @@ +"""Last-token pooling for full sequences and sequence-parallel shards.""" +from typing import Any + +import torch +import torch.distributed as dist +from torch.distributed.nn.functional import all_reduce + + +def _packed_last_indices(position_ids: torch.Tensor, total_len: int) -> torch.Tensor: + """For padding-free batches: per-segment last-token indices into a [1, total] sequence.""" + flat = position_ids.squeeze(0) if position_ids.dim() == 2 else position_ids + starts = (flat == 0).nonzero(as_tuple=False).squeeze(-1) + end_anchor = torch.tensor([total_len], device=flat.device, dtype=starts.dtype) + boundaries = torch.cat([starts, end_anchor]) + return (boundaries[1:] - 1).long() + + +def pool_last_valid_token(features: torch.Tensor, inputs: dict[str, Any], + sp_strategy=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)``. + """ + sp_enabled = ( + 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 + elif inputs.get('position_ids') is not None: + ref_pos = inputs['position_ids'] + else: + # Standard padded batch: HF omits position_ids (it derives them internally). + # Rebuild a position_ids-like reference (valid token -> its index, padding -> -1) + # so the last-valid-token pooling below, which keys off ``ref_pos >= 0`` for + # validity and a single ``== 0`` per row for packing detection, still works. + T = features.shape[1] + arange = torch.arange(T, device=features.device) + am = inputs.get('attention_mask') + if am is None: + ref_pos = arange.expand(features.shape[0], T) + else: + if am.dim() >= 3: + am = am.reshape(am.shape[0], -1)[:, :T] + ref_pos = torch.where(am.bool(), arange, arange.new_tensor(-1)) + if ref_pos.dim() == 3: + ref_pos = ref_pos[0] + # 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')) + + is_packed = ( + features.shape[0] == 1 and (cu_seq_lens_q is not None or int((ref_pos.reshape(-1) == 0).sum()) > 1)) + + device, dtype = features.device, features.dtype + T_real = ref_pos.shape[-1] + + if is_packed: + if torch.is_tensor(cu_seq_lens_q) and cu_seq_lens_q.numel() >= 2: + end_idx = (cu_seq_lens_q[1:].long() - 1).to(device) + else: + end_idx = _packed_last_indices(ref_pos, T_real).to(device) + n_seqs = end_idx.shape[0] + batch_idx = torch.zeros_like(end_idx) + target_idx = torch.arange(n_seqs, device=device) + B_mask, n_targets = 1, n_seqs + else: + n_seqs = ref_pos.shape[0] + # clamp_min(0) prevents negative indexing when a row is fully padded (sum == 0). + positions = torch.arange(T_real, device=ref_pos.device).expand_as(ref_pos) + end_idx = positions.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 + + 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) + + 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.') + + 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: + pooled = all_reduce(pooled, op=dist.ReduceOp.SUM, group=grp) + + return pooled, is_packed 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/task_outputs.py b/src/twinkle/model/transformers/task_outputs.py new file mode 100644 index 000000000..b65be4039 --- /dev/null +++ b/src/twinkle/model/transformers/task_outputs.py @@ -0,0 +1,29 @@ +"""Model-side task reductions, followed by sequence-parallel loss collection.""" +from copy import copy + +import torch.nn.functional as F + +from .strategy.sequence_parallel.pooling import pool_last_valid_token + + +def postprocess_task_outputs(inputs, outputs, *, task='causal_lm', sp_strategy=None, sp_gather_keys=()): + """Reduce task features before loss; processors only restore batch layouts.""" + if task in ('embedding', 'generative_reranker', 'seq_cls', 'reranker'): + features = outputs.get('features') if task == 'embedding' else outputs.get('logits') + if features is None: + raise ValueError(f'Missing per-token features for task={task!r}') + # Classification heads may already return one score per sequence. + if task in ('seq_cls', 'reranker') and features.dim() == 2: + return inputs, outputs + pool = sp_strategy.pool_last_valid_token if sp_strategy is not None else pool_last_valid_token + pooled, _ = pool(features, inputs) + outputs = copy(outputs) + if task == 'embedding': + outputs.pop('features', None) + outputs['embeddings'] = F.normalize(pooled, p=2, dim=-1).contiguous() + else: + outputs['logits'] = pooled.contiguous() + return inputs, outputs + if sp_strategy is not None: + return sp_strategy.gather_loss_tensors(inputs, outputs, sp_gather_keys=sp_gather_keys) + return inputs, outputs diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 85fa19b37..185eb1afe 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -20,6 +20,7 @@ from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel from transformers.models.auto.auto_factory import _BaseAutoModelClass from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union, overload +from weakref import ref import twinkle import twinkle.module.optimizer @@ -51,15 +52,45 @@ from twinkle.template.utils import construct_template from twinkle.utils import (construct_class, get_logger, prepare_replayed_selective_log_softmax, replayed_selective_log_softmax, selective_log_softmax, torch_util) + from twinkle.utils.framework import Torch from twinkle.utils.grad_clip import normalize_and_clip_grad_norm from twinkle.utils.transformers_utils import filter_from_config_kwargs +from twinkle.data_format.sequence_fields import partition_model_inputs +from twinkle.processor.fields import resolve_key_specs + +from .task_outputs import postprocess_task_outputs + logger = get_logger() CHECKPOINT_ADAPTER_NAME = 'default' +class _FieldSpecsCache: + """Process-local declarations cache; transport always starts with an empty cache.""" + + def __init__(self, owners=(), declarations=(), specs=None): + self.owners = [] + for owner in owners: + try: + self.owners.append(ref(owner)) + except TypeError: + self.owners.append(lambda owner=owner: owner) + self.declarations = declarations + self.specs = specs + + def matches(self, owners, declarations): + return (len(self.owners) == len(owners) + and all(previous() is current for previous, current in zip(self.owners, owners)) + and self.declarations == declarations) + + def __reduce__(self): + # Weak references and fallback closures belong to this process. Neither + # pickle nor cloudpickle should transport them or stale identity checks. + return type(self), () + + def _get_vocab_size(config: PretrainedConfig) -> int: """Resolve the LM vocabulary size without running a sharded model forward.""" vocab_size = getattr(config, 'vocab_size', None) @@ -111,17 +142,16 @@ 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', ...)``. + ``postprocess_task_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_task_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_task_outputs`` can do packed-aware last-token reduction (same as embedding). """ if task in (None, 'causal_lm'): return contextlib.nullcontext() if task in ('seq_cls', 'reranker'): - # Skip HF's internal pooling; surface per-token [B, T, num_labels] so the processor's - # _pool_last_valid_token handles packed/non-packed uniformly. + # Surface per-token scores for model-side packed/SP-aware pooling. from twinkle.patch.transformers_value import TransformersValuePatch return apply_context(model, TransformersValuePatch(squeeze_single=False)) if task == 'embedding': @@ -667,16 +697,18 @@ def _decide_strategy(self, def _ensure_sp_strategy(self) -> None: if not getattr(self, '_enable_sp', False): return - if self.sp_strategy is not None: - return - from .strategy.sequence_parallel import SequenceParallelStrategy - - self.sp_strategy = SequenceParallelStrategy( - self.device_mesh, - {}, - model=self.model, - tokenizer_id=self.tokenizer_id, - ) + if self.sp_strategy is None: + from .strategy.sequence_parallel import SequenceParallelStrategy + + self.sp_strategy = SequenceParallelStrategy( + self.device_mesh, + {}, + model=self.model, + tokenizer_id=self.tokenizer_id, + ) + for optimizer_group in self.optimizer_group.values(): + for loss_instance in (optimizer_group.loss_instance, *optimizer_group.loss_registry.values()): + self.sp_strategy.register_loss_key_specs(loss_instance) def _ensure_gdn_padding_free_patch(self) -> None: """Install the non-SP GatedDeltaNet patch before wrapping the model.""" @@ -799,7 +831,8 @@ def _router_replay(self, action, *, routed_experts=None, batch_size=1, manual_cl return RouterReplay(action=None) self._maybe_apply_router_replay() return RouterReplay( - action, + action=action, + sp_context=getattr(self.sp_strategy, '_ctx', None), routed_experts=routed_experts, batch_size=batch_size, manual_cleanup=manual_cleanup, @@ -826,15 +859,48 @@ def _construct_default_optimizer_group(self): acc_strategy=self._acc_strategy, ) - def _restore_loss_inputs(self, inputs, labels, loss_scale=None, channel=None, completion_mask=None): - """Put the loss-related tensors popped before the model call back onto ``inputs``.""" - inputs['labels'] = labels - if loss_scale is not None: - inputs['loss_scale'] = loss_scale - if channel is not None: - inputs['channel'] = channel - if completion_mask is not None: - inputs['completion_mask'] = completion_mask + def _prepare_sp_inputs(self, inputs, processor, task='causal_lm'): + if self.sp_strategy is None: + return inputs + payload = dict(inputs) + payload['padding_free'] = processor.is_padding_free(inputs) + # Sentence-level labels describe whole sequences rather than tokens. + sentence_task = task in ('seq_cls', 'reranker', 'embedding', 'generative_reranker') + sequence_labels = payload.pop('labels', None) if sentence_task else None + payload = self.sp_strategy.preprocess_inputs(payload) + if sequence_labels is not None: + payload['labels'] = sequence_labels + return payload + + def _get_field_specs(self, loss_instance): + """Cache model/loss declarations independently of SP and batch contents. + + Configure provider declarations before forwarding. Model/strategy/loss + replacement and direct registrations invalidate the single-entry cache; + set_loss also invalidates it when reconfiguring the same loss instance. + """ + owners = (self.model, self.strategy, self.sp_strategy, loss_instance) + direct_fields = self.sp_strategy.field_declarations if self.sp_strategy is not None else None + declarations = ((direct_fields.input_key_specs, direct_fields.output_key_specs) + if direct_fields is not None else ((), ())) + cached = getattr(self, '_field_specs_cache', None) + if cached is not None and cached.matches(owners, declarations): + return cached.specs + + from twinkle.utils.transformers_utils import get_llm_model + from .strategy.sequence_parallel.adapters import adapters_for_model + + model = self.strategy.unwrap_model(self.model) + text_model = get_llm_model(model) + models = (model,) if text_model is model else (model, text_model) + adapters = tuple(adapter for candidate in models for adapter in adapters_for_model(candidate)) + # Direct registration is a legacy declaration interface. Merge only + # those explicit declarations, not SP's derived execution registry. + specs = resolve_key_specs(*models, *adapters, loss_instance, direct_fields) + if self.sp_strategy is not None: + self.sp_strategy.register_key_specs(*specs) + self._field_specs_cache = _FieldSpecsCache(owners, declarations, specs) + return specs def _postprocess_forward_outputs(self, outputs, *, labels, task, temperature, sampling_masks, return_logits, loss_instance, replay=None): @@ -876,10 +942,16 @@ def _postprocess_forward_outputs(self, outputs, *, labels, task, temperature, sa outputs['logits'] = None def _finalize_forward_outputs(self, inputs, outputs, *, processor, task, optimizer_config, is_training, - forward_kwargs, loss_instance, return_logits, recorded_routing): + forward_kwargs, loss_instance, return_logits, recorded_routing, + input_key_specs, output_key_specs): """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 = postprocess_task_outputs( + inputs, outputs, sp_strategy=self.sp_strategy, task=task, + sp_gather_keys=getattr(loss_instance, 'sp_gather_keys', ())) + inputs, outputs = processor.unpack_packed_sequences( + inputs, outputs, task=task, + input_key_specs=input_key_specs, + output_key_specs=output_key_specs) status = optimizer_config.train_status if is_training else optimizer_config.eval_status status.inputs = inputs status.outputs = outputs @@ -934,18 +1006,17 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' + input_key_specs, output_key_specs = self._get_field_specs(loss_instance) inputs: Dict[str, Any] = processor( inputs, - sp_strategy=self.sp_strategy, + input_key_specs=input_key_specs, hf_config=self.hf_config, enable_sp=getattr(self, '_enable_sp', False), gdn_padding_free_enabled=self._gdn_padding_free_enabled, ) - labels: torch.Tensor = inputs.pop('labels', None) - loss_scale = inputs.pop('loss_scale', None) - channel = inputs.pop('channel', None) - # Not a model argument; the loss reads it back off `inputs` further down. - completion_mask = inputs.pop('completion_mask', None) + inputs = self._prepare_sp_inputs(inputs, processor, task=task) + inputs, loss_inputs = partition_model_inputs(inputs, input_key_specs) + labels = loss_inputs.get('labels') replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay( @@ -953,8 +1024,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec sampling_masks=sampling_masks, temperature=temperature, vocab_size=_get_vocab_size(self.hf_config), - allow_packed_masks=(processor.padding_free - or processor._is_packed_position_ids(inputs.get('position_ids'))), + allow_packed_masks=processor.is_padding_free(inputs), ) optimizer_config.accumulate_metrics(True) @@ -983,7 +1053,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec recorded_routing = rr.finish() - self._restore_loss_inputs(inputs, labels, loss_scale, channel, completion_mask) + inputs.update(loss_inputs) self._postprocess_forward_outputs( outputs, labels=labels, @@ -1005,6 +1075,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec loss_instance=loss_instance, return_logits=return_logits, recorded_routing=recorded_routing, + input_key_specs=input_key_specs, output_key_specs=output_key_specs, ) # lazy_collect follows the global default (True): a caller that reads the output itself calls @@ -1101,27 +1172,25 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs): cp_world_size = self.device_mesh.cp_world_size if self.device_mesh is not None else 1 if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') + input_key_specs, output_key_specs = self._get_field_specs(loss_instance) inputs: Dict[str, Any] = processor( inputs, - sp_strategy=self.sp_strategy, + input_key_specs=input_key_specs, hf_config=self.hf_config, enable_sp=getattr(self, '_enable_sp', False), gdn_padding_free_enabled=self._gdn_padding_free_enabled, ) - labels = inputs.pop('labels', None) - loss_scale = inputs.pop('loss_scale', None) - channel = inputs.pop('channel', None) - # Not a model argument; the loss reads it back off `inputs` further down. - completion_mask = inputs.pop('completion_mask', None) + inputs = self._prepare_sp_inputs(inputs, processor, task=task) + inputs, loss_inputs = partition_model_inputs(inputs, input_key_specs) + labels = loss_inputs.get('labels') replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: - packed_position_ids = processor._is_packed_position_ids(inputs.get('position_ids')) replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay( labels=labels, sampling_masks=sampling_masks, temperature=temperature, vocab_size=_get_vocab_size(self.hf_config), - allow_packed_masks=processor.padding_free or packed_position_ids, + allow_packed_masks=processor.is_padding_free(inputs), ) optimizer_config.accumulate_metrics(False) unwrapped_model = self.strategy.unwrap_model(self.model) @@ -1149,7 +1218,7 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs): recorded_routing = rr.finish() - self._restore_loss_inputs(inputs, labels, loss_scale, channel, completion_mask) + inputs.update(loss_inputs) self._postprocess_forward_outputs( outputs, labels=labels, @@ -1171,6 +1240,7 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs): loss_instance=loss_instance, return_logits=return_logits, recorded_routing=recorded_routing, + input_key_specs=input_key_specs, output_key_specs=output_key_specs, ) @remote_function(dispatch='slice_dp', collect=collect_dp_and_flatten, lazy_collect=False) @@ -1840,7 +1910,11 @@ def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature adapter_name = kwargs.pop('adapter_name', self._get_default_group()) loss_key = kwargs.pop('loss_key', None) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.set_loss(construct_class(loss_cls, Loss, twinkle.loss, **kwargs), loss_key=loss_key) + loss_instance = construct_class(loss_cls, Loss, twinkle.loss, **kwargs) + optimizer_config.set_loss(loss_instance, loss_key=loss_key) + self._field_specs_cache = None + if self.sp_strategy is not None: + self.sp_strategy.register_loss_key_specs(loss_instance) @remote_function() def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], **kwargs): diff --git a/src/twinkle/patch/transformers_emb.py b/src/twinkle/patch/transformers_emb.py index 2a14b852f..7d08519d0 100644 --- a/src/twinkle/patch/transformers_emb.py +++ b/src/twinkle/patch/transformers_emb.py @@ -10,7 +10,7 @@ ``[B, T, H]`` or ``[B, T_local, H]`` under SP). Last-token pooling (incl. padding-free, SP gather) is **deferred** to -``InputProcessor.postprocess_tensor_sp(task='embedding', ...)`` so this patch +``postprocess_task_outputs(task='embedding', ...)`` so this patch stays SP/CP/packed-agnostic and the dispatch sits in one place. Both mutations are reverted by ``unpatch``. @@ -44,7 +44,7 @@ def get_lm_head_model(module, lm_heads=None): def _output_features_hook(module, args, kwargs, output): # Replaces the model's structured output entirely; downstream HF utilities # (generate, loss, past_key_values, ...) won't work until ``unpatch``. - # Pooling + L2-normalize is deferred to ``InputProcessor._postprocess_embedding`` + # Pooling + L2-normalize is deferred to ``postprocess_task_outputs`` # so the hook stays pooling-agnostic (last-token / mean / CLS all work uniformly). hidden_states = output.logits return {'features': hidden_states.contiguous()} diff --git a/src/twinkle/patch/transformers_reranker.py b/src/twinkle/patch/transformers_reranker.py index 9b20e6987..4bc09ba6f 100644 --- a/src/twinkle/patch/transformers_reranker.py +++ b/src/twinkle/patch/transformers_reranker.py @@ -14,7 +14,7 @@ are introduced (unlike the ``num_labels=1`` sequence-classification head used by the plain ``reranker`` task, which trains a fresh ``score`` linear). 2. Last-valid-token reduction ``[B, T, 1] -> [B, 1]`` is DEFERRED to - ``InputProcessor.postprocess_tensor_sp(task='generative_reranker', ...)`` -- symmetric with how + ``postprocess_task_outputs(task='generative_reranker', ...)`` -- symmetric with how embedding defers pooling -- so this patch stays SP/CP/packed-agnostic and the dispatch lives in one place. diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index cbea6cd3b..cbee0aa38 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -5,12 +5,12 @@ import numpy as np import torch -import torch.distributed as dist -import torch.nn.functional as F +from typing import Dict, List from twinkle import DeviceMesh, Platform, remote_class, remote_function, torch_util from twinkle.data_format import InputFeature from twinkle.perf import timed +from .fields import PADDING_VALUES, VLM_CONCAT_FIELDS, TRANSFORMERS_INPUT_KEYS from twinkle.utils.transformers_utils import is_flash_attention_implementation @@ -30,37 +30,7 @@ class PackedSeqParams: @remote_class() class InputProcessor: - padding_map = { - "input_ids": 0, - "mm_token_type_ids": 0, - "inputs_embeds": 0.0, - "attention_mask": 0, - "labels": -100, - "completion_mask": 0, - "loss_scale": 0.0, - "position_ids": -1, - "length": -1, - "pixel_values": 0.0, - "image_grid_thw": 0, - "pixel_values_videos": 0.0, - "video_grid_thw": 0, - "input_features": 0.0, - "feature_attention_mask": 0, - "routed_experts": 0, - "replay_mask": 0, - } - - # VLM fields to concatenate (not pad) in batch - VLM_CONCAT_FIELDS = { - "pixel_values", - "image_grid_thw", - "pixel_values_videos", - "video_grid_thw", - "input_features", - "input_features_mask", - "feature_attention_mask", - "grid_thws", - } + padding_map = PADDING_VALUES def __init__( self, @@ -84,7 +54,6 @@ def __init__( self.prepare_gdn_padding_free_args, self.drop_causal_4d_mask, self.split_cp, - self.apply_transformers_sp, self.prepare_outputs, ] @@ -105,14 +74,22 @@ def prepare_outputs(self, inputs: list[InputFeature], **kwargs) -> list[InputFea _input["position_ids"] = _input["position_ids"][1:] return inputs - def prepare_inputs(self, inputs: list[InputFeature] | InputFeature, **kwargs) -> list[InputFeature]: + def prepare_inputs(self, inputs: Union[List[InputFeature], InputFeature], input_key_specs=None, + **kwargs) -> List[InputFeature]: + from .fields import COMMON_INPUT_SPECS + + specs = COMMON_INPUT_SPECS if input_key_specs is None else input_key_specs + registered_keys = {spec.name for spec in specs} + def to_tensor(_input): for key in list(_input.keys()): value = _input[key] # Ray/pyarrow can return numpy or list scalars; normalize to tensors. # After distributed/datasets.map, labels/completion_mask may become numpy arrays or lists, # so tensor ops like labels != ignore_index or .to(device) would fail without this. - if isinstance(value, np.ndarray): + if key in registered_keys and isinstance(value, (list, np.ndarray)): + value = torch.as_tensor(value) + elif isinstance(value, np.ndarray): value = torch.from_numpy(value) elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], (int, float, np.number)): value = torch.tensor(value) @@ -135,12 +112,14 @@ def to_tensor(_input): value = torch.tensor(value) elif (isinstance(value, list)) and key in ("completion_mask", "mm_token_type_ids", "routed_experts"): value = torch.tensor(value) - elif key in self.VLM_CONCAT_FIELDS: + elif key in VLM_CONCAT_FIELDS: if not isinstance(value[0], torch.Tensor): value = [torch.tensor(v) for v in value] value = torch.cat(value, dim=0) if isinstance(value, torch.Tensor): value = value.to(Platform.get_local_device()) + if key == 'labels' and value.dim() == 0: + value = value.reshape(1) if value.dim() == 1: value = value.unsqueeze(0) _input[key] = value @@ -148,218 +127,8 @@ 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 - - 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. - """ - 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 - - @staticmethod - def _packed_last_indices(position_ids: torch.Tensor, total_len: int) -> torch.Tensor: - """For padding-free batches: per-segment last-token indices into a [1, total] sequence.""" - flat = position_ids.squeeze(0) if position_ids.dim() == 2 else position_ids - starts = (flat == 0).nonzero(as_tuple=False).squeeze(-1) - end_anchor = torch.tensor([total_len], device=flat.device, dtype=starts.dtype) - boundaries = torch.cat([starts, end_anchor]) - return (boundaries[1:] - 1).long() - - def _postprocess_embedding( - self, inputs: dict[str, Any], outputs: dict[str, Any], sp_strategy=None - ) -> tuple[dict[str, Any], dict[str, Any]]: - """Pool per-token features to per-sequence embeddings (last-valid-token). - - 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, H]`` - tensor across SP × RP. No feature gather; uniform across - DP / Ulysses / zigzag-ring / padding-free. - """ - features = outputs.get("features") - assert features is not None, ( - "outputs['features'] missing; ensure TransformersEmbeddingPatch is applied for task='embedding'." - ) - - pooled, is_packed = self._pool_last_valid_token(features, inputs, sp_strategy=sp_strategy) + def pad_cp(self, inputs: List[InputFeature], **kwargs) -> List[InputFeature]: - # Normalize after pool+reduce so swapping pooling strategy (last/mean/CLS) stays mathematically correct. - embeddings = F.normalize(pooled, p=2, dim=-1) - - outputs = copy(outputs) - outputs.pop("features", None) - outputs["embeddings"] = embeddings.contiguous() - return inputs, outputs - - def _postprocess_generative_reranker( - self, inputs: dict[str, Any], outputs: dict[str, Any], sp_strategy=None - ) -> tuple[dict[str, Any], dict[str, Any]]: - """Reduce per-token generative-reranker scores ``[B, T, 1]`` to ``[n_seqs, 1]``. - - Symmetric with embedding pooling (same last-valid-token gather, SP/packed-aware), but the - source is ``outputs['logits']`` -- the yes-minus-no score TransformersGenerativeRerankerPatch - writes per token -- and there is NO L2 normalize: a relevance score is a scalar, not a - direction. Mirrors legacy RerankerTrainer, which reads ``logits[batch, last_valid]``. - """ - logits = outputs.get("logits") - assert logits is not None, ( - "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) - 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 - ) -> tuple[dict[str, Any], dict[str, Any]]: - """Pool per-token seq_cls/reranker scores ``[B, T, num_labels]`` to ``[n_seqs, num_labels]``. - - Same last-valid-token logic as embedding/generative_reranker, packed-aware. - TransformersValuePatch(squeeze_single=False) surfaces the per-token scores; - this method reduces them so the loss sees one score per sequence. - """ - logits = outputs.get("logits") - assert logits is not None, ( - "outputs['logits'] missing; ensure TransformersValuePatch is applied for task='seq_cls'/'reranker'." - ) - if logits.dim() == 2: - # 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) - 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 - ) -> 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)``. - """ - 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 - elif inputs.get("position_ids") is not None: - ref_pos = inputs["position_ids"] - else: - # Standard padded batch: HF omits position_ids (it derives them internally). - # Rebuild a position_ids-like reference (valid token -> its index, padding -> -1) - # so the last-valid-token pooling below, which keys off ``ref_pos >= 0`` for - # validity and a single ``== 0`` per row for packing detection, still works. - T = features.shape[1] - arange = torch.arange(T, device=features.device) - am = inputs.get("attention_mask") - if am is None: - ref_pos = arange.expand(features.shape[0], T) - else: - if am.dim() >= 3: - am = am.reshape(am.shape[0], -1)[:, :T] - ref_pos = torch.where(am.bool(), arange, arange.new_tensor(-1)) - if ref_pos.dim() == 3: - ref_pos = ref_pos[0] - # 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")) - - is_packed = features.shape[0] == 1 and (cu_seq_lens_q is not None or int((ref_pos.reshape(-1) == 0).sum()) > 1) - - device, dtype = features.device, features.dtype - T_real = ref_pos.shape[-1] - - if is_packed: - if torch.is_tensor(cu_seq_lens_q) and cu_seq_lens_q.numel() >= 2: - end_idx = (cu_seq_lens_q[1:].long() - 1).to(device) - else: - end_idx = self._packed_last_indices(ref_pos, T_real).to(device) - n_seqs = end_idx.shape[0] - batch_idx = torch.zeros_like(end_idx) - target_idx = torch.arange(n_seqs, device=device) - B_mask, n_targets = 1, n_seqs - 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) - 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 - - 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) - - 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." - ) - - 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) - - return pooled, is_packed - - def pad_cp(self, inputs: list[InputFeature], **kwargs) -> list[InputFeature]: if self.device_mesh is None: return inputs if self.framework == "transformers": @@ -526,7 +295,7 @@ def add_extra_padding_free_args(self, inputs: list[InputFeature], **kwargs) -> l def prepare_gdn_padding_free_args(self, inputs: list[InputFeature], **kwargs) -> list[InputFeature]: if self.framework != "transformers" or not kwargs.get("gdn_padding_free_enabled", False): return inputs - padding_free = bool(self.padding_free or self._any_packing(inputs)) + padding_free = bool(self.is_padding_free(inputs)) if not padding_free or bool(kwargs.get("enable_sp", False)): return inputs @@ -651,6 +420,11 @@ def _any_packing(inputs: list[InputFeature]): return True return False + def is_padding_free(self, inputs: InputFeature | list[InputFeature]) -> bool: + """Return the effective packed/padding-free layout without applying parallelism.""" + rows = inputs if isinstance(inputs, list) else [inputs] + return bool(self.padding_free or self._any_packing(rows)) + @staticmethod def _unpack_by_position_ids( position_ids: "torch.Tensor", @@ -699,10 +473,12 @@ def _unpack_by_position_ids( def unpack_packed_sequences( self, - inputs: dict[str, Any], - outputs: dict[str, Any] | None = None, - task: str = "causal_lm", - ) -> tuple[dict[str, Any], dict[str, Any] | None]: + inputs: Dict[str, Any], + outputs: Optional[Dict[str, Any]] = None, + task: str = 'causal_lm', + input_key_specs=None, + output_key_specs=None, + ) -> tuple[Dict[str, Any], Optional[Dict[str, Any]]]: """Unpack packed (padding_free) sequences into per-sequence batch format. Called after SP gather / CP gather, before loss computation. @@ -711,55 +487,53 @@ def unpack_packed_sequences( ``[num_sequences, max_seq_len, ...]``. Keys that are ``None`` are silently skipped. + Input/output key specs describe the tensor layout and padding values; + this restoration step does not require a parallel strategy. + For ``task='embedding'`` the outputs are already pooled to ``[n_seqs, H]`` - by ``postprocess_tensor_sp``, 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``. + by model-side task postprocessing, so this is a no-op. Same for + ``task='generative_reranker'``, whose logits are already reduced to ``[n_seqs, 1]``. """ if task in ("embedding", "generative_reranker", "seq_cls", "reranker"): return inputs, outputs - labels = inputs.get("labels") position_ids = inputs.get("position_ids") - if labels is None or position_ids is None: + if position_ids is None: return inputs, outputs if not self._is_packed_position_ids(position_ids): return inputs, outputs from copy import copy - # (key, tensor, pad_value) for everything that must come back as - # [num_sequences, max_seq_len]. ``channel`` stays sample-level metadata, - # while ``loss_scale`` and ``completion_mask`` share the labels' index - # space, so leaving them packed would hand the loss differently shaped - # views of the same sequence. - input_specs = [("labels", labels, -100)] - if inputs.get("loss_scale") is not None: - input_specs.append(("loss_scale", inputs["loss_scale"], self.padding_map["loss_scale"])) - if inputs.get("completion_mask") is not None: - input_specs.append(("completion_mask", inputs["completion_mask"], self.padding_map["completion_mask"])) - output_specs = [ - (key, outputs[key], 0) - for key in ("logps", "values", "entropies", "logits") - if outputs and outputs.get(key) is not None - ] - - specs = input_specs + output_specs - unpacked = iter( - self._unpack_by_position_ids( - position_ids, *[tensor for _, tensor, _ in specs], padding_values=[pad for _, _, pad in specs] - ) + from .fields import ( + COMMON_INPUT_SPECS, COMMON_OUTPUT_SPECS, ) + input_specs = COMMON_INPUT_SPECS if input_key_specs is None else input_key_specs + output_specs = COMMON_OUTPUT_SPECS if output_key_specs is None else output_key_specs + input_keys = [spec for spec in input_specs + if spec.stage == 'auxiliary' and spec.unpack_packed and inputs.get(spec.name) is not None] + output_keys = [spec for spec in output_specs + if spec.unpack_packed and outputs and outputs.get(spec.name) is not None] + all_specs = input_keys + output_keys + all_tensors = [inputs[spec.name] for spec in input_keys] + [outputs[spec.name] for spec in output_keys] + # The unpack primitive expects [1, T, ...]; restore the declared axis + # afterwards so tensors like [1, H, T] retain their public layout. + axes = [spec.sequence_dim(tensor.dim()) for spec, tensor in zip(all_specs, all_tensors)] + all_tensors = [tensor.movedim(axis, 1) for tensor, axis in zip(all_tensors, axes)] + all_pads = [spec.pad_value for spec in all_specs] + unpacked = self._unpack_by_position_ids(position_ids, *all_tensors, padding_values=all_pads) + unpacked = [tensor.movedim(1, axis) for tensor, axis in zip(unpacked, axes)] + inputs = copy(inputs) - for key, _, _ in input_specs: - inputs[key] = next(unpacked) + for i, spec in enumerate(input_keys): + inputs[spec.name] = unpacked[i] - if output_specs: + if output_keys: outputs = copy(outputs) - for key, _, _ in output_specs: - outputs[key] = next(unpacked) + offset = len(input_keys) + for i, spec in enumerate(output_keys): + outputs[spec.name] = unpacked[offset + i] return inputs, outputs @@ -772,31 +546,14 @@ def unpack_inputs(self, inputs: list[dict[str, Any]], task: str = "causal_lm") - return [self.unpack_packed_sequences(inp, task=task)[0] for inp in inputs] @staticmethod - def to_transformers_dict(inputs: list[InputFeature], **kwargs) -> list[InputFeature]: + def to_transformers_dict(inputs: List[InputFeature], input_key_specs=None, **kwargs) -> List[InputFeature]: + import torch + + registered_keys = [spec.name for spec in input_key_specs] if input_key_specs is not None else [] results = [] for _input in inputs: output = {} - _keys = [ - "input_ids", - "inputs_embeds", - "attention_mask", - "position_ids", - "labels", - "completion_mask", - "loss_scale", - "channel", - "cu_seq_lens_q", - "cu_seq_lens_k", - "cu_seqlens_q", - "cu_seqlens_kv", - "max_length_q", - "max_length_k", - "packed_seq_params", - "routed_experts", - "replay_mask", - "mm_token_type_ids", - "second_per_grid_ts", - ] + list(InputProcessor.VLM_CONCAT_FIELDS) + _keys = TRANSFORMERS_INPUT_KEYS | set(registered_keys) for key in list(_input.keys()): if key not in _keys: continue @@ -808,13 +565,43 @@ def to_transformers_dict(inputs: list[InputFeature], **kwargs) -> list[InputFeat results.append(InputFeature(**output)) return results + def _collate_sequence_key(self, values, spec, padding_free): + samples = [] + seq_dim = None + for value in values: + if not torch.is_tensor(value): + raise TypeError(f'Registered sequence input {spec.name!r} must be a tensor before collation') + if value.dim() == 1: + value = value.unsqueeze(0) + dim = spec.sequence_dim(value.dim()) + if dim == 0 or value.shape[0] != 1: + raise ValueError(f'Sequence input {spec.name!r} requires one sample with a leading batch dimension') + if seq_dim is not None and seq_dim != dim - 1: + raise ValueError(f'Inconsistent sequence axis for input {spec.name!r}') + seq_dim = dim - 1 + samples.append(value.squeeze(0)) + if padding_free: + return torch.cat(samples, dim=seq_dim).unsqueeze(0) + length = max(value.shape[seq_dim] for value in samples) + padded = [] + for value in samples: + shape = list(value.shape) + shape[seq_dim] = length + result = value.new_full(shape, spec.pad_value) + slices = [slice(None)] * value.dim() + start = length - value.shape[seq_dim] if self.padding_side == 'left' else 0 + slices[seq_dim] = slice(start, start + value.shape[seq_dim]) + result[tuple(slices)] = value + padded.append(result) + return torch.stack(padded) + # when training on mixed text + multimodal data, some fields (e.g. mm_token_type_ids) # only exist on multimodal samples. We pad those fields for the missing samples instead of # letting collate blow up with a KeyError. def _fill_optional_sequence_fields(self, batch: list[InputFeature]) -> None: if len(batch) < 2: return - seq_fields = set(self.padding_map) - set(self.VLM_CONCAT_FIELDS) + seq_fields = set(self.padding_map) - set(VLM_CONCAT_FIELDS) present = {k for feat in batch for k in feat if k in seq_fields} for key in present: missing = [feat for feat in batch if feat.get(key) is None] @@ -843,26 +630,20 @@ def _fill_optional_sequence_fields(self, batch: list[InputFeature]) -> None: continue feat[key] = torch.full(shape, pad_value, dtype=torch.long, device=device) - #: Audio fields whose leading dim is the audio count under the per-audio-batched contract - #: (see _collate_macro_batch); the blanket squeeze must not fold that dim away. The - #: plane-masked contract instead RELIES on the squeeze folding a single-audio leading dim. - #: - #: TODO (template-layer cleanup): this constant and the probe in _collate_macro_batch are - #: model knowledge that does not belong in a processor (task-related, model-agnostic). The - #: field-name probe decides which model family produced the batch -- today only the qwen-side - #: template emits ``feature_attention_mask`` and only the gemma-side template emits - #: ``input_features_mask``, so the name is a proxy for the model -- and the reshape branch - #: below bakes in the consuming model's layout (trailing singleton, mask expanded to the - #: feature plane). End state: every template emits batch-concat-ready audio fields (dim-0 - #: concat yields the layout its model consumes: the gemma template applies the trailing - #: singleton and the plane-expanded mask itself; the qwen-side shape already concatenates - #: correctly), and this class drops the probe, this constant and both reshape branches, - #: keeping only "audio fields skip the blanket squeeze and concat on dim 0". Land it with - #: the template migration: the legacy gemma collator and test_multimodal assertions consume - #: these fields and must change in the same commit. - _PER_AUDIO_BATCHED_FIELDS = ("input_features", "feature_attention_mask") - - def _collate_macro_batch(self, inputs: list[InputFeature]) -> InputFeature: + #: Audio fields whose leading dim is the audio count and must survive collate under the batched + #: Qwen-audio contract (see _collate_macro_batch). The gemma4 channels-last contract instead + #: relies on the blanket squeeze folding a single-audio leading dim away, so the squeeze is only + #: skipped for these fields when the batch is on the Qwen-audio contract. + _PER_AUDIO_BATCHED_FIELDS = ('input_features', 'feature_attention_mask') + + def _collate_macro_batch(self, inputs: List[InputFeature], input_key_specs=None) -> InputFeature: + import torch + from .fields import COMMON_INPUT_SPECS + + builtin_names = {spec.name for spec in COMMON_INPUT_SPECS} + specs = COMMON_INPUT_SPECS if input_key_specs is None else input_key_specs + key_specs = {spec.name: spec for spec in specs + if spec.stage == 'auxiliary' or (spec.stage == 'forward' and spec.name not in builtin_names)} # Two mutually-exclusive audio contracts share the `input_features` field name and are told # apart by the mask field that rides alongside it -- the producer (the model's template) # emits one pair or the other, so the field name IS the contract marker: @@ -885,17 +666,18 @@ def _collate_macro_batch(self, inputs: list[InputFeature]) -> InputFeature: for key in list(_input.keys()): if isinstance(_input[key], torch.Tensor): if per_audio_batched and key in self._PER_AUDIO_BATCHED_FIELDS: - continue # keep the audio-count leading dim intact - _input[key] = _input[key].squeeze() + continue # keep the audio-count leading dim intact for get_audio_features + if key not in key_specs: + _input[key] = _input[key].squeeze() squeezed.append(_input) inputs = squeezed self._fill_optional_sequence_fields(inputs) - vlm_fields = {k: [] for k in self.VLM_CONCAT_FIELDS} + vlm_fields = {k: [] for k in VLM_CONCAT_FIELDS} text_inputs = [] for inp in inputs: inp = dict(inp) - for field in self.VLM_CONCAT_FIELDS: + for field in VLM_CONCAT_FIELDS: if field in inp: vlm_fields[field].append(inp.pop(field)) text_inputs.append(inp) @@ -911,11 +693,14 @@ def is_mm_position_ids(position_ids): return False return position_ids.dim() > 1 and position_ids.shape[0] > 1 - padding_free = self.padding_free or self._any_packing(inputs) + padding_free = self.is_padding_free(inputs) if padding_free: for key in text_keys: values = [item[key] for item in text_inputs] - if key == "attention_mask": + if key in key_specs: + result[key] = self._collate_sequence_key(values, key_specs[key], True) + continue + elif key == 'attention_mask': # attention_mask is not needed continue if key == "position_ids" and is_mm_position_ids(values[0]): @@ -930,7 +715,9 @@ def is_mm_position_ids(position_ids): else: for key in text_keys: values = [item[key] for item in text_inputs] - if self.framework == "megatron" and key == "attention_mask": + if key in key_specs: + result[key] = self._collate_sequence_key(values, key_specs[key], False) + elif self.framework == 'megatron' and key == 'attention_mask': result[key] = self._create_4d_attention_mask(values) elif key == "position_ids" and is_mm_position_ids(values[0]): result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side) @@ -975,33 +762,37 @@ def is_mm_position_ids(position_ids): result[field] = _values return result - def collate_fn( - self, inputs: list[InputFeature], micro_batch_size: int | None = None, variable_seq_lengths=False, **kwargs - ) -> list[InputFeature]: - if len(inputs) == 1 and self.framework != "megatron": + def collate_fn(self, + inputs: List[InputFeature], + micro_batch_size: Optional[int] = None, + variable_seq_lengths=False, + input_key_specs=None, + **kwargs) -> List[InputFeature]: + if len(inputs) == 1 and self.framework != 'megatron': return inputs if micro_batch_size is None: # normal collate - outputs = self._collate_macro_batch(inputs) + outputs = self._collate_macro_batch(inputs, input_key_specs=input_key_specs) for key in outputs: - if key in self.VLM_CONCAT_FIELDS: + if key in VLM_CONCAT_FIELDS: outputs[key] = torch.cat(outputs[key], dim=0) return [outputs] - padding_free = self.padding_free or self._any_packing(inputs) + padding_free = self.is_padding_free(inputs) if variable_seq_lengths or padding_free: # each micro batch has its own packed length assert len(inputs) >= micro_batch_size outputs = [] for i in range(0, len(inputs), micro_batch_size): - _output = self._collate_macro_batch(inputs[i : i + micro_batch_size]) + _output = self._collate_macro_batch(inputs[i:i + micro_batch_size], + input_key_specs=input_key_specs) for key in _output: - if key in self.VLM_CONCAT_FIELDS: + if key in VLM_CONCAT_FIELDS: _output[key] = torch.cat(_output[key], dim=0) outputs.append(_output) return outputs else: # each macro batch shares the same length - res = self._collate_macro_batch(inputs) + res = self._collate_macro_batch(inputs, input_key_specs=input_key_specs) keys = list(res.keys()) outputs = [] for i in range(0, len(inputs), micro_batch_size): @@ -1010,7 +801,7 @@ def collate_fn( for key in keys: if key == "position_ids" and res[key].dim() > 2: output[key] = res[key][:, i:end, :] - elif key in self.VLM_CONCAT_FIELDS: + elif key in VLM_CONCAT_FIELDS: output[key] = torch.cat(res[key][i:end], dim=0) else: output[key] = res[key][i:end] diff --git a/src/twinkle/processor/fields.py b/src/twinkle/processor/fields.py new file mode 100644 index 000000000..853bc0739 --- /dev/null +++ b/src/twinkle/processor/fields.py @@ -0,0 +1,123 @@ +"""InputProcessor field handling: collation, filtering and packed restoration. + +These rules cover the processor's supported encoded inputs. Raw media is encoded +by templates first; model/loss providers declare any additional sequence keys. +SP's execution defaults live separately in its data_keys module. +""" + +from twinkle.data_format.sequence_fields import DataKeySpec, declared_key_specs + +VLM_CONCAT_FIELDS = { + "pixel_values", + "image_grid_thw", + "pixel_values_videos", + "video_grid_thw", + "input_features", + "input_features_mask", + "feature_attention_mask", + "grid_thws", +} + +PROCESSOR_PASSTHROUGH_KEYS = ("cache_position", "padding_free") + +# Single source for built-in sequence layouts. SP derives execution defaults +# from these declarations; the processor never schedules SP transforms. +COMMON_INPUT_SPECS = ( + DataKeySpec("input_ids", split_dim=1, pad_value="tokenizer", stage="forward"), + DataKeySpec("inputs_embeds", split_dim=1, pad_value=0.0, stage="forward"), + DataKeySpec("position_ids", pad_value=-1, stage="forward"), + DataKeySpec("attention_mask", split_dim=1, stage="forward"), + DataKeySpec("labels", pad_value=-100), + DataKeySpec("completion_mask"), + DataKeySpec("loss_scale", pad_value=0.0), +) +_FORWARD_KEYS = tuple(spec.name for spec in COMMON_INPUT_SPECS if spec.stage == "forward") +TRANSFORMERS_INPUT_KEYS = frozenset(spec.name for spec in COMMON_INPUT_SPECS) | frozenset( + { + "channel", + "cu_seq_lens_q", + "cu_seq_lens_k", + "cu_seqlens_q", + "cu_seqlens_kv", + "max_length_q", + "max_length_k", + "packed_seq_params", + "routed_experts", + "replay_mask", + "mm_token_type_ids", + "second_per_grid_ts", + *VLM_CONCAT_FIELDS, + *PROCESSOR_PASSTHROUGH_KEYS, + } +) + +# Numeric padding for existing processor/CP paths. Processor has no tokenizer: +# preserve its input_ids fallback of 0; SP resolves "tokenizer" in its context. +PADDING_VALUES = { + spec.name: 0 if spec.pad_value == "tokenizer" else spec.pad_value + for spec in COMMON_INPUT_SPECS +} +PADDING_VALUES.update({ + "mm_token_type_ids": 0, + "length": -1, + "pixel_values": 0.0, + "image_grid_thw": 0, + "pixel_values_videos": 0.0, + "video_grid_thw": 0, + "input_features": 0.0, + "feature_attention_mask": 0, + "routed_experts": 0, + "replay_mask": 0, +}) + +# Only default sequence-aligned outputs need unpacking here. Task-specific +# outputs (e.g. PPO values) are declared by their model/loss provider. +COMMON_OUTPUT_SPECS = tuple( + DataKeySpec(name, split_dim=1, side="output", stage="output") for name in ("logps", "entropies", "logits") +) + + +def validate_input_key_spec(spec): + """Validate provider and direct-registration inputs with the same rules.""" + if not isinstance(spec, DataKeySpec): + raise TypeError("Expected a DataKeySpec") + if spec.side != "input" or spec.stage not in ("auxiliary", "forward"): + raise ValueError("Register only auxiliary or model-forward input keys here") + if spec.name in _FORWARD_KEYS: + raise ValueError(f"{spec.name!r} is owned by the model forward hook") + if spec.name in PROCESSOR_PASSTHROUGH_KEYS: + raise ValueError(f"{spec.name!r} is a special processor/model argument, not an SP data key") + if not spec.model_input and "forward" in (spec.pad_at, spec.split_at): + raise ValueError( + f"Loss-only input {spec.name!r} cannot use the forward boundary; " + "use prepare, or set model_input=True for a model argument" + ) + + +def resolve_key_specs(*providers): + """Merge provider schemas independently of any SP strategy or execution registry.""" + inputs = {spec.name: spec for spec in COMMON_INPUT_SPECS} + outputs = {spec.name: spec for spec in COMMON_OUTPUT_SPECS} + declared = {"input": set(), "output": set()} + layout_fields = ("name", "split_dim", "pad_value", "side", "stage", "model_input") + for provider in providers: + for side, registry in (("input", inputs), ("output", outputs)): + for spec in declared_key_specs(provider, side): + if side == "input": + validate_input_key_spec(spec) + else: + if not isinstance(spec, DataKeySpec): + raise TypeError("Expected a DataKeySpec") + if spec.side != side: + raise ValueError(f"Expected an {side} key rule") + previous = registry.get(spec.name) + if previous is not None and previous != spec: + # Defaults describe layouts only. A provider may add its SP + # operations without changing a built-in field's layout. + if spec.name in declared[side] or any( + getattr(previous, field) != getattr(spec, field) for field in layout_fields + ): + raise ValueError(f"Conflicting {side} field rule for {spec.name!r}") + registry[spec.name] = spec + declared[side].add(spec.name) + return tuple(inputs.values()), tuple(outputs.values()) diff --git a/src/twinkle/protocol/types/data.py b/src/twinkle/protocol/types/data.py index b7936dbcd..89420d32e 100644 --- a/src/twinkle/protocol/types/data.py +++ b/src/twinkle/protocol/types/data.py @@ -72,7 +72,7 @@ # The VLM tensor fields batched by concatenation rather than padding. Declared here # because this module must stay free of Twinkle_Core's heavyweight imports; a -# consistency test asserts this set equals ``InputProcessor.VLM_CONCAT_FIELDS``, so a +# consistency test asserts this set equals ``processor.fields.VLM_CONCAT_FIELDS``, so a # future addition there fails loudly instead of being silently dropped on the wire. VLM_TENSOR_FIELDS: frozenset[str] = frozenset({ 'pixel_values', diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index 1ea97a44f..193ad51e5 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -136,9 +136,9 @@ def _build_sampling_params(num_samples: int, sampling_params: dict | None, force def _vision_feature_keys() -> frozenset: """The multimodal feature keys to lift from the engine's feature into ``encoded``.""" - from twinkle.processor.base import InputProcessor + from twinkle.processor.fields import VLM_CONCAT_FIELDS - return frozenset(InputProcessor.VLM_CONCAT_FIELDS) | frozenset(("mm_token_type_ids", "second_per_grid_ts")) + return frozenset(VLM_CONCAT_FIELDS) | frozenset(("mm_token_type_ids", "second_per_grid_ts")) def trajectory_to_rollout_sample(traj: dict[str, Any], prompt_id: str, extra: dict[str, Any]) -> RolloutSample: diff --git a/tests/loss/test_channel_loss.py b/tests/loss/test_channel_loss.py index 2d6c8f008..a39b3f421 100644 --- a/tests/loss/test_channel_loss.py +++ b/tests/loss/test_channel_loss.py @@ -125,9 +125,6 @@ def __init__(self): def __call__(self, inputs, **kwargs): return dict(inputs) - def postprocess_tensor_sp(self, inputs, outputs, **kwargs): - return inputs, outputs - def unpack_packed_sequences(self, inputs, outputs=None, **kwargs): return inputs, outputs diff --git a/tests/server/validation/test_wire_schema.py b/tests/server/validation/test_wire_schema.py index 09b18e7d7..19e64c325 100644 --- a/tests/server/validation/test_wire_schema.py +++ b/tests/server/validation/test_wire_schema.py @@ -12,7 +12,6 @@ from typing import Union, get_args, get_origin from twinkle.data_format.encoding import ENCODED_INPUT_KEYS, is_encoded -from twinkle.processor.base import InputProcessor from twinkle.protocol.types import data as wire _INPUTS = TypeAdapter(wire.WireInputBatch) @@ -202,7 +201,9 @@ def test_vlm_field_set_matches_the_processor(): A field added to ``VLM_CONCAT_FIELDS`` and not here would be silently absent from the wire while the batching code still expects it. """ - assert wire.VLM_TENSOR_FIELDS == frozenset(InputProcessor.VLM_CONCAT_FIELDS) + from twinkle.processor.fields import VLM_CONCAT_FIELDS + + assert wire.VLM_TENSOR_FIELDS == frozenset(VLM_CONCAT_FIELDS) if __name__ == '__main__': diff --git a/tests/transformers/test_dpo_sp_cpu.py b/tests/transformers/test_dpo_sp_cpu.py new file mode 100644 index 000000000..b8331b25a --- /dev/null +++ b/tests/transformers/test_dpo_sp_cpu.py @@ -0,0 +1,107 @@ +"""Offline DPO loss/gradient equivalence over two CPU/Gloo SP ranks.""" +import copy +import os +import pytest +import socket +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from datetime import timedelta +from types import SimpleNamespace +from transformers import LlamaConfig, LlamaForCausalLM + +from twinkle import DeviceMesh +from twinkle.loss import DPOLoss +from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy, _current_ctx +from twinkle.utils import selective_log_softmax + + +def _worker(rank, port): + os.environ['RANK'] = str(rank) + os.environ['WORLD_SIZE'] = '2' + torch.set_num_threads(1) + torch.manual_seed(42) + # torch_npu sets checkpoint's default device to NPU even for CPU-only + # tensors. This regression deliberately checkpoints on CPU. + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, + world_size=2, timeout=timedelta(seconds=90)) + try: + config = LlamaConfig(vocab_size=64, hidden_size=32, intermediate_size=64, + num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=4, + attention_dropout=0.0, pad_token_id=0) + config._attn_implementation = 'sdpa' + baseline = LlamaForCausalLM(config).eval() + reference = copy.deepcopy(baseline).eval() + with torch.no_grad(): + for param in baseline.parameters(): + param.add_(torch.randn_like(param) * 1e-3) + policy = copy.deepcopy(baseline) + ref_sp = copy.deepcopy(reference) + # Four rows = two interleaved chosen/rejected pairs. Length 9 exercises SP padding. + ids = torch.randint(1, 64, (4, 9)) + positions = torch.arange(9).expand(4, -1) + labels = torch.cat((ids[:, 1:], torch.full((4, 1), -100)), dim=1) + labels[:, :2] = -100 + loss_fn = DPOLoss(beta=0.1, loss_type='sigmoid') + + def score(model, strategy=None): + local_labels = labels + if strategy is not None: + local_labels = strategy.preprocess_inputs( + {'input_ids': ids, 'position_ids': positions, 'labels': labels})['labels'] + logits = model(input_ids=ids, position_ids=positions, use_cache=False).logits + outputs = {'logps': selective_log_softmax(logits, local_labels.masked_fill(local_labels == -100, 0))} + inputs = {'labels': local_labels} + if strategy is not None: + inputs, outputs = strategy.gather_loss_tensors(inputs, outputs) + return inputs, outputs + + with torch.no_grad(): + _, base_ref = score(reference) + base_inputs, base_outputs = score(baseline) + base_loss = loss_fn(base_inputs, base_outputs, ref_outputs=base_ref)['loss'] + base_loss.backward() + base_grads = {name: p.grad.clone() for name, p in baseline.named_parameters()} + + mesh = DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu') + strategies = [] + for model in (policy, ref_sp): + strategy = SequenceParallelStrategy(mesh, model=model) + strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0) + assert strategy.initialize() + strategies.append(strategy) + assert strategies[0]._ctx is not strategies[1]._ctx + policy.gradient_checkpointing_enable() + policy.train() + with torch.no_grad(): + _, parallel_ref = score(ref_sp, strategies[1]) + inputs, outputs = score(policy, strategies[0]) + parallel_loss = loss_fn(inputs, outputs, ref_outputs=parallel_ref)['loss'] + torch.testing.assert_close(parallel_loss, base_loss, atol=1e-6, rtol=1e-5) + parallel_loss.backward() + # Mirror the data-parallel gradient averaging performed by DDP. + for name, param in policy.named_parameters(): + assert param.grad is not None, name + dist.all_reduce(param.grad) + param.grad.div_(2) + torch.testing.assert_close(param.grad, base_grads[name], atol=1e-6, rtol=1e-4, msg=name) + assert _current_ctx() is None + # Avoid torch_npu's process-wide optimizer capability probe on this + # explicitly CPU-only test. + optimizer = torch.optim.AdamW(policy.parameters(), lr=1e-4, foreach=False, fused=False) + before = policy.lm_head.weight.detach().clone() + optimizer.step() + assert not torch.equal(before, policy.lm_head.weight) + assert all(param.grad is None for param in ref_sp.parameters()) + finally: + dist.destroy_process_group() + + +@pytest.mark.slow +def test_dpo_sp_loss_all_parameter_gradients_and_step(): + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + mp.spawn(_worker, args=(port,), nprocs=2, join=True) diff --git a/tests/transformers/test_qwen35_linear_attention_sp.py b/tests/transformers/test_qwen35_linear_attention_sp.py index d46c423cf..4289be684 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 @@ -215,7 +215,6 @@ def _make_strategy(model: torch.nn.Module, world_size: int) -> SequenceParallelS sp_config={ 'enabled': True, 'ulysses_size': world_size, - 'gather_logits': True, }, model=model, tokenizer_id=None, @@ -350,7 +349,7 @@ def _run_prefill_alignment_worker(rank: int, position_ids=position_ids, use_cache=False, ) - gathered_outputs = strategy.postprocess_outputs(copy.copy(sp_outputs)) + _, gathered_outputs = strategy.gather_loss_tensors({}, sp_outputs, sp_gather_keys=('logits',)) 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 +366,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}') diff --git a/tests/transformers/test_sequence_parallel_and_cp.py b/tests/transformers/test_sequence_parallel_and_cp.py index a2ce49e61..b0ecc4359 100644 --- a/tests/transformers/test_sequence_parallel_and_cp.py +++ b/tests/transformers/test_sequence_parallel_and_cp.py @@ -10,17 +10,17 @@ from datetime import timedelta from transformers import LlamaConfig, LlamaForCausalLM 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 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.') +pytestmark = [pytest.mark.slow, pytest.mark.accel(2)] LOGITS_RTOL = 1e-2 -LOGITS_ATOL = 5e-3 +LOGITS_ATOL = 1e-2 LOSS_ATOL = 5e-3 GRAD_RTOL = 2e-2 GRAD_ATOL = 1e-2 @@ -50,6 +50,7 @@ def _make_padded_labels(input_ids: torch.Tensor, valid_lengths: list[int]) -> to def _make_case(case_name: str) -> dict: cases = { 'sp_only': { + 'num_hidden_layers': 2, 'expected_mode': 'sp_only', 'world_size': 2, 'ulysses_size': 2, @@ -58,6 +59,7 @@ def _make_case(case_name: str) -> dict: 'seq_len': 769, }, 'sp_only_multi_sample': { + 'num_hidden_layers': 2, 'expected_mode': 'sp_only', 'world_size': 2, 'ulysses_size': 2, @@ -67,6 +69,7 @@ def _make_case(case_name: str) -> dict: 'batch_size': 3, }, 'sp_only_multi_sample_masked': { + 'num_hidden_layers': 2, 'expected_mode': 'sp_only', 'world_size': 2, 'ulysses_size': 2, @@ -250,7 +253,6 @@ def _make_strategy(model: LlamaForCausalLM, device_mesh: DeviceMesh, ulysses_siz sp_config={ 'enabled': True, 'ulysses_size': ulysses_size, - 'gather_logits': True, }, model=model, tokenizer_id=None, @@ -366,16 +368,22 @@ def _assert_logits_close(case_name: str, rank: int, baseline_logits: torch.Tenso seq_len: int, attention_mask: torch.Tensor | None): baseline_slice = baseline_logits[:, :seq_len] sp_slice = sp_logits[:, :seq_len] - if attention_mask is None: - if not torch.allclose(sp_slice, baseline_slice, rtol=LOGITS_RTOL, atol=LOGITS_ATOL): - diff = (sp_slice - baseline_slice).abs().max().item() - raise AssertionError(f'{case_name} logits mismatch on rank {rank}: max_diff={diff}') - return - - mask = attention_mask.to(dtype=torch.bool).unsqueeze(-1).expand_as(baseline_slice) - if not torch.allclose(sp_slice[mask], baseline_slice[mask], rtol=LOGITS_RTOL, atol=LOGITS_ATOL): - diff = (sp_slice[mask] - baseline_slice[mask]).abs().max().item() - raise AssertionError(f'{case_name} masked logits mismatch on rank {rank}: max_diff={diff}') + if attention_mask is not None: + mask = attention_mask.to(dtype=torch.bool).unsqueeze(-1).expand_as(baseline_slice) + baseline_slice, sp_slice = baseline_slice[mask], sp_slice[mask] + if not torch.allclose(sp_slice, baseline_slice, rtol=LOGITS_RTOL, atol=LOGITS_ATOL): + difference = (sp_slice - baseline_slice).abs() + tolerance = LOGITS_ATOL + LOGITS_RTOL * baseline_slice.abs() + failing = ~torch.isclose(sp_slice, baseline_slice, rtol=LOGITS_RTOL, atol=LOGITS_ATOL) + worst = (difference / tolerance).flatten().argmax() + baseline_value = baseline_slice.flatten()[worst].item() + sp_value = sp_slice.flatten()[worst].item() + raise AssertionError( + f'{case_name} logits mismatch on rank {rank}: max_diff={difference.max().item()}, ' + f'mean_diff={difference.mean().item()}, ' + f'failed_elements={failing.sum().item()}/{failing.numel()}, ' + f'worst_flat_index={worst.item()}, baseline={baseline_value}, sp={sp_value}, ' + f'allowed_error={tolerance.flatten()[worst].item()}') def _init_dist(rank: int, world_size: int, port: int, backend_config: dict): @@ -534,13 +542,13 @@ def _run_precision_worker(rank: int, use_cache=False, ) local_logits = sp_outputs.logits - gathered_outputs = strategy.postprocess_outputs(sp_outputs) + _, gathered_outputs = strategy.gather_loss_tensors({}, sp_outputs, sp_gather_keys=('logits',)) 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: @@ -720,6 +728,178 @@ def test_cp_sp_precision_alignment(self): join=True) +def _run_multi_model_worker(rank: int, world_size: int, port: int, backend_config: dict, + attn_implementation: str): + """Two SP-enabled models (policy/ref style) with interleaved forwards. + + Same architecture, different weights and different sequence lengths. Each + model must use its own SP context: a shared/singleton context shows up here + as cross-model position_ids pollution (logits mismatch or shape errors). + """ + device = _init_dist(rank, world_size, port, backend_config) + try: + device_type = backend_config['device_type'] + case = _make_case('sp_only') + vocab_size = 256 + + _seed_backend(1234, device_type) + policy_base = _build_tiny_llama(case, device, device_type, attn_implementation) + _seed_backend(5678, device_type) + ref_base = _build_tiny_llama(case, device, device_type, attn_implementation) + + policy_sp = _build_tiny_llama(case, device, device_type, attn_implementation) + policy_sp.load_state_dict(policy_base.state_dict()) + ref_sp = _build_tiny_llama(case, device, device_type, attn_implementation) + ref_sp.load_state_dict(ref_base.state_dict()) + + seq_len_p = int(case['seq_len']) + seq_len_r = seq_len_p // 2 + _seed_backend(42, device_type) + input_ids_p = torch.randint(low=0, high=vocab_size, size=(1, seq_len_p), device=device) + position_ids_p = torch.arange(seq_len_p, device=device).unsqueeze(0) + input_ids_r = torch.randint(low=0, high=vocab_size, size=(1, seq_len_r), device=device) + position_ids_r = torch.arange(seq_len_r, device=device).unsqueeze(0) + labels_p = _make_labels(input_ids_p) + labels_r = _make_labels(input_ids_r) + + # Non-SP baselines (computed before SP strategies exist). + policy_base.zero_grad(set_to_none=True) + base_outputs_p = policy_base( + input_ids=input_ids_p, position_ids=position_ids_p, attention_mask=None, use_cache=False) + base_logits_p = base_outputs_p.logits.detach().float() + base_loss_sum_p, base_num_tokens_p = _compute_training_path_loss(base_outputs_p.logits, labels_p) + base_loss_sum_p.backward() + base_grads_p = _normalize_grad_dict(_collect_attention_param_grads(policy_base), base_num_tokens_p) + + ref_base.zero_grad(set_to_none=True) + base_outputs_r = ref_base( + input_ids=input_ids_r, position_ids=position_ids_r, attention_mask=None, use_cache=False) + base_logits_r = base_outputs_r.logits.detach().float() + base_loss_sum_r, base_num_tokens_r = _compute_training_path_loss(base_outputs_r.logits, labels_r) + base_loss_sum_r.backward() + base_grads_r = _normalize_grad_dict(_collect_attention_param_grads(ref_base), base_num_tokens_r) + + # Two independent SP strategies on the same device mesh. + device_mesh = DeviceMesh.from_sizes( + fsdp_size=world_size, + dp_size=1, + ulysses_size=2, + device_type=device_type, + ) + policy_strategy = _make_strategy(policy_sp, device_mesh, 2) + ref_strategy = _make_strategy(ref_sp, device_mesh, 2) + + if policy_strategy._ctx is ref_strategy._ctx: + raise AssertionError('multi_model: strategies must own distinct SP contexts.') + + # Module registry must route each model's attention to its own context. + from twinkle.model.transformers.strategy.sequence_parallel import _MODULE_TO_CTX + policy_attn = next(m for m in policy_sp.modules() if m.__class__.__name__.endswith('Attention')) + ref_attn = next(m for m in ref_sp.modules() if m.__class__.__name__.endswith('Attention')) + if _MODULE_TO_CTX.get(policy_attn) is not policy_strategy._ctx: + raise AssertionError('multi_model: policy attention module not registered to policy ctx.') + if _MODULE_TO_CTX.get(ref_attn) is not ref_strategy._ctx: + raise AssertionError('multi_model: ref attention module not registered to ref ctx.') + + local_labels_p = _prepare_label_inputs(policy_strategy, input_ids_p, position_ids_p) + local_labels_r = _prepare_label_inputs(ref_strategy, input_ids_r, position_ids_r) + + # Interleaved forwards: policy -> ref -> policy. + policy_sp.zero_grad(set_to_none=True) + ref_sp.zero_grad(set_to_none=True) + sp_outputs_p1 = policy_sp( + input_ids=input_ids_p, position_ids=position_ids_p, attention_mask=None, use_cache=False) + sp_outputs_r = ref_sp( + input_ids=input_ids_r, position_ids=position_ids_r, attention_mask=None, use_cache=False) + sp_outputs_p2 = policy_sp( + input_ids=input_ids_p, position_ids=position_ids_p, attention_mask=None, use_cache=False) + + # Context isolation: each ctx must still hold its own sequence metadata. + if policy_strategy._ctx.real_position_ids.shape[-1] != seq_len_p: + raise AssertionError('multi_model: policy ctx position_ids polluted by ref forward.') + if ref_strategy._ctx.real_position_ids.shape[-1] != seq_len_r: + raise AssertionError('multi_model: ref ctx position_ids polluted by policy forward.') + + # Capture per-rank local logits (with grad) before loss-side gathering. + # Both paths must use the same registered gather implementation as training. + local_logits_p1 = sp_outputs_p1.logits + local_logits_r = sp_outputs_r.logits + + _, gathered = policy_strategy.gather_loss_tensors({}, sp_outputs_p1, sp_gather_keys=('logits',)) + sp_logits_p1 = gathered.logits.detach().float() + _, gathered = ref_strategy.gather_loss_tensors({}, sp_outputs_r, sp_gather_keys=('logits',)) + sp_logits_r = gathered.logits.detach().float() + _, gathered = policy_strategy.gather_loss_tensors({}, sp_outputs_p2, sp_gather_keys=('logits',)) + sp_logits_p2 = gathered.logits.detach().float() + + _assert_logits_close('multi_model policy fwd1', rank, base_logits_p, sp_logits_p1, seq_len_p, None) + _assert_logits_close('multi_model ref fwd', rank, base_logits_r, sp_logits_r, seq_len_r, None) + _assert_logits_close('multi_model policy fwd2', rank, base_logits_p, sp_logits_p2, seq_len_p, None) + + # Interleaved loss + backward (ref backward first, then policy): + # GatherLoss.backward must use the ctx captured at forward time. + sp_loss_sum_p, sp_num_tokens_p = _compute_training_path_loss(local_logits_p1, local_labels_p, + policy_strategy) + sp_loss_sum_r, sp_num_tokens_r = _compute_training_path_loss(local_logits_r, local_labels_r, + ref_strategy) + sp_loss_sum_r.backward() + sp_loss_sum_p.backward() + _average_model_grads_over_group(policy_sp, policy_strategy._ctx._data_rank_group) + _average_model_grads_over_group(ref_sp, ref_strategy._ctx._data_rank_group) + sp_grads_p = _normalize_grad_dict(_collect_attention_param_grads(policy_sp), sp_num_tokens_p) + sp_grads_r = _normalize_grad_dict(_collect_attention_param_grads(ref_sp), sp_num_tokens_r) + + if sp_num_tokens_p != base_num_tokens_p or sp_num_tokens_r != base_num_tokens_r: + raise AssertionError(f'multi_model num_tokens mismatch on rank {rank}: ' + f'policy sp={sp_num_tokens_p} base={base_num_tokens_p}, ' + f'ref sp={sp_num_tokens_r} base={base_num_tokens_r}') + if abs(sp_loss_sum_p.item() / max(sp_num_tokens_p, 1) - base_loss_sum_p.item() / max(base_num_tokens_p, + 1)) > LOSS_ATOL: + raise AssertionError(f'multi_model policy loss mismatch on rank {rank}: ' + f'sp={sp_loss_sum_p.item()} base={base_loss_sum_p.item()}') + if abs(sp_loss_sum_r.item() / max(sp_num_tokens_r, 1) - base_loss_sum_r.item() / max(base_num_tokens_r, + 1)) > LOSS_ATOL: + raise AssertionError(f'multi_model ref loss mismatch on rank {rank}: ' + f'sp={sp_loss_sum_r.item()} base={base_loss_sum_r.item()}') + _assert_grad_dict_close('multi_model policy', rank, base_grads_p, sp_grads_p) + _assert_grad_dict_close('multi_model ref', rank, base_grads_r, sp_grads_r) + dist.barrier() + finally: + dist.destroy_process_group() + + +class TestMultiModelSequenceParallel: + """Multi-model (policy/ref style) SP context isolation tests.""" + + def _get_backend_or_skip(self, world_size: int = 2) -> dict: + backend = _get_runtime_backend() + if backend is None: + pytest.skip('CUDA or NPU is required for multi-model SP tests.') + if backend['device_count'] < world_size: + pytest.skip(f'Requires at least {world_size} {backend["label"]} devices.') + return backend + + @pytest.mark.parametrize('attn_implementation', ['sdpa', 'flash_attention_2', 'flash_attention_3']) + def test_two_models_interleaved_forward(self, attn_implementation): + world_size = 2 + backend = self._get_backend_or_skip(world_size) + if attn_implementation == 'flash_attention_2': + # CUDA must detect FA2 specifically; another installed FA backend + # does not satisfy the model's flash_attention_2 dispatch check. + available = (is_flash_attn_2_available() if backend['device_type'] == 'cuda' + else is_flash_attn_available()) + if not available: + pytest.skip('flash_attention_2 is unavailable.') + elif attn_implementation == 'flash_attention_3': + if backend['device_type'] != 'cuda': + pytest.skip('This flash_attention_3 test requires CUDA.') + if not is_flash_attn_3_available(): + pytest.skip('flash_attention_3 is unavailable.') + port = _find_free_port() + mp.spawn(_run_multi_model_worker, args=(world_size, port, backend, attn_implementation), + nprocs=world_size, join=True) + + class TestDerivedRingMemoryProfile: def _get_backend_or_skip(self, world_size: int = 4) -> dict: diff --git a/tests/transformers/test_sp_context_isolation.py b/tests/transformers/test_sp_context_isolation.py new file mode 100644 index 000000000..5dd8c6900 --- /dev/null +++ b/tests/transformers/test_sp_context_isolation.py @@ -0,0 +1,244 @@ +"""CPU regression tests for process-wide SP patches and exception cleanup.""" +import pytest +import torch +import gc +import weakref +from types import SimpleNamespace + +import twinkle.model.transformers.strategy.sequence_parallel as sp +from twinkle.model.transformers.strategy.sequence_parallel import linear_attention_sp + + +@pytest.mark.parametrize('positions, expected', [ + ([[0, 1, 2, -1]], [[0, 1, 2, 3]]), + ([[0, 1, -1, -1]], [[0, 1, 2, 3]]), + ([[0, 1, -1, 0, 1, -1]], [[0, 1, 2, 0, 1, 2]]), + ([[5, 6, -1, 5, 6, -1]], [[5, 6, 7, 5, 6, 7]]), + ([[0, 1, -1], [0, 1, 2]], [[0, 1, 2], [0, 1, 2]]), + ([[-1, -1, -1]], [[0, 1, 2]]), +]) +def test_flash_boundary_positions_preserve_real_tokens_and_packed_starts(positions, expected): + positions = torch.tensor(positions) + original = positions.clone() + received = sp._call_flash_attention(lambda **kwargs: kwargs['position_ids'], position_ids=positions) + torch.testing.assert_close(received, torch.tensor(expected)) + torch.testing.assert_close(positions, original) + # Mirrors Transformers 5.15's minimum-based boundary detection. + for actual_row, expected_row in zip(received, torch.tensor(expected)): + actual_starts = (actual_row == actual_row.min()).nonzero().flatten() + expected_starts = (expected_row == expected_row.min()).nonzero().flatten() + torch.testing.assert_close(actual_starts, expected_starts) + assert actual_starts[0].item() == 0 + + +@pytest.mark.parametrize('backend', ['flash_attention_2', 'flash_attention_3']) +def test_flash_origin_receives_valid_boundary_metadata(monkeypatch, backend): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + positions = torch.tensor([[0, 1, 2, -1]]) + + def origin(module, query, key, value, *args, **kwargs): + torch.testing.assert_close(kwargs['position_ids'], torch.tensor([[0, 1, 2, 3]])) + assert kwargs['is_causal'] is True + return query.transpose(1, 2), None + + # Global wrapper installation keeps this compatibility alias even when + # the selected backend is FA3 and no FA2 kernel is installed. + monkeypatch.setitem(ALL_ATTENTION_FUNCTIONS, 'flash_attention_2_origin', origin) + monkeypatch.setitem(ALL_ATTENTION_FUNCTIONS, f'{backend}_origin', origin) + ctx = sp.SequenceParallel() + ctx.world_size = ctx.sp_world_size = ctx.rp_world_size = 1 + module = torch.nn.Linear(2, 2) + module.config = SimpleNamespace(_attn_implementation=backend) + query = torch.randn(1, 2, 4, 2) + actual, _ = ctx._local_flash_attn_impl( + module, query, query, query, None, position_ids=positions, is_causal=True) + torch.testing.assert_close(actual, query.transpose(1, 2)) + torch.testing.assert_close(positions, torch.tensor([[0, 1, 2, -1]])) + + +@pytest.mark.parametrize('backend', ['sdpa', 'flash_attention_2']) +def test_attention_cache_routes_current_module_and_releases_owners(monkeypatch, backend): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + seen = [] + + def origin(module, query, key, value, *args, **kwargs): + assert '_sp_module' not in kwargs + seen.append(weakref.ref(module)) + return query, None + + monkeypatch.setitem(ALL_ATTENTION_FUNCTIONS, f'{backend}_origin', origin) + ctx = sp.SequenceParallel() + ctx.world_size = ctx.sp_world_size = ctx.rp_world_size = 1 + first, second = torch.nn.Linear(2, 2), torch.nn.Linear(2, 2) + refs = [weakref.ref(first), weakref.ref(second)] + sp._MODULE_TO_CTX[first] = ctx + sp._MODULE_TO_CTX[second] = ctx + x = torch.zeros(1, 1, 2, 2) + method = ctx._local_sdpa_attn_impl if backend == 'sdpa' else ctx._local_flash_attn_impl + try: + method(first, x, x, x, None) + method(second, x, x, x, None) + assert [r() for r in seen] == [first, second] + del first, second, method, ctx + gc.collect() + assert all(r() is None for r in refs) + finally: + for ref in refs: + if ref() is not None: + sp._MODULE_TO_CTX.pop(ref(), None) + + +@pytest.fixture(autouse=True) +def empty_context_stack(): + saved = list(sp._ctx_stack()) + sp._ctx_stack().clear() + yield + sp._ctx_stack()[:] = saved + + +@pytest.mark.parametrize('failure', ['forward', 'split', 'other_pre_hook']) +def test_exception_restores_enclosing_context(failure): + outer = object() + sp._push_ctx(outer) + ctx = sp.SequenceParallel() + ctx.world_size = 2 if failure == 'split' else 1 + + class Model(torch.nn.Module): + def forward(self, **kwargs): + assert sp._current_ctx() is ctx + raise RuntimeError('forward failed') + + model = Model() + if failure == 'other_pre_hook': + def fail(*args): + assert sp._current_ctx() is ctx + raise RuntimeError('pre-hook failed') + model.register_forward_pre_hook(fail) + ctx._prepare_forward_hook(model) + with pytest.raises((RuntimeError, KeyError)): + model() + assert sp._ctx_stack() == [outer] + sp._pop_ctx(outer) + assert sp._current_ctx() is None + + +def test_input_split_pre_hook_failure_does_not_leak_context(): + ctx = sp.SequenceParallel() + ctx.world_size = 2 + + model = torch.nn.Linear(2, 2) + ctx._prepare_forward_hook(model) + with pytest.raises(KeyError, match='position_ids'): + model(torch.zeros(1, 2)) + assert sp._current_ctx() is None + + # A later model in this thread must not inherit the failed SP model's ctx. + plain_model = torch.nn.Linear(2, 2) + assert sp._current_ctx() is None + plain_model(torch.zeros(1, 2)) + assert sp._current_ctx() is None + + +def test_gdn_patch_routes_each_owner_and_preserves_non_sp(monkeypatch): + class Gate(torch.nn.Module): + def forward(self, hidden_states, **kwargs): + return ('original', hidden_states) + + monkeypatch.setattr(linear_attention_sp, '_iter_qwen35_gated_delta_net_classes', lambda: [Gate]) + patch_type = linear_attention_sp.Qwen3_5GatedDeltaNetUlyssesPatch + monkeypatch.setattr(patch_type, '_run_forward', staticmethod( + lambda mod, hidden_states, **kwargs: kwargs['sequence_parallel_context'])) + policy_ctx = SimpleNamespace(world_size=2, sp_world_size=2, rp_world_size=1) + ref_ctx = SimpleNamespace(world_size=2, sp_world_size=2, rp_world_size=1) + policy, ref, plain = Gate(), Gate(), Gate() + monkeypatch.setitem(sp._MODULE_TO_CTX, policy, policy_ctx) + monkeypatch.setitem(sp._MODULE_TO_CTX, ref, ref_ctx) + patch_type()(None, sequence_parallel=policy_ctx) + patch_type()(None, sequence_parallel=ref_ctx) + x = torch.zeros(1) + assert policy(x) is policy_ctx + assert ref(x) is ref_ctx + assert policy(x) is policy_ctx + assert plain(x)[0] == 'original' + assert plain(x, sequence_parallel_context=ref_ctx) is ref_ctx + + +def test_router_replay_uses_model_owner_outside_forward(monkeypatch): + from twinkle.model.transformers.moe.router_replay import RouterReplay as rr + + class Context: + world_size = sp_world_size = 2 + real_position_ids = None + + def __init__(self, offset): + self.offset = offset + + def gather(self, tensor, **kwargs): + return tensor + self.offset + + models = [torch.nn.Linear(1, 1) for _ in range(3)] + monkeypatch.setitem(sp._MODULE_TO_CTX, models[0], Context(10)) + monkeypatch.setitem(sp._MODULE_TO_CTX, models[1], Context(20)) + monkeypatch.setattr(rr, '_find_moe_blocks_with_names', lambda model: [('block', model)]) + monkeypatch.setattr(rr, '_registry', {'block': SimpleNamespace(recorded_indices=torch.zeros(2, 1))}) + assert sp._current_ctx() is None + assert rr.get_router_replay_data(models[0]).eq(10).all() + assert rr.get_router_replay_data(models[1]).eq(20).all() + assert rr.get_router_replay_data(models[2]).eq(0).all() + + +@pytest.mark.parametrize('rank', [0, 1]) +def test_router_replay_shards_mask_with_routing_indices(monkeypatch, rank): + from twinkle.model.transformers.moe.router_replay import RouterReplay + + model = torch.nn.Linear(1, 1) + ctx = sp.SequenceParallel() + ctx.world_size = ctx.sp_world_size = 2 + ctx.rp_world_size = 1 + ctx._sp_rank = rank + ctx.extra_kwargs['position_ids'] = torch.arange(5).unsqueeze(0) + monkeypatch.setattr(RouterReplay, '_find_moe_blocks_with_names', lambda _: [('block', model)]) + state = SimpleNamespace(target_indices=None, target_mask=None) + monkeypatch.setattr(RouterReplay, '_registry', {'block': state}) + mask = torch.tensor([[True, False, True, False, True]]) + routing = torch.arange(5).reshape(1, 5, 1, 1) + from twinkle import Platform + monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu')) + RouterReplay.set_router_replay_data(routing, model, mask, sp_context=ctx) + expected_indices = torch.tensor([0, 1, 2, 3, 4, 0]).chunk(2)[rank] + expected_mask = torch.tensor([True, False, True, False, True, False]).chunk(2)[rank] + torch.testing.assert_close(state.target_indices[:, 0], expected_indices) + torch.testing.assert_close(state.target_mask, expected_mask) + + +@pytest.mark.parametrize('nested', [False, True]) +def test_prepare_causal_mask_method_does_not_retain_weak_registry_owner(monkeypatch, nested): + class Backbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.config = SimpleNamespace(num_attention_heads=2, _attn_implementation='sdpa') + + def _update_causal_mask(self, mask, *args): + return mask + self.weight + + model = Backbone() + if nested: + wrapper = torch.nn.Module() + wrapper.language_model = model + wrapper.config = model.config + model = wrapper + del wrapper + refs = [weakref.ref(module) for module in model.modules()] + ctx = sp.SequenceParallel() + monkeypatch.setattr(sp, 'get_llm_model', lambda value: value) + monkeypatch.setattr(sp.SequenceParallel, '_global_inited', True) + monkeypatch.setattr(sp.SequenceParallel, '_init_device_mesh', lambda *args: None) + ctx.prepare(2, model, SimpleNamespace(pad_token_id=0)) + torch.testing.assert_close(ctx.causal_mask_func(torch.zeros(1)), torch.ones(1)) + del model + gc.collect() + assert all(ref() is None for ref in refs) + assert ctx.causal_mask_func is None diff --git a/tests/transformers/test_sp_custom_loss_cpu.py b/tests/transformers/test_sp_custom_loss_cpu.py new file mode 100644 index 000000000..e84b403ad --- /dev/null +++ b/tests/transformers/test_sp_custom_loss_cpu.py @@ -0,0 +1,220 @@ +"""End-to-end contract test for a user-defined sequence loss.""" +import importlib +import os +import socket +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM + +from twinkle import DeviceMesh, Platform +from twinkle.data_format import LossOutput +from twinkle.loss import Loss +from twinkle.model.transformers.transformers import TransformersModel +from twinkle.model.transformers.strategy.sequence_parallel import DataKeySpec, SequenceParallelStrategy + + +class TokenRegressionLoss(Loss): + """Example plugin loss with custom sequence input and output keys.""" + + require_logps = False + input_key_specs = (DataKeySpec('token_targets', split_dim=1, pad_value=0),) + output_key_specs = (DataKeySpec( + 'token_scores', + split_dim=1, + side='output', + stage='output', + gather_at='loss', + autograd=True, + ),) + + def __call__(self, inputs, outputs, **kwargs): + scores = outputs['token_scores'] + targets = inputs['token_targets'].to(dtype=scores.dtype) + loss = (scores - targets).square().mean() + return LossOutput(loss=loss, num_tokens=loss.new_tensor(targets.numel())) + + +class _ToyTokenRegressor(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(num_key_value_heads=2, num_attention_heads=2, + _attn_implementation='sdpa', model_type='toy_token_regressor') + self.scale = torch.nn.Parameter(torch.tensor(0.7)) + + def forward(self, input_ids, position_ids, attention_mask=None, **kwargs): + return {'token_scores': input_ids.to(self.scale.dtype) * self.scale} + + +class _NoopTrainStrategy: + """Keep this CPU E2E focused on Twinkle's model/loss/optimizer flow.""" + + def capture_pre_ep_state_if_needed(self, model, *, enable_ep): + return None + + def unwrap_model(self, model): + return model + + def wrap_model(self, model, *args): + return (model, *args) if args else model + + +def _worker(rank, port): + os.environ['RANK'] = str(rank) + os.environ['WORLD_SIZE'] = '2' + torch.set_num_threads(1) + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, + world_size=2, timeout=timedelta(seconds=90)) + try: + ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7], [2, 3, 4, 5, 6, 7, 8]]) + positions = torch.arange(ids.shape[1]).expand_as(ids) + targets = ids.float() * 0.4 + 0.25 + loss_fn = TokenRegressionLoss() + + baseline = _ToyTokenRegressor() + full_outputs = baseline(input_ids=ids, position_ids=positions) + full_loss = loss_fn({'token_targets': targets}, full_outputs)['loss'] + full_loss.backward() + expected_grad = baseline.scale.grad.clone() + + model = _ToyTokenRegressor() + strategy = SequenceParallelStrategy( + DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu'), model=model) + strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0) + strategy.register_loss_key_specs(loss_fn) + assert strategy.initialize() + + prepared = strategy.preprocess_inputs({'input_ids': ids, 'position_ids': positions, + 'token_targets': targets}) + local_targets = prepared.pop('token_targets') + local_outputs = model(input_ids=prepared['input_ids'], position_ids=prepared['position_ids']) + loss_inputs, loss_outputs = strategy.gather_loss_tensors({'token_targets': local_targets}, local_outputs) + sp_loss = loss_fn(loss_inputs, loss_outputs)['loss'] + torch.testing.assert_close(sp_loss, full_loss, atol=1e-7, rtol=1e-6) + sp_loss.backward() + dist.all_reduce(model.scale.grad) + model.scale.grad.div_(2) + torch.testing.assert_close(model.scale.grad, expected_grad, atol=1e-7, rtol=1e-6) + finally: + dist.destroy_process_group() + + +@pytest.mark.slow +def test_custom_token_regression_loss_matches_full_sequence_loss_and_gradient(): + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + mp.spawn(_worker, args=(port,), nprocs=2, join=True) + + +def _e2e_worker(rank, port, sp_size): + os.environ['RANK'] = str(rank) + os.environ['WORLD_SIZE'] = '2' + torch.set_num_threads(1) + Platform.get_local_device = staticmethod(lambda: torch.device('cpu')) + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, + world_size=2, timeout=timedelta(seconds=90)) + try: + config = LlamaConfig(vocab_size=32, hidden_size=16, intermediate_size=32, + num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=4, + attention_dropout=0.0, pad_token_id=0, use_cache=False) + config._attn_implementation = 'sdpa' + torch.manual_seed(1234) + baseline = LlamaForCausalLM(config).train() + + mesh = DeviceMesh.from_sizes(dp_size=2, ulysses_size=sp_size, device_type='cpu') + model_module = importlib.import_module('twinkle.model.transformers.transformers') + original_template = model_module.Template + model_module.Template = lambda *_args, **_kwargs: SimpleNamespace(is_mm=False) + try: + training_model = TransformersModel( + model_cls=AutoModelForCausalLM, + config=config, + device_mesh=mesh, + mixed_precision='no', + ) + finally: + model_module.Template = original_template + training_model.strategy = _NoopTrainStrategy() + # This local CPU harness supplies worker mesh state explicitly. + training_model.device_mesh = mesh + training_model._enable_sp = sp_size > 1 + model_forward = training_model.model.forward + + def forward_with_token_scores(*args, **kwargs): + outputs = model_forward(*args, **kwargs) + outputs['token_scores'] = outputs.logits[..., 0] + return outputs + + training_model.model.forward = forward_with_token_scores + training_model.model.load_state_dict(baseline.state_dict()) + training_model.model.gradient_checkpointing_disable() + training_model.model.train() + strategy = SequenceParallelStrategy(mesh, model=training_model.model) if sp_size > 1 else None + if strategy is not None: + strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0) + training_model.sp_strategy = strategy + + optimizer_group = training_model.optimizer_group[training_model._get_default_group()] + loss_fn = TokenRegressionLoss() + optimizer_group.loss_instance = loss_fn + assert optimizer_group.loss_instance.input_key_specs + training_model._ensure_sp_strategy() + if strategy is not None: + assert 'token_targets' in {spec.name for spec in strategy.input_key_specs} + + ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7], [2, 3, 4, 5, 6, 7, 8]]) + positions = torch.arange(ids.shape[1]).expand_as(ids) + targets = ids.float() * 0.4 + 0.25 + labels = torch.full_like(ids, -100) + batch = [ + {'input_ids': ids[row], 'position_ids': positions[row], 'labels': labels[row], + 'token_targets': targets[row]} + for row in range(ids.shape[0]) + ] + + # Baseline uses the same model weights and the same optimizer update, + # but computes the loss on the unsharded full sequence. + base_optimizer = torch.optim.SGD(baseline.parameters(), lr=1e-3) + baseline_output = baseline(input_ids=ids, position_ids=positions, use_cache=False) + baseline_output['token_scores'] = baseline_output.logits[..., 0] + baseline_loss = loss_fn({'token_targets': targets}, baseline_output)['loss'] + baseline_loss.backward() + base_optimizer.step() + + # End-to-end Twinkle path: processor -> SP setup/hook -> forward -> + # loss registration/collection -> calculate_loss -> backward -> step. + forward_outputs = training_model.forward(inputs=batch) + assert 'token_targets' in optimizer_group.train_status.inputs, optimizer_group.train_status.inputs.keys() + sp_loss = training_model.calculate_loss() + torch.testing.assert_close(torch.tensor(sp_loss), baseline_loss.detach(), atol=1e-6, rtol=1e-5) + training_model.backward() + optimizer = torch.optim.SGD(training_model.model.parameters(), lr=1e-3) + for parameter in training_model.model.parameters(): + assert parameter.grad is not None + dist.all_reduce(parameter.grad) + parameter.grad.div_(2) + optimizer_group.optimizer = optimizer + training_model.step() + assert 'token_scores' in forward_outputs + for (name, actual), (_, expected) in zip(training_model.model.named_parameters(), baseline.named_parameters()): + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5, msg=name) + finally: + dist.destroy_process_group() + + +@pytest.mark.slow +@pytest.mark.parametrize('sp_size', [1, 2]) +def test_custom_loss_through_transformers_training_pipeline_and_optimizer_step(sp_size): + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + mp.spawn(_e2e_worker, args=(port, sp_size), nprocs=2, join=True) diff --git a/tests/transformers/test_sp_data_keys.py b/tests/transformers/test_sp_data_keys.py new file mode 100644 index 000000000..01bdeb9e7 --- /dev/null +++ b/tests/transformers/test_sp_data_keys.py @@ -0,0 +1,877 @@ +"""Input-rule checks without model training or accelerator allocation.""" +import pytest +import torch + +from twinkle.model.transformers.strategy.sequence_parallel import DataKeySpec, SequenceParallelStrategy +from twinkle.data_format.sequence_fields import partition_model_inputs +from twinkle.processor.fields import resolve_key_specs +from twinkle.model.transformers.strategy.sequence_parallel.data_keys import SPKeyCollection, SPKeyValue + + +def _strategy(rank=0): + strategy = SequenceParallelStrategy(sp_config={'ulysses_size': 2}) + ctx = strategy._ctx + ctx.world_size = ctx.sp_world_size = 2 + ctx.rp_world_size = 1 + ctx._sp_rank = rank + return strategy + + +@pytest.mark.parametrize('rank', [0, 1]) +@pytest.mark.parametrize('length', [8, 9]) +def test_builtin_labels_match_existing_pad_split(rank, length): + strategy = _strategy(rank) + positions = torch.arange(length).expand(2, -1) + labels = torch.arange(2 * length).reshape(2, length) + spec = strategy._ctx.input_key_registry.get('labels') + padded_labels = strategy._ctx._pad_key(labels, spec, positions) + padded_positions = strategy._ctx.pad(positions, padding_value=-1, position_ids=positions) + old = strategy._ctx._split_key(padded_labels, spec, padded_positions) + inputs = {'labels': labels.clone(), 'position_ids': positions} + actual = strategy.preprocess_inputs(inputs)['labels'] + torch.testing.assert_close(actual, old, rtol=0, atol=0) + torch.testing.assert_close(labels, torch.arange(2 * length).reshape(2, length)) + + +@pytest.mark.parametrize('rank', [0, 1]) +def test_custom_key_padding_dimension_and_model_isolation(rank): + strategy = _strategy(rank) + other = _strategy(rank) + spec = DataKeySpec('token_weights', split_dim=1, pad_value=-2, pad_at='prepare', split_at='prepare') + strategy.register_input_key(spec) + strategy.register_input_key(spec) # Identical declarations are safe. + values = torch.arange(10).reshape(1, 5, 2) + positions = torch.arange(5).unsqueeze(0) + prepared = strategy.preprocess_inputs({'token_weights': values, 'position_ids': positions}) + expected = torch.cat((values, torch.full((1, 1, 2), -2)), dim=1) + torch.testing.assert_close(prepared['token_weights'], expected[:, rank * 3:(rank + 1) * 3]) + untouched = other.preprocess_inputs({'token_weights': values, 'position_ids': positions}) + assert untouched['token_weights'] is values + + +@pytest.mark.parametrize('rank', [0, 1]) +def test_key_collection_executes_only_at_declared_boundary(rank): + strategy = _strategy(rank) + ctx = strategy._ctx + ctx.tokenizer = type('Tokenizer', (), {'pad_token_id': 0})() + token_type_spec = DataKeySpec( + 'token_type_ids', + split_dim=1, + pad_value=-1, + stage='forward', + pad_at='forward', + split_at='forward', + ) + strategy.register_input_key(token_type_spec) + values = { + 'input_ids': torch.arange(5).expand(2, -1), + 'position_ids': torch.arange(5).expand(2, -1), + 'attention_mask': torch.ones((2, 5), dtype=torch.long), + 'token_type_ids': torch.arange(5).expand(2, -1), + } + specs = tuple(spec for spec in ctx.input_key_registry if spec.stage == 'forward') + keys = SPKeyCollection.from_mapping(values, specs) + ctx.pad_and_split_inputs(keys, boundary='forward', real_position_ids=values['position_ids']) + assert keys.get('input_ids').value.shape == (2, 3) + assert keys.get('position_ids').value.shape == (2, 3) + assert keys.get('token_type_ids').value.shape == (2, 3) + # attention_mask follows its own rule: pad to the global frame but do not shard. + torch.testing.assert_close(keys.get('attention_mask').value, torch.tensor([[1, 1, 1, 1, 1, 0]]).expand(2, -1)) + + +def test_key_collection_rejects_duplicate_names(): + spec = DataKeySpec('token_type_ids', stage='forward', pad_at='forward', split_at='forward') + keys = SPKeyCollection() + keys.add(SPKeyValue(spec, torch.ones(1, 2))) + with pytest.raises(ValueError, match='Duplicate SP key'): + keys.add(SPKeyValue(spec, torch.zeros(1, 2))) + + +@pytest.mark.parametrize('rank', [0, 1]) +def test_prepare_keeps_loss_scale_and_completion_mask_global(rank): + strategy = _strategy(rank) + positions = torch.arange(5).unsqueeze(0) + loss_scale = torch.arange(1, 6, dtype=torch.float).unsqueeze(0) + completion_mask = torch.tensor([[0, 0, 1, 1, 1]]) + inputs = { + 'position_ids': positions, + 'labels': torch.arange(5).unsqueeze(0), + 'loss_scale': loss_scale.clone(), + 'completion_mask': completion_mask.clone(), + } + prepared = strategy.preprocess_inputs(inputs) + assert prepared['labels'].shape == (1, 3) + torch.testing.assert_close(prepared['loss_scale'], loss_scale) + torch.testing.assert_close(prepared['completion_mask'], completion_mask) + + +def test_invalid_rules_and_non_tensor_payload_fail_clearly(): + strategy = _strategy() + with pytest.raises(ValueError, match='Invalid SP key stage'): + DataKeySpec('cache_position', stage='passthrough') + with pytest.raises(ValueError, match='Conflicting'): + strategy.register_input_key(DataKeySpec('labels', pad_value=0)) + with pytest.raises(ValueError, match='forward hook'): + strategy.register_input_key(DataKeySpec('input_ids')) + with pytest.raises(ValueError, match='not an SP data key'): + strategy.register_input_key(DataKeySpec('padding_free')) + with pytest.raises(TypeError, match='must be a tensor'): + strategy.preprocess_inputs({'labels': [1, 2]}) + + +def test_disabled_sp_and_absent_keys_are_unchanged(): + strategy = SequenceParallelStrategy(sp_config={'enabled': False, 'ulysses_size': 2}) + inputs = {'labels': torch.tensor([[1, 2]])} + assert strategy.preprocess_inputs(inputs) is inputs + assert _strategy().preprocess_inputs({'labels': None}) == {'labels': None} + + +@pytest.mark.parametrize('sp_rank', [0, 1]) +@pytest.mark.parametrize('rp_rank', [0, 1]) +def test_packed_ring_labels_preserve_existing_layout(sp_rank, rp_rank): + strategy = _strategy(sp_rank) + ctx = strategy._ctx + ctx.world_size = 4 + ctx.rp_world_size = 2 + ctx._rp_rank = rp_rank + positions = torch.tensor([[0, 1, 2, 3, 4, 0, 1, 2]]) + labels = torch.arange(8).unsqueeze(0) + spec = ctx.input_key_registry.get('labels') + padded_labels = ctx._pad_key(labels, spec, positions) + padded_positions = ctx.pad(positions, padding_value=-1, position_ids=positions) + old = ctx._split_key(padded_labels, spec, padded_positions) + new = strategy.preprocess_inputs({'labels': labels, 'position_ids': positions})['labels'] + torch.testing.assert_close(new, old, rtol=0, atol=0) + + +def test_packed_unpack_uses_explicit_per_key_unpack_rule(): + from twinkle.processor.base import InputProcessor + + processor = InputProcessor(padding_free=True) + positions = torch.tensor([[0, 1, 2, 0, 1]]) + labels = torch.tensor([[10, 11, 12, 20, 21]]) + completion_mask = torch.tensor([[0, 1, 1, 0, 1]]) + loss_scale = torch.tensor([[1., 2., 3., 4., 5.]]) + inputs = { + 'position_ids': positions, + 'labels': labels, + 'completion_mask': completion_mask, + 'loss_scale': loss_scale, + } + unpacked, _ = processor.unpack_packed_sequences(inputs, outputs=None) + torch.testing.assert_close(unpacked['labels'], torch.tensor([[10, 11, 12], [20, 21, -100]])) + torch.testing.assert_close(unpacked['completion_mask'], torch.tensor([[0, 1, 1], [0, 1, 0]])) + torch.testing.assert_close(unpacked['loss_scale'], torch.tensor([[1., 2., 3.], [4., 5., 0.]])) + + +@pytest.mark.parametrize('rank', [0, 1]) +def test_processor_preserves_registered_keys_before_sp_split(monkeypatch, rank): + from twinkle import Platform + from twinkle.processor.base import InputProcessor + + monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu')) + strategy = _strategy(rank) + strategy.register_input_key(DataKeySpec('token_weights', pad_value=-2, pad_at='prepare', split_at='prepare')) + processor = InputProcessor() + values = torch.arange(5).unsqueeze(0) + sample = {'input_ids': torch.ones(1, 5, dtype=torch.long), + 'position_ids': torch.arange(5).unsqueeze(0), + 'labels': torch.ones(1, 5, dtype=torch.long), + 'token_weights': values, 'unregistered': values} + result = processor([sample], input_key_specs=strategy.input_key_specs) + result = strategy.preprocess_inputs(result) + expected = torch.tensor([[0, 1, 2, 3, 4, -2]])[:, rank * 3:(rank + 1) * 3] + torch.testing.assert_close(result['token_weights'], expected) + assert 'unregistered' not in result + other = _strategy(rank) + assert 'token_weights' not in processor.to_transformers_dict([sample], input_key_specs=other.input_key_specs)[0] + + +@pytest.mark.parametrize('rank', [0, 1]) +def test_builtin_completion_mask_and_loss_scale_stay_global_during_prepare(rank): + strategy = _strategy(rank) + payload = {'position_ids': torch.arange(5).unsqueeze(0), + 'completion_mask': torch.tensor([[0, 0, 1, 1, 1]]), + 'loss_scale': torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]])} + prepared = strategy.preprocess_inputs(payload) + torch.testing.assert_close(prepared['completion_mask'], torch.tensor([[0, 0, 1, 1, 1]])) + torch.testing.assert_close(prepared['loss_scale'], torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]])) + + +def test_registered_multidimensional_key_collation_and_loss_boundary(monkeypatch): + from twinkle import Platform + from twinkle.processor import InputProcessor + + monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu')) + strategy = _strategy() + strategy.register_input_key(DataKeySpec( + 'token_features', + pad_value=-2, + gather_at='loss', + pad_at='prepare', + split_at='prepare', + )) + processor = InputProcessor() + samples = [] + for length in (3, 5): + samples.append({'input_ids': list(range(1, length + 1)), + 'position_ids': list(range(length)), 'labels': list(range(1, length + 1)), + 'token_features': [ [[1] * length, [2] * length] ], + 'completion_mask': [1] * length}) + result = processor(samples, input_key_specs=strategy.input_key_specs) + result = strategy.preprocess_inputs(result) + expected = torch.tensor([[[1, 1, 1], [2, 2, 2]], [[1, 1, 1], [2, 2, 2]]]) + torch.testing.assert_close(result['token_features'], expected) + result, auxiliary = partition_model_inputs(result, strategy.input_key_specs) + assert set(auxiliary) == {'labels', 'completion_mask', 'token_features'} + assert 'token_features' not in result + + +def test_model_specific_forward_key_is_split_only_in_model_hook(monkeypatch): + from types import SimpleNamespace + from twinkle import Platform + from twinkle.processor import InputProcessor + + monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu')) + strategy = _strategy(1) + strategy._ctx.tokenizer = SimpleNamespace(pad_token_id=0) + strategy.register_input_key(DataKeySpec('token_type_ids', stage='forward', pad_at='forward', split_at='forward')) + + class Model(torch.nn.Module): + def forward(self, token_type_ids, **kwargs): + return token_type_ids + + model = Model() + strategy._ctx._prepare_forward_hook(model) + ids = torch.arange(5).unsqueeze(0) + prepared = InputProcessor()([{'input_ids': ids, 'position_ids': ids, + 'token_type_ids': ids}], input_key_specs=strategy.input_key_specs) + prepared = strategy.preprocess_inputs(prepared) + assert prepared['token_type_ids'].shape == (1, 5) + torch.testing.assert_close(model(**prepared), torch.tensor([[3, 4, 0]])) + + +def test_declared_collection_uses_sequence_axis_and_collects_without_logps(monkeypatch): + import twinkle.model.transformers.strategy.sequence_parallel as sp + + strategy = _strategy() + strategy._ctx.extra_kwargs['position_ids'] = torch.arange(5).unsqueeze(0) + strategy.register_input_key(DataKeySpec( + 'token_features', + gather_at='loss', + pad_value=-2, + pad_at='prepare', + split_at='prepare', + )) + strategy.register_output_key(DataKeySpec( + 'scores', + split_dim=2, + side='output', + stage='output', + gather_at='loss', + autograd=True, + )) + monkeypatch.setattr(strategy._ctx, 'gather', lambda value, dim, position_ids: torch.cat((value, value), dim=dim)) + monkeypatch.setattr(sp.GatherLoss, 'apply', lambda ctx, value, labels, dim, pos: (ctx.gather(value, dim, pos), None)) + local = torch.tensor([[[1, 2, 3], [4, 5, 6]]]) + inputs, outputs = strategy.gather_loss_tensors({'token_features': local}, {'scores': local.float()}) + expected = torch.tensor([[[1, 2, 3, 1, 2], [4, 5, 6, 4, 5]]]) + torch.testing.assert_close(inputs['token_features'], expected) + torch.testing.assert_close(outputs['scores'], expected.float()) + + +def test_packing_unpacks_registered_keys_on_declared_axis(): + from twinkle.processor import InputProcessor + from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS + + input_specs = (*BUILTIN_INPUT_KEYS, DataKeySpec( + 'token_features', + pad_value=-2, + gather_at='loss', + pad_at='prepare', + split_at='prepare', + )) + output_specs = (DataKeySpec('scores', side='output', stage='output', pad_value=-3),) + inputs = {'position_ids': torch.tensor([[0, 1, 2, 0, 1]]), 'labels': torch.tensor([[1, 2, 3, 4, 5]]), + 'completion_mask': torch.ones(1, 5), 'token_features': torch.tensor([[[1, 2, 3, 4, 5]]])} + outputs = {'scores': torch.tensor([[[10, 20, 30, 40, 50]]])} + unpacked, unpacked_outputs = InputProcessor().unpack_packed_sequences( + inputs, outputs, input_key_specs=input_specs, output_key_specs=output_specs) + assert unpacked['token_features'].shape == (2, 1, 3) + assert unpacked['token_features'].tolist() == [[[1, 2, 3]], [[4, 5, -2]]] + assert unpacked['completion_mask'].tolist() == [[1, 1, 1], [1, 1, 0]] + assert unpacked_outputs['scores'].tolist() == [[[10, 20, 30]], [[40, 50, -3]]] + + +@pytest.mark.parametrize('padding_free', [False, True]) +def test_processor_preserves_custom_keys_without_sp_strategy(monkeypatch, padding_free): + from twinkle import Platform + from twinkle.processor import InputProcessor + from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS + + monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu')) + specs = (*BUILTIN_INPUT_KEYS, DataKeySpec('token_features', pad_value=-2)) + samples = [ + {'input_ids': list(range(length)), 'position_ids': list(range(length)), + 'token_features': [[[1] * length, [2] * length]]} + for length in (3, 5) + ] + processor = InputProcessor(padding_free=padding_free) + result = processor(samples, input_key_specs=specs) + expected = (torch.tensor([[[1] * 8, [2] * 8]]) if padding_free else + torch.tensor([[[1, 1, 1, -2, -2], [2, 2, 2, -2, -2]], [[1] * 5, [2] * 5]])) + torch.testing.assert_close(result['token_features'], expected) + result, auxiliary = partition_model_inputs(result, specs) + torch.testing.assert_close(auxiliary['token_features'], expected) + + +@pytest.mark.parametrize('padding_free', [False, True]) +@pytest.mark.parametrize('micro_batch_size,variable_seq_lengths', [(None, False), (2, False), (1, False), (1, True)]) +def test_collation_preserves_declared_sequence_axis(padding_free, micro_batch_size, variable_seq_lengths): + from twinkle.processor import InputProcessor + from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS + + specs = (*BUILTIN_INPUT_KEYS, DataKeySpec('token_features', split_dim=-1, pad_value=-2)) + processor = InputProcessor(padding_free=padding_free) + samples = [ + {'input_ids': torch.tensor([[1, 2, 3]]), + 'position_ids': torch.tensor([[0, 1, 2]]), + 'token_features': torch.tensor([[[10, 11, 12], [20, 21, 22]]])}, + {'input_ids': torch.tensor([[4, 5, 6, 7, 8]]), + 'position_ids': torch.tensor([[0, 1, 2, 3, 4]]), + 'token_features': torch.tensor([[[30, 31, 32, 33, 34], [40, 41, 42, 43, 44]]])}, + ] + + results = processor.collate_fn(samples, micro_batch_size=micro_batch_size, + variable_seq_lengths=variable_seq_lengths, input_key_specs=specs) + first = torch.tensor([[[10, 11, 12], [20, 21, 22]]]) + first_padded = torch.tensor([[[10, 11, 12, -2, -2], [20, 21, 22, -2, -2]]]) + second = torch.tensor([[[30, 31, 32, 33, 34], [40, 41, 42, 43, 44]]]) + if micro_batch_size == 1: + expected = [first if padding_free or variable_seq_lengths else first_padded, second] + elif padding_free: + expected = [torch.cat((first, second), dim=-1)] + else: + expected = [torch.cat((first_padded, second), dim=0)] + assert len(results) == len(expected) + for result, value in zip(results, expected): + torch.testing.assert_close(result['token_features'], value) + + +def test_processor_preserves_special_arguments_outside_data_key_registry(): + from twinkle.processor import InputProcessor + + strategy = _strategy() + spec_names = {spec.name for spec in strategy.input_key_specs} + assert 'cache_position' not in spec_names + assert 'padding_free' not in spec_names + + sample = {'input_ids': torch.tensor([[1, 2]]), + 'cache_position': torch.tensor([0, 1]), + 'padding_free': True} + converted_with_sp = InputProcessor.to_transformers_dict([sample], input_key_specs=strategy.input_key_specs)[0] + converted_without_sp = InputProcessor.to_transformers_dict([sample])[0] + for converted in (converted_with_sp, converted_without_sp): + torch.testing.assert_close(converted['cache_position'], sample['cache_position']) + assert converted['padding_free'] is True + + +@pytest.mark.parametrize('padding_free', [False, True]) +def test_common_loss_fields_collate_and_unpack_without_sp(monkeypatch, padding_free): + from types import SimpleNamespace + from twinkle import Platform + from twinkle.processor import InputProcessor + + monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu')) + loss = SimpleNamespace(input_key_specs=(DataKeySpec('token_targets', split_dim=2, pad_value=-9),)) + input_specs, output_specs = resolve_key_specs(loss) + spec = next(spec for spec in input_specs if spec.name == 'token_targets') + assert not spec.need_pad and not spec.need_split and not spec.gather_at + samples = [dict(input_ids=list(range(length)), position_ids=list(range(length)), + labels=list(range(length)), token_targets=[[[1] * length, [2] * length]]) + for length in (3, 5)] + processor = InputProcessor(padding_free=padding_free) + batch = processor(samples, input_key_specs=input_specs) + model_inputs, loss_inputs = partition_model_inputs(batch, input_specs) + assert 'token_targets' not in model_inputs and 'labels' not in model_inputs + assert loss_inputs['token_targets'].shape == ((1, 2, 8) if padding_free else (2, 2, 5)) + restored, _ = processor.unpack_packed_sequences( + batch, task='causal_lm', input_key_specs=input_specs, output_key_specs=output_specs) + torch.testing.assert_close(restored['token_targets'], torch.tensor([ + [[1, 1, 1, -9, -9], [2, 2, 2, -9, -9]], [[1] * 5, [2] * 5]])) + + +@pytest.mark.parametrize('ring_size', [1, 2]) +def test_global_loss_fields_are_never_gathered(monkeypatch, ring_size): + strategy = _strategy() + strategy._ctx.rp_world_size = ring_size + strategy._ctx.sp_world_size = 2 // ring_size + strategy._ctx.extra_kwargs['position_ids'] = torch.tensor([[0, 1, 2, 0, 1]]) + + def forbidden(*args, **kwargs): + raise AssertionError('Global loss fields must not enter SP collectives') + + monkeypatch.setattr(strategy._ctx, 'gather', forbidden) + mask = torch.tensor([[0, 1, 1, 0, 1]]) + scale = torch.tensor([[0., 1., 1., 0., 1.]]) + inputs, _ = strategy.gather_loss_tensors({'completion_mask': mask, 'loss_scale': scale}, {}) + assert inputs['completion_mask'] is mask and inputs['loss_scale'] is scale + + +def test_model_input_destination_is_independent_of_split_stage(): + spec = DataKeySpec('token_context', stage='auxiliary', model_input=True) + model_inputs, loss_inputs = partition_model_inputs({'token_context': 1, 'channel': ['a']}, (spec,)) + assert model_inputs == {'token_context': 1} + assert loss_inputs == {'channel': ['a']} + + + +def test_explicit_parallel_boundaries_opt_in_without_changing_loss_field_defaults(): + global_field = DataKeySpec('targets') + assert not global_field.need_pad and not global_field.need_split + local_field = DataKeySpec('local_targets', split_at='prepare', pad_at='prepare') + assert local_field.need_split and local_field.need_pad + padded_field = DataKeySpec('padded_targets', pad_at='prepare') + assert padded_field.need_pad and not padded_field.need_split + + +@pytest.mark.parametrize('rank', [0, 1]) +@pytest.mark.parametrize('stage', ['forward', 'auxiliary']) +def test_prepare_then_forward_honors_separate_pad_and_split_boundaries(rank, stage): + from types import SimpleNamespace + + strategy = _strategy(rank) + strategy._ctx.tokenizer = SimpleNamespace(pad_token_id=0) + spec = DataKeySpec( + 'token_type_ids', + split_dim=1, + pad_value=-7, + stage=stage, + model_input=True, + pad_at='prepare', + split_at='forward', + ) + strategy.register_input_key(spec) + + class Model(torch.nn.Module): + def forward(self, token_type_ids, **kwargs): + return token_type_ids + + model = Model() + strategy._ctx._prepare_forward_hook(model) + positions = torch.arange(5).unsqueeze(0) + prepared = strategy.preprocess_inputs(dict(input_ids=positions.clone(), position_ids=positions, + token_type_ids=positions.clone())) + torch.testing.assert_close(prepared['token_type_ids'], torch.tensor([[0, 1, 2, 3, 4, -7]])) + expected = torch.tensor([[0, 1, 2]]) if rank == 0 else torch.tensor([[3, 4, -7]]) + torch.testing.assert_close(model(**prepared), expected) + + +def test_direct_registration_survives_model_schema_resolution_and_processor(monkeypatch): + from types import SimpleNamespace + from twinkle import Platform + from twinkle.model.transformers.transformers import TransformersModel + from twinkle.processor import InputProcessor + + monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu')) + strategy = _strategy() + spec = DataKeySpec('token_weights', pad_value=-2) + strategy.register_input_key(spec) + output_spec = DataKeySpec('token_scores', side='output', stage='output', gather_at='loss') + strategy.register_output_key(output_spec) + model = torch.nn.Linear(2, 2) + wrapper = SimpleNamespace(model=model, strategy=SimpleNamespace(unwrap_model=lambda value: value), + sp_strategy=strategy) + input_specs, output_specs = TransformersModel._get_field_specs(wrapper, None) + assert spec in input_specs and output_spec in output_specs + processor = InputProcessor() + result = processor([dict(input_ids=[1, 2, 3], token_weights=[0.1, 0.2, 0.3])], + input_key_specs=input_specs) + torch.testing.assert_close(result['token_weights'], torch.tensor([[0.1, 0.2, 0.3]])) + _, loss_inputs = partition_model_inputs(result, input_specs) + assert 'token_weights' in loss_inputs + # Consuming provider schemas must not turn them into direct registrations. + strategy.register_key_specs((DataKeySpec('provider_targets'),), ()) + assert strategy.field_declarations.input_key_specs == (spec,) + + +@pytest.mark.parametrize('enable_sp', [False, True]) +def test_model_field_specs_cache_tracks_loss_model_and_direct_registrations(monkeypatch, enable_sp): + from types import SimpleNamespace + from unittest.mock import Mock + from twinkle.model.transformers import transformers as model_module + + resolve = Mock(wraps=model_module.resolve_key_specs) + monkeypatch.setattr(model_module, 'resolve_key_specs', resolve) + wrapper = SimpleNamespace(model=torch.nn.Linear(2, 2), + strategy=SimpleNamespace(unwrap_model=Mock(side_effect=lambda model: model)), + sp_strategy=_strategy() if enable_sp else None) + get_specs = model_module.TransformersModel._get_field_specs + loss = SimpleNamespace(input_key_specs=(DataKeySpec('targets'),)) + first = get_specs(wrapper, loss) + assert get_specs(wrapper, loss) is first + assert resolve.call_count == wrapper.strategy.unwrap_model.call_count == 1 + + # Replacing a loss object cannot reuse the previous loss's schema. + other_loss = SimpleNamespace(input_key_specs=(DataKeySpec('other_targets'),)) + second = get_specs(wrapper, other_loss) + assert any(spec.name == 'other_targets' for spec in second[0]) + assert not any(spec.name == 'targets' for spec in second[0]) + assert resolve.call_count == 2 + + wrapper.model = torch.nn.Linear(2, 2) + get_specs(wrapper, other_loss) + assert resolve.call_count == 3 + wrapper.strategy = SimpleNamespace(unwrap_model=lambda model: model) + get_specs(wrapper, other_loss) + assert resolve.call_count == 4 + if enable_sp: + wrapper.sp_strategy.register_input_key(DataKeySpec('late_targets')) + assert any(spec.name == 'late_targets' for spec in get_specs(wrapper, other_loss)[0]) + wrapper.sp_strategy.register_output_key(DataKeySpec('late_scores', side='output', stage='output')) + assert any(spec.name == 'late_scores' for spec in get_specs(wrapper, other_loss)[1]) + assert resolve.call_count == 6 + wrapper.sp_strategy = _strategy() + assert not any(spec.name == 'late_targets' for spec in get_specs(wrapper, other_loss)[0]) + assert resolve.call_count == 7 + + +def test_set_loss_invalidates_cached_declarations_of_the_same_instance(): + from types import SimpleNamespace + from unittest.mock import Mock + from twinkle.loss import CrossEntropyLoss + from twinkle.model.transformers.transformers import TransformersModel + + loss = CrossEntropyLoss() + loss.input_key_specs = (DataKeySpec('targets'),) + wrapper = SimpleNamespace(model=torch.nn.Linear(2, 2), + strategy=SimpleNamespace(unwrap_model=lambda model: model), sp_strategy=None, + optimizer_group={'default': SimpleNamespace(set_loss=Mock())}, + _get_default_group=lambda: 'default') + first = TransformersModel._get_field_specs(wrapper, loss) + loss.input_key_specs = (DataKeySpec('new_targets'),) + TransformersModel.set_loss.__wrapped__(wrapper, loss) + second = TransformersModel._get_field_specs(wrapper, loss) + assert second is not first + assert any(spec.name == 'new_targets' for spec in second[0]) + + +def test_field_specs_cache_does_not_keep_replaced_model_alive(): + import gc + import weakref + from types import SimpleNamespace + from twinkle.model.transformers.transformers import TransformersModel + + wrapper = SimpleNamespace(model=torch.nn.Linear(2, 2), + strategy=SimpleNamespace(unwrap_model=lambda model: model), sp_strategy=None) + previous_model = weakref.ref(wrapper.model) + TransformersModel._get_field_specs(wrapper, None) + wrapper.model = torch.nn.Linear(2, 2) + gc.collect() + assert previous_model() is None + + +@pytest.mark.parametrize('label_type', ['tensor', 'numpy', 'python']) +@pytest.mark.parametrize('framework', ['transformers', 'megatron']) +def test_scalar_classification_labels_collate_with_common_schema(monkeypatch, label_type, framework): + import numpy as np + from twinkle import Platform + from twinkle.processor import InputProcessor + + monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu')) + convert = {'tensor': torch.tensor, 'numpy': np.array, 'python': int}[label_type] + batch = [dict(input_ids=[1, 2, 3], position_ids=[0, 1, 2], labels=convert(label)) for label in (1, 0)] + result = InputProcessor(framework=framework)(batch) + if framework == 'megatron': + result = result[0] + torch.testing.assert_close(result['labels'], torch.tensor([[1], [0]])) + assert result['input_ids'].shape == (2, 3) + + +@pytest.mark.parametrize('with_labels', [False, True]) +def test_packed_custom_loss_restoration_does_not_require_token_labels(with_labels): + from types import SimpleNamespace + from twinkle.processor import InputProcessor + + loss = SimpleNamespace( + input_key_specs=(DataKeySpec('token_targets', pad_value=-9),), + output_key_specs=(DataKeySpec('token_scores', side='output', stage='output'),)) + input_specs, output_specs = resolve_key_specs(loss) + inputs = dict(position_ids=torch.tensor([[0, 1, 2, 0, 1]]), token_targets=torch.zeros(1, 5)) + if with_labels: + inputs['labels'] = torch.full((1, 5), -100) + outputs = {'token_scores': torch.tensor([[1., 1., 1., 3., 3.]])} + restored, outputs = InputProcessor(padding_free=True).unpack_packed_sequences( + inputs, outputs, input_key_specs=input_specs, output_key_specs=output_specs) + targets, scores = restored['token_targets'], outputs['token_scores'] + torch.testing.assert_close(targets, torch.tensor([[0., 0., 0.], [0., 0., -9.]])) + valid = targets != -9 + balanced_loss = (((scores - targets).square() * valid).sum(-1) / valid.sum(-1)).mean() + torch.testing.assert_close(balanced_loss, torch.tensor(5.)) + assert ('labels' in restored) == with_labels + + +def test_packed_output_logits_restore_without_labels(): + from twinkle.processor import InputProcessor + + logits = torch.arange(10).reshape(1, 5, 2) + _, outputs = InputProcessor(padding_free=True).unpack_packed_sequences( + {'position_ids': torch.tensor([[0, 1, 2, 0, 1]])}, {'logits': logits}) + torch.testing.assert_close(outputs['logits'], torch.tensor([[[0, 1], [2, 3], [4, 5]], + [[6, 7], [8, 9], [0, 0]]])) + + +@pytest.mark.parametrize('operation', ['pad', 'split']) +@pytest.mark.parametrize('entry_point', ['registry', 'schema']) +def test_loss_only_fields_reject_unreachable_forward_operations(operation, entry_point): + from dataclasses import replace + from types import SimpleNamespace + + register = (_strategy().register_input_key if entry_point == 'registry' else + lambda spec: resolve_key_specs(SimpleNamespace(input_key_specs=(spec,)))) + + spec = (DataKeySpec('token_targets', pad_at='forward') if operation == 'pad' else + DataKeySpec('token_targets', pad_at='prepare', split_at='forward')) + with pytest.raises(ValueError, match='Loss-only input.*forward boundary'): + register(spec) + # Model arguments can reach the forward hook, regardless of their stage name. + register(replace(spec, model_input=True)) + + +@pytest.mark.parametrize('pad_at,split_at', [(None, None), ('prepare', None), (None, 'forward'), + ('prepare', 'forward')]) +def test_parallel_switches_are_derived_from_boundaries(pad_at, split_at): + from dataclasses import fields + + spec = DataKeySpec('token_features', pad_at=pad_at, split_at=split_at, model_input=True) + assert spec.need_pad is (pad_at is not None) + assert spec.need_split is (split_at is not None) + assert not {'need_pad', 'need_split', 'need_gather', 'owner', 'roll'} & {field.name for field in fields(spec)} + + +@pytest.mark.parametrize('kwargs,message', [ + ({'pad_at': 'auxiliary'}, 'padding boundary'), + ({'split_at': 'output'}, 'split boundary'), + ({'gather_at': 'forward'}, 'gather boundary'), + ({'gather_at': 'loss'}, 'Only sharded input'), + ({'side': 'output', 'stage': 'output', 'pad_at': 'prepare'}, 'Output keys cannot'), +]) +def test_invalid_parallel_operation_boundaries_are_rejected(kwargs, message): + with pytest.raises(ValueError, match=message): + DataKeySpec('token_features', **kwargs) + + +@pytest.mark.parametrize('rank', [0, 1]) +@pytest.mark.parametrize('explicit_mask', [False, True]) +def test_single_sample_sdpa_builds_a_padded_global_attention_mask(rank, explicit_mask): + from types import SimpleNamespace + + ctx = _strategy(rank)._ctx + ctx.attn_implementation = 'sdpa' + ctx.tokenizer = SimpleNamespace(pad_token_id=0) + positions = torch.arange(5).unsqueeze(0) + values = {'input_ids': positions.clone(), 'position_ids': positions, + 'attention_mask': torch.ones(1, 5, dtype=torch.long) if explicit_mask else None} + specs = tuple(spec for spec in ctx.input_key_registry if spec.stage == 'forward') + keys = SPKeyCollection.from_mapping(values, specs) + ctx.pad_and_split_inputs(keys, boundary='forward', real_position_ids=positions) + torch.testing.assert_close(keys.get_value('attention_mask'), torch.tensor([[1, 1, 1, 1, 1, 0]])) + assert keys.get_value('input_ids').shape == (1, 3) + + +@pytest.mark.parametrize('configured,packed', [(False, False), (True, False), (False, True)]) +def test_processor_public_layout_query(configured, packed): + from twinkle.processor import InputProcessor + + positions = torch.tensor([[0, 1, 0, 1]]) if packed else torch.arange(4).unsqueeze(0) + batch = {'position_ids': positions} + processor = InputProcessor(padding_free=configured) + assert processor.is_padding_free(batch) is (configured or packed) + assert processor.is_padding_free([batch]) is (configured or packed) + assert batch['position_ids'] is positions + + +def test_processor_supported_fields_include_multimodal_replay_and_control_inputs(): + from twinkle.processor import InputProcessor + from twinkle.processor.fields import TRANSFORMERS_INPUT_KEYS + + required = {'pixel_values', 'pixel_values_videos', 'image_grid_thw', 'video_grid_thw', + 'input_features', 'input_features_mask', 'feature_attention_mask', 'grid_thws', + 'routed_experts', 'replay_mask', 'mm_token_type_ids', 'second_per_grid_ts', + 'cu_seq_lens_q', 'cu_seq_lens_k', 'cu_seqlens_q', 'cu_seqlens_kv', + 'cache_position', 'padding_free', 'packed_seq_params'} + assert required <= TRANSFORMERS_INPUT_KEYS + assert set(InputProcessor.padding_map) - TRANSFORMERS_INPUT_KEYS == {'length'} + value = torch.ones(1, 2) + sample = {key: value for key in required} + sample.update(length=2, unregistered=value) + output = InputProcessor.to_transformers_dict([sample])[0] + assert set(output) == required + assert all(output[key] is value for key in required) + + +def test_processor_defaults_do_not_schedule_parallel_operations(): + from twinkle.processor.fields import resolve_key_specs + + inputs, outputs = resolve_key_specs() + assert all(spec.pad_at is None and spec.split_at is None and spec.gather_at is None + for spec in (*inputs, *outputs)) + strategy = _strategy() + strategy.register_key_specs(inputs, outputs) + labels = strategy._ctx.input_key_registry.get('labels') + assert labels.pad_at == labels.split_at == 'prepare' + assert labels.gather_at == 'loss' + assert strategy._ctx.output_key_registry.get('logps').gather_at == 'loss' + + +def test_sp_builtin_execution_rules_preserve_common_layouts(): + from dataclasses import fields + from twinkle.processor.fields import COMMON_INPUT_SPECS, COMMON_OUTPUT_SPECS + from twinkle.model.transformers.strategy.sequence_parallel.data_keys import ( + BUILTIN_INPUT_KEYS, BUILTIN_OUTPUT_KEYS) + + for common, parallel in ((COMMON_INPUT_SPECS, BUILTIN_INPUT_KEYS), (COMMON_OUTPUT_SPECS, BUILTIN_OUTPUT_KEYS)): + assert [spec.name for spec in common] == [spec.name for spec in parallel] + for layout, execution in zip(common, parallel): + for item in fields(DataKeySpec): + if item.name not in ('pad_at', 'split_at', 'gather_at', 'autograd'): + assert getattr(layout, item.name) == getattr(execution, item.name) + + +def test_processor_padding_values_follow_common_layouts(): + from twinkle.processor.fields import COMMON_INPUT_SPECS, PADDING_VALUES + + for spec in COMMON_INPUT_SPECS: + # The processor/CP numeric fallback remains 0 without a tokenizer; + # only SP resolves the tokenizer sentinel to the model's pad ID. + expected = 0 if spec.pad_value == 'tokenizer' else spec.pad_value + assert PADDING_VALUES[spec.name] == expected + assert isinstance(PADDING_VALUES[spec.name], (int, float)) + assert PADDING_VALUES['mm_token_type_ids'] == 0 + assert PADDING_VALUES['length'] == -1 + assert PADDING_VALUES['pixel_values'] == 0.0 + assert PADDING_VALUES['routed_experts'] == PADDING_VALUES['replay_mask'] == 0 + + +def test_direct_builtin_output_declaration_keeps_model_schema_and_sp_registry_consistent(): + from twinkle.processor.fields import resolve_key_specs + + strategy = _strategy() + logits = strategy._ctx.output_key_registry.get('logits') + strategy.register_output_key(logits) + inputs, outputs = resolve_key_specs(strategy.field_declarations) + assert next(spec for spec in outputs if spec.name == 'logits') == logits + strategy.register_key_specs(inputs, outputs) + + +def _megatron_processor_cp_worker(rank, port): + import os + from datetime import timedelta + from types import SimpleNamespace + import torch.distributed as dist + from megatron.core import parallel_state as mpu + from twinkle.processor import InputProcessor + from twinkle.loss import CrossEntropyLoss + from twinkle.patch.megatron_emb import _output_embedding_hook + + os.environ['RANK'], os.environ['WORLD_SIZE'] = str(rank), '2' + torch.set_num_threads(1) + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, + world_size=2, timeout=timedelta(seconds=120)) + original = mpu.get_context_parallel_group + # Only resolve the CP group here; pad/split/gather and unpack are real code. + mpu.get_context_parallel_group = lambda: dist.group.WORLD + try: + mesh = SimpleNamespace(cp_world_size=2, cp_rank=rank, tp_world_size=1, sequence_parallel=False) + processor = InputProcessor(device_mesh=mesh, framework='megatron', padding_free=True) + processor.device_mesh = mesh + rows = [dict(input_ids=torch.arange(1, length + 1).unsqueeze(0), + position_ids=torch.arange(length).unsqueeze(0), + labels=torch.arange(1, length + 1).unsqueeze(0), + loss_scale=torch.ones(1, length), completion_mask=torch.ones(1, length, dtype=torch.long)) + for length in (3, 5)] + padded = processor.pad_cp(rows) + full = processor.collate_fn(padded)[0] + full = processor.add_extra_padding_free_args([full])[0] + cu = full['packed_seq_params'].cu_seqlens_q + torch.testing.assert_close(cu, torch.tensor([0, 4, 12], dtype=cu.dtype)) + local = processor.split_cp([dict(full)])[0] + reconstructed = dict(local) + for key in ('labels', 'position_ids', 'loss_scale', 'completion_mask'): + reconstructed[key] = processor.postprocess_tensor_cp(local[key], cu_seqlens=cu) + torch.testing.assert_close(reconstructed[key], full[key]) + # Match Megatron.forward_step_func: CP reconstruction followed by unpack. + full_logps = -full['labels'].clamp_min(0).float() / 10 + local_logps = -local['labels'].clamp_min(0).float() / 10 + gathered_logps = processor.postprocess_tensor_cp(local_logps, cu_seqlens=cu) + expected_inputs, expected_outputs = processor.unpack_packed_sequences(full, {'logps': full_logps}) + actual_inputs, actual_outputs = processor.unpack_packed_sequences(reconstructed, {'logps': gathered_logps}) + for key in ('labels', 'completion_mask', 'loss_scale'): + torch.testing.assert_close(actual_inputs[key], expected_inputs[key]) + torch.testing.assert_close(actual_inputs['labels'], + torch.tensor([[1, 2, 3, -100, -100, -100, -100, -100], + [1, 2, 3, 4, 5, -100, -100, -100]])) + loss = CrossEntropyLoss(reduction='sum') + torch.testing.assert_close(loss(actual_inputs, actual_outputs)['loss'], + loss(expected_inputs, expected_outputs)['loss']) + # Megatron's embedding hook remains independent of Transformers pooling. + full_features = torch.arange(24, dtype=torch.float32).reshape(1, 12, 2) + 1 + chunks = [] + for start, end in ((0, 4), (4, 12)): + pieces = full_features[:, start:end].chunk(4, dim=1) + chunks.extend((pieces[rank], pieces[3 - rank])) + local_features = torch.cat(chunks, dim=1) + pooled = _output_embedding_hook(SimpleNamespace(cp_group=dist.group.WORLD), (), + {'packed_seq_params': full['packed_seq_params']}, local_features) + expected = torch.nn.functional.normalize(full_features[0, torch.tensor([3, 11])], p=2, dim=1) + torch.testing.assert_close(pooled, expected) + finally: + mpu.get_context_parallel_group = original + dist.destroy_process_group() + + +@pytest.mark.slow +def test_megatron_cp_reconstruction_packed_loss_and_embedding_hook(): + import socket + import torch.multiprocessing as mp + + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + mp.spawn(_megatron_processor_cp_worker, args=(port,), nprocs=2, join=True) + + +def test_processor_filter_preserves_all_declared_vlm_fields(): + from twinkle.processor import InputProcessor + from twinkle.processor.fields import VLM_CONCAT_FIELDS + + value = torch.ones(1, 3, 4) + output = InputProcessor.to_transformers_dict([{key: value for key in VLM_CONCAT_FIELDS}])[0] + assert all(output[key] is value for key in VLM_CONCAT_FIELDS) + + +def test_model_forward_key_ownership_is_preserved_without_sp(): + from types import SimpleNamespace + from twinkle.processor.fields import resolve_key_specs + + with pytest.raises(ValueError, match='model forward hook'): + resolve_key_specs(SimpleNamespace(input_key_specs=(DataKeySpec('input_ids', stage='forward'),))) + + +@pytest.mark.parametrize('serializer_name', ['pickle', 'ray.cloudpickle']) +def test_field_specs_cache_serialization_rebuilds_without_retaining_owners(serializer_name): + from types import SimpleNamespace + from twinkle.model.transformers.transformers import TransformersModel + + serializer = pytest.importorskip(serializer_name) + wrapper = SimpleNamespace(model=torch.nn.Linear(2, 2), + strategy=SimpleNamespace(unwrap_model=torch.nn.Identity()), sp_strategy=None) + loss = None + expected = TransformersModel._get_field_specs(wrapper, loss) + restored = serializer.loads(serializer.dumps(wrapper)) + assert restored._field_specs_cache.specs is None + actual = TransformersModel._get_field_specs(restored, loss) + assert actual == expected + assert TransformersModel._get_field_specs(restored, loss) is actual + # Serializing leaves the sender's warm cache intact. + assert TransformersModel._get_field_specs(wrapper, loss) is expected diff --git a/tests/transformers/test_sp_loss_contract.py b/tests/transformers/test_sp_loss_contract.py new file mode 100644 index 000000000..a69a1e3ff --- /dev/null +++ b/tests/transformers/test_sp_loss_contract.py @@ -0,0 +1,410 @@ +"""Loss-declared output gathering for Sequence Parallel.""" +import pytest +import importlib +import os +import socket +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from datetime import timedelta +from types import SimpleNamespace + +from twinkle.data_format import LossOutput +from twinkle.loss import Loss, PPOValueLoss +from twinkle.model.transformers.strategy.sequence_parallel import DataKeySpec, SequenceParallelStrategy + + +def _strategy(): + strategy = SequenceParallelStrategy(sp_config={'ulysses_size': 2}) + ctx = strategy._ctx + ctx.world_size = ctx.sp_world_size = 2 + ctx.rp_world_size = 1 + ctx.extra_kwargs['position_ids'] = torch.arange(5).unsqueeze(0) + return strategy + + +def _mock_gather(monkeypatch): + from twinkle.model.transformers.strategy.sequence_parallel import GatherLoss + + monkeypatch.setattr(GatherLoss, 'apply', lambda ctx, value, labels, dim, positions: + (torch.cat((value, value), dim=dim), None)) + + +def test_loss_declared_logits_are_gathered_even_when_not_default(monkeypatch): + strategy = _strategy() + _mock_gather(monkeypatch) + local_values = torch.tensor([[[1.0], [2.0], [3.0]]]) + + inputs, outputs = strategy.gather_loss_tensors({}, {'logits': local_values}, sp_gather_keys=('logits',)) + + assert inputs == {} + torch.testing.assert_close(outputs['logits'], torch.tensor([[[1.0], [2.0], [3.0], [1.0], [2.0]]])) + + +def test_custom_loss_declares_and_consumes_a_registered_output_key(monkeypatch): + class ToyValueLoss(Loss): + require_logps = False + sp_gather_keys = ('values',) + + def __call__(self, inputs, outputs, **kwargs): + values = outputs['values'] + return LossOutput(loss=values.square().mean(), num_tokens=values.numel()) + + strategy = _strategy() + strategy.register_output_key( + DataKeySpec('values', split_dim=1, side='output', stage='output', autograd=True)) + _mock_gather(monkeypatch) + _, outputs = strategy.gather_loss_tensors({}, {'values': torch.tensor([[1.0, 2.0, 3.0]])}, + sp_gather_keys=ToyValueLoss.sp_gather_keys) + + assert outputs['values'].shape == (1, 5) + result = ToyValueLoss()({}, outputs) + torch.testing.assert_close(result['loss'], torch.tensor((1.0 + 4.0 + 9.0 + 1.0 + 4.0) / 5)) + + +def test_loss_can_declare_custom_sequence_inputs_and_outputs(monkeypatch): + class TokenAlignmentLoss(Loss): + require_logps = False + sp_input_key_specs = (DataKeySpec( + 'token_targets', + split_dim=1, + pad_value=-1, + gather_at='loss', + autograd=True, + pad_at='prepare', + split_at='prepare', + ),) + sp_output_key_specs = (DataKeySpec( + 'token_scores', + split_dim=1, + side='output', + stage='output', + gather_at='loss', + autograd=True, + ),) + + def __call__(self, inputs, outputs, **kwargs): + return LossOutput(loss=(outputs['token_scores'] - inputs['token_targets']).square().mean(), + num_tokens=inputs['token_targets'].numel()) + + strategy = _strategy() + strategy.register_loss_key_specs(TokenAlignmentLoss()) + _mock_gather(monkeypatch) + inputs, outputs = strategy.gather_loss_tensors( + {'token_targets': torch.tensor([[1.0, 2.0, 3.0]])}, + {'token_scores': torch.tensor([[1.0, 2.0, 3.0]])}) + + assert strategy._ctx.input_key_registry.get('token_targets') is not None + assert strategy._ctx.output_key_registry.get('token_scores') is not None + torch.testing.assert_close(inputs['token_targets'], torch.tensor([[1.0, 2.0, 3.0, 1.0, 2.0]])) + torch.testing.assert_close(outputs['token_scores'], torch.tensor([[1.0, 2.0, 3.0, 1.0, 2.0]])) + torch.testing.assert_close(TokenAlignmentLoss()(inputs, outputs)['loss'], torch.tensor(0.0)) + + +def test_declared_but_unregistered_or_missing_output_fails_loudly(): + strategy = _strategy() + with pytest.raises(KeyError, match='unregistered output key'): + strategy.gather_loss_tensors({}, {'values': torch.ones(1, 3)}, sp_gather_keys=('values',)) + with pytest.raises(KeyError, match='model output is missing'): + strategy.gather_loss_tensors({}, {}, sp_gather_keys=('logits',)) + + +def test_ppo_value_loss_declares_per_token_value_collection(): + assert PPOValueLoss.sp_gather_keys == ('values',) + assert PPOValueLoss.output_key_specs[0].name == 'values' + + +class _ToyValueModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(num_attention_heads=2, _attn_implementation='sdpa', model_type='toy_value') + self.scale = torch.nn.Parameter(torch.tensor(0.7)) + + def forward(self, input_ids, position_ids, **kwargs): + return {'logits': input_ids.to(self.scale.dtype) * self.scale} + + +def _ppo_value_worker(rank, port): + os.environ['RANK'] = str(rank) + os.environ['WORLD_SIZE'] = '2' + torch.set_num_threads(1) + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, + world_size=2, timeout=timedelta(seconds=90)) + try: + from twinkle import DeviceMesh + from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy + from twinkle.processor import InputProcessor + + ids = torch.tensor([[1, 2, 3, 4, 5, 6]]) + positions = torch.arange(6).unsqueeze(0) + labels = torch.tensor([[-100, -100, 3, 4, 5, 6]]) + returns = torch.tensor([[0.5, 1.0, 1.5, 2.0]]) + old_values = torch.tensor([[0.2, 0.4, 0.6, 0.8]]) + loss_fn = PPOValueLoss(cliprange_value=0.2, vf_coef=0.5) + + baseline = _ToyValueModel() + base_logits = baseline(ids, positions)['logits'] + base_loss = loss_fn({'labels': labels}, {'values': base_logits}, + returns=returns, old_values=old_values)['loss'] + base_loss.backward() + base_grad = baseline.scale.grad.clone() + + model = _ToyValueModel() + strategy = SequenceParallelStrategy( + DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu'), + sp_config={'ulysses_size': 2}, + model=model, + ) + strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0) + assert strategy.initialize() + strategy.register_loss_key_specs(loss_fn) + prepared = strategy.preprocess_inputs({'input_ids': ids, 'position_ids': positions, 'labels': labels}) + local_labels = prepared.pop('labels') + local_outputs = model(**{key: prepared[key] for key in ('input_ids', 'position_ids')}) + local_outputs['values'] = local_outputs.pop('logits') + from twinkle.model.transformers.task_outputs import postprocess_task_outputs + loss_inputs, loss_outputs = postprocess_task_outputs( + {'labels': local_labels}, + local_outputs, + sp_strategy=strategy, + sp_gather_keys=loss_fn.sp_gather_keys, + ) + parallel_loss = loss_fn(loss_inputs, loss_outputs, returns=returns, old_values=old_values)['loss'] + torch.testing.assert_close(parallel_loss, base_loss, atol=1e-7, rtol=1e-6) + parallel_loss.backward() + dist.all_reduce(model.scale.grad) + model.scale.grad.div_(2) + torch.testing.assert_close(model.scale.grad, base_grad, atol=1e-7, rtol=1e-6) + finally: + dist.destroy_process_group() + + +@pytest.mark.slow +def test_ppo_value_loss_sp_matches_full_sequence_loss_and_gradient(): + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + mp.spawn(_ppo_value_worker, args=(port,), nprocs=2, join=True) + + +class _NoopTrainStrategy: + """Keep the CPU E2E focused on TransformersModel and SP loss handling.""" + + def capture_pre_ep_state_if_needed(self, model, *, enable_ep): + return None + + def unwrap_model(self, model): + return model + + def wrap_model(self, model, *args): + return (model, *args) if args else model + + +def _ppo_training_pipeline_worker(rank, port): + os.environ['RANK'] = str(rank) + os.environ['WORLD_SIZE'] = '2' + torch.set_num_threads(1) + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, + world_size=2, timeout=timedelta(seconds=90)) + try: + from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM + from twinkle import DeviceMesh, Platform + from twinkle.model.transformers.transformers import TransformersModel + from twinkle.processor import InputProcessor + + Platform.get_local_device = staticmethod(lambda: torch.device('cpu')) + config = LlamaConfig(vocab_size=32, hidden_size=16, intermediate_size=32, + num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=4, + attention_dropout=0.0, pad_token_id=0, use_cache=False) + # Exercise a real Transformers attention implementation available in + # this environment; this E2E does not require flash-attn. + config._attn_implementation = 'sdpa' + torch.manual_seed(321) + baseline = LlamaForCausalLM(config).train() + + mesh = DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu') + model_module = importlib.import_module('twinkle.model.transformers.transformers') + original_template = model_module.Template + model_module.Template = lambda *_args, **_kwargs: SimpleNamespace(is_mm=False) + try: + training_model = TransformersModel( + model_cls=AutoModelForCausalLM, + config=config, + device_mesh=mesh, + mixed_precision='no', + ) + finally: + model_module.Template = original_template + training_model.strategy = _NoopTrainStrategy() + training_model.device_mesh = mesh + training_model._enable_sp = True + training_model.model.load_state_dict(baseline.state_dict()) + training_model.model.gradient_checkpointing_disable() + training_model.model.train() + + # The value head's [B, T] values travel through the standard `logits` + # output slot; retain a differentiable projection of real LM outputs. + model_forward = training_model.model.forward + + def forward_per_token_values(*args, **kwargs): + outputs = model_forward(*args, **kwargs) + outputs['logits'] = outputs.logits[..., 0] + return outputs + + training_model.model.forward = forward_per_token_values + training_model.sp_strategy = SequenceParallelStrategy(mesh, model=training_model.model) + training_model.sp_strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0) + training_model.set_loss(PPOValueLoss(cliprange_value=0.2, vf_coef=0.5)) + training_model.set_processor(InputProcessor) + training_model._ensure_sp_strategy() + + ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8], [2, 3, 4, 5, 6, 7, 8, 9]]) + positions = torch.arange(ids.shape[1]).expand_as(ids) + labels = torch.tensor([[-100, -100, -100, 4, 5, 6, 7, 8], + [-100, -100, -100, 5, 6, 7, 8, 9]]) + returns = torch.tensor([[0.5, 1.0, 1.5, 2.0, 2.5], [0.4, 0.8, 1.2, 1.6, 2.0]]) + old_values = torch.tensor([[0.2, 0.4, 0.6, 0.8, 1.0], [0.1, 0.3, 0.5, 0.7, 0.9]]) + batch = [ + {'input_ids': ids[i], 'position_ids': positions[i], 'labels': labels[i]} + for i in range(ids.shape[0]) + ] + + loss_fn = PPOValueLoss(cliprange_value=0.2, vf_coef=0.5) + baseline_values = baseline(input_ids=ids, position_ids=positions, use_cache=False).logits[..., 0] + baseline_loss = loss_fn({'labels': labels}, {'values': baseline_values}, + returns=returns, old_values=old_values)['loss'] + baseline_loss.backward() + baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=1e-3) + baseline_optimizer.step() + + training_model.forward_backward(inputs=batch, returns=returns, old_values=old_values) + optimizer_group = training_model.optimizer_group[training_model._get_default_group()] + assert optimizer_group.train_status.inputs['labels'].shape == labels.shape + assert optimizer_group.train_status.outputs['values'].shape == labels.shape + optimizer_group.optimizer = torch.optim.SGD(training_model.model.parameters(), lr=1e-3) + for parameter in training_model.model.parameters(): + assert parameter.grad is not None + dist.all_reduce(parameter.grad) + parameter.grad.div_(2) + training_model.step() + + for (name, actual), (_, expected) in zip(training_model.model.named_parameters(), baseline.named_parameters()): + torch.testing.assert_close(actual, expected, atol=3e-6, rtol=3e-5, msg=name) + finally: + dist.destroy_process_group() + + +@pytest.mark.slow +def test_ppo_value_loss_through_transformers_sdpa_sp_training_pipeline(): + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + mp.spawn(_ppo_training_pipeline_worker, args=(port,), nprocs=2, join=True) + + +@pytest.mark.parametrize('task', ['embedding', 'generative_reranker']) +def test_model_task_pooling_without_parallel_strategy(task): + from twinkle.model.transformers.task_outputs import postprocess_task_outputs + + features = torch.arange(28, dtype=torch.float32).reshape(2, 7, 2).requires_grad_() + # Left padding on row 0, right padding on row 1: use the rightmost valid index. + positions = torch.tensor([[-1, -1, 0, 1, 2, 3, 4], [0, 1, 2, -1, -1, -1, -1]]) + key = {'embedding': 'features', 'generative_reranker': 'logits'}[task] + _, outputs = postprocess_task_outputs({'position_ids': positions}, {key: features}, task=task) + expected = features[torch.arange(2), torch.tensor([6, 2])] + if task == 'embedding': + expected = torch.nn.functional.normalize(expected, dim=-1) + actual = outputs['embeddings' if task == 'embedding' else 'logits'] + torch.testing.assert_close(actual, expected) + actual.sum().backward() + assert features.grad is not None and torch.isfinite(features.grad).all() + assert 'sequence_scores' not in outputs + + +@pytest.mark.parametrize('task', ['seq_cls', 'reranker', 'embedding', 'generative_reranker']) +def test_sequence_labels_remain_global_in_model_sp_preparation(task): + from twinkle.model.transformers.transformers import TransformersModel + from twinkle.processor import InputProcessor + + model = SimpleNamespace(sp_strategy=_strategy()) + labels = torch.tensor([[0], [1]]) + positions = torch.arange(5).expand(2, -1) + inputs = {'input_ids': positions, 'position_ids': positions, 'labels': labels} + prepared = TransformersModel._prepare_sp_inputs(model, inputs, InputProcessor(), task=task) + assert prepared['labels'] is labels + assert 'padding_free' not in prepared + + +def _sp_pooling_worker(rank, port): + os.environ['RANK'], os.environ['WORLD_SIZE'] = str(rank), '2' + torch.set_num_threads(1) + from torch.utils.checkpoint import DefaultDeviceType + DefaultDeviceType.set_device_type('cpu') + dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank, + world_size=2, timeout=timedelta(seconds=120)) + try: + from twinkle.model.transformers.task_outputs import postprocess_task_outputs + + # Pooling values and averaged parameter gradients match a full sequence, + # for both ordinary batches and packed samples, with non-divisible lengths. + for topology, packed in (('ulysses', False), ('ulysses', True), ('ring', True)): + positions = (torch.tensor([[0, 1, 2, 0, 1, 2, 3]]) if packed else + torch.tensor([[0, 1, 2, -1, -1, -1, -1], [0, 1, 2, 3, 4, 5, 6]])) + values = torch.arange(positions.numel() * 2, dtype=torch.float32).reshape(*positions.shape, 2) + 1 + for task in ('embedding', 'generative_reranker'): + key = {'embedding': 'features', 'generative_reranker': 'logits'}[task] + strategy = _strategy() + strategy._ctx._sp_rank = rank + strategy._ctx._sp_group = dist.group.WORLD + strategy._ctx.extra_kwargs['position_ids'] = positions + if topology == 'ring': + strategy._ctx.sp_world_size = 1 + strategy._ctx.rp_world_size = 2 + strategy._ctx._sp_rank = 0 + strategy._ctx._rp_rank = rank + strategy._ctx._sp_group = None + strategy._ctx._rp_group = dist.group.WORLD + base_scale = torch.tensor([0.7, 1.2], requires_grad=True) + scale = base_scale.detach().clone().requires_grad_() + _, base_outputs = postprocess_task_outputs({'position_ids': positions}, {key: values * base_scale}, task=task) + padded_positions = strategy.pad(positions, padding_value=-1, position_ids=positions) + local = strategy.pad(values, padding_value=0, dim=1, position_ids=positions) + local = strategy.split(local, dim=1, position_ids=padded_positions) + _, outputs = postprocess_task_outputs({'position_ids': positions}, {key: local * scale}, + task=task, sp_strategy=strategy) + name = 'embeddings' if task == 'embedding' else 'logits' + torch.testing.assert_close(outputs[name], base_outputs[name]) + outputs[name].square().sum().backward() + base_outputs[name].square().sum().backward() + dist.all_reduce(scale.grad) + torch.testing.assert_close(scale.grad / 2, base_scale.grad, atol=2e-5, rtol=2e-5) + + finally: + dist.destroy_process_group() + + +@pytest.mark.slow +def test_sp_pooling_and_gradients_match_full_sequence(): + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + mp.spawn(_sp_pooling_worker, args=(port,), nprocs=2, join=True) + + +@pytest.mark.parametrize('task', ['seq_cls', 'reranker']) +def test_prepooled_classification_logits_skip_token_pooling(task): + from twinkle.model.transformers.task_outputs import postprocess_task_outputs + + def unexpected_pool(*args): + raise AssertionError('Already pooled classification logits must bypass token pooling') + + inputs = {'attention_mask': torch.ones(2, 5)} + outputs = {'logits': torch.tensor([[1., 2.], [3., 4.]])} + strategy = SimpleNamespace(pool_last_valid_token=unexpected_pool) + actual_inputs, actual_outputs = postprocess_task_outputs(inputs, outputs, task=task, sp_strategy=strategy) + assert actual_inputs is inputs + assert actual_outputs is outputs diff --git a/tests/transformers/test_sp_model_adapters.py b/tests/transformers/test_sp_model_adapters.py new file mode 100644 index 000000000..05dd3d82b --- /dev/null +++ b/tests/transformers/test_sp_model_adapters.py @@ -0,0 +1,100 @@ +"""Tests for model-family SP adapter registration and behavior.""" +import pytest +import torch +from types import SimpleNamespace + +import twinkle.model.transformers.strategy.sequence_parallel as sp +from twinkle.model.transformers.strategy.sequence_parallel import adapters as sp_adapters + + +@pytest.mark.parametrize(('model_type', 'expected'), [ + ('qwen3_vl', {'qwen3_vl'}), + ('qwen3_5', {'qwen3_5_linear_attention'}), + ('qwen3_5_moe', {'qwen3_5_linear_attention', 'moe_aux'}), + ('llama', set()), +]) +def test_builtin_model_adapters_are_selected_by_model_family(model_type, expected): + model = torch.nn.Linear(2, 2) + model.config = SimpleNamespace(model_type=model_type) + assert {adapter.name for adapter in sp_adapters.adapters_for_model(model)} == expected + + +def test_model_adapter_declared_key_specs(monkeypatch): + class CustomAdapter(sp_adapters.ModelSPAdapter): + name = 'test_custom_sp_adapter' + + def matches(self, model): + return getattr(getattr(model, 'config', None), 'model_type', None) == 'custom_sp_test' + + def input_key_specs(self): + return (sp.DataKeySpec('custom_features', pad_value=-3),) + + adapter = CustomAdapter() + monkeypatch.setattr(sp_adapters, '_ADAPTERS', [*sp_adapters._ADAPTERS, adapter]) + model = torch.nn.Linear(2, 2) + model.config = SimpleNamespace(model_type='custom_sp_test') + context = sp.SequenceParallel() + context._configure_model_adapters(model, model) + assert context.model_adapters == (adapter,) + assert context.input_key_registry.get('custom_features') == sp.DataKeySpec('custom_features', pad_value=-3) + + +def test_model_adapters_declare_mrope_and_keep_qwen35_ring_limit(monkeypatch): + vl_adapter = sp_adapters.Qwen3VLSPAdapter() + assert vl_adapter.position_ids_mode == 'mrope' + ctx = sp.SequenceParallel() + ctx.position_ids_mode = vl_adapter.position_ids_mode + real_positions = torch.arange(4).unsqueeze(0) + assert ctx._extract_real_position_ids(real_positions) is real_positions + mrope_positions = real_positions.unsqueeze(0).expand(3, -1, -1) + torch.testing.assert_close(ctx._extract_real_position_ids(mrope_positions), real_positions) + ctx.rp_world_size = 2 + with pytest.raises(NotImplementedError, match='rp_world_size > 1'): + sp_adapters.Qwen35LinearAttentionSPAdapter().configure(ctx, torch.nn.Linear(2, 2)) + + +@pytest.mark.parametrize(('adapter', 'patch_name'), [ + (sp_adapters.Qwen3VLSPAdapter(), 'qwen3_vl_deepstack'), + (sp_adapters.Qwen35LinearAttentionSPAdapter(), 'qwen35_gdn'), +]) +def test_builtin_adapters_install_model_specific_patches(monkeypatch, adapter, patch_name): + calls = [] + monkeypatch.setattr(sp_adapters, 'apply_patch', lambda model, patch, **kwargs: calls.append( + (model, patch, kwargs['sequence_parallel']))) + model = torch.nn.Linear(2, 2) + ctx = sp.SequenceParallel() + ctx.rp_world_size = 1 + adapter.configure(ctx, model) + assert len(calls) == 1 + assert calls[0][0] is (None if patch_name == 'qwen35_gdn' else model) + assert calls[0][2] is ctx + + +@pytest.mark.parametrize('mask_kind', ['omitted', 'none', 'explicit']) +def test_moe_adapter_registers_router_logits_gather_hook(monkeypatch, mask_kind): + hooks = [] + + class MoEConfig: + model_type = 'example_moe' + + class Model: + config = MoEConfig() + + def register_forward_hook(self, hook, with_kwargs=False): + hooks.append((hook, with_kwargs)) + + class Output(dict): + router_logits = (torch.arange(6, dtype=torch.float32).reshape(2, 3),) + + ctx = SimpleNamespace(real_position_ids=torch.tensor([[0, 1, 2]])) + monkeypatch.setattr(sp_adapters.GatherLoss, 'apply', + lambda context, value, labels, dim, positions: (torch.cat((value, value), dim), None)) + model = Model() + adapter = sp_adapters.MoEAuxSPAdapter() + assert adapter.matches(model) + adapter.configure(ctx, model) + assert hooks and hooks[0][1] + kwargs = {} if mask_kind == 'omitted' else {'attention_mask': None if mask_kind == 'none' else torch.ones(1, 3)} + output = hooks[0][0](model, (), kwargs, Output()) + assert len(output['router_logits']) == 1 + assert output['router_logits'][0].shape == (3, 3) diff --git a/tests/twinkle_agentic/test_multi_turn_rollout.py b/tests/twinkle_agentic/test_multi_turn_rollout.py index 0c06d0271..f18b914f3 100644 --- a/tests/twinkle_agentic/test_multi_turn_rollout.py +++ b/tests/twinkle_agentic/test_multi_turn_rollout.py @@ -1481,3 +1481,10 @@ def followup(traj, n_before): assert '<|im_end|>' not in last['content'] # And it was not dispatched: a dispatch appends a tool message. assert [m['role'] for m in out['messages'] if m['role'] == 'tool'] == [] + + +def test_vision_feature_keys_follow_processor_field_declarations(): + from twinkle.processor.fields import VLM_CONCAT_FIELDS + from twinkle_agentic.rollout.multi_turn import _vision_feature_keys + + assert _vision_feature_keys() == frozenset(VLM_CONCAT_FIELDS) | {'mm_token_type_ids', 'second_per_grid_ts'}