4. 可选模块¶
4.1 CoLLiE 的 Monitor 模块
4.2 CoLLiE 的 Callback 模块
4.3 CoLLiE 的 Server 模块
4.1 CoLLiE 的 Monitor 模块¶
Monitor 用来记录训练或者评测过程中的结果。CoLLiE 目前支持使用 tensorboard
或者 wandb
记录,也支持将结果直接存到 csv
文件中。
名称 |
描述 |
---|---|
BaseMonitor |
Monitor 的基类,其它 Monitor 继承此类 |
LossMonitor |
记录每个 step 的 loss |
LRMonitor |
记录每个 step 的 learning rate |
EvalMonitor |
传给 Evaluator,记录每次评测结果 |
StepTimeMonitor |
记录每个 step 的时间 |
TGSMonitor |
记录每秒每张 GPU 可训练的 token 数 (token / GPU / s) |
MemoryMonitor |
记录每个 step 的内存占用 |
NetworkIOMonitor |
记录每个 step 的网络带宽情况 |
DiskIOMonitor |
记录每个 step 的硬盘读写情况 |
CPUMemoryMonitor |
记录每个 step 的 CPU 内存占用 |
要使用 Monitor,需要指明这次实验的 tag
。记录到 wandb
、tensorboard
和 csv
中的方法如下:
config.ds_config = {
"monitor_config": {
"enabled": True,
"tag": f"your_tag", # job name
"wandb": {
"enabled": True,
"team": "collie_exp",
"project": "llama_alpaca",
"group": f"llama-7b",
},
"tensorboard": {
"enabled": True,
"output_path": "./ds_tb_logs/",
},
"csv_monitor": {
"enabled": True,
"output_path": "./ds_csv_logs/",
}
},
}
然后,将需要记录的指标对应的 Monitor 类传给 Evaluator 或者 Trainer。
from collie.utils.monitor import StepTimeMonitor, MemoryMonitor, LossMonitor, LRMonitor, EvalMonitor
monitors = [
LossMonitor(config),
LRMonitor(config),
StepTimeMonitor(config),
MemoryMonitor(config),
]
evaluator = EvaluatorForPer(
monitors=[EvalMonitor(config)],
**other_kwargs # 这里省略了其他传入参数
)
trainer = Trainer(
evaluators=[evaluator],
monitors=monitors,
**other_kwargs # 这里省略了其他传入参数
)
wandb
上记录的结果如下图所示,左边的 Name
即为设置的 tag
,右边的图表是记录下来的指标。
在finetune_moss_for_training.py
脚本中,示例代码
from collie.utils.monitor import StepTimeMonitor, TGSMonitor, MemoryMonitor, LossMonitor, EvalMonitor
monitors = [
StepTimeMonitor(config),
TGSMonitor(config),
MemoryMonitor(config),
LossMonitor(config),
EvalMonitor(config)
]
4.2 CoLLiE 的 Callback 模块¶
Callback 类在训练过程中的各个回调点被调用,方便用户自定义训练循环。回调类都是基于同一个基类 Callback
实现的,用户可以继承此基类继承,并重写不同的方法来选择回调时间及操作。
名称 |
功能 |
---|---|
Callback |
Callback 的基类,其它 Callback 继承此类 |
CheckpointCallback |
根据用户设置,保存模型权重或者训练 checkpoint |
LoadBestModelCallback |
保存 monitor 值最佳的模型,并在训练结束的时候重新加载模型 |
from collie.callbacks import CheckpointCallback, LoadBestModelCallback
callbacks = [
CheckpointCallback(
"./models",
every_n_epochs=1, # 每一个 epoch 保存一次
model_only=True, # 仅保存模型权重,不保存optimzer、训练步数等断点重训信息
),
LoadBestModelCallback(
"./best_model",
monitor="acc_metric#acc" # 根据这个 metric 保存并加载最优的模型
)
]
trainer = Trainer(
callbacks=callbacks,
**other_kwargs # 这里省略了其他传入参数
)
Callback 的更多用法和详细配置参见 callbacks 文档。
4.3 CoLLiE 的 Server 模块¶
Server 类用于模型的部署,通过 Gradio 或者 Dash 提供基于网页的交互式服务。
from collie.utils import GradioProvider
from transformers import GenerationConfig
server = Server(
model=model,
data_provider=GradioProvider(
tokenizer=tokenizer,
stream=True, # 流式生成
port=8080,
generation_config=GenerationConfig(max_new_tokens=250) # 生成的设置
)
)
server.run()
也可以将 server 传给 trainer,在训练过程中实时试验训练到当前步数模型的能力。
其原理是通过 data_provider 异步地将数据给模型来生成:等到当前 step 训练结束后,模型就会根据用户输入生成回复并且展示到网页上。生成结束后模型会自动继续训练。
trainer = Trainer(
server=server,
**other_kwargs # 这里省略了其他传入参数
)
基于 Gradio 的页面示例如下所示: