关于代码理解和显存占用的问题

[anbyaa]

您好，非常感谢您的工作，在查看您的代码和论文的过程中，我有几个不成熟的问题，能否请您解答一下我的疑惑。

1. 在您的代码目录 LOMO/src/lomo.py fused_backward函数第172行中loss.backward()的操作似乎没有做特别的优化，它是否会进行反向传播，同时计算出网络中所有结点的梯度数值？如果是这样，它似乎与论文中，只迭代性的计算上一层的梯度的想法相互矛盾了，请问这里应该怎样理解呢？
2. 在您的代码目录 LOMO/src/lomo_lora_trainer.py class LOMOLoRATrainer的初始化方法第90行中，您调用到了AdamW优化器，这是否说明您的LOMO方法，只是分层调用已有优化器的更高层次方法，是可以兼容包括SGD和AdamW在内的优化器的优化方法。只是这种方法在和SGD优化器结合使用时效果更好，和Adamw联合使用时也可以工作，是否可以这样理解呢？
3. 在https://github.com/OpenLMLab/LOMO/issues/38中，您提到table2中的memory占用是指各GPU在时间上的平均显存占用，那么请问下各GPU的峰值显存占用相比之前是否有明显优化？是否有相关的具体实验数据呢？ 期待您的回复，谢谢~

[QipengGuo]

1. loss.backward()不需要改动，因为我们通过register_hook的方法改变了backward函数的行为。
2. lora这个实验是lora（用Adam）加上其他参数使用LOMO，所以其实是两部分。LOMO这个方法理论上不限定于SGD，但是对Adam这种需要存储中间变量的方法，需要再做调整，并且也不是论文中所有的优化都能使用。在backward过程直接更新来规避存储gradient还是可行的，但存储中间变量的开销更大，节省一部分的收益有限。如果需要在Adam上进一步应用LOMO的话，一个选择是采用deepspeed里面的optimizer offload机制，把优化器的中间变量交给CPU处理，因为CPUAdam的效率还是可以的。在GPU上再采用LOMO的方法。
3. 峰值也会减小，我们后续会补充相关的实验数据。