训练中间件 | Twinkle

DeviceMesh/DeviceGroup

Mon, 01 Jan 0001 00:00:00 +0000

这两个类用于表达硬件资源分配和网络拓扑，Twinkle 的数据分发和收集也依赖它们。

DeviceGroup

@dataclass
class DeviceGroup:
 name: str
 ranks: Union[List[int], int]
 device_type: str
 visible_devices: Optional[str] = None # Optional: explicitly set visible devices (e.g., "8,9")
 gpus_per_worker: int = 1

name: 资源组名
ranks: 占用硬件列表，如果是CPU资源仅支持int类型
device_type: 硬件类型，例如 GPU/CPU/NPU 等
visible_devices: 可见资源列表，用于希望仅使用部分 rank 的硬件的情况
gpus_per_worker: 每个 worker 占用多少硬件

如果训练 RL，开发者可以构造多个这样的组，并将对应的模型、采样器分配进入其中。

DeviceMesh

DeviceMesh 承载了组件拓扑、分布式并行信息，这个类会在组件内传递，用于数据分发和数据收集。

@dataclass
class DeviceMesh:
 ...

 @staticmethod
 def from_sizes(*, world_size: int = 1, dp_size: int = 1, fsdp_size: int = None, tp_size: int = None,
 pp_size: int = None, ulysses_size: int = None, cp_size: int = None, ep_size: int = None,
 etp_size: int = None,vpp_size: int = None, device_type: str = 'cuda', sequence_parallel: bool = False) -> "DeviceMesh":
 ...

推荐使用 from_sizes 来构造它。

参数参考

参数	说明	默认值
`world_size`	总进程数	1
`dp_size`	数据并行度	1
`fsdp_size`	全分片数据并行度	None
`tp_size`	张量并行度	None
`pp_size`	流水线并行度	None
`ulysses_size`	Ulysses 序列并行度	None
`cp_size`	上下文并行度	None
`ep_size`	专家并行度（MoE 模型）	None
`etp_size`	专家张量并行度	None
`ep_fsdp_size`	每个 EP 组内的 FSDP 度	None
`vpp_size`	虚拟流水线并行度	None
`device_type`	设备类型（`cuda`、`npu` 等）	`cuda`
`sequence_parallel`	启用 Megatron 风格序列并行	False

我们举一个例子：

sampler_device_mesh = DeviceMesh.from_sizes(dp_size=4)
actor_device_mesh = DeviceMesh.from_sizes(dp_size=2, pp_size=2, tp_size=2)

dataloader = DataLoader(...)
sampler = vLLMSampler(..., device_mesh=sampler_device_mesh, remote_group=...)
actor = MegatronModel(..., device_mesh=actor_device_mesh, remote_group=...)

for data in dataloader:
 sampler_output = sampler.sample(data)
 input_data = [seq.new_input_feature for response in sampler_output for seq in response.sequences]
 ...
 model_output = actor.forward(input_data)

我们以上面的伪代码来分析数据传递情况。

dataloader 取出数据 -> 按照 dp_size=4 分发给 sampler -> 按照 dp_size=4 收集数据 -> 按照 dp_size=2 分发给模型 -> 按照 dp_size=2 收集输出

通过 DeviceMesh，可以将数据流平顺地在各个 group 和组件之间流转起来。

数据的分发判断由 DeviceMesh 的 get_slice 方法执行：

batch[device_mesh.get_slice(len(batch))]

get_slice 会根据当前 rank，计算出当前 worker 属于哪个 dp 组，并获取对应的数据。该过程发生在 DataLoader 的 DeviceMeshSampler 中，同样发生在 remote_class 的 dispatch 和 collect 中。

DeviceMesh/DeviceGroup

这两个类用于表达硬件资源分配和网络拓扑，Twinkle 的数据分发和收集也依赖它们。

DeviceGroup

@dataclass
class DeviceGroup:
 name: str
 ranks: Union[List[int], int]
 device_type: str
 visible_devices: Optional[str] = None # Optional: explicitly set visible devices (e.g., "8,9")
 gpus_per_worker: int = 1

name: 资源组名
ranks: 占用硬件列表，如果是CPU资源仅支持int类型
device_type: 硬件类型，例如 GPU/CPU/NPU 等
visible_devices: 可见资源列表，用于希望仅使用部分 rank 的硬件的情况
gpus_per_worker: 每个 worker 占用多少硬件

如果训练 RL，开发者可以构造多个这样的组，并将对应的模型、采样器分配进入其中。

DeviceMesh

DeviceMesh 承载了组件拓扑、分布式并行信息，这个类会在组件内传递，用于数据分发和数据收集。

