论文 Full Parameter Fine-Tuning for Large Language Models with Limited Resources 的实现.
在这个工作中，我们提出了一个新的优化器，LOw-Memory Optimization (LOMO)，它将梯度计算和参数更新融合在一步中，以减少内存使用。 我们的方法使得在单张 RTX 3090 上可以进行 7B 模型的全参数微调，或者在单个 8×RTX 3090 的机器上可以进行 65B 模型的全参数微调（RTX 3090 的内存为 24GB）。
torch
deepspeed
transformers
peft
wandbLOMO本身只依赖 PyTorch，其他依赖用于复现我们的论文结果。
我们提供了三种不同方法来微调大型语言模型（LLM）的代码：LOMO，LoRA和LoRA + LOMO。
- 对于使用LOMO进行完全参数微调的实现位于
src/lomo_trainer.py，您可以运行以下命令：
deepspeed --master_port "$port" --include localhost:"$CUDA_VISIBLE_DEVICES" src/train_lomo.py config/args_lomo.yaml- 对于LoRA和LoRA + LOMO的实现位于
src/lomo_lora_trainer.py，您可以运行以下命令：
deepspeed --master_port "$port" --include localhost:"$CUDA_VISIBLE_DEVICES" src/train_lomo_lora.py config/args_lomo_lora.yaml在代码中，我们在src/arguments.py中包含了lora_only参数，用于控制是否仅使用LoRA或LoRA + LOMO。请注意，当将lora_only设置为True时，与LOMO相关的参数将不起作用。
此外，我们还提供了一个简单的run.sh脚本以方便使用。您可以使用以下命令执行代码：
bash run.sh对于数据处理，我们目前只提供了论文中提到的SuperGLUE的六个数据集。如果您希望使用新的数据集，请相应地修改Dataset和DataCollator。
对于评估，我们目前仅为多项选择问答（multiple-choice QA）和生成（generation）任务提供了eval_step代码。如果您有其他需求，请相应地修改LOMOTrainer或LOMOLoRATrainer中的eval_step代码，并为训练器提供必要的compute_metrics函数。
我们在这里提供了我们实验中使用的采样数据集。
由于计算资源有限，我们报告了使用相同随机种子（42）进行的实验中获得的最高结果。
我们在我们的工作中承认了这个限制，并计划在下一个版本中进行重复实验来解决这个问题。
如果您有任何问题，请随时提出问题。
我们通过在PyTorch的反向传播过程中注入钩子函数实现我们的方法。如图中所示，我们为模型的每一个参数都注册了自定义的钩子函数。当一个参数的梯度计算完毕之后(但还没有写入到.grad)，它对应的钩子函数就被调用了。更多关于钩子函数和反向传播的介绍可以参考PyTorch的官方文档。简而言之，反向过程会从一个张量到它的梯度函数，然后把梯度写入.grad，再传递到下一个张量。
我们的自定义钩子函数会扫描所有的参数，如果发现有.grad不为空的参数就进行更新，然后清空并释放相应的.grad。因为一个参数的钩子函数会在它的.grad还未被赋值时调用，整个求导图的最后一个参数的钩子函数调用时，它的.grad还不可用。因此，我们额外进行一次扫描来更新最后一个参数。
@article{lv2023full,
title={Full Parameter Fine-tuning for Large Language Models with Limited Resources},
author={Lv, Kai and Yang, Yuqing and Liu, Tengxiao and Gao, Qinghui and Guo, Qipeng and Qiu, Xipeng},
journal={arXiv preprint arXiv:2306.09782},
year={2023}
}