Ulysses 模块化 by werwrewe · Pull Request #321 · modelscope/twinkle

werwrewe · 2026-10-16T08:20:09Z

PR type

Bug Fix
New Feature
Document Updates
More Models or Datasets Support

PR information

改动内容

Phase 1：隔离每个模型的 SP 上下文

每个 SequenceParallelStrategy 持有自己的 SequenceParallel 上下文，保存该模型的通信组、rank 和当前 batch 状态。
Transformers 的 attention/mask 全局包装器安装一次，在调用时按 attention 模块或 forward 上下文栈找到当前模型。
GatherLoss 显式保存前向所属的上下文，反向传播不再依赖可能已被其他模型覆盖的全局单例。
router replay 和 Qwen3.5 线性注意力路径使用模型对应的上下文。
forward 输入切分 hook 在自身失败时清理上下文栈，避免错误状态影响后续模型调用。

Phase 2：声明 SP 数据键规则

新增 DataKeySpec、输入键注册表和输出键注册表，描述键的序列维度、padding 值、是否补齐/切分及各自执行边界、是否 gather，以及 gather 是否保留梯度。常用键与模型/任务扩展键统一放入 SPKeyCollection，由 pad_and_split_inputs 按声明执行。
内置键（如 labels、completion_mask、loss_scale、logps、entropies）使用统一规则处理。
处理器按注册声明完成辅助输入的拼批、padding、切分、恢复和必要的 loss 前 gather；模型前向专属键由 forward hook 处理。
real_position_ids 与 cache_position 作为位置/模型调用上下文单独传入，不混入普通键集合；completion_mask、loss_scale 默认保持完整序列，避免未经声明就参与 SP 分片。
自定义键默认不 gather，调用方可以显式声明其行为；不兼容的重复声明会报错。

Phase 3：将模型专属行为移入适配器

新增 ModelSPAdapter 合约和模型适配器注册机制，由适配器负责模型家族识别、数据键声明及模型专属 hook/patch 配置。
将 Qwen3-VL DeepStack/mRoPE、Qwen3.5 GatedDeltaNet 线性注意力、MoE router auxiliary loss 的特殊处理迁入内置适配器。
允许同一模型组合多个适配器，也允许后续模型注册自己的适配器，不必把模型判断继续堆进 SP 主类。

Phase 4：由 loss 声明需要的序列数据

新增 loss 侧的 sp_gather_keys、sp_input_key_specs 和 sp_output_key_specs 声明，让 loss 指定计算前需要收集或对齐的序列张量。
模型初始化和运行时更换 loss 时都会登记声明；训练、评估和 processor 路径会把声明传递给 SP 收集逻辑。
已注册但默认不 gather 的输出，也可以由 loss 声明后在计算前收集；未知声明键或模型未产出所需输出时明确报错。
PPOValueLoss 声明收集逐 token value（存于 outputs['logits']）；新增自定义 token loss 用例验证输入/输出键声明和 optimizer 更新。
修复 SafeLossWrapper 丢失自定义 loss SP 声明的问题，确保经过安全包装后声明仍到达训练入口。
该契约处理序列张量的 padding、切分和 gather，不会自动推断任意 loss 的数学语义，也不会自动实现跨样本/跨 DP 统计或自定义通信归约。

Experiment results

测试环境

硬件/软件：Ascend 910B4；CANN 9.1.0、驱动 25.5.1、Python 3.11.15、PyTorch 2.10.0+cpu / torch_npu 2.10.0、Transformers 5.13.1

训练loss的回归测试

dpo的sp>1分支打通测试，以及异构模型持有不同sp状态的测试。

带补丁的模型回归测试（改动前 vs 改动后）：

…v-swift-v5' into resolve-pr-dev-swift-v5 # Conflicts: # src/twinkle/loss/base.py # src/twinkle/loss/value.py # src/twinkle/model/transformers/moe/router_replay.py # src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py # src/twinkle/model/transformers/transformers.py # src/twinkle/processor/base.py

meichangsu1 · 2026-10-20T01:33:30Z

            pass

-        if hasattr(base_model, 'language_model'):
-            text_model = base_model.language_model


这里为啥要删掉

这块是dpo训练时出现的问题，这个判断无法区分同架构的不同模型实例（虽然 text_model 取的是 policy，但列表里存的是各层的类型）。比如 policy 和 reference 的 attention 类型相同，就可能被当成同一个模型处理。

现在改为在每个模型的 prepare() 中选择文本侧模块树，并通过 _register_module_tree(text_model, self) 注册模块实例与 SP context 的对应关系；wrapper 调用时通过 _MODULE_TO_CTX.get(module) 找到所属 context。因此 _prepare_flash_attn() 中不再需要这段局部变量和模型判断。

