Adam-mini

本仓库包含了Adam-mini的PyTorch实现，这是Adam的迷你版本，能够以**45%到50%**更少的内存占用达到与AdamW相当或更好的性能。

Adam-mini通过减少Adam中的学习率(lr)资源(即$1/\sqrt{v}$)来减少内存使用。我们认为，如果我们：

(1) 按照我们提出的与Hessian结构相关的原则仔细将参数分成块。 (2) 为每个参数块分配一个单一但良好的lr。

那么$v$中>90%的这些lr可以被无害地移除。

我们找到了一种简单有效的方法来满足这些要求。结果算法如下面的算法1所示。更详细的描述请查看我们的论文：Adam-mini: 使用更少的学习率以获得更多。

如何使用

安装torch (>=1.8.0)并运行以下命令。

pip install adam-mini

或者如果你更喜欢从源代码导入

git clone https://github.com/zyushun/Adam-mini
cd Adam-mini
pip install -e .

然后按如下方式使用Adam-mini优化器。

from adam_mini import Adam_mini

optimizer = Adam_mini(
            named_parameters = model.named_parameters(), 
            lr = lr, 
            betas = (beta1,beta2), 
            eps = eps,
            weight_decay = weight_decay, 
            model_sharding = True,
            dim = model_config.dim,
            n_heads = model_config.n_heads,
            n_kv_heads = model_config.n_kv_heads,
            )

关于所有超参数，包括学习率(lr)、weight_decay、beta1、beta2、eps，我们建议使用与AdamW相同的值。

如果你正在训练语言模型，请向Adam-mini传递以下信息：

model_sharding：如果你使用超过1个GPU的模型并行，包括FSDP和Deepspeed中的zero_1,2,3，请设置为True。如果你使用DDP或单GPU训练，请设置为False。
dim：隐藏特征的维度。如果你训练的不是transformer模型，可以不指定。
n_heads：注意力头的数量。如果你训练的不是transformer模型，可以不指定。
n_kv_heads：Key和Value的头数。或等效地，Group query Attention中的查询组数。也称为"n_query_groups"。如果为None，它将与n_head的值相同。如果你训练的不是transformer模型，可以不指定。

支持

我们目前的Adam-mini实现支持流行的分布式框架和代码库，包括：

DDP分布式框架
FSDP分布式框架
DeepSpeed
Hugginface Trainer
Torchtitan
LLaMA-Factory。详细使用说明可以在examples中查看
更多正在开发中！如果Adam-mini不支持你的代码库，请随时联系我们！

示例

我们在此提供预训练、SFT和RLHF的示例代码。你需要2个A800-80GB或2个A100-80GB GPU来运行以下实验。

示例1：GPT2预训练

我们使用NanoGPT代码库在DDP框架下预训练GPT2系列(125M-1.5B)。通过pip安装依赖：

conda env create -f gpt2/environment.yml
conda activate gpt2
cd examples/gpt2

运行GPT2预训练的代码：

bash run_gpt2.sh

你将得到以下曲线。

示例2：Llama3-8B预训练

我们在此提供使用Torchtitan代码库在FSDP框架下预训练Llama3-8B的示例代码。该代码库还支持Llama系列的其他模型，如Llama2-7B。通过pip安装依赖：

cd examples/llama3_8b
pip install -r requirements.txt
pip3 install --pre torch==2.5.0.dev20240617  --index-url https://download.pytorch.org/whl/nightly/cu121 #或cu118
pip3 install --pre torchdata --index-url https://download.pytorch.org/whl/nightly

下载tokenizer.model。按照官方meta-llama仓库的说明确保你有权访问Llama模型权重。确认访问权限后，你可以运行以下命令将Llama 3 / Llama 2分词器下载到本地机器。

# 从https://huggingface.co/settings/tokens获取你的HF令牌

# llama3 tokenizer.model
python torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Meta-Llama-3-8B --tokenizer_path "original" --hf_token=...

