adamw优化器占用显存

### AdamW优化器高VRAM占用的原因

AdamW优化器相较于传统的SGD或Adagrad等方法，在实现过程中引入了权重衰减机制，这使得参数更新更加稳定并有助于防止过拟合。然而，这种改进也带来了额外的状态维护需求，具体来说：

- 对于每一个可训练参数，AdamW都需要保存一阶矩估计（mean）和二阶矩估计（uncentered variance），即m_t 和v_t两个缓存向量[^1]。

这些状态变量的存在显著增加了内存开销，尤其是在处理大规模神经网络时。

为了缓解这一问题，可以采取以下几种策略来降低AdamW优化器的显存消耗：

### 减少AdamW优化器显存占用的方法

#### 使用混合精度训练
通过采用半精度浮点数（FP16）代替单精度浮点数（FP32）来进行计算与存储梯度、参数和其他中间结果，可以在几乎不影响模型性能的情况下大幅减少所需GPU显存。PyTorch提供了`torch.cuda.amp.autocast()`上下文管理器用于自动切换到较低精度模式执行操作[^2]。

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

for input, target in data_loader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(input)
        loss = criterion(output, target)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

#### 应用稀疏化技术
对于某些类型的层（如Embedding Layers），可以通过设置特定比例的权值为零从而节省空间。这种方法不仅减少了需要更新的参数数量，还降低了每次迭代所需的临时缓冲区大小。

#### 利用分布式训练框架
当单个设备无法满足资源要求时，考虑利用多台机器或多张GPU卡共同完成任务不失为一种有效手段。像Horovod这样的库可以帮助开发者轻松搭建起高效的分布式环境，进而分摊各节点上的负载压力。

#### 替换为更轻量级版本
如果项目允许的话，可以选择其他更为精简但仍能保持良好效果的选择，比如paged_adamw系列中的adamw_bnb_8bit选项就专为此设计，旨在提供更低的内存足迹的同时维持相近的学习效率。