meichangsu1 · 2026-10-20T01:34:03Z

            pass

-        if hasattr(base_model, 'language_model'):
-            text_model = base_model.language_model


这里为啥要删掉

meichangsu1 · 2026-10-20T02:03:19Z

+    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


这个条件的含义是什么

就是取出pad value，由于注册的时候DataKeySpec('input_ids', split_dim=1, pad_value='tokenizer', stage='forward'),input_ids的pad值为tokenizer.pad_token_id，导入文件时没有模型实例，先用文本占位，在 _pad_key这个函数进行取出。对于别的键值都是取注册时候的pad value。

主要是不知道 spec.pad_value == 'tokenizer' 为什么需要有这个条件

是给input_ids的pad value开的特例判断，原因是input_ids的pad value需要根据实例化tokenizer的tokenizer.pad_token_id去取，不能在注册键值的时候固定一个pad value，因此在注册input_ids时后给的pad value是用字符串'tokenizer'占位，在实际使用到pad value的时候根据当前实例化的tokenizer.pad_token_id去取值。

meichangsu1 · 2026-10-20T02:03:56Z

-        tokenizer = self.tokenizer
+        if not isinstance(keys, SPKeyCollection):
+            raise TypeError('pad_and_split_inputs expects an SPKeyCollection')
+        if boundary not in ('prepare', 'forward', 'model_internal', 'router_replay'):


这三个boundary分别是什么含义

这块原子化改动主要围绕pad and split函数，这四个boundary对应pad and split函数调用的四个时机，每个键值都持有一个boundary代表它被pad and split的位置。

prepare：SequenceParallel.prepare_inputs()
forward：SequenceParallel._prepare_forward_hook()
model_internal：SequenceParallel.pad_and_split_mm_tokens()
router_replay：moe/router_replay.py::set_router_replay_data()

meichangsu1 · 2026-10-20T02:08:12Z

        if micro_batch_size is None:
            # normal collate
-            outputs = self._collate_macro_batch(inputs)
+            outputs = self._collate_macro_batch(inputs, sp_strategy=kwargs.get('sp_strategy'))


请评估一下sp_strategy能否传入processor,也许这一直是一个错误的做法

已经改正追加commit，从 strategy 读取键规则，改为直接接收 input_key_specs / output_key_specs。改正仅限本次提交的代码，不包括以前的情形。

以前的也可以评估下是否不传入

然后就是假设ppo所有模型要开启sp，要怎么写cookbook

已经改正追加commit，把processor里面所有SP相关准备移动到各自的位置去。

可以基于现有 cookbook/rl/ppo/ppo.py 扩展，，为每个模型分别创建 DeviceMesh，并设置 ulysses_size。

Move SP input preparation, loss tensor gathering and last-token pooling into the model/strategy layer. Keep processors responsible for batch layouts and preserve sentence-level labels during SP preprocessing. Cover pooling outputs and gradients for padded, packed and ring layouts with CPU/Gloo regression tests.

meichangsu1 · 2026-10-22T02:32:12Z

        inputs: Dict[str, Any] = processor(
            inputs,
-            sp_strategy=self.sp_strategy,
+            input_key_specs=input_key_specs,


sp 中有哪些额外字段是需要传入processor中特殊处理的呢，是不是所有字段本身就已经被processor处理过了

meichangsu1 · 2026-10-22T02:53:38Z

        labels: torch.Tensor = inputs.pop('labels', None)
        loss_scale = inputs.pop('loss_scale', None)
        channel = inputs.pop('channel', None)
+        sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, input_key_specs=input_key_specs)


pop_sp_auxiliary_inputs 的作用是什么，sp 有什么字段是需要特殊弹出的吗
？不开sp时，没有这些字段吗

clc added 4 commits October 13, 2026 16:16

sp: isolate sequence parallel context per model

3257307

sp: add data-key contracts and model adapters

7bca22b

sp: add loss-declared gather contracts

2f7e78d

meichangsu1 reviewed Oct 20, 2026

View reviewed changes

clc added 2 commits October 20, 2026 21:17

sp: decouple processor key rules from strategy

9ad5b7a

werwrewe force-pushed the dev-swift-v5 branch from eb404b8 to f787f85 Compare October 21, 2026 13:02

meichangsu1 reviewed Oct 22, 2026

View reviewed changes

Conversation

werwrewe commented Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR type

PR information

改动内容

Phase 1：隔离每个模型的 SP 上下文

Phase 2：声明 SP 数据键规则

Phase 3：将模型专属行为移入适配器

Phase 4：由 loss 声明需要的序列数据

Experiment results

测试环境

训练loss的回归测试

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

werwrewe Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

werwrewe commented Oct 16, 2026 •

edited

Loading

werwrewe Oct 21, 2026 •

edited

Loading