# llama2 tokenizer.model
python torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Llama-2-13b-hf --hf_token=...

在./train_configs/llama3_8b_mini.toml中更改你的数据路径。为了调试，你可以通过GoogleDrive下载小数据集"c4_mini"，并将其放在"./torchtitan/datasets/c4_mini/"路径下。

dataset = "c4" #调试时可以使用"c4_mini"
dataset_path = "your_path/c4" #调试时可以使用"./torchtitan/datasets/c4_mini/"

然后我们可以开始训练。

bash run_llama_training_mini.sh

你将得到以下曲线。 <img src="https://yellow-cdn.veclightyear.com/835a84d5/059be172-30d8-4e84-8938-1947826c2676.png" style="zoom:100%;" />

我们还使用TinyLlama代码库对Llama系列进行预训练。但在保存训练后的检查点时出现了一些意外错误。我们目前正在解决这个问题，希望能尽快完全支持TinyLlama代码库。如果您对预训练Llama系列感兴趣，请先尝试使用上述Torchtitan代码库。

示例2：Llama2-7B监督微调和RLHF

我们使用ReMax代码库在DeepSpeed框架下对Llama2-7B进行微调。通过pip安装依赖：

conda env create -f RLHF/environment.yml
conda activate rlhf
cd examples/RLHF

运行使用LoRA的SFT代码：

bash training_scripts/sft/run_sft_lora.sh

运行全参数SFT代码：

bash training_scripts/sft/run_sft_full.sh

运行RLHF中奖励模型训练的代码：

bash training_scripts/reward/run_reward.sh

运行使用ReMax进行RLHF奖励优化的代码：

bash training_scripts/po/remax/run_remax.sh

您将得到以下曲线。

备注

如何在Huggingface Trainer中使用Adam-mini。如果您使用的是Huggingface Trainer，请按如下方式覆盖"create_optimizer"以更改优化器：

 def create_optimizer(self) -> "torch.optim.Optimizer":
        if self.optimizer is None:
            if (self.finetuning_args.use_adammini):
                self.optimizer = Adam_mini(
            named_parameters = model.named_parameters(), 
            lr = lr, 
            betas = (beta1,beta2), 
            eps = eps,
            weight_decay = weight_decay, 
            model_sharding = True,
            dim = model_config.dim,
            n_heads = model_config.n_heads,
            n_kv_heads = model_config.n_kv_heads,
            )
        return super().create_optimizer()

关于FSDP下的检查点保存： 如果您使用FSDP分布式框架，我们对保存检查点时仍出现意外错误表示歉意。我们正在解决这个问题，并将尽快更新。

关于CPU卸载： 我们当前的Adam-mini实现支持FSDP中的CPU卸载，但不支持DeepSpeed中的CPU卸载。使用DeepSpeed时请关闭卸载功能。我们将尽快解决这个问题。

更新日志

[24/08/09] 我们现在支持在LLaMA-Factory中使用Adam-mini。

[24/07/21] 我们现在支持通过pip安装Adam-mini。

致谢

上述代码主要基于NanoGPT、Torchtitan、ReMax和DeepSpeed的代码库。
我们要感谢@lessw2020和@awgu对Torchtitan的支持以及对重构Adam-mini代码的宝贵建议！
我们要感谢@Mrw33554432提交的pip安装拉取请求！
我们要感谢@relic-yuexi提交的LLaMA-Factory拉取请求！

引用

如果您觉得这段代码有帮助，请按以下格式引用我们的论文。

@article{zhang2024adam,
  title     = {Adam-mini: Use Fewer Learning Rates To Gain More},
  author    = {Zhang, Yushun and Chen, Congliang  and Li, Ziniu and Ding, Tian and Wu, Chenwei and Ye, Yinyu and Luo, Zhi-Quan and Sun, Ruoyu},
  booktitle = {arXiv preprint arXiv:2406.16793},
  year      = {2024},
}