TorchSampler

TorchSampler 使用原生 PyTorch 和 transformers 进行推理,适合小规模采样或调试。

使用示例

from twinkle.sampler import TorchSampler
from twinkle import DeviceMesh

sampler = TorchSampler(
    model_id='ms://Qwen/Qwen3.5-4B',
    device_mesh=DeviceMesh.from_sizes(dp_size=1),
)

responses = sampler.sample(trajectories, sampling_params=params)

特性

  • 简单易用: 基于 transformers 的标准接口

  • 灵活性高: 容易定制和扩展

  • 内存占用小: 适合小规模采样

适用场景

TorchSampler 特别适合以下场景:

  • 调试和开发: 简单直接,容易调试

  • 小规模实验: 不需要高吞吐量的场景

  • 自定义需求: 需要修改采样逻辑的场景

  • 资源受限: 内存或GPU资源有限的环境

对于生产环境或大规模训练,建议使用 vLLMSampler 以获得更好的性能。