@dataclass
class DeviceMesh:
 ...

 @staticmethod
 def from_sizes(*, world_size: int = 1, dp_size: int = 1, fsdp_size: int = None, tp_size: int = None,
 pp_size: int = None, ulysses_size: int = None, cp_size: int = None, ep_size: int = None,
 etp_size: int = None,vpp_size: int = None, device_type: str = 'cuda', sequence_parallel: bool = False) -> "DeviceMesh":
 ...

推荐使用 from_sizes 来构造它。

我们举一个例子：

sampler_device_mesh = DeviceMesh.from_sizes(dp_size=4)
actor_device_mesh = DeviceMesh.from_sizes(dp_size=2, pp_size=2, tp_size=2)

dataloader = DataLoader(...)
sampler = vLLMSampler(..., device_mesh=sampler_device_mesh, remote_group=...)
actor = MegatronModel(..., device_mesh=actor_device_mesh, remote_group=...)

for data in dataloader:
 sampler_output = sampler.sample(data)
 input_data = [seq.new_input_feature for response in sampler_output for seq in response.sequences]
 ...
 model_output = actor.forward(input_data)

我们以上面的伪代码来分析数据传递情况。

dataloader 取出数据 -> 按照 dp_size=4 分发给 sampler -> 按照 dp_size=4 收集数据 -> 按照 dp_size=2 分发给模型 -> 按照 dp_size=2 收集输出

通过 DeviceMesh，可以将数据流平顺地在各个 group 和组件之间流转起来。

数据的分发判断由 DeviceMesh 的 get_slice 方法执行：

batch[device_mesh.get_slice(len(batch))]

专家并行 (EP)

Mon, 01 Jan 0001 00:00:00 +0000

专家并行将混合专家模型（MoE）的专家分布到多个 GPU 上，每个 rank 只持有部分专家。这降低了单卡显存占用，使大规模 MoE 模型的训练成为可能。

概览

概念	说明
ExpertParallelConfig	控制 EP 行为的配置数据类
apply_expert_parallel()	入口函数，负责分片专家并替换前向传播
shard_experts()	将专家均匀分配到各 EP rank
patch_forward()	将 MoE block 的 forward 替换为带 all-to-all 通信的 EP 版本

配置

from twinkle.model.transformers.moe.expert_parallel import ExpertParallelConfig

config = ExpertParallelConfig(
 enabled=True, # 启用专家并行
 router_dtype='fp32', # 路由计算精度：'fp32', 'bf16', 'fp16'
 keep_router_logits=True, # 在输出中保留路由 logits
 ignore_shared_experts=False,# 跳过共享专家计算（如 DeepSeek）
 ep_size=None, # EP 并行度（由 TransformersModel 使用）
)

配合 DeviceMesh 使用

在 DeviceMesh.from_sizes() 中设置 ep_size 即可激活 EP。框架会在模型初始化时自动调用 apply_expert_parallel()。

from twinkle.utils import DeviceMesh

# 8 卡：2 路 EP × 4 路数据并行
device_mesh = DeviceMesh.from_sizes(
 world_size=8,
 dp_size=4,
 ep_size=2,
)

EP + FSDP 组合分片：

# 8 卡：2 路 EP，每个 EP 组内 2 路 FSDP
device_mesh = DeviceMesh.from_sizes(
 world_size=8,
 dp_size=2,
 ep_size=2,
 ep_fsdp_size=2,
)

通信模式

EP 前向传播遵循 4 阶段流水线：

预处理 — 计算每个专家的 token 数量和分割大小
Token Pre-All2All — 按专家分配排列 token，然后在 EP rank 间执行 all-to-all 交换
专家计算 — 每个 rank 在接收到的 token 上运行本地专家
Token Post-All2All — all-to-all 交换结果，反排列并应用路由权重

输入 token → 路由器 → [预处理] → [pre_all2all] → [本地专家] → [post_all2all] → 输出

要求

num_experts 必须能被 ep_size 整除
torch.distributed 必须已初始化
MoE block 必须定义 gate/router 模块和 experts（支持 nn.ModuleList 或张量形式的 gate_up_proj/down_proj）
共享专家（如 DeepSeek MoE）会自动处理，除非设置 ignore_shared_experts=True

序列并行 (SP)

Mon, 01 Jan 0001 00:00:00 +0000

序列并行沿序列维度将长序列分割到多个 GPU 上，使训练能处理超出单卡显存的序列长度。Twinkle 实现了 Ulysses 风格的序列并行，并可选地支持派生环形注意力。

概览

概念	说明
SequenceParallelConfig	SP 配置数据类
SequenceParallelStrategy	封装 SP 生命周期的策略类
SequenceParallel	核心实现，处理填充/分割/聚合

配置

from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelConfig

config = SequenceParallelConfig(
 enabled=True, # 启用序列并行
 ulysses_size=None, # Ulysses SP 并行度（若为 None 则从 DeviceMesh 自动推导）
 gather_logits=True, # 前向后聚合 logits 用于损失计算
)

