模型 | Twinkle

支持的模型

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

Twinkle 支持任何兼容 HuggingFace Transformers 或 Megatron-LM 的模型。以下是经过测试的模型列表。

语言模型

模型系列	模型 ID	参数量	特性
Qwen 3.5	`Qwen/Qwen3.5-0.6B` ~ `Qwen/Qwen3.5-235B-A22B`	0.6B–235B	MoE、思考模式
Qwen 2.5	`Qwen/Qwen2.5-0.5B` ~ `Qwen/Qwen2.5-72B`	0.5B–72B	Dense
DeepSeek V4	`deepseek-ai/DeepSeek-V4`	685B MoE	自定义 DSML 编码
DeepSeek R1	`deepseek-ai/DeepSeek-R1`	685B MoE	推理
LLaMA 3	`meta-llama/Llama-3.3-70B-Instruct`	8B–70B	Dense
Mistral	`mistralai/Mistral-7B-v0.3`	7B	Dense
Yi	`01-ai/Yi-1.5-34B`	6B–34B	Dense
GLM-4	`THUDM/glm-4-9b-chat`	9B	Dense
InternLM 2.5	`internlm/internlm2_5-7b-chat`	7B–20B	Dense

视觉语言模型

模型系列	模型 ID	特性
Qwen 3.5 VL	`Qwen/Qwen3.5-VL-3B` ~ `Qwen/Qwen3.5-VL-72B`	图片、视频
Qwen 2.5 VL	`Qwen/Qwen2.5-VL-7B-Instruct`	图片、视频
InternVL 2.5	`OpenGVLab/InternVL2_5-8B`	图片

嵌入模型

模型系列	模型 ID	训练方法
Qwen3 Embedding	`Qwen/Qwen3-Embedding-0.6B`	InfoNCE 对比学习
GTE	`thenlper/gte-large-zh`	InfoNCE 对比学习

模型加载

from twinkle.model import TransformersModel

# 从 ModelScope 加载（ms:// 前缀）
model = TransformersModel(model_id='ms://Qwen/Qwen3.5-4B')

# 从 HuggingFace 加载（hf:// 前缀）
model = TransformersModel(model_id='hf://meta-llama/Llama-3.3-70B-Instruct')

# 本地路径
model = TransformersModel(model_id='/path/to/model')

框架支持

框架	类名	适用场景
Transformers	`TransformersModel`	通用训练（SFT、RLHF、DPO）
Transformers + Multi-LoRA	`MultiLoraTransformersModel`	多租户训练
Megatron-LM	`MegatronModel`	大规模分布式预训练
Megatron + Multi-LoRA	`MultiLoraMegatronModel`	大规模多租户

精度支持

模式	说明
`bf16`	BFloat16 混合精度（推荐 A100/H100）
`fp16`	Float16 混合精度（适用于旧 GPU）
`fp8`	FP8 精度（H100 + Transformer Engine）
`no`	全精度（仅用于调试）

并行策略

策略	配置键	说明
FSDP	`strategy=accelerate`	Accelerate 管理的 FSDP（默认）
原生 FSDP	`strategy=native_fsdp`	PyTorch 原生 FSDP
张量并行	`tp_size`	跨 GPU 切分层
流水线并行	`pp_size`	切分模型阶段
数据并行	`dp_size`	复制模型，切分数据
序列并行	`sequence_parallel`	切分长序列
专家并行	`ep_size`	MoE 专家分布

TwinkleModel

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

TwinkleModel是twinkle所有模型的基类。twinkle的模型不单单包含了模型本身，也包含了模型配套的训练组件。我们在其他文档中介绍的组件基本均在这里进行组合使用。

任何模型符合TwinkleModel的基类设定均可以配合框架的其他组件使用：

