CAME Optimizer

ACL 2023 Outstanding Paper Award
Confidence-guided Adaptive Memory Efficient Optimization

This is an official implementation of CAME optimizer in the "Confidence-guided Adaptive Memory Efficient Optimization". Please cite the paper and star this repo if you find CAME useful. Thanks!

Paper | Twitter | Blog | Pypi Package | zhihu

Method

In this work, we studied a confidence-guided strategy to reduce the instability of existing memory efficient optimizers. Based on this strategy, we proposed CAME to simultaneously achieve two goals: fast convergence as in traditional adaptive methods, and low memory usage as in memory-efficient methods.

The pseudo code is presented in the figure with difference with Adafactor in blue fonts.

Install

pip install came-pytorch

Usage

from came_pytorch import CAME
optimizer = CAME(
    model.parameters(),
    lr=2e-4,
    weight_decay=1e-2,
    betas=(0.9, 0.999, 0.9999),
    eps=(1e-30, 1e-16)
)

Hyper-parameter Tuning

• Pre-training: Based on our experiments on BERT-Large, GPT-2 and T5, it's suitable to choose a learning rate for CAME 3-1x smaller than that for AdamW.
• Consider choosing $\beta_3$ between $[0.9995, 0.99995]$ if setting $\beta_1, \beta_2=0.9, 0.999$. Due to computational resource constraints, we did not explore more combinations of three betas. Different training tasks may require different combinations of optimal performance.
• If you have any feedback or comments regarding hyper-parameter tuning, please do not hesitate to provide them to us!

Experiments

Apart from the BERT and T5 experiments shown in the paper, we conduct more and record the results here.

Fine-tuning LLaMA-7B

We fine-tuned LLaMA-7B with stanford-alpaca (52k instruction-tuning dataset). To replicate our result, first register the CAME optimizer to the transformer package. Then in Alpaca training script, change the default optimizer from "adamw" to "came".

Alpaca-7B and Alpaca-7B-CAME are evaluated using Instruct-eval and lm-evaluation-harness.

Pre-training GPT-2

The pre-training of GPT-2 (Medium, 345M) is based on Megatron-LM. To replicate our result, add the CAME optimizer in megatron/optimizer/__init__.py and set the args.optimizer to "came".

Memory Usage Comparison

To ensure a fair comparison, we set the batch size to 1 for the pre-training of GPT-2 (Medium) to examine the memory footprint of CAME and AdamW.

Citation

@inproceedings{luo2023came,
  title={CAME: Confidence-guided Adaptive Memory Efficient Optimization},
  author={Luo, Yang and Ren, Xiaozhe and Zheng, Zangwei and Jiang, Zhuo and Jiang, Xin and You, Yang},
  booktitle={Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  pages={4442--4453},
  year={2023}
}