配合 DeviceMesh 使用

在 DeviceMesh.from_sizes() 中设置 ulysses_size 即可激活 SP：

from twinkle.utils import DeviceMesh

# 8 卡：4 路 Ulysses SP × 2 路数据并行
device_mesh = DeviceMesh.from_sizes(
 world_size=8,
 dp_size=2,
 ulysses_size=4,
)

工作原理

填充 — 输入序列被填充到可被 SP 并行度整除的长度
分割 — 填充后的输入沿序列维度均匀分配到各 SP rank
分布式注意力 — FlashAttention2 被 patch 为在注意力计算前后执行 Ulysses all-to-all 通信
聚合 — 前向传播后，logits 被聚合回完整序列长度用于损失计算

支持的注意力后端

后端	状态
FlashAttention2	完全支持（包括打包/padding-free 序列）
SDPA	支持（仅非打包批次）
派生环形注意力	仅支持 FlashAttention2（`rp_world_size > 1`）

Qwen3.5 线性注意力

SP 自动检测 Qwen3.5 GatedDeltaNet 线性注意力层，并应用 Qwen3_5GatedDeltaNetUlyssesPatch，确保混合注意力架构下序列并行的正确性。

MoE 辅助损失

对于 MoE 模型，SP 自动安装前向 hook，在计算辅助损失前跨 SP rank 聚合路由 logits，确保负载均衡信号的正确性。

关键约束

num_key_value_heads 必须能被 ulysses_size 整除（Ulysses 模式），否则回退到环形注意力
打包/padding-free 批次需要 FlashAttention2
派生环形注意力要求 batch_size == 1（打包格式）
torch.distributed 必须已初始化

Padding-Free 训练

Mon, 01 Jan 0001 00:00:00 +0000

Padding-free（也称为"打包"）训练通过将多个序列拼接到一个打包批次中，消除了对 padding token 的无效计算。Twinkle 支持标准注意力和 Qwen3.5 GatedDeltaNet 线性注意力的 padding-free 训练。

工作原理

不同于将所有序列填充到 max_length，padding-free 将多个序列打包到一行中，并使用 position_ids 跟踪序列边界，从而避免在 padding token 上浪费算力。

标准方式: [tok tok tok PAD PAD PAD] [tok tok PAD PAD PAD PAD]
打包方式: [tok tok tok tok tok ...] ← 无 padding 浪费

使用方式

通过 PackingDataset 或 IterablePackingDataset 启用：

from twinkle.dataset import PackingDataset

dataset = PackingDataset(
 dataset=base_dataset,
 max_length=8192,
)

数据集会自动打包序列并生成正确的 position_ids，在序列边界处重置。

GatedDeltaNet 补丁（Qwen3.5）

Qwen3.5 使用混合架构，融合了标准注意力和 GatedDeltaNet 线性注意力。原生 GatedDeltaNet 实现不会在打包序列边界处重置线性注意力状态。

GatedDeltaNetPaddingFreePatch 通过以下方式修复：

Patch Qwen3_5DecoderLayer.forward，将 cu_seq_lens_q（累积序列长度）传递给线性注意力层
Patch Qwen3_5GatedDeltaNet.forward，使用支持 cu_seqlens 的 flash-linear-attention 内核（causal_conv1d、chunk_gated_delta_rule）

在 Qwen3.5 模型上检测到 padding-free 时，补丁会自动应用。

要求

需安装 flash-linear-attention 包
仅适用于含 GatedDeltaNet 层的 Qwen3.5 模型
启用序列并行时，会使用 Qwen3_5GatedDeltaNetUlyssesPatch 替代

注意力后端要求

注意力后端	Padding-Free 支持
FlashAttention2	完全支持
SDPA	支持（不兼容序列并行）
Eager	不支持

RemoteClass

Mon, 01 Jan 0001 00:00:00 +0000

所有 Twinkle 中支持 Ray 和 HTTP 中使用的组件均通过 @remote_class 和 @remote_function 进行了装饰。该装饰器会拦截类的构造，在 Ray 模式下，将类的构造转为 worker 执行。

from twinkle import remote_class, remote_function

@remote_class(execute='first')
class MyComponent:

 def __init__(self, **kwargs):
 ...

 @remote_function(dispatch='slice_dp', collect='first')
 def func(self, *args, **kwargs):
 ...
 return ...

开发者只需要编写上述代码，就可以将 MyComponent 类转入 worker 执行。其中：

remote_class: 将类标记为需要远端执行。如果 Twinkle 初始化设置为 local 模式，或者该类构造时没有传入 remote_group 设置，或者 remote_group 为当前 worker，都会在进程内构造该类。
remote_function: 将某个标记了 remote_class 的方法标记为可以在 Ray 中执行。其输入和输出均会被 Ray 压缩传递。