class TwinkleModel(ABC):

 @abstractmethod
 def forward(self, *, inputs: Dict[str, Any], **kwargs):
 # 进行一次forward，并返回logits
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def forward_only(self, *, inputs: Dict[str, Any], **kwargs):
 # 以推理模式进行一次forward，并返回logits
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def calculate_loss(self, **kwargs):
 # 使用Loss的子类完成loss计算
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def backward(self, **kwargs):
 # 进行一次backward
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def forward_backward(self, *, inputs: Dict[str, Any], **kwargs):
 # 组合了forward、loss计算、backward过程，并返回loss值
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs):
 # 梯度裁剪，发生在gradient_accumulation_steps完成的条件下，可以在kwargs传入gradient_accumulation_steps
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def step(self, **kwargs):
 # 梯度更新，发生在gradient_accumulation_steps完成的条件下，可以在kwargs传入gradient_accumulation_steps
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def zero_grad(self, **kwargs):
 # 梯度清理，发生在gradient_accumulation_steps完成的条件下，可以在kwargs传入gradient_accumulation_steps
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def lr_step(self, **kwargs):
 # lr更新，发生在gradient_accumulation_steps完成的条件下，可以在kwargs传入gradient_accumulation_steps
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def clip_grad_and_step(self, max_grad_norm: float=1.0, norm_type=2, **kwargs):
 # 组合了clip、step、zero_grad、lr_step
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature, ModelOutput, ...], torch.Tensor]], **kwargs):
 # 设置loss
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def set_optimizer(self, optimizer_cls: Union[Optimizer, Type[Optimizer], str], **kwargs):
 # 设置optimizer
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def set_lr_scheduler(self, scheduler_cls: Union[LRScheduler, Type[LRScheduler], str], **kwargs):
 # 设置lr_scheduler
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def save(self, name: str, output_dir: Optional[str] = None, **kwargs):
 # 保存checkpoint
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def load(self, name: str, output_dir: Optional[str] = None, **kwargs):
 # 加载checkpoint
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def get_state_dict(self, **kwargs):
 # 获取state_dict
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def apply_patch(self, patch_cls: Union[Patch, Type[Patch], str], **kwargs):
 # 对模型应用一个补丁

 @abstractmethod
 def add_metric(self, metric_cls: Union[Metric, str], is_training, **kwargs):
 # 增加一个训练指标，可以设置is_training参数，代表在forward/forward_only中累加。如果不设置，则对forward/forward_only分别生效
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def calculate_metric(self, is_training: bool, **kwargs):
 # 计算metric并返回
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def add_adapter_to_model(self, adapter_name: str, config_or_dir, **kwargs):
 # 增加一个lora

 @abstractmethod
 def set_template(self, template_cls: Union[Template, Type[Template], str], **kwargs):
 # 设置template
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def set_processor(self, processor_cls: Union[InputProcessor, Type[InputProcessor], str], **kwargs):
 # 设置任务数据处理
 # 支持adapter_name参数，对某个lora生效

 @abstractmethod
 def get_train_configs(self, **kwargs) -> str:
 # 获取模型训练配置，用于打印
 # 支持adapter_name参数，对某个lora生效

TransformersModel

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

这个模型封装了transformers的LLM，并可以使用FSDP2、DDP等方式启动并训练模型。

class TransformersModel:

 def __init__(self, # noqa
 model_cls: Optional[Union[Type[PreTrainedModel], str, Type[_BaseAutoModelClass]]] = AutoModelForCausalLM,
 model_id: Optional[str] = None,
 config: Optional[PretrainedConfig] = None,
 device_mesh: Optional[DeviceMesh] = None,
 mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] = 'bf16',
 strategy: Literal['accelerate', 'native_fsdp'] = 'accelerate',
 ddp_config: Dict[str, Any] = None,
 fsdp_config: Dict[str, Any] = None,
 grad_scaler_config: Dict[str, Any] = None,
 memory_efficient_init: bool = False,
 **kwargs):
 ...

 ...

