彻底疯狂！实操用Golare在单块4090上预训练大模型及实测效率

kirov1024

已于 2024-03-21 23:35:16 修改

阅读量2.9k

点赞数 26

文章标签： gpt 开源人工智能 python 神经网络

于 2024-03-21 23:29:25 首次发布

本文链接：https://blog.csdn.net/kirov1024/article/details/136923992

版权

在单块4090上预训练大模型？？

田渊栋等人新作：突破内存瓶颈，GaLore实现22G显存预训练大模型

最近正想着微调scGPT这种生物大模型，把那篇文章的小部分结果复现一下，结果24G显存还没推理就帕的一下装满了，看着手里的双4090只能本地部署部署的14B，微调微调还要并行运行，可怜的不要不要的，早知道多整几块2080Ti魔改了。

詹姆斯表示2080 22g是一张极为......

正当我纳闷的时候，想着有没有刚好4090就能跑下的模型，不浪费一点内存。还真不是开玩笑，大佬Meta FAIR 田渊栋的论文《 MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases》就提供了GaLore架构，不牺牲模型性能的前提下显著降低内存消耗，主打在移动设备上运行10亿以下参数LLM。最重要的是，它能实现实现22G显存预训练大模型。这种送上门来的好项目怎么能不去骗去偷袭呢？

文章的预印本：https://arxiv.org/pdf/2402.14905.pdf

论文背景

LLM（Large Language Models）在运行时对内存的要求极高，不仅涵盖了模型参数本身的海量存储需求（常常以亿计），还包括了诸如Adam优化器所维护的梯度动量和方差等额外状态信息。例如，在初始预训练阶段，LLaMA 7B模型每批次仅处理一个样本时，其内存消耗就相当可观：大约58GB的总内存使用量中，有14GB专门用于存储模型可训练参数，而高达42GB的空间被Adam优化器的状态变量及权重梯度占据，另外还有2GB用于临时的激活函数输出。这意味着，对于拥有像NVIDIA RTX 4090这种仅配备24GB显存的消费级GPU而言，直接进行LLM的训练显然是不够的，因其内存资源远不足以容纳如此庞大的计算需求。

而通过梯度低秩投影，GaLore可以在不牺牲模型性能的前提下显著降低内存消耗，使得在消费级GPU上训练大型语言模型成为可能。田大佬本人强调：我们没有像 LoRA 那样假设低秩权重结构，而是证明权重梯度自然是低秩的，因此可以投影到（变化的）低维空间中。因此，我们同时节省了梯度、Adam 动量和方差的内存。那么，到底是如何让显存大大降低呢？

算法策略

本篇论文核心提出了一种名为“GaLore”的梯度低秩投影训练策略，用于提高大型语言模型（LLMs）的内存效率。算法的基本思想是利用权重矩阵梯度在训练过程中表现出的缓慢变化的低秩结构，而不是直接将权重矩阵近似为低秩矩阵。具体做法如下：

首先，理论分析表明权重矩阵的梯度G具有低秩特征。基于这一发现，GaLore计算两个投影矩阵P∈Rm×r和Q∈Rn×r，用以将梯度矩阵G投影到低秩形式P⊤GQ，从而大幅减少依赖于梯度分量统计信息的优化器状态所占用的内存

对于链式可逆神经网络，其权重矩阵Wl在第l层的梯度Gl可以用特定形式表示定理3.2，其中包含前向传播中的雅可比矩阵及其乘积，以及与输出相关的项。

定理3.2（可逆模型的梯度形式）：在具有ℓ2损失函数φ:= 1/2 ||y - N(x)||^2的链式可逆神经网络N(x) := NL(NL-1(...N1(x)))中，在批量大小为1的情况下，第l层权重矩阵Wl的梯度Gl有以下形式：

其中：

请注意，这里假设了在softmax目标函数下梯度也有类似的结构，并适用于类似的情况。

https://github.com/jiaweizzhao/GaLore