调用 MyComponent：

import twinkle
from twinkle import DeviceGroup

device_groups = [
 DeviceGroup(
 name='default',
 ranks=4,
 device_type='cuda',
 )
]

twinkle.initialize('ray', groups=device_groups)

_my_component = MyComponent(remote_group='default')
_my_component.func(...)

通过这种方式，我们编写了一个 MyComponent，并在 Ray 集群中使用 4 张卡构造了一个叫 default 的组，把 MyComponent 构造在了该组中。

remote_class 在装饰类的时候的参数：

execute: 支持 first/all。first 仅会在该组的第 0 个设备上创建，一般用于 Dataset、DataLoader 的构造，all 会在所有设备上构造。

remote_function 在装饰方法的时候有下面的参数：

dispatch: 如何分发输入数据。支持 slice/all/slice_dp/函数四种。slice 会将 list 输入均匀分发（非 list 会全部分发），all 进行全部分发，slice_dp 会将输入数据按照 device_mesh 的 dp 组进行切分分发，来保障模型输入数据的正确性，函数方式支持以自己的实现来分发输入数据：

def _dispatcher(length, i, args, kwargs, device_mesh):
 # length 是 worker 数量，i 是当前 rank，args 和 kwargs 是输入数据，在这里具体执行分发逻辑
 # device_mesh是隶属于目标组件的device_mesh
 return _args_rank, _kwargs_rank

execute: 支持 first/all，仅在第一个 worker 上执行，还是全部执行
collect: 如何收集返回的数据，支持 none/flatten/mean/sum/first/last_pp/函数
- none: 不做任何处理
- flatten: 将所有 worker 数据进行拉平，模仿单一 worker 执行的返回结构
- mean/sum: 返回均值或累加值
- first: 仅返回第一个 worker 的结果。一般用于所有 worker 需要输入，但输出结果相同的情况
- last_pp: 返回最后一个 pipeline 的结果，用于 pp 并行的情况
- 函数: 支持自定义收集方法

def _collect(all_results: List, device_mesh):
 # device_mesh是隶属于目标组件的device_mesh
 return ...

sync: 是否以 Ray 的同步方式执行，默认为 False
lazy_collect: 默认为 True，在这种情况下，会不在 driver 进程中收集结果，而在需要这些结果的 worker 中延迟展开，对于具体方法来说，某些方法需要在 driver 中收集，例如收集 loss、metric 等网络负载不大的情况，可以设置为 False

TwinkleClient 客户端

Mon, 01 Jan 0001 00:00:00 +0000

TwinkleClient 是与 Twinkle REST API 交互的 Python 客户端，管理会话、训练任务和检查点。

初始化

from twinkle_client.manager import TwinkleClient

client = TwinkleClient(
 base_url='http://localhost:8000', # 或 TWINKLE_SERVER_URL 环境变量
 api_key='your-api-key', # 或 TWINKLE_SERVER_TOKEN 环境变量
 route_prefix='/twinkle', # API 路由前缀
 session_heartbeat_interval=10, # 心跳间隔（秒）
 session_metadata={'user': 'alice'}, # 可选的会话元数据
)

初始化时客户端会：

将 base_url 和 api_key 设置到共享上下文（所有客户端对象自动使用）
创建服务端会话
启动后台心跳线程保持会话活跃

健康检查

is_healthy = client.health_check() # 返回 True/False
capabilities = client.get_server_capabilities() # 支持的模型

训练任务

# 列出训练任务
runs = client.list_training_runs(limit=20, offset=0)

# 带分页游标列出
runs, cursor = client.list_training_runs_with_cursor(limit=20)

# 获取特定任务
run = client.get_training_run(run_id='run_abc123')

# 按基础模型查找
qwen_runs = client.find_training_run_by_model('Qwen/Qwen3.5-4B')

检查点

# 列出训练任务的检查点
checkpoints = client.list_checkpoints(run_id='run_abc123')

# 获取检查点路径
parsed = client.get_checkpoint_path(run_id, checkpoint_id)
# parsed.path → 文件系统路径
# parsed.twinkle_path → twinkle:// URI

# 获取最新检查点（用于恢复训练）
latest_path = client.get_latest_checkpoint_path(run_id)

# 删除检查点
client.delete_checkpoint(run_id, checkpoint_id)

容量与权重信息

# LoRA 容量
capacity = client.get_capacity_info()
# capacity.max_loras, capacity.used_loras, capacity.free_loras

# 权重元数据
info = client.get_weights_info('twinkle://run_id/weights/checkpoint')
# info.base_model, info.is_lora, info.lora_rank

清理

client.close() # 停止心跳线程（也通过 atexit 自动注册）