model_cls: 使用哪个类拉起模型，默认为AutoModelForCausalLM
model_id: 模型id
config: 拉起模型的配置
device_mesh: DeviceMesh信息
mixed_precision: 混合精度信息，默认bf16，如果有30系以上显卡建议维持不变
strategy: 如何封装模型使用多卡训练，默认使用accelerate框架。
ddp_config: strategy为accelerate时的DDP配置，参见：
fsdp_config: strategy为accelerate时的FSDP配置，参见：
grad_scaler_config: PyTorch的grad_scaler初始化配置，参见：
memory_efficient_init: 是否启用FSDP内存高效初始化。启用后仅rank 0加载完整权重，其余rank通过广播获取分片参数，降低初始化阶段的内存和显存峰值。默认False。注意：该优化目前仅适用于 transformers <= 4.57.6；对于 transformers >= 5.0.0，可能会导致负面性能影响。
kwargs:
- 如果你不希望传递模型config字段，可以把零星的配置从这里放置进去。后续这些参数会传递到from_pretrained或者from_config中。

TransformersModel支持@remote_class注解，并且支持device_mesh，这意味着它可以运行在ray的worker中。

使用样例：

from twinkle.model import TransformersModel
from twinkle import DeviceMesh
from twinkle.dataloader import DataLoader
dataloader = DataLoader(...)
model = TransformersModel(model_id='ms://Qwen/Qwen3.5-4B', device_mesh=DeviceMesh.from_sizes(dp_size=2, fsdp_size=2), remote_group='actor')
model.add_adapter_to_model(...)
model.set_optimizer(..., adapter_name='...')
for data in dataloader:
 model.forward_backward(...)
 model.clip_grad_and_step(..., gradient_accumulation_steps=16)

检查点保存与续训

TransformersModel.save() 既可以只保存权重，也可以保存可续训的训练检查点。

model.save(name, save_optimizer=True, consumed_train_samples=...) 保存权重、优化器、调度器、scaler、RNG 状态和 trainer_state.json。
model.resume_from_checkpoint(checkpoint_dir) 恢复完整训练状态（权重、优化器、调度器、scaler、RNG），返回 {'cur_step', 'consumed_train_samples', 'gradient_accumulation_steps'}。
model.resume_from_checkpoint(checkpoint_dir, resume_only_model=True) 仅加载权重并返回进度元数据，不恢复优化器状态。
dataloader.resume_from_checkpoint(consumed_train_samples) 跳过已消费的样本。
dataloader.get_state() 返回 {'consumed_train_samples': int} — DataLoader 会自动追踪已消费样本数，无需手动维护计数器。

对于全参训练，恢复模型权重时需要在创建 TransformersModel 时直接把 checkpoint 路径传给 model_id，例如 TransformersModel(model_id='./output/fsdp2/last-checkpoint')，随后再调用 resume_from_checkpoint(...) 恢复优化器和训练进度。

如果需要完整的断点续训流程，包括 dataloader 跳过已消费数据的逻辑，建议直接参考 cookbook/transformers/fsdp2.py。

MultiLoraTransformersModel

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

这个模型继承了TransformersModel，除提供了相同功能外，还提供了分时运行多个lora的能力，主要用于多租户训练。

class MultiLoraTransformersModel:

 def __init__(self, # noqa
 model_cls = AutoModelForCausalLM,
 model_id: Optional[str] = None,
 config: Optional[PretrainedConfig] = None,
 device_mesh: Optional[DeviceMesh] = None,
 mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] = 'bf16',
 grad_scaler_config: Dict[str, Any] = None,
 max_loras: int = 5,
 max_r: int = 32,
 max_length: int = 8192,
 **kwargs):
 ...

 ...

除了和基类相同的参数外，本类提供了几个额外参数用于多lora配置：

max_loras: 最大lora的数量
max_r: 最大的lora rank
max_length: 最大的支持训练长度

之所以存在max_loras和max_r参数，是因为twinkle的多lora技术方案是在DDP wrap之前增加lora到max_loras个，防止后添加的lora无法接受DDP的管理。正因如此，用户的r必须要小于等于max_r的配置，在实际训练时仅会使用lora的部分rank参与计算。

MultiLoraTransformersModel支持@remote_class注解，并且支持device_mesh，这意味着它可以运行在ray的worker中。

租户生命周期