GaLore的核心操作是在每个预设的频率T时刻重新计算投影矩阵P和Q，通过在不同的低秩子空间中进行更新来逐步逼近最优解。 例如，在时间步t1∈[0, T1-1]期间，权重W在由固定投影矩阵Pt1和Qt1确定的低秩子空间中根据投影梯度G˜t1进行更新； 当达到时间步T1后，会重新计算Pt2和Qt2进入下一个低秩子空间，并重复此过程直至收敛。

同时作者给出了梯度稳定秩随时间变化的不等式以及梯度随时间指数衰减至低秩的证明，这有助于理解为何GaLore能够在保证性能的同时减少内存需求。在附带的Adam优化器结合GaLore的算法描述中，即每经过一定频率T的时间步长，根据当前梯度Gt重新初始化投影矩阵，并基于此进行低秩投影后的梯度更新，以实现高效、全参数学习下的内存节省。

Adam with GaLore 结合了Adam优化器和梯度低秩投影技术，下面是结合这两种方法进行训练的简化算法流程：

首先，按照标准Adam更新规则初始化一阶动量M0和二阶动量V0，然后开始迭代训练过程：
获取当前时间步t的权重矩阵Wt对应的梯度Gt = -∇Wφt(Wt)。
如果t满足一定的子空间频率条件（例如t mod T= 0），则执行以下操作：
对梯度矩阵Gt进行奇异值分解（SVD），得到U, S, V三个矩阵。
初始化左投影矩阵Pt为U[:, :r]（若m <= n，否则选择合适的投影矩阵）。
否则，复用上一步的投影矩阵Pt = Pt-1。
将梯度投影到紧凑空间得到Rt = P⊤ t Gt。
使用Adam或其他相似优化器更新投影后的梯度Rt：
更新一阶动量Mt：Mt ← β1·Mt-1 + (1 - β1)·Rt
更新二阶动量Vt：Vt ← β2·Vt-1 + (1 - β2)·Rt²
计算归一化的低秩梯度Nt：Nt = Mt / (√Vt + ϵ)
投影回原空间的梯度更新量G˜t：G˜t ← α·P Nt
更新权重Wt：Wt ← Wt-1 + η·G˜t，其中η为学习率。
继续增加时间步t，直到满足收敛条件。

实验结果

本文首先使用 Adam 优化器将 GaLore 与现有的低秩方法进行了比较，研究者将 GaLore 应用于 AdamW、8 bit Adam 和 Adafactor 优化器。他们采用一阶统计的 Adafactor，以避免性能下降。

同时，在大多数任务中，GaLore 都能以更少的内存占用获得比 LoRA 更高的性能。这表明，GaLore 可以作为一种全栈内存高效训练策略，用于 LLM 预训练和微调。

实操Golare！单张4090预训练7B大模型

将近4个小时才跑了171/150000,如果按照不减速，显卡不损耗的话要将近3509个小时，折合146天5个月！ 估计再过一段时间学校都要来抓我挖矿了，这是真的以时间换空间...... 然后是跑上之后显卡信息，基本上在80度左右

数据集用的是C4，下面是一些运行示例。

