Shortcuts

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。记录到 wandbtensorboardcsv 中的方法如下:

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,右边的图表是记录下来的指标。

eb40add11c5c4c76b0cacb26b48a04be

    在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 的页面示例如下所示:

8d0ed8f29b334d7bb95a638296d0aaef

其它版本