检查点引擎 | Twinkle

CheckpointEngine

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

CheckpointEngine (检查点引擎) 是用于在训练器和推理进程之间同步模型权重的组件,主要用于 RLHF 训练中 Actor 模型和 Rollout 采样器之间的权重同步。

基本接口

class CheckpointEngine(ABC):
 """检查点引擎基类

 检查点引擎处理训练器和推理进程之间的权重同步。
 """

 @abstractmethod
 def prepare(self) -> dict[str, Any]:
 """准备权重同步前的准备工作"""
 ...

 @abstractmethod
 def init_process_group(self, rank: int, world_size: int, **kwargs):
 """初始化进程组"""
 ...

 @abstractmethod
 async def send_weights(self, weight_generator):
 """发送权重(在训练器进程中调用)"""
 ...

 @abstractmethod
 def receive_weights(self) -> AsyncGenerator:
 """接收权重(在推理进程中调用)"""
 ...

 @abstractmethod
 def finalize(self):
 """清理资源"""
 ...

可用的检查点引擎

Twinkle 提供了两种检查点引擎实现:

NCCLCheckpointEngine

使用 NCCL 进行 GPU 间高速权重传输的检查点引擎。

高速传输: 使用 NCCL 实现 GPU 间点对点高速传输
零拷贝: 直接在 GPU 内存间传输,无需经过 CPU
分桶传输: 支持大模型的分桶传输

详见:

HCCLCheckpointEngine

使用 HCCL 进行昇腾 NPU 间权重传输的检查点引擎。

NPU 优化: 专为昇腾 NPU 优化的权重传输
高效通信: 使用 HCCL 实现 NPU 间高速通信
兼容接口: 与 NCCLCheckpointEngine 保持一致的接口

详见:

如何选择

NCCLCheckpointEngine: 适用于 GPU 环境,提供最高的传输性能
HCCLCheckpointEngine: 适用于昇腾 NPU 环境

检查点引擎是 RLHF 训练基础设施的关键组件,确保训练器和采样器使用一致的模型权重。目前的同步分为merge_and_sync=True/False两种情况，为True时将lora合并仅基模并同步，为False时仅同步lora权重。另外，多租户直接附加lora文件到vLLM上，在merge_and_sync=False，或使用多租户时， vLLM的启动参数需要开启enable_lora=True，在merge_and_sync=True或全参时，该值设置为False.

NCCLCheckpointEngine

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

使用 NCCL 进行 GPU 间高速权重传输的检查点引擎。

使用示例

from twinkle.checkpoint_engine import NCCLCheckpointEngine

# 在训练进程 (rank 0)
engine = NCCLCheckpointEngine(bucket_size=512<<20) # 512MB bucket
engine.is_master = True
engine.prepare()
engine.init_process_group(rank=0, world_size=5)

# 发送权重
await engine.send_weights(model.named_parameters())
engine.finalize()

# 在推理进程 (rank 1-4)
engine = NCCLCheckpointEngine(bucket_size=512<<20)
engine.prepare()
engine.init_process_group(rank=1, world_size=5, master_metadata=metadata)

# 接收权重
async for name, tensor in engine.receive_weights():
 model.load_state_dict({name: tensor}, strict=False)
engine.finalize()

特性

高速传输: 使用 NCCL 实现 GPU 间点对点高速传输
零拷贝: 直接在 GPU 内存间传输,无需经过 CPU
分桶传输: 支持大模型的分桶传输

配置参数

bucket_size: 权重桶大小,控制每次传输的数据量。较大的桶可以提高传输效率,但会占用更多内存
timeout: 传输超时时间

NCCLCheckpointEngine 是 GPU 训练的推荐选择,提供最高的传输性能。

HCCLCheckpointEngine

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

使用 HCCL 进行昇腾 NPU 间权重传输的检查点引擎。

使用示例

from twinkle.checkpoint_engine import HCCLCheckpointEngine

engine = HCCLCheckpointEngine(bucket_size=512<<20)
# 使用方式与 NCCLCheckpointEngine 相同

特性

NPU 优化: 专为昇腾 NPU 优化的权重传输
高效通信: 使用 HCCL 实现 NPU 间高速通信
兼容接口: 与 NCCLCheckpointEngine 保持一致的接口

适用场景

HCCLCheckpointEngine 专门用于昇腾 NPU 环境:

使用华为昇腾 NPU 进行训练
需要在 NPU 间同步模型权重
大规模 NPU 集群部署

环境变量

TWINKLE_CKPT_HCCL_META_TIMEOUT_S: 控制 HCCL CheckpointEngine 元数据握手通道（ZMQ REQ/REP）的超时时间（秒）。默认值为 300。该值应设置为大于 0 的整数。

在昇腾 NPU 环境中,HCCLCheckpointEngine 提供了与 NCCL 相当的性能。