Official code for the paper "Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark ".
Authors: Yihua Zhang, Pingzhi Li, Junyuan Hong, Jiaxiang Li, Yimeng Zhang, Wenqing Zheng, Pin-Yu Chen, Jason D. Lee, Wotao Yin, Mingyi Hong, Zhangyang Wang, Sijia Liu, and Tianlong Chen
This repo contains the source code and reproducing guide of ZO-LLM. This research endeavor is designed to help researchers better understand the capabilities, limitations and principles associated with the BP-free, zeroth-order (ZO) optimization as a solution for reducing memory costs during Large Language Model (LLM) fine-tuning. Our study unveils previously overlooked optimization principles, highlighting the importance of task alignment, the role of the forward gradient method, and the balance between algorithm complexity and fine-tuning performance.
This project is organized around the following scopes, including:
- Five LLM families: Roberta, OPT, LLaMA, Vicuna, and Mistral.
- Three task complexities: binary classification, question-answering, and commonsense reasoning.
- Four fine-tuning schemes: full fine-tuning, LoRA, prefix tuning, and prompt tuning.
- Six BP-free optimization methods: ZO-SGD, ZO-SGD-Sign, ZO-SGD-MMT, ZO-SGD-Cons, ZO-Adam, and forward gradient.
- Three novel enhancements to ZO optimization: block-wise descent, hybrid training, and gradient sparsity.
This project is structured around the hyperparameter sweeping for various tasks & models & tuning schemes & optimization
methods. All optimization methods are implemented in zo-bench/trainer.py
. Task configurations are defined in
zo-bench/tasks.py
and zo-bench/templates.py
. The main entry point is zo-bench/run.py
.
.
โโโ zo-bench
โ โโโ modeling_mistral
โ โ โโโโ __init__.py
โ โ โโโ configuration_mistral.py
โ โ โโโ modleing_mistral.py
โ โโโ modeling_llama.py
โ โโโ modeling_opt.py
โ โโโ modeling_roberta.py
โ โโโ prefix_tuning.py
โ โโโ prompt_tuning.py
โ โโโ run.py
โ โโโ tasks.py
โ โโโ templates.py
โ โโโ test_fake_text_memory.py
โ โโโ trainer.py
โ โโโ utils.py
โ โโโ sweep
โ โ โโโ Copa_llama-7b
โ โ โ โโโ adam
โ โ โ โ โโโ adam_copa_ft.yml
โ โ โ โ โโโ adam_copa_lora.yml
โ โ โ โ โโโ adam_copa_prefix.yml
โ โ โ โ โโโ adam_copa_prompt.yml
โ โ โ โโโ forward_grad
โ โ โ โ โโโ forward_grad_copa_ft.yml
โ โ โ โ โโโ forward_grad_copa_lora.yml
โ โ โ โ โโโ forward_grad_copa_prefix.yml
โ โ โ โ โโโ forward_grad_copa_prompt.yml
โ โ โ โโโ sgd
โ โ โ โ โโโ sgd_copa_ft.yml
โ โ โ โ โโโ sgd_copa_lora.yml
โ โ โ โ โโโ sgd_copa_prefix.yml
โ โ โ โ โโโ sgd_copa_prompt.yml
โ โ โ โโโ sign_sgd
โ โ โ โ โโโ sign_sgd_copa_ft.yml
โ โ โ โ โโโ sign_sgd_copa_lora.yml
โ โ โ โ โโโ sign_sgd_copa_prefix.yml
โ โ โ โ โโโ sign_sgd_copa_prompt.yml
โ โ โ โโโ zo_adam
โ โ โ โ โโโ zo_adam_copa_ft.yml
โ โ โ โ โโโ zo_adam_copa_lora.yml
โ โ โ โ โโโ zo_adam_copa_prefix.yml
โ โ โ โ โโโ zo_adam_copa_prompt.yml
โ โ โ โโโ zo_sgd
โ โ โ โ โโโ zo_sgd_copa_ft.yml
โ โ โ โ โโโ zo_sgd_copa_lora.yml
โ โ โ โ โโโ zo_sgd_copa_prefix.yml
โ โ โ โ โโโ zo_sgd_copa_prompt.yml
โ โ โ โโโ zo_sgd_conserv
โ โ โ โ โโโ zo_sgd_conserv_copa_ft.yml
โ โ โ โ โโโ zo_sgd_conserv_copa_lora.yml
โ โ โ โ โโโ zo_sgd_conserv_copa_prefix.yml
โ โ โ โ โโโ zo_sgd_conserv_copa_prompt.yml
โ โ โ โโโ zo_sgd_momen
โ โ โ โ โโโ zo_sgd_momen_copa_ft.yml
โ โ โ โ โโโ zo_sgd_momen_copa_lora.yml
โ โ โ โ โโโ zo_sgd_momen_copa_prefix.yml
โ โ โ โ โโโ zo_sgd_momen_copa_prompt.yml
โ โ โโโ Copa_llama-13b
โ โ โ โโโ ...
โ โ โโโ Copa_mistral
โ โ โ โโโ ...
โ โ โโโ Copa_opt-13b
โ โ โ โโโ ...
โ โ โโโ Copa_vicuna
โ โ โ โโโ ...
โ โ โโโ SST2_opt-1.3b
โ โ โ โโโ ...
โ โ โโโ WinoGrande_llama-7b
โ โ โ โโโ ...
โ โ โโโ WinoGrande_llama-13b
โ โ โ โโโ ...
โ โ โโโ WinoGrande_mistral
โ โ โ โโโ ...
โ โ โโโ WinoGrande_opt-13b
โ โ โ โโโ ...
โ โ โโโ WinoGrande_vicuna
โ โ โ โโโ ...
โโโ environment.yml
All you need is:
conda env create -f environment.yml
We provide detailed hyperparameter settings in sweeps,
where the sweep configuration for tuning a MODEL on TASK under SCHEME with OPTIMIZER is organized as zo-bench/sweeps/TASK_MODEL/OPTIMIZER/SCHEME.yml
.
An example use of sweep for full fine-tuning LLaMA-7B with ZO-SGD on the COPA task is as follows:
~> wandb sweep zo-bench/sweeps/Copa_llama-7b/zo_sgd/zo_sgd_copa_ft.yml
wandb: Creating sweep from: zo-bench/sweeps/Copa_llama-7b/zo_sgd/zo_sgd_copa_ft.yml
wandb: Created sweep with ID: <ID>
wandb: View sweep at: https://wandb.ai/<unique ID>
wandb: Run sweep agent with: wandb agent <unique ID>
~> wandb agent <unique ID>
@misc{zhang2024revisiting,
title={Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark},
author={Yihua Zhang and Pingzhi Li and Junyuan Hong and Jiaxiang Li and Yimeng Zhang and Wenqing Zheng and Pin-Yu Chen and Jason D. Lee and Wotao Yin and Mingyi Hong and Zhangyang Wang and Sijia Liu and Tianlong Chen},
year={2024},
eprint={2402.11592},
archivePrefix={arXiv},
primaryClass={cs.LG}
}