Table of Contents
Shortcuts

5. 高效训练

5.1   CoLLiE 实现 LOMO 算法

5.2   CoLLiE 结合 PEFT 算法

5.3   其它高效优化器

5.1   CoLLiE 实现 LOMO 算法

LOMO 和 AdaLomo 提供了有限显存下的全量微调方案。其显存占用与 PEFT 方法相当，并且适用于包括指令微调和继续预训练在内的一系列任务。

LOMO (LOw Memory Optimization) 算法把梯度反向传播和参数更新放在 fused backward 一个过程中进行，在训练过程中将得到的梯度直接应用到参数更新上，从而大大减少了梯度的显存占用。AdaLomo 算法通过非负矩阵分解为每个参数提供了可变的学习率，在明显提高收敛效果的同时保持了同样低的显存使用。

要在 CoLLiE 中使用 LOMO 系列算法，只需要改变优化器即可：

from collie.optim import AdaLomo
optimizer = AdaLomo(model, lr=1e-3)

目前 LOMO 系列优化器支持 ZeRO-3 和张量并行。因为其 fused backward 过程，并不支持流水线并行。

5.2   CoLLiE 结合 PEFT 算法

PEFT (Parameter-efficient fine-tuning) 方法通过减少训练参数达到节省显存的效果。CoLLiE 通过集成 peft 库为用户提供易用且流畅的 PEFT 方法。

以 LoRA 为例，在 config 中指定 peft_config 即可使用对应的 PEFT 方法进行训练。

from peft import LoraConfig, TaskType
config.peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

peft_config 的参数设置与 peft 库相同，更多 PEFT 方法的配置可以参考 peft 文档。

同时，只需要将 PEFT 部分的参数传入优化器进行更新，保持预训练模型本身权重不变。

optimizer = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-4
)

5.3   其它高效优化器

除了上述内存高效的优化方法，CoLLiE 中还实现了包括 Adan、Lion 和 Sophia 在内的高效优化器，实现代码在 collie/optim 文件夹中。

其它版本