torchrun --standalone --nproc_per_node 1 torchrun_main.py     --model_config configs/llama_7b.json     --lr 0.005     --galore_scale 0.25     --rank 1024     --update_proj_gap 500     --batch_size 16     --total_batch_size 512     --activation_checkpointing     --num_training_steps 150000     --warmup_steps 15000     --weight_decay 0     --grad_clipping 1.0     --dtype bfloat16     --eval_every 1000     --single_gpu     --optimizer galore_adamw8bit_per_layer
Starting script
2024-03-21 18:00:59.614 | INFO     | __main__:main:140 - Global rank 0, local rank 0, device: 0
2024-03-21 18:00:59.615 | INFO     | __main__:main:144 - Process group initialized
wandb: Currently logged in as: 12017016 (zslnb). Use `wandb login --relogin` to force relogin
wandb: Tracking run with wandb version 0.16.4
wandb: Run data is saved locally in /media/ubuntu/20TB/Project/Deeplearning/GPT/pythonProject/GaLore-master/wandb/run-20240321_180100-lk9n5ytr
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run stellar-durian-3
wandb: ⭐️ View project at https://wandb.ai/zslnb/galore-c4
wandb: 🚀 View run at https://wandb.ai/zslnb/galore-c4/runs/lk9n5ytr
2024-03-21 18:01:02.072 | INFO     | __main__:main:163 - Using dist with rank 0 (only rank 0 will log)
2024-03-21 18:01:02.072 | INFO     | __main__:main:164 - ****************************************
2024-03-21 18:01:02.072 | INFO     | __main__:main:165 - Starting training with the arguments
2024-03-21 18:01:02.072 | INFO     | __main__:main:167 - model_config                   configs/llama_7b.json
2024-03-21 18:01:02.072 | INFO     | __main__:main:167 - use_hf_model                   False
2024-03-21 18:01:02.072 | INFO     | __main__:main:167 - continue_from                  None
2024-03-21 18:01:02.072 | INFO     | __main__:main:167 - batch_size                     16
2024-03-21 18:01:02.072 | INFO     | __main__:main:167 - gradient_accumulation          32
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - total_batch_size               512
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - max_length                     256
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - optimizer                      galore_adamw8bit_per_layer
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - lr                             0.005
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - scheduler                      cosine
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - min_lr_ratio                   0.1
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - activation_checkpointing       True
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - weight_decay                   0.0
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - warmup_steps                   15000
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - eval_every                     1000
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - num_training_steps             150000
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - max_train_tokens               None
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - save_every                     10000
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - save_dir                       checkpoints/llama_7b-2024-03-21-18-00-59
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - tags                           None
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - dtype                          bfloat16
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - workers                        8
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - seed                           0
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - name                           test
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - grad_clipping                  1.0
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - beta1                          0.0
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - rank                           1024
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - update_proj_gap                500
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - galore_scale                   0.25
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - proj_type                      std
2024-03-21 18:01:02.073 | INFO     | __main__:main:167 - single_gpu                     True
2024-03-21 18:01:02.073 | INFO     | __main__:main:168 - ****************************************
Downloading readme: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 41.1k/41.1k [00:00<00:00, 805kB/s]
2024-03-21 18:01:54.563 | INFO     | __main__:main:174 - Shuffling data with seed 42
Downloading config.json: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.21k/1.21k [00:00<00:00, 3.69MB/s]
Downloading spiece.model: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 792k/792k [00:00<00:00, 4.21MB/s]
Downloading tokenizer.json: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.39M/1.39M [00:00<00:00, 1.91MB/s]
wandb: WARNING Saving files without folders. If you want to preserve sub directories pass base_path to wandb.save, i.e. wandb.save("/mnt/folder/file.h5", base_path="/mnt")
Update steps:   0%|                                  | 0/150000 [00:00<?, ?it/s]enable GaLore for weights in module:  model.layers.0.self_attn.q_proj
enable GaLore for weights in module:  model.layers.0.self_attn.k_proj
enable GaLore for weights in module:  model.layers.0.self_attn.v_proj
enable GaLore for weights in module:  model.layers.0.self_attn.o_proj
enable GaLore for weights in module:  model.layers.0.mlp.gate_proj
enable GaLore for weights in module:  model.layers.0.mlp.down_proj
enable GaLore for weights in module:  model.layers.0.mlp.up_proj
enable GaLore for weights in module:  model.layers.1.self_attn.q_proj
enable GaLore for weights in module:  model.layers.1.self_attn.k_proj
enable GaLore for weights in module:  model.layers.1.self_attn.v_proj
enable GaLore for weights in module:  model.layers.1.self_attn.o_proj
enable GaLore for weights in module:  model.layers.1.mlp.gate_proj
enable GaLore for weights in module:  model.layers.1.mlp.down_proj
enable GaLore for weights in module:  model.layers.1.mlp.up_proj
enable GaLore for weights in module:  model.layers.2.self_attn.q_proj
enable GaLore for weights in module:  model.layers.2.self_attn.k_proj
enable GaLore for weights in module:  model.layers.2.self_attn.v_proj
enable GaLore for weights in module:  model.layers.2.self_attn.o_proj
enable GaLore for weights in module:  model.layers.2.mlp.gate_proj
enable GaLore for weights in module:  model.layers.2.mlp.down_proj
enable GaLore for weights in module:  model.layers.2.mlp.up_proj
enable GaLore for weights in module:  model.layers.3.self_attn.q_proj
enable GaLore for weights in module:  model.layers.3.self_attn.k_proj
enable GaLore for weights in module:  model.layers.3.self_attn.v_proj
enable GaLore for weights in module:  model.layers.3.self_attn.o_proj
enable GaLore for weights in module:  model.layers.3.mlp.gate_proj
enable GaLore for weights in module:  model.layers.3.mlp.down_proj
enable GaLore for weights in module:  model.layers.3.mlp.up_proj
enable GaLore for weights in module:  model.layers.4.self_attn.q_proj
enable GaLore for weights in module:  model.layers.4.self_attn.k_proj
enable GaLore for weights in module:  model.layers.4.self_attn.v_proj
enable GaLore for weights in module:  model.layers.4.self_attn.o_proj
enable GaLore for weights in module:  model.layers.4.mlp.gate_proj
enable GaLore for weights in module:  model.layers.4.mlp.down_proj
enable GaLore for weights in module:  model.layers.4.mlp.up_proj
enable GaLore for weights in module:  model.layers.5.self_attn.q_proj
enable GaLore for weights in module:  model.layers.5.self_attn.k_proj
enable GaLore for weights in module:  model.layers.5.self_attn.v_proj
enable GaLore for weights in module:  model.layers.5.self_attn.o_proj
enable GaLore for weights in module:  model.layers.5.mlp.gate_proj
enable GaLore for weights in module:  model.layers.5.mlp.down_proj
enable GaLore for weights in module:  model.layers.5.mlp.up_proj
enable GaLore for weights in module:  model.layers.6.self_attn.q_proj
enable GaLore for weights in module:  model.layers.6.self_attn.k_proj
enable GaLore for weights in module:  model.layers.6.self_attn.v_proj
enable GaLore for weights in module:  model.layers.6.self_attn.o_proj
enable GaLore for weights in module:  model.layers.6.mlp.gate_proj
enable GaLore for weights in module:  model.layers.6.mlp.down_proj
enable GaLore for weights in module:  model.layers.6.mlp.up_proj
enable GaLore for weights in module:  model.layers.7.self_attn.q_proj
enable GaLore for weights in module:  model.layers.7.self_attn.k_proj
enable GaLore for weights in module:  model.layers.7.self_attn.v_proj
enable GaLore for weights in module:  model.layers.7.self_attn.o_proj
enable GaLore for weights in module:  model.layers.7.mlp.gate_proj
enable GaLore for weights in module:  model.layers.7.mlp.down_proj
enable GaLore for weights in module:  model.layers.7.mlp.up_proj
enable GaLore for weights in module:  model.layers.8.self_attn.q_proj
enable GaLore for weights in module:  model.layers.8.self_attn.k_proj
enable GaLore for weights in module:  model.layers.8.self_attn.v_proj
enable GaLore for weights in module:  model.layers.8.self_attn.o_proj
enable GaLore for weights in module:  model.layers.8.mlp.gate_proj
enable GaLore for weights in module:  model.layers.8.mlp.down_proj
enable GaLore for weights in module:  model.layers.8.mlp.up_proj
enable GaLore for weights in module:  model.layers.9.self_attn.q_proj
enable GaLore for weights in module:  model.layers.9.self_attn.k_proj
enable GaLore for weights in module:  model.layers.9.self_attn.v_proj
enable GaLore for weights in module:  model.layers.9.self_attn.o_proj
enable GaLore for weights in module:  model.layers.9.mlp.gate_proj
enable GaLore for weights in module:  model.layers.9.mlp.down_proj
enable GaLore for weights in module:  model.layers.9.mlp.up_proj
enable GaLore for weights in module:  model.layers.10.self_attn.q_proj
enable GaLore for weights in module:  model.layers.10.self_attn.k_proj
enable GaLore for weights in module:  model.layers.10.self_attn.v_proj
enable GaLore for weights in module:  model.layers.10.self_attn.o_proj
enable GaLore for weights in module:  model.layers.10.mlp.gate_proj
enable GaLore for weights in module:  model.layers.10.mlp.down_proj
enable GaLore for weights in module:  model.layers.10.mlp.up_proj
enable GaLore for weights in module:  model.layers.11.self_attn.q_proj
enable GaLore for weights in module:  model.layers.11.self_attn.k_proj
enable GaLore for weights in module:  model.layers.11.self_attn.v_proj
enable GaLore for weights in module:  model.layers.11.self_attn.o_proj
enable GaLore for weights in module:  model.layers.11.mlp.gate_proj
enable GaLore for weights in module:  model.layers.11.mlp.down_proj
enable GaLore for weights in module:  model.layers.11.mlp.up_proj
enable GaLore for weights in module:  model.layers.12.self_attn.q_proj
enable GaLore for weights in module:  model.layers.12.self_attn.k_proj
enable GaLore for weights in module:  model.layers.12.self_attn.v_proj
enable GaLore for weights in module:  model.layers.12.self_attn.o_proj
enable GaLore for weights in module:  model.layers.12.mlp.gate_proj
enable GaLore for weights in module:  model.layers.12.mlp.down_proj
enable GaLore for weights in module:  model.layers.12.mlp.up_proj
enable GaLore for weights in module:  model.layers.13.self_attn.q_proj
enable GaLore for weights in module:  model.layers.13.self_attn.k_proj
enable GaLore for weights in module:  model.layers.13.self_attn.v_proj
enable GaLore for weights in module:  model.layers.13.self_attn.o_proj
enable GaLore for weights in module:  model.layers.13.mlp.gate_proj
enable GaLore for weights in module:  model.layers.13.mlp.down_proj
enable GaLore for weights in module:  model.layers.13.mlp.up_proj
enable GaLore for weights in module:  model.layers.14.self_attn.q_proj
enable GaLore for weights in module:  model.layers.14.self_attn.k_proj
enable GaLore for weights in module:  model.layers.14.self_attn.v_proj
enable GaLore for weights in module:  model.layers.14.self_attn.o_proj
enable GaLore for weights in module:  model.layers.14.mlp.gate_proj
enable GaLore for weights in module:  model.layers.14.mlp.down_proj
enable GaLore for weights in module:  model.layers.14.mlp.up_proj
enable GaLore for weights in module:  model.layers.15.self_attn.q_proj
enable GaLore for weights in module:  model.layers.15.self_attn.k_proj
enable GaLore for weights in module:  model.layers.15.self_attn.v_proj
enable GaLore for weights in module:  model.layers.15.self_attn.o_proj
enable GaLore for weights in module:  model.layers.15.mlp.gate_proj
enable GaLore for weights in module:  model.layers.15.mlp.down_proj
enable GaLore for weights in module:  model.layers.15.mlp.up_proj
enable GaLore for weights in module:  model.layers.16.self_attn.q_proj
enable GaLore for weights in module:  model.layers.16.self_attn.k_proj
enable GaLore for weights in module:  model.layers.16.self_attn.v_proj
enable GaLore for weights in module:  model.layers.16.self_attn.o_proj
enable GaLore for weights in module:  model.layers.16.mlp.gate_proj
enable GaLore for weights in module:  model.layers.16.mlp.down_proj
enable GaLore for weights in module:  model.layers.16.mlp.up_proj
enable GaLore for weights in module:  model.layers.17.self_attn.q_proj
enable GaLore for weights in module:  model.layers.17.self_attn.k_proj
enable GaLore for weights in module:  model.layers.17.self_attn.v_proj
enable GaLore for weights in module:  model.layers.17.self_attn.o_proj
enable GaLore for weights in module:  model.layers.17.mlp.gate_proj
enable GaLore for weights in module:  model.layers.17.mlp.down_proj
enable GaLore for weights in module:  model.layers.17.mlp.up_proj
enable GaLore for weights in module:  model.layers.18.self_attn.q_proj
enable GaLore for weights in module:  model.layers.18.self_attn.k_proj
enable GaLore for weights in module:  model.layers.18.self_attn.v_proj
enable GaLore for weights in module:  model.layers.18.self_attn.o_proj
enable GaLore for weights in module:  model.layers.18.mlp.gate_proj
enable GaLore for weights in module:  model.layers.18.mlp.down_proj
enable GaLore for weights in module:  model.layers.18.mlp.up_proj
enable GaLore for weights in module:  model.layers.19.self_attn.q_proj
enable GaLore for weights in module:  model.layers.19.self_attn.k_proj
enable GaLore for weights in module:  model.layers.19.self_attn.v_proj
enable GaLore for weights in module:  model.layers.19.self_attn.o_proj
enable GaLore for weights in module:  model.layers.19.mlp.gate_proj
enable GaLore for weights in module:  model.layers.19.mlp.down_proj
enable GaLore for weights in module:  model.layers.19.mlp.up_proj
enable GaLore for weights in module:  model.layers.20.self_attn.q_proj
enable GaLore for weights in module:  model.layers.20.self_attn.k_proj
enable GaLore for weights in module:  model.layers.20.self_attn.v_proj
enable GaLore for weights in module:  model.layers.20.self_attn.o_proj
enable GaLore for weights in module:  model.layers.20.mlp.gate_proj
enable GaLore for weights in module:  model.layers.20.mlp.down_proj
enable GaLore for weights in module:  model.layers.20.mlp.up_proj
enable GaLore for weights in module:  model.layers.21.self_attn.q_proj
enable GaLore for weights in module:  model.layers.21.self_attn.k_proj
enable GaLore for weights in module:  model.layers.21.self_attn.v_proj
enable GaLore for weights in module:  model.layers.21.self_attn.o_proj
enable GaLore for weights in module:  model.layers.21.mlp.gate_proj
enable GaLore for weights in module:  model.layers.21.mlp.down_proj
enable GaLore for weights in module:  model.layers.21.mlp.up_proj
enable GaLore for weights in module:  model.layers.22.self_attn.q_proj
enable GaLore for weights in module:  model.layers.22.self_attn.k_proj
enable GaLore for weights in module:  model.layers.22.self_attn.v_proj
enable GaLore for weights in module:  model.layers.22.self_attn.o_proj
enable GaLore for weights in module:  model.layers.22.mlp.gate_proj
enable GaLore for weights in module:  model.layers.22.mlp.down_proj
enable GaLore for weights in module:  model.layers.22.mlp.up_proj
enable GaLore for weights in module:  model.layers.23.self_attn.q_proj
enable GaLore for weights in module:  model.layers.23.self_attn.k_proj
enable GaLore for weights in module:  model.layers.23.self_attn.v_proj
enable GaLore for weights in module:  model.layers.23.self_attn.o_proj
enable GaLore for weights in module:  model.layers.23.mlp.gate_proj
enable GaLore for weights in module:  model.layers.23.mlp.down_proj
enable GaLore for weights in module:  model.layers.23.mlp.up_proj
enable GaLore for weights in module:  model.layers.24.self_attn.q_proj
enable GaLore for weights in module:  model.layers.24.self_attn.k_proj
enable GaLore for weights in module:  model.layers.24.self_attn.v_proj
enable GaLore for weights in module:  model.layers.24.self_attn.o_proj
enable GaLore for weights in module:  model.layers.24.mlp.gate_proj
enable GaLore for weights in module:  model.layers.24.mlp.down_proj
enable GaLore for weights in module:  model.layers.24.mlp.up_proj
enable GaLore for weights in module:  model.layers.25.self_attn.q_proj
enable GaLore for weights in module:  model.layers.25.self_attn.k_proj
enable GaLore for weights in module:  model.layers.25.self_attn.v_proj
enable GaLore for weights in module:  model.layers.25.self_attn.o_proj
enable GaLore for weights in module:  model.layers.25.mlp.gate_proj
enable GaLore for weights in module:  model.layers.25.mlp.down_proj
enable GaLore for weights in module:  model.layers.25.mlp.up_proj
enable GaLore for weights in module:  model.layers.26.self_attn.q_proj
enable GaLore for weights in module:  model.layers.26.self_attn.k_proj
enable GaLore for weights in module:  model.layers.26.self_attn.v_proj
enable GaLore for weights in module:  model.layers.26.self_attn.o_proj
enable GaLore for weights in module:  model.layers.26.mlp.gate_proj
enable GaLore for weights in module:  model.layers.26.mlp.down_proj
enable GaLore for weights in module:  model.layers.26.mlp.up_proj
enable GaLore for weights in module:  model.layers.27.self_attn.q_proj
enable GaLore for weights in module:  model.layers.27.self_attn.k_proj
enable GaLore for weights in module:  model.layers.27.self_attn.v_proj
enable GaLore for weights in module:  model.layers.27.self_attn.o_proj
enable GaLore for weights in module:  model.layers.27.mlp.gate_proj
enable GaLore for weights in module:  model.layers.27.mlp.down_proj
enable GaLore for weights in module:  model.layers.27.mlp.up_proj
enable GaLore for weights in module:  model.layers.28.self_attn.q_proj
enable GaLore for weights in module:  model.layers.28.self_attn.k_proj
enable GaLore for weights in module:  model.layers.28.self_attn.v_proj
enable GaLore for weights in module:  model.layers.28.self_attn.o_proj
enable GaLore for weights in module:  model.layers.28.mlp.gate_proj
enable GaLore for weights in module:  model.layers.28.mlp.down_proj
enable GaLore for weights in module:  model.layers.28.mlp.up_proj
enable GaLore for weights in module:  model.layers.29.self_attn.q_proj
enable GaLore for weights in module:  model.layers.29.self_attn.k_proj
enable GaLore for weights in module:  model.layers.29.self_attn.v_proj
enable GaLore for weights in module:  model.layers.29.self_attn.o_proj
enable GaLore for weights in module:  model.layers.29.mlp.gate_proj
enable GaLore for weights in module:  model.layers.29.mlp.down_proj
enable GaLore for weights in module:  model.layers.29.mlp.up_proj
enable GaLore for weights in module:  model.layers.30.self_attn.q_proj
enable GaLore for weights in module:  model.layers.30.self_attn.k_proj
enable GaLore for weights in module:  model.layers.30.self_attn.v_proj
enable GaLore for weights in module:  model.layers.30.self_attn.o_proj
enable GaLore for weights in module:  model.layers.30.mlp.gate_proj
enable GaLore for weights in module:  model.layers.30.mlp.down_proj
enable GaLore for weights in module:  model.layers.30.mlp.up_proj
enable GaLore for weights in module:  model.layers.31.self_attn.q_proj
enable GaLore for weights in module:  model.layers.31.self_attn.k_proj
enable GaLore for weights in module:  model.layers.31.self_attn.v_proj
enable GaLore for weights in module:  model.layers.31.self_attn.o_proj
enable GaLore for weights in module:  model.layers.31.mlp.gate_proj
enable GaLore for weights in module:  model.layers.31.mlp.down_proj
enable GaLore for weights in module:  model.layers.31.mlp.up_proj
2024-03-21 18:02:36.933 | INFO     | __main__:main:284 - 
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096, padding_idx=31999)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

2024-03-21 18:02:36.934 | INFO     | __main__:main:285 - Total params: 6738.42M
2024-03-21 18:02:36.934 | INFO     | __main__:main:286 - Trainable params: 6738.42M
2024-03-21 18:02:36.934 | INFO     | __main__:main:287 - Total params with GaLore enabled: 6476.01M
2024-03-21 18:02:36.934 | INFO     | __main__:main:288 - Saving model to checkpoints/llama_7b-2024-03-21-18-00-59 every 10000 update steps
`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...
/media/ubuntu/20TB/anaconda3/lib/python3.11/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.
  warnings.warn(
Update steps:   0%|                 | 171/150000 [3:39:26<2558:58:57, 61.49s/it]

下面是wandb的内容，可以看到显卡基本上是跑满的