collie.utils.setup_ds_engine

setup_ds_engine(config, model, optimizer=None, lr_scheduler=None)

启动 DeepSpeed 引擎。

参数:

• config (CollieConfig) – CoLLie 的配置

• model (Module) – 模型

• optimizer (Union[Optimizer, Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer], None], default: None) – 优化器

• lr_scheduler (Union[_LRScheduler, Callable[[Optimizer], _LRScheduler], None], default: None) – 学习率调度器

返回:

DeepSpeed 引擎、优化器、dataloader (为 None)、学习率调度器