底层使用 MultiLora 管理器来处理租户 LoRA 槽位。关键 API：

acquire_lora

为租户获取一个可用的 LoRA 槽位：

adapter_name = model.multi_lora.acquire_lora('tenant_a', LoraConfig(r=16, lora_alpha=32))

如果所有槽位已被占用或 config.r > max_r，则抛出 RuntimeError

release_lora

释放租户的 LoRA 槽位，权重重置为初始状态：

model.multi_lora.release_lora('tenant_a')

上下文管理器

使用 adapter() 进行作用域激活：

with model.multi_lora.adapter('tenant_a') as name:
 output = model.forward(inputs)

LoraTenant

每个槽位以 LoraTenant 数据类追踪：

@dataclass
class LoraTenant:
 index: int # 槽位索引 (0..max_loras-1)
 adapter_name: str # 内部名称（如 "lora_0"）
 config: LoraConfig # 预分配配置（max_r）
 tenant_adapter_name: str # 面向用户的租户名（空闲时为 None）
 tenant_config: LoraConfig # 租户实际配置（空闲时为 None）

MegatronModel

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

这个模型封装了Megatron的LLM，并可以使用TP/DP/CP/PP/EP组合启动模型。

注意：VPP的支持目前存在问题，请暂时不要配置使用。

class MegatronModel:

 def __init__(
 self,
 model_id: str,
 config: Optional[PretrainedConfig] = None,
 device_mesh: Optional[DeviceMesh] = None,
 mixed_precision: Literal['no', 'fp16', 'bf16'] = 'bf16',
 **kwargs,
 ):
 ...

 ...

model_id: 模型id
config: 拉起模型的配置
device_mesh: DeviceMesh信息
mixed_precision: 混合精度信息，默认bf16，如果有30系以上显卡建议维持不变
kwargs:
- 所有Megatron初始化的参数，即的配置均可以传递到kwargs中。

MegatronModel支持@remote_class注解，并且支持device_mesh，这意味着它可以运行在ray的worker中。

使用样例：

from twinkle.model import MegatronModel
from twinkle import DeviceMesh
from twinkle.dataloader import DataLoader
dataloader = DataLoader(...)
model = MegatronModel(model_id='ms://Qwen/Qwen3.5-4B', device_mesh=DeviceMesh.from_sizes(dp_size=2, tp_size=2, pp_size=2), remote_group='actor')
model.add_adapter_to_model(...)
model.set_optimizer('default', adapter_name='...')
for data in dataloader:
 model.forward_backward(...)
 model.clip_grad_and_step(..., gradient_accumulation_steps=16)

注意：

megatron模型不支持使用AdamW的原始optimizer，仅支持配置MegatronDistributedOptimizer, 你可以传递MegatronDistributedOptimizer, default来使用它

megatron模型不支持使用其他lr_scheduler，仅支持使用OptimizerParamScheduler，你可以传递OptimizerParamScheduler, default来使用它

你需要将tp/cp/dp/ep/pp/sequence_parallel配置传入device_mesh参数中，以方便twinkle管理数据分配。这些参数会由device_mesh代为传递到megatron初始化流程中

MultiLoraMegatronModel

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

这个模型继承了MegatronModel，除提供了相同功能外，还提供了分时运行多个lora的能力，主要用于多租户训练。

class MultiLoraMegatronModel:

 def __init__(self, # noqa
 model_id: str,
 config: Optional[PretrainedConfig] = None,
 device_mesh: Optional[DeviceMesh] = None,
 mixed_precision: Literal['no', 'fp16', 'bf16'] = 'bf16',
 max_loras: int = 5,
 max_r: int = 32,
 max_length: int = 8192,
 **kwargs):
 ...

 ...

除了和基类相同的参数外，本类提供了几个额外参数用于多lora配置：

max_loras: 最大lora的数量
max_r: 最大的lora rank
max_length: 最大的支持训练长度

MultiLoraMegatronModel支持@remote_class注解，并且支持device_mesh，这意味着它可以运行在ray的worker中。